[policy-alert] Replace non-atomic delete+add in ChatDatabase with transactional upsert

[dislovelhl]

Context

Follow-up from PR #952. _persist_policy_block_if_needed in src/gaia/ui/_chat_helpers.py updates a persisted message via non-atomic delete_message → add_message. The author acknowledged this with a comment: "ChatDatabase is single-writer SQLite today; rtk make this replacement transactional before moving the chat DB to a multi-writer backend." This pattern appears twice in the file.

Risk

Between DELETE and INSERT a crash, exception, or future multi-writer access can leave the session without the policy-block message — silent data loss. With SQLite WAL mode enabled the window is narrow but non-zero.

Fix

Add ChatDatabase.upsert_message(session_id, msg_id, content, ...) backed by a single INSERT OR REPLACE or BEGIN IMMEDIATE transaction. Replace both call sites in _chat_helpers.py:

# Before
db.delete_message(request.session_id, persisted_policy_block_msg_id)
persisted_policy_block_msg_id = db.add_message(...)

# After
persisted_policy_block_msg_id = db.upsert_message(
    request.session_id, persisted_policy_block_msg_id, ...
)

Files

• src/gaia/ui/chat_database.py — add upsert_message
• src/gaia/ui/_chat_helpers.py — replace both delete+add call sites
• tests/unit/chat/ — add concurrent-write and crash-recovery tests

[github-actions]

Thanks @dislovelhl — the analysis lines up with the code. Confirmed against PR #952 (still open) and src/gaia/ui/database.py:

🟡 Verified non-atomic pattern, two call sites. Both are inside _run_agent in src/gaia/ui/_chat_helpers.py (introduced by #952):

• _persist_policy_block_if_needed does db.delete_message(...) then db.add_message(...) (the inline comment in the PR diff acknowledges this is a placeholder).
• The end-of-stream finalization repeats the same pattern when a policy_block message needs to be replaced with the full assistant response.

🟡 Cleanest fix uses existing transaction primitives. ChatDatabase already has a _transaction() context manager and a single shared connection (src/gaia/ui/database.py:399, :472). A new upsert_message(session_id, msg_id, role, content, ..., agent_steps=None, inference_stats=None) that wraps both the DELETE (when msg_id is not None) and the INSERT in one with self._transaction(): ... block keeps the helpers short:

def upsert_message(self, session_id, msg_id, role, content, **kwargs):
    with self._transaction():
        if msg_id is not None:
            self._conn.execute(
                "DELETE FROM messages WHERE id = ? AND session_id = ?",
                (msg_id, session_id),
            )
        # ... inline the existing INSERT from add_message, then bump session timestamp
        return cursor.lastrowid

Note: add_message doesn't take a target id, so a plain INSERT OR REPLACE won't preserve the original message id (the FE may rely on stable ids when re-rendering after reload). Easier to keep the delete-then-insert semantics but make them atomic via the existing _transaction() (BEGIN/COMMIT/rollback on exception). That sidesteps the multi-writer concern the PR comment flagged without changing the id model.

🟢 Suggested next steps:

• Since PR feat(ui): policy alert cards, notifications, and durable receipts #952 is still open, the cleanest path is to fold this into that PR (the author already left a TODO-style comment) rather than open a separate follow-up. cc @aaronn-amd / whoever's on feat(ui): policy alert cards, notifications, and durable receipts #952.
• Tests: a unit test that simulates an exception between the DELETE and the INSERT (e.g., monkeypatch add_message to raise) and asserts the original row is still present is enough — concurrent-write tests against single-writer SQLite mostly serialize through _lock already, so the crash-recovery case is the higher-value test.

If the PR is too far along to retouch, happy to be the follow-up tracker — let me know which direction the maintainers prefer. cc @kovtcharov-amd for routing.

[github-actions]

Pushed a fix to branch autofix/issue-987 — ready for a maintainer to open as a PR (the Actions bot isn't permitted to create PRs directly).

What was fixed: the two non-atomic delete_message → add_message sequences in _chat_helpers.py (_persist_policy_block_if_needed and the end-of-stream finalization) now go through a new ChatDatabase.upsert_message, which wraps the DELETE and INSERT in the existing _transaction() (src/gaia/ui/database.py). The replacement is now all-or-nothing — a crash between the two writes can no longer leave the session without its message. Delete-then-insert id semantics are preserved (a new id is minted), per the note in the issue thread, so no FE id-stability change.

How to verify:

1. git fetch origin autofix/issue-987 && git checkout autofix/issue-987
2. pip install -e .
3. python -m pytest tests/unit/chat/ui/test_database.py -x — the new test_upsert_message_atomic_on_insert_failure simulates a crash between the DELETE and INSERT and asserts the original row survives.

Lint and the full tests/unit/ suite (5343 passed) are green on the branch.