fix(audit): raise on write-after-close + weakref.finalize + parametrized event coverage

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

tests/test_audit.py

@@ -78,3 +78,78 @@ def test_writer_no_op_when_path_none():
     w.write(RunEndEvent(run_id="x", elapsed_s=1.0,
                         total_generated_tokens=0, peak_memory_bytes=0))
     w.close()  # no exception
+
+
+import pytest
+from markovian_rsa_mlx.audit import (
+    GenerationStartEvent, ErrorEvent,
+)
+
+EVENT_CASES = [
+    (RunStartEvent, dict(run_id="r", model_id="m", quantization="q",
+                         config=RSAConfig(), prompt="p", created_at="t"), "run_start"),
+    (GenerationStartEvent, dict(run_id="r", round=0, trace_id="t", seed=1,
+                                 prompt_token_count=2, max_tokens=3,
+                                 parent_trace_ids=[]), "generation_start"),
+    (TraceCompleteEvent, dict(run_id="r", round=0, trace_id="t",
+                              text="x", token_ids=[1], generated_tokens=1,
+                              finish_reason="eos", elapsed_s=0.1), "trace_complete"),
+    (TailExtractedEvent, dict(run_id="r", round=0, trace_id="t",
+                              tail_token_ids=[1], tail_text="x", tail_tokens=1),
+     "tail_extracted"),
+    (AggregationPromptEvent, dict(run_id="r", round=0, trace_id="t",
+                                   selected_tail_trace_ids=[], prompt_text="p",
+                                   prompt_token_ids=[1]), "aggregation_prompt"),
+    (RoundCompleteEvent, dict(run_id="r", round=0, trace_ids=[],
+                              memory_estimate_bytes=0, elapsed_s=0.1), "round_complete"),
+    (FinalEvent, dict(run_id="r", final_trace_id="t", final_text="x",
+                      all_final_trace_ids=[], answer_selection="first_final_candidate"),
+     "final"),
+    (RunEndEvent, dict(run_id="r", elapsed_s=0.1, total_generated_tokens=1,
+                       peak_memory_bytes=1), "run_end"),
+    (ErrorEvent, dict(run_id="r", stage="x", message="m", recoverable=False),
+     "error"),
+]
+
+
+@pytest.mark.parametrize("event_cls,kwargs,expected_event_name", EVENT_CASES)
+def test_event_round_trip_through_writer(tmp_path, event_cls, kwargs, expected_event_name):
+    """All 9 event types serialize through AuditWriter with the right `event` field."""
+    audit_path = tmp_path / f"{expected_event_name}.jsonl"
+    with AuditWriter(audit_path) as w:
+        w.write(event_cls(**kwargs))
+    line = audit_path.read_text().strip()
+    payload = json.loads(line)
+    assert payload["event"] == expected_event_name
+
+
+def test_with_context_manager_closes_on_exit(tmp_path):
+    """Context manager closes the writer on block exit, even on exception."""
+    audit_path = tmp_path / "ctx.jsonl"
+    w = AuditWriter(audit_path)
+    with w:
+        w.write(RunEndEvent(run_id="r", elapsed_s=0.0,
+                            total_generated_tokens=0, peak_memory_bytes=0))
+    # Writer must be closed after exiting the with block
+    with pytest.raises(ValueError, match="closed"):
+        w.write(RunEndEvent(run_id="r", elapsed_s=0.0,
+                            total_generated_tokens=0, peak_memory_bytes=0))
+    assert audit_path.read_text().strip().count("\n") == 0  # exactly one event written
+
+
+def test_write_after_close_raises(tmp_path):
+    """Calling write() after close() raises a clear error (not silent no-op)."""
+    audit_path = tmp_path / "post_close.jsonl"
+    w = AuditWriter(audit_path)
+    w.close()
+    with pytest.raises(ValueError, match="closed"):
+        w.write(RunEndEvent(run_id="r", elapsed_s=0.0,
+                            total_generated_tokens=0, peak_memory_bytes=0))
+
+
+def test_close_is_idempotent(tmp_path):
+    """Calling close() twice is fine."""
+    audit_path = tmp_path / "double_close.jsonl"
+    w = AuditWriter(audit_path)
+    w.close()
+    w.close()  # no exception