From 0f322b89de7a2784b8936bf2cb4a6dcea112f40f Mon Sep 17 00:00:00 2001 From: Andrew Reid Date: Sun, 26 Jul 2026 12:17:40 +0930 Subject: [PATCH] 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 | 253 ++++++++++++++++-- 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, 828 insertions(+), 55 deletions(-) diff --git a/omnigent/runtime/policies/builder.py b/omnigent/runtime/policies/builder.py index 12320e7d6f..ce9b749dc1 100644 --- a/omnigent/runtime/policies/builder.py +++ b/omnigent/runtime/policies/builder.py @@ -704,17 +704,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 da6c280f87..c3112b5634 100644 --- a/omnigent/server/routes/_sessions/orchestration.py +++ b/omnigent/server/routes/_sessions/orchestration.py @@ -1237,7 +1237,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 d03e79cac4..40981231e9 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. """ @@ -841,6 +851,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, @@ -861,9 +911,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, @@ -945,11 +999,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"``. @@ -958,6 +1016,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, @@ -1035,6 +1124,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 e62f94b8a8..e200a43192 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, Protocol, cast from sqlalchemy import ( @@ -20,6 +21,7 @@ text, update, ) +from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import QueryableAttribute, Session, aliased, load_only from sqlalchemy.sql.selectable import Subquery @@ -469,6 +471,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, @@ -1279,6 +1341,62 @@ def set_labels( with self._conv_session("set_labels") 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("seed_labels_if_absent") 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, @@ -1306,6 +1424,109 @@ 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. + """ + with self._session_immediate("mutate_session_state") as session: + return self._mutate_metadata_json_in_session( + session, conversation_id, "session_state", mutate, what="session state" + ) + + def _mutate_metadata_json_in_session( + self, + session: Session, + 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. + + Callers open the session with a literal query name and pass it here, + keeping the ``_session_immediate`` call (which the query-naming lint + requires to carry a string literal) at the public call sites rather + than inside this helper. + + 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 session: An already-open SQLAlchemy session. + :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 + + 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, @@ -1393,32 +1614,22 @@ 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("increment_session_usage") 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 self._mutate_metadata_json_in_session( + session, + conversation_id, + "session_usage", + lambda current: apply_session_usage_delta(current, delta), + what="session usage", ) - return current 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..f4e1cb7225 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("drop_metadata_row") 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 ed2bc484e5..8708ed586a 100644 --- a/tests/server/integration/test_sessions_endpoints.py +++ b/tests/server/integration/test_sessions_endpoints.py @@ -5116,6 +5116,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 dc867d9399..854b43a3b5 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 ( @@ -5304,3 +5304,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("check_orphan_labels") 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" + )