fix(gmail-bulk-import): harden re-import against 8-bit headers + poison batches

Four audit findings (2026-07-14), each reproduced on crafted mboxes; the
Takeout corpus has proven 8-bit header bytes, so all are real re-import risks.

1. Whole-run crash on 8-bit Message-ID. compat32 .get() returns an
   email.header.Header (not str) for raw 8-bit bytes; the old
   Header.strip() raised AttributeError. _message_id ran BEFORE the
   per-message try, so one bad header killed the entire import.
   Fix: str() + sanitize_surrogates() before strip; and move
   _message_id/_parse_date/_parse_attachments INSIDE the per-message try —
   a broken message is now errors += 1, never run death.

2. Poison batch. pending.clear() ran only AFTER a successful insert, so a
   failed flush (DB down / bad row) left pending intact and every later
   message re-flushed the doomed batch; the final flush sat in try/finally
   with no except and propagated out, losing all stats. Fix: _flush always
   clears pending and counts a failed insert as db_insert_failed; the run
   always reaches import_complete.

3. Stats didn't reconcile with the DB. imported counts archive writes, not
   DB rows, so a partial-insert drift was invisible. Fix: separate
   db_inserted/db_insert_failed counters; main() exits non-zero on any
   error, DB drift, or a processed = imported + skipped + errors imbalance.

4. 8-bit Date → needless epoch_fallback. parsedate_to_datetime(Header)
   raised even when str(header) parses fine. Fix: str() before the epoch
   fallback.

Shared helper: _sanitize moved from gmail-header-backfill into
packages/kb-mail (kb_mail.text.sanitize_surrogates) and used by both jobs;
gmail-header-backfill now depends on kb-mail.

Tests: regression coverage for all four findings in gmail-bulk-import
(8-bit id/date, per-message guard, failed-insert non-poisoning, stats
balance) plus kb_mail.text unit tests. Full suites green:
kb-mail 27, gmail-bulk-import 33, gmail-header-backfill 43.

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

jobs/gmail-bulk-import/src/gmail_bulk_import/importer.py

@@ -43,6 +43,7 @@ import structlog
 
 from kb_mail.archive import save_eml
 from kb_mail.envelope import Envelope
+from kb_mail.text import sanitize_surrogates as _sanitize
 
 _log = structlog.get_logger(__name__)
 
@@ -60,16 +61,28 @@ def _eml_ref(envelope_id: str, ts: datetime) -> str:
 
 
 def _message_id(msg: mailbox.mboxMessage) -> str:
-    """Return a stable envelope id from Message-ID header or SHA-256 content hash."""
-    mid = msg.get("Message-ID", "").strip().strip("<>")
+    """Return a stable envelope id from Message-ID header or SHA-256 content hash.
+
+    compat32 .get() returns an email.header.Header (not str) when the raw value
+    holds 8-bit bytes — str() + _sanitize keeps the id jsonb-safe and gives
+    .strip() a real string to work on. An unguarded .strip() on a Header raised
+    AttributeError and, because _message_id runs before the per-message try,
+    killed the whole run (the archive has proven 8-bit header bytes).
+    """
+    mid = _sanitize(str(msg.get("Message-ID", ""))).strip().strip("<>")
     if mid:
         return mid
     return "sha256-" + hashlib.sha256(msg.as_bytes()).hexdigest()[:32]
 
 
 def _parse_date(msg: mailbox.mboxMessage) -> datetime:
-    """Parse Date header to UTC-aware datetime. Returns epoch as last resort."""
-    date_str = msg.get("Date", "")
+    """Parse Date header to UTC-aware datetime. Returns epoch as last resort.
+
+    str() first: a compat32 8-bit Date surfaces as an email.header.Header, and
+    parsedate_to_datetime() raises on a Header (→ needless epoch_fallback) even
+    when str(header) is perfectly parseable.
+    """
+    date_str = str(msg.get("Date", ""))
     if date_str:
         try:
             dt = parsedate_to_datetime(date_str)
