fix(cost): use BEGIN IMMEDIATE for SQLite in increment_session_usage

The previous implementation used a deferred session (self._session) for
all dialects, then added SELECT FOR UPDATE only for non-SQLite. On SQLite
a deferred SELECT-then-UPDATE races: two concurrent writers each take a
read snapshot, and the second writer's UPDATE upgrade hits
SQLITE_BUSY_SNAPSHOT — the delta is dropped and an exception propagates.

Fix: open increment_session_usage with self._session_immediate (a new
make_managed_session_maker(engine, immediate=True) session maker added in
__init__). On SQLite this issues BEGIN IMMEDIATE, acquiring the write lock
before the first read so concurrent writers are serialised at lock
acquisition time rather than failing mid-transaction. On PostgreSQL/MySQL
immediate=True is a no-op — SELECT FOR UPDATE (via _supports_for_update)
continues to handle those dialects.

omnigent/server/routes/sessions.py

@@ -2881,10 +2881,13 @@ def _accumulate_session_usage(
     Increment the session's cumulative token counters from a
     ``response.completed`` event's usage data.
 
-    Called synchronously from the relay loop. Reads the current
-    persisted ``session_usage``, adds the delta from the
-    response's ``usage`` field, and writes the updated totals
-    back. No-op when the response carries no usage data.
+    Called synchronously from the relay loop. Builds a usage delta from
+    the response's ``usage`` field and atomically applies it to the
+    persisted ``session_usage`` via a single database transaction
+    (``SELECT FOR UPDATE`` on PostgreSQL, SQLite's single-writer lock
+    otherwise). This prevents the read-modify-write race that caused
+    concurrent relay completions to silently drop each other's cost /
+    token deltas (#9). No-op when the response carries no usage data.
 
     Cost is computed when the model's per-token pricing is
     available from the MLflow catalog (looked up once per call
@@ -2920,20 +2923,10 @@ def _accumulate_session_usage(
     cache_read_input_tokens = usage_obj.get("cache_read_input_tokens", 0)
     cache_creation_input_tokens = usage_obj.get("cache_creation_input_tokens", 0)
 
-    # Load current cumulative usage from the store.
+    # Load conversation metadata for pricing only (NOT for reading session_usage —
+    # the atomic increment_session_usage call below handles that separately to
+    # avoid the read-modify-write race).
     conv = conversation_store.get_conversation(session_id)
-    current = dict(conv.session_usage) if conv else {}
-    current.setdefault("input_tokens", 0)
-    current.setdefault("output_tokens", 0)
-    current.setdefault("total_tokens", 0)
-    current.setdefault("cache_read_input_tokens", 0)
-    current.setdefault("cache_creation_input_tokens", 0)
-
-    current["input_tokens"] += input_tokens
-    current["output_tokens"] += output_tokens
-    current["total_tokens"] += total_tokens
-    current["cache_read_input_tokens"] += cache_read_input_tokens
-    current["cache_creation_input_tokens"] += cache_creation_input_tokens
 
     # Compute cost delta if pricing is available for the model. Resolve
     # the model to price with, most-specific first:
@@ -2948,6 +2941,7 @@ def _accumulate_session_usage(
     # created only on this priced branch, so an unpriced session never
     # gains a (misleading $0.00) cost key.
     cost_delta = 0.0
+    priced = False
     # Prefer an authoritative harness-reported cost over the catalog estimate.
     provider_cost = usage_obj.get("cost_usd")
     has_provider_cost = isinstance(provider_cost, (int, float))
@@ -2971,29 +2965,42 @@ def _accumulate_session_usage(
                 # token counts when the harness reports them; compute_llm_cost
                 # prices them at their own (cheaper read / pricier write) rates.
                 cost_delta = compute_llm_cost(usage_obj, pricing)
-        if priced:
-            current["total_cost_usd"] = current.get("total_cost_usd", 0.0) + cost_delta
-        # Per-model attribution (ADD). Tokens are attributed whenever the
-        # model is known — including unpriced turns — so the per-model token
-        # view is complete; cost is attributed only when this model's turn
-        # was priced (passing ``None`` otherwise keeps the model's cost key
-        # absent, matching the flat "priced ⟺ key present" contract).
-        _add_model_usage_delta(
-            _model_usage_bucket(current, llm_model),
-            {
-                "input_tokens": input_tokens,
-                "output_tokens": output_tokens,
-                "total_tokens": total_tokens,
-                "cache_read_input_tokens": cache_read_input_tokens,
-                "cache_creation_input_tokens": cache_creation_input_tokens,
-            },
-            cost_delta if priced else None,
-        )
 
-    conversation_store.set_session_usage(session_id, current)
+    # Build the delta dict and atomically apply it to the persisted
+    # session_usage in a single DB transaction (SELECT FOR UPDATE on
+    # PostgreSQL; SQLite's exclusive write lock on SQLite). This is the fix
+    # for the read-modify-write race that caused concurrent completions to
+    # overwrite each other's deltas (#9).
+    delta: dict[str, Any] = {
+        "input_tokens": input_tokens,
+        "output_tokens": output_tokens,
+        "total_tokens": total_tokens,
+        "cache_read_input_tokens": cache_read_input_tokens,
+        "cache_creation_input_tokens": cache_creation_input_tokens,
+    }
+    if priced:
+        delta["total_cost_usd"] = cost_delta
+    if llm_model:
+        # Per-model attribution. Tokens are attributed whenever the model is
+        # known — including unpriced turns — so the per-model token view is
+        # complete; cost is attributed only when this model's turn was priced
+        # (keeping the model's cost key absent otherwise, matching the flat
+        # "priced ⟺ key present" contract).
+        model_delta: dict[str, Any] = {
+            "input_tokens": input_tokens,
+            "output_tokens": output_tokens,
+            "total_tokens": total_tokens,
+            "cache_read_input_tokens": cache_read_input_tokens,
+            "cache_creation_input_tokens": cache_creation_input_tokens,
+        }
+        if priced:
+            model_delta["total_cost_usd"] = cost_delta
+        delta["by_model"] = {llm_model: model_delta}
+
+    new_current = conversation_store.increment_session_usage(session_id, delta)
     # 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(current)
+    return _priced_cost_for_display(new_current)
 
 
 def _persist_native_cumulative_usage(

omnigent/stores/conversation_store/__init__.py

@@ -178,6 +178,29 @@ class NameAlreadyExistsError(Exception):
     """
 
 
+def apply_session_usage_delta(current: dict[str, Any], delta: dict[str, Any]) -> None:
+    """
+    Apply a usage *delta* to *current* in place (add semantics, nested-aware).
+
+    Flat numeric keys are summed; ``"by_model"`` sub-dicts are merged by
+    model id, summing each model's sub-keys independently. Used by
+    :meth:`ConversationStore.increment_session_usage` implementations to
+    keep the merge logic in one place.
+
+    :param current: Existing ``session_usage`` dict (mutated in place).
+    :param delta: Increments to apply (same layout as ``session_usage``).
+    """
+    for key, value in delta.items():
+        if key == "by_model":
+            by_model = current.setdefault("by_model", {})
+            for model_id, model_delta in value.items():
+                bucket = by_model.setdefault(model_id, {})
+                for sub_key, sub_value in model_delta.items():
+                    bucket[sub_key] = bucket.get(sub_key, 0) + sub_value
+        else:
+            current[key] = current.get(key, 0) + value
+
+
 class ConversationStore(ABC):
     """
     Abstract base for conversation persistence.
@@ -773,6 +796,43 @@ class ConversationStore(ABC):
         """
         ...
 
+    @abstractmethod
+    def increment_session_usage(
+        self,
+        conversation_id: str,
+        delta: dict[str, Any],
+    ) -> dict[str, Any]:
+        """
+        Atomically apply a usage delta to a conversation's ``session_usage``.
+
+        Reads the current JSON, applies *delta* (adding each key's value to the
+        existing value, with ``by_model`` merged recursively), and writes back —
+        all within a single database transaction. Concurrent writers are
+        serialised via dialect-appropriate locking: ``SELECT FOR UPDATE`` on
+        PostgreSQL / MySQL / MariaDB; ``BEGIN IMMEDIATE`` (write lock before
+        the first read) on SQLite, which avoids ``SQLITE_BUSY_SNAPSHOT`` that
+        a plain deferred ``SELECT``-then-``UPDATE`` would raise under concurrent
+        writers. This prevents the read-modify-write
+        race that caused concurrent relay completions to silently drop each
+        other's cost / token deltas (#9).
+
+        *delta* uses the same key layout as ``session_usage``:
+        - flat numeric keys (``"input_tokens"``, ``"total_cost_usd"``, …) are
+          added to the existing value (``0`` when absent).
+        - ``"by_model"`` is a nested dict ``{model_id: {sub_key: value}}``; each
+          model's sub-keys are added independently, creating the bucket on first
+          use.
+
+        :param conversation_id: The conversation to update,
+            e.g. ``"conv_abc123"``.
+        :param delta: Usage increments to apply, e.g.
+            ``{"input_tokens": 1000, "total_cost_usd": 0.05,
+            "by_model": {"claude-sonnet-4-6": {"input_tokens": 1000,
+            "total_cost_usd": 0.05}}}``.
+        :returns: The updated ``session_usage`` dict after the increment.
+        """
+        ...
+
     @abstractmethod
     def add_daily_cost(self, user_id: str, day_utc: str, delta_usd: float) -> None:
         """

omnigent/stores/conversation_store/sqlalchemy_store.py

@@ -477,6 +477,12 @@ class SqlAlchemyConversationStore(ConversationStore):
         super().__init__(storage_location)
         self._engine = get_or_create_engine(storage_location)
         self._session = make_managed_session_maker(self._engine)
+        # Immediate session: used for read-modify-write operations that must be
+        # atomic. On SQLite, ``BEGIN IMMEDIATE`` acquires the write lock before
+        # the first read, preventing ``SQLITE_BUSY_SNAPSHOT`` under concurrent
+        # writers. On other dialects ``immediate=True`` is a no-op — those paths
+        # use ``SELECT … FOR UPDATE`` via ``_supports_for_update`` instead.
+        self._session_immediate = make_managed_session_maker(self._engine, immediate=True)
         self._supports_for_update = self._engine.dialect.name != "sqlite"
         # SQLite rowid is monotonically increasing absent deletions; it serves
         # as an insertion-ordered tiebreaker for timestamp ties. Note: without
@@ -932,6 +938,53 @@ class SqlAlchemyConversationStore(ConversationStore):
                 .values(session_usage=json.dumps(usage))
             )
 
+    def increment_session_usage(
+        self,
+        conversation_id: str,
+        delta: dict[str, Any],
+    ) -> dict[str, Any]:
+        """
+        Atomically increment the session usage for one conversation.
+
+        Runs the read-modify-write in a single database transaction, serialising
+        concurrent writers via two complementary mechanisms:
+
+        - **PostgreSQL / MySQL / MariaDB**: ``SELECT … FOR UPDATE`` acquires an
+          exclusive row lock for the duration of the transaction; a concurrent
+          second writer blocks until this one commits.
+        - **SQLite**: the session is opened with ``BEGIN IMMEDIATE``
+          (``self._session_immediate``), which acquires SQLite's write lock
+          *before* the first read. A plain ``SELECT``-then-``UPDATE`` in a
+          deferred transaction would expose concurrent writers to
+          ``SQLITE_BUSY_SNAPSHOT`` because each writer takes a read snapshot
+          first; ``BEGIN IMMEDIATE`` prevents that by serialising at lock
+          acquisition time.
+
+        :param conversation_id: The conversation to update.
+        :param delta: Usage increments (see
+            :meth:`ConversationStore.increment_session_usage`).
+        :returns: The updated ``session_usage`` dict.
+        """
+        import json
+
+        from omnigent.stores.conversation_store import apply_session_usage_delta
+
+        with self._session_immediate() as session:
+            q = select(SqlConversation).where(SqlConversation.id == conversation_id)
+            if self._supports_for_update:
+                q = q.with_for_update()
+            row = session.scalars(q).first()
+            current: dict[str, Any] = (
+                dict(json.loads(row.session_usage)) if row and row.session_usage else {}
+            )
+            apply_session_usage_delta(current, delta)
+            session.execute(
+                update(SqlConversation)
+                .where(SqlConversation.id == conversation_id)
+                .values(session_usage=json.dumps(current))
+            )
+            return current
+
     def add_daily_cost(self, user_id: str, day_utc: str, delta_usd: float) -> None:
         """
         Atomically add *delta_usd* to a user's spend for one UTC day.

tests/server/integration/test_sessions_endpoints.py

@@ -4407,6 +4407,68 @@ 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_concurrent_calls_accumulate_both_deltas(
+    client: httpx.AsyncClient,
+    db_uri: str,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Concurrent _accumulate_session_usage calls each persist their full delta.
+
+    Regression guard for the read-modify-write lost-update bug (#9): the old
+    implementation read session_usage in one transaction and wrote back in
+    another, so two concurrent completions could each read the same stale total,
+    compute their own delta, and overwrite the other's result — silently dropping
+    a delta. The fix (increment_session_usage — single atomic transaction with
+    SELECT FOR UPDATE on PostgreSQL/MySQL) serialises concurrent writers so both
+    deltas are always preserved.
+
+    The calls are dispatched from two threads simultaneously via
+    concurrent.futures.ThreadPoolExecutor so the race window genuinely exists:
+    the old non-atomic implementation would fail this test non-deterministically
+    (or always, if the two reads are forced to coincide).
+    """
+    import concurrent.futures
+
+    from omnigent.server.routes import sessions as sessions_routes
+
+    monkeypatch.setattr(
+        "omnigent.llms.context_window.fetch_model_pricing",
+        lambda model: ModelPricing(input_per_token=1e-6, output_per_token=2e-6),
+    )
+    agent = await create_test_agent(client)
+    session = await _create_session(client, agent["id"])
+
+    def _call(input_tokens: int, output_tokens: int) -> None:
+        # Each thread gets its own store instance (its own DB connection) so
+        # the concurrency is real — not short-circuited by a shared connection.
+        store = SqlAlchemyConversationStore(db_uri)
+        sessions_routes._accumulate_session_usage(
+            {
+                "usage": {
+                    "input_tokens": input_tokens,
+                    "output_tokens": output_tokens,
+                    "model": "m1",
+                }
+            },
+            session["id"],
+            store,
+        )
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool:
+        f1 = pool.submit(_call, 1000, 500)
+        f2 = pool.submit(_call, 200, 100)
+        f1.result()
+        f2.result()
+
+    usage = _read_session_usage(db_uri, session["id"])
+    assert usage["input_tokens"] == 1200  # 1000 + 200
+    assert usage["output_tokens"] == 600  # 500 + 100
+    # cost = (1000+200)*1e-6 + (500+100)*2e-6 = 0.0012 + 0.0012 = 0.0024
+    assert usage.get("total_cost_usd") == pytest.approx(0.0024)
+    assert usage["by_model"]["m1"]["input_tokens"] == 1200
+    assert usage["by_model"]["m1"]["total_cost_usd"] == pytest.approx(0.0024)
+
+
 async def test_external_session_usage_records_per_model_breakdown(
     client: httpx.AsyncClient,
     db_uri: str,