Skip to main content

yapi_core/agent_session/
mod.rs

1//! A conversation: the agent loop bound to a session file, tools, the system prompt
2//! and the model registry.
3//!
4//! Port of the core of `packages/coding-agent/src/core/agent-session.ts` in pi
5//! `v1.0.0`: prompts with template expansion, the system prompt as transcript
6//! messages, persistence of every message, queues, abort and model selection.
7
8use std::collections::VecDeque;
9use std::path::{Path, PathBuf};
10use std::sync::atomic::{AtomicUsize, Ordering};
11use std::sync::{Arc, Mutex, RwLock};
12
13use futures_util::future::BoxFuture;
14use indexmap::IndexMap;
15use serde_json::Value;
16use tokio_util::sync::CancellationToken;
17use yapi_agent::hooks::AgentHooks;
18use yapi_agent::{
19    AgentContext, ExecutionMode, LoopConfig, Tool, ToolCallOutcome, ToolCallScope, UpdateSink,
20};
21use yapi_ai::api::Apis;
22use yapi_ai::registry::ModelRegistry;
23use yapi_ai::stream::{Hook, RequestHeaders, RequestHooks, StreamOptions, ThinkingBudgets};
24use yapi_types::event::{AgentEvent, SummarySource, ToolResult};
25use yapi_types::extension_event::{ExtensionEvent, RegisteredMcpServer, TreePreparation};
26use yapi_types::message::{
27    AssistantMessage, Content, ContentBlock, ImageContent, Message, StopReason, SystemMessage,
28    ThinkingLevel, ToolCall, UserMessage,
29};
30use yapi_types::model::Model;
31use yapi_types::rpc::{PromptDisposition, StreamingBehavior};
32use yapi_types::session::FileEntry;
33use yapi_types::settings::Transport;
34use yapi_types::sync::{lock, read, write};
35use yapi_types::system_prompt::SystemPromptOptions;
36
37use crate::compaction::BranchSummary;
38use crate::extensions::{
39    BashOperations, Extension, ExtensionUi, Loadout, LoadoutChanges, Mode, NoUi, NotifyKind,
40    SessionActions, Tools,
41};
42use crate::resources::{ContextFile, PromptTemplate, Skill, expand_prompt_template};
43use crate::session::SessionManager;
44use crate::settings::SettingsManager;
45use crate::system_prompt::{build_sections, diff_sections};
46use crate::time::now_ms;
47use crate::tools::registry::ToolRegistry;
48use crate::tools::{BUILTIN_TOOLS, Described, Exposure, RegisteredTool, Runtime, ToolEnv, builtin};
49
50mod boundary;
51mod cache_warmer;
52mod extensions;
53mod models;
54mod recovery;
55mod stats;
56
57pub use boundary::Boundary;
58pub use cache_warmer::{
59    CacheWarmingDecision, CacheWarmingPhase, CacheWarmingState, CacheWarmingStatus,
60    format_cache_warming_status, format_cache_warming_usage,
61};
62use extensions::Hooks;
63pub use models::CatalogRefresh;
64pub use stats::{ContextUsage, SessionStats, UsageTotals};
65
66/// Where user input comes from, as pi's `InputSource` tells extensions'
67/// `input` handlers.
68#[derive(Clone, Copy, Debug, PartialEq, Eq)]
69pub enum InputSource {
70    /// Typed in the terminal.
71    Interactive,
72    /// An RPC client.
73    Rpc,
74    /// An extension's `sendUserMessage`.
75    Extension,
76}
77
78impl InputSource {
79    fn as_str(self) -> &'static str {
80        match self {
81            InputSource::Interactive => "interactive",
82            InputSource::Rpc => "rpc",
83            InputSource::Extension => "extension",
84        }
85    }
86}
87
88/// pi's `PromptOptions` for [`AgentSession::prompt_with`]. The default is
89/// pi's.
90#[derive(Clone, Debug)]
91pub struct PromptOptions {
92    /// Extension commands run, and `/skill:` commands and prompt templates
93    /// expand. Without it they stay literal text, but `input` handlers still
94    /// run.
95    pub expand_prompt_templates: bool,
96    /// How the message queues while a run streams. Without one, that is an
97    /// error.
98    pub streaming_behavior: Option<StreamingBehavior>,
99    /// Images sent with the text.
100    pub images: Vec<ImageContent>,
101    /// Where the message comes from, as `input` handlers see it.
102    pub source: InputSource,
103}
104
105impl Default for PromptOptions {
106    fn default() -> PromptOptions {
107        PromptOptions {
108            expand_prompt_templates: true,
109            streaming_behavior: None,
110            images: Vec::new(),
111            source: InputSource::Interactive,
112        }
113    }
114}
115
116/// Why a session takes over from another: pi's `session_start` and
117/// `session_shutdown` reasons other than `startup` and `quit`.
118#[derive(Clone, Copy, Debug, PartialEq, Eq)]
119pub enum Replacement {
120    /// `/new` or `newSession`.
121    New,
122    /// `/resume` or `switchSession`.
123    Resume,
124    /// `/fork`, `/clone` or `fork`.
125    Fork,
126    /// `/reload` or `reload`.
127    Reload,
128}
129
130impl Replacement {
131    /// The reason as pi's events name it.
132    pub fn as_str(self) -> &'static str {
133        match self {
134            Replacement::New => "new",
135            Replacement::Resume => "resume",
136            Replacement::Fork => "fork",
137            Replacement::Reload => "reload",
138        }
139    }
140
141    /// The other session's file as pi's `session_start` and
142    /// `session_shutdown` report it: not on reload.
143    fn reported(self, file: Option<String>) -> Option<String> {
144        file.filter(|_| self != Replacement::Reload)
145    }
146}
147
148/// A session change extensions may cancel, as pi's `session_before_switch`
149/// and `session_before_fork` report it.
150#[derive(Clone, Debug)]
151pub enum SessionChange {
152    /// A new session.
153    New,
154    /// A switch to this session file.
155    Resume(String),
156    /// A fork at an entry, keeping the entry when `at` (`/clone`) or ending
157    /// the fork before it.
158    Fork {
159        /// The entry.
160        entry_id: String,
161        /// Keep the entry.
162        at: bool,
163    },
164}
165
166impl SessionChange {
167    /// The pi event that may cancel the change.
168    pub fn event(&self) -> &'static str {
169        match self {
170            SessionChange::Fork { .. } => "session_before_fork",
171            _ => "session_before_switch",
172        }
173    }
174
175    /// Why the replacement session starts.
176    pub fn reason(&self) -> Replacement {
177        match self {
178            SessionChange::New => Replacement::New,
179            SessionChange::Resume(_) => Replacement::Resume,
180            SessionChange::Fork { .. } => Replacement::Fork,
181        }
182    }
183}
184
185/// What `user_bash` handlers decided for a `!` command.
186pub enum UserBash {
187    /// Run it with the session's shell.
188    Local,
189    /// Run it through an extension's operations.
190    Operations(BashOperations),
191    /// An extension ran it. The result is not yet recorded.
192    Done(crate::bash_executor::BashResult),
193}
194
195/// Receives every session event.
196pub type Listener = Box<dyn Fn(&AgentEvent) + Send + Sync>;
197
198/// Discovered resources the prompt uses.
199#[derive(Clone, Debug, Default)]
200pub struct Resources {
201    /// `AGENTS.md` and `CLAUDE.md` files.
202    pub context_files: Vec<ContextFile>,
203    /// Skills.
204    pub skills: Vec<Skill>,
205    /// Problems found while loading skills.
206    pub skill_diagnostics: Vec<crate::resources::Diagnostic>,
207    /// The skill files and directories `skills` came from, each path once.
208    pub skill_sources: Vec<crate::resources::SourceInfo>,
209    /// Prompt templates.
210    pub templates: Vec<PromptTemplate>,
211    /// Problems found while loading prompt templates.
212    pub template_diagnostics: Vec<crate::resources::Diagnostic>,
213    /// The template files and directories `templates` came from, each path
214    /// once.
215    pub template_sources: Vec<crate::resources::SourceInfo>,
216    /// Replaces the default prompt (`SYSTEM.md`, `--system-prompt`).
217    pub custom_prompt: Option<String>,
218    /// Appended to the prompt (`APPEND_SYSTEM.md`, `--append-system-prompt`).
219    pub append_prompt: Option<String>,
220    /// Theme files and directories with their sources, in precedence order:
221    /// the first file to declare a name wins.
222    pub themes: Vec<crate::resources::SourceInfo>,
223}
224
225/// What a session starts with.
226pub struct SessionConfig {
227    /// Working directory.
228    pub cwd: PathBuf,
229    /// Agent directory.
230    pub agent_dir: PathBuf,
231    /// Settings.
232    pub settings: SettingsManager,
233    /// Models and credentials.
234    pub registry: ModelRegistry,
235    /// Wire APIs.
236    pub apis: Apis,
237    /// The session file.
238    pub session: SessionManager,
239    /// The model, if one could be chosen.
240    pub model: Option<Model>,
241    /// The thinking level.
242    pub thinking_level: ThinkingLevel,
243    /// Active tool names.
244    pub tools: Vec<String>,
245    /// The session's extensions, which load before the active tools are set.
246    pub extensions: Vec<Arc<dyn Extension>>,
247    /// Activate the tools extensions register as they load, unless they opt
248    /// out. pi does unless `--tools` names the tools.
249    pub include_extension_tools: bool,
250    /// The only tools that may exist, from `--tools` (none for `--no-tools`),
251    /// or every tool when `None`.
252    pub allowed_tools: Option<Vec<String>>,
253    /// Tools that may never exist, from `--exclude-tools`.
254    pub excluded_tools: Vec<String>,
255    /// Prompt resources.
256    pub resources: Resources,
257    /// Where the model and the sign-in help find the docs.
258    pub docs: crate::docs::Locations,
259}
260
261struct State {
262    session: SessionManager,
263    model: Option<Model>,
264    thinking_level: ThinkingLevel,
265    /// A run is active, as pi's `isStreaming`. Queueing a message while it
266    /// is and the run's last look at the queues before it ends hold this
267    /// lock, so the run takes every message queued meanwhile.
268    streaming: bool,
269    /// `agent_settled` handlers are running. A prompt deferred while they
270    /// do and the run's taking of the deferred prompts hold this lock, so
271    /// the run starts every prompt deferred meanwhile.
272    settling: bool,
273}
274
275struct Inner {
276    cwd: PathBuf,
277    settings: Mutex<SettingsManager>,
278    registry: RwLock<Arc<ModelRegistry>>,
279    /// Serializes model catalog refreshes.
280    catalog_refresh: tokio::sync::Mutex<()>,
281    apis: RwLock<Apis>,
282    state: Mutex<State>,
283    /// The models cycling moves through. Empty means every available one.
284    scoped_models: Mutex<Vec<crate::model_resolver::ScopedModel>>,
285    tools: Tools,
286    extensions: Vec<Arc<dyn Extension>>,
287    /// The UI and mode extensions see, and how the mode changes sessions
288    /// for their commands.
289    binding: Mutex<(Arc<dyn ExtensionUi>, Mode, Option<SessionActions>)>,
290    /// Extension sections of the current run's system prompt.
291    run_options: Mutex<Option<SystemPromptOptions>>,
292    /// The base prompt options a command edited, with the options they were
293    /// made from. pi drops the edits when it rebuilds the base.
294    edited_base: Mutex<Option<(SystemPromptOptions, SystemPromptOptions)>>,
295    /// The system prompt a `before_agent_start` handler forced for the run.
296    forced_prompt: Mutex<Option<String>>,
297    /// Declared tools the last loadout left out of requests.
298    hidden_tools: Mutex<std::collections::HashSet<String>>,
299    /// Extension messages sent with the next prompt.
300    next_turn: Mutex<Vec<Message>>,
301    /// Extension messages sent during a turn, appended when it ends.
302    pending_custom: Mutex<Vec<Message>>,
303    resources: RwLock<Resources>,
304    runtime: Arc<RwLock<Runtime>>,
305    listeners: Mutex<Vec<Listener>>,
306    steering: Mutex<VecDeque<Message>>,
307    follow_up: Mutex<VecDeque<Message>>,
308    /// The steering and follow-up modes as pi's agent holds them: any value
309    /// last set, `None` once set to nothing.
310    queue_modes: Mutex<(Option<Value>, Option<Value>)>,
311    /// Texts of queued steering and follow-up messages until they start, as
312    /// `queue_update` reports them.
313    queued: Mutex<(Vec<String>, Vec<String>)>,
314    cancel: Mutex<Option<CancellationToken>>,
315    recovery: Mutex<Recovery>,
316    bash: Mutex<Vec<(u64, CancellationToken)>>,
317    pending_bash: Mutex<Vec<Message>>,
318    agent_dir: PathBuf,
319    /// Where the prompt points the model to the docs, fixed for the session
320    /// so a first-run download does not change it mid-session.
321    docs: crate::docs::Locations,
322    /// Woken when a run ends.
323    idle: tokio::sync::Notify,
324    /// Prompts started while `agent_settled` handlers ran, in order, as pi's
325    /// `_deferredSettledActions`: each learns when it may start, with a guard
326    /// it drops as it ends.
327    deferred: Mutex<Vec<tokio::sync::oneshot::Sender<tokio::sync::oneshot::Sender<()>>>>,
328    /// Compactions and branch summaries in progress.
329    compacting: AtomicUsize,
330    /// A manual compaction is running.
331    manual_compaction: std::sync::atomic::AtomicBool,
332    /// Cancels the wait before an automatic retry.
333    retry_cancel: Mutex<Option<CancellationToken>>,
334    /// Calls tools made through [`AgentSession::execute_tool`].
335    nested: crate::nested::NestedCalls,
336    /// Events for extensions that are delivered in the background, in order.
337    announcements: Mutex<VecDeque<Value>>,
338    /// Held while announcements are delivered.
339    announcing: tokio::sync::Mutex<()>,
340    /// How the last turn ended, as pi's boundary events report it.
341    outcome: Mutex<&'static str>,
342    /// Blocking extension dialogs open, and the kind and title of the
343    /// outermost.
344    ui_prompts: Mutex<(usize, Option<UiPrompt>)>,
345    /// MCP servers extensions registered, by name in registration order,
346    /// with the extension runtime that reported each, and whether the
347    /// extensions have started, so changes are announced.
348    mcp_servers: Mutex<McpServers>,
349    /// Registered MCP servers reported as connected by no extension.
350    unhandled_mcp_servers: Mutex<std::collections::HashSet<String>>,
351    /// Extensions that failed to load for the session, by path, with pi's
352    /// message.
353    extension_load_errors: Mutex<Vec<(String, String)>>,
354    /// Extensions that were not loaded for the session, by path, with pi's
355    /// warning.
356    extension_warnings: Mutex<Vec<(String, String)>>,
357    /// The extensions and events, by path and event type, whose handler
358    /// results were dropped for lack of a grant, which are warned of once.
359    dropped_results: Mutex<std::collections::HashSet<(String, String)>>,
360    /// Keeps the latest request's prompt cache entry warm.
361    cache_warmer: Mutex<cache_warmer::CacheWarmer>,
362    #[cfg(test)]
363    races: RacePoints,
364}
365
366/// A blocking extension dialog's kind and title.
367type UiPrompt = (String, Option<String>);
368
369/// See [`Inner::mcp_servers`].
370type McpServers = (IndexMap<String, (u64, RegisteredMcpServer)>, bool);
371
372/// Counts an operation in [`Inner::compacting`] while alive.
373struct Compacting<'a>(&'a AtomicUsize);
374
375impl<'a> Compacting<'a> {
376    fn start(counter: &'a AtomicUsize) -> Compacting<'a> {
377        counter.fetch_add(1, Ordering::SeqCst);
378        Compacting(counter)
379    }
380}
381
382impl Drop for Compacting<'_> {
383    fn drop(&mut self) {
384        self.0.fetch_sub(1, Ordering::SeqCst);
385    }
386}
387
388/// Code a test runs where a message may meet the end of a run.
389#[cfg(test)]
390#[derive(Default)]
391struct RacePoints {
392    /// After the run's last look at the queues, before it ends.
393    before_end: Mutex<Option<Box<dyn FnOnce() + Send>>>,
394    /// As a prompt asks to wait for `agent_settled` handlers, before it
395    /// takes its place among the deferred prompts.
396    before_defer: Mutex<Option<Box<dyn FnOnce() + Send>>>,
397}
398
399/// Runs the code a test set at `point`, once.
400#[cfg(test)]
401fn race(point: &Mutex<Option<Box<dyn FnOnce() + Send>>>) {
402    let hook = lock(point).take();
403    if let Some(hook) = hook {
404        hook();
405    }
406}
407
408/// Holds [`State::settling`], which the run set as it ended, until the run
409/// takes the prompts deferred meanwhile, so a run dropped while extensions
410/// hear `agent_settled` clears it, as pi's `finally` does.
411struct Settling<'a>(&'a Inner);
412
413impl Settling<'_> {
414    /// Ends it: the prompts deferred meanwhile, in order.
415    fn finish(self) -> Vec<tokio::sync::oneshot::Sender<tokio::sync::oneshot::Sender<()>>> {
416        let mut state = lock(&self.0.state);
417        state.settling = false;
418        let deferred = std::mem::take(&mut *lock(&self.0.deferred));
419        drop(state);
420        deferred
421    }
422}
423
424impl Drop for Settling<'_> {
425    fn drop(&mut self) {
426        lock(&self.0.state).settling = false;
427    }
428}
429
430/// Post-run bookkeeping for retries and overflow recovery, as pi keeps it.
431#[derive(Default)]
432struct Recovery {
433    retry_attempt: u32,
434    overflow_recovery_attempted: bool,
435    /// The last assistant message of the run and its entry id.
436    last_assistant: Option<(AssistantMessage, Option<String>)>,
437    /// Entry ids of tool results since the last assistant message.
438    turn_tool_results: Vec<String>,
439    /// Entry ids of the last completed turn's tool results.
440    last_tool_results: Vec<String>,
441}
442
443/// Options for [`AgentSession::navigate_tree`].
444#[derive(Clone, Debug, Default)]
445pub struct TreeNavigation {
446    /// Summarize the branch being left.
447    pub summarize: bool,
448    /// Extra or replacement summary instructions.
449    pub custom_instructions: Option<String>,
450    /// Use `custom_instructions` as the whole instruction.
451    pub replace_instructions: bool,
452    /// Label for the summary, or the target without one.
453    pub label: Option<String>,
454}
455
456/// What [`AgentSession::navigate_tree`] did.
457#[derive(Clone, Debug, Default, PartialEq)]
458pub struct TreeOutcome {
459    /// Text of a user or custom message target, for the editor.
460    pub editor_text: Option<String>,
461    /// Nothing changed.
462    pub cancelled: bool,
463    /// The summary request was cancelled.
464    pub aborted: bool,
465    /// The branch summary entry, when one was made.
466    pub summary_entry: Option<FileEntry>,
467}
468
469/// A running conversation. Cheap to clone, and clones share state.
470#[derive(Clone)]
471pub struct AgentSession {
472    inner: Arc<Inner>,
473}
474
475/// A handle to a session that does not keep it alive.
476#[derive(Clone, Default)]
477pub struct WeakSession(std::sync::Weak<Inner>);
478
479impl WeakSession {
480    /// The session, while it exists.
481    pub fn upgrade(&self) -> Option<AgentSession> {
482        self.0.upgrade().map(|inner| AgentSession { inner })
483    }
484}
485
486impl AgentSession {
487    /// Starts a session. A new session records its model and thinking level.
488    pub fn new(config: SessionConfig) -> AgentSession {
489        let SessionConfig {
490            cwd,
491            agent_dir,
492            settings,
493            registry,
494            apis,
495            mut session,
496            model,
497            thinking_level,
498            tools,
499            extensions,
500            include_extension_tools,
501            allowed_tools,
502            excluded_tools,
503            resources,
504            docs,
505        } = config;
506        let runtime = Arc::new(RwLock::new(Runtime::default()));
507        let env = ToolEnv {
508            cwd: cwd.clone(),
509            runtime: runtime.clone(),
510            bin_dir: crate::config::bin_dir(&agent_dir),
511        };
512        let available: Vec<RegisteredTool> = BUILTIN_TOOLS
513            .iter()
514            .filter_map(|name| builtin(name, &env))
515            .collect();
516        let mut tool_set = ToolRegistry::new(available, Vec::new());
517        tool_set.restrict(allowed_tools, excluded_tools);
518        let tool_registry = Tools::new(tool_set);
519        for extension in &extensions {
520            extension.load(&tool_registry);
521        }
522        let mut tools = tools;
523        if include_extension_tools {
524            tool_registry.with(|registry| {
525                for tool in registry.all() {
526                    let name = tool.name().to_owned();
527                    let builtin = BUILTIN_TOOLS.contains(&name.as_str());
528                    if !builtin && tool.exposure.declarable() && tool.default_active {
529                        tools.push(name);
530                    }
531                }
532            });
533        }
534        tool_registry.set_active(tools);
535
536        // As in pi, a session with messages gains a missing thinking level,
537        // and any other records its model and level.
538        let has_messages = !session.build_context().messages.is_empty();
539        if has_messages {
540            let has_thinking = session.branch_path(None).iter().any(|entry| {
541                matches!(
542                    entry,
543                    yapi_types::session::FileEntry::ThinkingLevelChange(_)
544                )
545            });
546            if !has_thinking {
547                let _ = session.append_thinking_level_change(thinking_level.as_str());
548            }
549        } else {
550            if let Some(model) = &model {
551                let _ = session.append_model_change(&model.provider, &model.id);
552            }
553            let _ = session.append_thinking_level_change(thinking_level.as_str());
554        }
555
556        // pi's agent takes its queue modes from the settings: a value that
557        // is not truthy reads as `one-at-a-time`.
558        let queue_mode = |key: &str| {
559            Some(
560                settings
561                    .merged_document()
562                    .get(key)
563                    .filter(|mode| yapi_ai::http::is_truthy(mode))
564                    .cloned()
565                    .unwrap_or_else(|| "one-at-a-time".into()),
566            )
567        };
568        let queue_modes = (queue_mode("steeringMode"), queue_mode("followUpMode"));
569        AgentSession {
570            inner: Arc::new(Inner {
571                cwd,
572                settings: Mutex::new(settings),
573                registry: RwLock::new(Arc::new(registry)),
574                catalog_refresh: tokio::sync::Mutex::new(()),
575                apis: RwLock::new(apis),
576                state: Mutex::new(State {
577                    session,
578                    model,
579                    thinking_level,
580                    streaming: false,
581                    settling: false,
582                }),
583                scoped_models: Mutex::new(Vec::new()),
584                tools: tool_registry,
585                extensions,
586                binding: Mutex::new((Arc::new(NoUi), Mode::Print, None)),
587                run_options: Mutex::new(None),
588                edited_base: Mutex::new(None),
589                forced_prompt: Mutex::new(None),
590                hidden_tools: Mutex::default(),
591                next_turn: Mutex::new(Vec::new()),
592                pending_custom: Mutex::new(Vec::new()),
593                resources: RwLock::new(resources),
594                runtime,
595                listeners: Mutex::new(Vec::new()),
596                steering: Mutex::new(VecDeque::new()),
597                follow_up: Mutex::new(VecDeque::new()),
598                queue_modes: Mutex::new(queue_modes),
599                queued: Mutex::new((Vec::new(), Vec::new())),
600                cancel: Mutex::new(None),
601                recovery: Mutex::new(Recovery::default()),
602                bash: Mutex::new(Vec::new()),
603                pending_bash: Mutex::new(Vec::new()),
604                docs,
605                agent_dir,
606                idle: tokio::sync::Notify::new(),
607                deferred: Mutex::new(Vec::new()),
608                #[cfg(test)]
609                races: RacePoints::default(),
610                compacting: AtomicUsize::new(0),
611                manual_compaction: std::sync::atomic::AtomicBool::new(false),
612                retry_cancel: Mutex::new(None),
613                announcements: Mutex::new(VecDeque::new()),
614                announcing: tokio::sync::Mutex::new(()),
615                outcome: Mutex::new("completed"),
616                ui_prompts: Mutex::new((0, None)),
617                mcp_servers: Mutex::default(),
618                unhandled_mcp_servers: Mutex::default(),
619                extension_load_errors: Mutex::default(),
620                extension_warnings: Mutex::default(),
621                dropped_results: Mutex::default(),
622                cache_warmer: Mutex::default(),
623                nested: crate::nested::NestedCalls::default(),
624            }),
625        }
626    }
627
628    /// Where the model and the sign-in help find the docs.
629    pub fn docs(&self) -> &crate::docs::Locations {
630        &self.inner.docs
631    }
632
633    /// Adds a listener for every event.
634    pub fn subscribe(&self, listener: Listener) {
635        lock(&self.inner.listeners).push(listener);
636    }
637
638    fn emit(&self, event: &AgentEvent) {
639        for listener in lock(&self.inner.listeners).iter() {
640            listener(event);
641        }
642    }
643
644    /// Runs `f` with the session file. The session's first failed write is
645    /// reported to the user, after `f`, by whichever call finds it.
646    pub fn with_session<T>(&self, f: impl FnOnce(&mut SessionManager) -> T) -> T {
647        let (value, error) = {
648            let session = &mut lock(&self.inner.state).session;
649            (f(session), session.take_write_error())
650        };
651        if let Some(error) = error {
652            let message = format!("Failed to save session: {error}");
653            let (ui, _) = self.extension_binding();
654            if ui.has_ui() {
655                ui.notify(&message, NotifyKind::Error);
656            } else {
657                use std::io::Write as _;
658                let _ = writeln!(std::io::stderr(), "{message}");
659            }
660        }
661        value
662    }
663
664    /// Moves the session file out, leaving an empty in-memory one, so a new
665    /// `AgentSession` can take it over.
666    pub fn take_session(&self) -> SessionManager {
667        let mut state = lock(&self.inner.state);
668        let placeholder = SessionManager::in_memory(state.session.cwd());
669        std::mem::replace(&mut state.session, placeholder)
670    }
671
672    /// The current model.
673    pub fn model(&self) -> Option<Model> {
674        lock(&self.inner.state).model.clone()
675    }
676
677    /// The current thinking level.
678    pub fn thinking_level(&self) -> ThinkingLevel {
679        lock(&self.inner.state).thinking_level
680    }
681
682    /// Whether a run is in progress.
683    pub fn is_streaming(&self) -> bool {
684        lock(&self.inner.state).streaming
685    }
686
687    /// The model context of the current branch.
688    pub fn messages(&self) -> Vec<Message> {
689        self.with_session(|session| session.build_context().messages)
690    }
691
692    /// The text of the last assistant message.
693    pub fn last_assistant(&self) -> Option<yapi_types::message::AssistantMessage> {
694        self.messages()
695            .into_iter()
696            .rev()
697            .find_map(|message| match message {
698                Message::Assistant(assistant) => Some(*assistant),
699                _ => None,
700            })
701    }
702
703    /// Names the session. Extensions hear of it as pi's
704    /// `session_info_changed`.
705    pub fn set_name(&self, name: &str) {
706        let name = self.with_session(|session| {
707            let _ = session.append_session_info(name);
708            session.name()
709        });
710        self.emit(&AgentEvent::SessionInfoChanged { name: name.clone() });
711        self.announce(ExtensionEvent::SessionInfoChanged {
712            name: name.as_deref(),
713        });
714    }
715
716    /// A copy of the session's settings manager, for work that reads both
717    /// scopes, such as package checks.
718    pub fn settings_manager(&self) -> SettingsManager {
719        lock(&self.inner.settings).clone()
720    }
721
722    /// A snapshot of the merged settings.
723    pub fn settings(&self) -> yapi_types::settings::Settings {
724        lock(&self.inner.settings).settings().clone()
725    }
726
727    /// The `images.autoResize` setting and the current model's resize
728    /// profile, for images entering the conversation.
729    fn image_options(&self) -> (bool, Option<yapi_types::models::ImageResize>) {
730        let auto_resize = lock(&self.inner.settings).settings().image_auto_resize();
731        let limits = lock(&self.inner.state)
732            .model
733            .as_ref()
734            .and_then(|model| model.image_resize().cloned());
735        (auto_resize, limits)
736    }
737
738    /// pi's `getHttpIdleTimeoutMs` for these settings. Startup rejects an
739    /// invalid value, so it reads as the default here.
740    pub fn http_idle_timeout_ms(&self) -> u64 {
741        lock(&self.inner.settings)
742            .http_idle_timeout_ms()
743            .unwrap_or(yapi_ai::http::DEFAULT_IDLE_TIMEOUT_MS)
744    }
745
746    /// Settings files that failed to load, as pi's warnings word them.
747    pub fn settings_errors(&self) -> Vec<String> {
748        lock(&self.inner.settings).errors().to_vec()
749    }
750
751    /// Whether this session loads the project's settings and resources.
752    pub fn project_trusted(&self) -> bool {
753        lock(&self.inner.settings).project_trusted()
754    }
755
756    /// The directory this session's project files come from.
757    pub fn project_dir(&self) -> PathBuf {
758        lock(&self.inner.settings).project_dir().to_path_buf()
759    }
760
761    /// Sets a global setting and writes `settings.json`.
762    pub fn set_global_setting(
763        &self,
764        key: &str,
765        value: Option<serde_json::Value>,
766    ) -> Result<(), String> {
767        lock(&self.inner.settings)
768            .set(crate::settings::Scope::Global, key, value)
769            .map_err(|err| err.to_string())
770    }
771
772    /// The prompt resources: those the session was built with, and those
773    /// its extensions discovered.
774    pub fn resources(&self) -> Resources {
775        read(&self.inner.resources).clone()
776    }
777
778    /// Empties the steering and follow-up queues and returns their texts.
779    pub fn clear_queues(&self) -> (Vec<String>, Vec<String>) {
780        lock(&self.inner.steering).clear();
781        lock(&self.inner.follow_up).clear();
782        let queued = std::mem::take(&mut *lock(&self.inner.queued));
783        self.emit_queue_update();
784        queued
785    }
786
787    /// Queued steering and follow-up messages.
788    pub fn pending_message_count(&self) -> usize {
789        let queued = lock(&self.inner.queued);
790        queued.0.len() + queued.1.len()
791    }
792
793    /// How queued steering messages are delivered: `all`, or one at a time
794    /// for any other value. It is the mode last set, as pi's agent keeps it,
795    /// and `None` once set to nothing.
796    pub fn steering_mode(&self) -> Option<Value> {
797        lock(&self.inner.queue_modes).0.clone()
798    }
799
800    /// How queued follow-up messages are delivered, as
801    /// [`AgentSession::steering_mode`].
802    pub fn follow_up_mode(&self) -> Option<Value> {
803        lock(&self.inner.queue_modes).1.clone()
804    }
805
806    /// Whether the steering and the follow-up queue each deliver all their
807    /// messages at once.
808    pub(super) fn drains_all(&self) -> (bool, bool) {
809        let all = |mode: &Option<Value>| mode.as_ref().and_then(Value::as_str) == Some("all");
810        let modes = lock(&self.inner.queue_modes);
811        (all(&modes.0), all(&modes.1))
812    }
813
814    /// Sets the steering mode, any value as pi takes it, and saves it in the
815    /// global settings. `None` removes the setting.
816    pub fn set_steering_mode(&self, mode: Option<Value>) -> Result<(), String> {
817        lock(&self.inner.queue_modes).0.clone_from(&mode);
818        self.set_global_setting("steeringMode", mode)
819    }
820
821    /// Sets the follow-up mode as [`AgentSession::set_steering_mode`] does.
822    pub fn set_follow_up_mode(&self, mode: Option<Value>) -> Result<(), String> {
823        lock(&self.inner.queue_modes).1.clone_from(&mode);
824        self.set_global_setting("followUpMode", mode)
825    }
826
827    /// pi's `autoCompactionEnabled`: the setting as saved, `true` when unset
828    /// or `null`. Compaction runs when it is truthy.
829    pub fn auto_compaction_enabled(&self) -> Value {
830        lock(&self.inner.settings)
831            .merged_document()
832            .get("compaction")
833            .and_then(|compaction| compaction.get("enabled"))
834            .filter(|enabled| !enabled.is_null())
835            .cloned()
836            .unwrap_or(Value::Bool(true))
837    }
838
839    /// Saves whether automatic compaction runs in the global settings, any
840    /// value as pi takes it. `None` removes the setting.
841    pub fn set_auto_compaction(&self, enabled: Option<Value>) -> Result<(), String> {
842        lock(&self.inner.settings)
843            .set_nested(
844                crate::settings::Scope::Global,
845                "compaction",
846                "enabled",
847                enabled,
848            )
849            .map_err(|err| err.to_string())
850    }
851
852    /// Saves whether automatic retries run, as
853    /// [`AgentSession::set_auto_compaction`] does.
854    pub fn set_auto_retry(&self, enabled: Option<Value>) -> Result<(), String> {
855        lock(&self.inner.settings)
856            .set_nested(crate::settings::Scope::Global, "retry", "enabled", enabled)
857            .map_err(|err| err.to_string())
858    }
859
860    /// Sets `key` inside the global object setting `field` and writes
861    /// `settings.json`.
862    pub fn set_nested_global_setting(
863        &self,
864        field: &str,
865        key: &str,
866        value: serde_json::Value,
867    ) -> Result<(), String> {
868        lock(&self.inner.settings)
869            .set_nested(crate::settings::Scope::Global, field, key, Some(value))
870            .map_err(|err| err.to_string())
871    }
872
873    /// Whether a compaction or branch summary is running.
874    pub fn is_compacting(&self) -> bool {
875        self.inner.compacting.load(Ordering::SeqCst) > 0
876    }
877
878    /// Waits until no run is in progress.
879    pub async fn wait_for_idle(&self) {
880        loop {
881            let idle = self.inner.idle.notified();
882            // As pi's `isIdle`, a compaction keeps the session busy.
883            if !self.is_streaming() && !self.is_compacting() {
884                return;
885            }
886            idle.await;
887        }
888    }
889
890    /// Runs a user `!` command in the session's directory, with the
891    /// session's shell or through an extension's `operations`, streaming
892    /// output to `on_chunk` and as `bash_execution_update` events tagged
893    /// `id`, and records it. `exclude_from_context` (`!!`) keeps the output
894    /// from the model and is recorded as given. While a run streams, the
895    /// record waits for its end.
896    pub async fn execute_bash(
897        &self,
898        command: &str,
899        exclude_from_context: Option<bool>,
900        id: Option<String>,
901        operations: Option<BashOperations>,
902        mut on_chunk: impl FnMut(&str),
903    ) -> Result<crate::bash_executor::BashResult, String> {
904        static NEXT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
905        let token = NEXT.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
906        let cancel = CancellationToken::new();
907        lock(&self.inner.bash).push((token, cancel.clone()));
908        let settings = self.settings();
909        let resolved =
910            crate::tools::bash::with_prefix(settings.shell_command_prefix.as_deref(), command);
911        let cwd = self.with_session(|session| session.cwd().to_path_buf());
912        let on_chunk = |delta: &str| {
913            on_chunk(delta);
914            self.emit(&AgentEvent::BashExecutionUpdate {
915                id: id.clone(),
916                delta: delta.to_owned(),
917            });
918        };
919        let result = match &operations {
920            None => {
921                crate::bash_executor::execute(
922                    &resolved,
923                    &cwd,
924                    settings.shell_path.as_deref(),
925                    &crate::config::bin_dir(&self.inner.agent_dir),
926                    cancel.clone(),
927                    on_chunk,
928                )
929                .await
930            }
931            Some(operations) => {
932                let run = |output| {
933                    operations.extension.run_bash(
934                        &operations.handle,
935                        &resolved,
936                        &cwd,
937                        output,
938                        cancel.clone(),
939                    )
940                };
941                crate::bash_executor::execute_with(run, &cancel, on_chunk).await
942            }
943        };
944        lock(&self.inner.bash).retain(|(running, _)| *running != token);
945        let result = result?;
946        self.record_bash(command, &result, exclude_from_context);
947        Ok(result)
948    }
949
950    /// pi's `user_bash` event for a `!` command: whether an extension ran
951    /// it, runs it through its operations, or leaves it to the session's
952    /// shell. Fails with the error of a handler, which extension error
953    /// reporting has already shown.
954    pub async fn user_bash(
955        &self,
956        command: &str,
957        exclude_from_context: bool,
958    ) -> Result<UserBash, String> {
959        let handlers = self.handlers_of("user_bash");
960        if handlers.is_empty() {
961            return Ok(UserBash::Local);
962        }
963        let cwd = self.with_session(|session| session.cwd().to_path_buf());
964        let event = ExtensionEvent::UserBash {
965            command,
966            exclude_from_context,
967            cwd: &cwd,
968        }
969        .to_value();
970        let ctx = self.extension_context(CancellationToken::new());
971        for extension in handlers {
972            let Some(result) = extension.handle(&ctx, &event).await else {
973                continue;
974            };
975            if let Some(error) = result["error"].as_str() {
976                return Err(error.to_owned());
977            }
978            // What the command runs and shows then comes from the extension,
979            // where the user expects their own shell.
980            let replaces = !result["operations"].is_null() || !result["result"].is_null();
981            if replaces
982                && !self.applies(
983                    extension,
984                    "user_bash",
985                    extension.grants().is_all(),
986                    "every grant",
987                )
988            {
989                continue;
990            }
991            if !result["operations"].is_null() {
992                return Ok(UserBash::Operations(BashOperations {
993                    extension: Arc::clone(extension),
994                    handle: result["operations"].clone(),
995                }));
996            }
997            if let Ok(done) = serde_json::from_value(result["result"].clone()) {
998                return Ok(UserBash::Done(done));
999            }
1000            let error = "Invalid user_bash handler result: return undefined for local execution or exactly one valid { operations } or { result } object";
1001            ctx.ui
1002                .extension_error(&extension.source().path, "user_bash", error, None);
1003            return Err(error.to_owned());
1004        }
1005        Ok(UserBash::Local)
1006    }
1007
1008    /// pi's `recordBashResult`: records a `!` command's result, after the
1009    /// current run when one streams.
1010    pub fn record_bash(
1011        &self,
1012        command: &str,
1013        result: &crate::bash_executor::BashResult,
1014        exclude_from_context: Option<bool>,
1015    ) {
1016        let message = Message::BashExecution(yapi_types::message::BashExecutionMessage {
1017            command: command.to_owned(),
1018            output: result.output.clone(),
1019            exit_code: result.exit_code.map(i64::from),
1020            cancelled: result.cancelled,
1021            truncated: result.truncated,
1022            full_output_path: result.full_output_path.clone(),
1023            timestamp: now_ms(),
1024            exclude_from_context,
1025        });
1026        if self.is_streaming() {
1027            lock(&self.inner.pending_bash).push(message);
1028        } else {
1029            let _ = self.with_session(|session| session.append_message(message));
1030        }
1031    }
1032
1033    fn flush_pending_bash(&self) {
1034        let pending: Vec<Message> = lock(&self.inner.pending_bash).drain(..).collect();
1035        for message in pending {
1036            let _ = self.with_session(|session| session.append_message(message));
1037        }
1038    }
1039
1040    /// Cancels running `!` commands.
1041    pub fn abort_bash(&self) {
1042        for (_, token) in lock(&self.inner.bash).iter() {
1043            token.cancel();
1044        }
1045    }
1046
1047    /// Whether a `!` command is running.
1048    pub fn is_bash_running(&self) -> bool {
1049        !lock(&self.inner.bash).is_empty()
1050    }
1051
1052    /// User messages of every branch with text, in file order, as fork points.
1053    pub fn user_messages_for_forking(&self) -> Vec<(String, String)> {
1054        self.with_session(|session| {
1055            session
1056                .entries()
1057                .filter_map(|entry| match entry {
1058                    FileEntry::Message(entry) => match &entry.message {
1059                        Message::User(user) => {
1060                            let text = user.content.text("");
1061                            (!text.is_empty()).then(|| (entry.meta.id.clone(), text))
1062                        }
1063                        _ => None,
1064                    },
1065                    _ => None,
1066                })
1067                .collect()
1068        })
1069    }
1070
1071    /// pi's `getLastAssistantText`: the text of the last assistant message,
1072    /// skipping aborted ones without content, or `None` when that message has
1073    /// no text.
1074    pub fn last_assistant_text(&self) -> Option<String> {
1075        let assistant = self
1076            .messages()
1077            .into_iter()
1078            .rev()
1079            .find_map(|message| match message {
1080                Message::Assistant(assistant)
1081                    if !(assistant.stop_reason == StopReason::Aborted
1082                        && assistant.content.is_empty()) =>
1083                {
1084                    Some(assistant)
1085                }
1086                _ => None,
1087            })?;
1088        let text = assistant_text(&assistant);
1089        let text = text.trim();
1090        (!text.is_empty()).then(|| text.to_owned())
1091    }
1092
1093    /// What [`AgentSession::abort`] cancels: the current run's token, or
1094    /// `None` while no run is active, as pi's `agent.signal`.
1095    pub fn run_cancel(&self) -> Option<CancellationToken> {
1096        lock(&self.inner.cancel).clone()
1097    }
1098
1099    /// Cancels the current run.
1100    pub fn abort(&self) {
1101        if let Some(cancel) = lock(&self.inner.cancel).as_ref() {
1102            cancel.cancel();
1103        }
1104    }
1105
1106    fn user_message(text: String, images: Vec<ImageContent>) -> Message {
1107        let mut content = vec![ContentBlock::text(text)];
1108        content.extend(images.into_iter().map(ContentBlock::Image));
1109        Message::User(UserMessage {
1110            content: Content::Blocks(content),
1111            timestamp: now_ms(),
1112        })
1113    }
1114
1115    /// The part of pi's `sendCustomMessage` that happens at once. The message
1116    /// is queued, held or recorded and reported before this returns:
1117    /// - `deliver_as` `nextTurn` holds the message for the next prompt.
1118    /// - While a run streams, it steers (or follows up) unless `trigger_turn`
1119    ///   is false, in which case it is appended when the turn ends.
1120    /// - Otherwise `trigger_turn` returns the run it starts, which the caller
1121    ///   defers as pi does (see [`AgentSession::settled_turn`]), and without
1122    ///   one it is appended at once.
1123    pub fn deliver_custom_message(
1124        &self,
1125        message: yapi_types::message::CustomMessage,
1126        trigger_turn: Option<bool>,
1127        deliver_as: Option<&str>,
1128    ) -> Option<impl Future<Output = ()> + Send + 'static> {
1129        let message = Message::Custom(message);
1130        if deliver_as == Some("nextTurn") {
1131            lock(&self.inner.next_turn).push(message);
1132            return None;
1133        }
1134        // Queued in one step with the check; see [`State::streaming`].
1135        let state = lock(&self.inner.state);
1136        if state.streaming && trigger_turn != Some(false) {
1137            let queue = if deliver_as == Some("followUp") {
1138                &self.inner.follow_up
1139            } else {
1140                &self.inner.steering
1141            };
1142            lock(queue).push_back(message);
1143        } else if trigger_turn == Some(true) {
1144            drop(state);
1145            let session = self.clone();
1146            return Some(async move {
1147                let _ = session.run(vec![message]).await;
1148            });
1149        } else if state.streaming {
1150            lock(&self.inner.pending_custom).push(message);
1151        } else {
1152            drop(state);
1153            self.append_custom_message(&message);
1154        }
1155        None
1156    }
1157
1158    /// pi defers a prompt started while `agent_settled` handlers run: the
1159    /// run that settled starts it once they return, and ends after it. `None`
1160    /// outside those handlers. Otherwise the prompt's place is taken now, and
1161    /// the future resolves when it may start, to a guard it holds until it
1162    /// ends.
1163    pub fn settled_turn(
1164        &self,
1165    ) -> Option<impl Future<Output = Option<tokio::sync::oneshot::Sender<()>>> + Send + 'static>
1166    {
1167        #[cfg(test)]
1168        race(&self.inner.races.before_defer);
1169        let state = lock(&self.inner.state);
1170        if !state.settling {
1171            return None;
1172        }
1173        let (turn, start) = tokio::sync::oneshot::channel();
1174        lock(&self.inner.deferred).push(turn);
1175        drop(state);
1176        Some(async move { start.await.ok() })
1177    }
1178
1179    /// Records an extension message and reports it as pi's `_appendCustomMessage` does.
1180    fn append_custom_message(&self, message: &Message) {
1181        let Message::Custom(custom) = message else {
1182            return;
1183        };
1184        self.with_session(|file| {
1185            let _ = file.append_custom_message(
1186                &custom.custom_type,
1187                custom.content.clone(),
1188                custom.display,
1189                custom.details.clone(),
1190            );
1191        });
1192        self.emit(&AgentEvent::MessageStart {
1193            message: message.clone(),
1194        });
1195        self.emit(&AgentEvent::MessageEnd {
1196            message: message.clone(),
1197        });
1198    }
1199
1200    /// Appends extension messages held while a turn ran.
1201    fn flush_pending_custom(&self) {
1202        let pending = std::mem::take(&mut *lock(&self.inner.pending_custom));
1203        for message in &pending {
1204            self.append_custom_message(message);
1205        }
1206    }
1207
1208    fn emit_queue_update(&self) {
1209        let (steering, follow_up) = lock(&self.inner.queued).clone();
1210        self.emit(&AgentEvent::QueueUpdate {
1211            steering,
1212            follow_up,
1213        });
1214    }
1215
1216    /// pi's queue display: a queued message leaves it when it starts.
1217    fn dequeue_started(&self, message: &Message) {
1218        let Message::User(user) = message else {
1219            return;
1220        };
1221        let text = user.content.text("");
1222        if text.is_empty() {
1223            return;
1224        }
1225        let removed = {
1226            let mut queued = lock(&self.inner.queued);
1227            let (steering, follow_up) = &mut *queued;
1228            [steering, follow_up].into_iter().any(|texts| {
1229                let index = texts.iter().position(|queued| *queued == text);
1230                index.map(|index| texts.remove(index)).is_some()
1231            })
1232        };
1233        if removed {
1234            self.emit_queue_update();
1235        }
1236    }
1237
1238    /// pi's `steer` and `followUp`: an extension command cannot be queued.
1239    /// Otherwise the `input` handlers see the message first, then it is
1240    /// queued with templates and skills expanded.
1241    pub async fn queue_input(
1242        &self,
1243        text: &str,
1244        images: Vec<ImageContent>,
1245        behavior: StreamingBehavior,
1246        source: InputSource,
1247    ) -> Result<PromptDisposition, String> {
1248        if self.is_extension_command(text) {
1249            let name = text[1..].split(' ').next().unwrap_or_default();
1250            return Err(format!(
1251                "Extension command \"/{name}\" cannot be queued. Use prompt() or execute the command when not streaming."
1252            ));
1253        }
1254        let streaming = self.is_streaming().then_some(behavior);
1255        let Some((text, images)) = self.input_handlers(text, images, source, streaming).await
1256        else {
1257            return Ok(PromptDisposition::Handled);
1258        };
1259        self.queue(behavior, self.expand(&text), images);
1260        Ok(PromptDisposition::Queued)
1261    }
1262
1263    /// Queues a message to steer the current run after its current tool calls.
1264    pub fn steer(&self, text: &str, images: Vec<ImageContent>) {
1265        self.queue(StreamingBehavior::Steer, self.expand(text), images);
1266    }
1267
1268    /// Queues a message for when the current run would otherwise end.
1269    pub fn follow_up(&self, text: &str, images: Vec<ImageContent>) {
1270        self.queue(StreamingBehavior::FollowUp, self.expand(text), images);
1271    }
1272
1273    fn queue(&self, behavior: StreamingBehavior, text: String, images: Vec<ImageContent>) {
1274        self.push_queued(behavior, text, images);
1275        self.emit_queue_update();
1276    }
1277
1278    /// [`AgentSession::queue`] while a run is active, in one step with the
1279    /// check (see [`State::streaming`]); otherwise the message back.
1280    fn queue_streaming(
1281        &self,
1282        behavior: StreamingBehavior,
1283        text: String,
1284        images: Vec<ImageContent>,
1285    ) -> Result<(), (String, Vec<ImageContent>)> {
1286        let state = lock(&self.inner.state);
1287        if !state.streaming {
1288            return Err((text, images));
1289        }
1290        self.push_queued(behavior, text, images);
1291        drop(state);
1292        self.emit_queue_update();
1293        Ok(())
1294    }
1295
1296    fn push_queued(&self, behavior: StreamingBehavior, text: String, images: Vec<ImageContent>) {
1297        let mut queued = lock(&self.inner.queued);
1298        let (texts, queue) = match behavior {
1299            StreamingBehavior::Steer => (&mut queued.0, &self.inner.steering),
1300            StreamingBehavior::FollowUp => (&mut queued.1, &self.inner.follow_up),
1301        };
1302        texts.push(text.clone());
1303        lock(queue).push_back(Self::user_message(text, images));
1304    }
1305
1306    fn expand(&self, text: &str) -> String {
1307        let text = self.expand_skill(text);
1308        expand_prompt_template(&text, &read(&self.inner.resources).templates)
1309    }
1310
1311    /// `/skill:name args` becomes the skill's content with the arguments. The
1312    /// name ends at the first space. A skill file that cannot be read is
1313    /// reported as an extension error and leaves the text as typed.
1314    fn expand_skill(&self, text: &str) -> String {
1315        let Some(rest) = text.strip_prefix("/skill:") else {
1316            return text.to_owned();
1317        };
1318        let (name, args) = match rest.split_once(' ') {
1319            Some((name, args)) => (name, args.trim()),
1320            None => (rest, ""),
1321        };
1322        let Some(skill) = read(&self.inner.resources)
1323            .skills
1324            .iter()
1325            .find(|skill| skill.name == name)
1326            .cloned()
1327        else {
1328            return text.to_owned();
1329        };
1330        let content = match std::fs::read_to_string(&skill.file_path) {
1331            Ok(content) => content,
1332            Err(error) => {
1333                let error = crate::tools::node_error(&error, "open", &skill.file_path);
1334                let ui = lock(&self.inner.binding).0.clone();
1335                ui.extension_error(
1336                    &skill.file_path.to_string_lossy(),
1337                    "skill_expansion",
1338                    &error,
1339                    None,
1340                );
1341                return text.to_owned();
1342            }
1343        };
1344        let (_, body) = crate::resources::split_frontmatter(&content);
1345        let mut expanded = format!(
1346            "<skill name=\"{}\" location=\"{}\">\nReferences are relative to {}.\n\n{}\n</skill>",
1347            skill.name,
1348            skill.file_path.display(),
1349            skill.base_dir.display(),
1350            body.trim()
1351        );
1352        if !args.is_empty() {
1353            expanded += &format!("\n\n{args}");
1354        }
1355        expanded
1356    }
1357
1358    /// pi's base prompt options, which `ctx.getSystemPromptOptions()` gives
1359    /// commands as a live object: the session's resources, the active tools,
1360    /// and the snippet and guidelines of every registered tool, without what
1361    /// a run adds, as a command last edited them.
1362    pub fn base_prompt_options(&self) -> SystemPromptOptions {
1363        let base = self.rebuilt_prompt_options();
1364        match &*lock(&self.inner.edited_base) {
1365            Some((edited_from, edited)) if *edited_from == base => edited.clone(),
1366            _ => base,
1367        }
1368    }
1369
1370    /// Keeps a command's edits to the base prompt options until the base
1371    /// changes, as pi's live options keep them until it rebuilds them.
1372    pub fn set_base_prompt_options(&self, options: SystemPromptOptions) {
1373        *lock(&self.inner.edited_base) = Some((self.rebuilt_prompt_options(), options));
1374    }
1375
1376    /// The base prompt options as pi's `_rebuildSystemPrompt` builds them.
1377    /// Tools a loadout hides from requests have no snippet.
1378    fn rebuilt_prompt_options(&self) -> SystemPromptOptions {
1379        let resources = self.resources();
1380        let selected_tools = self.inner.tools.active();
1381        let (_, changes) = self.loadout_changes(self.active_tools(&selected_tools));
1382        let mut options = SystemPromptOptions {
1383            custom_prompt: resources.custom_prompt,
1384            selected_tools,
1385            append_system_prompt: resources.append_prompt.unwrap_or_default(),
1386            cwd: self.inner.cwd.clone(),
1387            context_files: resources.context_files,
1388            skills: resources.skills,
1389            ..SystemPromptOptions::default()
1390        };
1391        for tool in self.inner.tools.with(|registry| registry.all()) {
1392            let name = tool.name().to_owned();
1393            if let Some(snippet) = tool
1394                .snippet
1395                .filter(|snippet| !snippet.is_empty() && !changes.hidden.contains(&name))
1396            {
1397                options.tool_snippets.insert(name.clone(), snippet);
1398            }
1399            if !tool.guidelines.is_empty() {
1400                options.tool_guidelines.insert(name, tool.guidelines);
1401            }
1402        }
1403        options
1404    }
1405
1406    /// The prompt's options for `active` tools: the current run's, as
1407    /// `before_agent_start` left them, with the snippets and guidelines of
1408    /// active tools it did not set.
1409    fn prompt_options(&self, active: &[String]) -> SystemPromptOptions {
1410        let mut options = lock(&self.inner.run_options)
1411            .clone()
1412            .unwrap_or_else(|| self.base_prompt_options());
1413        options.selected_tools = active.to_vec();
1414        let declared = self.active_tools(active);
1415        for tool in &declared {
1416            let name = tool.tool.declaration().name.clone();
1417            if let Some(snippet) = &tool.snippet {
1418                options
1419                    .tool_snippets
1420                    .entry(name.clone())
1421                    .or_insert_with(|| snippet.clone());
1422            }
1423            if !tool.guidelines.is_empty() {
1424                options
1425                    .tool_guidelines
1426                    .entry(name)
1427                    .or_insert_with(|| tool.guidelines.clone());
1428            }
1429        }
1430        // The prompt lists only the tools requests declare.
1431        let (_, changes) = self.loadout_changes(declared);
1432        options
1433            .tool_snippets
1434            .retain(|name, _| !changes.hidden.contains(name));
1435        options
1436    }
1437
1438    /// Shows the `declared` tools to the extensions' loadout hooks that
1439    /// answer later, before a request presents them.
1440    pub(super) async fn refresh_loadouts(&self, declared: Vec<RegisteredTool>) {
1441        let loadout = Loadout {
1442            declared,
1443            callable: self.callable_tools(),
1444        };
1445        for extension in &self.inner.extensions {
1446            extension.refresh_loadout(&loadout).await;
1447        }
1448    }
1449
1450    /// The extensions' loadout hooks' changes for `declared` tools, with the
1451    /// loadout they saw.
1452    fn loadout_changes(&self, declared: Vec<RegisteredTool>) -> (Loadout, LoadoutChanges) {
1453        let loadout = Loadout {
1454            declared,
1455            callable: self.callable_tools(),
1456        };
1457        let mut changes = LoadoutChanges::default();
1458        for extension in &self.inner.extensions {
1459            let next = extension.prepare_loadout(&loadout);
1460            // The descriptions and declarations are the model's context.
1461            let changed = !next.descriptions.is_empty() || !next.hidden.is_empty();
1462            if changed
1463                && !self.applies(
1464                    extension,
1465                    "prepareLoadout",
1466                    extension.grants().session,
1467                    extensions::SESSION,
1468                )
1469            {
1470                continue;
1471            }
1472            changes.descriptions.extend(next.descriptions);
1473            changes.hidden.extend(next.hidden);
1474        }
1475        (loadout, changes)
1476    }
1477
1478    /// pi's `_applyToolLoadout`: the `declared` tools as the model sees
1479    /// them, with the descriptions the extensions' loadout hooks give them.
1480    /// Records the tools the hooks hide from requests.
1481    fn loadout(&self, declared: Vec<RegisteredTool>) -> Vec<Arc<dyn Tool>> {
1482        let (loadout, changes) = self.loadout_changes(declared);
1483        let mut descriptions = changes.descriptions;
1484        *lock(&self.inner.hidden_tools) = changes.hidden;
1485        loadout
1486            .declared
1487            .into_iter()
1488            .map(|tool| match descriptions.remove(tool.name()) {
1489                Some(description) if description != tool.tool.declaration().description => {
1490                    Arc::new(Described::new(tool.tool, description)) as Arc<dyn Tool>
1491                }
1492                _ => tool.tool,
1493            })
1494            .collect()
1495    }
1496
1497    fn active_tools(&self, active: &[String]) -> Vec<RegisteredTool> {
1498        self.inner.tools.with(|registry| {
1499            active
1500                .iter()
1501                .filter_map(|name| registry.get(name).cloned())
1502                .collect()
1503        })
1504    }
1505
1506    /// The system message patch that brings the transcript's prompt up to date.
1507    fn system_update(
1508        &self,
1509        messages: &[Message],
1510        active: &[String],
1511    ) -> Result<Option<Message>, String> {
1512        let sections = build_sections(&self.prompt_options(active), &self.inner.docs)?;
1513        let current = yapi_ai::transcript::current_system_message(messages)
1514            .and_then(|system| system.sections)
1515            .unwrap_or_default();
1516        Ok(diff_sections(&current, &sections).map(|patch| {
1517            Message::System(SystemMessage {
1518                content: Content::Text(String::new()),
1519                sections: Some(patch),
1520                timestamp: now_ms(),
1521                tools_added: None,
1522                tools_removed: None,
1523            })
1524        }))
1525    }
1526
1527    /// The system prompt for `active` tools as the model reads it.
1528    fn system_prompt_text(&self, active: &[String]) -> Result<String, String> {
1529        self.prompt_text(&self.prompt_options(active))
1530    }
1531
1532    /// pi's `systemPrompt`: the prompt of the current run's options, as
1533    /// the run sends it, or of the base options while no run is active.
1534    pub fn system_prompt(&self) -> Result<String, String> {
1535        if !self.is_streaming() {
1536            return self.prompt_text(&self.base_prompt_options());
1537        }
1538        if let Some(forced) = lock(&self.inner.forced_prompt).clone() {
1539            return Ok(forced);
1540        }
1541        self.system_prompt_text(&self.inner.tools.active())
1542    }
1543
1544    /// pi's `buildSystemPrompt`: the system prompt `options` give, or their
1545    /// forced prompt if they have one.
1546    pub fn render_prompt(&self, options: &SystemPromptOptions) -> Result<String, String> {
1547        match &options.force_system_prompt {
1548            Some(forced) => Ok(forced.clone()),
1549            None => self.prompt_text(options),
1550        }
1551    }
1552
1553    /// The system prompt `options` describe, as the model reads it.
1554    fn prompt_text(&self, options: &SystemPromptOptions) -> Result<String, String> {
1555        let sections = build_sections(options, &self.inner.docs)?;
1556        Ok(SystemMessage {
1557            content: Content::Text(String::new()),
1558            sections: Some(
1559                sections
1560                    .into_iter()
1561                    .map(|(name, text)| (name, Some(text)))
1562                    .collect(),
1563            ),
1564            timestamp: 0,
1565            tools_added: None,
1566            tools_removed: None,
1567        }
1568        .text())
1569    }
1570
1571    /// [`AgentSession::prompt_with`] while a run streams, when the message
1572    /// has a streaming behavior and neither an `input` handler nor an
1573    /// extension command takes it. It is queued before this returns, as pi
1574    /// queues it before any later event. Returns whether it was.
1575    pub fn queue_now(&self, text: &str, options: &PromptOptions) -> bool {
1576        let Some(behavior) = options.streaming_behavior.filter(|_| self.is_streaming()) else {
1577            return false;
1578        };
1579        let expand = options.expand_prompt_templates;
1580        if self.has_handlers("input")
1581            || (expand && self.is_extension_command(text))
1582            || self.inner.manual_compaction.load(Ordering::SeqCst)
1583        {
1584            return false;
1585        }
1586        let text = if expand {
1587            self.expand(text)
1588        } else {
1589            text.to_owned()
1590        };
1591        self.queue_streaming(behavior, text, options.images.clone())
1592            .is_ok()
1593    }
1594
1595    /// Sends a user prompt and runs until the agent settles. Templates and
1596    /// `/skill:` commands expand first.
1597    pub async fn prompt(&self, text: &str, images: Vec<ImageContent>) -> Result<(), String> {
1598        let options = PromptOptions {
1599            images,
1600            ..PromptOptions::default()
1601        };
1602        self.prompt_with(text, options, |_| {}).await
1603    }
1604
1605    /// pi's `prompt`: while a run streams, the options' streaming behavior
1606    /// queues the message, and without one it is an error. Otherwise the
1607    /// message starts a run that lasts until the agent settles. `preflight`
1608    /// learns which, before any event of the run. A prompt started while
1609    /// `agent_settled` handlers run waits until they return.
1610    pub async fn prompt_with(
1611        &self,
1612        text: &str,
1613        options: PromptOptions,
1614        preflight: impl FnOnce(PromptDisposition),
1615    ) -> Result<(), String> {
1616        let _turn = match self.settled_turn() {
1617            Some(turn) => turn.await,
1618            None => None,
1619        };
1620        let PromptOptions {
1621            expand_prompt_templates: expand,
1622            streaming_behavior: behavior,
1623            images,
1624            source,
1625        } = options;
1626        if expand && self.run_extension_command(text).await {
1627            preflight(PromptDisposition::Handled);
1628            return Ok(());
1629        }
1630        if self.inner.manual_compaction.load(Ordering::SeqCst) {
1631            return Err("Cannot submit a prompt while compaction is in progress. Wait for compaction to finish and retry.".into());
1632        }
1633        let streaming = behavior.filter(|_| self.is_streaming());
1634        let Some((text, images)) = self.input_handlers(text, images, source, streaming).await
1635        else {
1636            preflight(PromptDisposition::Handled);
1637            return Ok(());
1638        };
1639        let expanded = if expand { self.expand(&text) } else { text };
1640        let (expanded, images) = match behavior {
1641            Some(behavior) => match self.queue_streaming(behavior, expanded, images) {
1642                Ok(()) => {
1643                    preflight(PromptDisposition::Queued);
1644                    return Ok(());
1645                }
1646                Err(message) => message,
1647            },
1648            None if self.is_streaming() => {
1649                return Err("Agent is already processing. Specify streamingBehavior ('steer' or 'followUp') to queue the message.".into());
1650            }
1651            None => (expanded, images),
1652        };
1653        self.flush_pending_bash();
1654        self.flush_pending_custom();
1655        // Without a model pi's agent holds a placeholder from provider
1656        // "unknown", which has no credential.
1657        let model = lock(&self.inner.state)
1658            .model
1659            .clone()
1660            .ok_or_else(|| crate::auth_guidance::no_api_key_found("unknown", &self.inner.docs))?;
1661        let registry = self.registry();
1662        if let Some(error) = registry.store_error(&model.provider) {
1663            return Err(error);
1664        }
1665        let mut has_auth = registry.has_auth(&model.provider);
1666        // pi checks a key a command supplies by running it: a command that
1667        // fails or prints nothing leaves the provider without a key.
1668        if has_auth && registry.uses_command_key(&model.provider) {
1669            let auth = registry.auth(&model).await;
1670            if let Some(error) = auth.error {
1671                return Err(error);
1672            }
1673            has_auth = auth.source.is_some();
1674        }
1675        if !has_auth {
1676            return Err(crate::auth_guidance::no_api_key_found(
1677                &model.provider,
1678                &self.inner.docs,
1679            ));
1680        }
1681        self.compact_before_prompt().await?;
1682        let mut sections = self.base_prompt_options().sections;
1683        let ctx = self.extension_context(CancellationToken::new());
1684        for extension in &self.inner.extensions {
1685            extension.before_agent_start(&ctx, &mut sections).await;
1686        }
1687        // After the extensions, which may have registered tools.
1688        *lock(&self.inner.run_options) = Some(SystemPromptOptions {
1689            sections,
1690            ..self.base_prompt_options()
1691        });
1692        let (custom, forced) = if self.has_handlers("before_agent_start") {
1693            let options = self.prompt_options(&self.inner.tools.active());
1694            self.before_agent_start_handlers(&expanded, &images, options)
1695                .await
1696        } else {
1697            let forced = lock(&self.inner.run_options)
1698                .as_mut()
1699                .and_then(|options| options.force_system_prompt.take());
1700            (Vec::new(), forced)
1701        };
1702        *lock(&self.inner.forced_prompt) = forced;
1703        // Handlers may have changed the active tools.
1704        let active = self.inner.tools.active();
1705        // After the handlers, so a model they select sets the resize profile.
1706        let (images, hints) = if images.is_empty() {
1707            (images, Vec::new())
1708        } else {
1709            let (auto_resize, limits) = self.image_options();
1710            tokio::task::spawn_blocking(move || {
1711                crate::images::normalize_prompt(images, auto_resize, limits.as_ref())
1712            })
1713            .await
1714            .map_err(|err| err.to_string())?
1715        };
1716        let expanded = if hints.is_empty() {
1717            expanded
1718        } else {
1719            format!("{expanded}\n\n{}", hints.join("\n"))
1720        };
1721        let messages = self.messages();
1722        // The prompt leaves out the snippets of tools the hooks hide.
1723        self.refresh_loadouts(self.active_tools(&active)).await;
1724        let mut prompts = Vec::new();
1725        if let Some(update) = self.system_update(&messages, &active)? {
1726            prompts.push(update);
1727        }
1728        prompts.push(Self::user_message(expanded, images));
1729        prompts.append(&mut lock(&self.inner.next_turn));
1730        prompts.extend(custom);
1731        preflight(PromptDisposition::Started);
1732        self.run(prompts).await
1733    }
1734
1735    /// Runs prompts, then pi's post-run loop: retries, overflow recovery and
1736    /// compaction, and queued messages, until the agent settles. An invalid
1737    /// compaction setting ends the loop with its error once the agent settles.
1738    async fn run(&self, prompts: Vec<Message>) -> Result<(), String> {
1739        let cancel = CancellationToken::new();
1740        *lock(&self.inner.cancel) = Some(cancel.clone());
1741        lock(&self.inner.state).streaming = true;
1742        self.run_agent(Some(prompts), &cancel).await;
1743        let mut result = Ok(());
1744        loop {
1745            self.settle_run(&cancel, &mut result).await;
1746            if cancel.is_cancelled() {
1747                self.finish_cancelled_retry();
1748            }
1749            #[cfg(test)]
1750            race(&self.inner.races.before_end);
1751            if self.end_run(cancel.is_cancelled() || result.is_err()) {
1752                break;
1753            }
1754            // A message queued since the last look at the queues.
1755            self.run_agent(None, &cancel).await;
1756        }
1757        let settling = Settling(&self.inner);
1758        self.flush_pending_bash();
1759        self.flush_pending_custom();
1760        *lock(&self.inner.cancel) = None;
1761        self.cache_warming_settled();
1762        // Extensions hear it first, as in pi.
1763        self.emit_extension_event(
1764            &ExtensionEvent::AgentSettled.to_value(),
1765            CancellationToken::new(),
1766        )
1767        .await;
1768        self.emit(&AgentEvent::AgentSettled);
1769        for turn in settling.finish() {
1770            let (guard, ended) = tokio::sync::oneshot::channel();
1771            if turn.send(guard).is_ok() {
1772                // Fails once the prompt drops the guard as it ends.
1773                let _ = ended.await;
1774            }
1775        }
1776        self.inner.idle.notify_waiters();
1777        result
1778    }
1779
1780    /// pi's post-run loop until the agent would settle: retries, overflow
1781    /// recovery and compaction, the queued messages, and the
1782    /// `agent_before_settle` boundary. An invalid compaction setting ends
1783    /// it with its error in `result`.
1784    async fn settle_run(&self, cancel: &CancellationToken, result: &mut Result<(), String>) {
1785        while !cancel.is_cancelled() {
1786            let again = match self.handle_post_run(cancel).await {
1787                Ok(again) => again,
1788                Err(error) => {
1789                    *result = Err(error);
1790                    break;
1791                }
1792            };
1793            if again {
1794                if cancel.is_cancelled() {
1795                    break;
1796                }
1797                self.run_agent(None, cancel).await;
1798                continue;
1799            }
1800            if cancel.is_cancelled() || !self.before_settle(cancel).await {
1801                break;
1802            }
1803            self.run_agent(None, cancel).await;
1804        }
1805    }
1806
1807    /// Ends the run, as pi's `_emitAgentSettled` starts: no longer active,
1808    /// and settling. Unless `stopped`, only when no message was queued since
1809    /// the run last looked, in one step with queueing (see
1810    /// [`State::streaming`]). Whether it ended.
1811    fn end_run(&self, stopped: bool) -> bool {
1812        let mut state = lock(&self.inner.state);
1813        if !stopped && self.has_queued() {
1814            return false;
1815        }
1816        state.streaming = false;
1817        state.settling = true;
1818        true
1819    }
1820
1821    fn has_queued(&self) -> bool {
1822        !lock(&self.inner.steering).is_empty() || !lock(&self.inner.follow_up).is_empty()
1823    }
1824
1825    /// pi's `buildRequestOptions`: the settings every provider request
1826    /// carries, compaction and summaries included, or the error an invalid
1827    /// one raises.
1828    pub(super) fn request_options(&self) -> Result<StreamOptions, String> {
1829        let manager = lock(&self.inner.settings);
1830        let idle = manager.http_idle_timeout_ms()?;
1831        let provider = manager
1832            .settings()
1833            .retry
1834            .as_ref()
1835            .and_then(|retry| retry.provider.clone())
1836            .unwrap_or_default();
1837        // pi uses the longest timer it can set in place of a disabled one.
1838        let idle = if idle == 0 { 2_147_483_647 } else { idle };
1839        Ok(StreamOptions {
1840            timeout_ms: Some(manager.provider_timeout_ms().unwrap_or(idle as f64)),
1841            websocket_connect_timeout_ms: manager.websocket_connect_timeout_ms()?,
1842            max_retries: provider.max_retries.unwrap_or(0),
1843            max_retry_delay_ms: provider.max_retry_delay_ms,
1844            ..StreamOptions::default()
1845        })
1846    }
1847
1848    /// The options of the run's requests: `request`'s, with the session's
1849    /// transport, thinking budgets and hooks.
1850    fn stream_options(
1851        &self,
1852        model: &Model,
1853        session_id: String,
1854        cancel: &CancellationToken,
1855        request: StreamOptions,
1856    ) -> StreamOptions {
1857        let settings = lock(&self.inner.settings).settings().clone();
1858        let budgets = settings.thinking_budgets.as_ref();
1859        StreamOptions {
1860            session_id: Some(session_id),
1861            transport: Some(settings.transport.unwrap_or(Transport::Auto)),
1862            thinking_budgets: ThinkingBudgets {
1863                minimal: budgets.and_then(|b| b.minimal),
1864                low: budgets.and_then(|b| b.low),
1865                medium: budgets.and_then(|b| b.medium),
1866                high: budgets.and_then(|b| b.high),
1867            },
1868            cancel: cancel.clone(),
1869            hooks: self.request_hooks(model, cancel),
1870            ..request
1871        }
1872    }
1873
1874    /// What extensions see of the run's provider requests to `model`: pi's
1875    /// `before_provider_request`, whose handlers each return the body the
1876    /// next one sees, `before_provider_headers`, `after_provider_response`
1877    /// and `provider_stream_event`.
1878    fn request_hooks(&self, model: &Model, cancel: &CancellationToken) -> RequestHooks {
1879        let model = model.clone();
1880        RequestHooks {
1881            payload: self.hook(
1882                "before_provider_request",
1883                cancel,
1884                |session, payload, cancel| {
1885                    Box::pin(async move {
1886                        session
1887                            .chain("before_provider_request", "payload", payload, cancel)
1888                            .await
1889                    })
1890                },
1891            ),
1892            headers: self.headers_hook(cancel),
1893            response: self.hook(
1894                "after_provider_response",
1895                cancel,
1896                |session, response: yapi_ai::stream::ProviderResponse, cancel| {
1897                    let event = ExtensionEvent::AfterProviderResponse {
1898                        status: response.status,
1899                        headers: &response.headers,
1900                    }
1901                    .to_value();
1902                    Box::pin(async move {
1903                        session.emit_extension_event(&event, cancel).await;
1904                    })
1905                },
1906            ),
1907            stream_event: self.hook(
1908                "provider_stream_event",
1909                cancel,
1910                move |session, data, cancel| {
1911                    let event = serde_json::json!({
1912                        "data": data,
1913                        "type": "provider_stream_event",
1914                        "provider": model.provider,
1915                        "api": model.api,
1916                        "model": model.id,
1917                    });
1918                    Box::pin(async move {
1919                        session.emit_extension_event(&event, cancel).await;
1920                    })
1921                },
1922            ),
1923        }
1924    }
1925
1926    /// pi's `before_provider_headers` for every provider request, compaction
1927    /// and branch summaries included.
1928    pub(super) fn headers_hook(
1929        &self,
1930        cancel: &CancellationToken,
1931    ) -> Option<Hook<RequestHeaders, RequestHeaders>> {
1932        self.hook(
1933            "before_provider_headers",
1934            cancel,
1935            |session, headers, cancel| {
1936                Box::pin(async move { session.before_provider_headers(headers, cancel).await })
1937            },
1938        )
1939    }
1940
1941    /// A request hook that runs `event` on this session, when extensions
1942    /// handle events of type `kind`.
1943    fn hook<A: 'static, R: 'static>(
1944        &self,
1945        kind: &str,
1946        cancel: &CancellationToken,
1947        event: impl Fn(AgentSession, A, CancellationToken) -> BoxFuture<'static, R>
1948        + Send
1949        + Sync
1950        + 'static,
1951    ) -> Option<Hook<A, R>> {
1952        if !self.has_handlers(kind) {
1953            return None;
1954        }
1955        let (session, cancel) = (self.clone(), cancel.clone());
1956        Some(Arc::new(move |argument| {
1957            event(session.clone(), argument, cancel.clone())
1958        }))
1959    }
1960
1961    /// One agent run: with `prompts`, or continuing the transcript (or running
1962    /// queued messages when it ends with an assistant message).
1963    async fn run_agent(&self, prompts: Option<Vec<Message>>, cancel: &CancellationToken) {
1964        let active = self.inner.tools.active();
1965        let (model, thinking_level, messages, session_id, session_file) = {
1966            let state = lock(&self.inner.state);
1967            (
1968                state.model.clone(),
1969                state.thinking_level,
1970                state.session.build_context().messages,
1971                state.session.id().to_owned(),
1972                state.session.file().map(Path::to_path_buf),
1973            )
1974        };
1975        let Some(model) = model else {
1976            return;
1977        };
1978        let settings = lock(&self.inner.settings).settings().clone();
1979        *write(&self.inner.runtime) = Runtime {
1980            model: Some(model.clone()),
1981            thinking_level: Some(thinking_level),
1982            session_id: Some(session_id.clone()),
1983            session_file,
1984            auto_resize_images: settings.image_auto_resize(),
1985            shell_path: settings.shell_path.clone(),
1986            shell_command_prefix: settings.shell_command_prefix.clone(),
1987        };
1988        let apis = self.apis();
1989        let session = self.downgrade();
1990        let request = self.request_options();
1991        // pi's stream function throws for an invalid setting, which ends the run.
1992        let invalid = request.as_ref().err().cloned();
1993        let config = LoopConfig {
1994            options: self.stream_options(&model, session_id, cancel, request.unwrap_or_default()),
1995            model,
1996            thinking_level,
1997            stream: Arc::new(move |request| {
1998                if let Some(error) = &invalid {
1999                    return Err(error.clone());
2000                }
2001                if let Some(session) = session.upgrade() {
2002                    session.start_cache_warming(&request);
2003                }
2004                Ok(apis.stream(request))
2005            }),
2006            tool_execution: ExecutionMode::Parallel,
2007        };
2008        let mut context = AgentContext {
2009            messages,
2010            tools: self.loadout(self.active_tools(&active)),
2011        };
2012        let cancel_token = cancel.clone();
2013        let hooks = Hooks {
2014            session: self.clone(),
2015            cancel: cancel_token,
2016            turn_index: std::sync::atomic::AtomicU64::new(0),
2017        };
2018        let prompts = match prompts {
2019            Some(prompts) => Some(prompts),
2020            None if matches!(context.messages.last(), Some(Message::Assistant(_))) => {
2021                let (steering_all, follow_up_all) = self.drains_all();
2022                let steering = drain(&self.inner.steering, steering_all);
2023                let queued = if steering.is_empty() {
2024                    drain(&self.inner.follow_up, follow_up_all)
2025                } else {
2026                    steering
2027                };
2028                if queued.is_empty() {
2029                    return;
2030                }
2031                Some(queued)
2032            }
2033            None => None,
2034        };
2035        match prompts {
2036            Some(prompts) => {
2037                yapi_agent::run(prompts, &mut context, config, &hooks).await;
2038            }
2039            None => {
2040                let _ = yapi_agent::run_continue(&mut context, config, &hooks).await;
2041            }
2042        }
2043    }
2044
2045    /// Moves to another point of the session tree, as `/tree` does. A user or
2046    /// custom message target puts the leaf on its parent and returns its text for
2047    /// the editor. With `summarize`, the abandoned branch is summarized at the new
2048    /// position. `label` labels the summary, or the target without one.
2049    pub async fn navigate_tree(
2050        &self,
2051        target_id: &str,
2052        options: TreeNavigation,
2053    ) -> Result<TreeOutcome, String> {
2054        if self.is_streaming() {
2055            return Err(
2056                "Wait for the current response to finish before navigating the session tree."
2057                    .into(),
2058            );
2059        }
2060        let (old_leaf, target) = self.with_session(|session| {
2061            (
2062                session.leaf_id().map(str::to_owned),
2063                session.entry(target_id).cloned(),
2064            )
2065        });
2066        if old_leaf.as_deref() == Some(target_id) {
2067            return Ok(TreeOutcome::default());
2068        }
2069        let model = self.model();
2070        if options.summarize && model.is_none() {
2071            return Err("No model available for summarization".into());
2072        }
2073        let Some(target) = target else {
2074            return Err(format!("Entry {target_id} not found"));
2075        };
2076        let (entries, common) = self.with_session(|session| {
2077            crate::compaction::collect_branch_entries(session, old_leaf.as_deref(), target_id)
2078        });
2079        let mut options = options;
2080        let mut summary = None;
2081        let mut from_extension = false;
2082        if self.has_handlers("session_before_tree") {
2083            let event = ExtensionEvent::SessionBeforeTree {
2084                preparation: TreePreparation {
2085                    target_id,
2086                    old_leaf_id: old_leaf.as_deref(),
2087                    common_ancestor_id: common.as_deref(),
2088                    entries_to_summarize: &entries,
2089                    user_wants_summary: options.summarize,
2090                    custom_instructions: options.custom_instructions.as_deref(),
2091                    replace_instructions: options.replace_instructions,
2092                    label: options.label.as_deref(),
2093                },
2094            };
2095            if let Some(result) = self
2096                .emit_extension_event(&event.to_value(), CancellationToken::new())
2097                .await
2098            {
2099                if result["cancel"] == true {
2100                    return Ok(TreeOutcome {
2101                        cancelled: true,
2102                        ..TreeOutcome::default()
2103                    });
2104                }
2105                if options.summarize
2106                    && let Some(text) = result["summary"]["summary"]
2107                        .as_str()
2108                        .filter(|text| !text.is_empty())
2109                {
2110                    let details = Some(result["summary"]["details"].clone())
2111                        .filter(|details| !details.is_null());
2112                    let usage = serde_json::from_value(result["summary"]["usage"].clone()).ok();
2113                    summary = Some((text.to_owned(), details, usage));
2114                    from_extension = true;
2115                }
2116                if let Some(instructions) = result["customInstructions"].as_str() {
2117                    options.custom_instructions = Some(instructions.to_owned());
2118                }
2119                if let Some(replace) = result["replaceInstructions"].as_bool() {
2120                    options.replace_instructions = replace;
2121                }
2122                if let Some(label) = result["label"].as_str() {
2123                    options.label = Some(label.to_owned());
2124                }
2125            }
2126        }
2127        if options.summarize
2128            && !entries.is_empty()
2129            && summary.is_none()
2130            && let Some(model) = &model
2131        {
2132            let cancel = CancellationToken::new();
2133            *lock(&self.inner.cancel) = Some(cancel.clone());
2134            let compacting = Compacting::start(&self.inner.compacting);
2135            let summarizer = self
2136                .summarizer(
2137                    model,
2138                    ThinkingLevel::Off,
2139                    SummarySource::BranchSummary,
2140                    None,
2141                    cancel,
2142                )
2143                .await;
2144            let reserve = lock(&self.inner.settings)
2145                .settings()
2146                .branch_summary
2147                .as_ref()
2148                .and_then(|settings| settings.reserve_tokens)
2149                .unwrap_or(16384);
2150            let result = summarizer
2151                .branch_summary(
2152                    &entries,
2153                    options.custom_instructions.as_deref(),
2154                    options.replace_instructions,
2155                    reserve,
2156                )
2157                .await;
2158            *lock(&self.inner.cancel) = None;
2159            drop(compacting);
2160            match result? {
2161                BranchSummary::Aborted => {
2162                    return Ok(TreeOutcome {
2163                        cancelled: true,
2164                        aborted: true,
2165                        ..TreeOutcome::default()
2166                    });
2167                }
2168                BranchSummary::Done {
2169                    summary: text,
2170                    usage,
2171                    read_files,
2172                    modified_files,
2173                } => {
2174                    summary = Some((
2175                        text,
2176                        Some(
2177                            serde_json::json!({"readFiles": read_files, "modifiedFiles": modified_files}),
2178                        ),
2179                        usage,
2180                    ));
2181                }
2182            }
2183        }
2184        let (new_leaf, editor_text) = match &target {
2185            FileEntry::Message(yapi_types::session::MessageEntry {
2186                meta,
2187                message: Message::User(user),
2188            }) => (meta.parent_id.clone(), Some(user.content.text(""))),
2189            FileEntry::CustomMessage(entry) => {
2190                (entry.meta.parent_id.clone(), Some(entry.content.text("")))
2191            }
2192            _ => (Some(target_id.to_owned()), None),
2193        };
2194        let summary_entry = self.with_session(|session| -> Result<Option<FileEntry>, String> {
2195            let mut summary_entry = None;
2196            match summary {
2197                Some((text, details, usage)) => {
2198                    let id = session
2199                        .branch_with_summary(
2200                            new_leaf.as_deref(),
2201                            text,
2202                            details,
2203                            Some(from_extension),
2204                            usage,
2205                        )
2206                        .map_err(|err| err.to_string())?;
2207                    summary_entry = session.entry(&id).cloned();
2208                    if let Some(label) = &options.label {
2209                        session
2210                            .append_label(&id, Some(label.clone()))
2211                            .map_err(|err| err.to_string())?;
2212                    }
2213                }
2214                None => {
2215                    match &new_leaf {
2216                        None => session.reset_leaf(),
2217                        Some(id) => session.branch(id).map_err(|err| err.to_string())?,
2218                    }
2219                    if let Some(label) = &options.label {
2220                        session
2221                            .append_label(target_id, Some(label.clone()))
2222                            .map_err(|err| err.to_string())?;
2223                    }
2224                }
2225            }
2226            Ok(summary_entry)
2227        })?;
2228        self.restore_tools_from_transcript();
2229        // pi's `session_tree` event, after the leaf has moved.
2230        let new_leaf = self.with_session(|session| session.leaf_id().map(str::to_owned));
2231        let event = ExtensionEvent::SessionTree {
2232            new_leaf_id: new_leaf.as_deref(),
2233            old_leaf_id: old_leaf.as_deref(),
2234            summary_entry: summary_entry.as_ref(),
2235            from_extension: summary_entry.is_some().then_some(from_extension),
2236        };
2237        self.emit_extension_event(&event.to_value(), CancellationToken::new())
2238            .await;
2239        Ok(TreeOutcome {
2240            editor_text,
2241            cancelled: false,
2242            aborted: false,
2243            summary_entry,
2244        })
2245    }
2246
2247    /// Activates the tools the current branch's transcript declares, if any.
2248    fn restore_tools_from_transcript(&self) {
2249        let messages = self.messages();
2250        let Some(system) = yapi_ai::transcript::current_system_message(&messages) else {
2251            return;
2252        };
2253        let names: Vec<String> = system
2254            .tools_added
2255            .iter()
2256            .flatten()
2257            .map(|tool| tool.name.clone())
2258            .collect();
2259        self.inner.tools.with(|registry| registry.restore(names));
2260    }
2261
2262    /// The session header line as JSON mode prints it.
2263    pub fn header_json(&self) -> Option<String> {
2264        self.with_session(|session| {
2265            session.header().and_then(|header| {
2266                yapi_types::json::to_string(&yapi_types::session::FileEntry::Session(
2267                    header.clone(),
2268                ))
2269                .ok()
2270            })
2271        })
2272    }
2273
2274    /// Sets the active tools, ignoring unknown and hidden names.
2275    pub fn set_active_tools(&self, names: Vec<String>) {
2276        self.inner.tools.set_active(names);
2277    }
2278
2279    /// The active tool names.
2280    pub fn active_tool_names(&self) -> Vec<String> {
2281        self.inner.tools.active()
2282    }
2283
2284    /// What pi's HTML export takes from a live session: the system prompt as
2285    /// the model reads it, and the active tools' names, descriptions and
2286    /// parameter schemas.
2287    pub fn export_context(&self) -> (Option<String>, Vec<serde_json::Value>) {
2288        let active = self.inner.tools.active();
2289        let all = self.inner.tools.all();
2290        let tools = active
2291            .iter()
2292            .filter_map(|name| all.iter().find(|tool| tool.name == *name))
2293            .map(|tool| {
2294                serde_json::json!({
2295                    "name": tool.name,
2296                    "description": tool.description,
2297                    "parameters": tool.parameters,
2298                })
2299            })
2300            .collect();
2301        (self.system_prompt_text(&active).ok(), tools)
2302    }
2303
2304    /// The session's tools.
2305    pub fn tools(&self) -> &Tools {
2306        &self.inner.tools
2307    }
2308
2309    /// pi's `appendEntry` for extensions: appends custom entry `custom_type`
2310    /// and reports it as `entry_appended`.
2311    pub fn append_custom_entry(
2312        &self,
2313        custom_type: &str,
2314        data: Option<Value>,
2315    ) -> Result<(), String> {
2316        let entry = self.with_session(|file| {
2317            let id = file
2318                .append_custom_entry(custom_type, data)
2319                .map_err(|err| err.to_string())?;
2320            Ok::<_, String>(file.entry(&id).cloned())
2321        })?;
2322        if let Some(entry) = entry {
2323            self.emit(&AgentEvent::EntryAppended { entry });
2324        }
2325        Ok(())
2326    }
2327
2328    /// The tools other tools may call, in registration order, as pi's
2329    /// `_getCallableTools`: every `codemode` and `deferred` tool and the
2330    /// active `direct` ones.
2331    pub fn callable_tools(&self) -> Vec<RegisteredTool> {
2332        self.inner.tools.with(|registry| {
2333            let active = registry.active();
2334            registry
2335                .all()
2336                .into_iter()
2337                .filter(|tool| match tool.exposure {
2338                    Exposure::Codemode | Exposure::Deferred => true,
2339                    Exposure::Direct => active.iter().any(|name| name == tool.name()),
2340                    Exposure::ModelOnly | Exposure::Hidden => false,
2341                })
2342                .collect()
2343        })
2344    }
2345
2346    /// Runs tool `name` on behalf of the tool call `caller`, as pi's
2347    /// `ctx.executeTool()`: the call gets the id `<caller>/<n>`, runs through
2348    /// the tool pipeline and hooks against [`AgentSession::callable_tools`],
2349    /// emits `tool_execution_*` events with `parentToolCallId`, and is recorded
2350    /// on the caller's tool result message. Failures come back as error
2351    /// results. `updates` also receives the tool's partial results. Only
2352    /// the tools `allowed` accepts are found, as the call resolves its tool.
2353    pub async fn execute_tool(
2354        &self,
2355        caller: &str,
2356        name: &str,
2357        args: Value,
2358        cancel: CancellationToken,
2359        updates: Option<UpdateSink>,
2360        allowed: impl Fn(&RegisteredTool) -> bool + Send,
2361    ) -> ToolCallOutcome {
2362        let mut call = ToolCall {
2363            id: String::new(),
2364            name: name.to_owned(),
2365            arguments: match args {
2366                Value::Object(arguments) => arguments,
2367                _ => serde_json::Map::new(),
2368            },
2369            thought_signature: None,
2370            namespace: None,
2371        };
2372        let nested = &self.inner.nested;
2373        let started = nested.start(caller, &mut call);
2374        let hooks = Hooks {
2375            session: self.clone(),
2376            cancel: cancel.clone(),
2377            turn_index: std::sync::atomic::AtomicU64::new(0),
2378        };
2379        let parent = Some(caller.to_owned());
2380        hooks
2381            .on_event(&AgentEvent::ToolExecutionStart {
2382                tool_call_id: call.id.clone(),
2383                tool_name: call.name.clone(),
2384                args: Value::Object(call.arguments.clone()),
2385                parent_tool_call_id: parent.clone(),
2386            })
2387            .await;
2388        // pi awaits the start event, which lets the caller's queued progress
2389        // update, such as codemode's running row, go out before the call runs.
2390        tokio::task::yield_now().await;
2391        let mut callable = self.callable_tools();
2392        callable.retain(|tool| allowed(tool));
2393        let exclusive = !started.holds_queue
2394            && callable.iter().any(|tool| {
2395                tool.name() == name && tool.tool.execution_mode() == ExecutionMode::Sequential
2396            });
2397        let queue = if exclusive {
2398            Some(nested.queue.clone().lock_owned().await)
2399        } else {
2400            None
2401        };
2402        nested.enter(&started, started.holds_queue || exclusive);
2403        let outcome = match self.last_assistant() {
2404            None => ToolCallOutcome {
2405                call: call.clone(),
2406                result: ToolResult {
2407                    content: vec![ContentBlock::text("No assistant message issued this call")],
2408                    details: Some(Value::Object(serde_json::Map::new())),
2409                    ..ToolResult::default()
2410                },
2411                is_error: true,
2412            },
2413            Some(assistant) => {
2414                let tools: Vec<Arc<dyn Tool>> =
2415                    callable.into_iter().map(|tool| tool.tool).collect();
2416                let messages = self.messages();
2417                let scope = ToolCallScope {
2418                    tools: &tools,
2419                    assistant: &assistant,
2420                    messages: &messages,
2421                    parent: Some(caller),
2422                    cancel: &cancel,
2423                };
2424                let (sender, mut receiver) = tokio::sync::mpsc::unbounded_channel();
2425                let sink: UpdateSink = Arc::new(move |partial: ToolResult| {
2426                    if let Some(updates) = &updates {
2427                        updates(partial.clone());
2428                    }
2429                    let _ = sender.send(partial);
2430                });
2431                let update = |partial_result| AgentEvent::ToolExecutionUpdate {
2432                    tool_call_id: call.id.clone(),
2433                    tool_name: call.name.clone(),
2434                    args: Value::Object(call.arguments.clone()),
2435                    partial_result,
2436                    parent_tool_call_id: parent.clone(),
2437                };
2438                let run = yapi_agent::run_tool_call(&scope, call.clone(), &hooks, sink);
2439                tokio::pin!(run);
2440                loop {
2441                    tokio::select! {
2442                        biased;
2443                        Some(partial) = receiver.recv() => hooks.on_event(&update(partial)).await,
2444                        outcome = &mut run => {
2445                            while let Ok(partial) = receiver.try_recv() {
2446                                hooks.on_event(&update(partial)).await;
2447                            }
2448                            break outcome;
2449                        }
2450                    }
2451                }
2452            }
2453        };
2454        drop(queue);
2455        nested.finish(
2456            started,
2457            outcome.is_error,
2458            &yapi_types::message::blocks_text(&outcome.result.content, "\n"),
2459            outcome.result.usage.as_ref(),
2460        );
2461        hooks
2462            .on_event(&AgentEvent::ToolExecutionEnd {
2463                tool_call_id: call.id.clone(),
2464                tool_name: call.name.clone(),
2465                result: outcome.result.clone(),
2466                is_error: outcome.is_error,
2467                parent_tool_call_id: parent,
2468            })
2469            .await;
2470        outcome
2471    }
2472
2473    /// A handle that does not keep the session alive.
2474    pub fn downgrade(&self) -> WeakSession {
2475        WeakSession(Arc::downgrade(&self.inner))
2476    }
2477
2478    /// The working directory.
2479    pub fn cwd(&self) -> &Path {
2480        &self.inner.cwd
2481    }
2482}
2483
2484/// pi's `ParsedSkillBlock`: a skill invocation as `/skill:<name>` expands it.
2485#[derive(Clone, Debug, PartialEq, Eq)]
2486pub struct SkillBlock {
2487    /// The skill's name.
2488    pub name: String,
2489    /// Where its file is.
2490    pub location: String,
2491    /// The text between the tags.
2492    pub content: String,
2493    /// What the user wrote after the command, if anything.
2494    pub user_message: Option<String>,
2495}
2496
2497/// pi's `parseSkillBlock`: the skill invocation `text` holds, when it is one.
2498pub fn parse_skill_block(text: &str) -> Option<SkillBlock> {
2499    let rest = text.strip_prefix("<skill name=\"")?;
2500    let (name, rest) = rest.split_once('"')?;
2501    let rest = rest.strip_prefix(" location=\"")?;
2502    let (location, rest) = rest.split_once('"')?;
2503    let body = rest.strip_prefix(">\n")?;
2504    if name.is_empty() || location.is_empty() {
2505        return None;
2506    }
2507    // The content ends at the first closing tag that leaves nothing or a
2508    // blank line and a message after it.
2509    let mut from = 0;
2510    while let Some(found) = body[from..].find("\n</skill>") {
2511        let end = from + found;
2512        let after = &body[end + "\n</skill>".len()..];
2513        let user_message = match after {
2514            "" => Some(None),
2515            _ => after
2516                .strip_prefix("\n\n")
2517                .filter(|message| !message.is_empty())
2518                .map(|message| Some(message.trim()).filter(|message| !message.is_empty())),
2519        };
2520        if let Some(user_message) = user_message {
2521            return Some(SkillBlock {
2522                name: name.to_owned(),
2523                location: location.to_owned(),
2524                content: body[..end].to_owned(),
2525                user_message: user_message.map(str::to_owned),
2526            });
2527        }
2528        from = end + 1;
2529    }
2530    None
2531}
2532
2533/// Text of an assistant message's text blocks, as print mode prints it.
2534pub fn assistant_text(message: &yapi_types::message::AssistantMessage) -> String {
2535    yapi_types::message::blocks_text(&message.content, "")
2536}
2537
2538/// Whether an assistant message ended in failure.
2539pub fn failed(message: &yapi_types::message::AssistantMessage) -> bool {
2540    matches!(message.stop_reason, StopReason::Error | StopReason::Aborted)
2541}
2542
2543fn drain(queue: &Mutex<VecDeque<Message>>, all: bool) -> Vec<Message> {
2544    let mut queue = lock(queue);
2545    if all {
2546        queue.drain(..).collect()
2547    } else {
2548        queue.pop_front().into_iter().collect()
2549    }
2550}
2551
2552#[cfg(test)]
2553mod tests {
2554    use super::*;
2555
2556    #[test]
2557    fn parses_skill_blocks() {
2558        let block = parse_skill_block(
2559            "<skill name=\"demo\" location=\"/s/SKILL.md\">\nbody\n</skill>\n\n  go now ",
2560        )
2561        .expect("a skill block");
2562        assert_eq!(block.name, "demo");
2563        assert_eq!(block.location, "/s/SKILL.md");
2564        assert_eq!(block.content, "body");
2565        assert_eq!(block.user_message.as_deref(), Some("go now"));
2566        let block =
2567            parse_skill_block("<skill name=\"a\" location=\"b\">\nx\n</skill>\ny\n</skill>")
2568                .expect("the later closing tag");
2569        assert_eq!(block.content, "x\n</skill>\ny");
2570        assert_eq!(block.user_message, None);
2571        assert!(
2572            parse_skill_block("<skill name=\"a\" location=\"b\">\nx\n</skill> trailing").is_none()
2573        );
2574        assert!(parse_skill_block("hello").is_none());
2575    }
2576
2577    #[test]
2578    fn poisoned_registry_keeps_credentials() {
2579        let model: Model = serde_json::from_value(serde_json::json!({
2580            "id": "m", "name": "M", "api": "faux", "provider": "p", "baseUrl": "",
2581            "reasoning": false, "input": ["text"],
2582            "cost": {"input": 0, "output": 0, "cacheRead": 0, "cacheWrite": 0},
2583            "contextWindow": 1000, "maxTokens": 100,
2584        }))
2585        .expect("a model");
2586        let mut registry = ModelRegistry::builtin();
2587        registry.register_provider("p", vec![model.clone()]);
2588        registry.set_runtime_key("p", "key".into());
2589        let session = AgentSession::new(SessionConfig {
2590            cwd: PathBuf::from("/work"),
2591            agent_dir: PathBuf::from("/agent"),
2592            settings: SettingsManager::in_memory(),
2593            registry,
2594            apis: Apis::default(),
2595            session: SessionManager::in_memory(Path::new("/work")),
2596            model: Some(model),
2597            thinking_level: ThinkingLevel::Off,
2598            tools: Vec::new(),
2599            extensions: Vec::new(),
2600            include_extension_tools: false,
2601            allowed_tools: None,
2602            excluded_tools: Vec::new(),
2603            resources: Resources::default(),
2604            docs: crate::docs::Locations::default(),
2605        });
2606        let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
2607            let _slot = session.inner.registry.write();
2608            panic!("poison the registry");
2609        }));
2610        assert!(session.inner.registry.is_poisoned());
2611        assert!(session.registry().has_auth("p"));
2612        assert!(
2613            session
2614                .available_models()
2615                .iter()
2616                .any(|model| model.provider == "p")
2617        );
2618        assert!(session.with_registry(|registry| registry.has_auth("p")));
2619    }
2620}