fix: account for generator handler allocations

crates/zapcode-core/src/sandbox.rs

@@ -50,10 +50,22 @@ impl ResourceTracker {
     }
 
     pub fn track_allocation(&mut self, limits: &ResourceLimits) -> crate::error::Result<()> {
-        self.allocations += 1;
-        if self.allocations > limits.max_allocations {
+        self.track_allocations(limits, 1)
+    }
+
+    pub fn track_allocations(
+        &mut self,
+        limits: &ResourceLimits,
+        count: usize,
+    ) -> crate::error::Result<()> {
+        let allocations = self
+            .allocations
+            .checked_add(count)
+            .ok_or(crate::ZapcodeError::AllocationLimitExceeded)?;
+        if allocations > limits.max_allocations {
             return Err(crate::ZapcodeError::AllocationLimitExceeded);
         }
+        self.allocations = allocations;
         Ok(())
     }
 

crates/zapcode-core/src/vm/mod.rs

@@ -235,6 +235,16 @@ impl Vm {
     }
 
     fn try_handle_error(&mut self, err: &ZapcodeError, min_frame_depth: usize) -> Result<bool> {
+        if matches!(
+            err,
+            ZapcodeError::MemoryLimitExceeded(_)
+                | ZapcodeError::TimeLimitExceeded
+                | ZapcodeError::StackOverflow(_)
+                | ZapcodeError::AllocationLimitExceeded
+        ) {
+            return Ok(false);
+        }
+
         let Some(try_info) = self.take_try_handler(min_frame_depth) else {
             return Ok(false);
         };
@@ -1015,6 +1025,7 @@ impl Vm {
             None => {
                 let func = &self.program.functions[func_idx];
                 self.tracker.push_frame();
+                self.tracker.check_stack(&self.limits)?;
                 let mut locals = Vec::with_capacity(func.local_count);
                 for param in func.params.iter() {
                     match param {
@@ -1053,7 +1064,14 @@ impl Vm {
                 self.run_generator_until_yield_or_return(gen_obj)
             }
             Some(suspended) => {
+                let handler_count = self
+                    .generator_try_handlers
+                    .get(&gen_obj.id)
+                    .map_or(0, Vec::len);
+                self.tracker
+                    .track_allocations(&self.limits, handler_count)?;
                 self.tracker.push_frame();
+                self.tracker.check_stack(&self.limits)?;
                 let stack_base = self.stack.len();
                 for val in &suspended.stack {
                     self.push(val.clone())?;
@@ -1134,6 +1152,21 @@ impl Vm {
             }
             let instr = instructions[frame.ip].clone();
             if matches!(instr, Instruction::Yield) {
+                let frame_depth = self.frames.len();
+                let frame_stack_base = self.current_frame().stack_base;
+                let handler_count = self
+                    .try_stack
+                    .iter()
+                    .filter(|try_info| try_info.frame_depth == frame_depth)
+                    .count();
+                let frame_stack_len = self.stack.len().saturating_sub(frame_stack_base + 1);
+                let map_entry_count = usize::from(handler_count > 0);
+                self.tracker.track_allocations(
+                    &self.limits,
+                    handler_count
+                        .saturating_add(frame_stack_len)
+                        .saturating_add(map_entry_count),
+                )?;
                 self.current_frame_mut().ip += 1;
                 let yielded_value = self.pop()?;
                 let frame_depth = self.frames.len();
@@ -2087,6 +2120,7 @@ impl Vm {
 
             // Error handling
             Instruction::SetupTry(catch_ip, _) => {
+                self.tracker.track_allocation(&self.limits)?;
                 self.try_stack.push(TryInfo {
                     catch_ip,
                     frame_depth: self.frames.len(),

crates/zapcode-core/tests/error_handling.rs

@@ -1,5 +1,94 @@
-use zapcode_core::vm::eval_ts;
-use zapcode_core::Value;
+use zapcode_core::vm::{eval_ts, VmState};
+use zapcode_core::{ResourceLimits, Value, ZapcodeError, ZapcodeRun};
+
+#[test]
+fn test_try_handler_respects_allocation_limit() {
+    let runner = ZapcodeRun::new(
+        "try {} catch (e) {}".to_string(),
+        Vec::new(),
+        Vec::new(),
+        ResourceLimits {
+            max_allocations: 0,
+            ..ResourceLimits::default()
+        },
+    )
+    .unwrap();
+
+    assert!(matches!(
+        runner.start(Vec::new()),
+        Err(ZapcodeError::AllocationLimitExceeded)
+    ));
+}
+
+#[test]
+fn test_generator_handler_allocation_limit_is_not_catchable() {
+    let code = r#"
+        function* values() {
+            try {
+                try {
+                    try {
+                        try {
+                            try { yield 1; } catch (e) {}
+                        } catch (e) {}
+                    } catch (e) {}
+                } catch (e) {}
+            } catch (e) {}
+        }
+        const iterator = values();
+        iterator.next();
+        let outcome = 1;
+        try { iterator.next(); } catch (e) { outcome = 999; }
+        outcome
+    "#;
+
+    let run_with_limit = |max_allocations| {
+        let runner = ZapcodeRun::new(
+            code.to_string(),
+            Vec::new(),
+            Vec::new(),
+            ResourceLimits {
+                max_allocations,
+                ..ResourceLimits::default()
+            },
+        )
+        .unwrap();
+        runner.start(Vec::new())
+    };
+
+    assert!(matches!(
+        run_with_limit(27),
+        Err(ZapcodeError::AllocationLimitExceeded)
+    ));
+    assert!(matches!(
+        run_with_limit(512),
+        Ok(VmState::Complete(Value::Int(1)))
+    ));
+}
+
+#[test]
+fn test_stack_overflow_is_not_catchable() {
+    let runner = ZapcodeRun::new(
+        r#"
+        function recurse() { recurse(); }
+        let caught = false;
+        try { recurse(); } catch (e) { caught = true; }
+        caught
+        "#
+        .to_string(),
+        Vec::new(),
+        Vec::new(),
+        ResourceLimits {
+            max_stack_depth: 8,
+            ..ResourceLimits::default()
+        },
+    )
+    .unwrap();
+
+    assert!(matches!(
+        runner.start(Vec::new()),
+        Err(ZapcodeError::StackOverflow(_))
+    ));
+}
 
 #[test]
 fn test_try_catch() {

crates/zapcode-core/tests/generators.rs

@@ -1,10 +1,59 @@
-use zapcode_core::Value;
+use zapcode_core::{ResourceLimits, Value, ZapcodeError, ZapcodeRun};
 
 /// Helper to run TS code and get stdout + value
 fn eval_with_output(code: &str) -> (Value, String) {
     zapcode_core::vm::eval_ts_with_output(code).unwrap()
 }
 
+#[test]
+fn test_generator_next_respects_stack_limit() {
+    let runner = ZapcodeRun::new(
+        r#"
+        function* values() { yield 1; }
+        values().next()
+        "#
+        .to_string(),
+        Vec::new(),
+        Vec::new(),
+        ResourceLimits {
+            max_stack_depth: 0,
+            ..ResourceLimits::default()
+        },
+    )
+    .unwrap();
+
+    assert!(matches!(
+        runner.start(Vec::new()),
+        Err(ZapcodeError::StackOverflow(_))
+    ));
+}
+
+#[test]
+fn test_resumed_generator_respects_stack_limit() {
+    let runner = ZapcodeRun::new(
+        r#"
+        function* values() { yield 1; yield 2; }
+        const iterator = values();
+        iterator.next();
+        function resume() { return iterator.next(); }
+        resume()
+        "#
+        .to_string(),
+        Vec::new(),
+        Vec::new(),
+        ResourceLimits {
+            max_stack_depth: 1,
+            ..ResourceLimits::default()
+        },
+    )
+    .unwrap();
+
+    assert!(matches!(
+        runner.start(Vec::new()),
+        Err(ZapcodeError::StackOverflow(_))
+    ));
+}
+
 #[test]
 fn test_basic_generator_yield() {
     let (val, _) = eval_with_output(