Fix dropped return events (#60)

Before if according to our filter we needed to drop the return value of
a function call, we would end up not emitting a "Return" event to the
trace at all. This lead to an unbalanced trace file. After this fix we
emit Return events but we replace the return value in the event with the
string "<dropped>"

Author: tzanko-matev

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

@@ -179,7 +179,7 @@ mod tests {
     use std::collections::BTreeMap;
     use std::ffi::CString;
     use std::fs;
-    use std::path::Path;
+    use std::path::{Path, PathBuf};
     use std::sync::Arc;
     use std::thread;
 
@@ -716,6 +716,30 @@ result = compute()\n"
         });
     }
 
+    #[pyfunction]
+    fn capture_py_start(py: Python<'_>, code: Bound<'_, PyCode>, offset: i32) -> PyResult<()> {
+        ffi::wrap_pyfunction("test_capture_py_start", || {
+            ACTIVE_TRACER.with(|cell| -> PyResult<()> {
+                let ptr = cell.get();
+                if ptr.is_null() {
+                    panic!("No active RuntimeTracer for capture_py_start");
+                }
+                unsafe {
+                    let tracer = &mut *ptr;
+                    let wrapper = CodeObjectWrapper::new(py, &code);
+                    match tracer.on_py_start(py, &wrapper, offset) {
+                        Ok(outcome) => {
+                            LAST_OUTCOME.with(|cell| cell.set(Some(outcome)));
+                            Ok(())
+                        }
+                        Err(err) => Err(err),
+                    }
+                }
+            })?;
+            Ok(())
+        })
+    }
+
     #[pyfunction]
     fn capture_line(py: Python<'_>, code: Bound<'_, PyCode>, lineno: u32) -> PyResult<()> {
         ffi::wrap_pyfunction("test_capture_line", || {
@@ -770,7 +794,7 @@ result = compute()\n"
 
     const PRELUDE: &str = r#"
 import inspect
-from test_tracer import capture_line, capture_return_event
+from test_tracer import capture_line, capture_return_event, capture_py_start
 
 def snapshot(line=None):
     frame = inspect.currentframe().f_back
@@ -786,6 +810,10 @@ def emit_return(value):
     frame = inspect.currentframe().f_back
     capture_return_event(frame.f_code, value)
     return value
+
+def start_call():
+    frame = inspect.currentframe().f_back
+    capture_py_start(frame.f_code, frame.f_lasti)
 "#;
 
     #[derive(Debug, Clone, PartialEq)]
@@ -862,9 +890,14 @@ def emit_return(value):
 
     fn ensure_test_module(py: Python<'_>) {
         let module = PyModule::new(py, "test_tracer").expect("create module");
+        module
+            .add_function(
+                wrap_pyfunction!(capture_py_start, &module).expect("wrap capture_py_start"),
+            )
+            .expect("add py_start capture function");
         module
             .add_function(wrap_pyfunction!(capture_line, &module).expect("wrap capture_line"))
-            .expect("add function");
+            .expect("add line capture function");
         module
             .add_function(
                 wrap_pyfunction!(capture_return_event, &module).expect("wrap capture_return_event"),
@@ -906,6 +939,25 @@ def emit_return(value):
         fs::write(path, contents.trim_start()).expect("write filter");
     }
 
+    fn install_drop_everything_filter(project_root: &Path) -> PathBuf {
+        let filters_dir = project_root.join(".codetracer");
+        fs::create_dir(&filters_dir).expect("create .codetracer");
+        let drop_filter_path = filters_dir.join("drop-filter.toml");
+        write_filter(
+            &drop_filter_path,
+            r#"
+            [meta]
+            name = "drop-all"
+            version = 1
+
+            [scope]
+            default_exec = "trace"
+            default_value_action = "drop"
+            "#,
+        );
+        drop_filter_path
+    }
+
     #[test]
     fn trace_filter_redacts_values() {
         Python::with_gil(|py| {
@@ -1125,21 +1177,7 @@ sensitive("s3cr3t")
 
             let project = tempfile::tempdir().expect("project dir");
             let project_root = project.path();
-            let filters_dir = project_root.join(".codetracer");
-            fs::create_dir(&filters_dir).expect("create .codetracer");
-            let drop_filter_path = filters_dir.join("drop-filter.toml");
-            write_filter(
-                &drop_filter_path,
-                r#"
-                [meta]
-                name = "drop-all"
-                version = 1
-
-                [scope]
-                default_exec = "trace"
-                default_value_action = "drop"
-                "#,
-            );
+            let drop_filter_path = install_drop_everything_filter(project_root);
 
             let config = TraceFilterConfig::from_inline_and_paths(
                 &[("builtin-default", BUILTIN_TRACE_FILTER)],
@@ -1186,11 +1224,13 @@ dropper()
             }
 
             let mut variable_names: Vec<String> = Vec::new();
-            let mut return_events = 0usize;
+            let mut return_values: Vec<ValueRecord> = Vec::new();
             for event in &tracer.writer.events {
                 match event {
                     TraceLowLevelEvent::VariableName(name) => variable_names.push(name.clone()),
-                    TraceLowLevelEvent::Return(_) => return_events += 1,
+                    TraceLowLevelEvent::Return(record) => {
+                        return_values.push(record.return_value.clone())
+                    }
                     _ => {}
                 }
             }
@@ -1199,9 +1239,89 @@ dropper()
                 "expected no variables captured, found {:?}",
                 variable_names
             );
+            assert_eq!(return_values.len(), 1, "return event should remain balanced");
+            match &return_values[0] {
+                ValueRecord::Error { msg, .. } => assert_eq!(msg, "<dropped>"),
+                other => panic!("expected dropped sentinel return value, got {other:?}"),
+            }
+        });
+    }
+
+    #[test]
+    fn drop_filters_keep_call_return_pairs_balanced() {
+        Python::with_gil(|py| {
+            ensure_test_module(py);
+
+            let project = tempfile::tempdir().expect("project dir");
+            let project_root = project.path();
+            let drop_filter_path = install_drop_everything_filter(project_root);
+
+            let config = TraceFilterConfig::from_inline_and_paths(
+                &[("builtin-default", BUILTIN_TRACE_FILTER)],
+                &[drop_filter_path.clone()],
+            )
+            .expect("load filter chain");
+            let engine = Arc::new(TraceFilterEngine::new(config));
+
+            let app_dir = project_root.join("app");
+            fs::create_dir_all(&app_dir).expect("create app dir");
+            let script_path = app_dir.join("classes.py");
+            let body = r#"
+def initializer(label):
+    start_call()
+    return emit_return(label.upper())
+
+class Alpha:
+    TOKEN = initializer("alpha")
+
+class Beta:
+    TOKEN = initializer("beta")
+
+class Gamma:
+    TOKEN = initializer("gamma")
+
+initializer("omega")
+"#;
+            let script = format!("{PRELUDE}\n{body}", PRELUDE = PRELUDE, body = body);
+            fs::write(&script_path, script).expect("write script");
+
+            let mut tracer = RuntimeTracer::new(
+                script_path.to_string_lossy().as_ref(),
+                &[],
+                TraceEventsFileFormat::Json,
+                None,
+                Some(engine),
+            );
+
+            {
+                let _guard = ScopedTracer::new(&mut tracer);
+                LAST_OUTCOME.with(|cell| cell.set(None));
+                let run_code = format!(
+                    "import runpy, sys\nsys.path.insert(0, r\"{}\")\nrunpy.run_path(r\"{}\")",
+                    project_root.display(),
+                    script_path.display()
+                );
+                let run_code_c = CString::new(run_code).expect("script contains nul byte");
+                py.run(run_code_c.as_c_str(), None, None)
+                    .expect("execute classes script");
+            }
+
+            let mut call_count = 0usize;
+            let mut return_count = 0usize;
+            for event in &tracer.writer.events {
+                match event {
+                    TraceLowLevelEvent::Call(_) => call_count += 1,
+                    TraceLowLevelEvent::Return(_) => return_count += 1,
+                    _ => {}
+                }
+            }
+            assert!(
+                call_count >= 4,
+                "expected at least four call events, saw {call_count}"
+            );
             assert_eq!(
-                return_events, 0,
-                "return value should be dropped instead of recorded"
+                call_count, return_count,
+                "drop filters must keep call/return pairs balanced"
             );
         });
     }

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

@@ -19,6 +19,7 @@ use crate::trace_filter::config::ValueAction;
 use crate::trace_filter::engine::{ValueKind, ValuePolicy};
 
 const REDACTED_SENTINEL: &str = "<redacted>";
+const DROPPED_SENTINEL: &str = "<dropped>";
 
 const VALUE_KIND_COUNT: usize = 5;
 
@@ -54,6 +55,14 @@ fn redacted_value(writer: &mut NonStreamingTraceWriter) -> ValueRecord {
     }
 }
 
+fn dropped_value(writer: &mut NonStreamingTraceWriter) -> ValueRecord {
+    let ty = TraceWriter::ensure_type_id(writer, TypeKind::Raw, "Dropped");
+    ValueRecord::Error {
+        msg: DROPPED_SENTINEL.to_string(),
+        type_id: ty,
+    }
+}
+
 fn record_redaction(kind: ValueKind, candidate: &str, telemetry: Option<&mut ValueFilterStats>) {
     if let Some(stats) = telemetry {
         stats.record_redaction(kind);
@@ -322,8 +331,7 @@ pub fn record_return_value(
         ValueKind::Return,
         name,
         telemetry.as_deref_mut(),
-    );
-    if let Some(encoded) = encoded {
-        TraceWriter::register_return(writer, encoded);
-    }
+    )
+    .unwrap_or_else(|| dropped_value(writer));
+    TraceWriter::register_return(writer, encoded);
 }

design-docs/US0028 - Configurable Python trace filters.md

@@ -159,7 +159,7 @@ Callers validate whether a parsed selector is legal in the current context (e.g.
 1. Initialize the execution policy to `scope.default_exec` (or the inherited value when composing filters).  
 2. Walk `scope.rules` from top to bottom. Each rule whose selector matches the current frame updates the execution policy (`trace` vs `skip`) and the active default for value capture. Later matching rules replace earlier decisions because the traversal never rewinds.  
 3. For value capture inside a scope, start from the applicable default (`scope.default_value_action`, overridden by the scope rule’s `value_default` when provided).  
-4. Apply each `value_patterns` entry in order. The first pattern whose selector matches the variable or payload sets the decision to `allow` (serialize), `redact` (replace with `<redacted>`), or `drop` (omit entirely) and stops further evaluation for that value.  
+4. Apply each `value_patterns` entry in order. The first pattern whose selector matches the variable or payload sets the decision to `allow` (serialize), `redact` (replace with `<redacted>`), or `drop` (omit entirely; return-value drops still emit a structural return edge with a `<dropped>` placeholder) and stops further evaluation for that value.  
 5. If no pattern matches, fall back to the current default value action.  
 
 ## Sample Filters (TOML)