use std::{
collections::{BTreeMap, VecDeque},
sync::{Arc, Mutex},
};
use async_trait::async_trait;
use runtime_extension_api::{
AgentDefinition, CompletionMode, DefinitionManifest, ExtensionError, SideEffect, Tool,
ToolContext, ToolDescriptor, ToolOutcome,
};
use runtime_kernel_api::{KernelEvent, KernelEventSink, KernelFailure};
use runtime_ports::{
CommandOutput, CommandRequest, ContextAuditDerivation, ContextAuditSession,
ContextAuditSnapshot, ConversationMessage, ConversationSession, ExecutionSessions,
InteractionRequest, InteractionResponse, InteractionSession, ModelInvocationAudit,
ModelRequest, ModelResponse, ModelSession, OperationContext, PortFailure, ResourceSession,
ResourceSnapshot, SubagentOutcome, SubagentRequest, SubagentSession, TraceSession,
WorkspaceEntry, WorkspaceSession,
};
use runtime_types::DefinitionId;
use serde_json::{Value, json};
pub struct ScriptedModelSession {
responses: Mutex<VecDeque<Result<ModelResponse, PortFailure>>>,
requests: Mutex<Vec<ModelRequest>>,
}
impl ScriptedModelSession {
pub fn new(responses: Vec<ModelResponse>) -> Self {
Self {
responses: Mutex::new(responses.into_iter().map(Ok).collect()),
requests: Mutex::new(Vec::new()),
}
}
pub fn requests(&self) -> Vec<ModelRequest> {
self.requests.lock().expect("model requests lock").clone()
}
}
#[async_trait]
impl ModelSession for ScriptedModelSession {
async fn invoke(
&self,
_: &OperationContext,
request: ModelRequest,
) -> Result<ModelResponse, PortFailure> {
self.requests
.lock()
.expect("model requests lock")
.push(request);
self.responses
.lock()
.expect("model responses lock")
.pop_front()
.unwrap_or_else(|| {
Err(PortFailure::new(
runtime_ports::PortFailureKind::Protocol,
"SCRIPT_EXHAUSTED",
"scripted model has no response",
))
})
}
}
#[derive(Default)]
pub struct MemoryConversationSession {
messages: Mutex<Vec<ConversationMessage>>,
}
impl MemoryConversationSession {
pub fn with_messages(messages: Vec<ConversationMessage>) -> Self {
Self {
messages: Mutex::new(messages),
}
}
pub fn messages(&self) -> Vec<ConversationMessage> {
self.messages.lock().expect("conversation lock").clone()
}
}
#[async_trait]
impl ConversationSession for MemoryConversationSession {
async fn load_recent(
&self,
_: &OperationContext,
limit: usize,
) -> Result<Vec<ConversationMessage>, PortFailure> {
let messages = self.messages();
Ok(messages[messages.len().saturating_sub(limit)..].to_vec())
}
async fn append(
&self,
_: &OperationContext,
messages: &[ConversationMessage],
) -> Result<(), PortFailure> {
self.messages
.lock()
.expect("conversation lock")
.extend_from_slice(messages);
Ok(())
}
}
#[derive(Default)]
pub struct MemoryWorkspaceSession {
files: Mutex<BTreeMap<String, Vec<u8>>>,
commands: Mutex<Vec<CommandRequest>>,
}
impl MemoryWorkspaceSession {
pub fn put(&self, path: &str, content: impl Into<Vec<u8>>) {
self.files
.lock()
.expect("files lock")
.insert(path.into(), content.into());
}
pub fn get(&self, path: &str) -> Option<Vec<u8>> {
self.files.lock().expect("files lock").get(path).cloned()
}
}
#[async_trait]
impl WorkspaceSession for MemoryWorkspaceSession {
async fn describe(&self, _: &OperationContext) -> Result<String, PortFailure> {
Ok("memory workspace".into())
}
async fn read(&self, _: &OperationContext, path: &str) -> Result<Vec<u8>, PortFailure> {
self.get(path).ok_or_else(|| {
PortFailure::new(
runtime_ports::PortFailureKind::NotFound,
"FILE_NOT_FOUND",
path,
)
})
}
async fn write(
&self,
_: &OperationContext,
path: &str,
content: &[u8],
) -> Result<(), PortFailure> {
self.put(path, content);
Ok(())
}
async fn list(
&self,
_: &OperationContext,
path: &str,
) -> Result<Vec<WorkspaceEntry>, PortFailure> {
let prefix = path.trim_matches('/');
Ok(self
.files
.lock()
.expect("files lock")
.iter()
.filter(|(name, _)| name.starts_with(prefix))
.map(|(name, content)| WorkspaceEntry {
path: name.clone(),
is_dir: false,
size: Some(content.len() as u64),
})
.collect())
}
async fn search(
&self,
_: &OperationContext,
path: &str,
query: &str,
limit: usize,
) -> Result<Vec<String>, PortFailure> {
let mut hits = Vec::new();
for (name, content) in self.files.lock().expect("files lock").iter() {
if name.starts_with(path) && String::from_utf8_lossy(content).contains(query) {
hits.push(name.clone());
}
if hits.len() == limit {
break;
}
}
Ok(hits)
}
async fn execute(
&self,
_: &OperationContext,
request: CommandRequest,
) -> Result<CommandOutput, PortFailure> {
self.commands.lock().expect("commands lock").push(request);
Ok(CommandOutput {
exit_code: 0,
stdout: String::new(),
stderr: String::new(),
truncated: false,
})
}
}
#[derive(Default)]
pub struct RecordingTraceSession {
pub events: Mutex<Vec<(String, Value)>>,
}
#[async_trait]
impl TraceSession for RecordingTraceSession {
async fn append(
&self,
_: &OperationContext,
event_type: &str,
payload: Value,
) -> Result<(), PortFailure> {
self.events
.lock()
.expect("trace lock")
.push((event_type.into(), payload));
Ok(())
}
}
pub struct StaticResourceSession {
pub snapshot: ResourceSnapshot,
}
pub struct RecordingContextAuditSession {
pub snapshots: Mutex<Vec<ContextAuditSnapshot>>,
pub derivations: Mutex<Vec<ContextAuditDerivation>>,
pub invocations: Mutex<Vec<ModelInvocationAudit>>,
}
impl Default for RecordingContextAuditSession {
fn default() -> Self {
Self {
snapshots: Mutex::new(Vec::new()),
derivations: Mutex::new(Vec::new()),
invocations: Mutex::new(Vec::new()),
}
}
}
#[async_trait]
impl ContextAuditSession for RecordingContextAuditSession {
async fn snapshot(
&self,
_: &OperationContext,
value: ContextAuditSnapshot,
) -> Result<(), PortFailure> {
self.snapshots.lock().expect("snapshots lock").push(value);
Ok(())
}
async fn derivation(
&self,
_: &OperationContext,
value: ContextAuditDerivation,
) -> Result<(), PortFailure> {
self.derivations
.lock()
.expect("derivations lock")
.push(value);
Ok(())
}
async fn model_invocation(
&self,
_: &OperationContext,
value: ModelInvocationAudit,
) -> Result<(), PortFailure> {
self.invocations
.lock()
.expect("invocations lock")
.push(value);
Ok(())
}
}
#[async_trait]
impl ResourceSession for StaticResourceSession {
async fn load(&self, _: &OperationContext) -> Result<ResourceSnapshot, PortFailure> {
Ok(self.snapshot.clone())
}
}
pub struct RejectInteractionSession;
#[async_trait]
impl InteractionSession for RejectInteractionSession {
async fn prepare(
&self,
_: &OperationContext,
_: InteractionRequest,
) -> Result<(), PortFailure> {
Err(PortFailure::new(
runtime_ports::PortFailureKind::Conflict,
"INTERACTION_UNAVAILABLE",
"test interaction was not scripted",
))
}
async fn wait(
&self,
_: &OperationContext,
_: &str,
) -> Result<InteractionResponse, PortFailure> {
Err(PortFailure::new(
runtime_ports::PortFailureKind::Unavailable,
"INTERACTION_DISABLED",
"fixture does not accept interaction",
))
}
}
pub struct RejectSubagentSession;
#[async_trait]
impl SubagentSession for RejectSubagentSession {
async fn execute(
&self,
_: &OperationContext,
_: SubagentRequest,
) -> Result<SubagentOutcome, PortFailure> {
Err(PortFailure::new(
runtime_ports::PortFailureKind::Conflict,
"SUBAGENT_UNAVAILABLE",
"test subagent was not scripted",
))
}
}
pub fn sessions(model: Arc<dyn ModelSession>) -> ExecutionSessions {
sessions_with_context_audit(model, Arc::new(RecordingContextAuditSession::default()))
}
pub fn sessions_with_context_audit(
model: Arc<dyn ModelSession>,
context_audit: Arc<dyn ContextAuditSession>,
) -> ExecutionSessions {
ExecutionSessions {
model,
conversation: Arc::new(MemoryConversationSession::default()),
context_audit,
workspace: Arc::new(MemoryWorkspaceSession::default()),
trace: Arc::new(RecordingTraceSession::default()),
interaction: Arc::new(RejectInteractionSession),
subagent: Arc::new(RejectSubagentSession),
resources: Arc::new(StaticResourceSession {
snapshot: ResourceSnapshot {
instructions: Vec::new(),
skills: BTreeMap::new(),
digest: "fixture".into(),
},
}),
}
}
#[derive(Default)]
pub struct RecordingKernelEvents {
events: Mutex<Vec<KernelEvent>>,
}
impl RecordingKernelEvents {
pub fn events(&self) -> Vec<KernelEvent> {
self.events.lock().expect("events lock").clone()
}
}
#[async_trait]
impl KernelEventSink for RecordingKernelEvents {
async fn emit(&self, event: KernelEvent) -> Result<(), KernelFailure> {
self.events.lock().expect("events lock").push(event);
Ok(())
}
}
pub struct CompleteTool;
#[async_trait]
impl Tool for CompleteTool {
fn descriptor(&self) -> ToolDescriptor {
ToolDescriptor {
name: "runtime.complete".into(),
description: "Finish the execution with a final answer.".into(),
input_schema: json!({"type":"object","properties":{"answer":{"type":"string"}},"required":["answer"],"additionalProperties":false}),
required_capabilities: Vec::new(),
side_effect: SideEffect::ReadOnly,
}
}
async fn invoke(
&self,
_: ToolContext,
arguments: Value,
) -> Result<ToolOutcome, ExtensionError> {
let answer = arguments
.get("answer")
.and_then(Value::as_str)
.ok_or_else(|| ExtensionError::Invalid("answer is required".into()))?;
Ok(ToolOutcome::complete(answer))
}
}
pub fn fixture_definition(completion_mode: CompletionMode) -> Arc<AgentDefinition> {
Arc::new(AgentDefinition {
manifest: DefinitionManifest {
id: DefinitionId::new("fixture-agent").expect("valid fixture definition id"),
version: "1.0.0".into(),
runtime_type: "fixture".into(),
schema_version: 1,
digest: "fixture-digest".into(),
},
system_prompt: "You are a deterministic fixture agent.".into(),
tools: vec![Arc::new(CompleteTool)],
context_hooks: Vec::new(),
guards: Vec::new(),
observers: Vec::new(),
completion_mode,
})
}