From ca02917712b5aa3771f1e71c5dec0dc673b69f37 Mon Sep 17 00:00:00 2001
From: Manohar Paturi <186662190+ManoharPaturi@users.noreply.github.com>
Date: Sun, 18 Oct 2026 09:29:30 +0530
Subject: [PATCH] FIX: derive csv header from the union of row keys

---
 pyrit/common/csv_helper.py           | 14 +++++++++++++-
 tests/unit/common/test_csv_helper.py | 24 ++++++++++++++++++++++++
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/pyrit/common/csv_helper.py b/pyrit/common/csv_helper.py
index 51633c123b..698718a453 100644
--- a/pyrit/common/csv_helper.py
+++ b/pyrit/common/csv_helper.py
@@ -20,6 +20,10 @@ def write_csv(file: IO[Any], examples: list[dict[str, str]]) -> None:
     """
     Write a list of dictionaries to a CSV file.
 
+    The header is the union of the keys across all rows, in first-seen order, so a
+    row carrying a key the first row lacks is written instead of failing halfway
+    through the file. Rows missing a header key get an empty cell.
+
     Args:
         file: A file-like object opened for writing CSV data.
         examples (list[dict[str, str]]): List of dictionaries to write as CSV rows.
@@ -27,6 +31,14 @@ def write_csv(file: IO[Any], examples: list[dict[str, str]]) -> None:
     if not examples:
         return
 
-    writer = csv.DictWriter(file, fieldnames=examples[0].keys())
+    fieldnames: list[str] = []
+    seen: set[str] = set()
+    for example in examples:
+        for key in example:
+            if key not in seen:
+                seen.add(key)
+                fieldnames.append(key)
+
+    writer = csv.DictWriter(file, fieldnames=fieldnames)
     writer.writeheader()
     writer.writerows(examples)
diff --git a/tests/unit/common/test_csv_helper.py b/tests/unit/common/test_csv_helper.py
index b210297a7f..7424641d56 100644
--- a/tests/unit/common/test_csv_helper.py
+++ b/tests/unit/common/test_csv_helper.py
@@ -57,3 +57,27 @@ def test_write_csv_writes_header_and_rows():
     lines = file.getvalue().strip().splitlines()
     assert lines[0] == "name,role"
     assert lines[1] == "alice,admin"
+
+
+def test_write_csv_covers_keys_only_later_rows_have():
+    output = StringIO()
+    examples = [
+        {"name": "Alice"},
+        {"name": "Bob", "note": "extra"},
+    ]
+    write_csv(output, examples)
+    output.seek(0)
+    result = read_csv(output)
+    assert result == [{"name": "Alice", "note": ""}, {"name": "Bob", "note": "extra"}]
+
+
+def test_write_csv_keeps_first_seen_column_order():
+    output = StringIO()
+    examples = [
+        {"a": "1", "b": "2"},
+        {"c": "3", "a": "4"},
+    ]
+    write_csv(output, examples)
+    output.seek(0)
+    lines = output.read().strip().splitlines()
+    assert lines[0] == "a,b,c"