luft_core/
testing.rs

1//! Test utilities for resume and integration testing.
2//!
3//! Provides instrumented backends that can simulate crashes, blocking,
4//! and call recording — useful for testing crash-and-resume scenarios.
5
6use crate::contract::backend::{
7    AgentBackend, AgentCapabilities, AgentResult, AgentStatus, AgentTask, BackendError, LogRef,
8    RunContext,
9};
10use crate::contract::ids::TokenUsage;
11use serde_json::Value;
12use std::path::Path;
13use std::sync::atomic::{AtomicU64, Ordering};
14use std::sync::{Arc, Mutex};
15use std::time::Duration;
16
17#[derive(Debug, Clone)]
18pub struct CallRecord {
19    pub seq: u64,
20    pub agent_name: Option<String>,
21    pub thread_id: Option<String>,
22    pub prompt: String,
23}
24
25/// Backend that calls `std::process::exit(1)` after N calls.
26pub struct CrashBackend {
27    canned: Value,
28    crash_after: u64,
29    count: AtomicU64,
30}
31
32impl CrashBackend {
33    pub fn new(canned: Value, crash_after: u64) -> Self {
34        Self {
35            canned,
36            crash_after,
37            count: AtomicU64::new(0),
38        }
39    }
40}
41
42#[async_trait::async_trait]
43impl AgentBackend for CrashBackend {
44    fn id(&self) -> &'static str {
45        "crash"
46    }
47    fn capabilities(&self) -> AgentCapabilities {
48        AgentCapabilities {
49            streaming: true,
50            mcp_injection: false,
51            structured_output: false,
52            models: vec![],
53        }
54    }
55    async fn run(&self, task: AgentTask, _ctx: RunContext) -> Result<AgentResult, BackendError> {
56        let n = self.count.fetch_add(1, Ordering::SeqCst) + 1;
57        eprintln!(
58            "[crash-backend] agent #{n}: {}",
59            task.name.as_deref().unwrap_or("?")
60        );
61        if n >= self.crash_after {
62            eprintln!("[crash-backend] exiting after {n} calls");
63            std::process::exit(1);
64        }
65        Ok(AgentResult {
66            agent_id: task.agent_id,
67            status: AgentStatus::Ok,
68            output: self.canned.clone(),
69            thread_id: task.thread_id.clone(),
70            findings: vec![],
71            tokens_used: TokenUsage::default(),
72            artifacts: vec![],
73            logs: LogRef::default(),
74        })
75    }
76    fn as_any(&self) -> &dyn std::any::Any {
77        self
78    }
79}
80
81/// Backend that records all dispatched agent names.
82#[derive(Clone)]
83pub struct CountingBackend {
84    canned: Value,
85    calls: Arc<Mutex<Vec<String>>>,
86}
87
88impl CountingBackend {
89    pub fn new(canned: Value) -> Self {
90        Self {
91            canned,
92            calls: Arc::new(Mutex::new(Vec::new())),
93        }
94    }
95    pub fn dispatched_names(&self) -> Vec<String> {
96        self.calls.lock().unwrap().clone()
97    }
98    pub fn total_calls(&self) -> usize {
99        self.calls.lock().unwrap().len()
100    }
101}
102
103#[async_trait::async_trait]
104impl AgentBackend for CountingBackend {
105    fn id(&self) -> &'static str {
106        "counting"
107    }
108    fn capabilities(&self) -> AgentCapabilities {
109        AgentCapabilities {
110            streaming: true,
111            mcp_injection: false,
112            structured_output: false,
113            models: vec![],
114        }
115    }
116    async fn run(&self, task: AgentTask, _ctx: RunContext) -> Result<AgentResult, BackendError> {
117        self.calls
118            .lock()
119            .unwrap()
120            .push(task.name.clone().unwrap_or_default());
121        Ok(AgentResult {
122            agent_id: task.agent_id,
123            status: AgentStatus::Ok,
124            output: self.canned.clone(),
125            thread_id: task.thread_id.clone(),
126            findings: vec![],
127            tokens_used: TokenUsage::default(),
128            artifacts: vec![],
129            logs: LogRef::default(),
130        })
131    }
132    fn as_any(&self) -> &dyn std::any::Any {
133        self
134    }
135}
136
137/// Backend with shared state via Arc for resume tests.
138#[derive(Clone)]
139pub struct SharedBackend {
140    canned: Value,
141    call_count: Arc<AtomicU64>,
142    pub block_on: Arc<Mutex<Option<u64>>>,
143    pub fail_on: Arc<Mutex<Option<u64>>>,
144    calls: Arc<Mutex<Vec<CallRecord>>>,
145}
146
147impl SharedBackend {
148    pub fn new(canned: Value) -> Self {
149        Self {
150            canned,
151            call_count: Arc::new(AtomicU64::new(0)),
152            block_on: Arc::new(Mutex::new(None)),
153            fail_on: Arc::new(Mutex::new(None)),
154            calls: Arc::new(Mutex::new(Vec::new())),
155        }
156    }
157    pub fn with_block_on(self, n: u64) -> Self {
158        *self.block_on.lock().unwrap() = Some(n);
159        self
160    }
161    pub fn with_fail_on(self, n: u64) -> Self {
162        *self.fail_on.lock().unwrap() = Some(n);
163        self
164    }
165    pub fn total_calls(&self) -> usize {
166        self.calls.lock().unwrap().len()
167    }
168    pub fn calls_snapshot(&self) -> Vec<CallRecord> {
169        self.calls.lock().unwrap().clone()
170    }
171    pub fn dispatched_names(&self) -> Vec<String> {
172        self.calls
173            .lock()
174            .unwrap()
175            .iter()
176            .map(|c| c.agent_name.clone().unwrap_or_default())
177            .collect()
178    }
179    pub fn mirror(&self) -> Self {
180        Self {
181            canned: self.canned.clone(),
182            call_count: self.call_count.clone(),
183            block_on: self.block_on.clone(),
184            fail_on: self.fail_on.clone(),
185            calls: self.calls.clone(),
186        }
187    }
188}
189
190#[async_trait::async_trait]
191impl AgentBackend for SharedBackend {
192    fn id(&self) -> &'static str {
193        "shared"
194    }
195    fn capabilities(&self) -> AgentCapabilities {
196        AgentCapabilities {
197            streaming: true,
198            mcp_injection: false,
199            structured_output: false,
200            models: vec![],
201        }
202    }
203    async fn run(&self, task: AgentTask, ctx: RunContext) -> Result<AgentResult, BackendError> {
204        let seq = self.call_count.fetch_add(1, Ordering::SeqCst) + 1;
205        self.calls.lock().unwrap().push(CallRecord {
206            seq,
207            agent_name: task.name.clone(),
208            thread_id: task.thread_id.clone(),
209            prompt: task.prompt.clone(),
210        });
211        if self
212            .fail_on
213            .lock()
214            .unwrap()
215            .map(|n| n == seq)
216            .unwrap_or(false)
217        {
218            return Err(BackendError::Execution("simulated failure".into()));
219        }
220        if self
221            .block_on
222            .lock()
223            .unwrap()
224            .map(|n| n == seq)
225            .unwrap_or(false)
226        {
227            ctx.cancel.cancelled().await;
228            return Err(BackendError::Cancelled);
229        }
230        Ok(AgentResult {
231            agent_id: task.agent_id,
232            status: AgentStatus::Ok,
233            output: self.canned.clone(),
234            thread_id: task.thread_id.clone(),
235            findings: vec![],
236            tokens_used: TokenUsage::default(),
237            artifacts: vec![],
238            logs: LogRef::default(),
239        })
240    }
241    fn as_any(&self) -> &dyn std::any::Any {
242        self
243    }
244}
245
246pub async fn wait_for_calls(backend: &SharedBackend, n: usize, timeout_ms: u64) {
247    let deadline = tokio::time::sleep(Duration::from_millis(timeout_ms));
248    tokio::pin!(deadline);
249    loop {
250        if backend.total_calls() >= n {
251            return;
252        }
253        tokio::select! {
254            _ = &mut deadline => panic!("timeout waiting for {n} agent calls (got {})", backend.total_calls()),
255            _ = tokio::time::sleep(Duration::from_millis(25)) => {}
256        }
257    }
258}
259
260pub async fn read_checkpoint(base: &Path, run_dir: &str) -> Value {
261    let path = base.join(run_dir).join("checkpoint.json");
262    match tokio::fs::read(&path).await {
263        Ok(bytes) => serde_json::from_slice(&bytes).unwrap_or(Value::Null),
264        Err(_) => Value::Null,
265    }
266}
267
268pub fn completed_span_names(cp: &Value) -> Vec<String> {
269    cp.get("phase_state")
270        .and_then(|ps| ps.get("completed_spans"))
271        .and_then(|cs| cs.as_array())
272        .map(|arr| {
273            arr.iter()
274                .filter_map(|v| v.get("name").and_then(|n| n.as_str()).map(String::from))
275                .collect()
276        })
277        .unwrap_or_default()
278}
luft_core/testing.rs

luft_core/
testing.rs
