feat: attach to all frames

Author: ablaszkiewicz

posthog/client.py

@@ -19,7 +19,7 @@
     handle_in_app,
     exception_is_already_captured,
     mark_exception_as_captured,
-    extract_code_variables_from_exception,
+    try_attach_code_variables_to_frames,
     DEFAULT_CODE_VARIABLES_MASK_PATTERNS,
     DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS,
 )
@@ -993,15 +993,6 @@ def capture_exception(
             )
             all_exceptions_with_trace_and_in_app = event["exception"]["values"]
 
-            properties = {
-                "$exception_type": all_exceptions_with_trace_and_in_app[0].get("type"),
-                "$exception_message": all_exceptions_with_trace_and_in_app[0].get(
-                    "value"
-                ),
-                "$exception_list": all_exceptions_with_trace_and_in_app,
-                **properties,
-            }
-
             context_enabled = get_code_variables_enabled_context()
             context_mask = get_code_variables_mask_patterns_context()
             context_ignore = get_code_variables_ignore_patterns_context()
@@ -1023,13 +1014,21 @@ def capture_exception(
             )
 
             if enabled:
-                code_variables = extract_code_variables_from_exception(
+                try_attach_code_variables_to_frames(
+                    all_exceptions_with_trace_and_in_app,
                     exc_info,
                     mask_patterns=mask_patterns,
                     ignore_patterns=ignore_patterns,
                 )
-                if code_variables:
-                    properties["$exception_code_variables"] = code_variables
+
+            properties = {
+                "$exception_type": all_exceptions_with_trace_and_in_app[0].get("type"),
+                "$exception_message": all_exceptions_with_trace_and_in_app[0].get(
+                    "value"
+                ),
+                "$exception_list": all_exceptions_with_trace_and_in_app,
+                **properties,
+            }
 
             if self.log_captured_exceptions:
                 self.log.exception(exception, extra=kwargs)

posthog/exception_utils.py

@@ -58,6 +58,8 @@
 
 DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS = [r"^__.*"]
 
+CODE_VARIABLES_REDACTED_VALUE = "$$_posthog_redacted_based_on_masking_rules_$$"
+
 LogLevelStr = Literal["fatal", "critical", "error", "warning", "info", "debug"]
 
 Event = TypedDict(
@@ -733,6 +735,11 @@ def set_in_app_in_frames(frames, in_app_exclude, in_app_include, project_root=No
             frame["in_app"] = True
             continue
 
+        # Mark __main__ module as in_app by default (user's main script)
+        if module == "__main__":
+            frame["in_app"] = True
+            continue
+
     return frames
 
 
@@ -991,7 +998,7 @@ def serialize_code_variables(
         value = all_vars[name]
 
         if _pattern_matches(name, compiled_mask):
-            result[name] = "***"
+            result[name] = CODE_VARIABLES_REDACTED_VALUE
         else:
             result[name] = _serialize_variable_value(value, max_length)
 
@@ -1000,28 +1007,37 @@ def serialize_code_variables(
     return result
 
 
-def extract_code_variables_from_exception(
-    exception_info, mask_patterns, ignore_patterns
+def try_attach_code_variables_to_frames(
+    all_exceptions, exc_info, mask_patterns, ignore_patterns
 ):
-    exc_type, exc_value, tb = exception_info
-
+    exc_type, exc_value, tb = exc_info
+    
     if tb is None:
-        return {}
-
-    innermost_tb = None
-    for tb_item in iter_stacks(tb):
-        innermost_tb = tb_item
-
-    if innermost_tb is None:
-        return {}
-
-    # Extract variables from the kaboom frame only
-    variables = serialize_code_variables(
-        innermost_tb.tb_frame,
-        mask_patterns=mask_patterns,
-        ignore_patterns=ignore_patterns,
-        max_vars=20,
-        max_length=1024,
-    )
-
-    return variables
+        return
+    
+    tb_frames = list(iter_stacks(tb))
+    
+    if not tb_frames:
+        return
+    
+    for exception in all_exceptions:
+        stacktrace = exception.get("stacktrace")
+        if not stacktrace or "frames" not in stacktrace:
+            continue
+        
+        serialized_frames = stacktrace["frames"]
+        
+        for serialized_frame, tb_item in zip(serialized_frames, tb_frames):
+            if not serialized_frame.get("in_app"):
+                continue
+            
+            variables = serialize_code_variables(
+                tb_item.tb_frame,
+                mask_patterns=mask_patterns,
+                ignore_patterns=ignore_patterns,
+                max_vars=20,
+                max_length=1024,
+            )
+            
+            if variables:
+                serialized_frame["code_variables"] = variables

posthog/test/test_exception_capture.py

@@ -52,7 +52,7 @@ class UnserializableObject:
         code_variables_enabled=True
     )
     
-    def process_data():
+    def trigger_error():
         my_string = "hello world"
         my_number = 42
         my_bool = True
@@ -63,6 +63,19 @@ def process_data():
         
         1/0  # Trigger exception
     
+    def intermediate_function():
+        request_id = "abc-123"
+        user_count = 100
+        is_active = True
+        
+        trigger_error()
+    
+    def process_data():
+        batch_size = 50
+        retry_count = 3
+        
+        intermediate_function()
+    
     process_data()
     """
         )
@@ -74,14 +87,25 @@ def process_data():
     output = excinfo.value.output
 
     assert b"ZeroDivisionError" in output
-    assert b"$exception_code_variables" in output
+    assert b"code_variables" in output
+    
+    # Variables from trigger_error frame
     assert b"'my_string': \"'hello world'\"" in output
     assert b"'my_number': '42'" in output
     assert b"'my_bool': 'True'" in output
     assert b'"my_dict": "{\\"name\\": \\"test\\", \\"value\\": 123}"' in output
     assert b'"my_obj": "<UnserializableObject>"' in output
-    assert b"'my_password': '***'" in output
+    assert b"'my_password': '$$_posthog_redacted_based_on_masking_rules_$$'" in output
     assert b"'__should_be_ignored':" not in output
+    
+    # Variables from intermediate_function frame
+    assert b"'request_id': \"'abc-123'\"" in output
+    assert b"'user_count': '100'" in output
+    assert b"'is_active': 'True'" in output
+    
+    # Variables from process_data frame
+    assert b"'batch_size': '50'" in output
+    assert b"'retry_count': '3'" in output
 
 
 def test_code_variables_context_override(tmpdir):
@@ -122,6 +146,6 @@ def process_data():
     output = excinfo.value.output
 
     assert b"ZeroDivisionError" in output
-    assert b"$exception_code_variables" in output
-    assert b"'bank': '***'" in output
+    assert b"code_variables" in output
+    assert b"'bank': '$$_posthog_redacted_based_on_masking_rules_$$'" in output
     assert b"'__dunder_var': \"'should_be_visible'\"" in output