Skip to main content

agent_runtime_testkit/
lib.rs

1//! Stable deterministic fakes for Agent package and Kernel contract tests.
2
3use std::{
4    collections::{BTreeMap, VecDeque},
5    sync::{Arc, Mutex},
6};
7
8use async_trait::async_trait;
9use runtime_extension_api::{
10    AgentPackage, CompletionMode, ExtensionError, PackageManifest, SideEffect, Tool, ToolContext,
11    ToolDescriptor, ToolOutcome,
12};
13use runtime_kernel_api::{KernelEvent, KernelEventSink, KernelFailure};
14use runtime_ports::{
15    CommandOutput, CommandRequest, ContextAuditDerivation, ContextAuditSession,
16    ContextAuditSnapshot, ConversationMessage, ConversationSession, ExecutionSessions,
17    InteractionRequest, InteractionResponse, InteractionSession, ModelInvocationAudit,
18    ModelRequest, ModelResponse, ModelSession, OperationContext, PortFailure, ResourceSession,
19    ResourceSnapshot, SubagentOutcome, SubagentRequest, SubagentSession, TraceSession,
20    WorkspaceEntry, WorkspaceSession,
21};
22use runtime_types::PackageId;
23use serde_json::{Value, json};
24
25pub struct ScriptedModelSession {
26    responses: Mutex<VecDeque<Result<ModelResponse, PortFailure>>>,
27    requests: Mutex<Vec<ModelRequest>>,
28}
29
30impl ScriptedModelSession {
31    pub fn new(responses: Vec<ModelResponse>) -> Self {
32        Self {
33            responses: Mutex::new(responses.into_iter().map(Ok).collect()),
34            requests: Mutex::new(Vec::new()),
35        }
36    }
37
38    pub fn requests(&self) -> Vec<ModelRequest> {
39        self.requests.lock().expect("model requests lock").clone()
40    }
41}
42
43#[async_trait]
44impl ModelSession for ScriptedModelSession {
45    async fn invoke(
46        &self,
47        _: &OperationContext,
48        request: ModelRequest,
49    ) -> Result<ModelResponse, PortFailure> {
50        self.requests
51            .lock()
52            .expect("model requests lock")
53            .push(request);
54        self.responses
55            .lock()
56            .expect("model responses lock")
57            .pop_front()
58            .unwrap_or_else(|| {
59                Err(PortFailure::new(
60                    runtime_ports::PortFailureKind::Protocol,
61                    "SCRIPT_EXHAUSTED",
62                    "scripted model has no response",
63                ))
64            })
65    }
66}
67
68#[derive(Default)]
69pub struct MemoryConversationSession {
70    messages: Mutex<Vec<ConversationMessage>>,
71}
72
73impl MemoryConversationSession {
74    pub fn with_messages(messages: Vec<ConversationMessage>) -> Self {
75        Self {
76            messages: Mutex::new(messages),
77        }
78    }
79    pub fn messages(&self) -> Vec<ConversationMessage> {
80        self.messages.lock().expect("conversation lock").clone()
81    }
82}
83
84#[async_trait]
85impl ConversationSession for MemoryConversationSession {
86    async fn load_recent(
87        &self,
88        _: &OperationContext,
89        limit: usize,
90    ) -> Result<Vec<ConversationMessage>, PortFailure> {
91        let messages = self.messages();
92        Ok(messages[messages.len().saturating_sub(limit)..].to_vec())
93    }
94    async fn append(
95        &self,
96        _: &OperationContext,
97        messages: &[ConversationMessage],
98    ) -> Result<(), PortFailure> {
99        self.messages
100            .lock()
101            .expect("conversation lock")
102            .extend_from_slice(messages);
103        Ok(())
104    }
105}
106
107#[derive(Default)]
108pub struct MemoryWorkspaceSession {
109    files: Mutex<BTreeMap<String, Vec<u8>>>,
110    commands: Mutex<Vec<CommandRequest>>,
111}
112
113impl MemoryWorkspaceSession {
114    pub fn put(&self, path: &str, content: impl Into<Vec<u8>>) {
115        self.files
116            .lock()
117            .expect("files lock")
118            .insert(path.into(), content.into());
119    }
120    pub fn get(&self, path: &str) -> Option<Vec<u8>> {
121        self.files.lock().expect("files lock").get(path).cloned()
122    }
123}
124
125#[async_trait]
126impl WorkspaceSession for MemoryWorkspaceSession {
127    async fn describe(&self, _: &OperationContext) -> Result<String, PortFailure> {
128        Ok("memory workspace".into())
129    }
130    async fn read(&self, _: &OperationContext, path: &str) -> Result<Vec<u8>, PortFailure> {
131        self.get(path).ok_or_else(|| {
132            PortFailure::new(
133                runtime_ports::PortFailureKind::NotFound,
134                "FILE_NOT_FOUND",
135                path,
136            )
137        })
138    }
139    async fn write(
140        &self,
141        _: &OperationContext,
142        path: &str,
143        content: &[u8],
144    ) -> Result<(), PortFailure> {
145        self.put(path, content);
146        Ok(())
147    }
148    async fn list(
149        &self,
150        _: &OperationContext,
151        path: &str,
152    ) -> Result<Vec<WorkspaceEntry>, PortFailure> {
153        let prefix = path.trim_matches('/');
154        Ok(self
155            .files
156            .lock()
157            .expect("files lock")
158            .iter()
159            .filter(|(name, _)| name.starts_with(prefix))
160            .map(|(name, content)| WorkspaceEntry {
161                path: name.clone(),
162                is_dir: false,
163                size: Some(content.len() as u64),
164            })
165            .collect())
166    }
167    async fn search(
168        &self,
169        _: &OperationContext,
170        path: &str,
171        query: &str,
172        limit: usize,
173    ) -> Result<Vec<String>, PortFailure> {
174        let mut hits = Vec::new();
175        for (name, content) in self.files.lock().expect("files lock").iter() {
176            if name.starts_with(path) && String::from_utf8_lossy(content).contains(query) {
177                hits.push(name.clone());
178            }
179            if hits.len() == limit {
180                break;
181            }
182        }
183        Ok(hits)
184    }
185    async fn execute(
186        &self,
187        _: &OperationContext,
188        request: CommandRequest,
189    ) -> Result<CommandOutput, PortFailure> {
190        self.commands.lock().expect("commands lock").push(request);
191        Ok(CommandOutput {
192            exit_code: 0,
193            stdout: String::new(),
194            stderr: String::new(),
195            truncated: false,
196        })
197    }
198}
199
200#[derive(Default)]
201pub struct RecordingTraceSession {
202    pub events: Mutex<Vec<(String, Value)>>,
203}
204
205#[async_trait]
206impl TraceSession for RecordingTraceSession {
207    async fn append(
208        &self,
209        _: &OperationContext,
210        event_type: &str,
211        payload: Value,
212    ) -> Result<(), PortFailure> {
213        self.events
214            .lock()
215            .expect("trace lock")
216            .push((event_type.into(), payload));
217        Ok(())
218    }
219}
220
221pub struct StaticResourceSession {
222    pub snapshot: ResourceSnapshot,
223}
224
225pub struct RecordingContextAuditSession {
226    pub snapshots: Mutex<Vec<ContextAuditSnapshot>>,
227    pub derivations: Mutex<Vec<ContextAuditDerivation>>,
228    pub invocations: Mutex<Vec<ModelInvocationAudit>>,
229}
230
231impl Default for RecordingContextAuditSession {
232    fn default() -> Self {
233        Self {
234            snapshots: Mutex::new(Vec::new()),
235            derivations: Mutex::new(Vec::new()),
236            invocations: Mutex::new(Vec::new()),
237        }
238    }
239}
240
241#[async_trait]
242impl ContextAuditSession for RecordingContextAuditSession {
243    async fn snapshot(
244        &self,
245        _: &OperationContext,
246        value: ContextAuditSnapshot,
247    ) -> Result<(), PortFailure> {
248        self.snapshots.lock().expect("snapshots lock").push(value);
249        Ok(())
250    }
251    async fn derivation(
252        &self,
253        _: &OperationContext,
254        value: ContextAuditDerivation,
255    ) -> Result<(), PortFailure> {
256        self.derivations
257            .lock()
258            .expect("derivations lock")
259            .push(value);
260        Ok(())
261    }
262    async fn model_invocation(
263        &self,
264        _: &OperationContext,
265        value: ModelInvocationAudit,
266    ) -> Result<(), PortFailure> {
267        self.invocations
268            .lock()
269            .expect("invocations lock")
270            .push(value);
271        Ok(())
272    }
273}
274
275#[async_trait]
276impl ResourceSession for StaticResourceSession {
277    async fn load(&self, _: &OperationContext) -> Result<ResourceSnapshot, PortFailure> {
278        Ok(self.snapshot.clone())
279    }
280}
281
282pub struct RejectInteractionSession;
283#[async_trait]
284impl InteractionSession for RejectInteractionSession {
285    async fn prepare(
286        &self,
287        _: &OperationContext,
288        _: InteractionRequest,
289    ) -> Result<(), PortFailure> {
290        Err(PortFailure::new(
291            runtime_ports::PortFailureKind::Conflict,
292            "INTERACTION_UNAVAILABLE",
293            "test interaction was not scripted",
294        ))
295    }
296
297    async fn wait(
298        &self,
299        _: &OperationContext,
300        _: &str,
301    ) -> Result<InteractionResponse, PortFailure> {
302        Err(PortFailure::new(
303            runtime_ports::PortFailureKind::Unavailable,
304            "INTERACTION_DISABLED",
305            "fixture does not accept interaction",
306        ))
307    }
308}
309
310pub struct RejectSubagentSession;
311#[async_trait]
312impl SubagentSession for RejectSubagentSession {
313    async fn execute(
314        &self,
315        _: &OperationContext,
316        _: SubagentRequest,
317    ) -> Result<SubagentOutcome, PortFailure> {
318        Err(PortFailure::new(
319            runtime_ports::PortFailureKind::Conflict,
320            "SUBAGENT_UNAVAILABLE",
321            "test subagent was not scripted",
322        ))
323    }
324}
325
326pub fn sessions(model: Arc<dyn ModelSession>) -> ExecutionSessions {
327    sessions_with_context_audit(model, Arc::new(RecordingContextAuditSession::default()))
328}
329
330pub fn sessions_with_context_audit(
331    model: Arc<dyn ModelSession>,
332    context_audit: Arc<dyn ContextAuditSession>,
333) -> ExecutionSessions {
334    ExecutionSessions {
335        model,
336        conversation: Arc::new(MemoryConversationSession::default()),
337        context_audit,
338        workspace: Arc::new(MemoryWorkspaceSession::default()),
339        trace: Arc::new(RecordingTraceSession::default()),
340        interaction: Arc::new(RejectInteractionSession),
341        subagent: Arc::new(RejectSubagentSession),
342        resources: Arc::new(StaticResourceSession {
343            snapshot: ResourceSnapshot {
344                instructions: Vec::new(),
345                skills: BTreeMap::new(),
346                digest: "fixture".into(),
347            },
348        }),
349    }
350}
351
352#[derive(Default)]
353pub struct RecordingKernelEvents {
354    events: Mutex<Vec<KernelEvent>>,
355}
356
357impl RecordingKernelEvents {
358    pub fn events(&self) -> Vec<KernelEvent> {
359        self.events.lock().expect("events lock").clone()
360    }
361}
362
363#[async_trait]
364impl KernelEventSink for RecordingKernelEvents {
365    async fn emit(&self, event: KernelEvent) -> Result<(), KernelFailure> {
366        self.events.lock().expect("events lock").push(event);
367        Ok(())
368    }
369}
370
371pub struct CompleteTool;
372
373#[async_trait]
374impl Tool for CompleteTool {
375    fn descriptor(&self) -> ToolDescriptor {
376        ToolDescriptor {
377            name: "runtime.complete".into(),
378            description: "Finish the execution with a final answer.".into(),
379            input_schema: json!({"type":"object","properties":{"answer":{"type":"string"}},"required":["answer"],"additionalProperties":false}),
380            required_capabilities: Vec::new(),
381            side_effect: SideEffect::ReadOnly,
382        }
383    }
384    async fn invoke(
385        &self,
386        _: ToolContext,
387        arguments: Value,
388    ) -> Result<ToolOutcome, ExtensionError> {
389        let answer = arguments
390            .get("answer")
391            .and_then(Value::as_str)
392            .ok_or_else(|| ExtensionError::Invalid("answer is required".into()))?;
393        Ok(ToolOutcome::complete(answer))
394    }
395}
396
397pub fn fixture_package(completion_mode: CompletionMode) -> Arc<AgentPackage> {
398    Arc::new(AgentPackage {
399        manifest: PackageManifest {
400            id: PackageId::new("fixture-agent").expect("valid fixture package id"),
401            version: "1.0.0".into(),
402            runtime_type: "fixture".into(),
403            schema_version: 1,
404            digest: "fixture-digest".into(),
405        },
406        system_prompt: "You are a deterministic fixture agent.".into(),
407        tools: vec![Arc::new(CompleteTool)],
408        context_sources: Vec::new(),
409        context_hooks: Vec::new(),
410        guards: Vec::new(),
411        observers: Vec::new(),
412        completion_mode,
413    })
414}