@@ -152,8 +165,17 @@ async def run_import(
 ) -> dict[str, int]:
     """Import all messages from an mbox file.
 
-    Returns stats: {processed, imported, skipped, errors, epoch_fallback,
-    msgs_with_attachments, total_attachments, total_attachment_bytes}.
+    Returns stats: {processed, imported, skipped, errors, db_inserted,
+    db_insert_failed, epoch_fallback, msgs_with_attachments, total_attachments,
+    total_attachment_bytes}.
+
+    Two invariants always hold (checked by main() for the exit code):
+      * processed == imported + skipped + errors
+      * with a DB connection: db_inserted + db_insert_failed == imported + skipped
+
+    `imported`/`skipped` count archive writes; `db_inserted`/`db_insert_failed`
+    count DB rows. They are tracked separately because a failed batch insert
+    leaves the archive ahead of the DB, and nothing else would catch the drift.
 
     dry_run=True parses and counts everything without writing.
     limit=N stops after processing N messages (for sampling).
@@ -163,6 +185,8 @@ async def run_import(
         "imported": 0,
         "skipped": 0,
         "errors": 0,
+        "db_inserted": 0,
+        "db_insert_failed": 0,
         "epoch_fallback": 0,
         "msgs_with_attachments": 0,
         "total_attachments": 0,
@@ -176,8 +200,17 @@ async def run_import(
     pending: list[Envelope] = []
 
     async def _flush() -> None:
+        # Always clear pending, even when the insert fails: a poisoned batch
+        # (DB down, one bad row) must never be re-flushed by every subsequent
+        # message. A failed insert leaves the archive ahead of the DB, counted
+        # as db_insert_failed so the drift is visible instead of silent.
         if conn is not None and pending:
+            try:
                 await _insert_batch(conn, pending)
+                stats["db_inserted"] += len(pending)
+            except Exception:
+                _log.exception("db_insert_failed", count=len(pending))
+                stats["db_insert_failed"] += len(pending)
         pending.clear()
 
     try:
@@ -188,6 +221,11 @@ async def run_import(
 
             stats["processed"] += 1
 
+            # The ENTIRE per-message body is guarded: a single malformed
+            # message (8-bit header, unwalkable MIME, bad save) counts as one
+            # error and never aborts the run.
+            envelope_id = "<unknown>"
+            try:
                 envelope_id = _message_id(msg)
                 ts = _parse_date(msg)
 
@@ -205,7 +243,6 @@ async def run_import(
                     stats["imported"] += 1
                     continue
 
-            try:
                 raw = msg.as_bytes(policy=email.policy.compat32)
                 try:
                     raw_ref = await save_eml(archive_root, envelope_id, "gmail", ts, raw)
@@ -228,6 +265,9 @@ async def run_import(
             if stats["processed"] % 500 == 0:
                 _log.info("progress", **stats)
 
+        # Final flush swallows its own insert error (see _flush) so the run
+        # always reaches import_complete with balanced stats instead of dying
+        # here with the counters lost.
         await _flush()
 
     finally:
@@ -262,7 +302,17 @@ def main() -> None:
     stats = asyncio.run(
         run_import(args.mbox, args.archive, args.dsn or None, args.dry_run, args.limit)
     )
-    sys.exit(1 if stats["errors"] > 0 else 0)
+
+    balanced = (
+        stats["processed"] == stats["imported"] + stats["skipped"] + stats["errors"]
+    )
+    if not balanced:
+        _log.error("stats_unbalanced", **stats)
+
+    # Non-zero exit on any message failure, any DB drift (archive ahead of DB),
+    # or a stats imbalance — the run must not report clean when it isn't.
+    failed = stats["errors"] > 0 or stats["db_insert_failed"] > 0 or not balanced
+    sys.exit(1 if failed else 0)
 
 
 if __name__ == "__main__":

jobs/gmail-bulk-import/tests/test_importer.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 import hashlib
+import json
 import mailbox
 from datetime import datetime, timezone
 from email import encoders
@@ -12,6 +13,7 @@ from pathlib import Path
 
 import pytest
 
+from gmail_bulk_import import importer
 from gmail_bulk_import.importer import (
     _message_id,
     _parse_attachments,
@@ -281,3 +283,211 @@ class TestRunImport:
 
         assert stats["imported"] == 1
         assert stats["total_attachments"] == 2
+
+
+# ---------------------------------------------------------------------------
+# Regression tests for the 2026-07-14 hardening audit (four fixes).
+#
+# The audit reproduced each of these on crafted mboxes; the corpus is proven to
+# contain raw 8-bit header bytes, so every scenario below is a real re-import
+# risk, not a hypothetical.
+# ---------------------------------------------------------------------------
+
+# Raw 8-bit byte (0xe9) inside Message-ID: compat32 .get() returns an
+# email.header.Header, whose .strip() raised AttributeError. Because
+# _message_id ran BEFORE the per-message try, one such header killed the run.
+_EIGHTBIT_MID_EML = (
+    b"From: alice@example.com\r\n"
+    b"Message-ID: <bad\xe9id@example.com>\r\n"
+    b"Date: Tue, 10 Jun 2025 12:00:00 +0000\r\n"
+    b"Subject: hi\r\n\r\nbody"
+)
+
+# Raw 8-bit byte inside the Date header's trailing comment: .get() returns a
+# Header, parsedate_to_datetime(Header) raises → needless epoch_fallback, even
+# though str(header) parses to a real date (2009-05-19 10:27:09 +0200).
+_EIGHTBIT_DATE_EML = (
+    b"From: alice@example.com\r\n"
+    b"Message-ID: <clean-date@example.com>\r\n"
+    b"Date: Tue, 19 May 2009 10:27:09 +0200 (Ho\xe9ra)\r\n"
+    b"Subject: hi\r\n\r\nbody"
+)
+
+
+def _make_mbox_raw(path: Path, raw_messages: list[bytes]) -> None:
+    mbox = mailbox.mbox(str(path), create=True)
+    for raw in raw_messages:
+        mbox.add(mailbox.mboxMessage(raw))
+    mbox.flush()
+    mbox.close()
+
+
+class _FakeConn:
+    """Stand-in asyncpg connection: records executemany rows or fails them."""
+
+    def __init__(self, fail: bool = False):
+        self.fail = fail
+        self.inserted: list = []
+        self.executemany_calls = 0
+
+    async def executemany(self, query, rows):
+        self.executemany_calls += 1
+        if self.fail:
+            raise RuntimeError("DB down")
+        self.inserted.extend(list(rows))
+
+    async def close(self):
+        pass
+
+
+def _patch_connect(monkeypatch, conn: _FakeConn) -> None:
+    async def _fake_connect(dsn):
+        return conn
+    monkeypatch.setattr(importer.asyncpg, "connect", _fake_connect)
+
+
+class TestEightBitMessageId:
+    def test_message_id_8bit_header_does_not_crash(self):
+        # Fix 1: str() + _sanitize gives .strip() a real string. Old code did
+        # Header.strip() → AttributeError.
+        msg = mailbox.mboxMessage(_EIGHTBIT_MID_EML)
+        mid = _message_id(msg)
+        assert isinstance(mid, str)
+        assert mid.endswith("id@example.com")
+        json.dumps(mid).encode("utf-8")  # jsonb-safe: no lone surrogate
+
+    async def test_run_survives_8bit_message_id(self, tmp_path):
+        # Old code raised out of run_import here; the run must now complete and
+        # count the message normally.
+        mbox_path = tmp_path / "mail.mbox"
+        archive = tmp_path / "archive"
+        _make_mbox_raw(mbox_path, [_EIGHTBIT_MID_EML, _msg_with_id("<ok@e.com>").as_bytes()])
+
+        stats = await run_import(mbox_path, archive)
+
+        assert stats["processed"] == 2
+        assert stats["imported"] == 2
+        assert stats["errors"] == 0
+
+
+class TestPerMessageGuard:
+    async def test_one_bad_message_counts_once_and_run_continues(self, tmp_path, monkeypatch):
+        # Fix 2: the parse helpers now live inside the per-message try, so a
+        # message that raises is one error, not the death of the run.
+        calls = {"n": 0}
+        real_message_id = importer._message_id
+
+        def flaky(msg):
+            calls["n"] += 1
+            if calls["n"] == 1:
+                raise RuntimeError("boom parsing first message")
+            return real_message_id(msg)
+
+        monkeypatch.setattr(importer, "_message_id", flaky)
+
+        mbox_path = tmp_path / "mail.mbox"
+        archive = tmp_path / "archive"
+        _make_mbox(mbox_path, [_msg_with_id("<a@e.com>"), _msg_with_id("<b@e.com>")])
+
+        stats = await run_import(mbox_path, archive)
+
+        assert stats["processed"] == 2
+        assert stats["errors"] == 1
+        assert stats["imported"] == 1
+
+
+class TestFailedInsertDoesNotPoison:
+    async def test_final_flush_failure_does_not_kill_run(self, tmp_path, monkeypatch):
+        # Fix 3: the final flush used to be uncaught (try/finally, no except),
+        # so a failing insert lost all stats. Now it is swallowed and counted.
+        conn = _FakeConn(fail=True)
+        _patch_connect(monkeypatch, conn)
+
+        mbox_path = tmp_path / "mail.mbox"
+        archive = tmp_path / "archive"
+        _make_mbox(mbox_path, [_msg_with_id(f"<m{i}@e.com>") for i in range(3)])
+
+        stats = await run_import(mbox_path, archive, dsn="postgresql://fake")
+
+        # Run completes, archive is ahead of DB, and the drift is visible.
+        assert stats["imported"] == 3
+        assert stats["db_inserted"] == 0
+        assert stats["db_insert_failed"] == 3
+        assert stats["errors"] == 0
+
+    async def test_poisoned_batch_does_not_re_flush_forever(self, tmp_path, monkeypatch):
+        # Fix 3: pending is cleared even on insert failure, so a poisoned batch
+        # is not re-flushed by every subsequent message. With BATCH_SIZE=1 each
+        # message flushes exactly once; old code re-flushed a growing pending
+        # list and died in the uncaught final flush.
+        monkeypatch.setattr(importer, "BATCH_SIZE", 1)
+        conn = _FakeConn(fail=True)
+        _patch_connect(monkeypatch, conn)
+
+        mbox_path = tmp_path / "mail.mbox"
+        archive = tmp_path / "archive"
+        _make_mbox(mbox_path, [_msg_with_id(f"<m{i}@e.com>") for i in range(3)])
+
+        stats = await run_import(mbox_path, archive, dsn="postgresql://fake")
+
+        assert conn.executemany_calls == 3  # one per message, not a growing batch
+        assert stats["imported"] == 3
+        assert stats["db_insert_failed"] == 3
+        assert stats["errors"] == 0
+
+
+class TestStatsBalance:
+    async def test_processed_balances_and_db_reconciles(self, tmp_path, monkeypatch):
+        # Fix 4: imported/skipped count archive writes; db_inserted counts DB
+        # rows. Both invariants must always hold.
+        conn = _FakeConn()
+        _patch_connect(monkeypatch, conn)
+        mbox_path = tmp_path / "mail.mbox"
+        archive = tmp_path / "archive"
+        _make_mbox(mbox_path, [_msg_with_id(f"<m{i}@e.com>") for i in range(3)])
+
+        stats1 = await run_import(mbox_path, archive, dsn="postgresql://fake")
+
+        # second run over the same archive: every message is skipped, but each
+        # still goes to the DB (ON CONFLICT DO NOTHING)
+        conn2 = _FakeConn()
+        _patch_connect(monkeypatch, conn2)
+        stats2 = await run_import(mbox_path, archive, dsn="postgresql://fake")
+
+        for s in (stats1, stats2):
+            assert s["processed"] == s["imported"] + s["skipped"] + s["errors"]
+            assert s["db_inserted"] + s["db_insert_failed"] == s["imported"] + s["skipped"]
+
+        assert stats1["imported"] == 3 and stats1["db_inserted"] == 3
+        assert stats2["skipped"] == 3 and stats2["db_inserted"] == 3
+
+    async def test_archive_only_mode_leaves_db_counters_zero(self, tmp_path):
+        mbox_path = tmp_path / "mail.mbox"
+        archive = tmp_path / "archive"
+        _make_mbox(mbox_path, [_msg_with_id(f"<m{i}@e.com>") for i in range(2)])
+
+        stats = await run_import(mbox_path, archive, dsn=None)
+
+        assert stats["imported"] == 2
+        assert stats["db_inserted"] == 0
+        assert stats["db_insert_failed"] == 0
+        assert stats["processed"] == stats["imported"] + stats["skipped"] + stats["errors"]
+
+
+class TestParseDate8Bit:
+    def test_8bit_date_header_parses_via_str_not_epoch(self):
+        # Fix 4 (date): str(header) before epoch_fallback. Old code fed the
+        # Header straight to parsedate_to_datetime → AttributeError → epoch.
+        msg = mailbox.mboxMessage(_EIGHTBIT_DATE_EML)
+        ts = _parse_date(msg)
+        assert ts == datetime(2009, 5, 19, 8, 27, 9, tzinfo=timezone.utc)
+
+    async def test_8bit_date_not_counted_as_epoch_fallback(self, tmp_path):
+        mbox_path = tmp_path / "mail.mbox"
+        archive = tmp_path / "archive"
+        _make_mbox_raw(mbox_path, [_EIGHTBIT_DATE_EML])
+
+        stats = await run_import(mbox_path, archive)
+
+        assert stats["epoch_fallback"] == 0
+        assert stats["imported"] == 1

jobs/gmail-header-backfill/pyproject.toml

@@ -9,6 +9,7 @@ requires-python = ">=3.11"
 dependencies = [
     "asyncpg>=0.29",
     "structlog>=24.1",
+    "kb-mail",
 ]
 
 [project.scripts]

jobs/gmail-header-backfill/src/gmail_header_backfill/backfill.py

@@ -46,6 +46,8 @@ from typing import Optional
 import asyncpg
 import structlog
 
+from kb_mail.text import sanitize_surrogates as _sanitize
+
 _log = structlog.get_logger(__name__)
 
 DEFAULT_ARCHIVE_ROOT = Path("/home/oskar/kb/mail/archive")
@@ -126,17 +128,6 @@ def parse_headers(raw: bytes) -> dict:
     }
 
 
-def _sanitize(value: Optional[str]) -> Optional[str]:
-    """Strip surrogates a compat32 bytes-parse can leave in header text.
-
-    Lone surrogates are not valid UTF-8 and postgres jsonb rejects them; each
-    undecodable byte degrades to '?' instead. Real Unicode passes through.
-    """
-    if value is None:
-        return None
-    return value.encode("utf-8", errors="replace").decode("utf-8")
-
-
 def parse_headers_fallback(raw: bytes) -> dict:
     """Degraded parse for messages the typed (policy.default) path rejects.
 

