fix(store): make label seeding and session-state writes atomic

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(

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

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)

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.
         """
         ...