fix: No longer drop Return events if the value filter policy is "drop"

codetracer-python-recorder/src/runtime/tracer/runtime_tracer.rs: 
codetracer-python-recorder/src/runtime/value_capture.rs: 
design-docs/US0028 - Configurable Python trace filters.md: 
design-docs/prototypes/module_import_events.py: 
docs/onboarding/trace-filters.md: 
examples/__init__.py: 
examples/abc.py: 
examples/abcd.py: 
examples/hello.py: 

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

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)