Skip to main content

split_brain_harness/
wasm_forge.rs

1/// Phase 4 of the Ephemeral Tool Forge — WASM/WASI sandboxed execution.
2///
3/// Takes a verified `GeneratedTool` (Phase 3 output), compiles it to
4/// `wasm32-wasip1`, executes it in the `wasmtime` CLI sandbox (no network,
5/// no filesystem access), captures stdout, then destroys the binary.
6///
7/// **Status**: production-quality, full test coverage.
8/// **Requires**: `rustup target add wasm32-wasip1` and `wasmtime` in PATH.
9///   Run `sbh doctor` to verify the toolchain is installed.
10/// **CLI entry point**: `sbh forge "<capability>" "<input>"` (via `regenerative_forge`)
11/// and stores fingerprint metrics only.
12///
13/// The supervisor code calls `std::process::Command` — that is the supervisor's
14/// own privilege. Model-generated source code is still forbidden from using it
15/// (enforced by static analysis in Phase 3).
16use std::io::Write as _;
17use std::process::{Command, Stdio};
18use std::time::{Duration, Instant};
19
20/// Hard wall-clock limit on rustc compilation. Model-generated code is stdlib-only
21/// so 60 s is generous; pathological macro/trait-resolution abuse gets cut off here.
22const COMPILE_TIMEOUT: Duration = Duration::from_secs(60);
23
24/// Hard wall-clock limit passed directly to the wasmtime `-W timeout=` flag.
25/// Must be >= the capability constraint max (10 000 ms) and expressed as ms string.
26const WASM_EXEC_TIMEOUT_MS: u64 = 15_000;
27
28use serde::{Deserialize, Serialize};
29
30use crate::capability::{Budget, CapabilityMemoryRecord, CapabilityRequest, ToolMetrics};
31use crate::code_gen::GeneratedTool;
32use crate::input_validation;
33use crate::policy::{self, PolicyState};
34use crate::tool_memory::CapabilityMemory;
35
36// ---------------------------------------------------------------------------
37// WASM main() wrapper
38//
39// Appended to the generated source before compilation. Reads from stdin,
40// calls `run()`, writes JSON to stdout.
41// ---------------------------------------------------------------------------
42
43const WASM_MAIN: &str = r#"
44
45fn main() {
46    use std::io::Read;
47    let mut input = String::new();
48    if std::io::stdin().read_to_string(&mut input).is_err() {
49        std::process::exit(2);
50    }
51    match run(&input) {
52        Ok(output) => print!("{}", output),
53        Err(e) => {
54            eprintln!("run error: {}", e);
55            std::process::exit(1);
56        }
57    }
58}
59"#;
60
61// ---------------------------------------------------------------------------
62// Outcome enums
63// ---------------------------------------------------------------------------
64
65/// Result of the `rustc --target wasm32-wasi` compilation step.
66#[derive(Debug, Clone, Serialize, Deserialize)]
67pub enum CompileOutcome {
68    Success {
69        wasm_bytes: Vec<u8>,
70        compilation_ms: u64,
71    },
72    /// The wasm32-wasi (or wasm32-wasip1) target is not installed.
73    TargetNotInstalled { attempted_target: String },
74    /// rustc ran but returned a non-zero exit code.
75    CompilationFailed { stderr: String, compilation_ms: u64 },
76    /// rustc could not be spawned (not on PATH, OS error, etc.).
77    CompilerNotFound { error: String },
78}
79
80/// Result of executing the compiled WASM binary.
81#[derive(Debug, Clone, Serialize, Deserialize)]
82pub enum ExecuteOutcome {
83    Success {
84        stdout: String,
85        execution_ms: u64,
86    },
87    /// wasmtime binary is not available on PATH.
88    RuntimeNotFound,
89    /// wasmtime ran but the WASM process exited non-zero.
90    ExecutionFailed {
91        stderr: String,
92        exit_code: i32,
93        execution_ms: u64,
94    },
95    /// wasmtime could not be spawned (OS error).
96    RuntimeError {
97        error: String,
98    },
99}
100
101// ---------------------------------------------------------------------------
102// Phase 4 report
103// ---------------------------------------------------------------------------
104
105#[derive(Debug, Serialize, Deserialize, Clone)]
106pub struct WasmMetrics {
107    pub compilation_ms: u64,
108    pub execution_ms: u64,
109    pub wasm_binary_bytes: usize,
110    pub input_bytes: usize,
111    pub output_bytes: usize,
112}
113
114#[derive(Debug, Serialize, Deserialize, Clone)]
115pub struct WasmExecutionReport {
116    pub accepted: bool,
117    pub rejection_reasons: Vec<String>,
118    /// True when compilation succeeded.
119    pub compiled: bool,
120    #[serde(skip_serializing_if = "Option::is_none")]
121    pub compilation_error: Option<String>,
122    /// True when the WASM executed successfully (exit code 0).
123    pub executed: bool,
124    #[serde(skip_serializing_if = "Option::is_none")]
125    pub execution_error: Option<String>,
126    /// Captured stdout from the WASM process.
127    #[serde(skip_serializing_if = "Option::is_none")]
128    pub output: Option<String>,
129    /// Binary destroyed after use: always true once compilation completes.
130    pub destroyed: bool,
131    pub metrics: WasmMetrics,
132    #[serde(skip_serializing_if = "Option::is_none")]
133    pub memory_update: Option<CapabilityMemoryRecord>,
134}
135
136// ---------------------------------------------------------------------------
137// Compiler and executor traits — injectable for testing
138// ---------------------------------------------------------------------------
139
140pub trait WasmCompiler: Send + Sync {
141    fn compile(&self, source: &str) -> CompileOutcome;
142}
143
144pub trait WasmExecutor: Send + Sync {
145    fn execute(&self, wasm_bytes: &[u8], input: &str) -> ExecuteOutcome;
146}
147
148// ---------------------------------------------------------------------------
149// Real compiler: rustc --target wasm32-wasi
150// ---------------------------------------------------------------------------
151
152pub struct RustcCompiler;
153
154impl WasmCompiler for RustcCompiler {
155    fn compile(&self, source: &str) -> CompileOutcome {
156        // Detect installed WASM target (wasm32-wasip1 is the new name, wasm32-wasi the old)
157        let target = match detect_wasm_target() {
158            Some(t) => t,
159            None => {
160                return CompileOutcome::TargetNotInstalled {
161                    attempted_target: "wasm32-wasip1 / wasm32-wasi".into(),
162                }
163            }
164        };
165
166        // Write source + wrapper to a temp file
167        let tmp_dir = std::env::temp_dir().join(format!("sbh-wasm-{}", monotonic_id()));
168        if std::fs::create_dir_all(&tmp_dir).is_err() {
169            return CompileOutcome::CompilationFailed {
170                stderr: "failed to create temp directory".into(),
171                compilation_ms: 0,
172            };
173        }
174
175        let src_path = tmp_dir.join("tool.rs");
176        let wasm_path = tmp_dir.join("tool.wasm");
177        let full_source = format!("{}\n{}", source, WASM_MAIN);
178
179        if std::fs::write(&src_path, &full_source).is_err() {
180            let _ = std::fs::remove_dir_all(&tmp_dir);
181            return CompileOutcome::CompilationFailed {
182                stderr: "failed to write source file".into(),
183                compilation_ms: 0,
184            };
185        }
186
187        let start = Instant::now();
188
189        // Spawn rustc as a child so we can enforce a wall-clock timeout.
190        let child = Command::new("rustc")
191            .args(["--target", target, "--edition", "2021", "-o"])
192            .arg(&wasm_path)
193            .arg(&src_path)
194            .stdout(Stdio::piped())
195            .stderr(Stdio::piped())
196            .spawn();
197
198        let child = match child {
199            Err(e) => {
200                let _ = std::fs::remove_file(&src_path);
201                let _ = std::fs::remove_dir_all(&tmp_dir);
202                return CompileOutcome::CompilerNotFound {
203                    error: format!("could not spawn rustc: {e}"),
204                };
205            }
206            Ok(c) => c,
207        };
208
209        // Move child into a thread; wait for output or kill after COMPILE_TIMEOUT.
210        let pid = child.id();
211        let (tx, rx) = std::sync::mpsc::channel::<std::io::Result<std::process::Output>>();
212        std::thread::spawn(move || {
213            let _ = tx.send(child.wait_with_output());
214        });
215
216        let result = match rx.recv_timeout(COMPILE_TIMEOUT) {
217            Ok(r) => r,
218            Err(_) => {
219                // Kill the rustc process by PID — it's already moved into the thread,
220                // so we signal it via the OS. Cleanup is best-effort.
221                let _ = Command::new("kill").args(["-9", &pid.to_string()]).status();
222                let _ = std::fs::remove_file(&src_path);
223                let _ = std::fs::remove_dir_all(&tmp_dir);
224                return CompileOutcome::CompilationFailed {
225                    stderr: format!("rustc timed out after {}s", COMPILE_TIMEOUT.as_secs()),
226                    compilation_ms: COMPILE_TIMEOUT.as_millis() as u64,
227                };
228            }
229        };
230
231        let compilation_ms = start.elapsed().as_millis() as u64;
232
233        // Remove source immediately (never persist model-generated code)
234        let _ = std::fs::remove_file(&src_path);
235
236        match result {
237            Err(e) => {
238                let _ = std::fs::remove_dir_all(&tmp_dir);
239                CompileOutcome::CompilerNotFound {
240                    error: format!("rustc wait error: {e}"),
241                }
242            }
243            Ok(out) if !out.status.success() => {
244                let _ = std::fs::remove_dir_all(&tmp_dir);
245                CompileOutcome::CompilationFailed {
246                    stderr: String::from_utf8_lossy(&out.stderr)
247                        .chars()
248                        .take(2048)
249                        .collect(),
250                    compilation_ms,
251                }
252            }
253            Ok(_) => {
254                let wasm_bytes = std::fs::read(&wasm_path).unwrap_or_default();
255                // Destroy the compiled WASM — metrics only survive
256                let _ = std::fs::remove_dir_all(&tmp_dir);
257                CompileOutcome::Success {
258                    wasm_bytes,
259                    compilation_ms,
260                }
261            }
262        }
263    }
264}
265
266/// Detect an installed WASM/WASI target via `rustup target list --installed`.
267/// Returns the preferred target name, or None if neither is installed.
268fn detect_wasm_target() -> Option<&'static str> {
269    let out = Command::new("rustup")
270        .args(["target", "list", "--installed"])
271        .output()
272        .ok()?;
273    let targets = String::from_utf8_lossy(&out.stdout);
274    if targets.contains("wasm32-wasip1") {
275        Some("wasm32-wasip1")
276    } else if targets.contains("wasm32-wasi") {
277        Some("wasm32-wasi")
278    } else {
279        None
280    }
281}
282
283// ---------------------------------------------------------------------------
284// Real executor: wasmtime CLI subprocess
285// ---------------------------------------------------------------------------
286
287pub struct WasmtimeCli;
288
289impl WasmExecutor for WasmtimeCli {
290    fn execute(&self, wasm_bytes: &[u8], input: &str) -> ExecuteOutcome {
291        // Write WASM to a temp file — destroyed immediately after the call
292        let tmp_wasm = std::env::temp_dir().join(format!("sbh-run-{}.wasm", monotonic_id()));
293        if std::fs::write(&tmp_wasm, wasm_bytes).is_err() {
294            return ExecuteOutcome::RuntimeError {
295                error: "failed to write wasm to temp file".into(),
296            };
297        }
298
299        let start = Instant::now();
300
301        // Spawn wasmtime and feed input via stdin. We wait for the process to
302        // finish BEFORE deleting the WASM file — spawning is non-blocking so
303        // deleting first would cause a race where wasmtime tries to open a
304        // file that no longer exists.
305        let spawn_result = Command::new("wasmtime")
306            .arg(format!("-W timeout={WASM_EXEC_TIMEOUT_MS}ms"))
307            .arg("--")
308            .arg(&tmp_wasm)
309            .stdin(Stdio::piped())
310            .stdout(Stdio::piped())
311            .stderr(Stdio::piped())
312            .spawn();
313
314        let output_result = match spawn_result {
315            Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
316                let _ = std::fs::remove_file(&tmp_wasm);
317                return ExecuteOutcome::RuntimeNotFound;
318            }
319            Err(e) => {
320                let _ = std::fs::remove_file(&tmp_wasm);
321                return ExecuteOutcome::RuntimeError {
322                    error: format!("spawn error: {e}"),
323                };
324            }
325            Ok(mut child) => {
326                if let Some(mut stdin) = child.stdin.take() {
327                    let _ = stdin.write_all(input.as_bytes());
328                }
329                child.wait_with_output()
330            }
331        };
332
333        // WASM binary destroyed after the process has exited and closed the fd
334        let _ = std::fs::remove_file(&tmp_wasm);
335
336        match output_result {
337            Err(e) => ExecuteOutcome::RuntimeError {
338                error: format!("wait error: {e}"),
339            },
340            Ok(out) => {
341                let execution_ms = start.elapsed().as_millis() as u64;
342                let stdout = String::from_utf8_lossy(&out.stdout)
343                    .chars()
344                    .take(65_536)
345                    .collect();
346                let exit_code = out.status.code().unwrap_or(-1);
347
348                if out.status.success() {
349                    ExecuteOutcome::Success {
350                        stdout,
351                        execution_ms,
352                    }
353                } else {
354                    let stderr = String::from_utf8_lossy(&out.stderr)
355                        .chars()
356                        .take(1024)
357                        .collect();
358                    ExecuteOutcome::ExecutionFailed {
359                        stderr,
360                        exit_code,
361                        execution_ms,
362                    }
363                }
364            }
365        }
366    }
367}
368
369// ---------------------------------------------------------------------------
370// Supervisor
371// ---------------------------------------------------------------------------
372
373pub struct WasmForge {
374    budget: Budget,
375    state: PolicyState,
376    pub memory: CapabilityMemory,
377    compiler: Box<dyn WasmCompiler>,
378    executor: Box<dyn WasmExecutor>,
379    session_log: Vec<WasmExecutionReport>,
380}
381
382impl WasmForge {
383    pub fn new() -> Self {
384        Self::with_deps(Box::new(RustcCompiler), Box::new(WasmtimeCli))
385    }
386
387    pub fn with_deps(compiler: Box<dyn WasmCompiler>, executor: Box<dyn WasmExecutor>) -> Self {
388        Self {
389            // Tighter budget for Phase 4: compilation is expensive
390            budget: Budget {
391                max_tools_per_session: 2,
392                max_total_runtime_ms: 120_000, // 2 minutes total
393                require_approval_after_failures: 1,
394            },
395            state: PolicyState::default(),
396            memory: CapabilityMemory::new(),
397            compiler,
398            executor,
399            session_log: vec![],
400        }
401    }
402
403    pub fn audit(&self) -> &[WasmExecutionReport] {
404        &self.session_log
405    }
406
407    /// Process a Phase 3 GeneratedTool:
408    ///   validate → policy check → verify tool passed Phase 3 →
409    ///   compile → execute → destroy → memory update
410    pub fn handle(
411        &mut self,
412        req: &CapabilityRequest,
413        tool: &GeneratedTool,
414        input: &str,
415    ) -> WasmExecutionReport {
416        let report = self.handle_inner(req, tool, input);
417        self.session_log.push(report.clone());
418        report
419    }
420
421    fn handle_inner(
422        &mut self,
423        req: &CapabilityRequest,
424        tool: &GeneratedTool,
425        input: &str,
426    ) -> WasmExecutionReport {
427        // Input validation
428        if let Err(e) = input_validation::validate_forge_input(input) {
429            return rejected(vec![format!("input validation: {e}")]);
430        }
431        if let Err(e) = input_validation::validate_capability_fields(req) {
432            return rejected(vec![format!("capability field validation: {e}")]);
433        }
434
435        // Budget check
436        if let Some(reason) = self.state.budget_exceeded(&self.budget) {
437            return rejected(vec![reason]);
438        }
439
440        // Policy checks
441        let violations = policy::check_request(req);
442        if !violations.is_empty() {
443            return rejected(violations.into_iter().map(|v| v.detail).collect());
444        }
445
446        // Require Phase 3 verification to have passed
447        if !tool.static_analysis.passed {
448            let reasons: Vec<String> = tool
449                .static_analysis
450                .violations
451                .iter()
452                .map(|v| format!("static_analysis: {} at line {}", v.kind, v.line))
453                .collect();
454            return rejected(reasons);
455        }
456        if !tool.tests_included {
457            return rejected(vec![
458                "generated tool does not include the required minimum of 2 #[test] functions"
459                    .into(),
460            ]);
461        }
462
463        // Compile
464        let compile_outcome = self.compiler.compile(&tool.source);
465        let (wasm_bytes, compilation_ms) = match compile_outcome {
466            CompileOutcome::Success {
467                wasm_bytes,
468                compilation_ms,
469            } => (wasm_bytes, compilation_ms),
470            CompileOutcome::TargetNotInstalled { attempted_target } => {
471                let metrics = zero_metrics(input);
472                self.state.record_run(&ToolMetrics {
473                    success: false,
474                    input_bytes: input.len(),
475                    ..Default::default()
476                });
477                return WasmExecutionReport {
478                    accepted: true,
479                    rejection_reasons: vec![],
480                    compiled: false,
481                    compilation_error: Some(format!(
482                        "WASM target not installed: {attempted_target} — \
483                         run: rustup target add wasm32-wasip1"
484                    )),
485                    executed: false,
486                    execution_error: None,
487                    output: None,
488                    destroyed: true,
489                    metrics,
490                    memory_update: None,
491                };
492            }
493            CompileOutcome::CompilationFailed {
494                stderr,
495                compilation_ms,
496            } => {
497                let metrics = WasmMetrics {
498                    compilation_ms,
499                    ..zero_metrics(input)
500                };
501                self.state.record_run(&ToolMetrics {
502                    success: false,
503                    input_bytes: input.len(),
504                    runtime_ms: compilation_ms,
505                    ..Default::default()
506                });
507                return WasmExecutionReport {
508                    accepted: true,
509                    rejection_reasons: vec![],
510                    compiled: false,
511                    compilation_error: Some(stderr),
512                    executed: false,
513                    execution_error: None,
514                    output: None,
515                    destroyed: true,
516                    metrics,
517                    memory_update: None,
518                };
519            }
520            CompileOutcome::CompilerNotFound { error } => {
521                let metrics = zero_metrics(input);
522                self.state.record_run(&ToolMetrics {
523                    success: false,
524                    input_bytes: input.len(),
525                    ..Default::default()
526                });
527                return WasmExecutionReport {
528                    accepted: true,
529                    rejection_reasons: vec![],
530                    compiled: false,
531                    compilation_error: Some(error),
532                    executed: false,
533                    execution_error: None,
534                    output: None,
535                    destroyed: true,
536                    metrics,
537                    memory_update: None,
538                };
539            }
540        };
541
542        // WASM binary size (before destruction by executor)
543        let wasm_binary_bytes = wasm_bytes.len();
544
545        // Execute
546        let execute_outcome = self.executor.execute(&wasm_bytes, input);
547        // wasm_bytes is no longer needed — drop it now
548        drop(wasm_bytes);
549
550        let (executed, stdout, execution_ms, execution_error) = match execute_outcome {
551            ExecuteOutcome::Success {
552                stdout,
553                execution_ms,
554            } => (true, Some(stdout), execution_ms, None),
555            ExecuteOutcome::RuntimeNotFound => (
556                false,
557                None,
558                0,
559                Some("wasmtime not found on PATH — install from https://wasmtime.dev".into()),
560            ),
561            ExecuteOutcome::ExecutionFailed {
562                stderr,
563                execution_ms,
564                ..
565            } => (false, None, execution_ms, Some(stderr)),
566            ExecuteOutcome::RuntimeError { error } => (false, None, 0, Some(error)),
567        };
568
569        let success = executed && execution_error.is_none();
570        let tool_metrics = ToolMetrics {
571            runtime_ms: compilation_ms + execution_ms,
572            input_bytes: input.len(),
573            output_bytes: stdout.as_deref().map(|s| s.len()).unwrap_or(0),
574            success,
575        };
576        self.state.record_run(&tool_metrics);
577
578        let metrics = WasmMetrics {
579            compilation_ms,
580            execution_ms,
581            wasm_binary_bytes,
582            input_bytes: input.len(),
583            output_bytes: stdout.as_deref().map(|s| s.len()).unwrap_or(0),
584        };
585
586        // Update capability memory on success
587        let memory_update = if success {
588            let signature = CapabilityMemory::derive_signature(req);
589            let record = CapabilityMemoryRecord {
590                problem_signature: signature,
591                solution_pattern: format!("wasm:{}", req.capability),
592                input_shape: shape_token(&req.input_contract),
593                output_shape: shape_token(&req.output_contract),
594                constraints: req.constraints.clone(),
595            };
596            self.memory.upsert(record.clone(), &tool_metrics);
597            Some(record)
598        } else {
599            None
600        };
601
602        WasmExecutionReport {
603            accepted: true,
604            rejection_reasons: vec![],
605            compiled: true,
606            compilation_error: None,
607            executed,
608            execution_error,
609            output: stdout,
610            destroyed: true,
611            metrics,
612            memory_update,
613        }
614    }
615
616    pub fn tools_invoked(&self) -> usize {
617        self.state.tools_invoked
618    }
619}
620
621impl Default for WasmForge {
622    fn default() -> Self {
623        Self::new()
624    }
625}
626
627// ---------------------------------------------------------------------------
628// Helpers
629// ---------------------------------------------------------------------------
630
631fn rejected(reasons: Vec<String>) -> WasmExecutionReport {
632    WasmExecutionReport {
633        accepted: false,
634        rejection_reasons: reasons,
635        compiled: false,
636        compilation_error: None,
637        executed: false,
638        execution_error: None,
639        output: None,
640        destroyed: false,
641        metrics: WasmMetrics {
642            compilation_ms: 0,
643            execution_ms: 0,
644            wasm_binary_bytes: 0,
645            input_bytes: 0,
646            output_bytes: 0,
647        },
648        memory_update: None,
649    }
650}
651
652fn zero_metrics(input: &str) -> WasmMetrics {
653    WasmMetrics {
654        compilation_ms: 0,
655        execution_ms: 0,
656        wasm_binary_bytes: 0,
657        input_bytes: input.len(),
658        output_bytes: 0,
659    }
660}
661
662fn shape_token(contract: &str) -> String {
663    contract
664        .split_whitespace()
665        .take(3)
666        .map(|w| {
667            w.to_lowercase()
668                .trim_matches(|c: char| !c.is_alphanumeric())
669                .to_string()
670        })
671        .filter(|s| !s.is_empty())
672        .collect::<Vec<_>>()
673        .join("_")
674}
675
676fn monotonic_id() -> String {
677    let pid = std::process::id();
678    std::time::SystemTime::now()
679        .duration_since(std::time::UNIX_EPOCH)
680        .map(|d| format!("{}-{}{}", pid, d.as_secs(), d.subsec_nanos()))
681        .unwrap_or_else(|_| format!("{}-0", pid))
682}
683
684// ---------------------------------------------------------------------------
685// Tests — use mock compiler + executor so CI works without wasm32-wasi
686// ---------------------------------------------------------------------------
687
688#[cfg(test)]
689mod tests {
690    use super::*;
691    use crate::capability::CapabilityConstraints;
692    use crate::static_analysis;
693
694    // --- Mock compiler ---
695
696    struct MockCompiler(CompileOutcome);
697
698    impl WasmCompiler for MockCompiler {
699        fn compile(&self, _source: &str) -> CompileOutcome {
700            self.0.clone()
701        }
702    }
703
704    // --- Mock executor ---
705
706    struct MockExecutor(ExecuteOutcome);
707
708    impl WasmExecutor for MockExecutor {
709        fn execute(&self, _bytes: &[u8], _input: &str) -> ExecuteOutcome {
710            self.0.clone()
711        }
712    }
713
714    // --- Test helpers ---
715
716    fn clean_req(cap: &str) -> CapabilityRequest {
717        CapabilityRequest {
718            kind: "capability_request".into(),
719            capability: cap.into(),
720            input_contract: "utf8 text".into(),
721            output_contract: "json object".into(),
722            constraints: CapabilityConstraints::default(),
723            reason: "text reasoning insufficient".into(),
724        }
725    }
726
727    fn verified_tool() -> GeneratedTool {
728        let source = r#"pub fn run(input: &str) -> Result<String, String> {
729    let c = input.split_whitespace().count();
730    Ok(format!("{\"count\":{}}", c))
731}
732#[test] fn t1() { assert!(run("a b").is_ok()); }
733#[test] fn t2() { assert!(run("").is_ok()); }"#;
734        GeneratedTool {
735            source: source.into(),
736            function_name: "run".into(),
737            tests_included: true,
738            test_count: 2,
739            static_analysis: static_analysis::check(source),
740        }
741    }
742
743    fn unverified_tool_unsafe() -> GeneratedTool {
744        let source = r#"pub fn run(input: &str) -> Result<String, String> {
745    unsafe { }
746    Ok("ok".into())
747}
748#[test] fn t1() {}
749#[test] fn t2() {}"#;
750        GeneratedTool {
751            source: source.into(),
752            function_name: "run".into(),
753            tests_included: true,
754            test_count: 2,
755            static_analysis: static_analysis::check(source),
756        }
757    }
758
759    fn unverified_tool_no_tests() -> GeneratedTool {
760        let source = "pub fn run(input: &str) -> Result<String, String> { Ok(\"ok\".into()) }";
761        GeneratedTool {
762            source: source.into(),
763            function_name: "run".into(),
764            tests_included: false,
765            test_count: 0,
766            static_analysis: static_analysis::check(source),
767        }
768    }
769
770    fn forge_with(compile: CompileOutcome, exec: ExecuteOutcome) -> WasmForge {
771        WasmForge::with_deps(
772            Box::new(MockCompiler(compile)),
773            Box::new(MockExecutor(exec)),
774        )
775    }
776
777    const MOCK_WASM: &[u8] = b"\x00asm\x01\x00\x00\x00"; // minimal WASM magic bytes
778
779    // --- Acceptance path ---
780
781    #[test]
782    fn successful_compile_and_execute() {
783        let mut forge = forge_with(
784            CompileOutcome::Success {
785                wasm_bytes: MOCK_WASM.to_vec(),
786                compilation_ms: 800,
787            },
788            ExecuteOutcome::Success {
789                stdout: r#"{"count":2}"#.into(),
790                execution_ms: 12,
791            },
792        );
793        let report = forge.handle(&clean_req("word_count"), &verified_tool(), "hello world");
794        assert!(report.accepted);
795        assert!(report.compiled);
796        assert!(report.executed);
797        assert!(report.destroyed, "binary must be destroyed after execution");
798        assert_eq!(report.output.as_deref(), Some(r#"{"count":2}"#));
799        assert!(report.memory_update.is_some());
800    }
801
802    // --- Rejection before compilation ---
803
804    #[test]
805    fn rejects_tool_that_failed_static_analysis() {
806        let mut forge = forge_with(
807            CompileOutcome::Success {
808                wasm_bytes: MOCK_WASM.to_vec(),
809                compilation_ms: 0,
810            },
811            ExecuteOutcome::Success {
812                stdout: "ok".into(),
813                execution_ms: 0,
814            },
815        );
816        let report = forge.handle(&clean_req("x"), &unverified_tool_unsafe(), "input");
817        assert!(!report.accepted);
818        assert!(report
819            .rejection_reasons
820            .iter()
821            .any(|r| r.contains("static_analysis")));
822    }
823
824    #[test]
825    fn rejects_tool_without_tests() {
826        let mut forge = forge_with(
827            CompileOutcome::Success {
828                wasm_bytes: MOCK_WASM.to_vec(),
829                compilation_ms: 0,
830            },
831            ExecuteOutcome::Success {
832                stdout: "ok".into(),
833                execution_ms: 0,
834            },
835        );
836        let report = forge.handle(&clean_req("x"), &unverified_tool_no_tests(), "input");
837        assert!(!report.accepted);
838        assert!(report.rejection_reasons[0].contains("#[test]"));
839    }
840
841    #[test]
842    fn rejects_policy_violation() {
843        let mut req = clean_req("fetch_url");
844        req.constraints.no_network = false;
845        let mut forge = forge_with(
846            CompileOutcome::Success {
847                wasm_bytes: MOCK_WASM.to_vec(),
848                compilation_ms: 0,
849            },
850            ExecuteOutcome::RuntimeNotFound,
851        );
852        let report = forge.handle(&req, &verified_tool(), "input");
853        assert!(!report.accepted);
854        assert!(report
855            .rejection_reasons
856            .iter()
857            .any(|r| r.contains("no_network")));
858    }
859
860    #[test]
861    fn rejects_oversized_input() {
862        let mut forge = forge_with(
863            CompileOutcome::Success {
864                wasm_bytes: MOCK_WASM.to_vec(),
865                compilation_ms: 0,
866            },
867            ExecuteOutcome::Success {
868                stdout: "ok".into(),
869                execution_ms: 0,
870            },
871        );
872        let big = "x".repeat(crate::input_validation::MAX_FORGE_INPUT_BYTES + 1);
873        let report = forge.handle(&clean_req("x"), &verified_tool(), &big);
874        assert!(!report.accepted);
875        assert!(report.rejection_reasons[0].contains("input validation"));
876    }
877
878    #[test]
879    fn rejects_when_budget_exhausted() {
880        let mut forge = WasmForge {
881            budget: Budget {
882                max_tools_per_session: 1,
883                ..Budget::default()
884            },
885            state: PolicyState::default(),
886            memory: CapabilityMemory::new(),
887            compiler: Box::new(MockCompiler(CompileOutcome::Success {
888                wasm_bytes: MOCK_WASM.to_vec(),
889                compilation_ms: 0,
890            })),
891            executor: Box::new(MockExecutor(ExecuteOutcome::Success {
892                stdout: "ok".into(),
893                execution_ms: 0,
894            })),
895            session_log: vec![],
896        };
897        forge.handle(&clean_req("x"), &verified_tool(), "a");
898        let report = forge.handle(&clean_req("x"), &verified_tool(), "b");
899        assert!(!report.accepted);
900        assert!(report.rejection_reasons[0].contains("session tool limit"));
901    }
902
903    // --- Compilation failure paths ---
904
905    #[test]
906    fn compilation_failure_reports_stderr() {
907        let mut forge = forge_with(
908            CompileOutcome::CompilationFailed {
909                stderr: "error[E0001]: syntax error".into(),
910                compilation_ms: 300,
911            },
912            ExecuteOutcome::RuntimeNotFound,
913        );
914        let report = forge.handle(&clean_req("x"), &verified_tool(), "input");
915        assert!(report.accepted);
916        assert!(!report.compiled);
917        assert!(!report.executed);
918        assert!(report.destroyed, "nothing to destroy but flag must be set");
919        assert!(report
920            .compilation_error
921            .as_deref()
922            .unwrap_or("")
923            .contains("syntax error"));
924    }
925
926    #[test]
927    fn target_not_installed_returns_clear_message() {
928        let mut forge = forge_with(
929            CompileOutcome::TargetNotInstalled {
930                attempted_target: "wasm32-wasip1".into(),
931            },
932            ExecuteOutcome::RuntimeNotFound,
933        );
934        let report = forge.handle(&clean_req("x"), &verified_tool(), "input");
935        assert!(report.accepted);
936        assert!(!report.compiled);
937        let err = report.compilation_error.unwrap_or_default();
938        assert!(err.contains("wasm32-wasip1") || err.contains("target not installed"));
939    }
940
941    // --- Execution failure paths ---
942
943    #[test]
944    fn runtime_not_found_does_not_panic() {
945        let mut forge = forge_with(
946            CompileOutcome::Success {
947                wasm_bytes: MOCK_WASM.to_vec(),
948                compilation_ms: 500,
949            },
950            ExecuteOutcome::RuntimeNotFound,
951        );
952        let report = forge.handle(&clean_req("x"), &verified_tool(), "input");
953        assert!(report.accepted);
954        assert!(report.compiled);
955        assert!(!report.executed);
956        assert!(report
957            .execution_error
958            .as_deref()
959            .unwrap_or("")
960            .contains("wasmtime"));
961    }
962
963    #[test]
964    fn execution_failure_reports_stderr() {
965        let mut forge = forge_with(
966            CompileOutcome::Success {
967                wasm_bytes: MOCK_WASM.to_vec(),
968                compilation_ms: 500,
969            },
970            ExecuteOutcome::ExecutionFailed {
971                stderr: "runtime trap".into(),
972                exit_code: 1,
973                execution_ms: 5,
974            },
975        );
976        let report = forge.handle(&clean_req("x"), &verified_tool(), "input");
977        assert!(report.compiled);
978        assert!(!report.executed);
979        assert!(report
980            .execution_error
981            .as_deref()
982            .unwrap_or("")
983            .contains("runtime trap"));
984    }
985
986    // --- Memory and audit ---
987
988    #[test]
989    fn memory_not_updated_when_execution_fails() {
990        let mut forge = forge_with(
991            CompileOutcome::Success {
992                wasm_bytes: MOCK_WASM.to_vec(),
993                compilation_ms: 0,
994            },
995            ExecuteOutcome::RuntimeNotFound,
996        );
997        forge.handle(&clean_req("x"), &verified_tool(), "input");
998        assert_eq!(forge.memory.len(), 0);
999    }
1000
1001    #[test]
1002    fn session_log_records_all_calls() {
1003        let mut forge = forge_with(
1004            CompileOutcome::Success {
1005                wasm_bytes: MOCK_WASM.to_vec(),
1006                compilation_ms: 0,
1007            },
1008            ExecuteOutcome::Success {
1009                stdout: "ok".into(),
1010                execution_ms: 0,
1011            },
1012        );
1013        forge.handle(&clean_req("x"), &verified_tool(), "a");
1014        let mut req2 = clean_req("y");
1015        req2.constraints.no_network = false;
1016        forge.handle(&req2, &verified_tool(), "b");
1017        assert_eq!(forge.audit().len(), 2);
1018    }
1019
1020    #[test]
1021    fn wasm_wrapper_contains_stdin_read_and_run_call() {
1022        assert!(WASM_MAIN.contains("read_to_string"));
1023        assert!(WASM_MAIN.contains("run(&input)"));
1024        assert!(WASM_MAIN.contains("process::exit"));
1025    }
1026
1027    #[test]
1028    fn metrics_record_binary_bytes() {
1029        let wasm = vec![0u8; 42_000];
1030        let mut forge = forge_with(
1031            CompileOutcome::Success {
1032                wasm_bytes: wasm,
1033                compilation_ms: 900,
1034            },
1035            ExecuteOutcome::Success {
1036                stdout: "result".into(),
1037                execution_ms: 15,
1038            },
1039        );
1040        let report = forge.handle(&clean_req("x"), &verified_tool(), "input");
1041        assert_eq!(report.metrics.wasm_binary_bytes, 42_000);
1042        assert_eq!(report.metrics.compilation_ms, 900);
1043        assert_eq!(report.metrics.execution_ms, 15);
1044    }
1045}