perf(sessions): reuse the authorized row for snapshot and usage totals by andrewreid · Pull Request #3421 · omnigent-ai/omnigent

andrewreid · 2026-07-28T09:19:50Z

Related issue

Stacked on #3419 (event-path-load-once), which is stacked on #3417 (policy-engine-load-once), which is stacked on #3407 (store-write-primitives) — this branch includes all three of those commits until they merge, so the diff below shows all four; only the last commit (perf(sessions): reuse the authorized row for snapshot and usage totals) is new here.

Summary

Fetching a single session's snapshot measured about nineteen queries in production — more than listing a page of twenty — with the conversation row read six times. The route handler already threads its authorized row into the snapshot builder, so that part was fine; the redundancy was a layer deeper. The snapshot's subtree-cost recompute re-read the conversation before walking the spawn tree, the same redundancy already removed from the policy engine (#3417) and the event path (#3419), just never applied to this endpoint.

ELI5: the handler already had the session's row in hand from checking the caller's permissions. But when it went to add up spend across the session's sub-agents, it threw that row away and fetched a fresh copy of the session just to find the same tree root it already had.

Before:                                    After:
  authorize -> read conversation             authorize -> read conversation
  ... pass row into snapshot builder ...     ... pass row into snapshot builder ...
  compute subtree cost:                      compute subtree cost:
    read conversation AGAIN (for its root)     reuse the root already in hand
    walk spawn tree                            walk spawn tree (still fresh)

The snapshot now passes the row's tree root through, so the subtree-cost helper skips its own read; the usage report does the same per listed session, since each listed row already carries its root; and the relay's turn-completion handler reads the conversation once and shares it between the turn's pricing lookup and the subtree roll-up, instead of reading it twice for two different purposes. The tree scan itself always stays fresh — that's how a just-persisted sub-agent's spend shows up immediately.

A real correctness bug rides along in the relay handler, found while making this change: that shared read decides for the whole completion, and passing an absent row onward meant "the session is gone" was read as "re-read the row" by the usage accumulator but as "skip" by the cost roll-up. A session deleted and recreated under the same id had the turn's usage silently billed to the new row while the event that would have reported it was suppressed. Both now agree: no row means nothing to bill and nothing to publish.

Reuse here is a different contract from the policy engine's, worth being explicit about since it's easy to conflate the two: an enforcement path (#3417) keeps only immutable identity from a preloaded row and re-derives every mutable field, because it's deciding whether to gate a tool call. A snapshot renders the row the caller already authorized against — a looser, point-in-time contract. The response projection is genuinely mixed-epoch as a result: most of it comes straight off the authorized row; a handful of fields (liveness, lifecycle status, elicitations, items, the subtree cost rollup) are read fresh and independent of it; agent name, model, and context window are read fresh but resolved through the row's agent binding, so they describe the agent the row named rather than necessarily the one bound now. Two fields don't fit either bucket cleanly: skills looks agent-keyed but is actually fetched live from the runner by session id, and harness can be either epoch depending on whether the row carries a per-session override. An earlier, more granular field-by-field version of this docstring had already drifted from the actual schema by the time it was reviewed (an invented field name, a wrongly-claimed keying, an omitted field) — this version describes shapes that survive the schema changing, not the current field list itself.

Test Plan

uv run --no-sync pytest tests/server/routes/test_sessions_snapshot.py tests/server/routes/test_sessions_runner_relay.py tests/server/routes/test_usage_report.py -q
uv run --no-sync ruff check omnigent/server/routes/_sessions/orchestration.py omnigent/server/routes/usage.py
uv run --no-sync ruff format --check omnigent/server/routes/_sessions/orchestration.py omnigent/server/routes/usage.py

Fresh benchmark pass (this round), dev/benchmarks/omnigent, get_session P50: 5.2→5.3 ms SQLite (flat, within run-to-run noise at this scale), 8.6→7.0 ms Postgres (-18%). Smaller than #3405's list-endpoint win, as expected — a single-session fetch was already cheap in absolute terms; removing one redundant round trip still moves Postgres's network cost measurably, less so SQLite's in-process one.

Demo

N/A — backend/snapshot change, no UI surface.

Type of change

Test coverage

Coverage notes

Measured on both dialects: the snapshot emits four conversation SELECTs and the usage report two, asserted exactly rather than as ceilings, because a ceiling passed even with the fix removed. The relay's oracle counts how many times the reporting session's own row is read specifically, so it measures this change rather than the query volume of the paths beneath it. A full-revert mutation of this change is caught precisely: pricing, the subtree roll-up, and the ancestor walk each fall back to resolving their own row once the shared read is gone — four reads, not three, since the relay-status read was always separate — and the oracle pins that exact count rather than a looser bound. The billed-silently correctness fix is pinned by a test that deletes and recreates a session under the same id mid-turn and asserts both the usage total and the published event agree it's gone. Manual verification is the fresh benchmark pass above, run against this branch specifically.

Changelog

Fixed a bug where a session deleted and recreated during an active turn could have that turn's cost silently billed without the corresponding cost update ever appearing in the UI.

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>

…tion The per-event endpoint fires once per streamed chunk and accounts for most of the server's query volume during an active turn; one observed thirty-eight-minute turn cost roughly thirteen thousand queries. The redundancy this removes is one read per event. Runner routing loaded the whole conversation — row, metadata and labels — to use a single column. It now reads just the binding. The resolution stays current, which an event path resolving a runner per chunk requires: a rebind or host handoff that has already landed routes to the new runner. It does not fence the window between resolving and forwarding; that needs a binding generation on the wire and is out of scope here. The narrowed lookup keeps the existence semantics the full read carried, in both directions. The binding lives on the Omnigent metadata row while the AP conversations row is the existence authority — deletion is ordered AP row first, and a failed second transaction leaves orphaned metadata as a documented trade-off, acceptable only because reachability checks consult the AP row. Creation writes the same two rows in the opposite order, so either can outlive the other, and the two states mean different things: a deleted conversation must be reported missing, an existing one without a binding must be reported unbound. A metadata-rooted lookup answered the first correctly and turned the second into a false not-found. The lookup is therefore rooted in the AP row, by outer join on one bind and by seeding existence then overlaying bindings on two. Usage events reuse the handler's row where the value is immutable: the tree root passed to the subtree sum. They deliberately do NOT reuse it for the ancestor fan-out. What an ancestor's badge should read is a fact about publication time, not about when the request arrived, so the fan-out loads its own tree — reusing the one summed for this session meant two siblings publishing from their own request-start snapshots delivered a newer total followed by an older one, leaving the parent showing the smaller figure. That costs one tree load per usage event with ancestors, and the honest measurement is below. Deliberately unchanged: the per-event access check stays fresh, so a caller who loses access mid-stream stays gated; the native anti-replay clamp keeps its own read, because it must merge against the newest stored usage; and the message and tool-result paths keep fresh runner resolution, since input policy evaluation can hold those requests for seconds and a wrong-runner tool result returns 200 while dropping the result. `_ancestor_session_ids` likewise refuses a caller-supplied starting row: those ids decide where events are published, and a row read earlier can name a chain the session has left. Measured on both dialects, per table rather than as a ceiling on one of them: an idle status event is 5 statements; a usage event is 4 conversations, 5 metadata and 4 labels. A ceiling on conversations alone hid the other tables and also passed with the anti-replay read removed. Oracles ship with the mutation that kills them. The status event is delivered through the app's real router — for four review rounds this test bound a runner without registering it, so routing reported it offline, delivery fell through to the process-wide client, and the assertion observed nothing about the read it protects; the process-wide client is now a trap that raises if used. The fan-out oracle constructs the interleave rather than hoping for it, committing a sibling's spend immediately after this request's first tree load. A leftover comment at the tree-scan call site claimed it fed both this session's own subtree total and the ancestor fan-out below, from back when the fan-out reused that tree — the fan-out has its own comment explaining why it now reads a fresh one instead, and the two disagreed. And the fan-out's regression test claimed a universal "badge never moves backwards" guarantee; narrowed to the specific request-start-snapshot regression it actually pins, since a sibling landing after the fan-out's own tree read but before publish can still be missed, identically to unmodified main. Signed-off-by: Andrew Reid <andrew@reid.ee>

Fetching a single session measured about nineteen queries in production — more than listing a page of twenty — with the conversation read six times. The handler already threaded its authorized row into the snapshot builder; the cost was a layer deeper. The subtree-cost recompute re-read the conversation triplet before walking the spawn tree, the same redundancy already removed from the policy engine and the event path but never applied here. The snapshot now passes the row's tree root so the helper skips its own read; the usage report does the same per listed session, since each listed row already carries its root; and the relay's completion handler reads the conversation once and shares it between the turn's pricing lookup and the subtree roll-up. The tree scan itself stays fresh, which is how just-persisted sub-agent spend appears. That shared read decides for the whole completion. An absent row means the session is gone: nothing to bill, nothing to publish. Passing the absent row onward instead meant it read as "re-read the row" to the accumulator and as "skip" to the roll-up, so a session deleted and recreated under the same id had the turn's usage persisted to the new row while the event reporting it was suppressed — billed silently. Reuse here is point-in-time, which is a different contract from the policy engine's: an enforcement path keeps only immutable identity and re-derives everything else, while a snapshot renders the row the caller authorized. The projection is mixed-epoch, and the docstring describes it by group rather than enumerating every response field: most of the projection comes straight off the authorized row; a handful of fields (liveness, lifecycle status, elicitations, items, the subtree cost rollup) are read fresh and independent of it; and the agent name, model and context window are read fresh but resolved through the row's agent binding, so they describe the agent the row named rather than necessarily the one bound now. Two fields don't fit either bucket cleanly: skills looks agent-keyed but is actually fetched live from the runner by session id, and harness can be either epoch depending on whether the row carries a per-session override. A field-by-field version of this docstring existed briefly and had already drifted from the schema by the time it was reviewed — an invented field name, a wrongly-claimed keying, an omitted list of newer fields — which is why this describes shapes that survive the schema changing, not the current field list itself. Measured on both dialects: the snapshot emits four conversation SELECTs and the usage report two, asserted exactly rather than as ceilings, because a ceiling passed with the fix removed. The relay's oracle counts how many times the reporting session's own row is read, so it measures this change rather than the query volume of the paths beneath it — and its docstring now names the mutation it actually catches, since rooting the tree at the reporter is repaired by the tree loader in the PR that owns that contract. That docstring undercounted a full revert too: pricing, the subtree roll-up and the ancestor walk each fall back to resolving their own row once the shared read is gone, which is three, plus the one relay-status read that was always separate — four, not the three previously claimed. Signed-off-by: Andrew Reid <andrew@reid.ee>

github-actions · 2026-07-28T09:34:13Z

@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:

A screenshot or screen recording of the change, or
A link to a hosted video or GIF showing the new behaviour.

Use N/A only when the change has no user-visible effect whatsoever (e.g. a pure refactor or test-only change). If that's the case, uncheck the relevant type box and check Refactor / chore or Test / CI instead.

andrewreid · 2026-08-03T00:58:41Z

Closing this in favour of a reworked approach, PR to follow.

andrewreid added 4 commits July 27, 2026 12:19

github-actions Bot added the size/XL Pull request size: XL label Jul 28, 2026

github-actions Bot requested a review from dhruv0811 July 28, 2026 09:20

github-actions Bot assigned dhruv0811 Jul 28, 2026

github-actions Bot added the needs-demo PR needs a demo screenshot or recording label Jul 28, 2026

andrewreid closed this Aug 3, 2026

Navigation Menu

Uh oh!

perf(sessions): reuse the authorized row for snapshot and usage totals - #3421

perf(sessions): reuse the authorized row for snapshot and usage totals#3421
andrewreid wants to merge 4 commits into
omnigent-ai:mainfrom
andrewreid:upstream/session-snapshot-load-once

andrewreid commented Jul 28, 2026

Uh oh!

github-actions Bot commented Jul 28, 2026

Uh oh!

andrewreid commented Aug 3, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants