1pub mod annotations;
40pub mod fidelity;
41pub mod mcp_mock;
42pub mod overlay_fs;
43pub mod process_tape;
44pub mod tape;
45#[cfg(feature = "testbench-wasi")]
46pub mod wasi_process;
47
48use std::path::PathBuf;
49use std::sync::Arc;
50
51use crate::clock_mock::leak_audit::{self, ClockLeak};
52use crate::clock_mock::{install_override, ClockOverrideGuard, MockClock};
53use crate::egress::reset_egress_policy_for_host;
54
55use overlay_fs::{install_overlay, OverlayFs, OverlayFsGuard};
56use process_tape::{install_process_tape, ProcessTape, ProcessTapeGuard, ProcessTapeMode};
57use tape::{install_recorder, TapeHeader, TapeRecorder, TapeRecorderGuard};
58
59#[derive(Debug, Default, Clone)]
62pub struct Testbench {
63 pub clock: ClockConfig,
64 pub llm: LlmConfig,
65 pub filesystem: FilesystemConfig,
66 pub subprocess: SubprocessConfig,
67 pub network: NetworkConfig,
68 pub tape: TapeConfig,
69}
70
71#[derive(Debug, Default, Clone)]
74pub enum ClockConfig {
75 #[default]
77 Real,
78 Paused { starting_at_ms: i64 },
81}
82
83#[derive(Debug, Default, Clone)]
89pub enum LlmConfig {
90 #[default]
92 Real,
93 Replay { fixture: PathBuf },
95 Record { fixture: PathBuf },
97}
98
99#[derive(Debug, Default, Clone)]
101pub enum FilesystemConfig {
102 #[default]
104 Real,
105 Overlay { worktree: PathBuf },
109}
110
111#[derive(Debug, Default, Clone)]
113pub enum SubprocessConfig {
114 #[default]
116 Real,
117 Record { tape: PathBuf },
120 Replay { tape: PathBuf },
123 WasiToolchain { dir: PathBuf },
131}
132
133#[derive(Debug, Default, Clone)]
136pub enum NetworkConfig {
137 #[default]
139 Real,
140 DenyByDefault {
145 allow: Vec<String>,
148 },
149}
150
151#[derive(Debug, Default, Clone)]
157pub enum TapeConfig {
158 #[default]
159 Off,
160 Emit {
161 path: PathBuf,
162 argv: Vec<String>,
166 script_path: Option<String>,
169 },
170}
171
172impl Testbench {
173 pub fn builder() -> TestbenchBuilder {
175 TestbenchBuilder::default()
176 }
177
178 pub fn activate(self) -> Result<TestbenchSession, TestbenchError> {
181 TestbenchSession::install(self)
182 }
183}
184
185#[derive(Debug, Default, Clone)]
187pub struct TestbenchBuilder {
188 bench: Testbench,
189}
190
191impl TestbenchBuilder {
192 pub fn paused_clock_at_ms(mut self, starting_at_ms: i64) -> Self {
193 self.bench.clock = ClockConfig::Paused { starting_at_ms };
194 self
195 }
196
197 pub fn replay_llm(mut self, fixture: impl Into<PathBuf>) -> Self {
198 self.bench.llm = LlmConfig::Replay {
199 fixture: fixture.into(),
200 };
201 self
202 }
203
204 pub fn record_llm(mut self, fixture: impl Into<PathBuf>) -> Self {
205 self.bench.llm = LlmConfig::Record {
206 fixture: fixture.into(),
207 };
208 self
209 }
210
211 pub fn fs_overlay(mut self, worktree: impl Into<PathBuf>) -> Self {
212 self.bench.filesystem = FilesystemConfig::Overlay {
213 worktree: worktree.into(),
214 };
215 self
216 }
217
218 pub fn record_subprocesses(mut self, tape: impl Into<PathBuf>) -> Self {
219 self.bench.subprocess = SubprocessConfig::Record { tape: tape.into() };
220 self
221 }
222
223 pub fn replay_subprocesses(mut self, tape: impl Into<PathBuf>) -> Self {
224 self.bench.subprocess = SubprocessConfig::Replay { tape: tape.into() };
225 self
226 }
227
228 pub fn wasi_toolchain(mut self, dir: impl Into<PathBuf>) -> Self {
231 self.bench.subprocess = SubprocessConfig::WasiToolchain { dir: dir.into() };
232 self
233 }
234
235 pub fn deny_network(mut self) -> Self {
236 self.bench.network = NetworkConfig::DenyByDefault { allow: Vec::new() };
237 self
238 }
239
240 pub fn allow_network(mut self, allow: impl IntoIterator<Item = String>) -> Self {
241 self.bench.network = NetworkConfig::DenyByDefault {
242 allow: allow.into_iter().collect(),
243 };
244 self
245 }
246
247 pub fn emit_tape(mut self, path: impl Into<PathBuf>) -> Self {
248 self.bench.tape = TapeConfig::Emit {
249 path: path.into(),
250 argv: Vec::new(),
251 script_path: None,
252 };
253 self
254 }
255
256 pub fn emit_tape_for(
257 mut self,
258 path: impl Into<PathBuf>,
259 script_path: Option<String>,
260 argv: Vec<String>,
261 ) -> Self {
262 self.bench.tape = TapeConfig::Emit {
263 path: path.into(),
264 argv,
265 script_path,
266 };
267 self
268 }
269
270 pub fn build(self) -> Testbench {
271 self.bench
272 }
273}
274
275#[must_use = "the testbench tears down on drop; bind the handle to a `_session` local"]
278pub struct TestbenchSession {
279 _clock_leak_scope: Option<leak_audit::ClockLeakScopeGuard>,
280 _clock: Option<ClockOverrideGuard>,
281 _process: Option<ProcessTapeGuard>,
282 _overlay: Option<OverlayFsGuard>,
283 _recorder: Option<TapeRecorderGuard>,
284 process_tape: Option<Arc<ProcessTape>>,
285 overlay: Option<Arc<OverlayFs>>,
286 recorder: Option<Arc<TapeRecorder>>,
287 tape_path: Option<PathBuf>,
288 tape_started_at_unix_ms: Option<i64>,
289 tape_script_path: Option<String>,
290 tape_argv: Vec<String>,
291 subprocess_mode: ProcessTapeMode,
292 subprocess_tape_path: Option<PathBuf>,
293 #[cfg(feature = "testbench-wasi")]
294 _wasi_toolchain: Option<wasi_process::WasiToolchainGuard>,
295 egress_policy_installed: bool,
298}
299
300impl TestbenchSession {
301 fn install(bench: Testbench) -> Result<Self, TestbenchError> {
302 let (clock_leak_scope, clock_guard, started_at_unix_ms) = match bench.clock {
303 ClockConfig::Real => (None, None, None),
304 ClockConfig::Paused { starting_at_ms } => (
305 Some(leak_audit::install_scope()),
306 Some(install_override(MockClock::at_wall_ms(starting_at_ms))),
307 Some(starting_at_ms),
308 ),
309 };
310
311 #[allow(clippy::no_effect_underscore_binding)]
315 let _llm_config = bench.llm;
316
317 #[cfg(feature = "testbench-wasi")]
318 let mut wasi_guard: Option<wasi_process::WasiToolchainGuard> = None;
319
320 let (process_tape, process_guard, subprocess_mode, subprocess_tape_path) =
321 match bench.subprocess {
322 SubprocessConfig::Real => (None, None, ProcessTapeMode::Replay, None),
323 SubprocessConfig::Record { tape } => {
324 let active = Arc::new(ProcessTape::recording());
325 let guard = install_process_tape(Arc::clone(&active));
326 (
327 Some(Arc::clone(&active)),
328 Some(guard),
329 ProcessTapeMode::Record,
330 Some(tape),
331 )
332 }
333 SubprocessConfig::Replay { tape } => {
334 let loaded = ProcessTape::load(&tape).map_err(TestbenchError::Subprocess)?;
335 let active = Arc::new(loaded);
336 let guard = install_process_tape(Arc::clone(&active));
337 (
338 Some(Arc::clone(&active)),
339 Some(guard),
340 ProcessTapeMode::Replay,
341 Some(tape),
342 )
343 }
344 #[cfg(feature = "testbench-wasi")]
345 SubprocessConfig::WasiToolchain { dir } => {
346 if !dir.exists() {
347 return Err(TestbenchError::Subprocess(format!(
348 "wasi toolchain directory does not exist: {}",
349 dir.display()
350 )));
351 }
352 wasi_guard = Some(wasi_process::install_wasi_toolchain(dir));
353 (None, None, ProcessTapeMode::Replay, None)
354 }
355 #[cfg(not(feature = "testbench-wasi"))]
356 SubprocessConfig::WasiToolchain { .. } => {
357 return Err(TestbenchError::Subprocess(
358 "WasiToolchain requires the `testbench-wasi` Cargo feature".to_string(),
359 ));
360 }
361 };
362
363 let (overlay, overlay_guard) = match bench.filesystem {
364 FilesystemConfig::Real => (None, None),
365 FilesystemConfig::Overlay { worktree } => {
366 let overlay = Arc::new(OverlayFs::rooted_at(worktree));
367 let guard = install_overlay(Arc::clone(&overlay));
368 (Some(overlay), Some(guard))
369 }
370 };
371
372 let egress_policy_installed = match bench.network {
373 NetworkConfig::Real => false,
374 NetworkConfig::DenyByDefault { allow } => {
375 crate::egress::install_deny_by_default_policy(&allow)
376 .map_err(|error| TestbenchError::Network(error.to_string()))?;
377 true
378 }
379 };
380
381 let (recorder, recorder_guard, tape_path, tape_argv, tape_script_path) = match bench.tape {
382 TapeConfig::Off => (None, None, None, Vec::new(), None),
383 TapeConfig::Emit {
384 path,
385 argv,
386 script_path,
387 } => {
388 let recorder = Arc::new(TapeRecorder::new());
389 let guard = install_recorder(Arc::clone(&recorder));
390 (
391 Some(Arc::clone(&recorder)),
392 Some(guard),
393 Some(path),
394 argv,
395 script_path,
396 )
397 }
398 };
399
400 Ok(Self {
401 _clock_leak_scope: clock_leak_scope,
402 _clock: clock_guard,
403 _process: process_guard,
404 _overlay: overlay_guard,
405 _recorder: recorder_guard,
406 process_tape,
407 overlay,
408 recorder,
409 tape_path,
410 tape_started_at_unix_ms: started_at_unix_ms,
411 tape_script_path,
412 tape_argv,
413 subprocess_mode,
414 subprocess_tape_path,
415 #[cfg(feature = "testbench-wasi")]
416 _wasi_toolchain: wasi_guard,
417 egress_policy_installed,
418 })
419 }
420
421 pub fn subprocess_mode(&self) -> ProcessTapeMode {
423 self.subprocess_mode
424 }
425
426 pub fn subprocess_tape_path(&self) -> Option<&std::path::Path> {
429 self.subprocess_tape_path.as_deref()
430 }
431
432 pub fn overlay(&self) -> Option<&Arc<OverlayFs>> {
434 self.overlay.as_ref()
435 }
436
437 pub fn process_tape(&self) -> Option<&Arc<ProcessTape>> {
439 self.process_tape.as_ref()
440 }
441
442 pub fn tape_recorder(&self) -> Option<&Arc<TapeRecorder>> {
444 self.recorder.as_ref()
445 }
446
447 pub fn finalize(self) -> Result<TestbenchFinalize, TestbenchError> {
452 let diff = self
453 .overlay
454 .as_ref()
455 .map(|overlay| overlay.diff())
456 .unwrap_or_default();
457 let recorded = if matches!(self.subprocess_mode, ProcessTapeMode::Record) {
458 if let (Some(tape), Some(path)) = (
459 self.process_tape.as_ref(),
460 self.subprocess_tape_path.as_ref(),
461 ) {
462 tape.persist(path).map_err(TestbenchError::Subprocess)?;
463 }
464 self.process_tape
465 .as_ref()
466 .map(|tape| tape.recorded())
467 .unwrap_or_default()
468 } else {
469 Vec::new()
470 };
471 let mut emitted_tape = None;
472 if let (Some(recorder), Some(path)) = (self.recorder.as_ref(), self.tape_path.as_ref()) {
473 let header = TapeHeader::current(
474 self.tape_started_at_unix_ms,
475 self.tape_script_path.clone(),
476 self.tape_argv.clone(),
477 );
478 let tape = recorder.snapshot(header);
479 tape.persist(path).map_err(TestbenchError::Tape)?;
480 emitted_tape = Some(EmittedTape {
481 path: path.clone(),
482 records: tape.records.len(),
483 });
484 }
485 let clock_leaks = leak_audit::drain();
490 Ok(TestbenchFinalize {
492 fs_diff: diff,
493 recorded_subprocesses: recorded,
494 tape: emitted_tape,
495 clock_leaks,
496 })
497 }
498}
499
500impl Drop for TestbenchSession {
501 fn drop(&mut self) {
502 if self.egress_policy_installed {
503 reset_egress_policy_for_host();
504 }
505 }
508}
509
510#[derive(Debug, Default, Clone)]
513pub struct TestbenchFinalize {
514 pub fs_diff: Vec<overlay_fs::DiffEntry>,
515 pub recorded_subprocesses: Vec<process_tape::TapeEntry>,
516 pub tape: Option<EmittedTape>,
517 pub clock_leaks: Vec<ClockLeak>,
522}
523
524#[derive(Debug, Clone)]
526pub struct EmittedTape {
527 pub path: PathBuf,
528 pub records: usize,
529}
530
531#[derive(Debug)]
533pub enum TestbenchError {
534 Subprocess(String),
535 Tape(String),
536 Network(String),
537}
538
539impl std::fmt::Display for TestbenchError {
540 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
541 match self {
542 Self::Subprocess(msg) => write!(f, "testbench subprocess: {msg}"),
543 Self::Tape(msg) => write!(f, "testbench tape: {msg}"),
544 Self::Network(msg) => write!(f, "testbench network: {msg}"),
545 }
546 }
547}
548
549impl std::error::Error for TestbenchError {}
550
551#[cfg(test)]
552mod tests {
553 use super::*;
554 use std::sync::Mutex;
555
556 static ENV_TEST_LOCK: Mutex<()> = Mutex::new(());
557
558 fn serial<F: FnOnce()>(body: F) {
562 let _guard = ENV_TEST_LOCK.lock().unwrap_or_else(|p| p.into_inner());
563 body();
564 }
565
566 #[test]
567 fn paused_clock_pins_now_ms_for_session_lifetime() {
568 serial(|| {
569 let bench = Testbench::builder()
570 .paused_clock_at_ms(1_700_000_000_000)
571 .build();
572 let session = bench.activate().expect("activate");
573 assert_eq!(crate::clock_mock::now_ms(), 1_700_000_000_000);
574 crate::clock_mock::advance(std::time::Duration::from_mins(1));
575 assert_eq!(crate::clock_mock::now_ms(), 1_700_000_060_000);
576 drop(session);
577 assert!(!crate::clock_mock::is_mocked());
579 });
580 }
581
582 #[test]
583 fn deny_by_default_blocks_egress_until_drop() {
584 serial(|| {
585 let _env = crate::egress::test_env_guard();
586 let bench = Testbench::builder().deny_network().build();
587 let session = bench.activate().expect("activate");
588 assert!(
589 crate::egress::check_url("testbench", "https://example.com/x")
590 .expect("policy check")
591 .is_some()
592 );
593 drop(session);
594 assert!(
595 crate::egress::check_url("testbench", "https://example.com/x")
596 .expect("policy check")
597 .is_none()
598 );
599 });
600 }
601
602 #[test]
603 fn finalize_surfaces_clock_leaks_for_contrived_capability() {
604 serial(|| {
605 let bench = Testbench::builder()
606 .paused_clock_at_ms(1_700_000_000_000)
607 .build();
608 let session = bench.activate().expect("activate");
609
610 let _ = leak_audit::wall_now("test/contrived_leak");
614 let _ = leak_audit::instant_now("test/contrived_instant");
615 let _ = leak_audit::wall_now("test/contrived_leak");
616
617 let finalize = session.finalize().expect("finalize");
618 let by_id: std::collections::BTreeMap<&str, &ClockLeak> = finalize
619 .clock_leaks
620 .iter()
621 .map(|leak| (leak.capability_id.as_str(), leak))
622 .collect();
623 let wall = by_id
624 .get("test/contrived_leak")
625 .expect("wall leak surfaced");
626 assert_eq!(wall.count, 2);
627 let inst = by_id
628 .get("test/contrived_instant")
629 .expect("instant leak surfaced");
630 assert_eq!(inst.count, 1);
631
632 let next_session = Testbench::builder()
634 .paused_clock_at_ms(1_700_000_000_000)
635 .build()
636 .activate()
637 .expect("activate next");
638 let next = next_session.finalize().expect("finalize next");
639 assert!(next.clock_leaks.is_empty());
640 });
641 }
642
643 #[test]
644 fn audit_quiet_when_no_mock_is_active() {
645 serial(|| {
646 leak_audit::reset();
647 let _ = leak_audit::wall_now("test/no_mock");
650 let _ = leak_audit::instant_now("test/no_mock");
651 assert!(leak_audit::snapshot().is_empty());
652 });
653 }
654}