WS3

codetracer-python-recorder/src/runtime/activation.rs: 
codetracer-python-recorder/src/runtime/tracer/events.rs: 
design-docs/balanced-call-stack-events-implementation-plan.status.md: 

Signed-off-by: Tzanko Matev <[email protected]>

Author: tzanko-matev

codetracer-python-recorder/src/runtime/activation.rs

@@ -16,6 +16,13 @@ pub struct ActivationController {
     activation_code_id: Option<usize>,
     activation_done: bool,
     started: bool,
+    suspended: bool,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum ActivationExitKind {
+    Suspended,
+    Completed,
 }
 
 impl ActivationController {
@@ -28,6 +35,7 @@ impl ActivationController {
             activation_code_id: None,
             activation_done: false,
             started,
+            suspended: false,
         }
     }
 
@@ -38,6 +46,7 @@ impl ActivationController {
     /// Ensure activation state reflects the current event and report whether
     /// tracing should continue processing it.
     pub fn should_process_event(&mut self, py: Python<'_>, code: &CodeObjectWrapper) -> bool {
+        self.resume_if_needed(code);
         self.ensure_started(py, code);
         self.is_active()
     }
@@ -73,15 +82,33 @@ impl ActivationController {
         }
     }
 
-    /// Handle return events and turn off tracing when the activation function
-    /// exits. Returns `true` when tracing was deactivated by this call.
-    pub fn handle_return_event(&mut self, code_id: usize) -> bool {
-        if self.activation_code_id == Some(code_id) {
-            self.started = false;
-            self.activation_done = true;
-            return true;
+    /// Handle activation exits, marking suspension or completion as appropriate.
+    /// Returns `true` when tracing was deactivated by this call.
+    pub fn handle_exit(&mut self, code_id: usize, exit: ActivationExitKind) -> bool {
+        if self.activation_code_id != Some(code_id) {
+            return false;
+        }
+        match exit {
+            ActivationExitKind::Suspended => {
+                self.suspended = true;
+                false
+            }
+            ActivationExitKind::Completed => {
+                self.started = false;
+                self.activation_done = true;
+                self.suspended = false;
+                true
+            }
+        }
+    }
+
+    fn resume_if_needed(&mut self, code: &CodeObjectWrapper) {
+        if self.started
+            && self.suspended
+            && self.activation_code_id == Some(code.id())
+        {
+            self.suspended = false;
         }
-        false
     }
 }
 
@@ -147,12 +174,28 @@ mod tests {
             let mut controller = ActivationController::new(Some(Path::new("/tmp/target.py")));
             assert!(controller.should_process_event(py, &code));
             assert!(controller.is_active());
-            assert!(controller.handle_return_event(code.id()));
+            assert!(controller.handle_exit(code.id(), ActivationExitKind::Completed));
             assert!(!controller.is_active());
             assert!(!controller.should_process_event(py, &code));
         });
     }
 