packages/kb-mail/src/kb_mail/__init__.py

@@ -1,5 +1,12 @@
 from .archive import save_eml
 from .db import get_envelope, insert_envelope
 from .envelope import Envelope
+from .text import sanitize_surrogates
 
-__all__ = ["Envelope", "insert_envelope", "get_envelope", "save_eml"]
+__all__ = [
+    "Envelope",
+    "insert_envelope",
+    "get_envelope",
+    "save_eml",
+    "sanitize_surrogates",
+]

packages/kb-mail/src/kb_mail/text.py

@@ -0,0 +1,20 @@
+"""Text helpers shared across mail-parsing jobs."""
+from __future__ import annotations
+
+from typing import Optional
+
+
+def sanitize_surrogates(value: Optional[str]) -> Optional[str]:
+    """Strip surrogates a compat32 bytes-parse can leave in header text.
+
+    Lone surrogates are not valid UTF-8 and postgres jsonb rejects them; each
+    undecodable byte degrades to '?' instead. Real Unicode passes through.
+
+    Shared by gmail-header-backfill (header text) and gmail-bulk-import
+    (Message-ID → envelope id): both parse Gmail Takeout .eml bytes under
+    compat32, where 8-bit header bytes surface as email.header.Header /
+    surrogate-laden strings that must not reach json.dumps or jsonb.
+    """
+    if value is None:
+        return None
+    return value.encode("utf-8", errors="replace").decode("utf-8")

