use std::collections::VecDeque;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex, RwLock};
use futures_util::future::BoxFuture;
use indexmap::IndexMap;
use serde_json::Value;
use tokio_util::sync::CancellationToken;
use yapi_agent::hooks::AgentHooks;
use yapi_agent::{
AgentContext, ExecutionMode, LoopConfig, Tool, ToolCallOutcome, ToolCallScope, UpdateSink,
};
use yapi_ai::api::Apis;
use yapi_ai::registry::ModelRegistry;
use yapi_ai::stream::{Hook, RequestHeaders, RequestHooks, StreamOptions, ThinkingBudgets};
use yapi_types::event::{AgentEvent, SummarySource, ToolResult};
use yapi_types::extension_event::{ExtensionEvent, RegisteredMcpServer, TreePreparation};
use yapi_types::message::{
AssistantMessage, Content, ContentBlock, ImageContent, Message, StopReason, SystemMessage,
ThinkingLevel, ToolCall, UserMessage,
};
use yapi_types::model::Model;
use yapi_types::rpc::{PromptDisposition, StreamingBehavior};
use yapi_types::session::FileEntry;
use yapi_types::settings::Transport;
use yapi_types::sync::{lock, read, write};
use yapi_types::system_prompt::SystemPromptOptions;
use crate::compaction::BranchSummary;
use crate::extensions::{
BashOperations, Extension, ExtensionUi, Loadout, LoadoutChanges, Mode, NoUi, NotifyKind,
SessionActions, Tools,
};
use crate::resources::{ContextFile, PromptTemplate, Skill, expand_prompt_template};
use crate::session::SessionManager;
use crate::settings::SettingsManager;
use crate::system_prompt::{build_sections, diff_sections};
use crate::time::now_ms;
use crate::tools::registry::ToolRegistry;
use crate::tools::{BUILTIN_TOOLS, Described, Exposure, RegisteredTool, Runtime, ToolEnv, builtin};
mod boundary;
mod cache_warmer;
mod extensions;
mod models;
mod recovery;
mod stats;
pub use boundary::Boundary;
pub use cache_warmer::{
CacheWarmingDecision, CacheWarmingPhase, CacheWarmingState, CacheWarmingStatus,
format_cache_warming_status, format_cache_warming_usage,
};
use extensions::Hooks;
pub use models::CatalogRefresh;
pub use stats::{ContextUsage, SessionStats, UsageTotals};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum InputSource {
Interactive,
Rpc,
Extension,
}
impl InputSource {
fn as_str(self) -> &'static str {
match self {
InputSource::Interactive => "interactive",
InputSource::Rpc => "rpc",
InputSource::Extension => "extension",
}
}
}
#[derive(Clone, Debug)]
pub struct PromptOptions {
pub expand_prompt_templates: bool,
pub streaming_behavior: Option<StreamingBehavior>,
pub images: Vec<ImageContent>,
pub source: InputSource,
}
impl Default for PromptOptions {
fn default() -> PromptOptions {
PromptOptions {
expand_prompt_templates: true,
streaming_behavior: None,
images: Vec::new(),
source: InputSource::Interactive,
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Replacement {
New,
Resume,
Fork,
Reload,
}
impl Replacement {
pub fn as_str(self) -> &'static str {
match self {
Replacement::New => "new",
Replacement::Resume => "resume",
Replacement::Fork => "fork",
Replacement::Reload => "reload",
}
}
fn reported(self, file: Option<String>) -> Option<String> {
file.filter(|_| self != Replacement::Reload)
}
}
#[derive(Clone, Debug)]
pub enum SessionChange {
New,
Resume(String),
Fork {
entry_id: String,
at: bool,
},
}
impl SessionChange {
pub fn event(&self) -> &'static str {
match self {
SessionChange::Fork { .. } => "session_before_fork",
_ => "session_before_switch",
}
}
pub fn reason(&self) -> Replacement {
match self {
SessionChange::New => Replacement::New,
SessionChange::Resume(_) => Replacement::Resume,
SessionChange::Fork { .. } => Replacement::Fork,
}
}
}
pub enum UserBash {
Local,
Operations(BashOperations),
Done(crate::bash_executor::BashResult),
}
pub type Listener = Box<dyn Fn(&AgentEvent) + Send + Sync>;
#[derive(Clone, Debug, Default)]
pub struct Resources {
pub context_files: Vec<ContextFile>,
pub skills: Vec<Skill>,
pub skill_diagnostics: Vec<crate::resources::Diagnostic>,
pub skill_sources: Vec<crate::resources::SourceInfo>,
pub templates: Vec<PromptTemplate>,
pub template_diagnostics: Vec<crate::resources::Diagnostic>,
pub template_sources: Vec<crate::resources::SourceInfo>,
pub custom_prompt: Option<String>,
pub append_prompt: Option<String>,
pub themes: Vec<crate::resources::SourceInfo>,
}
pub struct SessionConfig {
pub cwd: PathBuf,
pub agent_dir: PathBuf,
pub settings: SettingsManager,
pub registry: ModelRegistry,
pub apis: Apis,
pub session: SessionManager,
pub model: Option<Model>,
pub thinking_level: ThinkingLevel,
pub tools: Vec<String>,
pub extensions: Vec<Arc<dyn Extension>>,
pub include_extension_tools: bool,
pub allowed_tools: Option<Vec<String>>,
pub excluded_tools: Vec<String>,
pub resources: Resources,
pub docs: crate::docs::Locations,
}
struct State {
session: SessionManager,
model: Option<Model>,
thinking_level: ThinkingLevel,
streaming: bool,
settling: bool,
}
struct Inner {
cwd: PathBuf,
settings: Mutex<SettingsManager>,
registry: RwLock<Arc<ModelRegistry>>,
catalog_refresh: tokio::sync::Mutex<()>,
apis: RwLock<Apis>,
state: Mutex<State>,
scoped_models: Mutex<Vec<crate::model_resolver::ScopedModel>>,
tools: Tools,
extensions: Vec<Arc<dyn Extension>>,
binding: Mutex<(Arc<dyn ExtensionUi>, Mode, Option<SessionActions>)>,
run_options: Mutex<Option<SystemPromptOptions>>,
edited_base: Mutex<Option<(SystemPromptOptions, SystemPromptOptions)>>,
forced_prompt: Mutex<Option<String>>,
hidden_tools: Mutex<std::collections::HashSet<String>>,
next_turn: Mutex<Vec<Message>>,
pending_custom: Mutex<Vec<Message>>,
resources: RwLock<Resources>,
runtime: Arc<RwLock<Runtime>>,
listeners: Mutex<Vec<Listener>>,
steering: Mutex<VecDeque<Message>>,
follow_up: Mutex<VecDeque<Message>>,
queue_modes: Mutex<(Option<Value>, Option<Value>)>,
queued: Mutex<(Vec<String>, Vec<String>)>,
cancel: Mutex<Option<CancellationToken>>,
recovery: Mutex<Recovery>,
bash: Mutex<Vec<(u64, CancellationToken)>>,
pending_bash: Mutex<Vec<Message>>,
agent_dir: PathBuf,
docs: crate::docs::Locations,
idle: tokio::sync::Notify,
deferred: Mutex<Vec<tokio::sync::oneshot::Sender<tokio::sync::oneshot::Sender<()>>>>,
compacting: AtomicUsize,
manual_compaction: std::sync::atomic::AtomicBool,
retry_cancel: Mutex<Option<CancellationToken>>,
nested: crate::nested::NestedCalls,
announcements: Mutex<VecDeque<Value>>,
announcing: tokio::sync::Mutex<()>,
outcome: Mutex<&'static str>,
ui_prompts: Mutex<(usize, Option<UiPrompt>)>,
mcp_servers: Mutex<McpServers>,
unhandled_mcp_servers: Mutex<std::collections::HashSet<String>>,
extension_load_errors: Mutex<Vec<(String, String)>>,
extension_warnings: Mutex<Vec<(String, String)>>,
dropped_results: Mutex<std::collections::HashSet<(String, String)>>,
cache_warmer: Mutex<cache_warmer::CacheWarmer>,
#[cfg(test)]
races: RacePoints,
}
type UiPrompt = (String, Option<String>);
type McpServers = (IndexMap<String, (u64, RegisteredMcpServer)>, bool);
struct Compacting<'a>(&'a AtomicUsize);
impl<'a> Compacting<'a> {
fn start(counter: &'a AtomicUsize) -> Compacting<'a> {
counter.fetch_add(1, Ordering::SeqCst);
Compacting(counter)
}
}
impl Drop for Compacting<'_> {
fn drop(&mut self) {
self.0.fetch_sub(1, Ordering::SeqCst);
}
}
#[cfg(test)]
#[derive(Default)]
struct RacePoints {
before_end: Mutex<Option<Box<dyn FnOnce() + Send>>>,
before_defer: Mutex<Option<Box<dyn FnOnce() + Send>>>,
}
#[cfg(test)]
fn race(point: &Mutex<Option<Box<dyn FnOnce() + Send>>>) {
let hook = lock(point).take();
if let Some(hook) = hook {
hook();
}
}
struct Settling<'a>(&'a Inner);
impl Settling<'_> {
fn finish(self) -> Vec<tokio::sync::oneshot::Sender<tokio::sync::oneshot::Sender<()>>> {
let mut state = lock(&self.0.state);
state.settling = false;
let deferred = std::mem::take(&mut *lock(&self.0.deferred));
drop(state);
deferred
}
}
impl Drop for Settling<'_> {
fn drop(&mut self) {
lock(&self.0.state).settling = false;
}
}
#[derive(Default)]
struct Recovery {
retry_attempt: u32,
overflow_recovery_attempted: bool,
last_assistant: Option<(AssistantMessage, Option<String>)>,
turn_tool_results: Vec<String>,
last_tool_results: Vec<String>,
}
#[derive(Clone, Debug, Default)]
pub struct TreeNavigation {
pub summarize: bool,
pub custom_instructions: Option<String>,
pub replace_instructions: bool,
pub label: Option<String>,
}
#[derive(Clone, Debug, Default, PartialEq)]
pub struct TreeOutcome {
pub editor_text: Option<String>,
pub cancelled: bool,
pub aborted: bool,
pub summary_entry: Option<FileEntry>,
}
#[derive(Clone)]
pub struct AgentSession {
inner: Arc<Inner>,
}
#[derive(Clone, Default)]
pub struct WeakSession(std::sync::Weak<Inner>);
impl WeakSession {
pub fn upgrade(&self) -> Option<AgentSession> {
self.0.upgrade().map(|inner| AgentSession { inner })
}
}
impl AgentSession {
pub fn new(config: SessionConfig) -> AgentSession {
let SessionConfig {
cwd,
agent_dir,
settings,
registry,
apis,
mut session,
model,
thinking_level,
tools,
extensions,
include_extension_tools,
allowed_tools,
excluded_tools,
resources,
docs,
} = config;
let runtime = Arc::new(RwLock::new(Runtime::default()));
let env = ToolEnv {
cwd: cwd.clone(),
runtime: runtime.clone(),
bin_dir: crate::config::bin_dir(&agent_dir),
};
let available: Vec<RegisteredTool> = BUILTIN_TOOLS
.iter()
.filter_map(|name| builtin(name, &env))
.collect();
let mut tool_set = ToolRegistry::new(available, Vec::new());
tool_set.restrict(allowed_tools, excluded_tools);
let tool_registry = Tools::new(tool_set);
for extension in &extensions {
extension.load(&tool_registry);
}
let mut tools = tools;
if include_extension_tools {
tool_registry.with(|registry| {
for tool in registry.all() {
let name = tool.name().to_owned();
let builtin = BUILTIN_TOOLS.contains(&name.as_str());
if !builtin && tool.exposure.declarable() && tool.default_active {
tools.push(name);
}
}
});
}
tool_registry.set_active(tools);
let has_messages = !session.build_context().messages.is_empty();
if has_messages {
let has_thinking = session.branch_path(None).iter().any(|entry| {
matches!(
entry,
yapi_types::session::FileEntry::ThinkingLevelChange(_)
)
});
if !has_thinking {
let _ = session.append_thinking_level_change(thinking_level.as_str());
}
} else {
if let Some(model) = &model {
let _ = session.append_model_change(&model.provider, &model.id);
}
let _ = session.append_thinking_level_change(thinking_level.as_str());
}
let queue_mode = |key: &str| {
Some(
settings
.merged_document()
.get(key)
.filter(|mode| yapi_ai::http::is_truthy(mode))
.cloned()
.unwrap_or_else(|| "one-at-a-time".into()),
)
};
let queue_modes = (queue_mode("steeringMode"), queue_mode("followUpMode"));
AgentSession {
inner: Arc::new(Inner {
cwd,
settings: Mutex::new(settings),
registry: RwLock::new(Arc::new(registry)),
catalog_refresh: tokio::sync::Mutex::new(()),
apis: RwLock::new(apis),
state: Mutex::new(State {
session,
model,
thinking_level,
streaming: false,
settling: false,
}),
scoped_models: Mutex::new(Vec::new()),
tools: tool_registry,
extensions,
binding: Mutex::new((Arc::new(NoUi), Mode::Print, None)),
run_options: Mutex::new(None),
edited_base: Mutex::new(None),
forced_prompt: Mutex::new(None),
hidden_tools: Mutex::default(),
next_turn: Mutex::new(Vec::new()),
pending_custom: Mutex::new(Vec::new()),
resources: RwLock::new(resources),
runtime,
listeners: Mutex::new(Vec::new()),
steering: Mutex::new(VecDeque::new()),
follow_up: Mutex::new(VecDeque::new()),
queue_modes: Mutex::new(queue_modes),
queued: Mutex::new((Vec::new(), Vec::new())),
cancel: Mutex::new(None),
recovery: Mutex::new(Recovery::default()),
bash: Mutex::new(Vec::new()),
pending_bash: Mutex::new(Vec::new()),
docs,
agent_dir,
idle: tokio::sync::Notify::new(),
deferred: Mutex::new(Vec::new()),
#[cfg(test)]
races: RacePoints::default(),
compacting: AtomicUsize::new(0),
manual_compaction: std::sync::atomic::AtomicBool::new(false),
retry_cancel: Mutex::new(None),
announcements: Mutex::new(VecDeque::new()),
announcing: tokio::sync::Mutex::new(()),
outcome: Mutex::new("completed"),
ui_prompts: Mutex::new((0, None)),
mcp_servers: Mutex::default(),
unhandled_mcp_servers: Mutex::default(),
extension_load_errors: Mutex::default(),
extension_warnings: Mutex::default(),
dropped_results: Mutex::default(),
cache_warmer: Mutex::default(),
nested: crate::nested::NestedCalls::default(),
}),
}
}
pub fn docs(&self) -> &crate::docs::Locations {
&self.inner.docs
}
pub fn subscribe(&self, listener: Listener) {
lock(&self.inner.listeners).push(listener);
}
fn emit(&self, event: &AgentEvent) {
for listener in lock(&self.inner.listeners).iter() {
listener(event);
}
}
pub fn with_session<T>(&self, f: impl FnOnce(&mut SessionManager) -> T) -> T {
let (value, error) = {
let session = &mut lock(&self.inner.state).session;
(f(session), session.take_write_error())
};
if let Some(error) = error {
let message = format!("Failed to save session: {error}");
let (ui, _) = self.extension_binding();
if ui.has_ui() {
ui.notify(&message, NotifyKind::Error);
} else {
use std::io::Write as _;
let _ = writeln!(std::io::stderr(), "{message}");
}
}
value
}
pub fn take_session(&self) -> SessionManager {
let mut state = lock(&self.inner.state);
let placeholder = SessionManager::in_memory(state.session.cwd());
std::mem::replace(&mut state.session, placeholder)
}
pub fn model(&self) -> Option<Model> {
lock(&self.inner.state).model.clone()
}
pub fn thinking_level(&self) -> ThinkingLevel {
lock(&self.inner.state).thinking_level
}
pub fn is_streaming(&self) -> bool {
lock(&self.inner.state).streaming
}
pub fn messages(&self) -> Vec<Message> {
self.with_session(|session| session.build_context().messages)
}
pub fn last_assistant(&self) -> Option<yapi_types::message::AssistantMessage> {
self.messages()
.into_iter()
.rev()
.find_map(|message| match message {
Message::Assistant(assistant) => Some(*assistant),
_ => None,
})
}
pub fn set_name(&self, name: &str) {
let name = self.with_session(|session| {
let _ = session.append_session_info(name);
session.name()
});
self.emit(&AgentEvent::SessionInfoChanged { name: name.clone() });
self.announce(ExtensionEvent::SessionInfoChanged {
name: name.as_deref(),
});
}
pub fn settings_manager(&self) -> SettingsManager {
lock(&self.inner.settings).clone()
}
pub fn settings(&self) -> yapi_types::settings::Settings {
lock(&self.inner.settings).settings().clone()
}
fn image_options(&self) -> (bool, Option<yapi_types::models::ImageResize>) {
let auto_resize = lock(&self.inner.settings).settings().image_auto_resize();
let limits = lock(&self.inner.state)
.model
.as_ref()
.and_then(|model| model.image_resize().cloned());
(auto_resize, limits)
}
pub fn http_idle_timeout_ms(&self) -> u64 {
lock(&self.inner.settings)
.http_idle_timeout_ms()
.unwrap_or(yapi_ai::http::DEFAULT_IDLE_TIMEOUT_MS)
}
pub fn settings_errors(&self) -> Vec<String> {
lock(&self.inner.settings).errors().to_vec()
}
pub fn project_trusted(&self) -> bool {
lock(&self.inner.settings).project_trusted()
}
pub fn project_dir(&self) -> PathBuf {
lock(&self.inner.settings).project_dir().to_path_buf()
}
pub fn set_global_setting(
&self,
key: &str,
value: Option<serde_json::Value>,
) -> Result<(), String> {
lock(&self.inner.settings)
.set(crate::settings::Scope::Global, key, value)
.map_err(|err| err.to_string())
}
pub fn resources(&self) -> Resources {
read(&self.inner.resources).clone()
}
pub fn clear_queues(&self) -> (Vec<String>, Vec<String>) {
lock(&self.inner.steering).clear();
lock(&self.inner.follow_up).clear();
let queued = std::mem::take(&mut *lock(&self.inner.queued));
self.emit_queue_update();
queued
}
pub fn pending_message_count(&self) -> usize {
let queued = lock(&self.inner.queued);
queued.0.len() + queued.1.len()
}
pub fn steering_mode(&self) -> Option<Value> {
lock(&self.inner.queue_modes).0.clone()
}
pub fn follow_up_mode(&self) -> Option<Value> {
lock(&self.inner.queue_modes).1.clone()
}
pub(super) fn drains_all(&self) -> (bool, bool) {
let all = |mode: &Option<Value>| mode.as_ref().and_then(Value::as_str) == Some("all");
let modes = lock(&self.inner.queue_modes);
(all(&modes.0), all(&modes.1))
}
pub fn set_steering_mode(&self, mode: Option<Value>) -> Result<(), String> {
lock(&self.inner.queue_modes).0.clone_from(&mode);
self.set_global_setting("steeringMode", mode)
}
pub fn set_follow_up_mode(&self, mode: Option<Value>) -> Result<(), String> {
lock(&self.inner.queue_modes).1.clone_from(&mode);
self.set_global_setting("followUpMode", mode)
}
pub fn auto_compaction_enabled(&self) -> Value {
lock(&self.inner.settings)
.merged_document()
.get("compaction")
.and_then(|compaction| compaction.get("enabled"))
.filter(|enabled| !enabled.is_null())
.cloned()
.unwrap_or(Value::Bool(true))
}
pub fn set_auto_compaction(&self, enabled: Option<Value>) -> Result<(), String> {
lock(&self.inner.settings)
.set_nested(
crate::settings::Scope::Global,
"compaction",
"enabled",
enabled,
)
.map_err(|err| err.to_string())
}
pub fn set_auto_retry(&self, enabled: Option<Value>) -> Result<(), String> {
lock(&self.inner.settings)
.set_nested(crate::settings::Scope::Global, "retry", "enabled", enabled)
.map_err(|err| err.to_string())
}
pub fn set_nested_global_setting(
&self,
field: &str,
key: &str,
value: serde_json::Value,
) -> Result<(), String> {
lock(&self.inner.settings)
.set_nested(crate::settings::Scope::Global, field, key, Some(value))
.map_err(|err| err.to_string())
}
pub fn is_compacting(&self) -> bool {
self.inner.compacting.load(Ordering::SeqCst) > 0
}
pub async fn wait_for_idle(&self) {
loop {
let idle = self.inner.idle.notified();
if !self.is_streaming() && !self.is_compacting() {
return;
}
idle.await;
}
}
pub async fn execute_bash(
&self,
command: &str,
exclude_from_context: Option<bool>,
id: Option<String>,
operations: Option<BashOperations>,
mut on_chunk: impl FnMut(&str),
) -> Result<crate::bash_executor::BashResult, String> {
static NEXT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
let token = NEXT.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let cancel = CancellationToken::new();
lock(&self.inner.bash).push((token, cancel.clone()));
let settings = self.settings();
let resolved =
crate::tools::bash::with_prefix(settings.shell_command_prefix.as_deref(), command);
let cwd = self.with_session(|session| session.cwd().to_path_buf());
let on_chunk = |delta: &str| {
on_chunk(delta);
self.emit(&AgentEvent::BashExecutionUpdate {
id: id.clone(),
delta: delta.to_owned(),
});
};
let result = match &operations {
None => {
crate::bash_executor::execute(
&resolved,
&cwd,
settings.shell_path.as_deref(),
&crate::config::bin_dir(&self.inner.agent_dir),
cancel.clone(),
on_chunk,
)
.await
}
Some(operations) => {
let run = |output| {
operations.extension.run_bash(
&operations.handle,
&resolved,
&cwd,
output,
cancel.clone(),
)
};
crate::bash_executor::execute_with(run, &cancel, on_chunk).await
}
};
lock(&self.inner.bash).retain(|(running, _)| *running != token);
let result = result?;
self.record_bash(command, &result, exclude_from_context);
Ok(result)
}
pub async fn user_bash(
&self,
command: &str,
exclude_from_context: bool,
) -> Result<UserBash, String> {
let handlers = self.handlers_of("user_bash");
if handlers.is_empty() {
return Ok(UserBash::Local);
}
let cwd = self.with_session(|session| session.cwd().to_path_buf());
let event = ExtensionEvent::UserBash {
command,
exclude_from_context,
cwd: &cwd,
}
.to_value();
let ctx = self.extension_context(CancellationToken::new());
for extension in handlers {
let Some(result) = extension.handle(&ctx, &event).await else {
continue;
};
if let Some(error) = result["error"].as_str() {
return Err(error.to_owned());
}
let replaces = !result["operations"].is_null() || !result["result"].is_null();
if replaces
&& !self.applies(
extension,
"user_bash",
extension.grants().is_all(),
"every grant",
)
{
continue;
}
if !result["operations"].is_null() {
return Ok(UserBash::Operations(BashOperations {
extension: Arc::clone(extension),
handle: result["operations"].clone(),
}));
}
if let Ok(done) = serde_json::from_value(result["result"].clone()) {
return Ok(UserBash::Done(done));
}
let error = "Invalid user_bash handler result: return undefined for local execution or exactly one valid { operations } or { result } object";
ctx.ui
.extension_error(&extension.source().path, "user_bash", error, None);
return Err(error.to_owned());
}
Ok(UserBash::Local)
}
pub fn record_bash(
&self,
command: &str,
result: &crate::bash_executor::BashResult,
exclude_from_context: Option<bool>,
) {
let message = Message::BashExecution(yapi_types::message::BashExecutionMessage {
command: command.to_owned(),
output: result.output.clone(),
exit_code: result.exit_code.map(i64::from),
cancelled: result.cancelled,
truncated: result.truncated,
full_output_path: result.full_output_path.clone(),
timestamp: now_ms(),
exclude_from_context,
});
if self.is_streaming() {
lock(&self.inner.pending_bash).push(message);
} else {
let _ = self.with_session(|session| session.append_message(message));
}
}
fn flush_pending_bash(&self) {
let pending: Vec<Message> = lock(&self.inner.pending_bash).drain(..).collect();
for message in pending {
let _ = self.with_session(|session| session.append_message(message));
}
}
pub fn abort_bash(&self) {
for (_, token) in lock(&self.inner.bash).iter() {
token.cancel();
}
}
pub fn is_bash_running(&self) -> bool {
!lock(&self.inner.bash).is_empty()
}
pub fn user_messages_for_forking(&self) -> Vec<(String, String)> {
self.with_session(|session| {
session
.entries()
.filter_map(|entry| match entry {
FileEntry::Message(entry) => match &entry.message {
Message::User(user) => {
let text = user.content.text("");
(!text.is_empty()).then(|| (entry.meta.id.clone(), text))
}
_ => None,
},
_ => None,
})
.collect()
})
}
pub fn last_assistant_text(&self) -> Option<String> {
let assistant = self
.messages()
.into_iter()
.rev()
.find_map(|message| match message {
Message::Assistant(assistant)
if !(assistant.stop_reason == StopReason::Aborted
&& assistant.content.is_empty()) =>
{
Some(assistant)
}
_ => None,
})?;
let text = assistant_text(&assistant);
let text = text.trim();
(!text.is_empty()).then(|| text.to_owned())
}
pub fn run_cancel(&self) -> Option<CancellationToken> {
lock(&self.inner.cancel).clone()
}
pub fn abort(&self) {
if let Some(cancel) = lock(&self.inner.cancel).as_ref() {
cancel.cancel();
}
}
fn user_message(text: String, images: Vec<ImageContent>) -> Message {
let mut content = vec![ContentBlock::text(text)];
content.extend(images.into_iter().map(ContentBlock::Image));
Message::User(UserMessage {
content: Content::Blocks(content),
timestamp: now_ms(),
})
}
pub fn deliver_custom_message(
&self,
message: yapi_types::message::CustomMessage,
trigger_turn: Option<bool>,
deliver_as: Option<&str>,
) -> Option<impl Future<Output = ()> + Send + 'static> {
let message = Message::Custom(message);
if deliver_as == Some("nextTurn") {
lock(&self.inner.next_turn).push(message);
return None;
}
let state = lock(&self.inner.state);
if state.streaming && trigger_turn != Some(false) {
let queue = if deliver_as == Some("followUp") {
&self.inner.follow_up
} else {
&self.inner.steering
};
lock(queue).push_back(message);
} else if trigger_turn == Some(true) {
drop(state);
let session = self.clone();
return Some(async move {
let _ = session.run(vec![message]).await;
});
} else if state.streaming {
lock(&self.inner.pending_custom).push(message);
} else {
drop(state);
self.append_custom_message(&message);
}
None
}
pub fn settled_turn(
&self,
) -> Option<impl Future<Output = Option<tokio::sync::oneshot::Sender<()>>> + Send + 'static>
{
#[cfg(test)]
race(&self.inner.races.before_defer);
let state = lock(&self.inner.state);
if !state.settling {
return None;
}
let (turn, start) = tokio::sync::oneshot::channel();
lock(&self.inner.deferred).push(turn);
drop(state);
Some(async move { start.await.ok() })
}
fn append_custom_message(&self, message: &Message) {
let Message::Custom(custom) = message else {
return;
};
self.with_session(|file| {
let _ = file.append_custom_message(
&custom.custom_type,
custom.content.clone(),
custom.display,
custom.details.clone(),
);
});
self.emit(&AgentEvent::MessageStart {
message: message.clone(),
});
self.emit(&AgentEvent::MessageEnd {
message: message.clone(),
});
}
fn flush_pending_custom(&self) {
let pending = std::mem::take(&mut *lock(&self.inner.pending_custom));
for message in &pending {
self.append_custom_message(message);
}
}
fn emit_queue_update(&self) {
let (steering, follow_up) = lock(&self.inner.queued).clone();
self.emit(&AgentEvent::QueueUpdate {
steering,
follow_up,
});
}
fn dequeue_started(&self, message: &Message) {
let Message::User(user) = message else {
return;
};
let text = user.content.text("");
if text.is_empty() {
return;
}
let removed = {
let mut queued = lock(&self.inner.queued);
let (steering, follow_up) = &mut *queued;
[steering, follow_up].into_iter().any(|texts| {
let index = texts.iter().position(|queued| *queued == text);
index.map(|index| texts.remove(index)).is_some()
})
};
if removed {
self.emit_queue_update();
}
}
pub async fn queue_input(
&self,
text: &str,
images: Vec<ImageContent>,
behavior: StreamingBehavior,
source: InputSource,
) -> Result<PromptDisposition, String> {
if self.is_extension_command(text) {
let name = text[1..].split(' ').next().unwrap_or_default();
return Err(format!(
"Extension command \"/{name}\" cannot be queued. Use prompt() or execute the command when not streaming."
));
}
let streaming = self.is_streaming().then_some(behavior);
let Some((text, images)) = self.input_handlers(text, images, source, streaming).await
else {
return Ok(PromptDisposition::Handled);
};
self.queue(behavior, self.expand(&text), images);
Ok(PromptDisposition::Queued)
}
pub fn steer(&self, text: &str, images: Vec<ImageContent>) {
self.queue(StreamingBehavior::Steer, self.expand(text), images);
}
pub fn follow_up(&self, text: &str, images: Vec<ImageContent>) {
self.queue(StreamingBehavior::FollowUp, self.expand(text), images);
}
fn queue(&self, behavior: StreamingBehavior, text: String, images: Vec<ImageContent>) {
self.push_queued(behavior, text, images);
self.emit_queue_update();
}
fn queue_streaming(
&self,
behavior: StreamingBehavior,
text: String,
images: Vec<ImageContent>,
) -> Result<(), (String, Vec<ImageContent>)> {
let state = lock(&self.inner.state);
if !state.streaming {
return Err((text, images));
}
self.push_queued(behavior, text, images);
drop(state);
self.emit_queue_update();
Ok(())
}
fn push_queued(&self, behavior: StreamingBehavior, text: String, images: Vec<ImageContent>) {
let mut queued = lock(&self.inner.queued);
let (texts, queue) = match behavior {
StreamingBehavior::Steer => (&mut queued.0, &self.inner.steering),
StreamingBehavior::FollowUp => (&mut queued.1, &self.inner.follow_up),
};
texts.push(text.clone());
lock(queue).push_back(Self::user_message(text, images));
}
fn expand(&self, text: &str) -> String {
let text = self.expand_skill(text);
expand_prompt_template(&text, &read(&self.inner.resources).templates)
}
fn expand_skill(&self, text: &str) -> String {
let Some(rest) = text.strip_prefix("/skill:") else {
return text.to_owned();
};
let (name, args) = match rest.split_once(' ') {
Some((name, args)) => (name, args.trim()),
None => (rest, ""),
};
let Some(skill) = read(&self.inner.resources)
.skills
.iter()
.find(|skill| skill.name == name)
.cloned()
else {
return text.to_owned();
};
let content = match std::fs::read_to_string(&skill.file_path) {
Ok(content) => content,
Err(error) => {
let error = crate::tools::node_error(&error, "open", &skill.file_path);
let ui = lock(&self.inner.binding).0.clone();
ui.extension_error(
&skill.file_path.to_string_lossy(),
"skill_expansion",
&error,
None,
);
return text.to_owned();
}
};
let (_, body) = crate::resources::split_frontmatter(&content);
let mut expanded = format!(
"<skill name=\"{}\" location=\"{}\">\nReferences are relative to {}.\n\n{}\n</skill>",
skill.name,
skill.file_path.display(),
skill.base_dir.display(),
body.trim()
);
if !args.is_empty() {
expanded += &format!("\n\n{args}");
}
expanded
}
pub fn base_prompt_options(&self) -> SystemPromptOptions {
let base = self.rebuilt_prompt_options();
match &*lock(&self.inner.edited_base) {
Some((edited_from, edited)) if *edited_from == base => edited.clone(),
_ => base,
}
}
pub fn set_base_prompt_options(&self, options: SystemPromptOptions) {
*lock(&self.inner.edited_base) = Some((self.rebuilt_prompt_options(), options));
}
fn rebuilt_prompt_options(&self) -> SystemPromptOptions {
let resources = self.resources();
let selected_tools = self.inner.tools.active();
let (_, changes) = self.loadout_changes(self.active_tools(&selected_tools));
let mut options = SystemPromptOptions {
custom_prompt: resources.custom_prompt,
selected_tools,
append_system_prompt: resources.append_prompt.unwrap_or_default(),
cwd: self.inner.cwd.clone(),
context_files: resources.context_files,
skills: resources.skills,
..SystemPromptOptions::default()
};
for tool in self.inner.tools.with(|registry| registry.all()) {
let name = tool.name().to_owned();
if let Some(snippet) = tool
.snippet
.filter(|snippet| !snippet.is_empty() && !changes.hidden.contains(&name))
{
options.tool_snippets.insert(name.clone(), snippet);
}
if !tool.guidelines.is_empty() {
options.tool_guidelines.insert(name, tool.guidelines);
}
}
options
}
fn prompt_options(&self, active: &[String]) -> SystemPromptOptions {
let mut options = lock(&self.inner.run_options)
.clone()
.unwrap_or_else(|| self.base_prompt_options());
options.selected_tools = active.to_vec();
let declared = self.active_tools(active);
for tool in &declared {
let name = tool.tool.declaration().name.clone();
if let Some(snippet) = &tool.snippet {
options
.tool_snippets
.entry(name.clone())
.or_insert_with(|| snippet.clone());
}
if !tool.guidelines.is_empty() {
options
.tool_guidelines
.entry(name)
.or_insert_with(|| tool.guidelines.clone());
}
}
let (_, changes) = self.loadout_changes(declared);
options
.tool_snippets
.retain(|name, _| !changes.hidden.contains(name));
options
}
pub(super) async fn refresh_loadouts(&self, declared: Vec<RegisteredTool>) {
let loadout = Loadout {
declared,
callable: self.callable_tools(),
};
for extension in &self.inner.extensions {
extension.refresh_loadout(&loadout).await;
}
}
fn loadout_changes(&self, declared: Vec<RegisteredTool>) -> (Loadout, LoadoutChanges) {
let loadout = Loadout {
declared,
callable: self.callable_tools(),
};
let mut changes = LoadoutChanges::default();
for extension in &self.inner.extensions {
let next = extension.prepare_loadout(&loadout);
let changed = !next.descriptions.is_empty() || !next.hidden.is_empty();
if changed
&& !self.applies(
extension,
"prepareLoadout",
extension.grants().session,
extensions::SESSION,
)
{
continue;
}
changes.descriptions.extend(next.descriptions);
changes.hidden.extend(next.hidden);
}
(loadout, changes)
}
fn loadout(&self, declared: Vec<RegisteredTool>) -> Vec<Arc<dyn Tool>> {
let (loadout, changes) = self.loadout_changes(declared);
let mut descriptions = changes.descriptions;
*lock(&self.inner.hidden_tools) = changes.hidden;
loadout
.declared
.into_iter()
.map(|tool| match descriptions.remove(tool.name()) {
Some(description) if description != tool.tool.declaration().description => {
Arc::new(Described::new(tool.tool, description)) as Arc<dyn Tool>
}
_ => tool.tool,
})
.collect()
}
fn active_tools(&self, active: &[String]) -> Vec<RegisteredTool> {
self.inner.tools.with(|registry| {
active
.iter()
.filter_map(|name| registry.get(name).cloned())
.collect()
})
}
fn system_update(
&self,
messages: &[Message],
active: &[String],
) -> Result<Option<Message>, String> {
let sections = build_sections(&self.prompt_options(active), &self.inner.docs)?;
let current = yapi_ai::transcript::current_system_message(messages)
.and_then(|system| system.sections)
.unwrap_or_default();
Ok(diff_sections(¤t, §ions).map(|patch| {
Message::System(SystemMessage {
content: Content::Text(String::new()),
sections: Some(patch),
timestamp: now_ms(),
tools_added: None,
tools_removed: None,
})
}))
}
fn system_prompt_text(&self, active: &[String]) -> Result<String, String> {
self.prompt_text(&self.prompt_options(active))
}
pub fn system_prompt(&self) -> Result<String, String> {
if !self.is_streaming() {
return self.prompt_text(&self.base_prompt_options());
}
if let Some(forced) = lock(&self.inner.forced_prompt).clone() {
return Ok(forced);
}
self.system_prompt_text(&self.inner.tools.active())
}
pub fn render_prompt(&self, options: &SystemPromptOptions) -> Result<String, String> {
match &options.force_system_prompt {
Some(forced) => Ok(forced.clone()),
None => self.prompt_text(options),
}
}
fn prompt_text(&self, options: &SystemPromptOptions) -> Result<String, String> {
let sections = build_sections(options, &self.inner.docs)?;
Ok(SystemMessage {
content: Content::Text(String::new()),
sections: Some(
sections
.into_iter()
.map(|(name, text)| (name, Some(text)))
.collect(),
),
timestamp: 0,
tools_added: None,
tools_removed: None,
}
.text())
}
pub fn queue_now(&self, text: &str, options: &PromptOptions) -> bool {
let Some(behavior) = options.streaming_behavior.filter(|_| self.is_streaming()) else {
return false;
};
let expand = options.expand_prompt_templates;
if self.has_handlers("input")
|| (expand && self.is_extension_command(text))
|| self.inner.manual_compaction.load(Ordering::SeqCst)
{
return false;
}
let text = if expand {
self.expand(text)
} else {
text.to_owned()
};
self.queue_streaming(behavior, text, options.images.clone())
.is_ok()
}
pub async fn prompt(&self, text: &str, images: Vec<ImageContent>) -> Result<(), String> {
let options = PromptOptions {
images,
..PromptOptions::default()
};
self.prompt_with(text, options, |_| {}).await
}
pub async fn prompt_with(
&self,
text: &str,
options: PromptOptions,
preflight: impl FnOnce(PromptDisposition),
) -> Result<(), String> {
let _turn = match self.settled_turn() {
Some(turn) => turn.await,
None => None,
};
let PromptOptions {
expand_prompt_templates: expand,
streaming_behavior: behavior,
images,
source,
} = options;
if expand && self.run_extension_command(text).await {
preflight(PromptDisposition::Handled);
return Ok(());
}
if self.inner.manual_compaction.load(Ordering::SeqCst) {
return Err("Cannot submit a prompt while compaction is in progress. Wait for compaction to finish and retry.".into());
}
let streaming = behavior.filter(|_| self.is_streaming());
let Some((text, images)) = self.input_handlers(text, images, source, streaming).await
else {
preflight(PromptDisposition::Handled);
return Ok(());
};
let expanded = if expand { self.expand(&text) } else { text };
let (expanded, images) = match behavior {
Some(behavior) => match self.queue_streaming(behavior, expanded, images) {
Ok(()) => {
preflight(PromptDisposition::Queued);
return Ok(());
}
Err(message) => message,
},
None if self.is_streaming() => {
return Err("Agent is already processing. Specify streamingBehavior ('steer' or 'followUp') to queue the message.".into());
}
None => (expanded, images),
};
self.flush_pending_bash();
self.flush_pending_custom();
let model = lock(&self.inner.state)
.model
.clone()
.ok_or_else(|| crate::auth_guidance::no_api_key_found("unknown", &self.inner.docs))?;
let registry = self.registry();
if let Some(error) = registry.store_error(&model.provider) {
return Err(error);
}
let mut has_auth = registry.has_auth(&model.provider);
if has_auth && registry.uses_command_key(&model.provider) {
let auth = registry.auth(&model).await;
if let Some(error) = auth.error {
return Err(error);
}
has_auth = auth.source.is_some();
}
if !has_auth {
return Err(crate::auth_guidance::no_api_key_found(
&model.provider,
&self.inner.docs,
));
}
self.compact_before_prompt().await?;
let mut sections = self.base_prompt_options().sections;
let ctx = self.extension_context(CancellationToken::new());
for extension in &self.inner.extensions {
extension.before_agent_start(&ctx, &mut sections).await;
}
*lock(&self.inner.run_options) = Some(SystemPromptOptions {
sections,
..self.base_prompt_options()
});
let (custom, forced) = if self.has_handlers("before_agent_start") {
let options = self.prompt_options(&self.inner.tools.active());
self.before_agent_start_handlers(&expanded, &images, options)
.await
} else {
let forced = lock(&self.inner.run_options)
.as_mut()
.and_then(|options| options.force_system_prompt.take());
(Vec::new(), forced)
};
*lock(&self.inner.forced_prompt) = forced;
let active = self.inner.tools.active();
let (images, hints) = if images.is_empty() {
(images, Vec::new())
} else {
let (auto_resize, limits) = self.image_options();
tokio::task::spawn_blocking(move || {
crate::images::normalize_prompt(images, auto_resize, limits.as_ref())
})
.await
.map_err(|err| err.to_string())?
};
let expanded = if hints.is_empty() {
expanded
} else {
format!("{expanded}\n\n{}", hints.join("\n"))
};
let messages = self.messages();
self.refresh_loadouts(self.active_tools(&active)).await;
let mut prompts = Vec::new();
if let Some(update) = self.system_update(&messages, &active)? {
prompts.push(update);
}
prompts.push(Self::user_message(expanded, images));
prompts.append(&mut lock(&self.inner.next_turn));
prompts.extend(custom);
preflight(PromptDisposition::Started);
self.run(prompts).await
}
async fn run(&self, prompts: Vec<Message>) -> Result<(), String> {
let cancel = CancellationToken::new();
*lock(&self.inner.cancel) = Some(cancel.clone());
lock(&self.inner.state).streaming = true;
self.run_agent(Some(prompts), &cancel).await;
let mut result = Ok(());
loop {
self.settle_run(&cancel, &mut result).await;
if cancel.is_cancelled() {
self.finish_cancelled_retry();
}
#[cfg(test)]
race(&self.inner.races.before_end);
if self.end_run(cancel.is_cancelled() || result.is_err()) {
break;
}
self.run_agent(None, &cancel).await;
}
let settling = Settling(&self.inner);
self.flush_pending_bash();
self.flush_pending_custom();
*lock(&self.inner.cancel) = None;
self.cache_warming_settled();
self.emit_extension_event(
&ExtensionEvent::AgentSettled.to_value(),
CancellationToken::new(),
)
.await;
self.emit(&AgentEvent::AgentSettled);
for turn in settling.finish() {
let (guard, ended) = tokio::sync::oneshot::channel();
if turn.send(guard).is_ok() {
let _ = ended.await;
}
}
self.inner.idle.notify_waiters();
result
}
async fn settle_run(&self, cancel: &CancellationToken, result: &mut Result<(), String>) {
while !cancel.is_cancelled() {
let again = match self.handle_post_run(cancel).await {
Ok(again) => again,
Err(error) => {
*result = Err(error);
break;
}
};
if again {
if cancel.is_cancelled() {
break;
}
self.run_agent(None, cancel).await;
continue;
}
if cancel.is_cancelled() || !self.before_settle(cancel).await {
break;
}
self.run_agent(None, cancel).await;
}
}
fn end_run(&self, stopped: bool) -> bool {
let mut state = lock(&self.inner.state);
if !stopped && self.has_queued() {
return false;
}
state.streaming = false;
state.settling = true;
true
}
fn has_queued(&self) -> bool {
!lock(&self.inner.steering).is_empty() || !lock(&self.inner.follow_up).is_empty()
}
pub(super) fn request_options(&self) -> Result<StreamOptions, String> {
let manager = lock(&self.inner.settings);
let idle = manager.http_idle_timeout_ms()?;
let provider = manager
.settings()
.retry
.as_ref()
.and_then(|retry| retry.provider.clone())
.unwrap_or_default();
let idle = if idle == 0 { 2_147_483_647 } else { idle };
Ok(StreamOptions {
timeout_ms: Some(manager.provider_timeout_ms().unwrap_or(idle as f64)),
websocket_connect_timeout_ms: manager.websocket_connect_timeout_ms()?,
max_retries: provider.max_retries.unwrap_or(0),
max_retry_delay_ms: provider.max_retry_delay_ms,
..StreamOptions::default()
})
}
fn stream_options(
&self,
model: &Model,
session_id: String,
cancel: &CancellationToken,
request: StreamOptions,
) -> StreamOptions {
let settings = lock(&self.inner.settings).settings().clone();
let budgets = settings.thinking_budgets.as_ref();
StreamOptions {
session_id: Some(session_id),
transport: Some(settings.transport.unwrap_or(Transport::Auto)),
thinking_budgets: ThinkingBudgets {
minimal: budgets.and_then(|b| b.minimal),
low: budgets.and_then(|b| b.low),
medium: budgets.and_then(|b| b.medium),
high: budgets.and_then(|b| b.high),
},
cancel: cancel.clone(),
hooks: self.request_hooks(model, cancel),
..request
}
}
fn request_hooks(&self, model: &Model, cancel: &CancellationToken) -> RequestHooks {
let model = model.clone();
RequestHooks {
payload: self.hook(
"before_provider_request",
cancel,
|session, payload, cancel| {
Box::pin(async move {
session
.chain("before_provider_request", "payload", payload, cancel)
.await
})
},
),
headers: self.headers_hook(cancel),
response: self.hook(
"after_provider_response",
cancel,
|session, response: yapi_ai::stream::ProviderResponse, cancel| {
let event = ExtensionEvent::AfterProviderResponse {
status: response.status,
headers: &response.headers,
}
.to_value();
Box::pin(async move {
session.emit_extension_event(&event, cancel).await;
})
},
),
stream_event: self.hook(
"provider_stream_event",
cancel,
move |session, data, cancel| {
let event = serde_json::json!({
"data": data,
"type": "provider_stream_event",
"provider": model.provider,
"api": model.api,
"model": model.id,
});
Box::pin(async move {
session.emit_extension_event(&event, cancel).await;
})
},
),
}
}
pub(super) fn headers_hook(
&self,
cancel: &CancellationToken,
) -> Option<Hook<RequestHeaders, RequestHeaders>> {
self.hook(
"before_provider_headers",
cancel,
|session, headers, cancel| {
Box::pin(async move { session.before_provider_headers(headers, cancel).await })
},
)
}
fn hook<A: 'static, R: 'static>(
&self,
kind: &str,
cancel: &CancellationToken,
event: impl Fn(AgentSession, A, CancellationToken) -> BoxFuture<'static, R>
+ Send
+ Sync
+ 'static,
) -> Option<Hook<A, R>> {
if !self.has_handlers(kind) {
return None;
}
let (session, cancel) = (self.clone(), cancel.clone());
Some(Arc::new(move |argument| {
event(session.clone(), argument, cancel.clone())
}))
}
async fn run_agent(&self, prompts: Option<Vec<Message>>, cancel: &CancellationToken) {
let active = self.inner.tools.active();
let (model, thinking_level, messages, session_id, session_file) = {
let state = lock(&self.inner.state);
(
state.model.clone(),
state.thinking_level,
state.session.build_context().messages,
state.session.id().to_owned(),
state.session.file().map(Path::to_path_buf),
)
};
let Some(model) = model else {
return;
};
let settings = lock(&self.inner.settings).settings().clone();
*write(&self.inner.runtime) = Runtime {
model: Some(model.clone()),
thinking_level: Some(thinking_level),
session_id: Some(session_id.clone()),
session_file,
auto_resize_images: settings.image_auto_resize(),
shell_path: settings.shell_path.clone(),
shell_command_prefix: settings.shell_command_prefix.clone(),
};
let apis = self.apis();
let session = self.downgrade();
let request = self.request_options();
let invalid = request.as_ref().err().cloned();
let config = LoopConfig {
options: self.stream_options(&model, session_id, cancel, request.unwrap_or_default()),
model,
thinking_level,
stream: Arc::new(move |request| {
if let Some(error) = &invalid {
return Err(error.clone());
}
if let Some(session) = session.upgrade() {
session.start_cache_warming(&request);
}
Ok(apis.stream(request))
}),
tool_execution: ExecutionMode::Parallel,
};
let mut context = AgentContext {
messages,
tools: self.loadout(self.active_tools(&active)),
};
let cancel_token = cancel.clone();
let hooks = Hooks {
session: self.clone(),
cancel: cancel_token,
turn_index: std::sync::atomic::AtomicU64::new(0),
};
let prompts = match prompts {
Some(prompts) => Some(prompts),
None if matches!(context.messages.last(), Some(Message::Assistant(_))) => {
let (steering_all, follow_up_all) = self.drains_all();
let steering = drain(&self.inner.steering, steering_all);
let queued = if steering.is_empty() {
drain(&self.inner.follow_up, follow_up_all)
} else {
steering
};
if queued.is_empty() {
return;
}
Some(queued)
}
None => None,
};
match prompts {
Some(prompts) => {
yapi_agent::run(prompts, &mut context, config, &hooks).await;
}
None => {
let _ = yapi_agent::run_continue(&mut context, config, &hooks).await;
}
}
}
pub async fn navigate_tree(
&self,
target_id: &str,
options: TreeNavigation,
) -> Result<TreeOutcome, String> {
if self.is_streaming() {
return Err(
"Wait for the current response to finish before navigating the session tree."
.into(),
);
}
let (old_leaf, target) = self.with_session(|session| {
(
session.leaf_id().map(str::to_owned),
session.entry(target_id).cloned(),
)
});
if old_leaf.as_deref() == Some(target_id) {
return Ok(TreeOutcome::default());
}
let model = self.model();
if options.summarize && model.is_none() {
return Err("No model available for summarization".into());
}
let Some(target) = target else {
return Err(format!("Entry {target_id} not found"));
};
let (entries, common) = self.with_session(|session| {
crate::compaction::collect_branch_entries(session, old_leaf.as_deref(), target_id)
});
let mut options = options;
let mut summary = None;
let mut from_extension = false;
if self.has_handlers("session_before_tree") {
let event = ExtensionEvent::SessionBeforeTree {
preparation: TreePreparation {
target_id,
old_leaf_id: old_leaf.as_deref(),
common_ancestor_id: common.as_deref(),
entries_to_summarize: &entries,
user_wants_summary: options.summarize,
custom_instructions: options.custom_instructions.as_deref(),
replace_instructions: options.replace_instructions,
label: options.label.as_deref(),
},
};
if let Some(result) = self
.emit_extension_event(&event.to_value(), CancellationToken::new())
.await
{
if result["cancel"] == true {
return Ok(TreeOutcome {
cancelled: true,
..TreeOutcome::default()
});
}
if options.summarize
&& let Some(text) = result["summary"]["summary"]
.as_str()
.filter(|text| !text.is_empty())
{
let details = Some(result["summary"]["details"].clone())
.filter(|details| !details.is_null());
let usage = serde_json::from_value(result["summary"]["usage"].clone()).ok();
summary = Some((text.to_owned(), details, usage));
from_extension = true;
}
if let Some(instructions) = result["customInstructions"].as_str() {
options.custom_instructions = Some(instructions.to_owned());
}
if let Some(replace) = result["replaceInstructions"].as_bool() {
options.replace_instructions = replace;
}
if let Some(label) = result["label"].as_str() {
options.label = Some(label.to_owned());
}
}
}
if options.summarize
&& !entries.is_empty()
&& summary.is_none()
&& let Some(model) = &model
{
let cancel = CancellationToken::new();
*lock(&self.inner.cancel) = Some(cancel.clone());
let compacting = Compacting::start(&self.inner.compacting);
let summarizer = self
.summarizer(
model,
ThinkingLevel::Off,
SummarySource::BranchSummary,
None,
cancel,
)
.await;
let reserve = lock(&self.inner.settings)
.settings()
.branch_summary
.as_ref()
.and_then(|settings| settings.reserve_tokens)
.unwrap_or(16384);
let result = summarizer
.branch_summary(
&entries,
options.custom_instructions.as_deref(),
options.replace_instructions,
reserve,
)
.await;
*lock(&self.inner.cancel) = None;
drop(compacting);
match result? {
BranchSummary::Aborted => {
return Ok(TreeOutcome {
cancelled: true,
aborted: true,
..TreeOutcome::default()
});
}
BranchSummary::Done {
summary: text,
usage,
read_files,
modified_files,
} => {
summary = Some((
text,
Some(
serde_json::json!({"readFiles": read_files, "modifiedFiles": modified_files}),
),
usage,
));
}
}
}
let (new_leaf, editor_text) = match &target {
FileEntry::Message(yapi_types::session::MessageEntry {
meta,
message: Message::User(user),
}) => (meta.parent_id.clone(), Some(user.content.text(""))),
FileEntry::CustomMessage(entry) => {
(entry.meta.parent_id.clone(), Some(entry.content.text("")))
}
_ => (Some(target_id.to_owned()), None),
};
let summary_entry = self.with_session(|session| -> Result<Option<FileEntry>, String> {
let mut summary_entry = None;
match summary {
Some((text, details, usage)) => {
let id = session
.branch_with_summary(
new_leaf.as_deref(),
text,
details,
Some(from_extension),
usage,
)
.map_err(|err| err.to_string())?;
summary_entry = session.entry(&id).cloned();
if let Some(label) = &options.label {
session
.append_label(&id, Some(label.clone()))
.map_err(|err| err.to_string())?;
}
}
None => {
match &new_leaf {
None => session.reset_leaf(),
Some(id) => session.branch(id).map_err(|err| err.to_string())?,
}
if let Some(label) = &options.label {
session
.append_label(target_id, Some(label.clone()))
.map_err(|err| err.to_string())?;
}
}
}
Ok(summary_entry)
})?;
self.restore_tools_from_transcript();
let new_leaf = self.with_session(|session| session.leaf_id().map(str::to_owned));
let event = ExtensionEvent::SessionTree {
new_leaf_id: new_leaf.as_deref(),
old_leaf_id: old_leaf.as_deref(),
summary_entry: summary_entry.as_ref(),
from_extension: summary_entry.is_some().then_some(from_extension),
};
self.emit_extension_event(&event.to_value(), CancellationToken::new())
.await;
Ok(TreeOutcome {
editor_text,
cancelled: false,
aborted: false,
summary_entry,
})
}
fn restore_tools_from_transcript(&self) {
let messages = self.messages();
let Some(system) = yapi_ai::transcript::current_system_message(&messages) else {
return;
};
let names: Vec<String> = system
.tools_added
.iter()
.flatten()
.map(|tool| tool.name.clone())
.collect();
self.inner.tools.with(|registry| registry.restore(names));
}
pub fn header_json(&self) -> Option<String> {
self.with_session(|session| {
session.header().and_then(|header| {
yapi_types::json::to_string(&yapi_types::session::FileEntry::Session(
header.clone(),
))
.ok()
})
})
}
pub fn set_active_tools(&self, names: Vec<String>) {
self.inner.tools.set_active(names);
}
pub fn active_tool_names(&self) -> Vec<String> {
self.inner.tools.active()
}
pub fn export_context(&self) -> (Option<String>, Vec<serde_json::Value>) {
let active = self.inner.tools.active();
let all = self.inner.tools.all();
let tools = active
.iter()
.filter_map(|name| all.iter().find(|tool| tool.name == *name))
.map(|tool| {
serde_json::json!({
"name": tool.name,
"description": tool.description,
"parameters": tool.parameters,
})
})
.collect();
(self.system_prompt_text(&active).ok(), tools)
}
pub fn tools(&self) -> &Tools {
&self.inner.tools
}
pub fn append_custom_entry(
&self,
custom_type: &str,
data: Option<Value>,
) -> Result<(), String> {
let entry = self.with_session(|file| {
let id = file
.append_custom_entry(custom_type, data)
.map_err(|err| err.to_string())?;
Ok::<_, String>(file.entry(&id).cloned())
})?;
if let Some(entry) = entry {
self.emit(&AgentEvent::EntryAppended { entry });
}
Ok(())
}
pub fn callable_tools(&self) -> Vec<RegisteredTool> {
self.inner.tools.with(|registry| {
let active = registry.active();
registry
.all()
.into_iter()
.filter(|tool| match tool.exposure {
Exposure::Codemode | Exposure::Deferred => true,
Exposure::Direct => active.iter().any(|name| name == tool.name()),
Exposure::ModelOnly | Exposure::Hidden => false,
})
.collect()
})
}
pub async fn execute_tool(
&self,
caller: &str,
name: &str,
args: Value,
cancel: CancellationToken,
updates: Option<UpdateSink>,
allowed: impl Fn(&RegisteredTool) -> bool + Send,
) -> ToolCallOutcome {
let mut call = ToolCall {
id: String::new(),
name: name.to_owned(),
arguments: match args {
Value::Object(arguments) => arguments,
_ => serde_json::Map::new(),
},
thought_signature: None,
namespace: None,
};
let nested = &self.inner.nested;
let started = nested.start(caller, &mut call);
let hooks = Hooks {
session: self.clone(),
cancel: cancel.clone(),
turn_index: std::sync::atomic::AtomicU64::new(0),
};
let parent = Some(caller.to_owned());
hooks
.on_event(&AgentEvent::ToolExecutionStart {
tool_call_id: call.id.clone(),
tool_name: call.name.clone(),
args: Value::Object(call.arguments.clone()),
parent_tool_call_id: parent.clone(),
})
.await;
tokio::task::yield_now().await;
let mut callable = self.callable_tools();
callable.retain(|tool| allowed(tool));
let exclusive = !started.holds_queue
&& callable.iter().any(|tool| {
tool.name() == name && tool.tool.execution_mode() == ExecutionMode::Sequential
});
let queue = if exclusive {
Some(nested.queue.clone().lock_owned().await)
} else {
None
};
nested.enter(&started, started.holds_queue || exclusive);
let outcome = match self.last_assistant() {
None => ToolCallOutcome {
call: call.clone(),
result: ToolResult {
content: vec![ContentBlock::text("No assistant message issued this call")],
details: Some(Value::Object(serde_json::Map::new())),
..ToolResult::default()
},
is_error: true,
},
Some(assistant) => {
let tools: Vec<Arc<dyn Tool>> =
callable.into_iter().map(|tool| tool.tool).collect();
let messages = self.messages();
let scope = ToolCallScope {
tools: &tools,
assistant: &assistant,
messages: &messages,
parent: Some(caller),
cancel: &cancel,
};
let (sender, mut receiver) = tokio::sync::mpsc::unbounded_channel();
let sink: UpdateSink = Arc::new(move |partial: ToolResult| {
if let Some(updates) = &updates {
updates(partial.clone());
}
let _ = sender.send(partial);
});
let update = |partial_result| AgentEvent::ToolExecutionUpdate {
tool_call_id: call.id.clone(),
tool_name: call.name.clone(),
args: Value::Object(call.arguments.clone()),
partial_result,
parent_tool_call_id: parent.clone(),
};
let run = yapi_agent::run_tool_call(&scope, call.clone(), &hooks, sink);
tokio::pin!(run);
loop {
tokio::select! {
biased;
Some(partial) = receiver.recv() => hooks.on_event(&update(partial)).await,
outcome = &mut run => {
while let Ok(partial) = receiver.try_recv() {
hooks.on_event(&update(partial)).await;
}
break outcome;
}
}
}
}
};
drop(queue);
nested.finish(
started,
outcome.is_error,
&yapi_types::message::blocks_text(&outcome.result.content, "\n"),
outcome.result.usage.as_ref(),
);
hooks
.on_event(&AgentEvent::ToolExecutionEnd {
tool_call_id: call.id.clone(),
tool_name: call.name.clone(),
result: outcome.result.clone(),
is_error: outcome.is_error,
parent_tool_call_id: parent,
})
.await;
outcome
}
pub fn downgrade(&self) -> WeakSession {
WeakSession(Arc::downgrade(&self.inner))
}
pub fn cwd(&self) -> &Path {
&self.inner.cwd
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SkillBlock {
pub name: String,
pub location: String,
pub content: String,
pub user_message: Option<String>,
}
pub fn parse_skill_block(text: &str) -> Option<SkillBlock> {
let rest = text.strip_prefix("<skill name=\"")?;
let (name, rest) = rest.split_once('"')?;
let rest = rest.strip_prefix(" location=\"")?;
let (location, rest) = rest.split_once('"')?;
let body = rest.strip_prefix(">\n")?;
if name.is_empty() || location.is_empty() {
return None;
}
let mut from = 0;
while let Some(found) = body[from..].find("\n</skill>") {
let end = from + found;
let after = &body[end + "\n</skill>".len()..];
let user_message = match after {
"" => Some(None),
_ => after
.strip_prefix("\n\n")
.filter(|message| !message.is_empty())
.map(|message| Some(message.trim()).filter(|message| !message.is_empty())),
};
if let Some(user_message) = user_message {
return Some(SkillBlock {
name: name.to_owned(),
location: location.to_owned(),
content: body[..end].to_owned(),
user_message: user_message.map(str::to_owned),
});
}
from = end + 1;
}
None
}
pub fn assistant_text(message: &yapi_types::message::AssistantMessage) -> String {
yapi_types::message::blocks_text(&message.content, "")
}
pub fn failed(message: &yapi_types::message::AssistantMessage) -> bool {
matches!(message.stop_reason, StopReason::Error | StopReason::Aborted)
}
fn drain(queue: &Mutex<VecDeque<Message>>, all: bool) -> Vec<Message> {
let mut queue = lock(queue);
if all {
queue.drain(..).collect()
} else {
queue.pop_front().into_iter().collect()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parses_skill_blocks() {
let block = parse_skill_block(
"<skill name=\"demo\" location=\"/s/SKILL.md\">\nbody\n</skill>\n\n go now ",
)
.expect("a skill block");
assert_eq!(block.name, "demo");
assert_eq!(block.location, "/s/SKILL.md");
assert_eq!(block.content, "body");
assert_eq!(block.user_message.as_deref(), Some("go now"));
let block =
parse_skill_block("<skill name=\"a\" location=\"b\">\nx\n</skill>\ny\n</skill>")
.expect("the later closing tag");
assert_eq!(block.content, "x\n</skill>\ny");
assert_eq!(block.user_message, None);
assert!(
parse_skill_block("<skill name=\"a\" location=\"b\">\nx\n</skill> trailing").is_none()
);
assert!(parse_skill_block("hello").is_none());
}
#[test]
fn poisoned_registry_keeps_credentials() {
let model: Model = serde_json::from_value(serde_json::json!({
"id": "m", "name": "M", "api": "faux", "provider": "p", "baseUrl": "",
"reasoning": false, "input": ["text"],
"cost": {"input": 0, "output": 0, "cacheRead": 0, "cacheWrite": 0},
"contextWindow": 1000, "maxTokens": 100,
}))
.expect("a model");
let mut registry = ModelRegistry::builtin();
registry.register_provider("p", vec![model.clone()]);
registry.set_runtime_key("p", "key".into());
let session = AgentSession::new(SessionConfig {
cwd: PathBuf::from("/work"),
agent_dir: PathBuf::from("/agent"),
settings: SettingsManager::in_memory(),
registry,
apis: Apis::default(),
session: SessionManager::in_memory(Path::new("/work")),
model: Some(model),
thinking_level: ThinkingLevel::Off,
tools: Vec::new(),
extensions: Vec::new(),
include_extension_tools: false,
allowed_tools: None,
excluded_tools: Vec::new(),
resources: Resources::default(),
docs: crate::docs::Locations::default(),
});
let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
let _slot = session.inner.registry.write();
panic!("poison the registry");
}));
assert!(session.inner.registry.is_poisoned());
assert!(session.registry().has_auth("p"));
assert!(
session
.available_models()
.iter()
.any(|model| model.provider == "p")
);
assert!(session.with_registry(|registry| registry.has_auth("p")));
}
}