perf(policies): load the conversation and spawn tree once per engine build - #3417
Closed
andrewreid wants to merge 2 commits into
Closed
perf(policies): load the conversation and spawn tree once per engine build#3417andrewreid wants to merge 2 commits into
andrewreid wants to merge 2 commits 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>
…build Policy evaluation sits on the PreToolUse critical path — the hook blocks on the verdict — and spent most of its time re-reading the same rows. build_policy_engine fetched the conversation about four times (root resolution, labels, session state, model override) and walked the spawn tree twice, because the session-wide gating seed and the per-node subtree seed each called load_session_usage, which does its own conversation read plus a full paged tree scan. One conversation read and one tree scan now feed everything. Both usage seeds derive from that list through a pure aggregation, so they stay semantically distinct: cost gating remains tree-wide, so a sub-agent gates against the whole session's spend, while the subtree total remains the per-node display figure. A caller that already holds the row can pass it and skip the read. A row the caller supplies is a HINT, not a fact. It names a tree, and loading that tree verifies the claim: if the conversation is not in it, the root is resolved again. Everything downstream — the rows, the root id, the policies attached to that root, the accounting sums — comes from the tree that verification produced. Deriving the root from the caller's row while taking rows from a corrected tree mixes two epochs, and a conversation deleted and recreated under a different root then seeded the old tree's spend. Mutable state is likewise re-derived rather than trusted: labels, session state, model override and agent binding all come from the verified tree, whoever read the row first, because a caller's preload and this function's own read are equally stale by the time a decision is made. A row absent from the tree is confirmed with one re-read and then fails closed. A tree that needed more than one page cannot vouch for its own rows — page one was read before page two — so identity is confirmed once in that case, which single-page trees never pay for. Also here, because it is the same tree: the ancestor cost re-publish used to do a conversation read plus a full tree scan PER ancestor, and derived the chain from a row read earlier in the request. It now walks the verified tree, so the whole fan-out costs one load and cannot publish to a chain that has since changed. A chain that cannot be walked to the root yields nothing rather than a prefix, since the caller publishes to every id returned. The tree also stopped excluding archived conversations. Archiving is a listing concern; the tree is an accounting structure. Excluding them let an archived root — or an archived mid-tree node, which orphaned its descendants from the walk — seed the enforcement total as $0 and allow a tool call over budget. Archived spend consequently appears in displayed totals too, which is the intended reading: the badge should agree with the gate. Measured on both dialects: 30 queries per build to 6, or 3 when the caller supplies the row. The whole authenticated route, by (tree size, whether the caller supplies the row): 11 on a one-page tree when supplied, 14 when not; 17 on a 101-node tree when supplied, 20 when not. The tree load pages, so cost is not independent of tree size, and the extra 3 on a paged tree over the one-page count are the paging confirmation above, a full conversation read — consistent at both tree sizes and both supplied/not-supplied. Counted as SQL statements rather than store calls, because a store-call count cannot see a helper that issues three statements per call. The route-level oracle below covers only the one-page shape; the 101-node figures are measured, not pinned by a test yet. Every oracle here is paired with the mutation that kills it, including the two that pin this round's fixes: deriving the root from the pre-refresh row fails the recreated-child test, and skipping the paged-tree confirmation fails the switch-during-paging test. Signed-off-by: Andrew Reid <andrew@reid.ee>
This was referenced Jul 28, 2026
Contributor
|
@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 |
TomeHirata
requested review from
TomeHirata
and removed request for
aravind-segu
August 7, 2026 03:46
13 tasks
Contributor
13 tasks
Contributor
|
@andrewreid this PR was closed by a maintainer. If you think that was a mistake, reply here and ask them to reopen it. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Related issue
Closes #3003
Stacked on #3407 (store-write-primitives) — this branch includes that PR's commit until it merges, so the diff below shows both; only the last commit (
perf(policies): load the conversation and spawn tree once per engine build) is new here.Summary
Policy evaluation sits on the PreToolUse critical path — the hook blocks on the verdict before a tool call can proceed.
build_policy_enginefetched the conversation about four times per call (root resolution, labels, session state, model override) and walked the spawn tree twice, because the session-wide gating seed and the per-node subtree seed each independently calledload_session_usage, which does its own conversation read plus a full paged tree scan.ELI5: instead of asking "who's the root, what are the labels, what's the state, what's the model, and what has everyone spent" as five separate trips to the database, the engine now makes one trip, reads everything off that one snapshot, and derives every other answer from it in memory.
One conversation read and one tree scan now feed everything. Both usage seeds derive from that same list through a pure aggregation, so they stay semantically distinct: cost gating remains tree-wide (a sub-agent gates against the whole session's spend), while the subtree total remains the per-node display figure. A caller that already holds the row can pass it and skip the read.
Two real correctness fixes ship in the same change, because they touch the same tree load:
Also included, because it's the same tree: the ancestor cost re-publish used to do a conversation read plus a full tree scan per ancestor, deriving the chain from a row read earlier in the request. It now walks the verified tree once, so the whole fan-out costs one load and can't publish to a chain that has since changed. A chain that can't be walked to the root yields nothing rather than a prefix, since the caller publishes to every id returned.
Test Plan
Measured on both dialects: 30 queries per build down to 6, or 3 when the caller supplies the row. The whole authenticated route, by (tree size, whether the caller supplies the row): 11 statements on a one-page tree when supplied, 14 when not; 17 on a 101-node tree when supplied, 20 when not — the tree load pages, so cost isn't independent of tree size, and the extra 3 on a paged tree are the paging confirmation, a full conversation read, consistent at both sizes. Counted as SQL statements rather than store calls, since a store-call count can't see a helper that issues three statements per call. The 101-node figures are measured, not pinned by an automated test yet — no wall-clock benchmark was run for this PR specifically (that harness doesn't seed spawn trees or exercise internal hook routes today); the statement-count matrix above is the evidence.
Demo
N/A — backend/policy-engine change, no UI surface.
Type of change
Test coverage
Coverage notes
Every oracle here is paired with the mutation that kills it, including the two that pin this round's fixes: deriving the root from the pre-refresh row fails a recreated-child test (a conversation deleted and recreated under a different root), and skipping the paged-tree confirmation fails a switch-during-paging test, parametrized over all four preload/no-preload × switch/delete combinations so a new variant is covered by construction. The archived-conversation fix is pinned by a test asserting an archived root's (and an archived mid-tree node's) spend still counts toward the gating total. The route-level statement-count oracle covers the one-page shape at both supplied/not-supplied; the 101-node figures in the Test Plan are measured but not yet pinned by an automated test.
Changelog
Fixed a bug where an archived session (or an archived sub-agent) could be gated as if it had spent nothing, letting a tool call proceed over its actual budget.