packages/kb-mail/tests/test_text.py

@@ -0,0 +1,32 @@
+"""Unit tests for kb_mail.text.sanitize_surrogates."""
+from __future__ import annotations
+
+import json
+
+from kb_mail.text import sanitize_surrogates
+
+
+def test_none_stays_none():
+    assert sanitize_surrogates(None) is None
+
+
+def test_plain_ascii_unchanged():
+    assert sanitize_surrogates("hello") == "hello"
+
+
+def test_real_unicode_passes_through():
+    assert sanitize_surrogates("Kąpała") == "Kąpała"
+
+
+def test_lone_surrogate_degraded_and_json_safe():
+    # A compat32 bytes-parse can leave a lone surrogate (undecodable byte);
+    # postgres jsonb and json.dumps().encode('utf-8') both reject it.
+    dirty = "Pr\udce9sent"  # 0xe9 smuggled in as a surrogate escape
+    cleaned = sanitize_surrogates(dirty)
+    assert cleaned == "Pr?sent"
+    json.dumps(cleaned, ensure_ascii=False).encode("utf-8")  # must not raise
+
+
+def test_replacement_char_preserved():
+    # U+FFFD is already valid UTF-8 and must survive untouched.
+    assert sanitize_surrogates("bad<EFBFBD>id") == "bad<EFBFBD>id"