Skip to main content

agent_runtime_extension_api/
lib.rs

1//! Public extension contract used by first-party and third-party Agent packages.
2
3use 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    /// Stable content identity used by the Kernel to avoid treating an
249    /// unchanged resource as a refresh.
250    pub digest: String,
251    /// Agent-side estimate for diagnostics. The Kernel always measures the
252    /// final serialized request independently and remains authoritative.
253    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}