Skip to main content

agent_runtime_extension_api/
lib.rs

1//! Public extension contract used by first-party and third-party Agent definitions.
2
3use 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    /// Stable content identity used by the Kernel to avoid treating an
243    /// unchanged resource as a refresh.
244    pub digest: String,
245    /// Agent-side estimate for diagnostics. The Kernel always measures the
246    /// final serialized request independently and remains authoritative.
247    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    /// Stable authoring fingerprint for the Hook implementation/configuration.
317    ///
318    /// The Definition builder incorporates this value into its immutable
319    /// digest. Change it whenever Hook behavior or embedded static resources
320    /// change. Per-execution dynamic resource content still carries its own
321    /// `ContextResource.digest` at runtime.
322    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}