Skip to main content

oxi/extensions/
context.rs

1//! Extension context and builder.
2
3use crate::extensions::types::ExtensionErrorRecord;
4use anyhow::{bail, Context, Result};
5use oxi_store::settings::Settings;
6use parking_lot::RwLock;
7use serde_json::Value;
8use std::path::{Path, PathBuf};
9use std::sync::Arc;
10
11/// Alias for a dynamically-dispatched agent tool.
12pub type ExtensionTool = dyn oxi_agent::AgentTool;
13/// Alias for a reference-counted agent tool.
14pub type ExtensionToolArc = Arc<ExtensionTool>;
15
16// ═══════════════════════════════════════════════════════════════════════════
17// Extension Context
18// ═══════════════════════════════════════════════════════════════════════════
19
20/// Runtime context provided to extension hooks.
21///
22/// Gives extensions access to the session's working directory, settings,
23/// configuration, and various control surfaces (tool registration,
24/// message sending, model switching, etc.).
25#[allow(clippy::type_complexity)]
26pub struct ExtensionContext {
27    /// Current working directory.
28    pub cwd: PathBuf,
29    settings: Arc<RwLock<Settings>>,
30    /// Extension-specific configuration value.
31    pub config: Value,
32    /// Active session ID, if any.
33    pub session_id: Option<String>,
34    idle: Arc<RwLock<bool>>,
35    tool_registrar: Arc<dyn Fn(ExtensionToolArc) + Send + Sync>,
36    message_sender: Arc<dyn Fn(&str) + Send + Sync>,
37    errors: Arc<RwLock<Vec<ExtensionErrorRecord>>>,
38    tool_getter: Arc<dyn Fn() -> Vec<ExtensionToolArc> + Send + Sync>,
39    tool_setter: Arc<dyn Fn(Vec<ExtensionToolArc>) + Send + Sync>,
40    model_setter: Arc<dyn Fn(&str) + Send + Sync>,
41    thinking_level_setter: Arc<dyn Fn(&str) + Send + Sync>,
42    system_prompt_appender: Arc<dyn Fn(&str) + Send + Sync>,
43    session_name_setter: Arc<dyn Fn(&str) + Send + Sync>,
44    session_entries_getter: Arc<dyn Fn() -> Vec<Value> + Send + Sync>,
45    session_fork: Arc<dyn Fn(&str) -> Result<String> + Send + Sync>,
46}
47
48impl std::fmt::Debug for ExtensionContext {
49    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
50        f.debug_struct("ExtensionContext")
51            .field("cwd", &self.cwd)
52            .field("session_id", &self.session_id)
53            .field("idle", &*self.idle.read())
54            .finish()
55    }
56}
57
58impl ExtensionContext {
59    /// Create a new extension context.
60    #[allow(clippy::too_many_arguments)]
61    pub fn new(
62        cwd: PathBuf,
63        settings: Arc<RwLock<Settings>>,
64        config: Value,
65        session_id: Option<String>,
66        idle: Arc<RwLock<bool>>,
67        tool_registrar: Arc<dyn Fn(ExtensionToolArc) + Send + Sync>,
68        message_sender: Arc<dyn Fn(&str) + Send + Sync>,
69        errors: Arc<RwLock<Vec<ExtensionErrorRecord>>>,
70    ) -> Self {
71        Self {
72            cwd,
73            settings,
74            config,
75            session_id,
76            idle,
77            tool_registrar,
78            message_sender,
79            errors,
80            tool_getter: Arc::new(std::vec::Vec::new),
81            tool_setter: Arc::new(|_| {}),
82            model_setter: Arc::new(|_| {}),
83            thinking_level_setter: Arc::new(|_| {}),
84            system_prompt_appender: Arc::new(|_| {}),
85            session_name_setter: Arc::new(|_| {}),
86            session_entries_getter: Arc::new(std::vec::Vec::new),
87            session_fork: Arc::new(|_| bail!("session fork not configured")),
88        }
89    }
90
91    /// Get a snapshot of the current settings.
92    pub fn settings(&self) -> Settings {
93        self.settings.read().clone()
94    }
95    /// Whether the agent is currently idle (not streaming).
96    pub fn is_idle(&self) -> bool {
97        *self.idle.read()
98    }
99
100    /// Register a tool with the agent.
101    pub fn register_tool(&self, tool: ExtensionToolArc) {
102        (self.tool_registrar)(tool);
103    }
104    /// Send a message as the extension.
105    pub fn send_message(&self, text: &str) {
106        (self.message_sender)(text);
107    }
108
109    /// Record an error from an extension.
110    pub fn record_error(&self, extension_name: &str, event: &str, error: &str) {
111        let record = ExtensionErrorRecord::new(extension_name, event, error);
112        tracing::warn!(
113            extension = extension_name,
114            event = event,
115            error = error,
116            "Extension error recorded"
117        );
118        self.errors.write().push(record);
119    }
120
121    /// Get all recorded errors.
122    pub fn errors(&self) -> Vec<ExtensionErrorRecord> {
123        self.errors.read().clone()
124    }
125    /// Clear all recorded errors.
126    pub fn clear_errors(&self) {
127        self.errors.write().clear();
128    }
129
130    /// Look up a value in the extension configuration by dot-separated path.
131    pub fn config_get(&self, path: &str) -> Option<Value> {
132        let mut current = &self.config;
133        for key in path.split('.') {
134            match current {
135                Value::Object(map) => current = map.get(key)?,
136                _ => return None,
137            }
138        }
139        Some(current.clone())
140    }
141
142    /// Read a file relative to the working directory.
143    pub fn read_file(&self, relative_path: &Path) -> Result<String> {
144        let full_path = self.cwd.join(relative_path);
145        std::fs::read_to_string(&full_path)
146            .with_context(|| format!("Failed to read file: {}", full_path.display()))
147    }
148
149    /// Get all currently registered tools.
150    pub fn get_tools(&self) -> Vec<ExtensionToolArc> {
151        (self.tool_getter)()
152    }
153    /// Replace the full set of registered tools.
154    pub fn set_tools(&self, tools: Vec<ExtensionToolArc>) {
155        (self.tool_setter)(tools);
156    }
157    /// Switch the active model.
158    pub fn set_model(&self, model: &str) {
159        (self.model_setter)(model);
160    }
161    /// Set the thinking level.
162    pub fn set_thinking_level(&self, level: &str) {
163        (self.thinking_level_setter)(level);
164    }
165    /// Append text to the system prompt.
166    pub fn append_system_prompt(&self, text: &str) {
167        (self.system_prompt_appender)(text);
168    }
169    /// Set the session display name.
170    pub fn set_session_name(&self, name: &str) {
171        (self.session_name_setter)(name);
172    }
173    /// Get all session entries.
174    pub fn get_session_entries(&self) -> Vec<Value> {
175        (self.session_entries_getter)()
176    }
177    /// Fork the session at the given entry ID.
178    pub fn fork_session(&self, entry_id: &str) -> Result<String> {
179        (self.session_fork)(entry_id)
180    }
181}
182
183// ═══════════════════════════════════════════════════════════════════════════
184// Extension Context Builder
185// ═══════════════════════════════════════════════════════════════════════════
186
187/// Builder for [`ExtensionContext`].
188#[allow(clippy::type_complexity)]
189pub struct ExtensionContextBuilder {
190    cwd: PathBuf,
191    settings: Option<Arc<RwLock<Settings>>>,
192    config: Value,
193    session_id: Option<String>,
194    idle: Arc<RwLock<bool>>,
195    tool_registrar: Option<Arc<dyn Fn(ExtensionToolArc) + Send + Sync>>,
196    message_sender: Option<Arc<dyn Fn(&str) + Send + Sync>>,
197    errors: Option<Arc<RwLock<Vec<ExtensionErrorRecord>>>>,
198    tool_getter: Option<Arc<dyn Fn() -> Vec<ExtensionToolArc> + Send + Sync>>,
199    tool_setter: Option<Arc<dyn Fn(Vec<ExtensionToolArc>) + Send + Sync>>,
200    model_setter: Option<Arc<dyn Fn(&str) + Send + Sync>>,
201    thinking_level_setter: Option<Arc<dyn Fn(&str) + Send + Sync>>,
202    system_prompt_appender: Option<Arc<dyn Fn(&str) + Send + Sync>>,
203    session_name_setter: Option<Arc<dyn Fn(&str) + Send + Sync>>,
204    session_entries_getter: Option<Arc<dyn Fn() -> Vec<Value> + Send + Sync>>,
205    session_fork: Option<Arc<dyn Fn(&str) -> Result<String> + Send + Sync>>,
206}
207
208impl ExtensionContextBuilder {
209    /// Create a new builder for the given working directory.
210    pub fn new(cwd: PathBuf) -> Self {
211        Self {
212            cwd,
213            settings: None,
214            config: Value::Null,
215            session_id: None,
216            idle: Arc::new(RwLock::new(true)),
217            tool_registrar: None,
218            message_sender: None,
219            errors: None,
220            tool_getter: None,
221            tool_setter: None,
222            model_setter: None,
223            thinking_level_setter: None,
224            system_prompt_appender: None,
225            session_name_setter: None,
226            session_entries_getter: None,
227            session_fork: None,
228        }
229    }
230
231    /// Set the shared settings handle.
232    pub fn settings(mut self, settings: Arc<RwLock<Settings>>) -> Self {
233        self.settings = Some(settings);
234        self
235    }
236    /// Set the extension configuration value.
237    pub fn config(mut self, config: Value) -> Self {
238        self.config = config;
239        self
240    }
241    /// Set the session ID.
242    pub fn session_id(mut self, id: impl Into<String>) -> Self {
243        self.session_id = Some(id.into());
244        self
245    }
246    /// Set the idle flag handle.
247    pub fn idle(mut self, idle: Arc<RwLock<bool>>) -> Self {
248        self.idle = idle;
249        self
250    }
251    /// Set the tool registration callback.
252    pub fn tool_registrar(
253        mut self,
254        registrar: Arc<dyn Fn(ExtensionToolArc) + Send + Sync>,
255    ) -> Self {
256        self.tool_registrar = Some(registrar);
257        self
258    }
259    /// Set the message sender callback.
260    pub fn message_sender(mut self, sender: Arc<dyn Fn(&str) + Send + Sync>) -> Self {
261        self.message_sender = Some(sender);
262        self
263    }
264    /// Set the shared error collection.
265    pub fn errors(mut self, errors: Arc<RwLock<Vec<ExtensionErrorRecord>>>) -> Self {
266        self.errors = Some(errors);
267        self
268    }
269    /// Set the tool getter callback.
270    pub fn tool_getter(
271        mut self,
272        getter: Arc<dyn Fn() -> Vec<ExtensionToolArc> + Send + Sync>,
273    ) -> Self {
274        self.tool_getter = Some(getter);
275        self
276    }
277    /// Set the tool setter callback.
278    pub fn tool_setter(mut self, setter: Arc<dyn Fn(Vec<ExtensionToolArc>) + Send + Sync>) -> Self {
279        self.tool_setter = Some(setter);
280        self
281    }
282    /// Set the model setter callback.
283    pub fn model_setter(mut self, setter: Arc<dyn Fn(&str) + Send + Sync>) -> Self {
284        self.model_setter = Some(setter);
285        self
286    }
287    /// Set the thinking-level setter callback.
288    pub fn thinking_level_setter(mut self, setter: Arc<dyn Fn(&str) + Send + Sync>) -> Self {
289        self.thinking_level_setter = Some(setter);
290        self
291    }
292    /// Set the system-prompt appender callback.
293    pub fn system_prompt_appender(mut self, appender: Arc<dyn Fn(&str) + Send + Sync>) -> Self {
294        self.system_prompt_appender = Some(appender);
295        self
296    }
297    /// Set the session-name setter callback.
298    pub fn session_name_setter(mut self, setter: Arc<dyn Fn(&str) + Send + Sync>) -> Self {
299        self.session_name_setter = Some(setter);
300        self
301    }
302    /// Set the session-entries getter callback.
303    pub fn session_entries_getter(
304        mut self,
305        getter: Arc<dyn Fn() -> Vec<Value> + Send + Sync>,
306    ) -> Self {
307        self.session_entries_getter = Some(getter);
308        self
309    }
310    /// Set the session-fork callback.
311    #[allow(clippy::type_complexity)]
312    pub fn session_fork(mut self, fork: Arc<dyn Fn(&str) -> Result<String> + Send + Sync>) -> Self {
313        self.session_fork = Some(fork);
314        self
315    }
316
317    /// Build the [`ExtensionContext`].
318    pub fn build(self) -> ExtensionContext {
319        ExtensionContext {
320            cwd: self.cwd,
321            settings: self
322                .settings
323                .unwrap_or_else(|| Arc::new(RwLock::new(Settings::default()))),
324            config: self.config,
325            session_id: self.session_id,
326            idle: self.idle,
327            tool_registrar: self.tool_registrar.unwrap_or_else(|| {
328                Arc::new(|_tool| {
329                    tracing::debug!("Tool registration attempted with no registrar");
330                })
331            }),
332            message_sender: self.message_sender.unwrap_or_else(|| {
333                Arc::new(|_msg| {
334                    tracing::debug!("Message send attempted with no sender");
335                })
336            }),
337            errors: self
338                .errors
339                .unwrap_or_else(|| Arc::new(RwLock::new(Vec::new()))),
340            tool_getter: self.tool_getter.unwrap_or_else(|| Arc::new(Vec::new)),
341            tool_setter: self.tool_setter.unwrap_or_else(|| Arc::new(|_| {})),
342            model_setter: self.model_setter.unwrap_or_else(|| Arc::new(|_| {})),
343            thinking_level_setter: self
344                .thinking_level_setter
345                .unwrap_or_else(|| Arc::new(|_| {})),
346            system_prompt_appender: self
347                .system_prompt_appender
348                .unwrap_or_else(|| Arc::new(|_| {})),
349            session_name_setter: self.session_name_setter.unwrap_or_else(|| Arc::new(|_| {})),
350            session_entries_getter: self
351                .session_entries_getter
352                .unwrap_or_else(|| Arc::new(Vec::new)),
353            session_fork: self
354                .session_fork
355                .unwrap_or_else(|| Arc::new(|_| bail!("session fork not configured"))),
356        }
357    }
358}