fix(safety): add try/finally cleanup and exception state resets

- Wrap validate_saved_activations in try/finally for cleanup (user_funcs.py)
- Wrap show_model_graph render_graph in try/finally with cleanup (user_funcs.py)
- Reset _track_tensors and _pause_logging in exception handler (trace_model.py)
- Update test to expect [] instead of None for cleared parent_params

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: johnmarktaylor91

tests/test_param_log.py

@@ -602,7 +602,7 @@ def test_parent_params_cleared_from_tle(self):
         """After postprocessing, parent_params references should be cleared."""
         mh = log_forward_pass(_make_simple_model(), _simple_input())
         for entry in mh:
-            assert entry.parent_params is None
+            assert entry.parent_params == []
 
     def test_vis_renders_without_error(self):
         """Basic smoke test that visualization renders for each param scenario."""

torchlens/user_funcs.py

@@ -307,21 +307,24 @@ def show_model_graph(
         save_gradients=False,
         random_seed=random_seed,
     )
-    model_log.render_graph(
-        vis_opt,
-        vis_nesting_depth,
-        vis_outpath,
-        vis_graph_overrides,
-        vis_node_overrides,
-        vis_nested_node_overrides,
-        vis_edge_overrides,
-        vis_gradient_edge_overrides,
-        vis_module_overrides,
-        save_only,
-        vis_fileformat,
-        vis_buffer_layers,
-        vis_direction,
-    )
+    try:
+        model_log.render_graph(
+            vis_opt,
+            vis_nesting_depth,
+            vis_outpath,
+            vis_graph_overrides,
+            vis_node_overrides,
+            vis_nested_node_overrides,
+            vis_edge_overrides,
+            vis_gradient_edge_overrides,
+            vis_module_overrides,
+            save_only,
+            vis_fileformat,
+            vis_buffer_layers,
+            vis_direction,
+        )
+    finally:
+        model_log.cleanup()
 
 
 def validate_saved_activations(
@@ -386,12 +389,13 @@ def validate_saved_activations(
         save_function_args=True,
         random_seed=random_seed,
     )
-    activations_are_valid = model_log.validate_saved_activations(
-        ground_truth_output_tensors, verbose
-    )
-
-    model_log.cleanup()
-    del model_log
+    try:
+        activations_are_valid = model_log.validate_saved_activations(
+            ground_truth_output_tensors, verbose
+        )
+    finally:
+        model_log.cleanup()
+        del model_log
     return activations_are_valid