Toplevel return event (#62)

We now emit a Return event at the end of the trace which contains the
exit code of the script which we are tracing.

We have also consolidated the logic which decides whether to handle an
event or not. We have combined the activation and filtering logic into a

Author: tzanko-matev

codetracer-python-recorder/codetracer_python_recorder/cli.py

@@ -287,7 +287,7 @@ def main(argv: Iterable[str] | None = None) -> int:
         try:
             flush()
         finally:
-            stop()
+            stop(exit_code=exit_code)
             sys.argv = old_argv
 
     _serialise_metadata(trace_dir, script=script_path)

codetracer-python-recorder/codetracer_python_recorder/session.py

@@ -38,10 +38,17 @@ def __init__(self, path: Path, format: str) -> None:
         self.path = path
         self.format = format
 
-    def stop(self) -> None:
-        """Stop this trace session."""
+    def stop(self, *, exit_code: int | None = None) -> None:
+        """Stop this trace session.
+
+        Parameters
+        ----------
+        exit_code:
+            Optional process exit status to forward to the recorder backend.
+            When ``None``, the session shutdown reason remains unspecified.
+        """
         if _active_session is self:
-            stop()
+            stop(exit_code=exit_code)
 
     def flush(self) -> None:
         """Flush buffered trace data for this session."""
@@ -51,6 +58,7 @@ def __enter__(self) -> "TraceSession":
         return self
 
     def __exit__(self, exc_type, exc, tb) -> None:  # pragma: no cover - thin wrapper
+        # Exit codes are not tracked for context-managed sessions; report unknown.
         self.stop()
 
 
@@ -121,12 +129,18 @@ def start(
     return session
 
 
-def stop() -> None:
-    """Stop the active trace session if one is running."""
+def stop(*, exit_code: int | None = None) -> None:
+    """Stop the active trace session if one is running.
+
+    Parameters
+    ----------
+    exit_code:
+        Optional process exit status to forward to the backend.
+    """
     global _active_session
     if not _is_tracing_backend():
         return
-    _stop_backend()
+    _stop_backend(exit_code)
     _active_session = None
 
 

codetracer-python-recorder/src/monitoring/api.rs

@@ -107,6 +107,15 @@ pub trait Tracer: Send + Any {
         Ok(())
     }
 
+    /// Provide the process exit status ahead of tracer teardown.
+    fn set_exit_status(
+        &mut self,
+        _py: Python<'_>,
+        _exit_code: Option<i32>,
+    ) -> PyResult<()> {
+        Ok(())
+    }
+
     /// Called on resumption of a generator/coroutine (not via throw()).
     fn on_py_resume(
         &mut self,

codetracer-python-recorder/src/monitoring/install.rs

@@ -82,3 +82,11 @@ pub fn flush_installed_tracer(py: Python<'_>) -> PyResult<()> {
     }
     Ok(())
 }
+
+/// Provide the session exit status to the active tracer if one is installed.
+pub fn update_exit_status(py: Python<'_>, exit_code: Option<i32>) -> PyResult<()> {
+    if let Some(global) = GLOBAL.lock().unwrap().as_mut() {
+        global.tracer.set_exit_status(py, exit_code)?;
+    }
+    Ok(())
+}

codetracer-python-recorder/src/monitoring/mod.rs

@@ -10,7 +10,12 @@ pub(crate) mod install;
 pub mod tracer;
 
 pub use api::Tracer;
-pub use install::{flush_installed_tracer, install_tracer, uninstall_tracer};
+pub use install::{
+    flush_installed_tracer,
+    install_tracer,
+    uninstall_tracer,
+    update_exit_status,
+};
 
 const MONITORING_TOOL_NAME: &str = "codetracer";
 

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

@@ -1,6 +1,5 @@
 //! Event handling pipeline for `RuntimeTracer`.
 
-use super::filtering::TraceDecision;
 use super::runtime_tracer::RuntimeTracer;
 use crate::code_object::CodeObjectWrapper;
 use crate::ffi;
@@ -177,18 +176,8 @@ impl Tracer for RuntimeTracer {
         code: &CodeObjectWrapper,
         _offset: i32,
     ) -> CallbackResult {
-        let is_active = self
-            .lifecycle
-            .activation_mut()
-            .should_process_event(py, code);
-        if matches!(
-            self.should_trace_code(py, code),
-            TraceDecision::SkipAndDisable
-        ) {
-            return Ok(CallbackOutcome::Continue);
-        }
-        if !is_active {
-            return Ok(CallbackOutcome::Continue);
+        if let Some(outcome) = self.evaluate_gate(py, code, true) {
+            return Ok(outcome);
         }
 
         if should_inject_failure(FailureStage::PyStart) {
@@ -244,18 +233,8 @@ impl Tracer for RuntimeTracer {
         code: &CodeObjectWrapper,
         _offset: i32,
     ) -> CallbackResult {
-        let is_active = self
-            .lifecycle
-            .activation_mut()
-            .should_process_event(py, code);
-        if matches!(
-            self.should_trace_code(py, code),
-            TraceDecision::SkipAndDisable
-        ) {
-            return Ok(CallbackOutcome::Continue);
-        }
-        if !is_active {
-            return Ok(CallbackOutcome::Continue);
+        if let Some(outcome) = self.evaluate_gate(py, code, false) {
+            return Ok(outcome);
         }
 
         log_event(py, code, "on_py_resume", None);
@@ -264,18 +243,8 @@ impl Tracer for RuntimeTracer {
     }
 
     fn on_line(&mut self, py: Python<'_>, code: &CodeObjectWrapper, lineno: u32) -> CallbackResult {
-        let is_active = self
-            .lifecycle
-            .activation_mut()
-            .should_process_event(py, code);
-        if matches!(
-            self.should_trace_code(py, code),
-            TraceDecision::SkipAndDisable
-        ) {
-            return Ok(CallbackOutcome::Continue);
-        }
-        if !is_active {
-            return Ok(CallbackOutcome::Continue);
+        if let Some(outcome) = self.evaluate_gate(py, code, false) {
+            return Ok(outcome);
         }
 
         if should_inject_failure(FailureStage::Line) {
@@ -378,18 +347,8 @@ impl Tracer for RuntimeTracer {
         _offset: i32,
         exception: &Bound<'_, PyAny>,
     ) -> CallbackResult {
-        let is_active = self
-            .lifecycle
-            .activation_mut()
-            .should_process_event(py, code);
-        if matches!(
-            self.should_trace_code(py, code),
-            TraceDecision::SkipAndDisable
-        ) {
-            return Ok(CallbackOutcome::Continue);
-        }
-        if !is_active {
-            return Ok(CallbackOutcome::Continue);
+        if let Some(outcome) = self.evaluate_gate(py, code, false) {
+            return Ok(outcome);
         }
 
         log_event(py, code, "on_py_throw", None);
@@ -439,8 +398,17 @@ impl Tracer for RuntimeTracer {
         )
     }
 
+    fn set_exit_status(
+        &mut self,
+        _py: Python<'_>,
+        exit_code: Option<i32>,
+    ) -> PyResult<()> {
+        self.record_exit_status(exit_code);
+        Ok(())
+    }
+
     fn notify_failure(&mut self, _py: Python<'_>) -> PyResult<()> {
-        self.mark_failure();
+        self.mark_disabled();
         Ok(())
     }
 
@@ -484,9 +452,14 @@ impl Tracer for RuntimeTracer {
             self.mark_event();
         }
 
+        self.emit_session_exit(py);
+
+        let exit_summary = self.exit_summary();
+
         if self.lifecycle.encountered_failure() {
             if policy.keep_partial_trace {
-                if let Err(err) = self.lifecycle.finalise(&mut self.writer, &self.filter) {
+                if let Err(err) = self.lifecycle.finalise(&mut self.writer, &self.filter, &exit_summary)
+                {
                     with_error_code(ErrorCode::TraceIncomplete, || {
                         log::warn!(
                             "failed to finalise partial trace after disable: {}",
@@ -519,7 +492,7 @@ impl Tracer for RuntimeTracer {
             .require_trace_or_fail(&policy)
             .map_err(ffi::map_recorder_error)?;
         self.lifecycle
-            .finalise(&mut self.writer, &self.filter)
+            .finalise(&mut self.writer, &self.filter, &exit_summary)
             .map_err(ffi::map_recorder_error)?;
         self.function_ids.clear();
         self.module_names.clear();
@@ -553,22 +526,8 @@ impl RuntimeTracer {
         exit_kind: Option<ActivationExitKind>,
         allow_disable: bool,
     ) -> CallbackResult {
-        let is_active = self
-            .lifecycle
-            .activation_mut()
-            .should_process_event(py, code);
-        if matches!(
-            self.should_trace_code(py, code),
-            TraceDecision::SkipAndDisable
-        ) {
-            return Ok(if allow_disable {
-                CallbackOutcome::DisableLocation
-            } else {
-                CallbackOutcome::Continue
-            });
-        }
-        if !is_active {
-            return Ok(CallbackOutcome::Continue);
+        if let Some(outcome) = self.evaluate_gate(py, code, allow_disable) {
+            return Ok(outcome);
         }
 
         log_event(py, code, label, None);

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

@@ -6,6 +6,8 @@ use crate::runtime::activation::ActivationController;
 use crate::runtime::io_capture::ScopedMuteIoCapture;
 use crate::runtime::output_paths::TraceOutputPaths;
 use crate::runtime::tracer::filtering::FilterCoordinator;
+use crate::runtime::tracer::runtime_tracer::ExitSummary;
+use log::debug;
 use recorder_errors::{enverr, usage, ErrorCode, RecorderResult};
 use runtime_tracing::{NonStreamingTraceWriter, TraceWriter};
 use serde_json::{self, json};
@@ -105,12 +107,12 @@ impl LifecycleController {
         &mut self,
         writer: &mut NonStreamingTraceWriter,
         filter: &FilterCoordinator,
+        exit_summary: &ExitSummary,
     ) -> RecorderResult<()> {
         TraceWriter::finish_writing_trace_metadata(writer).map_err(|err| {
             enverr!(ErrorCode::Io, "failed to finalise trace metadata")
                 .with_context("source", err.to_string())
         })?;
-        self.append_filter_metadata(filter)?;
         TraceWriter::finish_writing_trace_paths(writer).map_err(|err| {
             enverr!(ErrorCode::Io, "failed to finalise trace paths")
                 .with_context("source", err.to_string())
@@ -119,6 +121,9 @@ impl LifecycleController {
             enverr!(ErrorCode::Io, "failed to finalise trace events")
                 .with_context("source", err.to_string())
         })?;
+        debug!("[Lifecycle] writing exit metadata: code={:?}, label={:?}", exit_summary.code, exit_summary.label);
+        self.append_filter_metadata(filter)?;
+        self.append_exit_metadata(exit_summary)?;
         Ok(())
     }
 
@@ -132,6 +137,49 @@ impl LifecycleController {
         self.encountered_failure = false;
     }
 
+    fn append_exit_metadata(&self, exit_summary: &ExitSummary) -> RecorderResult<()> {
+        let Some(outputs) = &self.output_paths else {
+            return Ok(());
+        };
+
+        let path = outputs.metadata();
+        let original = fs::read_to_string(path).map_err(|err| {
+            enverr!(ErrorCode::Io, "failed to read trace metadata")
+                .with_context("path", path.display().to_string())
+                .with_context("source", err.to_string())
+        })?;
+
+        let mut metadata: serde_json::Value = serde_json::from_str(&original).map_err(|err| {
+            enverr!(ErrorCode::Io, "failed to parse trace metadata JSON")
+                .with_context("path", path.display().to_string())
+                .with_context("source", err.to_string())
+        })?;
+
+        if let serde_json::Value::Object(ref mut obj) = metadata {
+            let status = json!({
+                "code": exit_summary.code,
+                "label": exit_summary.label,
+            });
+            obj.insert("process_exit_status".to_string(), status);
+            let serialised = serde_json::to_string(&metadata).map_err(|err| {
+                enverr!(ErrorCode::Io, "failed to serialise trace metadata")
+                    .with_context("path", path.display().to_string())
+                    .with_context("source", err.to_string())
+            })?;
+            fs::write(path, serialised).map_err(|err| {
+                enverr!(ErrorCode::Io, "failed to write trace metadata")
+                    .with_context("path", path.display().to_string())
+                    .with_context("source", err.to_string())
+            })?;
+            Ok(())
+        } else {
+            Err(
+                enverr!(ErrorCode::Io, "trace metadata must be a JSON object")
+                    .with_context("path", path.display().to_string()),
+            )
+        }
+    }
+
     fn append_filter_metadata(&self, filter: &FilterCoordinator) -> RecorderResult<()> {
         let Some(outputs) = &self.output_paths else {
             return Ok(());