pi/core/agent_session/
extension_runner.rs

1//! Extension runner seam and session-shared hook handle.
2//!
3//! `AgentSession` never depends on `pi-ext` directly. All extension interaction
4//! goes through [`ExtensionRunner`]. [`NullExtensionRunner`] is the default so
5//! the product layer compiles, unit-tests, and ships before the host exists.
6//!
7//! [`SessionHooks`] is captured by the `pi-agent` tool and next-turn closures at
8//! Agent construction time. Reload swaps the runner under an `RwLock` without
9//! reinstalling those closures. Sync locks are never held across `.await`.
10
11use std::collections::HashMap;
12use std::sync::{Arc, RwLock};
13
14use crate::core::resources::ResourceExtensionPaths;
15use futures::future::BoxFuture;
16use pi_agent::{
17    AfterToolCallContext, AfterToolCallResult, AgentLoopError, AgentLoopTurnUpdate, AgentMessage,
18    AgentTool, AgentToolResult, BeforeToolCallContext, BeforeToolCallResult,
19    PrepareNextTurnContext,
20};
21use pi_ai::{AssistantMessageEvent, Model, ModelThinkingLevel, ToolCall, ToolResultContent};
22use serde_json::{Map, Value};
23use tokio_util::sync::CancellationToken;
24
25use super::events::AgentSessionEvent;
26
27/// Result of an extension `input` transform.
28#[derive(Clone, Debug, Default, PartialEq)]
29pub struct InputTransformResult {
30    /// When true the extension handled the input and `AgentSession` must not run.
31    pub handled: bool,
32    /// Replacement prompt text.
33    pub text: Option<String>,
34    /// Replacement image attachments (opaque JSON until image wiring lands).
35    pub images: Option<Value>,
36}
37
38/// Optional system-prompt / custom-message injection from `before_agent_start`.
39#[derive(Clone, Debug, Default, PartialEq)]
40pub struct BeforeAgentStartResult {
41    /// Custom messages to inject before the user prompt.
42    pub messages: Vec<AgentMessage>,
43    /// Per-turn system prompt override.
44    pub system_prompt: Option<String>,
45}
46
47/// Result of a cancellable extension lifecycle event.
48#[derive(Clone, Debug, Default, PartialEq, Eq)]
49pub struct CancelResult {
50    /// When true the operation must abort.
51    pub cancel: bool,
52    /// Optional human-readable reason.
53    pub reason: Option<String>,
54}
55
56/// Errors produced while dispatching extension hooks.
57#[derive(Debug, thiserror::Error)]
58pub enum ExtensionRunnerError {
59    /// Extension hook failed.
60    #[error("extension error: {0}")]
61    Failed(String),
62    /// Extension host was invalidated after session replacement.
63    #[error("extension context invalidated")]
64    Invalidated,
65}
66
67/// Extension seam consumed by `AgentSession`.
68///
69/// Methods cover every hook this session layer calls. Implementations must be
70/// cheap no-ops when no handlers are registered for a given event name.
71pub trait ExtensionRunner: Send + Sync {
72    /// Returns true when at least one handler is registered for `event`.
73    fn has_handlers(&self, event: &str) -> bool;
74
75    /// Emit a generic session lifecycle event (`agent_start`, turn_*, tool_*, etc.).
76    fn emit(
77        &self,
78        event: AgentSessionEvent,
79    ) -> BoxFuture<'_, Result<Option<CancelResult>, ExtensionRunnerError>>;
80
81    /// Emit a compact streaming assistant delta.
82    ///
83    /// Product runners may override this to avoid serializing the full
84    /// assistant snapshot on every token. The default is a compatibility
85    fn emit_message_update_delta<'a>(
86        &'a self,
87        _event: &'a AssistantMessageEvent,
88    ) -> BoxFuture<'a, Result<Option<CancelResult>, ExtensionRunnerError>> {
89        Box::pin(async { Ok(None) })
90    }
91
92    /// Emit `message_end` and optionally return a replacement message.
93    fn emit_message_end(
94        &self,
95        message: AgentMessage,
96    ) -> BoxFuture<'_, Result<Option<AgentMessage>, ExtensionRunnerError>>;
97
98    /// Emit `tool_call` (before execution). Returns an optional block result.
99    fn emit_tool_call<'a>(
100        &'a self,
101        tool_name: &'a str,
102        tool_call_id: &'a str,
103        input: Map<String, Value>,
104    ) -> BoxFuture<'a, Result<Option<BeforeToolCallResult>, ExtensionRunnerError>>;
105
106    /// Emit `tool_result` (after execution). Returns an optional override.
107    fn emit_tool_result<'a>(
108        &'a self,
109        tool_name: &'a str,
110        tool_call_id: &'a str,
111        input: Map<String, Value>,
112        content: Vec<ToolResultContent>,
113        details: Value,
114        is_error: bool,
115    ) -> BoxFuture<'a, Result<Option<AfterToolCallResult>, ExtensionRunnerError>>;
116
117    /// Emit the `input` transform event.
118    fn emit_input<'a>(
119        &'a self,
120        text: &'a str,
121        images: Option<Value>,
122        source: &'a str,
123        streaming_behavior: Option<&'a str>,
124    ) -> BoxFuture<'a, Result<InputTransformResult, ExtensionRunnerError>>;
125
126    /// Emit `before_agent_start` and return optional message/system prompt injection.
127    fn emit_before_agent_start<'a>(
128        &'a self,
129        prompt: &'a str,
130        images: Option<Value>,
131    ) -> BoxFuture<'a, Result<Option<BeforeAgentStartResult>, ExtensionRunnerError>>;
132
133    /// Discover additional resource paths from extensions.
134    fn emit_resources_discover<'a>(
135        &'a self,
136        cwd: &'a str,
137        reason: &'a str,
138    ) -> BoxFuture<'a, Result<ResourceExtensionPaths, ExtensionRunnerError>>;
139
140    /// Registered slash-command names (extension source).
141    fn get_registered_commands(&self) -> Vec<String>;
142
143    /// Whether a slash command named `name` is registered.
144    fn has_command(&self, name: &str) -> bool {
145        self.get_registered_commands().iter().any(|c| c == name)
146    }
147
148    /// Execute an extension slash command by name.
149    ///
150    /// Returns `Ok(true)` when the command was found and dispatched,
151    /// `Ok(false)` when no such command is registered. Errors during
152    /// execution are reported via [`ExtensionRunner::emit_error`] and the
153    /// return value stays `Ok(true)` (the command was still "handled").
154    fn execute_command<'a>(
155        &'a self,
156        name: &'a str,
157        args: &'a str,
158    ) -> BoxFuture<'a, Result<bool, ExtensionRunnerError>>;
159
160    /// Registered extension tools by name.
161    fn get_all_registered_tools(&self) -> HashMap<String, Arc<dyn AgentTool>>;
162
163    /// Flag values provided by extensions.
164    fn get_flag_values(&self) -> HashMap<String, Value>;
165
166    /// Mark the runner invalid after session replacement.
167    fn invalidate(&self);
168
169    /// Report an extension error to the host error listener.
170    fn emit_error(&self, message: String);
171}
172
173/// No-op extension runner used before `pi-ext` lands and in unit tests.
174#[derive(Clone, Debug, Default)]
175pub struct NullExtensionRunner;
176
177impl ExtensionRunner for NullExtensionRunner {
178    fn has_handlers(&self, _event: &str) -> bool {
179        false
180    }
181
182    fn emit(
183        &self,
184        _event: AgentSessionEvent,
185    ) -> BoxFuture<'_, Result<Option<CancelResult>, ExtensionRunnerError>> {
186        Box::pin(async { Ok(None) })
187    }
188
189    fn emit_message_end(
190        &self,
191        _message: AgentMessage,
192    ) -> BoxFuture<'_, Result<Option<AgentMessage>, ExtensionRunnerError>> {
193        Box::pin(async { Ok(None) })
194    }
195
196    fn emit_tool_call(
197        &self,
198        _tool_name: &str,
199        _tool_call_id: &str,
200        _input: Map<String, Value>,
201    ) -> BoxFuture<'_, Result<Option<BeforeToolCallResult>, ExtensionRunnerError>> {
202        Box::pin(async { Ok(None) })
203    }
204
205    fn emit_tool_result(
206        &self,
207        _tool_name: &str,
208        _tool_call_id: &str,
209        _input: Map<String, Value>,
210        _content: Vec<ToolResultContent>,
211        _details: Value,
212        _is_error: bool,
213    ) -> BoxFuture<'_, Result<Option<AfterToolCallResult>, ExtensionRunnerError>> {
214        Box::pin(async { Ok(None) })
215    }
216
217    fn emit_input(
218        &self,
219        _text: &str,
220        _images: Option<Value>,
221        _source: &str,
222        _streaming_behavior: Option<&str>,
223    ) -> BoxFuture<'_, Result<InputTransformResult, ExtensionRunnerError>> {
224        Box::pin(async { Ok(InputTransformResult::default()) })
225    }
226
227    fn emit_before_agent_start(
228        &self,
229        _prompt: &str,
230        _images: Option<Value>,
231    ) -> BoxFuture<'_, Result<Option<BeforeAgentStartResult>, ExtensionRunnerError>> {
232        Box::pin(async { Ok(None) })
233    }
234
235    fn emit_resources_discover(
236        &self,
237        _cwd: &str,
238        _reason: &str,
239    ) -> BoxFuture<'_, Result<ResourceExtensionPaths, ExtensionRunnerError>> {
240        Box::pin(async { Ok(ResourceExtensionPaths::default()) })
241    }
242
243    fn get_registered_commands(&self) -> Vec<String> {
244        Vec::new()
245    }
246
247    fn execute_command(
248        &self,
249        _name: &str,
250        _args: &str,
251    ) -> BoxFuture<'_, Result<bool, ExtensionRunnerError>> {
252        Box::pin(async { Ok(false) })
253    }
254
255    fn get_all_registered_tools(&self) -> HashMap<String, Arc<dyn AgentTool>> {
256        HashMap::new()
257    }
258
259    fn get_flag_values(&self) -> HashMap<String, Value> {
260        HashMap::new()
261    }
262
263    fn invalidate(&self) {}
264
265    fn emit_error(&self, _message: String) {}
266}
267
268/// System-prompt snapshot shared with the agent `prepare_next_turn` closure.
269#[derive(Clone, Debug, Default)]
270pub struct SystemPromptState {
271    /// Base system prompt rebuilt from resources/tools.
272    pub base: String,
273    /// Per-turn override from `before_agent_start` (cleared after use by callers).
274    pub override_prompt: Option<String>,
275}
276
277/// Shared handle captured by agent hook closures.
278///
279/// Lock order (never hold more than one at a time, never across `.await`):
280/// 1. `runner` (`RwLock`)
281/// 2. `system_prompt` (`RwLock`)
282/// 3. `tools` (`RwLock`)
283///
284/// The agent event pump and public `AgentSession` methods coordinate through
285/// [`super::AgentSessionInner`]; `SessionHooks` only exposes runner/prompt/tool
286/// snapshots for the pi-agent hook closures.
287#[derive(Clone)]
288pub struct SessionHooks {
289    runner: Arc<RwLock<Arc<dyn ExtensionRunner>>>,
290    system_prompt: Arc<RwLock<SystemPromptState>>,
291    tools: Arc<RwLock<Vec<Arc<dyn AgentTool>>>>,
292}
293
294impl SessionHooks {
295    /// Create hooks with the given extension runner.
296    #[must_use]
297    pub fn new(runner: Arc<dyn ExtensionRunner>) -> Self {
298        Self {
299            runner: Arc::new(RwLock::new(runner)),
300            system_prompt: Arc::new(RwLock::new(SystemPromptState::default())),
301            tools: Arc::new(RwLock::new(Vec::new())),
302        }
303    }
304
305    /// Create hooks with a null runner.
306    #[must_use]
307    pub fn null() -> Self {
308        Self::new(Arc::new(NullExtensionRunner))
309    }
310
311    /// Swap the extension runner (reload path).
312    pub fn set_runner(&self, runner: Arc<dyn ExtensionRunner>) {
313        if let Ok(mut guard) = self.runner.write() {
314            *guard = runner;
315        }
316    }
317
318    /// Snapshot the current runner.
319    #[must_use]
320    pub fn runner(&self) -> Arc<dyn ExtensionRunner> {
321        self.runner.read().map_or_else(
322            |poisoned| Arc::clone(&*poisoned.into_inner()),
323            |guard| Arc::clone(&*guard),
324        )
325    }
326
327    /// Replace the base system prompt.
328    pub fn set_base_system_prompt(&self, prompt: String) {
329        if let Ok(mut guard) = self.system_prompt.write() {
330            guard.base = prompt;
331        }
332    }
333
334    /// Set or clear the per-turn system prompt override.
335    pub fn set_system_prompt_override(&self, prompt: Option<String>) {
336        if let Ok(mut guard) = self.system_prompt.write() {
337            guard.override_prompt = prompt;
338        }
339    }
340
341    /// Snapshot system-prompt state.
342    #[must_use]
343    pub fn system_prompt_snapshot(&self) -> SystemPromptState {
344        self.system_prompt.read().map_or_else(
345            |poisoned| poisoned.into_inner().clone(),
346            |guard| guard.clone(),
347        )
348    }
349
350    /// Effective system prompt (`override` or `base`).
351    #[must_use]
352    pub fn effective_system_prompt(&self) -> String {
353        let snap = self.system_prompt_snapshot();
354        snap.override_prompt.unwrap_or(snap.base)
355    }
356
357    /// Replace the tools snapshot used by `prepare_next_turn`.
358    pub fn set_tools(&self, tools: Vec<Arc<dyn AgentTool>>) {
359        if let Ok(mut guard) = self.tools.write() {
360            *guard = tools;
361        }
362    }
363
364    /// Clone the current tools snapshot.
365    #[must_use]
366    pub fn tools_snapshot(&self) -> Vec<Arc<dyn AgentTool>> {
367        self.tools.read().map_or_else(
368            |poisoned| poisoned.into_inner().clone(),
369            |guard| guard.clone(),
370        )
371    }
372
373    /// Build the `before_tool_call` closure for `AgentLoopConfig`.
374    #[must_use]
375    pub fn before_tool_call_hook(self: &Arc<Self>) -> pi_agent::BeforeToolCall {
376        let hooks = Arc::clone(self);
377        Arc::new(
378            move |ctx: BeforeToolCallContext, cancel: CancellationToken| {
379                let hooks = Arc::clone(&hooks);
380                Box::pin(async move {
381                    let runner = hooks.runner();
382                    if !runner.has_handlers("tool_call") {
383                        return Ok(None);
384                    }
385                    let tool_name = ctx.tool_call.name.clone();
386                    let tool_call_id = ctx.tool_call.id.clone();
387                    let hook_result = tokio::select! {
388                        biased;
389                        () = cancel.cancelled() => return Ok(None),
390                        result = runner.emit_tool_call(&tool_name, &tool_call_id, ctx.args) => result,
391                    };
392                    match hook_result {
393                        Ok(result) => Ok(result),
394                        Err(err) => Err(AgentLoopError::message(err.to_string())),
395                    }
396                })
397            },
398        )
399    }
400
401    /// Build the `after_tool_call` closure for `AgentLoopConfig`.
402    #[must_use]
403    pub fn after_tool_call_hook(self: &Arc<Self>) -> pi_agent::AfterToolCall {
404        let hooks = Arc::clone(self);
405        Arc::new(
406            move |ctx: AfterToolCallContext, cancel: CancellationToken| {
407                let hooks = Arc::clone(&hooks);
408                Box::pin(async move {
409                    let runner = hooks.runner();
410                    if !runner.has_handlers("tool_result") {
411                        return Ok(None);
412                    }
413                    let tool_name = ctx.tool_call.name.clone();
414                    let tool_call_id = ctx.tool_call.id.clone();
415                    let content = ctx.result.content.clone();
416                    let details = ctx.result.details.clone();
417                    let hook_result = tokio::select! {
418                        biased;
419                        () = cancel.cancelled() => return Ok(None),
420                        result = runner.emit_tool_result(
421                            &tool_name,
422                            &tool_call_id,
423                            ctx.args,
424                            content,
425                            details,
426                            ctx.is_error,
427                        ) => result,
428                    };
429                    match hook_result {
430                        Ok(result) => Ok(result),
431                        Err(err) => Err(AgentLoopError::message(err.to_string())),
432                    }
433                })
434            },
435        )
436    }
437
438    /// Build the `prepare_next_turn` closure that refreshes system prompt + tools.
439    #[must_use]
440    pub fn prepare_next_turn_hook(self: &Arc<Self>) -> pi_agent::PrepareNextTurn {
441        let hooks = Arc::clone(self);
442        Arc::new(move |turn: PrepareNextTurnContext| {
443            let hooks = Arc::clone(&hooks);
444            Box::pin(async move {
445                let system_prompt = hooks.effective_system_prompt();
446                let tools = hooks.tools_snapshot();
447                let mut context = turn.context;
448                context.system_prompt = system_prompt;
449                if !tools.is_empty() {
450                    context.tools = tools;
451                }
452                Ok(Some(AgentLoopTurnUpdate {
453                    context: Some(context),
454                    model: None,
455                    thinking_level: None,
456                }))
457            })
458        })
459    }
460}
461
462/// Helper: unused imports kept for sibling modules via re-exports.
463#[allow(dead_code)]
464fn _keep_types(
465    _: ToolCall,
466    _: AssistantMessageEvent,
467    _: Model,
468    _: ModelThinkingLevel,
469    _: AgentToolResult,
470) {
471}
pi/core/agent_session/extension_runner.rs

pi/core/agent_session/
extension_runner.rs
