perf(policies): load the conversation and spawn tree once per engine build

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

@@ -6230,7 +6230,16 @@ def _build_policy_engine_from_spec_impl(
     spec: AgentSpec,
     session_id: str,
     conversation_store: ConversationStore,
+    conversation: Conversation | None = None,
 ) -> PolicyEngine:
+    """Build an engine for *spec*, reusing a conversation row when held.
+
+    Every caller of this wrapper already loaded the conversation to
+    resolve *spec*; passing it lets the builder skip its own read. Only
+    the row's immutable identity is reused — the builder re-derives
+    labels, session_state and model from a fresh read (see
+    :func:`build_policy_engine`).
+    """
     caps = get_caps()
     host_connection = (
         caps.policy_llm_connection_factory() if caps.policy_llm_connection_factory else None
@@ -6239,6 +6248,11 @@ def _build_policy_engine_from_spec_impl(
         spec=spec,
         conversation_id=session_id,
         conversation_store=conversation_store,
+        conversation=conversation,
+        # The spec was resolved from this row's agent binding; the builder
+        # confirms it against its own fresh read and fails closed if a
+        # switch-agent landed in between.
+        expected_agent_id=conversation.agent_id if conversation is not None else None,
         default_policies=caps.default_policies,
         policy_store=get_policy_store(),
         server_llm=caps.llm,
@@ -6302,7 +6316,7 @@ async def _apply_pending_policy_ask_writes(
     if spec is None:
         return
     engine = await asyncio.to_thread(
-        _build_policy_engine_from_spec, spec, session_id, conversation_store
+        _build_policy_engine_from_spec, spec, session_id, conversation_store, conv
     )
     # The label/state writes hit the DB synchronously too — keep them
     # off the loop.
@@ -6410,11 +6424,13 @@ def _build_evaluation_context(
             harness=hook_harness,
         )
     # REQUEST / RESPONSE — content is the user/assistant text. The wire ``data``
-    # is a dict for the native command hooks (``{"text"|"content": ...}``), but
-    # may be a bare string — opencode's policy plugin sends the prompt text
-    # directly for ``PHASE_REQUEST``. Accept both, and NEVER raise here: a crash
-    # 500s the evaluate endpoint, which silently fails the request/result gate
-    # OPEN (the exact symptom that let cost-over-budget terminal prompts through).
+    # is a dict for every current first-party producer (``{"text"|"content":
+    # ...}``, including OpenCode's plugin, which sends ``{"text": ...}``), but
+    # a bare string is still accepted for ``PHASE_REQUEST`` for compatibility
+    # with older or third-party callers that send the prompt text directly.
+    # Accept both, and NEVER raise here: a crash 500s the evaluate endpoint,
+    # which silently fails the request/result gate OPEN (the exact symptom
+    # that let cost-over-budget terminal prompts through).
     if isinstance(data, str):
         text = data
     elif isinstance(data, dict):
@@ -6690,7 +6706,7 @@ async def _evaluate_output_policy(
         return None
 
     engine = await asyncio.to_thread(
-        _build_policy_engine_from_spec, spec, session_id, conversation_store
+        _build_policy_engine_from_spec, spec, session_id, conversation_store, conv
     )
     ctx = EvaluationContext(
         phase=Phase.RESPONSE,

omnigent/server/routes/_sessions/orchestration.py

@@ -67,6 +67,9 @@
     resolve_ask_timeout,
 )
 from omnigent.runtime.policies.builder import (
+    _sum_subtree_usage,
+    ancestor_ids_from_tree,
+    load_session_tree,
     load_session_usage,
 )
 from omnigent.runtime.policies.engine import PolicyEngine
@@ -596,6 +599,7 @@ def _build_session_list_item(
 def _publish_subtree_cost_to_ancestors(
     conv_store: ConversationStore,
     session_id: str,
+    conv: Conversation | None = None,
 ) -> None:
     """
     Re-publish each ancestor's subtree-summed cost after a child usage update.
@@ -608,18 +612,35 @@ def _publish_subtree_cost_to_ancestors(
     display side.) For each ancestor of *session_id*, recompute its subtree
     priced cost and publish a ``session.usage`` event carrying it.
 
+    One tree load serves the whole walk. Every ancestor of this session is in
+    the same tree, so both the chain and each ancestor's sum are derived from
+    one set of freshly-read rows — previously each ancestor paged the tree
+    again, and the chain came from a conversation row that may have been read
+    before a concurrent delete/recreate moved this session elsewhere.
+
     Sync (does store reads + SSE fan-out); call via
     :func:`asyncio.to_thread`, mirroring the elicitation ancestor-publish
     helpers. ``session_stream.publish`` is safe to call from a worker thread.
 
-    :param conv_store: Store used to discover ancestors and sum each
-        ancestor's subtree usage.
+    :param conv_store: Store used to load the tree.
     :param session_id: The child session whose usage just changed, e.g.
         ``"conv_child123"``.
+    :param conv: The child's already-loaded conversation row, when the caller
+        holds one. Only its ``root_conversation_id`` is used, and only as a
+        hint: :func:`load_session_tree` verifies the tree it names actually
+        contains this session and resolves the root itself when it does not.
+        The parameter belongs to this function, not to whichever caller first
+        needed it, so that no caller can be removed and leave a signature
+        behind that its remaining callers already depend on.
     :returns: None.
     """
-    for ancestor_id in _ancestor_session_ids(conv_store, session_id):
-        ancestor_usage = load_session_usage(ancestor_id, conv_store)
+    tree = load_session_tree(
+        session_id,
+        conv_store,
+        conv.root_conversation_id if conv is not None else None,
+    )
+    for ancestor_id in ancestor_ids_from_tree(tree, session_id):
+        ancestor_usage = _sum_subtree_usage(tree, ancestor_id)
         subtree_cost = _priced_cost_for_display(ancestor_usage)
         usage_by_model = _usage_by_model_for_display(ancestor_usage)
         if subtree_cost is None and usage_by_model is None:
@@ -961,7 +982,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)
@@ -4742,7 +4771,7 @@ async def _evaluate_tool_call_policy(
     if spec is None:
         return None
     engine = await asyncio.to_thread(
-        _build_policy_engine_from_spec, spec, session_id, conversation_store
+        _build_policy_engine_from_spec, spec, session_id, conversation_store, conv
     )
 
     try:
@@ -4863,7 +4892,7 @@ async def _evaluate_input_policy(
         return None
 
     engine = await asyncio.to_thread(
-        _build_policy_engine_from_spec, spec, session_id, conversation_store
+        _build_policy_engine_from_spec, spec, session_id, conversation_store, conv
     )
     ctx = EvaluationContext(
         phase=Phase.REQUEST,
@@ -5775,7 +5804,7 @@ async def _handle_mcp_tools_call(
     # only) and TOOL_RESULT (both paths). Engine construction reads
     # session-policy specs and labels from the DB, so keep it off-loop too.
     engine = await asyncio.to_thread(
-        _build_policy_engine_from_spec, spec, session_id, conversation_store
+        _build_policy_engine_from_spec, spec, session_id, conversation_store, conv
     )
 
     if is_retry:

omnigent/server/routes/sessions/routes_hooks.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import json
 from typing import Any
 
@@ -14,6 +15,7 @@
 from fastapi.responses import Response
 
 from omnigent.codex_native_elicitation import codex_elicitation_id
+from omnigent.entities import Conversation
 from omnigent.errors import ElicitationDeclinedError, ErrorCode, OmnigentError
 from omnigent.policies.types import (
     PolicyAction,
@@ -510,7 +512,14 @@ async def evaluate_policy(
             hook_elicitation_id = raw_elicitation_id
         data = event.get("data") or {}
 
-        conv = conversation_store.get_conversation(session_id)
+        # Reuse the row the ACL check already fetched — same point in the
+        # request, so no less fresh than reading it again here, one query
+        # fewer on the blocking PreToolUse path. Absent for admin callers
+        # (who bypass the conversation lookup) and when permissions are
+        # disabled, which fall back to their own read.
+        conv = access.conversation
+        if conv is None:
+            conv = await asyncio.to_thread(conversation_store.get_conversation, session_id)
         if conv is None:
             raise OmnigentError(
                 f"Session {session_id!r} not found.",
@@ -570,7 +579,7 @@ async def evaluate_policy(
             _caps.policy_llm_connection_factory() if _caps.policy_llm_connection_factory else None
         )
 
-        def _build_engine() -> PolicyEngine:
+        def _build_engine(preloaded_conv: Conversation | None = None) -> PolicyEngine:
             """
             Build a policy engine for this session from the loaded spec.
 
@@ -579,27 +588,41 @@ def _build_engine() -> PolicyEngine:
             does not re-query it during ``evaluate``, so a fresh build is the
             only way to observe a concurrent sibling's just-recorded approval.
 
+            :param preloaded_conv: The conversation row this handler already
+                loaded, passed on the FIRST build only to skip the builder's
+                re-read. Rebuilds that must observe concurrent writes (the
+                ASK-gate re-evaluation) pass ``None`` for a fresh read.
             :returns: A :class:`PolicyEngine` seeded with the latest
                 persisted state for ``session_id``.
             """
             return build_policy_engine(
                 spec=loaded.spec,
                 conversation_id=session_id,
                 conversation_store=conversation_store,
+                conversation=preloaded_conv,
+                # ``agent`` below was resolved from conv.agent_id; the builder
+                # re-reads the row and fails closed if it was rebound since.
+                expected_agent_id=agent.id,
                 default_policies=_caps.default_policies,
                 policy_store=get_policy_store(),
                 server_llm=_caps.llm,
                 host_connection=_host_conn,
             )
 
-        engine = _build_engine()
+        engine = _build_engine(conv)
         # Use the turn-initiating human's identity (persisted at forward time)
         # so per-user policies gate on the correct actor even when the HTTP
         # caller is the runner's service-account credential.  Falls back to
         # user_id for direct API callers and native-terminal sessions (whose
         # turns go via _dispatch_session_event_to_runner, which does not write
         # this label).
-        turn_actor = conv.labels.get(_TURN_ACTOR_LABEL)
+        # Read the actor from the engine's label snapshot, not from the row
+        # fetched at the top of this handler: the engine's labels come from a
+        # read taken after the agent/spec load, so a turn-actor label written
+        # in that window still gates on the right principal. (``agent_id``
+        # cannot be treated the same way — it selects the spec the engine is
+        # built from, so it is necessarily read first.)
+        turn_actor = engine.labels.get(_TURN_ACTOR_LABEL)
         ctx = _build_evaluation_context(
             phase, data, event, actor=_build_actor(turn_actor or user_id)
         )