+    #[test]
+    fn suspension_keeps_tracing_active() {
+        Python::with_gil(|py| {
+            let code = build_code(py, "target", "/tmp/target.py");
+            let mut controller = ActivationController::new(Some(Path::new("/tmp/target.py")));
+            assert!(controller.should_process_event(py, &code));
+            assert!(controller.is_active());
+            assert!(!controller.handle_exit(code.id(), ActivationExitKind::Suspended));
+            // While suspended, subsequent events should keep tracing alive and clear suspension
+            assert!(controller.should_process_event(py, &code));
+            assert!(controller.is_active());
+            assert!(controller.handle_exit(code.id(), ActivationExitKind::Completed));
+            assert!(!controller.is_active());
+        });
+    }
+
     #[test]
     fn start_path_prefers_activation_path() {
         let controller = ActivationController::new(Some(Path::new("/tmp/target.py")));
@@ -164,6 +207,6 @@ mod tests {
 impl ActivationController {
     #[allow(dead_code)]
     pub fn handle_return(&mut self, code_id: usize) -> bool {
-        self.handle_return_event(code_id)
+        self.handle_exit(code_id, ActivationExitKind::Completed)
     }
 }

codetracer-python-recorder/src/runtime/tracer/events.rs

@@ -9,6 +9,7 @@ use crate::monitoring::{
     events_union, CallbackOutcome, CallbackResult, EventSet, MonitoringEvents, Tracer,
 };
 use crate::policy::policy_snapshot;
+use crate::runtime::activation::ActivationExitKind;
 use crate::runtime::frame_inspector::capture_frame;
 use crate::runtime::io_capture::ScopedMuteIoCapture;
 use crate::runtime::line_snapshots::FrameId;
@@ -338,7 +339,15 @@ impl Tracer for RuntimeTracer {
         _offset: i32,
         retval: &Bound<'_, PyAny>,
     ) -> CallbackResult {
-        self.handle_return_edge(py, code, "on_py_return", retval, None, true, true)
+        self.handle_return_edge(
+            py,
+            code,
+            "on_py_return",
+            retval,
+            None,
+            Some(ActivationExitKind::Completed),
+            true,
+        )
     }
 
     fn on_py_yield(
@@ -354,7 +363,7 @@ impl Tracer for RuntimeTracer {
             "on_py_yield",
             retval,
             Some("<yield>"),
-            false,
+            Some(ActivationExitKind::Suspended),
             false,
         )
     }
@@ -422,7 +431,7 @@ impl Tracer for RuntimeTracer {
             "on_py_unwind",
             exception,
             Some("<unwind>"),
-            true,
+            Some(ActivationExitKind::Completed),
             false,
         )
     }
@@ -536,7 +545,7 @@ impl RuntimeTracer {
         label: &'static str,
         retval: &Bound<'_, PyAny>,
         capture_label: Option<&'static str>,
-        deactivate_on_exit: bool,
+        exit_kind: Option<ActivationExitKind>,
         allow_disable: bool,
     ) -> CallbackResult {
         let is_active = self
@@ -587,14 +596,15 @@ impl RuntimeTracer {
         );
         self.mark_event();
 
-        if deactivate_on_exit
-            && self
+        if let Some(kind) = exit_kind {
+            if self
                 .lifecycle
                 .activation_mut()
-                .handle_return_event(code.id())
-        {
-            let _mute = ScopedMuteIoCapture::new();
-            log::debug!("[RuntimeTracer] deactivated on activation return");
+                .handle_exit(code.id(), kind)
+            {
+                let _mute = ScopedMuteIoCapture::new();
+                log::debug!("[RuntimeTracer] deactivated on activation return");
+            }
         }
 
         Ok(CallbackOutcome::Continue)

design-docs/balanced-call-stack-events-implementation-plan.status.md

@@ -31,12 +31,16 @@
   - Verification: `just dev test` (maturin develop + cargo nextest + pytest) passes.
 
 ### WS3 – Activation & Lifecycle Behaviour
-- **Status:** Not started.
+- **Status:** _Completed_
+  - `ActivationController` now tracks a suspended state and exposes `handle_exit(code_id, ActivationExitKind)`, so `PY_YIELD` transitions into suspension without disabling the activation while `PY_RETURN`/`PY_UNWIND` mark completion.
+  - Resume events clear suspension via `should_process_event`, ensuring activation gating stays engaged until the generator/coroutine finishes.
+  - Added Rust unit tests covering the suspension/resume flow, and the runtime now routes return-edge handling through the new enum to keep lifecycle state consistent.
+  - Verification: `just dev test` passes end-to-end.
 
 ### WS4 – Testing & Validation
 - **Status:** Not started.
 
 ## Next Checkpoints
-1. Begin WS3 by teaching `ActivationController` about suspension/resume semantics.
-2. Plan and implement lifecycle tests ensuring activation gating stays consistent across yields/unwinds.
-3. Evaluate whether additional telemetry/logging is needed before landing WS3/WS4.
+1. Expand WS4 coverage per plan (async awaits, throw/resume, unwind) and update rust/python integration tests accordingly.
+2. Add rust-side assertions (e.g., `print_tracer`) to validate the expanded event mask.
+3. Document any telemetry updates or metadata changes before shipping the feature.