1use std::{collections::BTreeMap, sync::Arc, time::Instant};
4
5use async_trait::async_trait;
6use runtime_types::{DefinitionId, ExecutionId, 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#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
174pub enum ContextSlot {
175 System,
176 Instructions,
177 ToolCatalog,
178 Skills,
179 Memory,
180 Plan,
181 RuntimeState,
182 History,
183 Observations,
184 Checkpoint,
185}
186
187#[derive(Debug, Clone, Copy, PartialEq, Eq)]
188pub enum ContextSensitivity {
189 Public,
190 Internal,
191 Sensitive,
192}
193
194#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
195pub struct ContextInvalidation(u16);
196
197impl ContextInvalidation {
198 pub const NONE: Self = Self(0);
199 pub const SYSTEM: Self = Self(1 << 0);
200 pub const TOOL_CATALOG: Self = Self(1 << 1);
201 pub const SKILLS: Self = Self(1 << 2);
202 pub const MEMORY: Self = Self(1 << 3);
203 pub const PLAN: Self = Self(1 << 4);
204 pub const RUNTIME_STATE: Self = Self(1 << 5);
205
206 pub const fn union(self, other: Self) -> Self {
207 Self(self.0 | other.0)
208 }
209
210 pub const fn contains(self, other: Self) -> bool {
211 self.0 & other.0 == other.0
212 }
213
214 pub const fn is_empty(self) -> bool {
215 self.0 == 0
216 }
217}
218
219impl std::ops::BitOr for ContextInvalidation {
220 type Output = Self;
221
222 fn bitor(self, rhs: Self) -> Self::Output {
223 self.union(rhs)
224 }
225}
226
227impl std::ops::BitOrAssign for ContextInvalidation {
228 fn bitor_assign(&mut self, rhs: Self) {
229 *self = self.union(rhs);
230 }
231}
232
233#[derive(Debug, Clone, PartialEq, Eq)]
234pub struct ContextResource {
235 pub id: String,
236 pub slot: ContextSlot,
237 pub source: String,
238 pub content: String,
239 pub required: bool,
240 pub sensitivity: ContextSensitivity,
241 pub version: String,
242 pub digest: String,
245 pub estimated_tokens: u64,
248}
249
250#[derive(Debug, Clone, PartialEq, Eq)]
251pub struct ContextResourceKey {
252 pub slot: ContextSlot,
253 pub id: String,
254}
255
256#[derive(Debug, Clone, Default, PartialEq, Eq)]
257pub struct ContextPatch {
258 pub resources: Vec<ContextResource>,
259 pub removals: Vec<ContextResourceKey>,
260 pub invalidations: ContextInvalidation,
261}
262
263#[derive(Debug, Clone, PartialEq, Eq)]
264pub enum ContextEvent {
265 ExecutionStart,
266 BeforeModel {
267 turn: usize,
268 },
269 AfterTool {
270 turn: usize,
271 tool_name: String,
272 failed: bool,
273 invalidation: ContextInvalidation,
274 },
275 BeforeCompaction {
276 turn: usize,
277 },
278 AfterCompaction {
279 turn: usize,
280 },
281}
282
283#[derive(Debug, Clone, PartialEq, Eq)]
284pub struct ContextFrame {
285 pub execution_id: ExecutionId,
286 pub turn: usize,
287 pub user_prompt: String,
288 pub visible_tools: Vec<String>,
289 pub pending_invalidations: ContextInvalidation,
290 pub resources: Vec<ContextResource>,
291 pub budget: Option<ContextBudgetFrame>,
292 pub latest_tool: Option<ContextToolFrame>,
293 pub transcript_summary: String,
294}
295
296#[derive(Debug, Clone, PartialEq, Eq)]
297pub struct ContextBudgetFrame {
298 pub max_input_tokens: u64,
299 pub message_tokens: u64,
300 pub tool_schema_tokens: u64,
301 pub protocol_overhead_tokens: u64,
302 pub reserved_output_tokens: u64,
303 pub total_tokens: u64,
304}
305
306#[derive(Debug, Clone, PartialEq, Eq)]
307pub struct ContextToolFrame {
308 pub name: String,
309 pub failed: bool,
310}
311
312#[async_trait]
313pub trait ContextHook: Send + Sync {
314 fn name(&self) -> &str;
315
316 fn definition_fingerprint(&self) -> String;
323
324 async fn apply(
325 &self,
326 context: &ExtensionContext,
327 event: &ContextEvent,
328 frame: &ContextFrame,
329 ) -> Result<ContextPatch, ExtensionError>;
330}
331
332#[derive(Debug, Clone)]
333pub struct ToolInvocation {
334 pub execution_id: ExecutionId,
335 pub call_id: ToolCallId,
336 pub descriptor: ToolDescriptor,
337 pub arguments: Value,
338}
339
340#[async_trait]
341pub trait Guard: Send + Sync {
342 async fn authorize(&self, invocation: &ToolInvocation) -> Result<(), ExtensionError>;
343}
344
345#[derive(Debug, Clone)]
346pub enum ObserverEvent {
347 ExecutionStarted {
348 execution_id: ExecutionId,
349 },
350 ModelCompleted {
351 execution_id: ExecutionId,
352 turn: usize,
353 },
354 ToolCompleted {
355 execution_id: ExecutionId,
356 call_id: ToolCallId,
357 name: String,
358 failed: bool,
359 },
360 ExecutionFinished {
361 execution_id: ExecutionId,
362 succeeded: bool,
363 },
364}
365
366#[async_trait]
367pub trait Observer: Send + Sync {
368 async fn observe(&self, event: ObserverEvent) -> Result<(), ExtensionError>;
369}
370
371#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
372#[serde(rename_all = "camelCase", deny_unknown_fields)]
373pub struct DefinitionManifest {
374 pub id: DefinitionId,
375 pub version: String,
376 pub runtime_type: String,
377 pub schema_version: u32,
378 pub digest: String,
379}
380
381pub struct AgentDefinition {
382 pub manifest: DefinitionManifest,
383 pub system_prompt: String,
384 pub tools: Vec<Arc<dyn Tool>>,
385 pub context_hooks: Vec<Arc<dyn ContextHook>>,
386 pub guards: Vec<Arc<dyn Guard>>,
387 pub observers: Vec<Arc<dyn Observer>>,
388 pub completion_mode: CompletionMode,
389}
390
391#[derive(Debug, Clone, PartialEq, Eq)]
392pub enum CompletionMode {
393 ModelOrTool,
394 RequiredTool(String),
395}
396
397impl AgentDefinition {
398 pub fn tool(&self, name: &str) -> Option<Arc<dyn Tool>> {
399 self.tools
400 .iter()
401 .find(|tool| tool.descriptor().name == name)
402 .cloned()
403 }
404}