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