1use 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}