1use std::{collections::BTreeMap, sync::Arc, time::Instant};
4
5use async_trait::async_trait;
6use runtime_types::{ExecutionId, PackageId, ToolCallId};
7use serde::{Deserialize, Serialize};
8use serde_json::Value;
9use thiserror::Error;
10use tokio_util::sync::CancellationToken;
11
12#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
13#[serde(rename_all = "camelCase", deny_unknown_fields)]
14pub struct ToolDescriptor {
15 pub name: String,
16 pub description: String,
17 pub input_schema: Value,
18 #[serde(default)]
19 pub required_capabilities: Vec<String>,
20 #[serde(default)]
21 pub side_effect: SideEffect,
22}
23
24#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq)]
25#[serde(rename_all = "snake_case")]
26pub enum SideEffect {
27 #[default]
28 ReadOnly,
29 WorkspaceWrite,
30 Process,
31 External,
32}
33
34#[derive(Debug, Error, Clone)]
35pub enum ExtensionError {
36 #[error("invalid extension input: {0}")]
37 Invalid(String),
38 #[error("extension permission denied: {0}")]
39 Denied(String),
40 #[error("extension operation failed: {0}")]
41 Failed(String),
42 #[error("extension operation canceled")]
43 Canceled,
44 #[error("extension operation timed out")]
45 Timeout,
46}
47
48#[derive(Debug, Clone)]
49pub struct ExtensionContext {
50 pub execution_id: ExecutionId,
51 pub deadline: Instant,
52 pub cancellation: CancellationToken,
53}
54
55#[derive(Clone)]
56pub struct ToolContext {
57 pub extension: ExtensionContext,
58 pub call_id: ToolCallId,
59 pub workspace: Arc<dyn WorkspaceFacade>,
60 pub interaction: Arc<dyn InteractionFacade>,
61 pub subagent: Arc<dyn SubagentFacade>,
62}
63
64#[derive(Debug, Clone, PartialEq)]
65pub struct ToolOutcome {
66 pub content: String,
67 pub failed: bool,
68 pub completion: Option<String>,
69 pub metadata: BTreeMap<String, Value>,
70 pub context_invalidation: ContextInvalidation,
71}
72
73impl ToolOutcome {
74 pub fn success(content: impl Into<String>) -> Self {
75 Self {
76 content: content.into(),
77 failed: false,
78 completion: None,
79 metadata: BTreeMap::new(),
80 context_invalidation: ContextInvalidation::NONE,
81 }
82 }
83 pub fn failure(content: impl Into<String>) -> Self {
84 Self {
85 content: content.into(),
86 failed: true,
87 completion: None,
88 metadata: BTreeMap::new(),
89 context_invalidation: ContextInvalidation::NONE,
90 }
91 }
92
93 pub fn complete(answer: impl Into<String>) -> Self {
94 let answer = answer.into();
95 Self {
96 content: answer.clone(),
97 failed: false,
98 completion: Some(answer),
99 metadata: BTreeMap::new(),
100 context_invalidation: ContextInvalidation::NONE,
101 }
102 }
103
104 pub fn with_context_invalidation(mut self, invalidation: ContextInvalidation) -> Self {
105 self.context_invalidation = invalidation;
106 self
107 }
108}
109
110#[async_trait]
111pub trait Tool: Send + Sync {
112 fn descriptor(&self) -> ToolDescriptor;
113 async fn invoke(
114 &self,
115 context: ToolContext,
116 arguments: Value,
117 ) -> Result<ToolOutcome, ExtensionError>;
118}
119
120#[derive(Debug, Clone, PartialEq, Eq)]
121pub struct WorkspaceEntry {
122 pub path: String,
123 pub is_dir: bool,
124 pub size: Option<u64>,
125}
126
127#[derive(Debug, Clone, PartialEq, Eq)]
128pub struct CommandOutcome {
129 pub exit_code: i32,
130 pub stdout: String,
131 pub stderr: String,
132 pub truncated: bool,
133}
134
135#[async_trait]
136pub trait WorkspaceFacade: Send + Sync {
137 async fn describe(&self) -> Result<String, ExtensionError>;
138 async fn read_text(&self, path: &str) -> Result<String, ExtensionError>;
139 async fn write_text(&self, path: &str, content: &str) -> Result<(), ExtensionError>;
140 async fn list(&self, path: &str) -> Result<Vec<WorkspaceEntry>, ExtensionError>;
141 async fn search(
142 &self,
143 path: &str,
144 query: &str,
145 limit: usize,
146 ) -> Result<Vec<String>, ExtensionError>;
147 async fn execute(
148 &self,
149 command: &str,
150 cwd: Option<&str>,
151 ) -> Result<CommandOutcome, ExtensionError>;
152}
153
154#[async_trait]
155pub trait InteractionFacade: Send + Sync {
156 async fn request_text(&self, prompt: &str) -> Result<String, ExtensionError>;
157}
158
159#[async_trait]
160pub trait SubagentFacade: Send + Sync {
161 async fn execute(&self, prompt: &str, max_model_turns: usize)
162 -> Result<String, ExtensionError>;
163}
164
165#[derive(Debug, Clone, PartialEq, Eq)]
166pub struct ContextItem {
167 pub source: String,
168 pub content: String,
169 pub priority: i32,
170 pub required: bool,
171}
172
173#[async_trait]
174pub trait ContextSource: Send + Sync {
175 fn name(&self) -> &str;
176 async fn load(&self, context: &ExtensionContext) -> Result<Vec<ContextItem>, ExtensionError>;
177}
178
179#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
180pub enum ContextSlot {
181 System,
182 Instructions,
183 ToolCatalog,
184 Skills,
185 Memory,
186 Plan,
187 RuntimeState,
188 History,
189 Observations,
190 Checkpoint,
191}
192
193#[derive(Debug, Clone, Copy, PartialEq, Eq)]
194pub enum ContextSensitivity {
195 Public,
196 Internal,
197 Sensitive,
198}
199
200#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
201pub struct ContextInvalidation(u16);
202
203impl ContextInvalidation {
204 pub const NONE: Self = Self(0);
205 pub const SYSTEM: Self = Self(1 << 0);
206 pub const TOOL_CATALOG: Self = Self(1 << 1);
207 pub const SKILLS: Self = Self(1 << 2);
208 pub const MEMORY: Self = Self(1 << 3);
209 pub const PLAN: Self = Self(1 << 4);
210 pub const RUNTIME_STATE: Self = Self(1 << 5);
211
212 pub const fn union(self, other: Self) -> Self {
213 Self(self.0 | other.0)
214 }
215
216 pub const fn contains(self, other: Self) -> bool {
217 self.0 & other.0 == other.0
218 }
219
220 pub const fn is_empty(self) -> bool {
221 self.0 == 0
222 }
223}
224
225impl std::ops::BitOr for ContextInvalidation {
226 type Output = Self;
227
228 fn bitor(self, rhs: Self) -> Self::Output {
229 self.union(rhs)
230 }
231}
232
233impl std::ops::BitOrAssign for ContextInvalidation {
234 fn bitor_assign(&mut self, rhs: Self) {
235 *self = self.union(rhs);
236 }
237}
238
239#[derive(Debug, Clone, PartialEq, Eq)]
240pub struct ContextResource {
241 pub id: String,
242 pub slot: ContextSlot,
243 pub source: String,
244 pub content: String,
245 pub required: bool,
246 pub sensitivity: ContextSensitivity,
247 pub version: String,
248 pub digest: String,
251 pub estimated_tokens: u64,
254}
255
256#[derive(Debug, Clone, PartialEq, Eq)]
257pub struct ContextResourceKey {
258 pub slot: ContextSlot,
259 pub id: String,
260}
261
262#[derive(Debug, Clone, Default, PartialEq, Eq)]
263pub struct ContextPatch {
264 pub resources: Vec<ContextResource>,
265 pub removals: Vec<ContextResourceKey>,
266 pub invalidations: ContextInvalidation,
267}
268
269#[derive(Debug, Clone, PartialEq, Eq)]
270pub enum ContextEvent {
271 ExecutionStart,
272 BeforeModel {
273 turn: usize,
274 },
275 AfterTool {
276 turn: usize,
277 tool_name: String,
278 failed: bool,
279 invalidation: ContextInvalidation,
280 },
281 BeforeCompaction {
282 turn: usize,
283 },
284 AfterCompaction {
285 turn: usize,
286 },
287}
288
289#[derive(Debug, Clone, PartialEq, Eq)]
290pub struct ContextFrame {
291 pub execution_id: ExecutionId,
292 pub turn: usize,
293 pub user_prompt: String,
294 pub visible_tools: Vec<String>,
295 pub pending_invalidations: ContextInvalidation,
296 pub resources: Vec<ContextResource>,
297 pub budget: Option<ContextBudgetFrame>,
298 pub latest_tool: Option<ContextToolFrame>,
299 pub transcript_summary: String,
300}
301
302#[derive(Debug, Clone, PartialEq, Eq)]
303pub struct ContextBudgetFrame {
304 pub max_input_tokens: u64,
305 pub message_tokens: u64,
306 pub tool_schema_tokens: u64,
307 pub protocol_overhead_tokens: u64,
308 pub reserved_output_tokens: u64,
309 pub total_tokens: u64,
310}
311
312#[derive(Debug, Clone, PartialEq, Eq)]
313pub struct ContextToolFrame {
314 pub name: String,
315 pub failed: bool,
316}
317
318#[async_trait]
319pub trait ContextHook: Send + Sync {
320 fn name(&self) -> &str;
321
322 async fn apply(
323 &self,
324 context: &ExtensionContext,
325 event: &ContextEvent,
326 frame: &ContextFrame,
327 ) -> Result<ContextPatch, ExtensionError>;
328}
329
330#[derive(Debug, Clone)]
331pub struct ToolInvocation {
332 pub execution_id: ExecutionId,
333 pub call_id: ToolCallId,
334 pub descriptor: ToolDescriptor,
335 pub arguments: Value,
336}
337
338#[async_trait]
339pub trait Guard: Send + Sync {
340 async fn authorize(&self, invocation: &ToolInvocation) -> Result<(), ExtensionError>;
341}
342
343#[derive(Debug, Clone)]
344pub enum ObserverEvent {
345 ExecutionStarted {
346 execution_id: ExecutionId,
347 },
348 ModelCompleted {
349 execution_id: ExecutionId,
350 turn: usize,
351 },
352 ToolCompleted {
353 execution_id: ExecutionId,
354 call_id: ToolCallId,
355 name: String,
356 failed: bool,
357 },
358 ExecutionFinished {
359 execution_id: ExecutionId,
360 succeeded: bool,
361 },
362}
363
364#[async_trait]
365pub trait Observer: Send + Sync {
366 async fn observe(&self, event: ObserverEvent) -> Result<(), ExtensionError>;
367}
368
369#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
370#[serde(rename_all = "camelCase", deny_unknown_fields)]
371pub struct PackageManifest {
372 pub id: PackageId,
373 pub version: String,
374 pub runtime_type: String,
375 pub schema_version: u32,
376 pub digest: String,
377}
378
379pub struct AgentPackage {
380 pub manifest: PackageManifest,
381 pub system_prompt: String,
382 pub tools: Vec<Arc<dyn Tool>>,
383 pub context_sources: Vec<Arc<dyn ContextSource>>,
384 pub context_hooks: Vec<Arc<dyn ContextHook>>,
385 pub guards: Vec<Arc<dyn Guard>>,
386 pub observers: Vec<Arc<dyn Observer>>,
387 pub completion_mode: CompletionMode,
388}
389
390#[derive(Debug, Clone, PartialEq, Eq)]
391pub enum CompletionMode {
392 ModelOrTool,
393 RequiredTool(String),
394}
395
396impl AgentPackage {
397 pub fn tool(&self, name: &str) -> Option<Arc<dyn Tool>> {
398 self.tools
399 .iter()
400 .find(|tool| tool.descriptor().name == name)
401 .cloned()
402 }
403}