perf(events): resolve the runner binding without loading the conversation

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:

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/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"``.
@@ -6230,7 +6241,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 +6259,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 +6327,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 +6435,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 +6717,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,