From 4a4095144f949ca66faaf3d85a78acb5b021d42a Mon Sep 17 00:00:00 2001 From: Andrew Reid Date: Sun, 26 Jul 2026 12:17:40 +0930 Subject: [PATCH 1/3] fix(store): make label seeding and session-state writes atomic MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- omnigent/runtime/policies/builder.py | 20 +- omnigent/runtime/policies/engine.py | 60 ++++- .../server/routes/_sessions/orchestration.py | 10 +- .../stores/conversation_store/__init__.py | 118 ++++++++- .../conversation_store/sqlalchemy_store.py | 248 ++++++++++++++++-- tests/runtime/policies/test_builder.py | 235 +++++++++++++++++ .../policies/test_session_cost_ask_routing.py | 64 +++++ .../integration/test_sessions_endpoints.py | 25 ++ tests/stores/test_conversation_store.py | 98 ++++++- 9 files changed, 821 insertions(+), 57 deletions(-) diff --git a/omnigent/runtime/policies/builder.py b/omnigent/runtime/policies/builder.py index 314a50e57c..eac290d43b 100644 --- a/omnigent/runtime/policies/builder.py +++ b/omnigent/runtime/policies/builder.py @@ -702,17 +702,15 @@ def _seed_and_load_labels( labels. """ existing = _load_existing_labels(conversation_id, conversation_store) - to_seed = { - key: ldef.initial - for key, ldef in label_defs.items() - if ldef.initial is not None and key not in existing - } - if to_seed: - conversation_store.set_labels(conversation_id, to_seed) - # Re-read to pick up the freshly seeded values plus any - # writes that landed concurrently from another workflow. - existing = _load_existing_labels(conversation_id, conversation_store) - return existing + declared = {key: ldef.initial for key, ldef in label_defs.items() if ldef.initial is not None} + if not declared: + return existing + # Insert-if-absent in one statement, NOT "diff against the snapshot then + # upsert": a policy write landing between the snapshot and the seed would + # be overwritten back to the initial value by an upsert. The database + # decides which keys are missing, and the returned snapshot is read in + # the same transaction as the insert. + return conversation_store.seed_labels_if_absent(conversation_id, declared) def _load_existing_labels( diff --git a/omnigent/runtime/policies/engine.py b/omnigent/runtime/policies/engine.py index 82409f3488..79da69b9a8 100644 --- a/omnigent/runtime/policies/engine.py +++ b/omnigent/runtime/policies/engine.py @@ -13,6 +13,7 @@ from __future__ import annotations +import contextlib from dataclasses import replace from typing import Any @@ -27,7 +28,7 @@ StateUpdate, StateUpdateAction, ) -from omnigent.stores.conversation_store import ConversationStore +from omnigent.stores.conversation_store import ConversationNotFoundError, ConversationStore # Number of recent conversation items the engine fetches from # the conversation store and threads onto :class:`EvaluationContext` @@ -557,9 +558,43 @@ def apply_state_updates(self, updates: list[StateUpdate]) -> None: else: session_ops.append(op) if session_ops: - for op in session_ops: - _apply_one(self._session_state, op) - self._store.set_session_state(self._conversation_id, self._session_state) + # Merge under a row lock rather than overwriting the whole blob + # from this engine's snapshot: concurrent evaluations (parallel + # tool calls on one session) each hold their own snapshot, and a + # blind write dropped the other's counter increments / appends. + def _merge(state: dict[str, Any]) -> None: + for op in session_ops: + _apply_one(state, op) + + # Closed over today's two-member StateUpdateAction enum (SET, DELETE): + # a future third action would need its own inclusion/exclusion here. + deleted_keys = {op.key for op in session_ops if op.action == StateUpdateAction.DELETE} + merged = self._store.mutate_session_state(self._conversation_id, _merge) + # Hot cache: persisted truth (``merged``) is authoritative for + # every key it holds. For a key it DOESN'T hold, though, absence + # is ambiguous — a blanket union with the old cache + # (``{**old, **merged}``) resolved it one way unconditionally, + # which is wrong for a genuine delete: a key this call just + # deleted is correctly missing from ``merged``, and re-adding it + # from the old cache silently undoes the delete. + # + # The two cases are told apart by what this call actually asked + # for, not by key identity: a key missing from ``merged`` that + # THIS BATCH deleted stays gone; a key missing from ``merged`` + # that this batch never touched was never part of this + # conversation's own persisted row to begin with (the two + # sub-agent root-inherited cost-approval keys, which are routed + # to the ROOT's row above and never reach ``session_ops`` at + # all; or any value seeded straight into the engine's + # constructor without ever being written through the store) and + # survives from the old cache, exactly as it did before this + # call. + preserved = { + key: value + for key, value in self._session_state.items() + if key not in merged and key not in deleted_keys + } + self._session_state = {**merged, **preserved} def _record_root_cost_ask_approved(self, op: StateUpdate) -> None: """ @@ -573,13 +608,20 @@ def _record_root_cost_ask_approved(self, op: StateUpdate) -> None: sub-agent); a top-level session writes through the normal per-conversation path (root == self). - :param op: The ``SET`` op carrying the approved checkpoint value, e.g. + :param op: The ``SET`` or ``DELETE`` op on one of the two reserved + cost-approval keys, e.g. ``StateUpdate(key=..., action=StateUpdateAction.SET, value=0.05)``. """ - root_conv = self._store.get_conversation(self._root_conversation_id) - root_state = dict(root_conv.session_state) if root_conv is not None else {} - _apply_one(root_state, op) - self._store.set_session_state(self._root_conversation_id, root_state) + # Same atomic merge as the per-conversation path: a sibling sub-agent + # approving concurrently must not lose this checkpoint (or vice versa). + # A root deleted while this sub-agent runs has nowhere to hold the + # checkpoint. Recording it here was always best-effort (a lost approval + # re-prompts, it does not overspend), so keep the in-memory mirror below + # and let the turn continue rather than failing the tool call. + with contextlib.suppress(ConversationNotFoundError): + self._store.mutate_session_state( + self._root_conversation_id, lambda state: _apply_one(state, op) + ) # Also mirror into this engine's hot in-memory state so a subsequent # evaluate() within the same sub-agent turn sees the approval (its # session_state was seeded from the root at construction, but a fresh diff --git a/omnigent/server/routes/_sessions/orchestration.py b/omnigent/server/routes/_sessions/orchestration.py index ecd9f74804..92ff50762b 100644 --- a/omnigent/server/routes/_sessions/orchestration.py +++ b/omnigent/server/routes/_sessions/orchestration.py @@ -961,7 +961,15 @@ def _accumulate_session_usage( model_delta["total_cost_usd"] = cost_delta delta["by_model"] = {llm_model: model_delta} - new_current = conversation_store.increment_session_usage(session_id, delta) + try: + new_current = conversation_store.increment_session_usage(session_id, delta) + except ConversationNotFoundError: + # The session was deleted while its turn was still streaming. There is + # no row to bill and nothing to publish; the store now refuses to report + # a total it did not persist, and failing the relay loop over a deleted + # session would be worse than dropping the increment. The primitive that + # raises ships this handler, so it can land without regressing callers. + return None # Per-user daily rollup (policy-gated; this is the per-turn delta). _record_daily_cost(conv, cost_delta, conversation_store) return _priced_cost_for_display(new_current) diff --git a/omnigent/stores/conversation_store/__init__.py b/omnigent/stores/conversation_store/__init__.py index 07cc4ea2c8..fde2a43436 100644 --- a/omnigent/stores/conversation_store/__init__.py +++ b/omnigent/stores/conversation_store/__init__.py @@ -3,6 +3,7 @@ import hashlib import time from abc import ABC, abstractmethod +from collections.abc import Callable from dataclasses import dataclass from typing import Any @@ -253,10 +254,19 @@ def runner_seen_is_fresh(last_seen: int | None, now: int | None = None) -> bool: class ConversationNotFoundError(Exception): """ - Raised when a required conversation row is missing. - - Store methods use this when absence is not a benign - no-op and the route layer must return a typed 404. + Raised when a store method needed a conversation row that isn't there. + + What the CALLER does with it varies by call site, deliberately: a route + handler resolving a conversation the user named should surface this as a + typed 404, since the caller asked about a specific id and got a genuine + answer of "no such thing". A best-effort write on a background path — + accumulating usage on a relay completion, mirroring a cost-ask + checkpoint to a tree root — instead catches it and treats it as a no-op: + the session vanished mid-turn, there is nothing left to record, and + failing the turn over a write that can no longer land would be worse + than silently dropping it. The exception's meaning ("this write did not + happen") doesn't change; whether that is 404-worthy is a decision each + caller makes for itself. """ @@ -832,6 +842,46 @@ def rename_conversation_if_title_matches( """ ... + @abstractmethod + def seed_labels_if_absent( + self, + conversation_id: str, + defaults: dict[str, str], + updated_at: int | None = None, + ) -> dict[str, str]: + """ + Insert declared initial labels for keys with none, then read back. + + Insert-if-absent, NOT upsert: seeding must never overwrite a value + another writer already persisted. Deciding "is this key missing?" + from a snapshot taken before the write would reset a concurrent + policy update to its initial value, so the check belongs in the + same statement as the insert. + + :param conversation_id: The conversation to seed, + e.g. ``"conv_abc123"``. + :param defaults: ``key -> initial value`` for declared labels. + Empty performs no write and just returns the snapshot — INCLUDING + for a conversation that does not exist, since there is nothing to + insert and therefore nothing that needs an existence check. + :param updated_at: Timestamp for inserted rows; ``None`` uses the + current wall-clock. + :returns: The conversation's full label snapshot after seeding, or + ``{}`` when *defaults* is empty and the conversation does not + exist (see above). + :raises ConversationNotFoundError: When *defaults* is non-empty and + the conversation does not exist. Implementations must check + existence in the same transaction as the insert whenever they + are about to insert: labels are not foreign-keyed, so a check + skipped entirely leaves orphan rows for a conversation that is + gone, and reports a write nothing can read back. This narrows + the exposure to a race window (an unlocked check followed by a + concurrent delete) rather than eliminating it outright — still + strictly better than ``set_labels``, which performs no check at + all. + """ + ... + @abstractmethod def set_labels( self, @@ -852,9 +902,13 @@ def set_labels( validation lives in ``PolicyEngine.apply_label_writes``). Callers that need "insert only if missing" semantics - (initial-value seeding — POLICIES.md §10) should check - ``conversation.labels`` first and filter the updates - to keys not already present; this method always + (initial-value seeding — POLICIES.md §10) must use + :meth:`seed_labels_if_absent`, NOT this method. Reading + ``conversation.labels`` and filtering the updates to the keys + that look missing is a check-then-write race: a value persisted + between the read and the write is overwritten by the initial + value. The insert-if-absent path leaves that decision to the + database, inside the same statement. This method always overwrites. :param conversation_id: The conversation to update, @@ -936,11 +990,15 @@ def set_session_state( """ Persist the full session-state snapshot for a conversation. - Overwrites the existing ``session_state`` JSON column with - the serialized *state* dict. Called by - :meth:`PolicyEngine.apply_state_updates` after applying - structured :class:`StateUpdate` operations to the hot - cache. + Overwrites the existing ``session_state`` JSON column with the + serialized *state* dict — last writer wins, including over keys the + caller never read. Use it to install a state wholesale (seeding, + import, tests). + + A caller applying updates to state it read earlier wants + :meth:`mutate_session_state` instead: two writers each holding their + own snapshot lose each other's changes here, which is why policy + evaluation no longer uses this method. :param conversation_id: The conversation to update, e.g. ``"conv_abc123"``. @@ -949,6 +1007,37 @@ def set_session_state( """ ... + @abstractmethod + def mutate_session_state( + self, + conversation_id: str, + mutate: Callable[[dict[str, Any]], None], + ) -> dict[str, Any]: + """ + Apply *mutate* to the persisted session state atomically. + + Read-merge-write under a row lock, so two concurrent writers + cannot lose each other's updates. Prefer this over + :meth:`set_session_state` whenever the new value depends on the + old one (counters, appends, checkpoints); a snapshot-then-write + pair drops anything persisted in between. + + :param conversation_id: The conversation to update. + :param mutate: Callable applied in place to the freshly read + state, inside the locked transaction. Must not make store + calls of its own. + :returns: The merged state as persisted. + :raises ConversationNotFoundError: When the conversation has no + row to merge into. Implementations must not treat an absent + row as empty state: applying the mutation, issuing an update + that matches nothing and returning the result reports a write + that did not happen. This is the shared contract for every + write on this store that reads before it writes — see also + :meth:`increment_session_usage` and + :meth:`seed_labels_if_absent`. + """ + ... + @abstractmethod def set_session_usage( self, @@ -1026,6 +1115,11 @@ def increment_session_usage( "by_model": {"claude-sonnet-4-6": {"input_tokens": 1000, "total_cost_usd": 0.05}}}``. :returns: The updated ``session_usage`` dict after the increment. + :raises ConversationNotFoundError: When the conversation has no row + to increment — the same row-missing contract as + :meth:`mutate_session_state`. Callers on streaming paths, where + a session can be deleted mid-turn, should treat this as nothing + to record rather than as a failure. """ ... diff --git a/omnigent/stores/conversation_store/sqlalchemy_store.py b/omnigent/stores/conversation_store/sqlalchemy_store.py index 3031c11728..258500e7a0 100644 --- a/omnigent/stores/conversation_store/sqlalchemy_store.py +++ b/omnigent/stores/conversation_store/sqlalchemy_store.py @@ -4,6 +4,7 @@ import json import logging +from collections.abc import Callable from typing import Any from sqlalchemy import ( @@ -20,6 +21,7 @@ text, update, ) +from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import QueryableAttribute, Session, aliased from sqlalchemy.sql.selectable import Subquery @@ -461,6 +463,66 @@ def _dialect_upsert_labels( session.execute(stmt) +def _insert_labels_if_absent( + session: Session, + conversation_id: str, + defaults: dict[str, str], + updated_at: int, +) -> None: + """ + Insert label rows only where the key is not already present. + + Distinct from :func:`_upsert_labels`, which overwrites. Used for + declared-initial seeding, where overwriting is precisely wrong: a + policy write that lands between a reader's snapshot and its seed + would be reset to the initial value. ``ON CONFLICT DO NOTHING`` + makes the seed a no-op for keys that already exist, so the decision + is made by the database rather than by a stale snapshot. + + Falls back to per-key ``INSERT`` attempts on dialects without + ``ON CONFLICT`` support, each in a nested transaction so a losing + race is absorbed rather than failing the caller. + + :param session: Active SQLAlchemy session. + :param conversation_id: Owning conversation id. + :param defaults: Non-empty ``key -> initial value`` mapping. + :param updated_at: Timestamp written on inserted rows. + """ + dialect = session.bind.dialect.name if session.bind is not None else "" + rows = [ + { + "conversation_id": conversation_id, + "key": key, + "value": value[:LABEL_VALUE_MAX_LEN], + "updated_at": updated_at, + } + for key, value in defaults.items() + ] + stmt: Any + if dialect in ("sqlite", "postgresql"): + if dialect == "sqlite": + from sqlalchemy.dialects.sqlite import insert as sqlite_insert + + stmt = sqlite_insert(SqlConversationLabel).values(rows) + else: + from sqlalchemy.dialects.postgresql import insert as pg_insert + + stmt = pg_insert(SqlConversationLabel).values(rows) + session.execute( + stmt.on_conflict_do_nothing( + index_elements=["workspace_id", "conversation_id", "key"], + ) + ) + return + for row in rows: + try: + with session.begin_nested(): + session.add(SqlConversationLabel(**row)) + except IntegrityError: + # Another writer inserted this key first — that value wins. + continue + + def _fetch_labels( session: Session, conversation_id: str, @@ -1250,6 +1312,62 @@ def set_labels( with self._conv_session() as session: _upsert_labels(session, conversation_id, updates, stamp) + def seed_labels_if_absent( + self, + conversation_id: str, + defaults: dict[str, str], + updated_at: int | None = None, + ) -> dict[str, str]: + """ + Insert declared initial labels for keys that have none, then read back. + + Seeding must not overwrite: a concurrent policy write landing + between a caller's snapshot and its seed would otherwise be reset + to the initial value. The insert is ``ON CONFLICT DO NOTHING``, so + existing keys keep their persisted value, and the post-seed + snapshot is read in the same transaction. + + :param conversation_id: The conversation to seed, + e.g. ``"conv_abc123"``. + :param defaults: ``key -> initial value``. Empty skips the write + and just returns the current snapshot — including for a + conversation that does not exist: with nothing to insert, + there is nothing that needs the existence check below. + :param updated_at: Timestamp for inserted rows (``None`` → now). + :returns: The conversation's labels after seeding, or ``{}`` when + *defaults* is empty and the conversation does not exist. + :raises ConversationNotFoundError: When *defaults* is non-empty and + the conversation does not exist. Labels are not foreign-keyed, + so an unchecked insert leaves orphan rows behind for a + conversation that was never there or has since been deleted — + the same phantom write the state and usage paths refuse, on + the third write path. The check below is an unlocked ``SELECT``, + not a locking read: it narrows the orphan-row window to a + concurrent delete landing between this check and the insert, + rather than eliminating it — still strictly better than + ``set_labels``, which performs no check at all. + """ + stamp = updated_at if updated_at is not None else now_epoch() + with self._conv_session() as session: + if defaults: + # Existence is checked inside the same transaction as the + # insert, for the same reason the insert is + # insert-if-absent: a check the caller made earlier can be + # stale by now. Unlocked, so it narrows rather than closes + # the race against a concurrent delete of this same row. + exists = session.execute( + select(SqlConversation.id).where( + SqlConversation.workspace_id == current_workspace_id(), + SqlConversation.id == conversation_id, + ) + ).first() + if exists is None: + raise ConversationNotFoundError( + f"Cannot seed labels for {conversation_id!r}: no conversation row exists." + ) + _insert_labels_if_absent(session, conversation_id, defaults, stamp) + return _fetch_labels(session, conversation_id) + def set_session_state( self, conversation_id: str, @@ -1277,6 +1395,102 @@ def set_session_state( .values(session_state=json.dumps(state)) ) + def mutate_session_state( + self, + conversation_id: str, + mutate: Callable[[dict[str, Any]], None], + ) -> dict[str, Any]: + """ + Apply *mutate* to the persisted session state under a row lock. + + Read-merge-write in ONE locked transaction, so concurrent writers + cannot lose each other's updates. ``set_session_state`` overwrites + the whole blob from a snapshot the caller read earlier, which drops + anything written in between — two policies each incrementing a + counter by 1 persisted 1, not 2. + + Locking mirrors :meth:`increment_session_usage`: ``SELECT … FOR + UPDATE`` where supported, and ``BEGIN IMMEDIATE`` on SQLite (via + the immediate session maker) so the read cannot race a second + writer's read. + + :param conversation_id: The conversation to update, + e.g. ``"conv_abc123"``. + :param mutate: Callable applied in place to the freshly read state + dict. Runs inside the locked transaction, so it must not + perform store calls of its own. + :returns: The merged state, as persisted. + :raises ConversationNotFoundError: When the conversation has no + metadata row. Nothing can be persisted in that case, so + returning a merged state would report a write that never + happened. + """ + return self._mutate_metadata_json( + conversation_id, "session_state", mutate, what="session state" + ) + + def _mutate_metadata_json( + self, + conversation_id: str, + column: str, + mutate: Callable[[dict[str, Any]], None], + *, + what: str, + ) -> dict[str, Any]: + """ + Locked read-merge-write of one JSON column on conversation metadata. + + The single row-missing contract for every read-modify-write primitive + on this table: no row means nothing can be persisted, so raise rather + than return a merged dict the database does not hold. Each caller used + to carry its own copy of this loop, and a fix applied to one of them + left the other reporting phantom writes. + + Locking is dialect-complementary: ``SELECT … FOR UPDATE`` where + supported, and ``BEGIN IMMEDIATE`` on SQLite (via the immediate + session maker), which takes the write lock before the first read so + two writers cannot each read a pre-merge snapshot. + + :param conversation_id: The conversation to update, + e.g. ``"conv_abc123"``. + :param column: Metadata column holding the JSON blob, either + ``"session_state"`` or ``"session_usage"``. + :param mutate: Callable applied in place to the freshly read dict. + Runs inside the locked transaction, so it must not perform store + calls of its own. + :param what: Human-readable name of the thing being written, used in + the error message, e.g. ``"session usage"``. + :returns: The merged dict, as persisted. + :raises ConversationNotFoundError: When the conversation has no + metadata row. + """ + import json + + with self._session_immediate() as session: + q = select(SqlConversationMetadata).where( + SqlConversationMetadata.workspace_id == current_workspace_id(), + SqlConversationMetadata.id == conversation_id, + ) + if self._meta_supports_for_update: + q = q.with_for_update() + meta = session.scalars(q).first() + if meta is None: + raise ConversationNotFoundError( + f"Cannot update {what} for {conversation_id!r}: no metadata row exists." + ) + raw = getattr(meta, column) + current: dict[str, Any] = dict(json.loads(raw)) if raw else {} + mutate(current) + session.execute( + update(SqlConversationMetadata) + .where( + SqlConversationMetadata.workspace_id == current_workspace_id(), + SqlConversationMetadata.id == conversation_id, + ) + .values(**{column: json.dumps(current)}) + ) + return current + def set_session_usage( self, conversation_id: str, @@ -1361,32 +1575,20 @@ def increment_session_usage( :param delta: Usage increments (see :meth:`ConversationStore.increment_session_usage`). :returns: The updated ``session_usage`` dict. + :raises ConversationNotFoundError: When the conversation has no + metadata row — the increment cannot be persisted, so reporting a + new total would be a lie. Callers on streaming paths where a + conversation can vanish mid-turn should treat this as "nothing to + accumulate", not as a failure. """ - import json - from omnigent.stores.conversation_store import apply_session_usage_delta - with self._session_immediate() as session: - q = select(SqlConversationMetadata).where( - SqlConversationMetadata.workspace_id == current_workspace_id(), - SqlConversationMetadata.id == conversation_id, - ) - if self._meta_supports_for_update: - q = q.with_for_update() - meta = session.scalars(q).first() - current: dict[str, Any] = ( - dict(json.loads(meta.session_usage)) if meta and meta.session_usage else {} - ) - apply_session_usage_delta(current, delta) - session.execute( - update(SqlConversationMetadata) - .where( - SqlConversationMetadata.workspace_id == current_workspace_id(), - SqlConversationMetadata.id == conversation_id, - ) - .values(session_usage=json.dumps(current)) - ) - return current + return self._mutate_metadata_json( + conversation_id, + "session_usage", + lambda current: apply_session_usage_delta(current, delta), + what="session usage", + ) def add_daily_cost(self, user_id: str, day_utc: str, delta_usd: float) -> None: """ diff --git a/tests/runtime/policies/test_builder.py b/tests/runtime/policies/test_builder.py index d03b73f8d0..934813cc96 100644 --- a/tests/runtime/policies/test_builder.py +++ b/tests/runtime/policies/test_builder.py @@ -1001,3 +1001,238 @@ def test_normalize_usage_for_engine_drops_display_fields() -> None: assert "by_model" not in normalized3 assert "policy_cost_usd" not in normalized3 assert normalized3["input_tokens"] == 0 + + +def test_initial_label_seed_does_not_clobber_a_concurrent_write( + conversation_store: SqlAlchemyConversationStore, +) -> None: + """ + Seeding declared initials must not overwrite a value written between the + snapshot read and the seed. + + The race is constructed, not simulated: a store proxy commits a policy + write while the seeding helper is reading its snapshot, so the helper + holds a view that is already stale by the time it decides what is + missing. Insert-if-absent leaves the persisted value alone because the + database makes that decision inside the same statement; the previous + diff-then-upsert path recomputed "missing" from the stale snapshot and + reset a live value to its initial. + """ + from omnigent.runtime.policies.builder import _seed_and_load_labels + from omnigent.spec.types import LabelDef + + conv = conversation_store.create_conversation() + + class _WriteDuringSnapshot: + """Commits a competing label write during the snapshot read.""" + + def __init__(self, inner: SqlAlchemyConversationStore) -> None: + self._inner = inner + self._fired = False + + def get_conversation(self, conversation_id: str): + row = self._inner.get_conversation(conversation_id) + if not self._fired: + self._fired = True + self._inner.set_labels(conversation_id, {"integrity": "7"}) + return row + + def __getattr__(self, name: str): + return getattr(self._inner, name) + + result = _seed_and_load_labels( + conversation_id=conv.id, + label_defs={"integrity": LabelDef(initial="0")}, + conversation_store=_WriteDuringSnapshot(conversation_store), # type: ignore[arg-type] + ) + + assert result["integrity"] == "7", "seed overwrote a concurrent write" + persisted = dict(conversation_store.get_conversation(conv.id).labels) + assert persisted["integrity"] == "7" + + +def test_increments_from_independent_snapshots_merge( + conversation_store: SqlAlchemyConversationStore, +) -> None: + """ + Two engines holding independent snapshots each INCREMENT one key, and + both increments survive. A blind whole-blob write persisted 1. + + Sequential by construction, and named for what it proves: the MERGE, not + the race. Serialisation of overlapping transactions is a store-level + property and is pinned there, against each dialect's own mechanism + (``test_two_real_writers_race_on_one_metadata_row``) — this test stays + green with the row lock removed and should not be read as covering it. + """ + from omnigent.spec.types import StateUpdate, StateUpdateAction + + conv = conversation_store.create_conversation(title="state-race") + spec = AgentSpec(spec_version=1, name="x") + engine_a = build_policy_engine( + spec=spec, conversation_id=conv.id, conversation_store=conversation_store + ) + engine_b = build_policy_engine( + spec=spec, conversation_id=conv.id, conversation_store=conversation_store + ) + + op = StateUpdate(key="risk", action=StateUpdateAction.INCREMENT, value=1) + engine_a.apply_state_updates([op]) + engine_b.apply_state_updates([op]) + + persisted = conversation_store.get_conversation(conv.id) + assert dict(persisted.session_state)["risk"] == 2 + + +def test_sets_from_independent_snapshots_preserve_other_keys( + conversation_store: SqlAlchemyConversationStore, +) -> None: + """A second engine's SET must not drop a key the first one wrote. + + Sequential, like its sibling above: the merge is the contract here. + """ + from omnigent.spec.types import StateUpdate, StateUpdateAction + + conv = conversation_store.create_conversation(title="state-keys") + spec = AgentSpec(spec_version=1, name="x") + engine_a = build_policy_engine( + spec=spec, conversation_id=conv.id, conversation_store=conversation_store + ) + engine_b = build_policy_engine( + spec=spec, conversation_id=conv.id, conversation_store=conversation_store + ) + + engine_a.apply_state_updates( + [StateUpdate(key="from_a", action=StateUpdateAction.SET, value="a")] + ) + engine_b.apply_state_updates( + [StateUpdate(key="from_b", action=StateUpdateAction.SET, value="b")] + ) + + state = dict(conversation_store.get_conversation(conv.id).session_state) + assert state["from_a"] == "a", state + assert state["from_b"] == "b", state + + +def test_delete_survives_the_hot_cache_overlay( + conversation_store: SqlAlchemyConversationStore, +) -> None: + """ + A DELETE must actually remove the key, not just from the persisted row. + + ``apply_state_updates`` merges under a store-side lock and then folds the + result onto the engine's in-memory cache. The persisted row was already + correct after a delete — ``mutate_session_state``'s callback pops the key + from the fresh state it is handed. The hot cache was not: a blanket union + ``{**old_cache, **merged}`` cannot express "this key is now gone", since a + key ``merged`` no longer has is simply missing from the right-hand side of + the union and the union keeps whatever the left-hand (stale) cache still + holds. The next evaluation reads that stale cache, not the store. + """ + from omnigent.spec.types import StateUpdate, StateUpdateAction + + conv = conversation_store.create_conversation(title="state-delete") + spec = AgentSpec(spec_version=1, name="x") + engine = build_policy_engine( + spec=spec, conversation_id=conv.id, conversation_store=conversation_store + ) + + engine.apply_state_updates( + [ + StateUpdate(key="risk", action=StateUpdateAction.SET, value=1), + StateUpdate(key="keep", action=StateUpdateAction.SET, value=2), + ] + ) + engine.apply_state_updates([StateUpdate(key="risk", action=StateUpdateAction.DELETE)]) + + assert "risk" not in engine.session_state, engine.session_state + assert engine.session_state["keep"] == 2, engine.session_state + + persisted = dict(conversation_store.get_conversation(conv.id).session_state) + assert "risk" not in persisted, persisted + assert persisted["keep"] == 2, persisted + + +def test_delete_of_a_root_inherited_key_does_not_resurrect_from_the_snapshot( + conversation_store: SqlAlchemyConversationStore, +) -> None: + """ + A sub-agent deleting its inherited approval key must not have the overlay + bring it back from the construction-time snapshot. + + This key never reaches ``session_ops`` at all for a sub-agent — a write + to it (SET or DELETE) is always diverted to + ``_record_root_cost_ask_approved``, which applies the op straight to the + hot cache itself. So this test doesn't exercise the merge/overlay + machinery; it pins that the diversion still produces the right answer for + a DELETE, since only SET was previously exercised anywhere. + """ + from omnigent.policies.schema import SESSION_COST_ASK_APPROVED_STATE_KEY + from omnigent.spec.types import StateUpdate, StateUpdateAction + + parent = conversation_store.create_conversation() + conversation_store.set_session_state(parent.id, {SESSION_COST_ASK_APPROVED_STATE_KEY: 0.05}) + child = conversation_store.create_conversation( + kind="sub_agent", parent_conversation_id=parent.id + ) + spec = AgentSpec(spec_version=1, name="x") + engine = build_policy_engine( + spec=spec, conversation_id=child.id, conversation_store=conversation_store + ) + # Inherited at construction, per build_policy_engine's root-seeding. + assert engine.session_state[SESSION_COST_ASK_APPROVED_STATE_KEY] == 0.05 + + engine.apply_state_updates( + [StateUpdate(key=SESSION_COST_ASK_APPROVED_STATE_KEY, action=StateUpdateAction.DELETE)] + ) + + assert SESSION_COST_ASK_APPROVED_STATE_KEY not in engine.session_state, engine.session_state + + # The diversion writes straight to the root's row, not just the child's + # cache — assert the persisted root, or a wrongly-local DELETE would pass. + root_state = dict(conversation_store.get_conversation(parent.id).session_state) + assert SESSION_COST_ASK_APPROVED_STATE_KEY not in root_state, root_state + + +def test_delete_of_the_same_key_name_on_a_top_level_session_removes_it( + conversation_store: SqlAlchemyConversationStore, +) -> None: + """ + The same key name is ordinary state on a top-level session, and a delete + of it must stick — it must not be mistaken for the sub-agent inheritance + case just because the name matches. + + For a top-level session (root == self) an op on + ``SESSION_COST_ASK_APPROVED_STATE_KEY`` goes through the same + ``session_ops``/merge path as any other key, never through + ``_record_root_cost_ask_approved`` (sub-agent only). The overlay tells + "genuinely never persisted" apart from "just deleted" by which keys THIS + call's ops named, not by a fixed key list — an earlier, key-list-based + version of this fix got exactly this case wrong. + """ + from omnigent.policies.schema import SESSION_COST_ASK_APPROVED_STATE_KEY + from omnigent.spec.types import StateUpdate, StateUpdateAction + + root = conversation_store.create_conversation(title="root-own-approval-key") + spec = AgentSpec(spec_version=1, name="x") + engine = build_policy_engine( + spec=spec, conversation_id=root.id, conversation_store=conversation_store + ) + + engine.apply_state_updates( + [ + StateUpdate( + key=SESSION_COST_ASK_APPROVED_STATE_KEY, + action=StateUpdateAction.SET, + value=0.05, + ) + ] + ) + assert engine.session_state[SESSION_COST_ASK_APPROVED_STATE_KEY] == 0.05 + + engine.apply_state_updates( + [StateUpdate(key=SESSION_COST_ASK_APPROVED_STATE_KEY, action=StateUpdateAction.DELETE)] + ) + + assert SESSION_COST_ASK_APPROVED_STATE_KEY not in engine.session_state, engine.session_state + persisted = dict(conversation_store.get_conversation(root.id).session_state) + assert SESSION_COST_ASK_APPROVED_STATE_KEY not in persisted, persisted diff --git a/tests/runtime/policies/test_session_cost_ask_routing.py b/tests/runtime/policies/test_session_cost_ask_routing.py index 581c2d2d4f..ca01e30c29 100644 --- a/tests/runtime/policies/test_session_cost_ask_routing.py +++ b/tests/runtime/policies/test_session_cost_ask_routing.py @@ -62,6 +62,18 @@ def _engine_on( ) +def _drop_metadata_row(store: SqlAlchemyConversationStore, conversation_id: str) -> None: + """Delete a conversation's metadata row, leaving nothing to mutate.""" + from sqlalchemy import delete as sa_delete + + from omnigent.db.db_models import SqlConversationMetadata + + with store._session() as session: + session.execute( + sa_delete(SqlConversationMetadata).where(SqlConversationMetadata.id == conversation_id) + ) + + def _set_approved(value: float) -> StateUpdate: return StateUpdate( key=SESSION_COST_ASK_APPROVED_STATE_KEY, @@ -246,3 +258,55 @@ async def test_subagent_gate_sees_parent_spend_not_just_own( ) assert await _evaluate_bash(engine) == PolicyAction.ASK + + +def test_deleted_root_does_not_fail_the_subagent_turn( + conversation_store: SqlAlchemyConversationStore, +) -> None: + """A root deleted mid-turn must not turn an approval into a failed call. + + The store refuses to report a state mutation it could not persist, so the + root mirror can now raise where the old snapshot-then-write path silently + wrote nothing. Recording the checkpoint on the root is best-effort — + losing it re-prompts, it does not overspend — so the sub-agent's turn + continues with the approval still visible to this engine. + + Both halves are observed, because the visible outcome alone is identical + under the old implementation: the write goes through the locked primitive + (not a get/apply/set pair), and the missing-row error it raises is + swallowed here rather than propagating. + """ + parent = conversation_store.create_conversation() + child = conversation_store.create_conversation( + kind="sub_agent", parent_conversation_id=parent.id + ) + + used: list[str] = [] + + class _RecordingStore: + """Records which write primitive the engine reaches for.""" + + def __init__(self, inner: SqlAlchemyConversationStore) -> None: + self._inner = inner + + def mutate_session_state(self, conversation_id: str, mutate): + used.append("mutate_session_state") + return self._inner.mutate_session_state(conversation_id, mutate) + + def set_session_state(self, conversation_id: str, state): + used.append("set_session_state") + return self._inner.set_session_state(conversation_id, state) + + def __getattr__(self, name: str): + return getattr(self._inner, name) + + engine = _engine_on(_RecordingStore(conversation_store), child.id, parent.id) # type: ignore[arg-type] + _drop_metadata_row(conversation_store, parent.id) + + engine.apply_state_updates([_set_approved(0.05)]) + + assert "mutate_session_state" in used, used + assert "set_session_state" not in used, ( + "the approval must go through the locked merge, not a snapshot write" + ) + assert engine._session_state.get(SESSION_COST_ASK_APPROVED_STATE_KEY) == 0.05 diff --git a/tests/server/integration/test_sessions_endpoints.py b/tests/server/integration/test_sessions_endpoints.py index 09b0a140b6..4439d860b4 100644 --- a/tests/server/integration/test_sessions_endpoints.py +++ b/tests/server/integration/test_sessions_endpoints.py @@ -4918,6 +4918,31 @@ async def test_accumulate_session_usage_unpriced_model_has_tokens_no_cost( assert "total_cost_usd" not in usage["by_model"]["free-model"] +async def test_accumulate_session_usage_deleted_session_returns_none( + client: httpx.AsyncClient, + db_uri: str, +) -> None: + """A session deleted mid-stream is dropped, not raised, from the relay loop. + + ``increment_session_usage`` raises ``ConversationNotFoundError`` once the + row is gone; ``_accumulate_session_usage`` must swallow it and return + ``None`` rather than let it propagate and fail the in-flight turn. + """ + from omnigent.server.routes import sessions as sessions_routes + + agent = await create_test_agent(client) + session = await _create_session(client, agent["id"]) + store = SqlAlchemyConversationStore(db_uri) + assert await store.delete_conversation(session["id"]) is True + + result = sessions_routes._accumulate_session_usage( + {"usage": {"input_tokens": 1000, "output_tokens": 500, "model": "harness-model"}}, + session["id"], + store, + ) + assert result is None + + async def test_accumulate_session_usage_concurrent_calls_accumulate_both_deltas( client: httpx.AsyncClient, db_uri: str, diff --git a/tests/stores/test_conversation_store.py b/tests/stores/test_conversation_store.py index ad27389e3d..cf78a3244c 100644 --- a/tests/stores/test_conversation_store.py +++ b/tests/stores/test_conversation_store.py @@ -3,7 +3,7 @@ from __future__ import annotations import pytest -from sqlalchemy import text +from sqlalchemy import select, text from omnigent.db.utils import get_or_create_engine from omnigent.entities import ( @@ -5264,3 +5264,99 @@ def _item_search_text(self, item: NewConversationItem) -> str: ).scalars() ) assert stored == ["custom-search-text"] + + +def test_read_modify_write_primitives_report_a_missing_row( + conversation_store: SqlAlchemyConversationStore, +) -> None: + """No metadata row means no phantom write — for EVERY such primitive. + + Both used to treat the absent row as ``{}``, apply the change, run an + UPDATE that matched nothing, and return the mutated dict as if + persisted. Fixing one and leaving the other is how the second kept + reporting writes that never happened, under a comment describing the + bug as fixed. Driven off the shared primitive so a third caller is + covered by construction. + """ + from omnigent.stores.conversation_store import ConversationNotFoundError + + missing = "0" * 32 + writers = { + "session state": lambda: conversation_store.mutate_session_state( + missing, lambda state: state.__setitem__("counter", 1) + ), + "session usage": lambda: conversation_store.increment_session_usage( + missing, {"total_tokens": 10} + ), + # Labels are not foreign-keyed, so an unchecked insert here leaves + # orphan rows rather than failing — the same phantom write, on the + # write path that was missed the first time. + "labels": lambda: conversation_store.seed_labels_if_absent(missing, {"risk": "seed"}), + } + for what, write in writers.items(): + with pytest.raises(ConversationNotFoundError, match=r"no .* row exists"): + write() + assert conversation_store.get_conversation(missing) is None, what + # Nothing was written on the way out, by any of them. + with conversation_store._conv_session() as session: + from omnigent.db.db_models import SqlConversationLabel + + orphans = session.execute( + select(SqlConversationLabel.key).where(SqlConversationLabel.conversation_id == missing) + ).all() + assert orphans == [], f"orphan label rows left behind: {orphans}" + + +def test_two_real_writers_race_on_one_metadata_row( + conversation_store: SqlAlchemyConversationStore, +) -> None: + """ + Two threads, two connections, overlapping read windows — a real race. + + Tests that merge two snapshots sequentially prove the merge but never + construct the race, so they stay green with the row lock removed. This + one holds each writer inside its own locked transaction (the mutate + callback runs there) and only releases when both have arrived, so the + writers' read windows overlap unless the store serialises them. + + Without the lock both writers read the same pre-state and the second + overwrites the first: ``risk`` ends at 1. With it, the second blocks at + the lock — never reaching the barrier, which is why the barrier has a + timeout rather than deadlocking — reads 1 and persists 2. + + Locking is dialect-complementary, so this is meaningful on both: + ``SELECT … FOR UPDATE`` on PostgreSQL, ``BEGIN IMMEDIATE`` on SQLite. + """ + import contextlib as _contextlib + import threading + + conv = conversation_store.create_conversation(title="real-race") + # Both threads try to meet here from inside their transactions. When the + # store serialises correctly only one can arrive, so the wait must expire + # rather than block forever. + barrier = threading.Barrier(2, timeout=0.5) + errors: list[BaseException] = [] + + def _increment() -> None: + def _mutate(state: dict) -> None: + with _contextlib.suppress(threading.BrokenBarrierError): + barrier.wait() + state["risk"] = state.get("risk", 0) + 1 + + try: + conversation_store.mutate_session_state(conv.id, _mutate) + except BaseException as exc: + errors.append(exc) + + threads = [threading.Thread(target=_increment) for _ in range(2)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=30) + + assert not errors, errors + persisted = dict(conversation_store.get_conversation(conv.id).session_state) + assert persisted["risk"] == 2, ( + f"one writer's increment was lost: {persisted} — the two transactions " + f"were not serialised on the metadata row" + ) From 7335d88e52256fbbbea3e11951499159bbe75130 Mon Sep 17 00:00:00 2001 From: Andrew Reid Date: Sun, 26 Jul 2026 12:20:48 +0930 Subject: [PATCH 2/3] perf(policies): load the conversation and spawn tree once per engine build MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- omnigent/runtime/policies/builder.py | 504 +++++++++--- omnigent/server/routes/_sessions/helpers.py | 30 +- .../server/routes/_sessions/orchestration.py | 35 +- .../server/routes/sessions/routes_hooks.py | 31 +- tests/runtime/policies/test_builder.py | 749 +++++++++++++++++- .../test_sessions_policy_evaluate.py | 170 ++++ .../test_sessions_mcp_proxy_policy_retry.py | 39 +- tests/server/routes/test_sessions_snapshot.py | 4 +- 8 files changed, 1442 insertions(+), 120 deletions(-) diff --git a/omnigent/runtime/policies/builder.py b/omnigent/runtime/policies/builder.py index eac290d43b..2cbe3f1566 100644 --- a/omnigent/runtime/policies/builder.py +++ b/omnigent/runtime/policies/builder.py @@ -4,11 +4,13 @@ Called at the top of ``_run_agent_loop``. Seeds any ``LabelDef.initial`` values that are not already present in -``conversation_labels`` using an -``INSERT ... ON CONFLICT DO NOTHING`` semantic so that two -concurrent workflows on the same conversation (the v2 case -tracked in POLICIES.md Open Q #6) never clobber each other's -view of a label's first value. +``conversation_labels`` via +:meth:`ConversationStore.seed_labels_if_absent`, whose +``INSERT ... ON CONFLICT DO NOTHING`` semantics mean two concurrent +workflows on the same conversation (the v2 case tracked in +POLICIES.md Open Q #6) never clobber each other's view of a label's +first value — the database decides which keys are missing, not a +snapshot taken before the write. Phase 2 scope: zero-policy and declared-policy paths both work; concrete Policy subclasses land in Phases 3+, and this builder @@ -18,6 +20,7 @@ from __future__ import annotations import logging +from dataclasses import dataclass from typing import Any import cachetools @@ -160,28 +163,6 @@ def _normalize_usage_for_engine(usage: dict[str, float]) -> dict[str, float]: return usage -def _subtree_usage_seed( - conversation_id: str, - conversation_store: ConversationStore, -) -> dict[str, float]: - """ - SUBTREE-scoped usage seed for the per-subagent cost budget. - - Unlike :func:`_policy_usage_seed` (which seeds from the whole session - tree via ``root_conversation_id``), this seeds from ``conversation_id`` - itself — so the budget gates on this conversation's own subtree cost - (itself + its descendants), not the whole session. - - :param conversation_id: Conversation to seed the subtree usage for, - e.g. ``"conv_child"``. - :param conversation_store: Store to read the subtree usage from. - :returns: Subtree usage seed dict; when an enforcement cost exists its - ``total_cost_usd`` is the enforcement total. - """ - usage = load_session_usage(conversation_id, conversation_store) - return _normalize_usage_for_engine(usage) - - def _resolve_session_owner_cached( conversation_id: str, conversation_store: ConversationStore, @@ -284,6 +265,8 @@ def build_policy_engine( spec: AgentSpec, conversation_id: str, conversation_store: ConversationStore, + conversation: Conversation | None = None, + expected_agent_id: str | None = None, connection_override: dict[str, str] | None = None, default_policies: list[PolicySpec] | None = None, policy_store: PolicyStore | None = None, @@ -301,10 +284,11 @@ def build_policy_engine( When declared labels have an ``initial`` value and no row exists yet in ``conversation_labels``, seeds via - ``ConversationStore.set_labels`` — but only for keys not - already persisted, so existing label state is never - clobbered. The hot cache is built from the freshly seeded - snapshot. + :meth:`ConversationStore.seed_labels_if_absent`, whose + insert-if-absent semantics leave any already-persisted value + untouched. The decision about which keys are missing is made by + the database inside the same statement, never from a snapshot read + beforehand. The hot cache is built from the post-seed snapshot. Policy run order: session policies (from the CRUD API) first, then agent spec policies, then *default_policies* @@ -324,6 +308,25 @@ def build_policy_engine( :param spec: The parsed agent spec. :param conversation_id: The conversation this workflow is running on, e.g. ``"conv_abc123"``. + :param conversation: The already-loaded conversation row for + ``conversation_id``, when the caller holds a current one — skips + the builder's own read. + + **Preload contract** (one rule, applied at every preload site): + a preloaded row may supply only IMMUTABLE identity — its ``id`` + and ``root_conversation_id``. Every mutable field the engine + depends on (labels, session_state, model_override) is re-derived + from a fresh read taken here. A row that has since disappeared + fails closed rather than authorizing from the snapshot. Callers + that rebuild an engine specifically to observe concurrent writes + (the native ASK gate's post-lock re-evaluation) pass ``None``. + :param expected_agent_id: The ``agent_id`` the caller resolved *spec* + from. ``agent_id`` is mutable (switch-agent) but selects the spec, + so it must be read before the engine exists and cannot be + re-derived here. Passing it lets the builder confirm it against + the fresh row and fail closed on a mismatch, instead of + authorizing an evaluation under the previous agent's guardrails. + ``None`` skips the check (callers with no spec/agent coupling). :param conversation_store: The store used for label reads and writes. Held by the engine for the life of the workflow. @@ -358,21 +361,37 @@ def build_policy_engine( guardrails = spec.guardrails agent_policy_specs: list[PolicySpec] = list(guardrails.policies or []) if guardrails else [] session_policy_specs = _load_session_policy_specs(conversation_id, policy_store) - # Session policies are per-conversation, but sub-agents must inherit - # the root conversation's policies so that guardrails set on the - # top-level session (e.g. via sys_add_policy) also govern spawned - # children. Load root policies and prepend them (root policies run - # first, then any child-specific overrides, matching the cost-budget - # root-seeding pattern below). - conv = conversation_store.get_conversation(conversation_id) - root_conversation_id = conv.root_conversation_id if conv is not None else conversation_id - if root_conversation_id != conversation_id: - root_policy_specs = _load_session_policy_specs(root_conversation_id, policy_store) - # Deduplicate: skip root policies already present on the child - # (keyed by policy name) to avoid double-evaluation. - child_names = {p.name for p in session_policy_specs} - root_policy_specs = [p for p in root_policy_specs if p.name not in child_names] - session_policy_specs = root_policy_specs + session_policy_specs + if conversation is not None and conversation.id != conversation_id: + # A misrouted preload would mix one session's labels, state, usage + # and model into another session's authorization decision — fail + # closed rather than build an engine from the wrong row. + raise OmnigentError( + f"preloaded conversation {conversation.id!r} does not match " + f"conversation_id {conversation_id!r}", + code=ErrorCode.INVALID_INPUT, + ) + conv = ( + conversation + if conversation is not None + else conversation_store.get_conversation(conversation_id) + ) + # The row in hand only SUGGESTS a tree root. Loading the tree verifies it + # and reports the root actually used, so everything downstream — the rows, + # the root's own policies, the accounting sums — comes from one snapshot. + # Deriving the root from the pre-refresh row while taking rows from a + # corrected tree was the defect: a conversation deleted and recreated under + # another root seeded the OLD tree's spend. + verified = ( + load_verified_session_tree( + conversation_id, + conversation_store, + conv.root_conversation_id if conv is not None else None, + ) + if conv is not None + else VerifiedSessionTree([], conversation_id, False) + ) + tree = verified.rows + root_conversation_id = verified.root_conversation_id db_default_policy_specs = _load_default_policy_specs(policy_store) admin_policy_specs: list[PolicySpec] = db_default_policy_specs + list(default_policies or []) all_policy_specs = session_policy_specs + agent_policy_specs + admin_policy_specs @@ -384,22 +403,123 @@ def build_policy_engine( all_policy_specs.append(_ASK_ON_ADD_POLICY_SPEC) label_defs = (guardrails.labels or {}) if guardrails else {} + # One conversation read (``conv``, resolved above for policy + # inheritance) and ONE spawn-tree load feed everything below: labels, + # session state (own + inherited root keys), both usage seeds, and the + # model override — on the single-page happy path, which is nearly every + # build; a paged tree pays one further confirming read (see below). The + # helpers each re-fetched the same rows before — ~4x conversation reads + # plus two identical tree loads per build, the dominant cost of a + # policies/evaluate call. + # Freshness contract, applied to EVERY row this function decides from, + # regardless of how it arrived: ONLY immutable identity (id, + # root_conversation_id) survives from the row read above. Every mutable + # field — labels, session_state, model_override, agent_id — is taken from + # the tree load, which happened later and is therefore the newest read. + # + # The provenance of the earlier row does not change the hazard. A caller's + # preload and this function's own ``get_conversation`` are both snapshots + # taken before the tree scan, so both can be stale by the time a decision + # is made; gating the refresh on ``conversation is not None`` closed the + # window on one path and left the identical window open on the other. + # + # The tree load includes archived rows (archived conversations still hold + # spend), so a row missing from the tree has genuinely been deleted. + # Re-read once to confirm, then fail closed — never fall back to the + # earlier copy, which would authorize from state captured before whatever + # removed the row. + if conv is not None: + fresh_self = next((c for c in tree if c.id == conversation_id), None) + if fresh_self is None: + fresh_self = conversation_store.get_conversation(conversation_id) + if fresh_self is None: + # The row existed moments ago and is now gone (deleted + # mid-request). Authorizing from the earlier copy would decide + # against state that no longer exists; authorizing from empty + # state would seed a $0 budget and ALLOW. Fail closed. + raise OmnigentError( + f"Conversation {conversation_id!r} disappeared while building its " + f"policy engine; refusing to authorize from a stale snapshot.", + code=ErrorCode.CONFLICT, + ) + conv = fresh_self + # Agent/spec confirmation — deliberately AFTER the refresh above, and + # nowhere else. Comparing against the earlier row (as a previous revision + # did) validated the very snapshot whose staleness is the hazard, so a + # switch-agent in the window was accepted. Exact equality: a fresh row + # whose binding is ``None``, or no fresh row at all, is a mismatch too — + # not a reason to skip the check. + if expected_agent_id is not None: + fresh_agent_id = conv.agent_id if conv is not None else None + if fresh_agent_id != expected_agent_id: + raise OmnigentError( + f"Session {conversation_id!r} no longer resolves to agent " + f"{expected_agent_id!r} (now {fresh_agent_id!r}); the spec this " + f"engine would enforce is stale. Re-resolve and retry.", + code=ErrorCode.CONFLICT, + ) + # A paged tree does not share one read instant: rows on page one were read + # before page two, so "the tree read is newer than the caller's row" holds + # for the tree but not for any row inside it. Confirm identity once when + # that is actually the case — single-page trees, which is nearly all of + # them, pay nothing. + if conv is not None and verified.paged: + confirmed = conversation_store.get_conversation(conversation_id) + if ( + confirmed is None + or confirmed.root_conversation_id != root_conversation_id + or confirmed.agent_id != conv.agent_id + ): + raise OmnigentError( + f"Conversation {conversation_id!r} moved while its spawn tree was " + f"being paged; refusing to authorize against a tree assembled " + f"across the change.", + code=ErrorCode.CONFLICT, + ) + conv = confirmed + # Session policies are per-conversation, but sub-agents inherit the root + # conversation's policies so guardrails set on the top-level session (e.g. + # via sys_add_policy) also govern spawned children. Loaded from the + # VERIFIED root, after the refresh: reading them from the caller's + # suggested root inherited another tree's guardrails. + if root_conversation_id != conversation_id: + root_policy_specs = _load_session_policy_specs(root_conversation_id, policy_store) + # Deduplicate: skip root policies already present on the child + # (keyed by policy name) to avoid double-evaluation. + child_names = {p.name for p in session_policy_specs} + root_policy_specs = [p for p in root_policy_specs if p.name not in child_names] + session_policy_specs = root_policy_specs + session_policy_specs + all_policy_specs = ( + session_policy_specs + + agent_policy_specs + + admin_policy_specs + + [_ASK_ON_ADD_POLICY_SPEC] + ) + root_conv = ( + conv + if root_conversation_id == conversation_id + else next((c for c in tree if c.id == root_conversation_id), None) + ) + if root_conv is None and conv is not None and root_conversation_id != conversation_id: + # The tree now includes archived rows, so a missing root means the row + # is genuinely gone (deleted mid-request). Read once to confirm before + # deciding — never silently proceed on an absent root. + root_conv = conversation_store.get_conversation(root_conversation_id) initial_labels = _seed_and_load_labels( conversation_id=conversation_id, label_defs=label_defs, conversation_store=conversation_store, + existing=dict(conv.labels) if conv is not None else {}, ) - initial_session_state = _load_session_state(conversation_id, conversation_store) + initial_session_state = dict(conv.session_state) if conv is not None else {} # The cost-budget approval is per-SESSION: the whole spawn tree shares one # soft-threshold gate. A sub-agent runs as its own conversation, so seed its # approved-checkpoint from the ROOT conversation — otherwise approving on the # parent wouldn't carry to the sub-agent and it would re-ask at the same # threshold. Other session_state stays per-conversation; the matching # write-back is routed to the root by PolicyEngine.apply_state_updates. - # (conv and root_conversation_id already resolved above for policy - # inheritance — reuse them here.) if root_conversation_id != conversation_id: - root_state = _load_session_state(root_conversation_id, conversation_store) + root_state = dict(root_conv.session_state) if root_conv is not None else {} for _root_key in ( SESSION_COST_ASK_APPROVED_STATE_KEY, SESSION_COST_UNPRICED_APPROVED_KEY, @@ -409,15 +529,24 @@ def build_policy_engine( # Gating is SESSION-wide: seed from the whole spawn-tree total so a # sub-agent gates against the session's full spend (parent + siblings), # not just its own subtree. The cost read is the enforcement total - # (in-flight sub-agent spend); see _policy_usage_seed. - initial_usage = _policy_usage_seed(conversation_id, conversation_store) - # Conditional injection (#1a): only compute subtree usage when a - # subagent_cost_budget policy is present. - initial_subtree_usage = ( - _subtree_usage_seed(conversation_id, conversation_store) - if _needs_subtree_usage(all_policy_specs) - else None + # (in-flight sub-agent spend); see _policy_usage_seed, whose semantics + # (including the empty seed when the root row is missing) this + # preserves while reusing the single tree load. + initial_usage = ( + _normalize_usage_for_engine(_sum_subtree_usage(tree, root_conversation_id)) + if conv is not None and root_conv is not None + else {} ) + # Conditional injection (#1a): only compute subtree usage when a + # subagent_cost_budget policy is present. Per-node DISPLAY-rooted seed: + # same tree, rooted at the evaluated node instead of the root. + initial_subtree_usage: dict[str, float] | None = None + if _needs_subtree_usage(all_policy_specs): + initial_subtree_usage = ( + _normalize_usage_for_engine(_sum_subtree_usage(tree, conversation_id)) + if conv is not None + else {} + ) # Conditional injection (#1): only pay the owner + daily-cost lookups # when a per-user daily cost-budget policy is actually present. initial_user_daily_cost = ( @@ -425,7 +554,14 @@ def build_policy_engine( if _needs_user_daily_cost(all_policy_specs) else None ) - initial_model = _resolve_session_model(conversation_id, conversation_store, spec) + # Session model: the conversation's model_override (set when a user + # picks a model mid-session) wins over the spec's llm.model; None when + # neither is available and cost policies treat it as undeterminable. + initial_model = ( + conv.model_override + if conv is not None and conv.model_override + else (spec.llm.model if spec.llm else None) + ) # Pass the full ModelPricing so the engine can price cache-read and # cache-write tokens at their own rates via compute_llm_cost(). token_pricing = fetch_model_pricing(spec.llm.model) if spec.llm else None @@ -682,6 +818,7 @@ def _seed_and_load_labels( conversation_id: str, label_defs: dict[str, LabelDef], conversation_store: ConversationStore, + existing: dict[str, str] | None = None, ) -> dict[str, str]: """ Seed declared initial values and return the current snapshot. @@ -698,10 +835,14 @@ def _seed_and_load_labels( labels start unset until a policy writes them). :param conversation_store: Target for both the read and the seed UPSERT. + :param existing: Pre-loaded current label snapshot, passed by callers + that already hold the conversation row (saves a re-read). + ``None`` loads it here. :returns: Full post-seed snapshot of the conversation's labels. """ - existing = _load_existing_labels(conversation_id, conversation_store) + if existing is None: + existing = _load_existing_labels(conversation_id, conversation_store) declared = {key: ldef.initial for key, ldef in label_defs.items() if ldef.initial is not None} if not declared: return existing @@ -759,36 +900,6 @@ def _load_session_state( return dict(conv.session_state) -def _resolve_session_model( - conversation_id: str, - conversation_store: ConversationStore, - spec: AgentSpec, -) -> str | None: - """ - Resolve the model the session is currently using. - - Prefers the conversation's ``model_override`` (set when a user - picks a model mid-session via ``/model`` or the web model picker) - and falls back to the agent spec's ``llm.model``. ``None`` when - neither is available — the conversation does not exist yet, has no - override, and the spec declares no ``llm`` block — in which case - cost policies treat the model as undeterminable. - - :param conversation_id: Conversation to read the override from, - e.g. ``"conv_abc123"``. - :param conversation_store: Store to read the conversation from. - :param spec: The parsed agent spec (its ``llm.model`` is the - fallback when no override is set). - :returns: The active model id, e.g. ``"databricks-claude-opus-4-8"`` - or the native tier alias ``"opus"``; ``None`` when - undeterminable. - """ - conv = conversation_store.get_conversation(conversation_id) - if conv is not None and conv.model_override: - return conv.model_override - return spec.llm.model if spec.llm else None - - # Page size for walking a spawn tree when summing sub-agent usage. # Sub-agent trees are small in practice, but we still paginate so a # large tree is not silently truncated (see load_session_usage). @@ -839,6 +950,8 @@ def _merge_by_model( def load_session_usage( conversation_id: str, conversation_store: ConversationStore, + *, + root_conversation_id: str | None = None, ) -> dict[str, Any]: """ Load cumulative session usage for a conversation **plus all of its @@ -862,6 +975,14 @@ def load_session_usage( :param conversation_id: Conversation to load, e.g. ``"conv_abc123"``. :param conversation_store: Store to read from. + :param root_conversation_id: The conversation's tree root, when the + caller already holds the row. Skips the internal conversation read. + The root binding is immutable **per row**, not per conversation id: + a conversation deleted and recreated under the same id gets a new + row, whose root may differ. A supplied root is therefore validated + against the tree it produces — if this conversation is not in that + tree, the caller's row is stale and the root is resolved here + instead. ``None`` resolves it here from the start. :returns: Summed usage dict with keys ``input_tokens``, ``output_tokens``, ``total_tokens``, ``total_cost_usd`` (the DISPLAY cost sum — statusLine ``S`` for claude-native), and @@ -875,10 +996,123 @@ def load_session_usage( the policy seed (:func:`_policy_usage_seed`) reads ``policy_cost_usd`` (both unaffected by ``by_model``). """ + tree = load_session_tree(conversation_id, conversation_store, root_conversation_id) + return _sum_subtree_usage(tree, conversation_id) + + +def load_session_tree( + conversation_id: str, + conversation_store: ConversationStore, + root_conversation_id: str | None = None, +) -> list[Conversation]: + """ + Load the spawn tree *conversation_id* belongs to, verifying the root. + + One place owns the reuse rule for a caller-supplied tree root, so every + consumer gets the same guarantee: the tree comes back containing this + conversation, or the supplied root was stale and is resolved again. + + A caller's ``root_conversation_id`` is immutable **per row**. Deleting a + conversation and recreating it under the same id produces a new row that + may sit in a different tree, so a row read earlier in the request can + name a root this conversation no longer belongs to. Rather than trust it + or re-read unconditionally, the supplied root is checked against the + tree it produced — a membership test on rows already in memory, so the + happy path costs nothing and the stale path costs one read. + + :param conversation_id: The conversation whose tree is wanted, + e.g. ``"conv_abc123"``. + :param conversation_store: Store to read from. + :param root_conversation_id: Caller-supplied tree root, validated as + above. ``None`` resolves the root here. + :returns: Every conversation in the tree (root plus all descendants, + archived included). Empty when the conversation does not exist. + """ + return load_verified_session_tree( + conversation_id, conversation_store, root_conversation_id + ).rows + + +@dataclass(frozen=True) +class VerifiedSessionTree: + """A spawn tree together with what is known about how it was loaded. + + :param rows: Every conversation in the tree, archived included. Empty + when the conversation does not exist. + :param root_conversation_id: The root the rows were actually loaded + from, which is not necessarily the one the caller suggested. + :param paged: Whether the listing needed more than one page. Rows on an + earlier page were read before rows on a later one, so for a paged + tree "the tree read is newer than the caller's row" holds for the + tree as a whole but not for any individual row in it. + """ + + rows: list[Conversation] + root_conversation_id: str + paged: bool + + +def load_verified_session_tree( + conversation_id: str, + conversation_store: ConversationStore, + root_conversation_id: str | None = None, +) -> VerifiedSessionTree: + """ + Load a spawn tree and report the root it came from. + + Same verification as :func:`load_session_tree`, but callers that derive + more than the sums from a tree — the tree root itself, the policies + attached to that root — need to know which root was used, because a + supplied one may have been discarded. Deriving those from the caller's + root while taking the rows from a corrected tree mixes two epochs. + + :param conversation_id: The conversation whose tree is wanted. + :param conversation_store: Store to read from. + :param root_conversation_id: Caller-supplied root, treated as a hint. + :returns: The rows, the root they came from, and whether it paged. + """ + if root_conversation_id is None: + conv = conversation_store.get_conversation(conversation_id) + if conv is None: + return VerifiedSessionTree([], conversation_id, False) + root_conversation_id = conv.root_conversation_id + rows, paged = _load_tree_pages(root_conversation_id, conversation_store) + return VerifiedSessionTree(rows, root_conversation_id, paged) + + rows, paged = _load_tree_pages(root_conversation_id, conversation_store) + if any(c.id == conversation_id for c in rows): + return VerifiedSessionTree(rows, root_conversation_id, paged) + # Not in the tree the supplied root produced: either this conversation + # is gone, or it now lives in a different tree. Resolve it once. conv = conversation_store.get_conversation(conversation_id) if conv is None: - return {} - tree = _load_tree_conversations(conv.root_conversation_id, conversation_store) + return VerifiedSessionTree([], root_conversation_id, paged) + if conv.root_conversation_id == root_conversation_id: + return VerifiedSessionTree(rows, root_conversation_id, paged) + rows, paged = _load_tree_pages(conv.root_conversation_id, conversation_store) + return VerifiedSessionTree(rows, conv.root_conversation_id, paged) + + +def _sum_subtree_usage( + tree: list[Conversation], + conversation_id: str, +) -> dict[str, Any]: + """ + Sum usage across the subtree of *tree* rooted at *conversation_id*. + + Pure aggregation over an already-loaded spawn tree — no store reads. + :func:`build_policy_engine` loads the tree once and derives both the + session-wide gating seed (rooted at the tree root) and the per-node + subtree seed (rooted at the evaluated node) from the same list; + :func:`load_session_usage` wraps this for callers that start from a + conversation id. See :func:`load_session_usage` for the shape of the + returned dict. + + :param tree: All conversations in the spawn tree (from + :func:`_load_tree_conversations`); order-independent. + :param conversation_id: The subtree root to sum from. + :returns: Summed usage dict (see :func:`load_session_usage`). + """ subtree_ids = _subtree_conversation_ids(tree, conversation_id) totals: dict[str, Any] = {} # Per-model breakdown summed across the subtree, parallel to the flat sums. @@ -988,6 +1222,13 @@ def _load_tree_conversations( # (not just "default") are included in the tree. kind=None, root_conversation_id=root_conversation_id, + # Archived conversations still hold spend, and archiving must not + # move a budget gate: excluding them let an archive-after-preload + # (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. The tree is an accounting structure, not a + # user-facing listing. + include_archived=True, ) convs.extend(page.data) if not page.has_more or page.last_id is None: @@ -996,6 +1237,38 @@ def _load_tree_conversations( return convs +def _load_tree_pages( + root_conversation_id: str, + conversation_store: ConversationStore, +) -> tuple[list[Conversation], bool]: + """ + Page through a spawn tree, reporting whether more than one page was read. + + :param root_conversation_id: The tree's root conversation id. + :param conversation_store: Store to read from. + :returns: ``(rows, paged)`` — ``paged`` is ``True`` when the listing + needed a second page, which means the rows do not share one read + instant. + """ + convs: list[Conversation] = [] + after: str | None = None + pages = 0 + while True: + page = conversation_store.list_conversations( + limit=_SUBTREE_USAGE_PAGE_SIZE, + after=after, + kind=None, + root_conversation_id=root_conversation_id, + include_archived=True, + ) + pages += 1 + convs.extend(page.data) + if not page.has_more or page.last_id is None: + break + after = page.last_id + return convs, pages > 1 + + def _subtree_conversation_ids( tree: list[Conversation], conversation_id: str, @@ -1033,6 +1306,53 @@ def _subtree_conversation_ids( return subtree +def ancestor_ids_from_tree( + tree: list[Conversation], + conversation_id: str, +) -> list[str]: + """ + Walk a conversation's ancestor chain inside an already-loaded tree. + + The mirror of :func:`_subtree_conversation_ids`, and public for the + same reason the tree loader is: the ancestor chain must come from the + same freshly-read rows as the sums, not from a conversation row the + caller read earlier. ``parent_conversation_id`` is immutable per row, + but a conversation deleted and recreated under the same id gets a new + row with a new parent, so a caller's copy can name a chain that no + longer exists — and walking it publishes to the wrong sessions. + + Pure: no store reads, and no reads are needed, because a tree already + contains every row on the chain by construction. + + :param tree: All conversations in the spawn tree (from + :func:`load_session_tree`); order-independent. + :param conversation_id: The node to walk upward from, + e.g. ``"conv_child123"``. + :returns: Ancestor ids nearest-parent-first. Empty when the node is + top-level, absent from the tree, or its chain is cyclic. + """ + by_id = {c.id: c for c in tree} + ancestors: list[str] = [] + seen = {conversation_id} + current = by_id.get(conversation_id) + while current is not None and current.parent_conversation_id is not None: + parent_id = current.parent_conversation_id + if parent_id in seen: + # A cycle makes the whole chain untrustworthy, not just the rest + # of it: returning the part walked so far would publish a + # descendant's usage to whichever ids happened to come first. + return [] + parent = by_id.get(parent_id) + if parent is None: + # The parent link points outside this tree. Appending before + # checking published to a conversation that is not there. + return [] + ancestors.append(parent_id) + seen.add(parent_id) + current = parent + return ancestors + + def _load_default_policy_specs( policy_store: PolicyStore | None, ) -> list[PolicySpec]: diff --git a/omnigent/server/routes/_sessions/helpers.py b/omnigent/server/routes/_sessions/helpers.py index b284fe55eb..5bcad9e05c 100644 --- a/omnigent/server/routes/_sessions/helpers.py +++ b/omnigent/server/routes/_sessions/helpers.py @@ -6230,7 +6230,16 @@ def _build_policy_engine_from_spec_impl( spec: AgentSpec, session_id: str, conversation_store: ConversationStore, + conversation: Conversation | None = None, ) -> PolicyEngine: + """Build an engine for *spec*, reusing a conversation row when held. + + Every caller of this wrapper already loaded the conversation to + resolve *spec*; passing it lets the builder skip its own read. Only + the row's immutable identity is reused — the builder re-derives + labels, session_state and model from a fresh read (see + :func:`build_policy_engine`). + """ caps = get_caps() host_connection = ( caps.policy_llm_connection_factory() if caps.policy_llm_connection_factory else None @@ -6239,6 +6248,11 @@ def _build_policy_engine_from_spec_impl( spec=spec, conversation_id=session_id, conversation_store=conversation_store, + conversation=conversation, + # The spec was resolved from this row's agent binding; the builder + # confirms it against its own fresh read and fails closed if a + # switch-agent landed in between. + expected_agent_id=conversation.agent_id if conversation is not None else None, default_policies=caps.default_policies, policy_store=get_policy_store(), server_llm=caps.llm, @@ -6302,7 +6316,7 @@ async def _apply_pending_policy_ask_writes( if spec is None: return engine = await asyncio.to_thread( - _build_policy_engine_from_spec, spec, session_id, conversation_store + _build_policy_engine_from_spec, spec, session_id, conversation_store, conv ) # The label/state writes hit the DB synchronously too — keep them # off the loop. @@ -6410,11 +6424,13 @@ def _build_evaluation_context( harness=hook_harness, ) # REQUEST / RESPONSE — content is the user/assistant text. The wire ``data`` - # is a dict for the native command hooks (``{"text"|"content": ...}``), but - # may be a bare string — opencode's policy plugin sends the prompt text - # directly for ``PHASE_REQUEST``. Accept both, and NEVER raise here: a crash - # 500s the evaluate endpoint, which silently fails the request/result gate - # OPEN (the exact symptom that let cost-over-budget terminal prompts through). + # is a dict for every current first-party producer (``{"text"|"content": + # ...}``, including OpenCode's plugin, which sends ``{"text": ...}``), but + # a bare string is still accepted for ``PHASE_REQUEST`` for compatibility + # with older or third-party callers that send the prompt text directly. + # Accept both, and NEVER raise here: a crash 500s the evaluate endpoint, + # which silently fails the request/result gate OPEN (the exact symptom + # that let cost-over-budget terminal prompts through). if isinstance(data, str): text = data elif isinstance(data, dict): @@ -6690,7 +6706,7 @@ async def _evaluate_output_policy( return None engine = await asyncio.to_thread( - _build_policy_engine_from_spec, spec, session_id, conversation_store + _build_policy_engine_from_spec, spec, session_id, conversation_store, conv ) ctx = EvaluationContext( phase=Phase.RESPONSE, diff --git a/omnigent/server/routes/_sessions/orchestration.py b/omnigent/server/routes/_sessions/orchestration.py index 92ff50762b..60a46c9b22 100644 --- a/omnigent/server/routes/_sessions/orchestration.py +++ b/omnigent/server/routes/_sessions/orchestration.py @@ -67,6 +67,9 @@ resolve_ask_timeout, ) from omnigent.runtime.policies.builder import ( + _sum_subtree_usage, + ancestor_ids_from_tree, + load_session_tree, load_session_usage, ) from omnigent.runtime.policies.engine import PolicyEngine @@ -596,6 +599,7 @@ def _build_session_list_item( def _publish_subtree_cost_to_ancestors( conv_store: ConversationStore, session_id: str, + conv: Conversation | None = None, ) -> None: """ Re-publish each ancestor's subtree-summed cost after a child usage update. @@ -608,18 +612,35 @@ def _publish_subtree_cost_to_ancestors( display side.) For each ancestor of *session_id*, recompute its subtree priced cost and publish a ``session.usage`` event carrying it. + One tree load serves the whole walk. Every ancestor of this session is in + the same tree, so both the chain and each ancestor's sum are derived from + one set of freshly-read rows — previously each ancestor paged the tree + again, and the chain came from a conversation row that may have been read + before a concurrent delete/recreate moved this session elsewhere. + Sync (does store reads + SSE fan-out); call via :func:`asyncio.to_thread`, mirroring the elicitation ancestor-publish helpers. ``session_stream.publish`` is safe to call from a worker thread. - :param conv_store: Store used to discover ancestors and sum each - ancestor's subtree usage. + :param conv_store: Store used to load the tree. :param session_id: The child session whose usage just changed, e.g. ``"conv_child123"``. + :param conv: The child's already-loaded conversation row, when the caller + holds one. Only its ``root_conversation_id`` is used, and only as a + hint: :func:`load_session_tree` verifies the tree it names actually + contains this session and resolves the root itself when it does not. + The parameter belongs to this function, not to whichever caller first + needed it, so that no caller can be removed and leave a signature + behind that its remaining callers already depend on. :returns: None. """ - for ancestor_id in _ancestor_session_ids(conv_store, session_id): - ancestor_usage = load_session_usage(ancestor_id, conv_store) + tree = load_session_tree( + session_id, + conv_store, + conv.root_conversation_id if conv is not None else None, + ) + for ancestor_id in ancestor_ids_from_tree(tree, session_id): + ancestor_usage = _sum_subtree_usage(tree, ancestor_id) subtree_cost = _priced_cost_for_display(ancestor_usage) usage_by_model = _usage_by_model_for_display(ancestor_usage) if subtree_cost is None and usage_by_model is None: @@ -4750,7 +4771,7 @@ async def _evaluate_tool_call_policy( if spec is None: return None engine = await asyncio.to_thread( - _build_policy_engine_from_spec, spec, session_id, conversation_store + _build_policy_engine_from_spec, spec, session_id, conversation_store, conv ) try: @@ -4871,7 +4892,7 @@ async def _evaluate_input_policy( return None engine = await asyncio.to_thread( - _build_policy_engine_from_spec, spec, session_id, conversation_store + _build_policy_engine_from_spec, spec, session_id, conversation_store, conv ) ctx = EvaluationContext( phase=Phase.REQUEST, @@ -5783,7 +5804,7 @@ async def _handle_mcp_tools_call( # only) and TOOL_RESULT (both paths). Engine construction reads # session-policy specs and labels from the DB, so keep it off-loop too. engine = await asyncio.to_thread( - _build_policy_engine_from_spec, spec, session_id, conversation_store + _build_policy_engine_from_spec, spec, session_id, conversation_store, conv ) if is_retry: diff --git a/omnigent/server/routes/sessions/routes_hooks.py b/omnigent/server/routes/sessions/routes_hooks.py index d46f5f5582..99abafeaaa 100644 --- a/omnigent/server/routes/sessions/routes_hooks.py +++ b/omnigent/server/routes/sessions/routes_hooks.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio import json from typing import Any @@ -14,6 +15,7 @@ from fastapi.responses import Response from omnigent.codex_native_elicitation import codex_elicitation_id +from omnigent.entities import Conversation from omnigent.errors import ElicitationDeclinedError, ErrorCode, OmnigentError from omnigent.policies.types import ( PolicyAction, @@ -510,7 +512,14 @@ async def evaluate_policy( hook_elicitation_id = raw_elicitation_id data = event.get("data") or {} - conv = conversation_store.get_conversation(session_id) + # Reuse the row the ACL check already fetched — same point in the + # request, so no less fresh than reading it again here, one query + # fewer on the blocking PreToolUse path. Absent for admin callers + # (who bypass the conversation lookup) and when permissions are + # disabled, which fall back to their own read. + conv = access.conversation + if conv is None: + conv = await asyncio.to_thread(conversation_store.get_conversation, session_id) if conv is None: raise OmnigentError( f"Session {session_id!r} not found.", @@ -570,7 +579,7 @@ async def evaluate_policy( _caps.policy_llm_connection_factory() if _caps.policy_llm_connection_factory else None ) - def _build_engine() -> PolicyEngine: + def _build_engine(preloaded_conv: Conversation | None = None) -> PolicyEngine: """ Build a policy engine for this session from the loaded spec. @@ -579,6 +588,10 @@ def _build_engine() -> PolicyEngine: does not re-query it during ``evaluate``, so a fresh build is the only way to observe a concurrent sibling's just-recorded approval. + :param preloaded_conv: The conversation row this handler already + loaded, passed on the FIRST build only to skip the builder's + re-read. Rebuilds that must observe concurrent writes (the + ASK-gate re-evaluation) pass ``None`` for a fresh read. :returns: A :class:`PolicyEngine` seeded with the latest persisted state for ``session_id``. """ @@ -586,20 +599,30 @@ def _build_engine() -> PolicyEngine: spec=loaded.spec, conversation_id=session_id, conversation_store=conversation_store, + conversation=preloaded_conv, + # ``agent`` below was resolved from conv.agent_id; the builder + # re-reads the row and fails closed if it was rebound since. + expected_agent_id=agent.id, default_policies=_caps.default_policies, policy_store=get_policy_store(), server_llm=_caps.llm, host_connection=_host_conn, ) - engine = _build_engine() + engine = _build_engine(conv) # Use the turn-initiating human's identity (persisted at forward time) # so per-user policies gate on the correct actor even when the HTTP # caller is the runner's service-account credential. Falls back to # user_id for direct API callers and native-terminal sessions (whose # turns go via _dispatch_session_event_to_runner, which does not write # this label). - turn_actor = conv.labels.get(_TURN_ACTOR_LABEL) + # Read the actor from the engine's label snapshot, not from the row + # fetched at the top of this handler: the engine's labels come from a + # read taken after the agent/spec load, so a turn-actor label written + # in that window still gates on the right principal. (``agent_id`` + # cannot be treated the same way — it selects the spec the engine is + # built from, so it is necessarily read first.) + turn_actor = engine.labels.get(_TURN_ACTOR_LABEL) ctx = _build_evaluation_context( phase, data, event, actor=_build_actor(turn_actor or user_id) ) diff --git a/tests/runtime/policies/test_builder.py b/tests/runtime/policies/test_builder.py index 934813cc96..3c2458c32a 100644 --- a/tests/runtime/policies/test_builder.py +++ b/tests/runtime/policies/test_builder.py @@ -25,10 +25,13 @@ from __future__ import annotations import uuid +from collections.abc import Callable, Iterator +from contextlib import contextmanager from pathlib import Path import pytest +from omnigent.entities import Conversation from omnigent.runtime.policies.builder import build_policy_engine from omnigent.spec.parser import parse from omnigent.spec.types import ( @@ -961,7 +964,7 @@ def test_normalize_usage_for_engine_drops_display_fields() -> None: """ _normalize_usage_for_engine removes by_model and promotes policy_cost_usd. - Both _policy_usage_seed and _subtree_usage_seed use this helper to + Both the session-wide and subtree usage seeds use this helper to prepare usage for the engine: strip the display-only ``by_model`` breakdown, and swap ``policy_cost_usd`` to ``total_cost_usd`` for enforcement cost (falling back to ``total_cost_usd`` when no enforcement @@ -1003,6 +1006,207 @@ def test_normalize_usage_for_engine_drops_display_fields() -> None: assert normalized3["input_tokens"] == 0 +@contextmanager +def _count_sql(store: SqlAlchemyConversationStore) -> Iterator[list[str]]: + """Capture every statement the store executes, on either bind.""" + from sqlalchemy import event as sa_event + + seen: list[str] = [] + + def _on(conn, cursor, statement, params, context, many): + seen.append(statement) + + engines = {store._engine, store._conv_engine} + for engine in engines: + sa_event.listen(engine, "before_cursor_execute", _on) + try: + yield seen + finally: + for engine in engines: + sa_event.remove(engine, "before_cursor_execute", _on) + + +@pytest.mark.parametrize("preloaded", [False, True], ids=["no-preload", "preload"]) +def test_build_issues_one_read_and_one_tree_scan( + conversation_store: SqlAlchemyConversationStore, + preloaded: bool, +) -> None: + """ + One engine build costs one conversation read plus ONE spawn-tree scan, + and nothing at all for the read when the caller supplies the row. + + Counts SQL STATEMENTS, not store-method calls: the redundancy this + change removes was measured in queries, and a store-call count cannot + see a helper that issues three statements per call. The builder used to + re-fetch the conversation ~4x and walk the tree twice, once per usage + seed. + + Both seeds must stay correct and identical either way: session-wide + gating from the whole tree, subtree display from the node's own subtree. + """ + from omnigent.spec.types import FunctionPolicySpec, FunctionRef + + parent = conversation_store.create_conversation() + child = conversation_store.create_conversation( + kind="sub_agent", parent_conversation_id=parent.id, title=_sub_agent_title() + ) + sibling = conversation_store.create_conversation( + kind="sub_agent", parent_conversation_id=parent.id, title=_sub_agent_title() + ) + conversation_store.set_session_usage(parent.id, {"total_cost_usd": 0.10}) + conversation_store.set_session_usage(child.id, {"total_cost_usd": 0.05}) + conversation_store.set_session_usage(sibling.id, {"total_cost_usd": 0.03}) + child_row = conversation_store.get_conversation(child.id) if preloaded else None + + subagent_budget = FunctionPolicySpec( + name="subtree_budget", + on=None, + function=FunctionRef( + path="omnigent.policies.builtins.cost.subagent_cost_budget", + arguments={"max_cost_usd": 10.0}, + ), + ) + with _count_sql(conversation_store) as statements: + engine = build_policy_engine( + spec=AgentSpec(spec_version=1, name="child"), + conversation_id=child.id, + conversation_store=conversation_store, + conversation=child_row, + default_policies=[subagent_budget], + ) + + # Semantics preserved: session-wide gate vs per-subtree display. + assert engine.usage["total_cost_usd"] == pytest.approx(0.18) + assert engine._subtree_usage is not None + assert engine._subtree_usage["total_cost_usd"] == pytest.approx(0.05) + + # The tree scan is one paged listing; the conversation read is a triplet + # (row + metadata + labels) and disappears entirely with a preload. + # Shape first, then the total. SQLite's connection PRAGMAs are setup, not + # work the builder asked for. + executed = [" ".join(q.split()) for q in statements if not q.startswith("PRAGMA")] + tree_scans = [q for q in executed if "root_conversation_id = " in q] + point_reads = [q for q in executed if "FROM conversations" in q and "conversations.id = " in q] + assert len(tree_scans) == 1, executed + assert len(point_reads) == (0 if preloaded else 1), executed + # A conversation read is a triplet (row + metadata + labels), and so is the + # tree scan; the preload removes one whole triplet. + assert len(executed) == (3 if preloaded else 6), [q[:80] for q in executed] + + +def test_build_counts_archived_spend_and_inherits_approval( + conversation_store: SqlAlchemyConversationStore, +) -> None: + """ + Archiving must not move a budget gate. An archived root's own spend + still counts toward the session total, and its approval checkpoint is + still inherited — archiving is a listing concern, not an accounting + one. (This previously asserted the opposite, codifying an omission + that let an archive-after-preload seed $0 and ALLOW over budget.) + """ + from omnigent.policies.schema import SESSION_COST_ASK_APPROVED_STATE_KEY + + parent = conversation_store.create_conversation() + child = conversation_store.create_conversation( + kind="sub_agent", parent_conversation_id=parent.id, title=_sub_agent_title() + ) + conversation_store.set_session_usage(parent.id, {"total_cost_usd": 0.10}) + conversation_store.set_session_usage(child.id, {"total_cost_usd": 0.05}) + conversation_store.set_session_state(parent.id, {SESSION_COST_ASK_APPROVED_STATE_KEY: 9.99}) + conversation_store.update_conversation(parent.id, archived=True) + + engine = build_policy_engine( + spec=AgentSpec(spec_version=1, name="child"), + conversation_id=child.id, + conversation_store=conversation_store, + ) + + # Whole-tree total including the archived root's own spend. + assert engine.usage["total_cost_usd"] == pytest.approx(0.15) + assert engine.session_state[SESSION_COST_ASK_APPROVED_STATE_KEY] == pytest.approx(9.99) + + +def test_build_rejects_mismatched_preloaded_conversation( + conversation_store: SqlAlchemyConversationStore, +) -> None: + """A preloaded row for a different session must fail closed, not mix + one session's labels/state/usage into another's policy decision.""" + from omnigent.errors import OmnigentError + + a = conversation_store.create_conversation(title="a") + b = conversation_store.create_conversation(title="b") + row_b = conversation_store.get_conversation(b.id) + + with pytest.raises(OmnigentError, match="does not match"): + build_policy_engine( + spec=AgentSpec(spec_version=1, name="x"), + conversation_id=a.id, + conversation_store=conversation_store, + conversation=row_b, + ) + + +def test_build_with_preloaded_row_sees_concurrent_mutable_writes( + conversation_store: SqlAlchemyConversationStore, +) -> None: + """ + The preloaded row contributes only immutable identity: labels, + session_state, and model written AFTER the row was captured must + still reach the engine (re-derived from the fresh tree row), so a + concurrent guard-label write can't be skipped by a stale snapshot. + """ + conv = conversation_store.create_conversation(title="fresh-check") + stale_row = conversation_store.get_conversation(conv.id) + + # Writes that land between the handler's read and the engine build. + conversation_store.set_labels(conv.id, {"guard": "tripped"}) + conversation_store.set_session_state(conv.id, {"checkpoint": 1.5}) + conversation_store.update_conversation(conv.id, model_override="claude-sonnet-4-6") + + engine = build_policy_engine( + spec=AgentSpec(spec_version=1, name="x"), + conversation_id=conv.id, + conversation_store=conversation_store, + conversation=stale_row, + ) + + assert engine.labels.get("guard") == "tripped" + assert engine.session_state.get("checkpoint") == 1.5 + assert engine.model == "claude-sonnet-4-6" + + +def test_build_uses_fresh_state_for_a_row_archived_after_preload( + conversation_store: SqlAlchemyConversationStore, +) -> None: + """ + A row archived after the preload still authorizes from FRESH state. + + The tree now includes archived rows, so the fresh row is found there + and the preload contributes only identity. (This test previously + described the tree as EXCLUDING archived rows and asserted the + re-read branch — a premise the ``include_archived=True`` fix + inverted, which made the test unable to fail for its stated reason.) + """ + conv = conversation_store.create_conversation(title="archive-race") + stale_row = conversation_store.get_conversation(conv.id) + + conversation_store.set_labels(conv.id, {"guard": "tripped"}) + conversation_store.update_conversation( + conv.id, model_override="claude-opus-4-8", archived=True + ) + + engine = build_policy_engine( + spec=AgentSpec(spec_version=1, name="x"), + conversation_id=conv.id, + conversation_store=conversation_store, + conversation=stale_row, + ) + + # Mutable fields come from the fresh (archived) row, not the preload. + assert engine.model == "claude-opus-4-8" + assert engine.labels.get("guard") == "tripped" + + def test_initial_label_seed_does_not_clobber_a_concurrent_write( conversation_store: SqlAlchemyConversationStore, ) -> None: @@ -1051,6 +1255,170 @@ def __getattr__(self, name: str): assert persisted["integrity"] == "7" +def test_initial_label_seed_uses_the_atomic_store_operation( + conversation_store: SqlAlchemyConversationStore, +) -> None: + """ + Seeding must go through ``seed_labels_if_absent``, never ``set_labels``. + + A mechanism assertion on top of the semantic one above: the race is + only impossible while the decision about which keys are missing is made + inside the insert statement, so the choice of store operation is itself + part of the contract. + """ + from omnigent.runtime.policies.builder import _seed_and_load_labels + from omnigent.spec.types import LabelDef + + conv = conversation_store.create_conversation() + calls: list[str] = [] + + class _RecordingStore: + def __init__(self, inner: SqlAlchemyConversationStore) -> None: + self._inner = inner + + def __getattr__(self, name: str): + attr = getattr(self._inner, name) + if not callable(attr): + return attr + + def _wrapper(*args: object, **kwargs: object) -> object: + calls.append(name) + return attr(*args, **kwargs) + + return _wrapper + + _seed_and_load_labels( + conversation_id=conv.id, + label_defs={"integrity": LabelDef(initial="0")}, + conversation_store=_RecordingStore(conversation_store), # type: ignore[arg-type] + existing={}, + ) + + assert "seed_labels_if_absent" in calls, calls + assert "set_labels" not in calls, calls + + +def test_deleted_after_preload_fails_closed( + conversation_store: SqlAlchemyConversationStore, +) -> None: + """A row deleted after the preload must fail closed, not authorize + from the stale snapshot nor from empty ($0) state.""" + import asyncio + + from omnigent.errors import OmnigentError + + conv = conversation_store.create_conversation(title="deleted") + conversation_store.set_session_usage(conv.id, {"total_cost_usd": 5.0}) + stale_row = conversation_store.get_conversation(conv.id) + asyncio.run(conversation_store.delete_conversation(conv.id)) + + with pytest.raises(OmnigentError, match="disappeared"): + build_policy_engine( + spec=AgentSpec(spec_version=1, name="x"), + conversation_id=conv.id, + conversation_store=conversation_store, + conversation=stale_row, + ) + + +def test_agent_rebind_after_spec_resolution_fails_closed( + conversation_store: SqlAlchemyConversationStore, +) -> None: + """ + ``agent_id`` selects the spec, so it is read before the engine exists + and cannot be re-derived. The builder confirms it against the row it + freshly read — NOT against the caller's preload, which is the snapshot + whose staleness is the whole hazard. + + Four cases, all of which an earlier revision accepted because the + comparison ran before the fresh read (and skipped on ``None``): + + 1. rebind with no preload + 2. rebind WITH a stale preload naming the old agent <- the reproduction + 3. fresh row whose binding is ``None`` + 4. no fresh row at all (deleted) + """ + import asyncio + + from omnigent.errors import OmnigentError + + spec = AgentSpec(spec_version=1, name="x") + agent_a = "1" * 32 + + def _switch(conv_id: str, new_agent_id: str) -> None: + # switch_conversation_agent inserts the target agent row, so each + # switch needs its own id. + conversation_store.switch_conversation_agent( + conv_id, + new_agent_id=new_agent_id, + new_agent_name="other", + new_agent_bundle_location="other/bundle", + new_agent_description=None, + copy_model_settings=False, + carry_history_into_native=False, + presentation_labels={}, + previous_builtin_id=None, + ) + + # Matching agent proceeds (guard must not be a blanket refusal). + ok_conv = conversation_store.create_conversation(title="match", agent_id=agent_a) + assert build_policy_engine( + spec=spec, + conversation_id=ok_conv.id, + conversation_store=conversation_store, + expected_agent_id=agent_a, + ) + + # 1. Rebind, no preload. + c1 = conversation_store.create_conversation(title="rebind-nopreload", agent_id=agent_a) + _switch(c1.id, uuid.uuid4().hex) + with pytest.raises(OmnigentError, match="no longer resolves to agent"): + build_policy_engine( + spec=spec, + conversation_id=c1.id, + conversation_store=conversation_store, + expected_agent_id=agent_a, + ) + + # 2. Rebind WITH a stale preload still naming agent A. Comparing against + # that preload would have found agent A == expected and accepted. + c2 = conversation_store.create_conversation(title="rebind-preload", agent_id=agent_a) + stale_row = conversation_store.get_conversation(c2.id) + assert stale_row.agent_id == agent_a + _switch(c2.id, uuid.uuid4().hex) + with pytest.raises(OmnigentError, match="no longer resolves to agent"): + build_policy_engine( + spec=spec, + conversation_id=c2.id, + conversation_store=conversation_store, + conversation=stale_row, + expected_agent_id=agent_a, + ) + + # 3. Fresh row with NO binding: previously skipped by an ``is not None`` + # guard, so an unbound session accepted any spec. + c3 = conversation_store.create_conversation(title="unbound") + assert conversation_store.get_conversation(c3.id).agent_id is None + with pytest.raises(OmnigentError, match="no longer resolves to agent"): + build_policy_engine( + spec=spec, + conversation_id=c3.id, + conversation_store=conversation_store, + expected_agent_id=agent_a, + ) + + # 4. No fresh row at all (deleted, no preload): also a mismatch. + c4 = conversation_store.create_conversation(title="gone", agent_id=agent_a) + asyncio.run(conversation_store.delete_conversation(c4.id)) + with pytest.raises(OmnigentError, match="no longer resolves to agent"): + build_policy_engine( + spec=spec, + conversation_id=c4.id, + conversation_store=conversation_store, + expected_agent_id=agent_a, + ) + + def test_increments_from_independent_snapshots_merge( conversation_store: SqlAlchemyConversationStore, ) -> None: @@ -1236,3 +1604,382 @@ def test_delete_of_the_same_key_name_on_a_top_level_session_removes_it( assert SESSION_COST_ASK_APPROVED_STATE_KEY not in engine.session_state, engine.session_state persisted = dict(conversation_store.get_conversation(root.id).session_state) assert SESSION_COST_ASK_APPROVED_STATE_KEY not in persisted, persisted + + +def test_supplied_root_is_a_hint_that_gets_verified( + conversation_store: SqlAlchemyConversationStore, +) -> None: + """ + A caller's tree root saves a read when right and is corrected when wrong. + + ``root_conversation_id`` is immutable per ROW, not per conversation id. + Delete a conversation and recreate it under the same id and the new row + can sit in a different tree, so a row read earlier in the request names a + root this conversation no longer belongs to. Trusting it silently summed + the wrong tree — the recreated child reported nothing at all. + + Three properties, because the argument has to be observably load-bearing + AND observably safe: + + 1. right root → same answer as self-resolving, one conversation read + fewer (a build that ignored the argument fails on the read count); + 2. wrong tree's root → still the right answer (a build that trusts it + fails here); + 3. the delete/recreate case that produced the regression. + """ + from omnigent.runtime.policies.builder import load_session_usage + + parent = conversation_store.create_conversation() + child = conversation_store.create_conversation( + kind="sub_agent", parent_conversation_id=parent.id, title=_sub_agent_title() + ) + conversation_store.set_session_usage(parent.id, {"total_cost_usd": 0.10}) + conversation_store.set_session_usage(child.id, {"total_cost_usd": 0.05}) + other_tree = conversation_store.create_conversation(title="unrelated") + + with _count_sql(conversation_store) as self_resolved_sql: + resolved = load_session_usage(child.id, conversation_store) + assert resolved["total_cost_usd"] == pytest.approx(0.05) + + with _count_sql(conversation_store) as supplied_sql: + supplied = load_session_usage(child.id, conversation_store, root_conversation_id=parent.id) + assert supplied == resolved + # The saved read is the point of the argument: a triplet fewer. + executed = lambda caught: [q for q in caught if not q.startswith("PRAGMA")] # noqa: E731 + assert len(executed(supplied_sql)) == len(executed(self_resolved_sql)) - 3, ( + executed(supplied_sql), + executed(self_resolved_sql), + ) + + # Wrong tree's root: verified against the tree it produced, so the sum is + # still right. Trusting it returned {} — the shape of the regression. + assert ( + load_session_usage(child.id, conversation_store, root_conversation_id=other_tree.id) + == resolved + ) + + +def test_recreated_conversation_is_summed_in_its_new_tree( + conversation_store: SqlAlchemyConversationStore, +) -> None: + """ + A row read before a delete/recreate must not decide which tree to sum. + + The reported regression: a child is deleted and recreated under the same + id in a different tree while a caller holds the old row. Summing from the + stale row's root found no such conversation and emitted nothing, so the + session's cost silently stopped updating. + """ + import asyncio + + from omnigent.runtime.policies.builder import load_session_usage + + old_parent = conversation_store.create_conversation(title="old-root") + child = conversation_store.create_conversation( + kind="sub_agent", parent_conversation_id=old_parent.id, title=_sub_agent_title() + ) + stale_row = conversation_store.get_conversation(child.id) + assert stale_row.root_conversation_id == old_parent.root_conversation_id + + # Recreate the same id under a different tree, as the delete/recreate + # boundary does. + asyncio.run(conversation_store.delete_conversation(child.id)) + new_parent = conversation_store.create_conversation(title="new-root") + conversation_store.create_conversation( + conversation_id=child.id, + kind="sub_agent", + parent_conversation_id=new_parent.id, + title=_sub_agent_title(), + ) + conversation_store.set_session_usage(child.id, {"total_cost_usd": 2.0}) + + summed = load_session_usage( + child.id, + conversation_store, + root_conversation_id=stale_row.root_conversation_id, + ) + assert summed["total_cost_usd"] == pytest.approx(2.0), ( + "a stale root must not silence a recreated conversation's spend" + ) + + +class _MutateOnTreeLoad: + """Store proxy that commits *hazard* the first time the tree is scanned. + + The builder reads the conversation, then scans the spawn tree. Anything + committed in between makes the earlier read stale — which is the window + the freshness refresh exists to close, whoever performed that read. + """ + + def __init__(self, inner: object, hazard: Callable[[], None]) -> None: + self._inner = inner + self._hazard = hazard + self._fired = False + + def list_conversations(self, *args: object, **kwargs: object) -> object: + if not self._fired: + self._fired = True + self._hazard() + return self._inner.list_conversations(*args, **kwargs) # type: ignore[attr-defined] + + def __getattr__(self, name: str) -> object: + return getattr(self._inner, name) + + +@pytest.mark.parametrize("preloaded", [True, False], ids=["preload", "no-preload"]) +@pytest.mark.parametrize("hazard", ["switch", "delete"]) +def test_mid_build_change_fails_closed_on_every_provenance( + conversation_store: SqlAlchemyConversationStore, + preloaded: bool, + hazard: str, +) -> None: + """ + A change landing between the conversation read and the tree scan must + fail closed — regardless of who performed that read. + + Parametrized over provenance on purpose. An earlier revision gated the + refresh on ``conversation is not None``, so the caller-preload path was + closed and the builder's own read had the identical window wide open: a + switch enforced the old agent's spec, a deletion seeded empty usage and + authorized a $0 budget. Provenance is a row here, not a separate test, + so a third way of acquiring the row is covered by construction. + """ + import asyncio + + from omnigent.errors import OmnigentError + + agent_a = "2" * 32 + conv = conversation_store.create_conversation(title=f"mid-build-{hazard}", agent_id=agent_a) + conversation_store.set_session_usage(conv.id, {"total_cost_usd": 5.0}) + row = conversation_store.get_conversation(conv.id) + + if hazard == "switch": + + def _apply() -> None: + conversation_store.switch_conversation_agent( + conv.id, + new_agent_id=uuid.uuid4().hex, + new_agent_name="other", + new_agent_bundle_location="other/bundle", + new_agent_description=None, + copy_model_settings=False, + carry_history_into_native=False, + presentation_labels={}, + previous_builtin_id=None, + ) + + expected = "no longer resolves to agent" + else: + + def _apply() -> None: + asyncio.run(conversation_store.delete_conversation(conv.id)) + + expected = "disappeared" + + store = _MutateOnTreeLoad(conversation_store, _apply) + with pytest.raises(OmnigentError, match=expected): + build_policy_engine( + spec=AgentSpec(spec_version=1, name="x"), + conversation_id=conv.id, + conversation_store=store, # type: ignore[arg-type] + conversation=row if preloaded else None, + expected_agent_id=agent_a, + ) + + +def test_archived_descendant_spend_counts_toward_the_displayed_total( + conversation_store: SqlAlchemyConversationStore, +) -> None: + """ + Including archived rows in the tree changes DISPLAY as well as gating. + + The change was made so archiving cannot reset a cost gate, but + ``load_session_usage`` is also the display path for the session badge + (``session.usage`` SSE) and the usage report, so an archived + descendant's spend now shows in the total a user sees. That is the + deliberate choice — the alternative, separate tree semantics for + display and enforcement, lets the badge disagree with the gate that + blocks the user — and it belongs with the change that caused it rather + than in a later PR. + """ + from omnigent.runtime.policies.builder import load_session_usage + + parent = conversation_store.create_conversation() + child = conversation_store.create_conversation( + kind="sub_agent", parent_conversation_id=parent.id, title=_sub_agent_title() + ) + conversation_store.set_session_usage(parent.id, {"total_cost_usd": 1.0}) + conversation_store.set_session_usage(child.id, {"total_cost_usd": 2.0}) + conversation_store.update_conversation(child.id, archived=True) + + displayed = load_session_usage(parent.id, conversation_store) + assert displayed["total_cost_usd"] == pytest.approx(3.0), ( + "archived descendant spend must reach the displayed subtree total" + ) + + # The same number the badge and the enforcement seed derive from, so the + # two cannot diverge. + engine = build_policy_engine( + spec=AgentSpec(spec_version=1, name="parent"), + conversation_id=parent.id, + conversation_store=conversation_store, + ) + assert engine.usage["total_cost_usd"] == pytest.approx(3.0) + + +def test_engine_seeds_the_new_tree_when_a_child_is_recreated_elsewhere( + conversation_store: SqlAlchemyConversationStore, +) -> None: + """ + Everything the build derives must come from the verified tree. + + The preload only *suggests* a root. Deriving the tree root from the + pre-refresh row while taking rows from a corrected tree mixed two + epochs: a child deleted and recreated beneath another root, with a + caller still holding the old row, seeded the OLD tree's spend — the + refresh landed on the row and left the root and the tree pointing at + the previous tree. (Root-inherited session *policies* take the same + refreshed root and are exposed to the same risk, but only spend is + asserted below — this test does not exercise policy inheritance.) + """ + import asyncio + + old_root = conversation_store.create_conversation(title="old-root") + child = conversation_store.create_conversation( + kind="sub_agent", parent_conversation_id=old_root.id, title=_sub_agent_title() + ) + conversation_store.set_session_usage(old_root.id, {"total_cost_usd": 10.0}) + stale_row = conversation_store.get_conversation(child.id) + + asyncio.run(conversation_store.delete_conversation(child.id)) + new_root = conversation_store.create_conversation(title="new-root") + conversation_store.set_session_usage(new_root.id, {"total_cost_usd": 3.0}) + conversation_store.create_conversation( + conversation_id=child.id, + kind="sub_agent", + parent_conversation_id=new_root.id, + title=_sub_agent_title(), + ) + + engine = build_policy_engine( + spec=AgentSpec(spec_version=1, name="x"), + conversation_id=child.id, + conversation_store=conversation_store, + conversation=stale_row, + ) + + # The gate seeds from the whole tree the child is in NOW. + assert engine.usage["total_cost_usd"] == pytest.approx(3.0), ( + "the engine seeded a tree the conversation no longer belongs to" + ) + + +@pytest.mark.parametrize("preload", [True, False], ids=["preloaded", "no-preload"]) +@pytest.mark.parametrize("mutation", ["switch", "delete"]) +def test_engine_refuses_a_tree_assembled_across_a_change( + conversation_store: SqlAlchemyConversationStore, + preload: bool, + mutation: str, +) -> None: + """ + A paged tree does not share one read instant, so it cannot vouch for + itself. + + Rows on page one are read before page two. The refresh takes the + evaluated row from the tree, so for a paged tree that row may predate + the last page — "the tree read is newer" is true of the tree, not of any + row in it. A switch OR a delete landing during paging must be caught + the same way, whether or not the caller preloaded the row — all four + preload/no-preload x switch/delete combinations previously proceeded. + """ + import asyncio + + from omnigent.errors import OmnigentError + from omnigent.runtime.policies import builder as builder_mod + + root = conversation_store.create_conversation(title="paged-root") + child = conversation_store.create_conversation( + kind="sub_agent", parent_conversation_id=root.id, title=_sub_agent_title() + ) + row = conversation_store.get_conversation(child.id) + + class _MutateDuringPaging: + """Switches or deletes the evaluated conversation between tree pages.""" + + def __init__(self, inner: SqlAlchemyConversationStore) -> None: + self._inner = inner + self._pages = 0 + self._fired = False + + def list_conversations(self, *args: object, **kwargs: object): + page = self._inner.list_conversations(*args, **kwargs) # type: ignore[arg-type] + self._pages += 1 + # Fire once the page holding the evaluated row has been read — + # which page that is depends on listing order, so keying on "page + # one" made the interleave dialect-dependent. + if not self._fired and any(c.id == child.id for c in page.data): + self._fired = True + if mutation == "switch": + self._inner.switch_conversation_agent( + child.id, + new_agent_id=uuid.uuid4().hex, + new_agent_name="other", + new_agent_bundle_location="other/bundle", + new_agent_description=None, + copy_model_settings=False, + carry_history_into_native=False, + presentation_labels={}, + previous_builtin_id=None, + ) + else: + asyncio.run(self._inner.delete_conversation(child.id)) + return page + + def __getattr__(self, name: str): + return getattr(self._inner, name) + + # One row per page, so the tree genuinely pages. + original_page_size = builder_mod._SUBTREE_USAGE_PAGE_SIZE + builder_mod._SUBTREE_USAGE_PAGE_SIZE = 1 + try: + with pytest.raises(OmnigentError, match="moved while its spawn tree"): + build_policy_engine( + spec=AgentSpec(spec_version=1, name="x"), + conversation_id=child.id, + conversation_store=_MutateDuringPaging(conversation_store), # type: ignore[arg-type] + conversation=row if preload else None, + ) + finally: + builder_mod._SUBTREE_USAGE_PAGE_SIZE = original_page_size + + +@pytest.mark.parametrize( + ("shape", "links"), + [ + ("cycle", {"a": "b", "b": "a"}), + ("parent outside the tree", {"a": "missing"}), + ], +) +def test_ancestor_walk_discards_an_untrustworthy_chain(shape: str, links: dict[str, str]) -> None: + """ + A chain that cannot be walked to the root yields nothing, not a prefix. + + The docstring promised empty for a cyclic or broken chain while the loop + returned the part it had already walked — and the caller publishes each + returned id a cost event, so `A → B → A` notified B off a cycle and + `C → missing-D` notified an id that is not in the tree at all. + """ + from omnigent.runtime.policies.builder import ancestor_ids_from_tree + + tree = [ + Conversation( + id=node, + created_at=0, + updated_at=0, + root_conversation_id="a", + parent_conversation_id=parent, + ) + for node, parent in links.items() + ] + + assert ancestor_ids_from_tree(tree, "a") == [], shape diff --git a/tests/server/integration/test_sessions_policy_evaluate.py b/tests/server/integration/test_sessions_policy_evaluate.py index 2d4f44f457..8253718fee 100644 --- a/tests/server/integration/test_sessions_policy_evaluate.py +++ b/tests/server/integration/test_sessions_policy_evaluate.py @@ -24,10 +24,13 @@ from __future__ import annotations import asyncio +from pathlib import Path from typing import Any import httpx import pytest +import pytest_asyncio +from fastapi import FastAPI from omnigent.runtime import get_caps, session_stream from omnigent.runtime.caps import RuntimeCaps @@ -1162,3 +1165,170 @@ async def test_llm_response_allow_when_no_matching_policy( assert body["result"] in ("POLICY_ACTION_ALLOW", "POLICY_ACTION_UNSPECIFIED"), ( f"Expected ALLOW or UNSPECIFIED for no-policy session, got {body['result']}." ) + + +_EVALUATE_USER = "alice@example.com" + + +@pytest.fixture() +def auth_app(runtime_init: None, db_uri: str, tmp_path: Path) -> FastAPI: + """App with permissions + header auth, for route-level SQL budgeting.""" + from omnigent.runtime.agent_cache import AgentCache + from omnigent.server.app import create_app + from omnigent.server.auth import UnifiedAuthProvider + from omnigent.stores.agent_store.sqlalchemy_store import SqlAlchemyAgentStore + from omnigent.stores.artifact_store.local import LocalArtifactStore + from omnigent.stores.comment_store.sqlalchemy_store import SqlAlchemyCommentStore + from omnigent.stores.file_store.sqlalchemy_store import SqlAlchemyFileStore + from omnigent.stores.permission_store.sqlalchemy_store import SqlAlchemyPermissionStore + + artifact_store = LocalArtifactStore(str(tmp_path / "artifacts")) + return create_app( + agent_store=SqlAlchemyAgentStore(db_uri), + file_store=SqlAlchemyFileStore(db_uri), + conversation_store=SqlAlchemyConversationStore(db_uri), + artifact_store=artifact_store, + agent_cache=AgentCache(artifact_store=artifact_store, cache_dir=tmp_path / "cache"), + comment_store=SqlAlchemyCommentStore(db_uri), + permission_store=SqlAlchemyPermissionStore(db_uri), + auth_provider=UnifiedAuthProvider(source="header"), + ) + + +@pytest_asyncio.fixture() +async def auth_client(auth_app: FastAPI, mock_llm: Any, tmp_path: Path): + """Async client against the auth-enabled app.""" + from omnigent.runtime import set_harness_process_manager + from omnigent.runtime.harnesses.process_manager import HarnessProcessManager + + pm = HarnessProcessManager(tmp_parent=tmp_path / "harness_pm") + await pm.start() + set_harness_process_manager(pm) + transport = httpx.ASGITransport(app=auth_app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as c: + yield c + mock_llm.release_all() + set_harness_process_manager(None) + await pm.shutdown() + + +async def _seed_authenticated_session(auth_client: httpx.AsyncClient, db_uri: str) -> str: + """Create a real agent-bound session and grant it to the test user. + + Uses the real create path (a genuine uploaded bundle) so the agent + cache can load the spec — a fabricated agent row with a fake bundle + location cannot be loaded and the route would bail before the work + this budget measures. + """ + import json as _json + + from omnigent.server.auth import LEVEL_OWNER + from omnigent.stores.permission_store.sqlalchemy_store import SqlAlchemyPermissionStore + from tests.server.helpers import build_agent_bundle + + perms = SqlAlchemyPermissionStore(db_uri) + perms.ensure_user(_EVALUATE_USER) + headers = {"X-Forwarded-Email": _EVALUATE_USER} + resp = await auth_client.post( + "/v1/sessions", + data={"metadata": _json.dumps({"title": "budget"})}, + files={ + "bundle": ( + "agent.tar.gz", + # The agent MUST declare a policy: with no policies the route + # short-circuits via ``any_policies_apply`` before building an + # engine, so a budget measured without one would not cover the + # preload path at all (it silently didn't, first time round). + build_agent_bundle( + name="budget-agent", + guardrails={ + "policies": { + "allow_all": { + "type": "function", + "on": ["tool_call"], + "function": "tests.runtime.policies.conftest._always_allow", + } + } + }, + ), + "application/gzip", + ) + }, + headers=headers, + ) + assert resp.status_code == 201, resp.text + session_id = resp.json()["session_id"] + perms.grant(_EVALUATE_USER, session_id, LEVEL_OWNER) + return session_id + + +# Measured SQL-statement budget for the AUTHENTICATED evaluate route — the +# blocking PreToolUse hook. Counted as executed statements, not store-method +# calls: a store-call oracle cannot see that one call fans out to three +# statements (conversation + metadata + labels), which is exactly how an +# earlier "6 queries" claim survived being wrong by five. +# +# Measured (both dialects) for the seed below: one owner grant, one agent +# declaring one tool_call policy, one conversation, no declared labels, no +# children. Composition: ACL resolution 3, the handler's conversation load +# 3, session-policy lookup, agent row, spawn-tree scan 3. +# +# This is the number the PR must quote — an earlier description claimed 6. +# It is also the oracle for the preload optimisation: dropping the +# ``conversation=`` argument makes the builder re-read the conversation and +# pushes this up, which no behavioural assertion can see because both +# variants return the same verdict. +_EVALUATE_ROUTE_SQL_BUDGET = 11 + + +@pytest.mark.asyncio +async def test_authenticated_evaluate_route_sql_budget( + auth_client: httpx.AsyncClient, + db_uri: str, +) -> None: + """ + Pin the route's SQL count so the description cannot drift from the code, + and so removing the preload is detectable. + + This is the oracle for the preload optimisation: dropping the + ``conversation=`` argument in ``_build_engine`` makes the builder read + the conversation again, which changes this count. A behavioural + assertion alone cannot see that, because both variants return the same + verdict. + """ + from sqlalchemy import event as sa_event + + from omnigent.db.utils import _engine_cache + + session_id = await _seed_authenticated_session(auth_client, db_uri) + headers = {"X-Forwarded-Email": _EVALUATE_USER} + payload = { + "event": { + "type": "PHASE_TOOL_CALL", + "data": {"name": "Bash", "arguments": {"command": "ls"}}, + } + } + # Warm the spec / policy caches so the count is the steady-state one. + await auth_client.post( + f"/v1/sessions/{session_id}/policies/evaluate", json=payload, headers=headers + ) + + engine = _engine_cache[db_uri] + statements: list[str] = [] + + def _on_exec(conn, cursor, statement, parameters, context, executemany): + if not statement.lstrip().upper().startswith("PRAGMA"): + statements.append(statement) + + sa_event.listen(engine, "before_cursor_execute", _on_exec) + try: + resp = await auth_client.post( + f"/v1/sessions/{session_id}/policies/evaluate", json=payload, headers=headers + ) + finally: + sa_event.remove(engine, "before_cursor_execute", _on_exec) + + assert resp.status_code == 200, resp.text + assert len(statements) == _EVALUATE_ROUTE_SQL_BUDGET, [ + s.split("\n")[0][:80] for s in statements + ] diff --git a/tests/server/routes/test_sessions_mcp_proxy_policy_retry.py b/tests/server/routes/test_sessions_mcp_proxy_policy_retry.py index 67c9e10180..dbbca6614b 100644 --- a/tests/server/routes/test_sessions_mcp_proxy_policy_retry.py +++ b/tests/server/routes/test_sessions_mcp_proxy_policy_retry.py @@ -76,6 +76,29 @@ def get_conversation(self, session_id: str) -> Conversation | None: return None +def _engine_factory_expecting_preload(engine: object): + """Return a `_build_policy_engine_from_spec` double that requires the row. + + Every production caller of that wrapper has already loaded the + conversation to resolve the spec, and passes it so the builder can skip + its own read and confirm the agent binding. A double that merely + *accepts* ``conversation=None`` asserts nothing: dropping the preload at + the call site leaves every test here green, which is exactly what + happened. This one fails if the row does not arrive, so the contract is + observed by all seven tests rather than described by none. + """ + + def _factory(spec, session_id, conversation_store, conversation=None): + assert conversation is not None, ( + "the handler must pass the conversation row it already read; " + "without it the builder re-reads and loses the agent-binding check" + ) + assert conversation.id == session_id, (conversation.id, session_id) + return engine + + return _factory + + @dataclass class _StubAgentStore: """ @@ -209,7 +232,7 @@ async def test_forged_retry_with_deny_policy_is_rejected( monkeypatch.setattr( sessions_mod, "_build_policy_engine_from_spec", - lambda spec, session_id, conversation_store: deny_engine, + _engine_factory_expecting_preload(deny_engine), ) response = await _handle_mcp_tools_call( @@ -266,7 +289,7 @@ async def test_forged_retry_with_ask_policy_rejects_unknown_elicitation( monkeypatch.setattr( sessions_mod, "_build_policy_engine_from_spec", - lambda spec, session_id, conversation_store: ask_engine, + _engine_factory_expecting_preload(ask_engine), ) forged_eid = "elicit_FORGED_never_issued" @@ -318,7 +341,7 @@ async def test_retry_with_allow_policy_falls_through( monkeypatch.setattr( sessions_mod, "_build_policy_engine_from_spec", - lambda spec, session_id, conversation_store: allow_engine, + _engine_factory_expecting_preload(allow_engine), ) response = await _handle_mcp_tools_call( @@ -374,7 +397,7 @@ async def test_retry_session_mismatch_still_rejected( monkeypatch.setattr( sessions_mod, "_build_policy_engine_from_spec", - lambda spec, session_id, conversation_store: deny_engine, + _engine_factory_expecting_preload(deny_engine), ) params = _forged_retry_params() @@ -432,7 +455,7 @@ async def test_legitimate_retry_with_pending_entry_proceeds( monkeypatch.setattr( sessions_mod, "_build_policy_engine_from_spec", - lambda spec, session_id, conversation_store: ask_engine, + _engine_factory_expecting_preload(ask_engine), ) eid = "elicit_LEGITIMATE_server_issued" @@ -533,8 +556,8 @@ async def test_non_mcp_entry_popped_by_events_handler_on_accept( monkeypatch.setattr( sessions_mod, "_build_policy_engine_from_spec", - lambda spec, session_id, conversation_store: _FixedPolicyEngine( - result=PolicyResult(action=PolicyAction.ALLOW, reason=None) + _engine_factory_expecting_preload( + _FixedPolicyEngine(result=PolicyResult(action=PolicyAction.ALLOW, reason=None)) ), ) @@ -612,7 +635,7 @@ async def test_mcp_proxy_runner_supplied_actor_reaches_policy_engine( monkeypatch.setattr( sessions_mod, "_build_policy_engine_from_spec", - lambda spec, session_id, conversation_store: engine, + _engine_factory_expecting_preload(engine), ) params = {"name": "sys_os_shell", "arguments": {"command": "echo hi"}} diff --git a/tests/server/routes/test_sessions_snapshot.py b/tests/server/routes/test_sessions_snapshot.py index ac4a095f48..859ebdb2a7 100644 --- a/tests/server/routes/test_sessions_snapshot.py +++ b/tests/server/routes/test_sessions_snapshot.py @@ -108,11 +108,13 @@ def list_conversations( after: str | None = None, kind: str | None = "default", root_conversation_id: str | None = None, + include_archived: bool = False, ) -> PagedList[Conversation]: """Return the spawn tree sharing ``root_conversation_id``. ``load_session_usage`` walks the tree via this method to sum a - parent's subtree usage. With an explicit graph, return every + parent's subtree usage, and passes ``include_archived=True`` — + archived conversations still hold spend. With an explicit graph, return every conversation sharing the root; otherwise synthesize the single childless conversation the legacy tests expect. """ From 05717549e8375e318d9026ab9eee7c2d06938256 Mon Sep 17 00:00:00 2001 From: Andrew Reid Date: Sun, 26 Jul 2026 12:24:06 +0930 Subject: [PATCH 3/3] perf(events): resolve the runner binding without loading the conversation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- omnigent/runner/routing.py | 35 +- omnigent/server/routes/_sessions/helpers.py | 15 +- .../server/routes/_sessions/orchestration.py | 28 +- .../server/routes/sessions/routes_events.py | 1 + .../conversation_store/sqlalchemy_store.py | 72 +++- tests/runner/test_routing.py | 77 ++++ .../routes/test_event_path_load_once.py | 328 ++++++++++++++++++ tests/stores/test_conversation_store.py | 50 +++ .../test_conversation_store_split_db.py | 33 ++ 9 files changed, 615 insertions(+), 24 deletions(-) create mode 100644 tests/server/routes/test_event_path_load_once.py diff --git a/omnigent/runner/routing.py b/omnigent/runner/routing.py index f38b9b2988..efa7af88e6 100644 --- a/omnigent/runner/routing.py +++ b/omnigent/runner/routing.py @@ -130,19 +130,28 @@ def client_for_session_resources(self, conversation_id: str) -> RoutedRunner: :raises OmnigentError: If the conversation is missing, the pinned runner is offline, or no online runner is available. """ - conv = self._conversation_store.get_conversation(conversation_id) - if conv is None: + # Only the runner binding is needed here, and it must be current — + # the event hot path resolves a runner per streamed chunk, so a + # rebind or host handoff that has already landed must route to the + # NEW runner. One indexed single-column read instead of the full + # conversation + metadata + labels load. Note this makes the + # RESOLUTION current; it does not fence the window between + # resolving and the caller's request, which would need a binding + # generation carried on the wire. + runner_ids = self._conversation_store.get_runner_ids([conversation_id]) + if conversation_id not in runner_ids: raise OmnigentError("conversation not found", code=ErrorCode.NOT_FOUND) - if conv.runner_id: - session = self._registry.get(conv.runner_id) + runner_id = runner_ids[conversation_id] + if runner_id: + session = self._registry.get(runner_id) if session is None: raise OmnigentError( - f"runner {conv.runner_id!r} is offline for conversation {conversation_id!r}", + f"runner {runner_id!r} is offline for conversation {conversation_id!r}", code=ErrorCode.RUNNER_UNAVAILABLE, ) return RoutedRunner( - runner_id=conv.runner_id, - client=self._client_for_runner(conv.runner_id), + runner_id=runner_id, + client=self._client_for_runner(runner_id), ) raise OmnigentError( @@ -167,18 +176,18 @@ def client_for_existing_conversation(self, conversation_id: str) -> RoutedRunner ``None`` when it is not pinned or not found. :raises OmnigentError: If the pinned runner is offline. """ - conv = self._conversation_store.get_conversation(conversation_id) - if conv is None or not conv.runner_id: + runner_id = self._conversation_store.get_runner_ids([conversation_id]).get(conversation_id) + if not runner_id: return None - session = self._registry.get(conv.runner_id) + session = self._registry.get(runner_id) if session is None: raise OmnigentError( - f"runner {conv.runner_id!r} is offline for conversation {conversation_id!r}", + f"runner {runner_id!r} is offline for conversation {conversation_id!r}", code=ErrorCode.RUNNER_UNAVAILABLE, ) return RoutedRunner( - runner_id=conv.runner_id, - client=self._client_for_runner(conv.runner_id), + runner_id=runner_id, + client=self._client_for_runner(runner_id), ) def runner_is_online(self, runner_id: str) -> bool: diff --git a/omnigent/server/routes/_sessions/helpers.py b/omnigent/server/routes/_sessions/helpers.py index 5bcad9e05c..371835e85c 100644 --- a/omnigent/server/routes/_sessions/helpers.py +++ b/omnigent/server/routes/_sessions/helpers.py @@ -1241,6 +1241,12 @@ def _ancestor_session_ids( """ Return ancestor session ids for a session, nearest parent first. + Reads each link fresh. The cost-usage fan-out walks a loaded tree + instead (:func:`ancestor_ids_from_tree`); this remains for callers that + hold no tree, and a caller-supplied starting row is deliberately not + accepted — a row read earlier in the request can name a parent chain the + session has since left, and these ids decide where events are published. + :param conv_store: Store used to read conversation parent links. :param session_id: Session to walk upward from, e.g. ``"conv_child123"``. @@ -3701,8 +3707,13 @@ async def _get_runner_client_impl( """ Get an HTTP client for the runner bound to a session. - Uses the ``RunnerRouter`` to resolve the pinned runner. Falls - back to the in-process runner client for test setups. + Uses the ``RunnerRouter`` to resolve the pinned runner via a fresh + single-column binding read, so the RESOLUTION reflects a rebind that + landed before this call. The returned client is bound to that runner: + a rebind between this resolution and the caller's POST is not fenced + (the resolve→forward window is unchanged by the read narrowing, and + fencing it needs a binding generation on the wire). Falls back to the + in-process runner client for test setups. :param session_id: Session/conversation identifier, e.g. ``"conv_abc123"``. diff --git a/omnigent/server/routes/_sessions/orchestration.py b/omnigent/server/routes/_sessions/orchestration.py index 60a46c9b22..61d756ad35 100644 --- a/omnigent/server/routes/_sessions/orchestration.py +++ b/omnigent/server/routes/_sessions/orchestration.py @@ -634,6 +634,11 @@ def _publish_subtree_cost_to_ancestors( behind that its remaining callers already depend on. :returns: None. """ + # Loaded HERE, not taken from the caller. What each ancestor's badge + # should read is a fact about now, at publication time — a tree captured + # when the request arrived can already describe a smaller total, and two + # siblings publishing from their own request-start snapshots delivered a + # newer figure followed by an older one, leaving the parent's badge stale. tree = load_session_tree( session_id, conv_store, @@ -1187,6 +1192,7 @@ def _persist_native_cumulative_usage( async def _persist_external_session_usage( session_id: str, + conv: Conversation, body: SessionEventInput, conversation_store: ConversationStore, ) -> int | None: @@ -1198,6 +1204,9 @@ async def _persist_external_session_usage( (:func:`_persist_native_cumulative_usage`) must be present. :param session_id: Session/conversation identifier. + :param conv: The already-loaded conversation row; only its immutable + ``root_conversation_id`` is read (so a request-start row is safe), + letting the subtree recompute skip one conversation read. :param body: External session-usage event body. :param conversation_store: Store used to upsert the labels. :returns: The persisted ``context_tokens`` when present, else ``None``. @@ -1260,7 +1269,16 @@ async def _persist_external_session_usage( # hide in-flight sub-agent spend until the next child flush (the badge would # oscillate own ⇄ subtree). For a childless session the subtree is just # itself, so this equals own cost — one indexed tree query per flush. - subtree_usage = await asyncio.to_thread(load_session_usage, session_id, conversation_store) + # This tree serves only this session's own subtree total. The ancestor + # re-publish below loads its OWN tree instead of reusing this one — see + # the comment down there for why sharing it would be wrong. + tree = await asyncio.to_thread( + load_session_tree, + session_id, + conversation_store, + conv.root_conversation_id, + ) + subtree_usage = _sum_subtree_usage(tree, session_id) subtree_cost = _priced_cost_for_display(subtree_usage) usage_by_model = _usage_by_model_for_display(subtree_usage) # Only include fields that were sent; the client treats absent @@ -1285,12 +1303,16 @@ async def _persist_external_session_usage( # This session's usage also moves its ANCESTORS' subtree cost (its spend # rolls up into every ancestor), so re-publish each ancestor's subtree cost # too — otherwise a grandparent's badge wouldn't reflect a deep descendant. - # No-op for a top-level session (no ancestors). Threaded: it pages the - # conversation tree per ancestor. + # Publishes nothing for a top-level session (no ancestors to walk), but + # still pays its own tree load to find that out — it reads its own tree + # rather than reusing the one summed above: what an ancestor's badge + # should show is a fact about publication time, and a sibling's event may + # have landed since this request started. await asyncio.to_thread( _publish_subtree_cost_to_ancestors, conversation_store, session_id, + conv, ) return raw_tokens diff --git a/omnigent/server/routes/sessions/routes_events.py b/omnigent/server/routes/sessions/routes_events.py index 2ae921c807..af5f7801f9 100644 --- a/omnigent/server/routes/sessions/routes_events.py +++ b/omnigent/server/routes/sessions/routes_events.py @@ -848,6 +848,7 @@ async def post_event( # post-hoc here — a logged output cannot be un-logged.) await _persist_external_session_usage( session_id, + conv, body, conversation_store, ) diff --git a/omnigent/stores/conversation_store/sqlalchemy_store.py b/omnigent/stores/conversation_store/sqlalchemy_store.py index 258500e7a0..73ddfe7543 100644 --- a/omnigent/stores/conversation_store/sqlalchemy_store.py +++ b/omnigent/stores/conversation_store/sqlalchemy_store.py @@ -1102,22 +1102,82 @@ def find_imported_conversation( def get_runner_ids(self, conversation_ids: list[str]) -> dict[str, str | None]: """ - Single ``SELECT id, runner_id WHERE id IN (...)`` — bulk - variant of :meth:`get_conversation` for the runner-dot path. - Missing ids are omitted; ids without a bound runner map to - ``None``. + Bulk ``runner_id`` lookup, gated on the conversation still existing. + + The binding lives on the Omnigent metadata row, but the AP + ``conversations`` row is the EXISTENCE authority: deletion is + deliberately ordered AP-row-first, metadata-second, and a failed + second transaction leaves orphaned metadata for a conversation that + no longer exists (documented as an acceptable best-effort tradeoff + in :meth:`delete_conversation`). That tradeoff is only acceptable + while every reachability check consults the AP row — a + metadata-only read would route a deleted conversation to its old + runner instead of reporting it missing. + + The lookup is therefore rooted in the AP row and reaches sideways for + the binding, never the other way round: one outer-joined query when + both live on the same bind, two when they are split (still narrower + than the three a full :meth:`get_conversation` costs). Rooting it in + metadata instead answered only one of the two failure directions — + a deleted conversation with orphaned metadata was correctly dropped, + while an existing conversation whose metadata row is absent was + *also* dropped, turning "exists but unbound" into a false + ``NOT_FOUND``. Both states are constructible: creation commits the AP + row and the metadata row in separate transactions, the same + two-transaction shape as deletion in the opposite order. + + :param conversation_ids: Conversation ids to look up. + :returns: ``{id: runner_id or None}`` for ids whose AP row exists. + Ids with no AP row are omitted, exactly as a missing + :meth:`get_conversation` would report them; ids that exist + without a binding map to ``None``. """ if not conversation_ids: return {} unique_ids = list(set(conversation_ids)) + same_bind = self._conv_engine is self._engine + if same_bind: + with self._session() as session: + rows = session.execute( + select(SqlConversation.id, SqlConversationMetadata.runner_id) + .outerjoin( + SqlConversationMetadata, + and_( + SqlConversationMetadata.workspace_id == SqlConversation.workspace_id, + SqlConversationMetadata.id == SqlConversation.id, + ), + ) + .where( + SqlConversation.workspace_id == current_workspace_id(), + SqlConversation.id.in_(unique_ids), + ) + ).all() + return {row.id: row.runner_id for row in rows} + # Split-DB: no cross-bind join, so read existence from the AP bind and + # overlay the bindings. Seeding every existing id with ``None`` keeps + # the two states distinct without a second existence check. + with self._conv_session() as ap_session: + existing = set( + ap_session.execute( + select(SqlConversation.id).where( + SqlConversation.workspace_id == current_workspace_id(), + SqlConversation.id.in_(unique_ids), + ) + ).scalars() + ) + if not existing: + return {} + bindings: dict[str, str | None] = dict.fromkeys(existing) with self._session() as session: rows = session.execute( select(SqlConversationMetadata.id, SqlConversationMetadata.runner_id).where( SqlConversationMetadata.workspace_id == current_workspace_id(), - SqlConversationMetadata.id.in_(unique_ids), + SqlConversationMetadata.id.in_(list(existing)), ) ).all() - return {row.id: row.runner_id for row in rows} + for row in rows: + bindings[row.id] = row.runner_id + return bindings def get_session_connectivity( self, conversation_ids: list[str] diff --git a/tests/runner/test_routing.py b/tests/runner/test_routing.py index ba8561d74a..b6be3b3cd8 100644 --- a/tests/runner/test_routing.py +++ b/tests/runner/test_routing.py @@ -44,6 +44,14 @@ def __init__(self, conversations: dict[str, Conversation]) -> None: """ self._conversations = conversations + def get_runner_ids(self, conversation_ids: list[str]) -> dict[str, str | None]: + """Bulk binding read mirroring the real store's contract.""" + return { + cid: self._conversations[cid].runner_id + for cid in conversation_ids + if cid in self._conversations + } + def get_conversation(self, conversation_id: str) -> Conversation | None: """ Return a conversation by id. @@ -249,3 +257,72 @@ async def test_runner_router_existing_conversation_returns_none_when_unpinned() assert router.client_for_existing_conversation("conv_test") is None finally: await router.aclose() + + +@pytest.mark.parametrize( + ("method", "missing", "unpinned"), + [ + ("client_for_session_resources", ErrorCode.NOT_FOUND, ErrorCode.CONFLICT), + # The optional variant reports both as "no runner to route to". + ("client_for_existing_conversation", None, None), + ], +) +@pytest.mark.asyncio +async def test_narrowed_binding_routing_is_fresh_and_keeps_its_semantics( + method: str, + missing: ErrorCode | None, + unpinned: ErrorCode | None, +) -> None: + """ + Every per-event route that narrowed its conversation read to the binding + must stay fresh, stay narrow, and keep its own missing/unpinned answers. + + Parametrized over the methods rather than written per method: the first + of the two was covered and the second was not, so reverting it to a + full-row read left the suite green. + + Narrow: the full-row read explodes, so any regression to + ``get_conversation`` fails here. Fresh: a rebind landing between + resolutions is picked up by the next one. Semantics: a missing + conversation and an existing-but-unbound one stay distinguishable, which + is what the store's two-state binding lookup exists to preserve. + """ + registry = TunnelRegistry() + registry.register("runner_new", _FakeWebSocket(), _hello(harnesses=["codex"])) + + class _BindingOnlyStore: + def __init__(self) -> None: + self.binding: dict[str, str | None] = { + "conv_test": "runner_old", + "conv_unpinned": None, + } + + def get_runner_ids(self, conversation_ids: list[str]) -> dict[str, str | None]: + return {c: self.binding[c] for c in conversation_ids if c in self.binding} + + def get_conversation(self, conversation_id: str) -> None: + raise AssertionError("routing must not load the full conversation row") + + store = _BindingOnlyStore() + router = RunnerRouter(registry=registry, conversation_store=store) # type: ignore[arg-type] + resolve = getattr(router, method) + try: + # Bound to an offline runner: both methods report it the same way. + with pytest.raises(OmnigentError) as excinfo: + resolve("conv_test") + _assert_omnigent_error(excinfo, code=ErrorCode.RUNNER_UNAVAILABLE) + + # A concurrent rebind must be visible to the very next resolution. + store.binding["conv_test"] = "runner_new" + routed = resolve("conv_test") + assert routed is not None and routed.runner_id == "runner_new" + + for conv_id, expected in (("conv_missing", missing), ("conv_unpinned", unpinned)): + if expected is None: + assert resolve(conv_id) is None, conv_id + else: + with pytest.raises(OmnigentError) as excinfo: + resolve(conv_id) + _assert_omnigent_error(excinfo, code=expected) + finally: + await router.aclose() diff --git a/tests/server/routes/test_event_path_load_once.py b/tests/server/routes/test_event_path_load_once.py new file mode 100644 index 0000000000..752879d145 --- /dev/null +++ b/tests/server/routes/test_event_path_load_once.py @@ -0,0 +1,328 @@ +"""Regression tests: the event hot path loads the conversation once. + +POST /v1/sessions/{id}/events fires per streamed chunk during active +turns (~98% of server query volume), so redundant conversation/ACL +reloads multiply hard. These pin two things at once, because a +query-count ceiling on its own would stay green if the work it is +counting disappeared entirely: + +1. the per-table query counts for an event, and +2. that the event's real effect still happens — the status forward + reaches the bound runner, the usage recompute publishes the right + subtree total. + +The routing read is deliberately NOT threaded from the handler's row: +runner resolution takes its own fresh single-column binding read so a +concurrent rebind is picked up. The rebind semantics themselves are +pinned in ``tests/runner/test_routing.py``, because the app builds its +own ``RunnerRouter`` over a real tunnel registry with no injection seam — +what IS testable here, and asserted below, is that every event issues its +own binding read rather than reusing a cached resolution. +""" + +from __future__ import annotations + +import re +from collections import Counter +from contextlib import contextmanager +from typing import Any + +import httpx +import pytest +from sqlalchemy import event + +from omnigent.db.utils import _engine_cache +from omnigent.runtime import session_stream, set_runner_client, set_runner_router +from omnigent.stores.conversation_store.sqlalchemy_store import ( + SqlAlchemyConversationStore, +) + +# Attribute every statement to its table so the counter cannot miss the +# metadata / labels / UPDATE traffic a conversations-only filter ignored. +_TABLE_RE = re.compile( + r"\bFROM\s+([a-z_][a-z0-9_]*)|\bINSERT\s+INTO\s+([a-z_][a-z0-9_]*)|\bUPDATE\s+([a-z_][a-z0-9_]*)", + re.IGNORECASE, +) + + +class _QueryCounts(Counter): + """Per-table statement counts, retaining the statements themselves.""" + + statements: list[str] + + +@contextmanager +def _count_queries_by_table(engine: Any): + """Count every non-PRAGMA statement, keyed by the table it touches.""" + counts: _QueryCounts = _QueryCounts() + counts.statements = [] + + def _on_exec(conn, cursor, statement, parameters, context, executemany): + if statement.lstrip().upper().startswith("PRAGMA"): + return + counts.statements.append(statement) + match = _TABLE_RE.search(statement) + table = next((g for g in match.groups() if g), "other") if match else "other" + counts[table.lower()] += 1 + + event.listen(engine, "before_cursor_execute", _on_exec) + try: + yield counts + finally: + event.remove(engine, "before_cursor_execute", _on_exec) + + +@contextmanager +def _record_published(session_id: str): + """Capture events published to *session_id*'s stream.""" + captured: list[dict[str, Any]] = [] + real_publish = session_stream.publish + + def _spy(conv_id: str, payload: dict[str, Any]) -> None: + if conv_id == session_id: + captured.append(payload) + real_publish(conv_id, payload) + + session_stream.publish = _spy # type: ignore[assignment] + try: + yield captured + finally: + session_stream.publish = real_publish # type: ignore[assignment] + + +class _RecordingRunnerClient: + """Minimal runner client that records the events posted to it.""" + + def __init__(self) -> None: + self.posts: list[tuple[str, Any]] = [] + + async def post( + self, + url: str, + *, + json: Any = None, + timeout: float | None = None, + ) -> httpx.Response: + del timeout + self.posts.append((url, json)) + return httpx.Response(200, json={}, request=httpx.Request("POST", url)) + + +class _ExplodingRunnerClient: + """Process-wide fallback client that must never be used. + + Delivery falling back to this client is what let the routed-binding + assertion pass for four review rounds. + """ + + async def post(self, url: str, **kwargs: Any) -> httpx.Response: + raise AssertionError( + f"event was delivered through the fallback client, not the bound runner: POST {url}" + ) + + +@pytest.fixture() +def runner_router_reset(): + """Install/remove runner globals around a test.""" + yield + set_runner_router(None) + set_runner_client(None) + + +@pytest.mark.asyncio +async def test_status_event_forwards_and_reads_each_table_once( + app: Any, + client: httpx.AsyncClient, + db_uri: str, + monkeypatch: pytest.MonkeyPatch, + runner_router_reset: None, +) -> None: + """ + A status event reaches the runner it is BOUND to, and reads each + conversation table once — not once per routing lookup on top of the + handler's own read. + + The event is delivered through a real :class:`RunnerRouter` over a real + registry, so the binding read under test actually selects the client + that receives the POST. Earlier versions of this test bound + ``runner_one`` without registering it: routing reported it offline, + delivery silently fell back to the process-wide client, and the + assertion passed while observing nothing about routing. The fallback + client is installed here as a trap that fails if it is ever used. + """ + from tests.server.helpers import register_test_runner + + store = SqlAlchemyConversationStore(db_uri) + conv = store.create_conversation(title="event-hot-path") + store.set_runner_id(conv.id, "runner_one") + + # Bind AND register: the app's own router resolves the binding through + # the read under test and then looks the runner up in the app's + # registry. Binding without registering is what made routing report the + # runner offline and delivery fall through to the global client. + register_test_runner(app, "runner_one") + runner = _RecordingRunnerClient() + routed_ids: list[str] = [] + + def _client_for_runner(runner_id: str) -> Any: + routed_ids.append(runner_id) + return runner + + # Real routing, stubbed transport: everything up to and including the + # binding lookup is production code; only the tunnel socket is replaced. + monkeypatch.setattr(app.state.runner_router, "_client_for_runner", _client_for_runner) + set_runner_client(_ExplodingRunnerClient()) # type: ignore[arg-type] + + payload = {"type": "external_session_status", "data": {"status": "idle"}} + await client.post(f"/v1/sessions/{conv.id}/events", json=payload) # warm + + engine = _engine_cache[db_uri] + baseline = len(runner.posts) + with _count_queries_by_table(engine) as counts: + resp = await client.post(f"/v1/sessions/{conv.id}/events", json=payload) + + assert resp.status_code == 202, resp.text + # Delivered to the BOUND runner, not merely delivered: the routed client + # is the one that received it, and the global fallback would have raised. + assert routed_ids and set(routed_ids) == {"runner_one"}, routed_ids + delivered = runner.posts[baseline:] + assert [body["type"] for _url, body in delivered] == ["external_session_status"], delivered + + # The handler's single conversation load (conversations + metadata + + # labels) plus routing's one narrow binding read — and nothing else. + # Before this change routing re-loaded the whole row, so metadata, + # conversations and labels were each read twice. + assert counts["conversation_labels"] == 1, dict(counts) + full_row_reads = [ + st + for st in counts.statements + if st.lstrip().upper().startswith("SELECT") + and "FROM conversations" in st + and "conversations.title" in st + ] + assert len(full_row_reads) == 1, full_row_reads + # Routing's read is the NARROW binding lookup: it projects the id and the + # runner binding only, never the row's other columns. + binding_reads = [ + st for st in counts.statements if "OUTER JOIN" in st.upper() and "runner_id" in st + ] + assert len(binding_reads) == 1, binding_reads + assert "conversations.title" not in binding_reads[0], binding_reads + + +@pytest.mark.asyncio +async def test_usage_event_publishes_subtree_total_with_exact_reads( + client: httpx.AsyncClient, + db_uri: str, +) -> None: + """ + The usage event recomputes the subtree total from the handler's row + (no root re-resolution) and publishes the correct value — a count + ceiling alone would pass if the recompute were deleted. + """ + store = SqlAlchemyConversationStore(db_uri) + conv = store.create_conversation(title="usage-hot-path") + child = store.create_conversation( + kind="sub_agent", parent_conversation_id=conv.id, title="usage:child" + ) + store.set_session_usage(child.id, {"total_cost_usd": 0.25}) + + payload = {"type": "external_session_usage", "data": {"cumulative_cost_usd": 0.5}} + await client.post(f"/v1/sessions/{conv.id}/events", json=payload) # warm + + engine = _engine_cache[db_uri] + with _record_published(conv.id) as published: + with _count_queries_by_table(engine) as counts: + resp = await client.post(f"/v1/sessions/{conv.id}/events", json=payload) + + assert resp.status_code == 202, resp.text + usage_events = [e for e in published if e.get("type") == "session.usage"] + assert usage_events, f"no session.usage published: {published}" + # Parent's own 0.5 plus the child's 0.25 — proves the subtree recompute + # ran and used the whole tree, not just this row. + assert usage_events[-1]["total_cost_usd"] == pytest.approx(0.75) + + # Per-table totals, not a ceiling on one table: a ceiling hid extra + # metadata and label statements entirely, and "<= 3" also passed with the + # anti-replay read removed. conversations = handler row, the anti-replay + # merge read (which must stay fresh) and the tree page scan; no root + # re-resolution in any of them. + assert dict(counts) == { + "conversations": 4, + "omnigent_conversation_metadata": 5, + "conversation_labels": 4, + }, dict(counts) + # The fourth conversations read is the ancestor fan-out loading its own + # tree at publication time. Reusing the tree summed above saved it and was + # wrong: two siblings then published from their own request-start + # snapshots, delivering a newer total followed by an older one. + + +@pytest.mark.asyncio +async def test_ancestor_publish_reads_the_tree_after_a_sibling_lands( + client: httpx.AsyncClient, + db_uri: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """ + A parent's badge must not go stale because it was computed from a tree + read at the START of the request that happened to publish it. + + Each usage event re-publishes every ancestor's subtree total. Taking that + total from the tree loaded at the start of the current request made the + figure a fact about request-arrival time, so a sibling's spend landing + mid-request was published as if it had not happened — two siblings + reporting around each other delivered a newer total and then an older one, + and the parent was left showing the smaller number. + + This pins that specific regression, not a universal "never moves + backwards" guarantee: a sibling's spend landing AFTER the ancestor + fan-out's own tree read but before publish can still be missed, exactly + as on unmodified `main` — that window is pre-existing and unrelated to + the request-start snapshot this PR removed. + + The interleave is constructed rather than hoped for: a sibling's spend is + committed immediately after this request's first tree load, which is + exactly the window. The fan-out reads its own tree afterwards, so the + published total includes it. + """ + from omnigent.server.routes._sessions import orchestration as orch + + store = SqlAlchemyConversationStore(db_uri) + parent = store.create_conversation(title="fanout-parent") + reporter = store.create_conversation( + kind="sub_agent", parent_conversation_id=parent.id, title="usage:reporter" + ) + sibling = store.create_conversation( + kind="sub_agent", parent_conversation_id=parent.id, title="usage:sibling" + ) + + payload = {"type": "external_session_usage", "data": {"cumulative_cost_usd": 1.0}} + await client.post(f"/v1/sessions/{reporter.id}/events", json=payload) # warm + + real_load = orch.load_session_tree + landed = {"done": False} + + def _load_then_land_sibling(*args: Any, **kwargs: Any): + tree = real_load(*args, **kwargs) + if not landed["done"]: + landed["done"] = True + # The sibling's spend commits here: after this request read a tree, + # before anything is published from it. + store.set_session_usage(sibling.id, {"total_cost_usd": 2.0}) + return tree + + monkeypatch.setattr(orch, "load_session_tree", _load_then_land_sibling) + + with _record_published(parent.id) as parent_events: + resp = await client.post(f"/v1/sessions/{reporter.id}/events", json=payload) + assert resp.status_code == 202, resp.text + assert landed["done"], "the interleave never happened; the test proves nothing" + + totals = [e["total_cost_usd"] for e in parent_events if e.get("type") == "session.usage"] + assert totals, f"the parent received no usage event: {parent_events}" + # reporter 1.0 + sibling 2.0, read after the sibling landed. + assert totals[-1] == pytest.approx(3.0), ( + f"the parent's published total {totals} was computed from a tree read " + f"before the sibling's spend landed" + ) diff --git a/tests/stores/test_conversation_store.py b/tests/stores/test_conversation_store.py index cf78a3244c..187f929f44 100644 --- a/tests/stores/test_conversation_store.py +++ b/tests/stores/test_conversation_store.py @@ -5360,3 +5360,53 @@ def _mutate(state: dict) -> None: f"one writer's increment was lost: {persisted} — the two transactions " f"were not serialised on the metadata row" ) + + +def test_get_runner_ids_distinguishes_missing_from_unbound( + conversation_store: SqlAlchemyConversationStore, +) -> None: + """ + The two half-written states are distinct, and both must report as such. + + The AP row is the existence authority and the binding lives on the + Omnigent metadata row, written in a separate transaction — so either row + can outlive the other: + + - AP row gone, metadata left behind (the documented deletion tradeoff): + the conversation no longer exists, so it must be OMITTED. Reporting + its old runner routes events to a session that is gone instead of + 404-ing. + - AP row present, metadata absent (the same shape during creation, in + the opposite order): the conversation exists but is unbound, so it + must map to ``None`` — which the route turns into CONFLICT. Omitting + it produces a false NOT_FOUND. + + A metadata-rooted lookup answers the first correctly and the second + wrongly, which is why both directions are asserted from one table. + """ + from sqlalchemy import delete as sa_delete + + from omnigent.db.db_models import SqlConversation as _Conv + from omnigent.db.db_models import SqlConversationMetadata as _Meta + + bound = conversation_store.create_conversation(title="bound") + conversation_store.set_runner_id(bound.id, "runner_live") + orphaned_metadata = conversation_store.create_conversation(title="ap-row-gone") + conversation_store.set_runner_id(orphaned_metadata.id, "runner_old") + missing_metadata = conversation_store.create_conversation(title="metadata-gone") + + with conversation_store._conv_session() as session: + session.execute(sa_delete(_Conv).where(_Conv.id == orphaned_metadata.id)) + with conversation_store._session() as session: + session.execute(sa_delete(_Meta).where(_Meta.id == missing_metadata.id)) + + ids = [bound.id, orphaned_metadata.id, missing_metadata.id] + assert conversation_store.get_runner_ids(ids) == { + bound.id: "runner_live", + # present but unbound -> None, never omitted + missing_metadata.id: None, + # AP row gone -> omitted entirely, never its stale runner + } + # The full read agrees about existence in both directions. + assert conversation_store.get_conversation(orphaned_metadata.id) is None + assert conversation_store.get_conversation(missing_metadata.id) is not None diff --git a/tests/stores/test_conversation_store_split_db.py b/tests/stores/test_conversation_store_split_db.py index fb56e459fa..0086b9ca9a 100644 --- a/tests/stores/test_conversation_store_split_db.py +++ b/tests/stores/test_conversation_store_split_db.py @@ -596,3 +596,36 @@ def test_delete_conversation_keeps_template_agent( asyncio.run(store.delete_conversation(conv.id)) assert _col(omnigent_db, "agents", "id") == ["191cbf904e3223e9e00ac9a1abfe79a5"] + + +def test_get_runner_ids_distinguishes_missing_from_unbound_split_db( + store: SqlAlchemyConversationStore, +) -> None: + """Split-DB must draw the same two-state distinction as one bind. + + With the tables on separate binds there is no join to root the lookup + in, so existence comes from the AP bind and bindings are overlaid onto + it. Skipping that seeding drops an existing-but-unbound conversation and + reports a false NOT_FOUND — the same defect as a metadata-rooted lookup, + reached a different way. + """ + from sqlalchemy import delete as sa_delete + + from omnigent.db.db_models import SqlConversation as _Conv + from omnigent.db.db_models import SqlConversationMetadata as _Meta + + bound = store.create_conversation(title="bound") + store.set_runner_id(bound.id, "runner_live") + orphaned_metadata = store.create_conversation(title="ap-row-gone") + store.set_runner_id(orphaned_metadata.id, "runner_old") + missing_metadata = store.create_conversation(title="metadata-gone") + + with store._conv_session() as session: + session.execute(sa_delete(_Conv).where(_Conv.id == orphaned_metadata.id)) + with store._session() as session: + session.execute(sa_delete(_Meta).where(_Meta.id == missing_metadata.id)) + + assert store.get_runner_ids([bound.id, orphaned_metadata.id, missing_metadata.id]) == { + bound.id: "runner_live", + missing_metadata.id: None, + }