Summary

Clean, well-scoped fix that closes a real data-loss bug — approve. The chat policy-block path previously replaced a stored message with a separate delete_message + add_message; a crash or exception between the two left the session with no message at all. This PR folds both into a single ChatDatabase.upsert_message that runs the DELETE and INSERT inside one transaction, so a failed replacement rolls back as a unit and the original message always survives. The change is tightly focused (one new method, two call-site swaps, five tests) and the test coverage is genuinely good.

I verified the core behavior by exercising upsert_message in isolation: insert-when-None, replace-existing (new id minted), wrong-session safety (an id from another session is never deleted), and the atomic-rollback path under a simulated INSERT crash all behave as the tests assert. The single most important thing: the fix correctly resolves the pre-existing "make this transactional before a multi-writer backend" TODO it replaces — no follow-up debt left behind.

Issues Found

🟢 Minor — add_message/upsert_message duplication (database.py:461, 505)
upsert_message is add_message plus an optional leading DELETE — the INSERT block, JSON serialization, and session-timestamp touch are copied verbatim. Not blocking (the duplication is small and each method reads clearly on its own), but add_message could delegate to keep them in sync if the INSERT columns ever change:

        """Add a message to a session. Returns message ID."""
        return self.upsert_message(
            session_id,
            None,
            role,
            content,
            rag_sources=rag_sources,
            agent_steps=agent_steps,
            tokens_prompt=tokens_prompt,
            tokens_completion=tokens_completion,
            inference_stats=inference_stats,
        )

(Apply only if you want the consolidation — keeping them separate is a defensible call.)

Strengths

• The fix is the right shape. A single transaction is exactly how to make delete-then-insert atomic, and it reuses the existing _transaction() context manager (commit-on-success, rollback-and-reraise) rather than hand-rolling transaction handling.
• Excellent test coverage for a bug fix. test_upsert_message_atomic_on_insert_failure simulates a crash between DELETE and INSERT with a flaky connection and asserts the original row survives — that's the actual regression being fixed, tested directly. The wrong-session-safety and metadata round-trip cases are also well chosen.
• No silent fallbacks. The failure path re-raises; callers reassign the returned id, so the "new id is minted" semantics documented in the docstring are honored at both call sites.

Verdict

Approve. No blocking issues — the one nit above is optional. Two notes for the merge, not blockers: I couldn't run pytest/black in this environment (heavy deps unavailable), so please confirm the test-plan checkboxes locally before merge; the diff itself reads black-clean. No documentation change is warranted — upsert_message is an internal DB helper, not a public SDK/CLI surface.