fix(store): make label seeding and session-state writes atomic - #3407
fix(store): make label seeding and session-state writes atomic#3407andrewreid wants to merge 1 commit into
Conversation
Three writes on the conversation store decided from a snapshot taken before the write, and lost or invented data when anything landed in between. Label seeding read the current labels, worked out which declared initials were missing, and upserted those. A policy write landing between the read and the upsert was overwritten back to its initial value. seed_labels_if_absent does insert-if-absent in one statement instead, so the database decides which keys are missing and a concurrent write survives. Session state was persisted by writing the whole blob back from the caller's snapshot. Two parallel tool calls on one session each held their own copy, so the second overwrote the first — two policies each incrementing a counter by one persisted one. mutate_session_state applies the caller's change inside a locked read-merge-write, serialised by SELECT ... FOR UPDATE where the dialect supports it and by BEGIN IMMEDIATE on SQLite, which takes the write lock before the first read. All three share one row-missing contract: a write that reads before it writes must not treat an absent row as empty. increment_session_usage did — absent row treated as empty, an UPDATE matching nothing, and the mutated total returned as persisted. Label seeding did too, differently: labels carry no foreign key, so seeding for a conversation that is gone left orphan rows and reported a snapshot nothing can read back. Both now raise, existence being checked in the same transaction as the write for the same reason the insert is insert-if-absent — a check the caller made earlier can already be stale. The contract is stated on the abstraction, so another implementation cannot phantom-write while nominally conforming. Two callers absorb that exception rather than propagate it, and they ship here because the primitive that raises owns its callers' handling: usage accumulation on the relay path and the cost-ask checkpoint mirrored to a tree root. A session deleted mid-turn has nothing to bill and a lost checkpoint re-prompts, so neither should fail a streaming turn. The engine's hot-state cache had the same "absence means something, but which thing" problem one layer up. apply_state_updates merges a caller's change into the persisted row and folds the authoritative result onto the in-memory cache — but a key the caller just DELETEd is, correctly, simply absent from that result, and a blanket union with the old cache put it right back: the delete took effect in the database and nowhere else, and the next evaluation in the same engine instance saw the pre-delete value again. The fix tracks which keys this batch of operations actually deleted, and only those are excluded from what survives from the old cache — a key missing from the merged result for any other reason (a sub-agent's inherited root-approval key, which is never part of its own row; a value seeded straight into the engine without ever being persisted) still survives, exactly as before. BREAKING: ConversationStore gains two abstract methods, so an out-of-tree subclass must implement them before upgrading. There is no compatible default — the whole point of both is that the decision happens inside one statement, and a base-class fallback built from the existing primitives would reintroduce exactly the race and the lost update being removed. SqlAlchemyConversationStore is the only in-tree implementation. Each oracle here is paired with the mutation that kills it. The seeding race is constructed rather than simulated: a store proxy commits a competing write while the helper reads its snapshot. The state race runs two threads whose transactions overlap; removing PostgreSQL's row lock loses an increment, and so does replacing SQLite's BEGIN IMMEDIATE with a deferred session. Skipping the existence check leaves orphan label rows the test then finds. Reverting the cost-ask mirror to snapshot-then-write is caught by observing which primitive the engine reaches for, since the visible outcome is identical either way. A delete of a session's own state key, a delete of a sub-agent's inherited approval key, and a delete of that same key name on a top-level session are each pinned separately, since an exemption scoped by key name instead of by which keys this call deleted gets the last of those three wrong. Signed-off-by: Andrew Reid <andrew@reid.ee>
|
@andrewreid This PR is a Bug fix, Feature, or UI / frontend change but the Demo section is missing or only contains a placeholder. These change types require a screenshot or screen recording so reviewers can see the new behaviour without checking out the branch. Please update the Demo section with:
Use |
github-actions
Bot
added
P3-low
|
/review |
|
|
@andrewreid this PR was closed by a maintainer. If you think that was a mistake, reply here and ask them to reopen it. |
|
Superseded by #4329 — rebased cleanly onto main with CI fixes (query-name lint, test helpers). |
6 participants
Related issue
Closes #3402
Summary
Two writes on
ConversationStoredecided what to write from a snapshot read before the write, rather than inside the write's own transaction, and each silently loses a concurrent update as a result.ELI5: two tool calls on the same session each read "the counter is 3", each add 1 in their head, each write back "4" — the database ends up at 4, not 5, and nothing ever complains.
seed_labels_if_absentnow does insert-if-absent in one statement, so the database — not a prior Python snapshot — decides which keys are missing, and a concurrent write survives.mutate_session_stateapplies the caller's change inside a locked read-merge-write —SELECT ... FOR UPDATEwhere the dialect supports it,BEGIN IMMEDIATEon SQLite to take the write lock before the first read — so overlapping writers serialize instead of losing an update. The policy engine's in-memory hot-cache merge was reworked alongside this: it now tracks which keys the current batch of operations actually deleted, and only excludes those from what survives the old cache, rather than a blanket union (which can't tell "missing because deleted" from "missing because never part of this row" — a sub-agent's root-inherited approval key is never part of its own row at all).increment_session_usagetreated an absent conversation row as an empty one — an UPDATE matching nothing, with the mutated total returned as if persisted. Label seeding had the sharper version: labels carry no foreign key, so seeding against a gone conversation left an orphan row behind. Both now raiseConversationNotFoundError, existence checked in the same transaction as the write, for the same reason the insert is insert-if-absent — a check made earlier can already be stale. The contract is stated on the abstraction itself, so an out-of-tree implementation can't phantom-write while nominally conforming.Two callers absorb that new exception rather than propagate it, and they ship here because the primitive that raises owns its callers' handling: usage accumulation on the relay path, and the cost-ask checkpoint mirrored to a session's tree root. A session deleted mid-turn has nothing to bill, and a lost approval checkpoint just re-prompts — neither should fail an in-flight streaming turn.
BREAKING:
ConversationStoregains two abstract methods (seed_labels_if_absent,mutate_session_state), so an out-of-tree subclass must implement them before upgrading. There's no compatible default — the whole point of both is that the decision happens inside one statement, and a base-class fallback built from the existing primitives would reintroduce exactly the race being removed.SqlAlchemyConversationStoreis the only in-tree implementation.Test Plan
Also ran the fresh benchmark pass from this round (
dev/benchmarks/omnigent) on the write-adjacent journeys, to confirm the atomicity fix doesn't cost meaningful latency:create_session45.3→48.0 ms SQLite / 21.5→22.8 ms Postgres,add_commentflat on both. Small, consistent, and expected — one existence-check query where the old unconditional upsert had none.Demo
N/A — backend/store change, no UI surface.
Type of change
Test coverage
Coverage notes
Every oracle here is paired with the mutation that kills it. The seeding race is constructed rather than simulated — a store proxy commits a competing write while the helper reads its snapshot. The state race runs two threads whose transactions overlap via a
threading.Barrier; removing PostgreSQL's row lock loses an increment, and so does replacing SQLite'sBEGIN IMMEDIATEwith a deferred session. Skipping the existence check leaves an orphan label row the test then finds. Reverting the cost-ask mirror to snapshot-then-write is caught by observing which primitive the engine reaches for, since the visible outcome is otherwise identical. A delete of a session's own state key, a delete of a sub-agent's inherited approval key, and a delete of that same key name on a top-level session are each pinned separately, since an exemption scoped by key name instead of by which keys the call actually deleted gets the last of those three wrong. Manual verification is the benchmark pass noted above, run against this branch specifically.Changelog
Fixed a race where concurrent tool calls or policy writes on the same session could silently lose a label or session-state update instead of both landing.