diff --git a/omnigent/runner/routing.py b/omnigent/runner/routing.py index f38b9b2988..efa7af88e6 100644 --- a/omnigent/runner/routing.py +++ b/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: diff --git a/omnigent/runtime/policies/builder.py b/omnigent/runtime/policies/builder.py index 314a50e57c..2cbe3f1566 100644 --- a/omnigent/runtime/policies/builder.py +++ b/omnigent/runtime/policies/builder.py @@ -4,11 +4,13 @@ Called at the top of ``_run_agent_loop``. Seeds any ``LabelDef.initial`` values that are not already present in -``conversation_labels`` using an -``INSERT ... ON CONFLICT DO NOTHING`` semantic so that two -concurrent workflows on the same conversation (the v2 case -tracked in POLICIES.md Open Q #6) never clobber each other's -view of a label's first value. +``conversation_labels`` via +:meth:`ConversationStore.seed_labels_if_absent`, whose +``INSERT ... ON CONFLICT DO NOTHING`` semantics mean two concurrent +workflows on the same conversation (the v2 case tracked in +POLICIES.md Open Q #6) never clobber each other's view of a label's +first value — the database decides which keys are missing, not a +snapshot taken before the write. Phase 2 scope: zero-policy and declared-policy paths both work; concrete Policy subclasses land in Phases 3+, and this builder @@ -18,6 +20,7 @@ from __future__ import annotations import logging +from dataclasses import dataclass from typing import Any import cachetools @@ -160,28 +163,6 @@ def _normalize_usage_for_engine(usage: dict[str, float]) -> dict[str, float]: return usage -def _subtree_usage_seed( - conversation_id: str, - conversation_store: ConversationStore, -) -> dict[str, float]: - """ - SUBTREE-scoped usage seed for the per-subagent cost budget. - - Unlike :func:`_policy_usage_seed` (which seeds from the whole session - tree via ``root_conversation_id``), this seeds from ``conversation_id`` - itself — so the budget gates on this conversation's own subtree cost - (itself + its descendants), not the whole session. - - :param conversation_id: Conversation to seed the subtree usage for, - e.g. ``"conv_child"``. - :param conversation_store: Store to read the subtree usage from. - :returns: Subtree usage seed dict; when an enforcement cost exists its - ``total_cost_usd`` is the enforcement total. - """ - usage = load_session_usage(conversation_id, conversation_store) - return _normalize_usage_for_engine(usage) - - def _resolve_session_owner_cached( conversation_id: str, conversation_store: ConversationStore, @@ -284,6 +265,8 @@ def build_policy_engine( spec: AgentSpec, conversation_id: str, conversation_store: ConversationStore, + conversation: Conversation | None = None, + expected_agent_id: str | None = None, connection_override: dict[str, str] | None = None, default_policies: list[PolicySpec] | None = None, policy_store: PolicyStore | None = None, @@ -301,10 +284,11 @@ def build_policy_engine( When declared labels have an ``initial`` value and no row exists yet in ``conversation_labels``, seeds via - ``ConversationStore.set_labels`` — but only for keys not - already persisted, so existing label state is never - clobbered. The hot cache is built from the freshly seeded - snapshot. + :meth:`ConversationStore.seed_labels_if_absent`, whose + insert-if-absent semantics leave any already-persisted value + untouched. The decision about which keys are missing is made by + the database inside the same statement, never from a snapshot read + beforehand. The hot cache is built from the post-seed snapshot. Policy run order: session policies (from the CRUD API) first, then agent spec policies, then *default_policies* @@ -324,6 +308,25 @@ def build_policy_engine( :param spec: The parsed agent spec. :param conversation_id: The conversation this workflow is running on, e.g. ``"conv_abc123"``. + :param conversation: The already-loaded conversation row for + ``conversation_id``, when the caller holds a current one — skips + the builder's own read. + + **Preload contract** (one rule, applied at every preload site): + a preloaded row may supply only IMMUTABLE identity — its ``id`` + and ``root_conversation_id``. Every mutable field the engine + depends on (labels, session_state, model_override) is re-derived + from a fresh read taken here. A row that has since disappeared + fails closed rather than authorizing from the snapshot. Callers + that rebuild an engine specifically to observe concurrent writes + (the native ASK gate's post-lock re-evaluation) pass ``None``. + :param expected_agent_id: The ``agent_id`` the caller resolved *spec* + from. ``agent_id`` is mutable (switch-agent) but selects the spec, + so it must be read before the engine exists and cannot be + re-derived here. Passing it lets the builder confirm it against + the fresh row and fail closed on a mismatch, instead of + authorizing an evaluation under the previous agent's guardrails. + ``None`` skips the check (callers with no spec/agent coupling). :param conversation_store: The store used for label reads and writes. Held by the engine for the life of the workflow. @@ -358,21 +361,37 @@ def build_policy_engine( guardrails = spec.guardrails agent_policy_specs: list[PolicySpec] = list(guardrails.policies or []) if guardrails else [] session_policy_specs = _load_session_policy_specs(conversation_id, policy_store) - # Session policies are per-conversation, but sub-agents must inherit - # the root conversation's policies so that guardrails set on the - # top-level session (e.g. via sys_add_policy) also govern spawned - # children. Load root policies and prepend them (root policies run - # first, then any child-specific overrides, matching the cost-budget - # root-seeding pattern below). - conv = conversation_store.get_conversation(conversation_id) - root_conversation_id = conv.root_conversation_id if conv is not None else conversation_id - if root_conversation_id != conversation_id: - root_policy_specs = _load_session_policy_specs(root_conversation_id, policy_store) - # Deduplicate: skip root policies already present on the child - # (keyed by policy name) to avoid double-evaluation. - child_names = {p.name for p in session_policy_specs} - root_policy_specs = [p for p in root_policy_specs if p.name not in child_names] - session_policy_specs = root_policy_specs + session_policy_specs + if conversation is not None and conversation.id != conversation_id: + # A misrouted preload would mix one session's labels, state, usage + # and model into another session's authorization decision — fail + # closed rather than build an engine from the wrong row. + raise OmnigentError( + f"preloaded conversation {conversation.id!r} does not match " + f"conversation_id {conversation_id!r}", + code=ErrorCode.INVALID_INPUT, + ) + conv = ( + conversation + if conversation is not None + else conversation_store.get_conversation(conversation_id) + ) + # The row in hand only SUGGESTS a tree root. Loading the tree verifies it + # and reports the root actually used, so everything downstream — the rows, + # the root's own policies, the accounting sums — comes from one snapshot. + # Deriving the root from the pre-refresh row while taking rows from a + # corrected tree was the defect: a conversation deleted and recreated under + # another root seeded the OLD tree's spend. + verified = ( + load_verified_session_tree( + conversation_id, + conversation_store, + conv.root_conversation_id if conv is not None else None, + ) + if conv is not None + else VerifiedSessionTree([], conversation_id, False) + ) + tree = verified.rows + root_conversation_id = verified.root_conversation_id db_default_policy_specs = _load_default_policy_specs(policy_store) admin_policy_specs: list[PolicySpec] = db_default_policy_specs + list(default_policies or []) all_policy_specs = session_policy_specs + agent_policy_specs + admin_policy_specs @@ -384,22 +403,123 @@ def build_policy_engine( all_policy_specs.append(_ASK_ON_ADD_POLICY_SPEC) label_defs = (guardrails.labels or {}) if guardrails else {} + # One conversation read (``conv``, resolved above for policy + # inheritance) and ONE spawn-tree load feed everything below: labels, + # session state (own + inherited root keys), both usage seeds, and the + # model override — on the single-page happy path, which is nearly every + # build; a paged tree pays one further confirming read (see below). The + # helpers each re-fetched the same rows before — ~4x conversation reads + # plus two identical tree loads per build, the dominant cost of a + # policies/evaluate call. + # Freshness contract, applied to EVERY row this function decides from, + # regardless of how it arrived: ONLY immutable identity (id, + # root_conversation_id) survives from the row read above. Every mutable + # field — labels, session_state, model_override, agent_id — is taken from + # the tree load, which happened later and is therefore the newest read. + # + # The provenance of the earlier row does not change the hazard. A caller's + # preload and this function's own ``get_conversation`` are both snapshots + # taken before the tree scan, so both can be stale by the time a decision + # is made; gating the refresh on ``conversation is not None`` closed the + # window on one path and left the identical window open on the other. + # + # The tree load includes archived rows (archived conversations still hold + # spend), so a row missing from the tree has genuinely been deleted. + # Re-read once to confirm, then fail closed — never fall back to the + # earlier copy, which would authorize from state captured before whatever + # removed the row. + if conv is not None: + fresh_self = next((c for c in tree if c.id == conversation_id), None) + if fresh_self is None: + fresh_self = conversation_store.get_conversation(conversation_id) + if fresh_self is None: + # The row existed moments ago and is now gone (deleted + # mid-request). Authorizing from the earlier copy would decide + # against state that no longer exists; authorizing from empty + # state would seed a $0 budget and ALLOW. Fail closed. + raise OmnigentError( + f"Conversation {conversation_id!r} disappeared while building its " + f"policy engine; refusing to authorize from a stale snapshot.", + code=ErrorCode.CONFLICT, + ) + conv = fresh_self + # Agent/spec confirmation — deliberately AFTER the refresh above, and + # nowhere else. Comparing against the earlier row (as a previous revision + # did) validated the very snapshot whose staleness is the hazard, so a + # switch-agent in the window was accepted. Exact equality: a fresh row + # whose binding is ``None``, or no fresh row at all, is a mismatch too — + # not a reason to skip the check. + if expected_agent_id is not None: + fresh_agent_id = conv.agent_id if conv is not None else None + if fresh_agent_id != expected_agent_id: + raise OmnigentError( + f"Session {conversation_id!r} no longer resolves to agent " + f"{expected_agent_id!r} (now {fresh_agent_id!r}); the spec this " + f"engine would enforce is stale. Re-resolve and retry.", + code=ErrorCode.CONFLICT, + ) + # A paged tree does not share one read instant: rows on page one were read + # before page two, so "the tree read is newer than the caller's row" holds + # for the tree but not for any row inside it. Confirm identity once when + # that is actually the case — single-page trees, which is nearly all of + # them, pay nothing. + if conv is not None and verified.paged: + confirmed = conversation_store.get_conversation(conversation_id) + if ( + confirmed is None + or confirmed.root_conversation_id != root_conversation_id + or confirmed.agent_id != conv.agent_id + ): + raise OmnigentError( + f"Conversation {conversation_id!r} moved while its spawn tree was " + f"being paged; refusing to authorize against a tree assembled " + f"across the change.", + code=ErrorCode.CONFLICT, + ) + conv = confirmed + # Session policies are per-conversation, but sub-agents inherit the root + # conversation's policies so guardrails set on the top-level session (e.g. + # via sys_add_policy) also govern spawned children. Loaded from the + # VERIFIED root, after the refresh: reading them from the caller's + # suggested root inherited another tree's guardrails. + if root_conversation_id != conversation_id: + root_policy_specs = _load_session_policy_specs(root_conversation_id, policy_store) + # Deduplicate: skip root policies already present on the child + # (keyed by policy name) to avoid double-evaluation. + child_names = {p.name for p in session_policy_specs} + root_policy_specs = [p for p in root_policy_specs if p.name not in child_names] + session_policy_specs = root_policy_specs + session_policy_specs + all_policy_specs = ( + session_policy_specs + + agent_policy_specs + + admin_policy_specs + + [_ASK_ON_ADD_POLICY_SPEC] + ) + root_conv = ( + conv + if root_conversation_id == conversation_id + else next((c for c in tree if c.id == root_conversation_id), None) + ) + if root_conv is None and conv is not None and root_conversation_id != conversation_id: + # The tree now includes archived rows, so a missing root means the row + # is genuinely gone (deleted mid-request). Read once to confirm before + # deciding — never silently proceed on an absent root. + root_conv = conversation_store.get_conversation(root_conversation_id) initial_labels = _seed_and_load_labels( conversation_id=conversation_id, label_defs=label_defs, conversation_store=conversation_store, + existing=dict(conv.labels) if conv is not None else {}, ) - initial_session_state = _load_session_state(conversation_id, conversation_store) + initial_session_state = dict(conv.session_state) if conv is not None else {} # The cost-budget approval is per-SESSION: the whole spawn tree shares one # soft-threshold gate. A sub-agent runs as its own conversation, so seed its # approved-checkpoint from the ROOT conversation — otherwise approving on the # parent wouldn't carry to the sub-agent and it would re-ask at the same # threshold. Other session_state stays per-conversation; the matching # write-back is routed to the root by PolicyEngine.apply_state_updates. - # (conv and root_conversation_id already resolved above for policy - # inheritance — reuse them here.) if root_conversation_id != conversation_id: - root_state = _load_session_state(root_conversation_id, conversation_store) + root_state = dict(root_conv.session_state) if root_conv is not None else {} for _root_key in ( SESSION_COST_ASK_APPROVED_STATE_KEY, SESSION_COST_UNPRICED_APPROVED_KEY, @@ -409,15 +529,24 @@ def build_policy_engine( # Gating is SESSION-wide: seed from the whole spawn-tree total so a # sub-agent gates against the session's full spend (parent + siblings), # not just its own subtree. The cost read is the enforcement total - # (in-flight sub-agent spend); see _policy_usage_seed. - initial_usage = _policy_usage_seed(conversation_id, conversation_store) - # Conditional injection (#1a): only compute subtree usage when a - # subagent_cost_budget policy is present. - initial_subtree_usage = ( - _subtree_usage_seed(conversation_id, conversation_store) - if _needs_subtree_usage(all_policy_specs) - else None + # (in-flight sub-agent spend); see _policy_usage_seed, whose semantics + # (including the empty seed when the root row is missing) this + # preserves while reusing the single tree load. + initial_usage = ( + _normalize_usage_for_engine(_sum_subtree_usage(tree, root_conversation_id)) + if conv is not None and root_conv is not None + else {} ) + # Conditional injection (#1a): only compute subtree usage when a + # subagent_cost_budget policy is present. Per-node DISPLAY-rooted seed: + # same tree, rooted at the evaluated node instead of the root. + initial_subtree_usage: dict[str, float] | None = None + if _needs_subtree_usage(all_policy_specs): + initial_subtree_usage = ( + _normalize_usage_for_engine(_sum_subtree_usage(tree, conversation_id)) + if conv is not None + else {} + ) # Conditional injection (#1): only pay the owner + daily-cost lookups # when a per-user daily cost-budget policy is actually present. initial_user_daily_cost = ( @@ -425,7 +554,14 @@ def build_policy_engine( if _needs_user_daily_cost(all_policy_specs) else None ) - initial_model = _resolve_session_model(conversation_id, conversation_store, spec) + # Session model: the conversation's model_override (set when a user + # picks a model mid-session) wins over the spec's llm.model; None when + # neither is available and cost policies treat it as undeterminable. + initial_model = ( + conv.model_override + if conv is not None and conv.model_override + else (spec.llm.model if spec.llm else None) + ) # Pass the full ModelPricing so the engine can price cache-read and # cache-write tokens at their own rates via compute_llm_cost(). token_pricing = fetch_model_pricing(spec.llm.model) if spec.llm else None @@ -682,6 +818,7 @@ def _seed_and_load_labels( conversation_id: str, label_defs: dict[str, LabelDef], conversation_store: ConversationStore, + existing: dict[str, str] | None = None, ) -> dict[str, str]: """ Seed declared initial values and return the current snapshot. @@ -698,21 +835,23 @@ def _seed_and_load_labels( labels start unset until a policy writes them). :param conversation_store: Target for both the read and the seed UPSERT. + :param existing: Pre-loaded current label snapshot, passed by callers + that already hold the conversation row (saves a re-read). + ``None`` loads it here. :returns: Full post-seed snapshot of the conversation's 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. + if existing is None: 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( @@ -761,36 +900,6 @@ def _load_session_state( return dict(conv.session_state) -def _resolve_session_model( - conversation_id: str, - conversation_store: ConversationStore, - spec: AgentSpec, -) -> str | None: - """ - Resolve the model the session is currently using. - - Prefers the conversation's ``model_override`` (set when a user - picks a model mid-session via ``/model`` or the web model picker) - and falls back to the agent spec's ``llm.model``. ``None`` when - neither is available — the conversation does not exist yet, has no - override, and the spec declares no ``llm`` block — in which case - cost policies treat the model as undeterminable. - - :param conversation_id: Conversation to read the override from, - e.g. ``"conv_abc123"``. - :param conversation_store: Store to read the conversation from. - :param spec: The parsed agent spec (its ``llm.model`` is the - fallback when no override is set). - :returns: The active model id, e.g. ``"databricks-claude-opus-4-8"`` - or the native tier alias ``"opus"``; ``None`` when - undeterminable. - """ - conv = conversation_store.get_conversation(conversation_id) - if conv is not None and conv.model_override: - return conv.model_override - return spec.llm.model if spec.llm else None - - # Page size for walking a spawn tree when summing sub-agent usage. # Sub-agent trees are small in practice, but we still paginate so a # large tree is not silently truncated (see load_session_usage). @@ -841,6 +950,8 @@ def _merge_by_model( def load_session_usage( conversation_id: str, conversation_store: ConversationStore, + *, + root_conversation_id: str | None = None, ) -> dict[str, Any]: """ Load cumulative session usage for a conversation **plus all of its @@ -864,6 +975,14 @@ def load_session_usage( :param conversation_id: Conversation to load, e.g. ``"conv_abc123"``. :param conversation_store: Store to read from. + :param root_conversation_id: The conversation's tree root, when the + caller already holds the row. Skips the internal conversation read. + The root binding is immutable **per row**, not per conversation id: + a conversation deleted and recreated under the same id gets a new + row, whose root may differ. A supplied root is therefore validated + against the tree it produces — if this conversation is not in that + tree, the caller's row is stale and the root is resolved here + instead. ``None`` resolves it here from the start. :returns: Summed usage dict with keys ``input_tokens``, ``output_tokens``, ``total_tokens``, ``total_cost_usd`` (the DISPLAY cost sum — statusLine ``S`` for claude-native), and @@ -877,10 +996,123 @@ def load_session_usage( the policy seed (:func:`_policy_usage_seed`) reads ``policy_cost_usd`` (both unaffected by ``by_model``). """ + tree = load_session_tree(conversation_id, conversation_store, root_conversation_id) + return _sum_subtree_usage(tree, conversation_id) + + +def load_session_tree( + conversation_id: str, + conversation_store: ConversationStore, + root_conversation_id: str | None = None, +) -> list[Conversation]: + """ + Load the spawn tree *conversation_id* belongs to, verifying the root. + + One place owns the reuse rule for a caller-supplied tree root, so every + consumer gets the same guarantee: the tree comes back containing this + conversation, or the supplied root was stale and is resolved again. + + A caller's ``root_conversation_id`` is immutable **per row**. Deleting a + conversation and recreating it under the same id produces a new row that + may sit in a different tree, so a row read earlier in the request can + name a root this conversation no longer belongs to. Rather than trust it + or re-read unconditionally, the supplied root is checked against the + tree it produced — a membership test on rows already in memory, so the + happy path costs nothing and the stale path costs one read. + + :param conversation_id: The conversation whose tree is wanted, + e.g. ``"conv_abc123"``. + :param conversation_store: Store to read from. + :param root_conversation_id: Caller-supplied tree root, validated as + above. ``None`` resolves the root here. + :returns: Every conversation in the tree (root plus all descendants, + archived included). Empty when the conversation does not exist. + """ + return load_verified_session_tree( + conversation_id, conversation_store, root_conversation_id + ).rows + + +@dataclass(frozen=True) +class VerifiedSessionTree: + """A spawn tree together with what is known about how it was loaded. + + :param rows: Every conversation in the tree, archived included. Empty + when the conversation does not exist. + :param root_conversation_id: The root the rows were actually loaded + from, which is not necessarily the one the caller suggested. + :param paged: Whether the listing needed more than one page. Rows on an + earlier page were read before rows on a later one, so for a paged + tree "the tree read is newer than the caller's row" holds for the + tree as a whole but not for any individual row in it. + """ + + rows: list[Conversation] + root_conversation_id: str + paged: bool + + +def load_verified_session_tree( + conversation_id: str, + conversation_store: ConversationStore, + root_conversation_id: str | None = None, +) -> VerifiedSessionTree: + """ + Load a spawn tree and report the root it came from. + + Same verification as :func:`load_session_tree`, but callers that derive + more than the sums from a tree — the tree root itself, the policies + attached to that root — need to know which root was used, because a + supplied one may have been discarded. Deriving those from the caller's + root while taking the rows from a corrected tree mixes two epochs. + + :param conversation_id: The conversation whose tree is wanted. + :param conversation_store: Store to read from. + :param root_conversation_id: Caller-supplied root, treated as a hint. + :returns: The rows, the root they came from, and whether it paged. + """ + if root_conversation_id is None: + conv = conversation_store.get_conversation(conversation_id) + if conv is None: + return VerifiedSessionTree([], conversation_id, False) + root_conversation_id = conv.root_conversation_id + rows, paged = _load_tree_pages(root_conversation_id, conversation_store) + return VerifiedSessionTree(rows, root_conversation_id, paged) + + rows, paged = _load_tree_pages(root_conversation_id, conversation_store) + if any(c.id == conversation_id for c in rows): + return VerifiedSessionTree(rows, root_conversation_id, paged) + # Not in the tree the supplied root produced: either this conversation + # is gone, or it now lives in a different tree. Resolve it once. conv = conversation_store.get_conversation(conversation_id) if conv is None: - return {} - tree = _load_tree_conversations(conv.root_conversation_id, conversation_store) + return VerifiedSessionTree([], root_conversation_id, paged) + if conv.root_conversation_id == root_conversation_id: + return VerifiedSessionTree(rows, root_conversation_id, paged) + rows, paged = _load_tree_pages(conv.root_conversation_id, conversation_store) + return VerifiedSessionTree(rows, conv.root_conversation_id, paged) + + +def _sum_subtree_usage( + tree: list[Conversation], + conversation_id: str, +) -> dict[str, Any]: + """ + Sum usage across the subtree of *tree* rooted at *conversation_id*. + + Pure aggregation over an already-loaded spawn tree — no store reads. + :func:`build_policy_engine` loads the tree once and derives both the + session-wide gating seed (rooted at the tree root) and the per-node + subtree seed (rooted at the evaluated node) from the same list; + :func:`load_session_usage` wraps this for callers that start from a + conversation id. See :func:`load_session_usage` for the shape of the + returned dict. + + :param tree: All conversations in the spawn tree (from + :func:`_load_tree_conversations`); order-independent. + :param conversation_id: The subtree root to sum from. + :returns: Summed usage dict (see :func:`load_session_usage`). + """ subtree_ids = _subtree_conversation_ids(tree, conversation_id) totals: dict[str, Any] = {} # Per-model breakdown summed across the subtree, parallel to the flat sums. @@ -990,6 +1222,13 @@ def _load_tree_conversations( # (not just "default") are included in the tree. kind=None, root_conversation_id=root_conversation_id, + # Archived conversations still hold spend, and archiving must not + # move a budget gate: excluding them let an archive-after-preload + # (or an archived mid-tree node, which orphaned its descendants + # from the walk) seed the enforcement total as $0 and ALLOW a tool + # call over budget. The tree is an accounting structure, not a + # user-facing listing. + include_archived=True, ) convs.extend(page.data) if not page.has_more or page.last_id is None: @@ -998,6 +1237,38 @@ def _load_tree_conversations( return convs +def _load_tree_pages( + root_conversation_id: str, + conversation_store: ConversationStore, +) -> tuple[list[Conversation], bool]: + """ + Page through a spawn tree, reporting whether more than one page was read. + + :param root_conversation_id: The tree's root conversation id. + :param conversation_store: Store to read from. + :returns: ``(rows, paged)`` — ``paged`` is ``True`` when the listing + needed a second page, which means the rows do not share one read + instant. + """ + convs: list[Conversation] = [] + after: str | None = None + pages = 0 + while True: + page = conversation_store.list_conversations( + limit=_SUBTREE_USAGE_PAGE_SIZE, + after=after, + kind=None, + root_conversation_id=root_conversation_id, + include_archived=True, + ) + pages += 1 + convs.extend(page.data) + if not page.has_more or page.last_id is None: + break + after = page.last_id + return convs, pages > 1 + + def _subtree_conversation_ids( tree: list[Conversation], conversation_id: str, @@ -1035,6 +1306,53 @@ def _subtree_conversation_ids( return subtree +def ancestor_ids_from_tree( + tree: list[Conversation], + conversation_id: str, +) -> list[str]: + """ + Walk a conversation's ancestor chain inside an already-loaded tree. + + The mirror of :func:`_subtree_conversation_ids`, and public for the + same reason the tree loader is: the ancestor chain must come from the + same freshly-read rows as the sums, not from a conversation row the + caller read earlier. ``parent_conversation_id`` is immutable per row, + but a conversation deleted and recreated under the same id gets a new + row with a new parent, so a caller's copy can name a chain that no + longer exists — and walking it publishes to the wrong sessions. + + Pure: no store reads, and no reads are needed, because a tree already + contains every row on the chain by construction. + + :param tree: All conversations in the spawn tree (from + :func:`load_session_tree`); order-independent. + :param conversation_id: The node to walk upward from, + e.g. ``"conv_child123"``. + :returns: Ancestor ids nearest-parent-first. Empty when the node is + top-level, absent from the tree, or its chain is cyclic. + """ + by_id = {c.id: c for c in tree} + ancestors: list[str] = [] + seen = {conversation_id} + current = by_id.get(conversation_id) + while current is not None and current.parent_conversation_id is not None: + parent_id = current.parent_conversation_id + if parent_id in seen: + # A cycle makes the whole chain untrustworthy, not just the rest + # of it: returning the part walked so far would publish a + # descendant's usage to whichever ids happened to come first. + return [] + parent = by_id.get(parent_id) + if parent is None: + # The parent link points outside this tree. Appending before + # checking published to a conversation that is not there. + return [] + ancestors.append(parent_id) + seen.add(parent_id) + current = parent + return ancestors + + def _load_default_policy_specs( policy_store: PolicyStore | None, ) -> list[PolicySpec]: 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/helpers.py b/omnigent/server/routes/_sessions/helpers.py index b284fe55eb..371835e85c 100644 --- a/omnigent/server/routes/_sessions/helpers.py +++ b/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, diff --git a/omnigent/server/routes/_sessions/orchestration.py b/omnigent/server/routes/_sessions/orchestration.py index ecd9f74804..3ab82d18f4 100644 --- a/omnigent/server/routes/_sessions/orchestration.py +++ b/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,40 @@ 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) + # Loaded HERE, not taken from the caller. What each ancestor's badge + # should read is a fact about now, at publication time — a tree captured + # when the request arrived can already describe a smaller total, and two + # siblings publishing from their own request-start snapshots delivered a + # newer figure followed by an older one, leaving the parent's badge stale. + 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: @@ -840,6 +866,7 @@ def _accumulate_session_usage( resp_obj: dict[str, Any], session_id: str, conversation_store: ConversationStore, + conversation: Conversation | None = None, ) -> float | None: """ Increment the session's cumulative token counters from a @@ -870,6 +897,17 @@ def _accumulate_session_usage( e.g. ``"conv_abc123"``. :param conversation_store: Store for reading and writing the ``session_usage`` column. + :param conversation: The relay's already-read conversation row for + this session, reused here for the pricing lookup instead of + re-reading the same row on every ``response.completed``. This is + point-in-time reuse, not the policy engine's immutable-fields-only + rule: the row supplies the mutable ``model_override``, so a + ``/model`` switch landing between the read and this call prices + the turn with the previous model — bounded to one event, and + ``usage.model`` (the model the harness actually used) takes + precedence when the harness reports it. The usage increment itself + is a locked read-modify-write in the store, so a stale row cannot + skew the persisted total. ``None`` reads the row here. :returns: The session's cumulative priced cost in USD after this update (for the caller to broadcast on a ``session.usage`` event), or ``None`` when the session is unpriced or carries no @@ -889,8 +927,14 @@ def _accumulate_session_usage( # 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) + # avoid the read-modify-write race). The caller may pass the row it just + # read (the relay roll-up needs the same row immediately afterwards), in + # which case pricing reuses it rather than reading it again. + conv = ( + conversation + if conversation is not None + else conversation_store.get_conversation(session_id) + ) # Compute cost delta if pricing is available for the model. Resolve # the model to price with, most-specific first: @@ -961,7 +1005,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) @@ -1158,6 +1210,7 @@ def _persist_native_cumulative_usage( async def _persist_external_session_usage( session_id: str, + conv: Conversation, body: SessionEventInput, conversation_store: ConversationStore, ) -> int | None: @@ -1169,6 +1222,9 @@ async def _persist_external_session_usage( (:func:`_persist_native_cumulative_usage`) must be present. :param session_id: Session/conversation identifier. + :param conv: The already-loaded conversation row; only its immutable + ``root_conversation_id`` is read (so a request-start row is safe), + letting the subtree recompute skip one conversation read. :param body: External session-usage event body. :param conversation_store: Store used to upsert the labels. :returns: The persisted ``context_tokens`` when present, else ``None``. @@ -1231,7 +1287,16 @@ async def _persist_external_session_usage( # hide in-flight sub-agent spend until the next child flush (the badge would # oscillate own ⇄ subtree). For a childless session the subtree is just # itself, so this equals own cost — one indexed tree query per flush. - subtree_usage = await asyncio.to_thread(load_session_usage, session_id, conversation_store) + # This tree serves only this session's own subtree total. The ancestor + # re-publish below loads its OWN tree instead of reusing this one — see + # the comment down there for why sharing it would be wrong. + tree = await asyncio.to_thread( + load_session_tree, + session_id, + conversation_store, + conv.root_conversation_id, + ) + subtree_usage = _sum_subtree_usage(tree, session_id) subtree_cost = _priced_cost_for_display(subtree_usage) usage_by_model = _usage_by_model_for_display(subtree_usage) # Only include fields that were sent; the client treats absent @@ -1256,12 +1321,16 @@ async def _persist_external_session_usage( # This session's usage also moves its ANCESTORS' subtree cost (its spend # rolls up into every ancestor), so re-publish each ancestor's subtree cost # too — otherwise a grandparent's badge wouldn't reflect a deep descendant. - # No-op for a top-level session (no ancestors). Threaded: it pages the - # conversation tree per ancestor. + # Publishes nothing for a top-level session (no ancestors to walk), but + # still pays its own tree load to find that out — it reads its own tree + # rather than reusing the one summed above: what an ancestor's badge + # should show is a fact about publication time, and a sibling's event may + # have landed since this request started. await asyncio.to_thread( _publish_subtree_cost_to_ancestors, conversation_store, session_id, + conv, ) return raw_tokens @@ -4352,11 +4421,27 @@ async def _relay_runner_stream( # policy callables can read # event["context"]["usage"]["total_cost_usd"] and the # subtree roll-up below sees the new totals. - _accumulate_session_usage( - event.get("response", {}), - session_id, - conversation_store, + # One conversation read serves both this pricing + # lookup and the subtree roll-up below (which needs + # the immutable tree root), instead of one each. + _conv_row = await asyncio.to_thread( + conversation_store.get_conversation, session_id ) + # That one read decides for the whole completion. An + # absent row means this session is gone: nothing to + # bill, nothing to publish. Passing ``None`` onward + # instead meant it read as "re-read the row" to the + # accumulator and as "skip" to the roll-up below — so a + # session deleted and recreated under the same id had + # the turn's usage persisted to the new row while the + # event reporting it was suppressed. + if _conv_row is not None: + _accumulate_session_usage( + event.get("response", {}), + session_id, + conversation_store, + _conv_row, + ) # Push the server-computed cost AND token breakdown # to the web client's session indicator, rolled up # over the spawn subtree. The session's own event @@ -4373,10 +4458,15 @@ async def _relay_runner_stream( # surfaces tokens). context_tokens/window already ride # on the response.completed event. Threaded: store # reads + SSE fan-out. - _subtree_usage = await asyncio.to_thread( - load_session_usage, - session_id, - conversation_store, + _subtree_usage = ( + await asyncio.to_thread( + load_session_usage, + session_id, + conversation_store, + root_conversation_id=_conv_row.root_conversation_id, + ) + if _conv_row is not None + else {} ) _subtree_cost = _priced_cost_for_display(_subtree_usage) _usage_by_model = _usage_by_model_for_display(_subtree_usage) @@ -4397,6 +4487,7 @@ async def _relay_runner_stream( _publish_subtree_cost_to_ancestors, conversation_store, session_id, + _conv_row, ) # Reset the turn-scoped response_id on any @@ -4742,7 +4833,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 +4954,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 +5866,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: @@ -6256,6 +6347,56 @@ async def _get_session_snapshot( skipping the ``get_conversation`` read. Pass it when the caller just authorized the session (which fetched the same row) so the snapshot doesn't re-read it. ``None`` reads it here as before. + + Reuse here is POINT-IN-TIME, not the policy-engine rule. The + engine re-derives every mutable field and keeps only the + immutable ``id`` / ``root_conversation_id`` from a preloaded row, + because it gates tool calls. A snapshot is a read projection, and + a MIXED-epoch one. A field-by-field table here previously drifted + from the response schema (an invented field name, a field claimed + agent-keyed that isn't, a source claimed single that isn't), and a + prior group-based rewrite still assigned several fields to a + single group when their actual value can come from either — + a symptom of the same "table copies the code, then the code moves + on" problem this reuse itself was written to avoid. What follows + is deliberately loose rather than exhaustively exact per field, + since that precision is exactly what keeps drifting: + + - Most of :class:`~omnigent.server.schemas.SessionResponse` — + identity, timestamps, the runner/host binding, archive state, + the model/cost overrides, external-session and workspace fields + — comes straight off THIS row. ``labels`` is also sourced from + this row alone, but is not passed through raw: it's + viewer-scoped and gets a closed-status marker derived from the + row's title, so "off this row" does not mean "unexamined". + - A handful of fields are read fresh after this row: runner/host + liveness, elicitations, items, the subtree cost rollup, and + process-local caches (todos, sandbox status, MCP startup, the + active response id). ``skills`` is fetched live from the + runner, keyed by session id, with no dependency on + ``agent_id`` at all. + - Several fields are genuine hybrids, not cleanly one epoch or + the other, and each blends its two sources differently: + ``agent_name``, ``llm_model`` and (absent a persisted + per-session override, which wins outright) ``harness`` are + read fresh but resolved THROUGH this row's ``agent_id`` — they + describe the agent this row named, not necessarily the one + bound now. Lifecycle ``status`` is read fresh (from the + relay-fed cache, or a live runner query on a cache miss) but + can be overridden to ``"failed"`` by a runner-crash report + looked up by THIS row's ``runner_id``. ``context_window`` is + resolved fresh through the agent's spec, but a raw override + value stored in THIS row's labels wins when present. + - ``permission_level`` is neither: it is the caller's ACL grant, + resolved before this function — and before the row itself, when + the caller preloads it — ever runs. + + None of this is an enforcement decision. A rebind landing + mid-request can be queried on the NEW runner while the response + reports the OLD ``runner_id`` and the OLD agent's resolved fields. + Making the whole projection single-epoch is a change to this + endpoint's contract, not to a redundant read. Do not reuse a + snapshot row to make an enforcement decision. :param liveness_lookup: Bulk session-liveness lookup (the server's ``_bulk_session_liveness``) used to populate ``runner_online`` and ``host_online`` on the snapshot. ``None`` (e.g. focused @@ -6460,7 +6601,14 @@ async def _get_session_snapshot( # is persisted on its own child conversation, not the parent's, so the # parent's own session_usage would under-report. Off the event loop # because it pages the conversation tree from the store. - subtree_usage = await asyncio.to_thread(load_session_usage, conv.id, conv_store) + # The row in hand supplies the (immutable) tree root, so the helper + # skips its own conversation re-read; the tree scan itself stays fresh. + subtree_usage = await asyncio.to_thread( + load_session_usage, + conv.id, + conv_store, + root_conversation_id=conv.root_conversation_id, + ) # Static signal telling the open view a host-bound, host-down session is a # resumable managed host it can wake by sending a message, vs a terminal # host_offline dead-end. Computed independently of liveness_lookup (the web diff --git a/omnigent/server/routes/sessions/routes_events.py b/omnigent/server/routes/sessions/routes_events.py index 2ae921c807..af5f7801f9 100644 --- a/omnigent/server/routes/sessions/routes_events.py +++ b/omnigent/server/routes/sessions/routes_events.py @@ -848,6 +848,7 @@ async def post_event( # post-hoc here — a logged output cannot be un-logged.) await _persist_external_session_usage( session_id, + conv, body, conversation_store, ) diff --git a/omnigent/server/routes/sessions/routes_hooks.py b/omnigent/server/routes/sessions/routes_hooks.py index d46f5f5582..99abafeaaa 100644 --- a/omnigent/server/routes/sessions/routes_hooks.py +++ b/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,6 +588,10 @@ 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``. """ @@ -586,20 +599,30 @@ def _build_engine() -> PolicyEngine: 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) ) diff --git a/omnigent/server/routes/usage.py b/omnigent/server/routes/usage.py index b3573d9f8d..270df2249a 100644 --- a/omnigent/server/routes/usage.py +++ b/omnigent/server/routes/usage.py @@ -122,7 +122,13 @@ def _build_usage_report( for conv in page.data: if conv.agent_id is None: continue - usage = load_session_usage(conv.id, conversation_store) + # The listed row supplies the (immutable) tree root, so the + # helper skips one conversation re-read per listed session. + usage = load_session_usage( + conv.id, + conversation_store, + root_conversation_id=conv.root_conversation_id, + ) sessions.append( SessionUsage( id=conv.id, diff --git a/omnigent/stores/conversation_store/__init__.py b/omnigent/stores/conversation_store/__init__.py index 07cc4ea2c8..fde2a43436 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. """ @@ -832,6 +842,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, @@ -852,9 +902,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, @@ -936,11 +990,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"``. @@ -949,6 +1007,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, @@ -1026,6 +1115,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 3031c11728..73ddfe7543 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 from sqlalchemy import ( @@ -20,6 +21,7 @@ text, update, ) +from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import QueryableAttribute, Session, aliased from sqlalchemy.sql.selectable import Subquery @@ -461,6 +463,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, @@ -1040,22 +1102,82 @@ def find_imported_conversation( def get_runner_ids(self, conversation_ids: list[str]) -> dict[str, str | None]: """ - Single ``SELECT id, runner_id WHERE id IN (...)`` — bulk - variant of :meth:`get_conversation` for the runner-dot path. - Missing ids are omitted; ids without a bound runner map to - ``None``. + Bulk ``runner_id`` lookup, gated on the conversation still existing. + + The binding lives on the Omnigent metadata row, but the AP + ``conversations`` row is the EXISTENCE authority: deletion is + deliberately ordered AP-row-first, metadata-second, and a failed + second transaction leaves orphaned metadata for a conversation that + no longer exists (documented as an acceptable best-effort tradeoff + in :meth:`delete_conversation`). That tradeoff is only acceptable + while every reachability check consults the AP row — a + metadata-only read would route a deleted conversation to its old + runner instead of reporting it missing. + + The lookup is therefore rooted in the AP row and reaches sideways for + the binding, never the other way round: one outer-joined query when + both live on the same bind, two when they are split (still narrower + than the three a full :meth:`get_conversation` costs). Rooting it in + metadata instead answered only one of the two failure directions — + a deleted conversation with orphaned metadata was correctly dropped, + while an existing conversation whose metadata row is absent was + *also* dropped, turning "exists but unbound" into a false + ``NOT_FOUND``. Both states are constructible: creation commits the AP + row and the metadata row in separate transactions, the same + two-transaction shape as deletion in the opposite order. + + :param conversation_ids: Conversation ids to look up. + :returns: ``{id: runner_id or None}`` for ids whose AP row exists. + Ids with no AP row are omitted, exactly as a missing + :meth:`get_conversation` would report them; ids that exist + without a binding map to ``None``. """ if not conversation_ids: return {} unique_ids = list(set(conversation_ids)) + same_bind = self._conv_engine is self._engine + if same_bind: + with self._session() as session: + rows = session.execute( + select(SqlConversation.id, SqlConversationMetadata.runner_id) + .outerjoin( + SqlConversationMetadata, + and_( + SqlConversationMetadata.workspace_id == SqlConversation.workspace_id, + SqlConversationMetadata.id == SqlConversation.id, + ), + ) + .where( + SqlConversation.workspace_id == current_workspace_id(), + SqlConversation.id.in_(unique_ids), + ) + ).all() + return {row.id: row.runner_id for row in rows} + # Split-DB: no cross-bind join, so read existence from the AP bind and + # overlay the bindings. Seeding every existing id with ``None`` keeps + # the two states distinct without a second existence check. + with self._conv_session() as ap_session: + existing = set( + ap_session.execute( + select(SqlConversation.id).where( + SqlConversation.workspace_id == current_workspace_id(), + SqlConversation.id.in_(unique_ids), + ) + ).scalars() + ) + if not existing: + return {} + bindings: dict[str, str | None] = dict.fromkeys(existing) with self._session() as session: rows = session.execute( select(SqlConversationMetadata.id, SqlConversationMetadata.runner_id).where( SqlConversationMetadata.workspace_id == current_workspace_id(), - SqlConversationMetadata.id.in_(unique_ids), + SqlConversationMetadata.id.in_(list(existing)), ) ).all() - return {row.id: row.runner_id for row in rows} + for row in rows: + bindings[row.id] = row.runner_id + return bindings def get_session_connectivity( self, conversation_ids: list[str] @@ -1250,6 +1372,62 @@ def set_labels( with self._conv_session() 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() 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, @@ -1277,6 +1455,102 @@ 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. + """ + return self._mutate_metadata_json( + conversation_id, "session_state", mutate, what="session state" + ) + + def _mutate_metadata_json( + self, + 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. + + 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 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 + + with self._session_immediate() 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() + 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, @@ -1361,32 +1635,20 @@ 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() 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 current + return self._mutate_metadata_json( + conversation_id, + "session_usage", + lambda current: apply_session_usage_delta(current, delta), + what="session usage", + ) def add_daily_cost(self, user_id: str, day_utc: str, delta_usd: float) -> None: """ diff --git a/tests/runner/test_routing.py b/tests/runner/test_routing.py index ba8561d74a..b6be3b3cd8 100644 --- a/tests/runner/test_routing.py +++ b/tests/runner/test_routing.py @@ -44,6 +44,14 @@ def __init__(self, conversations: dict[str, Conversation]) -> None: """ self._conversations = conversations + def get_runner_ids(self, conversation_ids: list[str]) -> dict[str, str | None]: + """Bulk binding read mirroring the real store's contract.""" + return { + cid: self._conversations[cid].runner_id + for cid in conversation_ids + if cid in self._conversations + } + def get_conversation(self, conversation_id: str) -> Conversation | None: """ Return a conversation by id. @@ -249,3 +257,72 @@ async def test_runner_router_existing_conversation_returns_none_when_unpinned() assert router.client_for_existing_conversation("conv_test") is None finally: await router.aclose() + + +@pytest.mark.parametrize( + ("method", "missing", "unpinned"), + [ + ("client_for_session_resources", ErrorCode.NOT_FOUND, ErrorCode.CONFLICT), + # The optional variant reports both as "no runner to route to". + ("client_for_existing_conversation", None, None), + ], +) +@pytest.mark.asyncio +async def test_narrowed_binding_routing_is_fresh_and_keeps_its_semantics( + method: str, + missing: ErrorCode | None, + unpinned: ErrorCode | None, +) -> None: + """ + Every per-event route that narrowed its conversation read to the binding + must stay fresh, stay narrow, and keep its own missing/unpinned answers. + + Parametrized over the methods rather than written per method: the first + of the two was covered and the second was not, so reverting it to a + full-row read left the suite green. + + Narrow: the full-row read explodes, so any regression to + ``get_conversation`` fails here. Fresh: a rebind landing between + resolutions is picked up by the next one. Semantics: a missing + conversation and an existing-but-unbound one stay distinguishable, which + is what the store's two-state binding lookup exists to preserve. + """ + registry = TunnelRegistry() + registry.register("runner_new", _FakeWebSocket(), _hello(harnesses=["codex"])) + + class _BindingOnlyStore: + def __init__(self) -> None: + self.binding: dict[str, str | None] = { + "conv_test": "runner_old", + "conv_unpinned": None, + } + + def get_runner_ids(self, conversation_ids: list[str]) -> dict[str, str | None]: + return {c: self.binding[c] for c in conversation_ids if c in self.binding} + + def get_conversation(self, conversation_id: str) -> None: + raise AssertionError("routing must not load the full conversation row") + + store = _BindingOnlyStore() + router = RunnerRouter(registry=registry, conversation_store=store) # type: ignore[arg-type] + resolve = getattr(router, method) + try: + # Bound to an offline runner: both methods report it the same way. + with pytest.raises(OmnigentError) as excinfo: + resolve("conv_test") + _assert_omnigent_error(excinfo, code=ErrorCode.RUNNER_UNAVAILABLE) + + # A concurrent rebind must be visible to the very next resolution. + store.binding["conv_test"] = "runner_new" + routed = resolve("conv_test") + assert routed is not None and routed.runner_id == "runner_new" + + for conv_id, expected in (("conv_missing", missing), ("conv_unpinned", unpinned)): + if expected is None: + assert resolve(conv_id) is None, conv_id + else: + with pytest.raises(OmnigentError) as excinfo: + resolve(conv_id) + _assert_omnigent_error(excinfo, code=expected) + finally: + await router.aclose() diff --git a/tests/runtime/policies/test_builder.py b/tests/runtime/policies/test_builder.py index d03b73f8d0..3c2458c32a 100644 --- a/tests/runtime/policies/test_builder.py +++ b/tests/runtime/policies/test_builder.py @@ -25,10 +25,13 @@ from __future__ import annotations import uuid +from collections.abc import Callable, Iterator +from contextlib import contextmanager from pathlib import Path import pytest +from omnigent.entities import Conversation from omnigent.runtime.policies.builder import build_policy_engine from omnigent.spec.parser import parse from omnigent.spec.types import ( @@ -961,7 +964,7 @@ def test_normalize_usage_for_engine_drops_display_fields() -> None: """ _normalize_usage_for_engine removes by_model and promotes policy_cost_usd. - Both _policy_usage_seed and _subtree_usage_seed use this helper to + Both the session-wide and subtree usage seeds use this helper to prepare usage for the engine: strip the display-only ``by_model`` breakdown, and swap ``policy_cost_usd`` to ``total_cost_usd`` for enforcement cost (falling back to ``total_cost_usd`` when no enforcement @@ -1001,3 +1004,982 @@ 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 + + +@contextmanager +def _count_sql(store: SqlAlchemyConversationStore) -> Iterator[list[str]]: + """Capture every statement the store executes, on either bind.""" + from sqlalchemy import event as sa_event + + seen: list[str] = [] + + def _on(conn, cursor, statement, params, context, many): + seen.append(statement) + + engines = {store._engine, store._conv_engine} + for engine in engines: + sa_event.listen(engine, "before_cursor_execute", _on) + try: + yield seen + finally: + for engine in engines: + sa_event.remove(engine, "before_cursor_execute", _on) + + +@pytest.mark.parametrize("preloaded", [False, True], ids=["no-preload", "preload"]) +def test_build_issues_one_read_and_one_tree_scan( + conversation_store: SqlAlchemyConversationStore, + preloaded: bool, +) -> None: + """ + One engine build costs one conversation read plus ONE spawn-tree scan, + and nothing at all for the read when the caller supplies the row. + + Counts SQL STATEMENTS, not store-method calls: the redundancy this + change removes was measured in queries, and a store-call count cannot + see a helper that issues three statements per call. The builder used to + re-fetch the conversation ~4x and walk the tree twice, once per usage + seed. + + Both seeds must stay correct and identical either way: session-wide + gating from the whole tree, subtree display from the node's own subtree. + """ + from omnigent.spec.types import FunctionPolicySpec, FunctionRef + + parent = conversation_store.create_conversation() + child = conversation_store.create_conversation( + kind="sub_agent", parent_conversation_id=parent.id, title=_sub_agent_title() + ) + sibling = conversation_store.create_conversation( + kind="sub_agent", parent_conversation_id=parent.id, title=_sub_agent_title() + ) + conversation_store.set_session_usage(parent.id, {"total_cost_usd": 0.10}) + conversation_store.set_session_usage(child.id, {"total_cost_usd": 0.05}) + conversation_store.set_session_usage(sibling.id, {"total_cost_usd": 0.03}) + child_row = conversation_store.get_conversation(child.id) if preloaded else None + + subagent_budget = FunctionPolicySpec( + name="subtree_budget", + on=None, + function=FunctionRef( + path="omnigent.policies.builtins.cost.subagent_cost_budget", + arguments={"max_cost_usd": 10.0}, + ), + ) + with _count_sql(conversation_store) as statements: + engine = build_policy_engine( + spec=AgentSpec(spec_version=1, name="child"), + conversation_id=child.id, + conversation_store=conversation_store, + conversation=child_row, + default_policies=[subagent_budget], + ) + + # Semantics preserved: session-wide gate vs per-subtree display. + assert engine.usage["total_cost_usd"] == pytest.approx(0.18) + assert engine._subtree_usage is not None + assert engine._subtree_usage["total_cost_usd"] == pytest.approx(0.05) + + # The tree scan is one paged listing; the conversation read is a triplet + # (row + metadata + labels) and disappears entirely with a preload. + # Shape first, then the total. SQLite's connection PRAGMAs are setup, not + # work the builder asked for. + executed = [" ".join(q.split()) for q in statements if not q.startswith("PRAGMA")] + tree_scans = [q for q in executed if "root_conversation_id = " in q] + point_reads = [q for q in executed if "FROM conversations" in q and "conversations.id = " in q] + assert len(tree_scans) == 1, executed + assert len(point_reads) == (0 if preloaded else 1), executed + # A conversation read is a triplet (row + metadata + labels), and so is the + # tree scan; the preload removes one whole triplet. + assert len(executed) == (3 if preloaded else 6), [q[:80] for q in executed] + + +def test_build_counts_archived_spend_and_inherits_approval( + conversation_store: SqlAlchemyConversationStore, +) -> None: + """ + Archiving must not move a budget gate. An archived root's own spend + still counts toward the session total, and its approval checkpoint is + still inherited — archiving is a listing concern, not an accounting + one. (This previously asserted the opposite, codifying an omission + that let an archive-after-preload seed $0 and ALLOW over budget.) + """ + from omnigent.policies.schema import SESSION_COST_ASK_APPROVED_STATE_KEY + + parent = conversation_store.create_conversation() + child = conversation_store.create_conversation( + kind="sub_agent", parent_conversation_id=parent.id, title=_sub_agent_title() + ) + conversation_store.set_session_usage(parent.id, {"total_cost_usd": 0.10}) + conversation_store.set_session_usage(child.id, {"total_cost_usd": 0.05}) + conversation_store.set_session_state(parent.id, {SESSION_COST_ASK_APPROVED_STATE_KEY: 9.99}) + conversation_store.update_conversation(parent.id, archived=True) + + engine = build_policy_engine( + spec=AgentSpec(spec_version=1, name="child"), + conversation_id=child.id, + conversation_store=conversation_store, + ) + + # Whole-tree total including the archived root's own spend. + assert engine.usage["total_cost_usd"] == pytest.approx(0.15) + assert engine.session_state[SESSION_COST_ASK_APPROVED_STATE_KEY] == pytest.approx(9.99) + + +def test_build_rejects_mismatched_preloaded_conversation( + conversation_store: SqlAlchemyConversationStore, +) -> None: + """A preloaded row for a different session must fail closed, not mix + one session's labels/state/usage into another's policy decision.""" + from omnigent.errors import OmnigentError + + a = conversation_store.create_conversation(title="a") + b = conversation_store.create_conversation(title="b") + row_b = conversation_store.get_conversation(b.id) + + with pytest.raises(OmnigentError, match="does not match"): + build_policy_engine( + spec=AgentSpec(spec_version=1, name="x"), + conversation_id=a.id, + conversation_store=conversation_store, + conversation=row_b, + ) + + +def test_build_with_preloaded_row_sees_concurrent_mutable_writes( + conversation_store: SqlAlchemyConversationStore, +) -> None: + """ + The preloaded row contributes only immutable identity: labels, + session_state, and model written AFTER the row was captured must + still reach the engine (re-derived from the fresh tree row), so a + concurrent guard-label write can't be skipped by a stale snapshot. + """ + conv = conversation_store.create_conversation(title="fresh-check") + stale_row = conversation_store.get_conversation(conv.id) + + # Writes that land between the handler's read and the engine build. + conversation_store.set_labels(conv.id, {"guard": "tripped"}) + conversation_store.set_session_state(conv.id, {"checkpoint": 1.5}) + conversation_store.update_conversation(conv.id, model_override="claude-sonnet-4-6") + + engine = build_policy_engine( + spec=AgentSpec(spec_version=1, name="x"), + conversation_id=conv.id, + conversation_store=conversation_store, + conversation=stale_row, + ) + + assert engine.labels.get("guard") == "tripped" + assert engine.session_state.get("checkpoint") == 1.5 + assert engine.model == "claude-sonnet-4-6" + + +def test_build_uses_fresh_state_for_a_row_archived_after_preload( + conversation_store: SqlAlchemyConversationStore, +) -> None: + """ + A row archived after the preload still authorizes from FRESH state. + + The tree now includes archived rows, so the fresh row is found there + and the preload contributes only identity. (This test previously + described the tree as EXCLUDING archived rows and asserted the + re-read branch — a premise the ``include_archived=True`` fix + inverted, which made the test unable to fail for its stated reason.) + """ + conv = conversation_store.create_conversation(title="archive-race") + stale_row = conversation_store.get_conversation(conv.id) + + conversation_store.set_labels(conv.id, {"guard": "tripped"}) + conversation_store.update_conversation( + conv.id, model_override="claude-opus-4-8", archived=True + ) + + engine = build_policy_engine( + spec=AgentSpec(spec_version=1, name="x"), + conversation_id=conv.id, + conversation_store=conversation_store, + conversation=stale_row, + ) + + # Mutable fields come from the fresh (archived) row, not the preload. + assert engine.model == "claude-opus-4-8" + assert engine.labels.get("guard") == "tripped" + + +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_initial_label_seed_uses_the_atomic_store_operation( + conversation_store: SqlAlchemyConversationStore, +) -> None: + """ + Seeding must go through ``seed_labels_if_absent``, never ``set_labels``. + + A mechanism assertion on top of the semantic one above: the race is + only impossible while the decision about which keys are missing is made + inside the insert statement, so the choice of store operation is itself + part of the contract. + """ + from omnigent.runtime.policies.builder import _seed_and_load_labels + from omnigent.spec.types import LabelDef + + conv = conversation_store.create_conversation() + calls: list[str] = [] + + class _RecordingStore: + def __init__(self, inner: SqlAlchemyConversationStore) -> None: + self._inner = inner + + def __getattr__(self, name: str): + attr = getattr(self._inner, name) + if not callable(attr): + return attr + + def _wrapper(*args: object, **kwargs: object) -> object: + calls.append(name) + return attr(*args, **kwargs) + + return _wrapper + + _seed_and_load_labels( + conversation_id=conv.id, + label_defs={"integrity": LabelDef(initial="0")}, + conversation_store=_RecordingStore(conversation_store), # type: ignore[arg-type] + existing={}, + ) + + assert "seed_labels_if_absent" in calls, calls + assert "set_labels" not in calls, calls + + +def test_deleted_after_preload_fails_closed( + conversation_store: SqlAlchemyConversationStore, +) -> None: + """A row deleted after the preload must fail closed, not authorize + from the stale snapshot nor from empty ($0) state.""" + import asyncio + + from omnigent.errors import OmnigentError + + conv = conversation_store.create_conversation(title="deleted") + conversation_store.set_session_usage(conv.id, {"total_cost_usd": 5.0}) + stale_row = conversation_store.get_conversation(conv.id) + asyncio.run(conversation_store.delete_conversation(conv.id)) + + with pytest.raises(OmnigentError, match="disappeared"): + build_policy_engine( + spec=AgentSpec(spec_version=1, name="x"), + conversation_id=conv.id, + conversation_store=conversation_store, + conversation=stale_row, + ) + + +def test_agent_rebind_after_spec_resolution_fails_closed( + conversation_store: SqlAlchemyConversationStore, +) -> None: + """ + ``agent_id`` selects the spec, so it is read before the engine exists + and cannot be re-derived. The builder confirms it against the row it + freshly read — NOT against the caller's preload, which is the snapshot + whose staleness is the whole hazard. + + Four cases, all of which an earlier revision accepted because the + comparison ran before the fresh read (and skipped on ``None``): + + 1. rebind with no preload + 2. rebind WITH a stale preload naming the old agent <- the reproduction + 3. fresh row whose binding is ``None`` + 4. no fresh row at all (deleted) + """ + import asyncio + + from omnigent.errors import OmnigentError + + spec = AgentSpec(spec_version=1, name="x") + agent_a = "1" * 32 + + def _switch(conv_id: str, new_agent_id: str) -> None: + # switch_conversation_agent inserts the target agent row, so each + # switch needs its own id. + conversation_store.switch_conversation_agent( + conv_id, + new_agent_id=new_agent_id, + new_agent_name="other", + new_agent_bundle_location="other/bundle", + new_agent_description=None, + copy_model_settings=False, + carry_history_into_native=False, + presentation_labels={}, + previous_builtin_id=None, + ) + + # Matching agent proceeds (guard must not be a blanket refusal). + ok_conv = conversation_store.create_conversation(title="match", agent_id=agent_a) + assert build_policy_engine( + spec=spec, + conversation_id=ok_conv.id, + conversation_store=conversation_store, + expected_agent_id=agent_a, + ) + + # 1. Rebind, no preload. + c1 = conversation_store.create_conversation(title="rebind-nopreload", agent_id=agent_a) + _switch(c1.id, uuid.uuid4().hex) + with pytest.raises(OmnigentError, match="no longer resolves to agent"): + build_policy_engine( + spec=spec, + conversation_id=c1.id, + conversation_store=conversation_store, + expected_agent_id=agent_a, + ) + + # 2. Rebind WITH a stale preload still naming agent A. Comparing against + # that preload would have found agent A == expected and accepted. + c2 = conversation_store.create_conversation(title="rebind-preload", agent_id=agent_a) + stale_row = conversation_store.get_conversation(c2.id) + assert stale_row.agent_id == agent_a + _switch(c2.id, uuid.uuid4().hex) + with pytest.raises(OmnigentError, match="no longer resolves to agent"): + build_policy_engine( + spec=spec, + conversation_id=c2.id, + conversation_store=conversation_store, + conversation=stale_row, + expected_agent_id=agent_a, + ) + + # 3. Fresh row with NO binding: previously skipped by an ``is not None`` + # guard, so an unbound session accepted any spec. + c3 = conversation_store.create_conversation(title="unbound") + assert conversation_store.get_conversation(c3.id).agent_id is None + with pytest.raises(OmnigentError, match="no longer resolves to agent"): + build_policy_engine( + spec=spec, + conversation_id=c3.id, + conversation_store=conversation_store, + expected_agent_id=agent_a, + ) + + # 4. No fresh row at all (deleted, no preload): also a mismatch. + c4 = conversation_store.create_conversation(title="gone", agent_id=agent_a) + asyncio.run(conversation_store.delete_conversation(c4.id)) + with pytest.raises(OmnigentError, match="no longer resolves to agent"): + build_policy_engine( + spec=spec, + conversation_id=c4.id, + conversation_store=conversation_store, + expected_agent_id=agent_a, + ) + + +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 + + +def test_supplied_root_is_a_hint_that_gets_verified( + conversation_store: SqlAlchemyConversationStore, +) -> None: + """ + A caller's tree root saves a read when right and is corrected when wrong. + + ``root_conversation_id`` is immutable per ROW, not per conversation id. + Delete a conversation and recreate it under the same id and the new row + can sit in a different tree, so a row read earlier in the request names a + root this conversation no longer belongs to. Trusting it silently summed + the wrong tree — the recreated child reported nothing at all. + + Three properties, because the argument has to be observably load-bearing + AND observably safe: + + 1. right root → same answer as self-resolving, one conversation read + fewer (a build that ignored the argument fails on the read count); + 2. wrong tree's root → still the right answer (a build that trusts it + fails here); + 3. the delete/recreate case that produced the regression. + """ + from omnigent.runtime.policies.builder import load_session_usage + + parent = conversation_store.create_conversation() + child = conversation_store.create_conversation( + kind="sub_agent", parent_conversation_id=parent.id, title=_sub_agent_title() + ) + conversation_store.set_session_usage(parent.id, {"total_cost_usd": 0.10}) + conversation_store.set_session_usage(child.id, {"total_cost_usd": 0.05}) + other_tree = conversation_store.create_conversation(title="unrelated") + + with _count_sql(conversation_store) as self_resolved_sql: + resolved = load_session_usage(child.id, conversation_store) + assert resolved["total_cost_usd"] == pytest.approx(0.05) + + with _count_sql(conversation_store) as supplied_sql: + supplied = load_session_usage(child.id, conversation_store, root_conversation_id=parent.id) + assert supplied == resolved + # The saved read is the point of the argument: a triplet fewer. + executed = lambda caught: [q for q in caught if not q.startswith("PRAGMA")] # noqa: E731 + assert len(executed(supplied_sql)) == len(executed(self_resolved_sql)) - 3, ( + executed(supplied_sql), + executed(self_resolved_sql), + ) + + # Wrong tree's root: verified against the tree it produced, so the sum is + # still right. Trusting it returned {} — the shape of the regression. + assert ( + load_session_usage(child.id, conversation_store, root_conversation_id=other_tree.id) + == resolved + ) + + +def test_recreated_conversation_is_summed_in_its_new_tree( + conversation_store: SqlAlchemyConversationStore, +) -> None: + """ + A row read before a delete/recreate must not decide which tree to sum. + + The reported regression: a child is deleted and recreated under the same + id in a different tree while a caller holds the old row. Summing from the + stale row's root found no such conversation and emitted nothing, so the + session's cost silently stopped updating. + """ + import asyncio + + from omnigent.runtime.policies.builder import load_session_usage + + old_parent = conversation_store.create_conversation(title="old-root") + child = conversation_store.create_conversation( + kind="sub_agent", parent_conversation_id=old_parent.id, title=_sub_agent_title() + ) + stale_row = conversation_store.get_conversation(child.id) + assert stale_row.root_conversation_id == old_parent.root_conversation_id + + # Recreate the same id under a different tree, as the delete/recreate + # boundary does. + asyncio.run(conversation_store.delete_conversation(child.id)) + new_parent = conversation_store.create_conversation(title="new-root") + conversation_store.create_conversation( + conversation_id=child.id, + kind="sub_agent", + parent_conversation_id=new_parent.id, + title=_sub_agent_title(), + ) + conversation_store.set_session_usage(child.id, {"total_cost_usd": 2.0}) + + summed = load_session_usage( + child.id, + conversation_store, + root_conversation_id=stale_row.root_conversation_id, + ) + assert summed["total_cost_usd"] == pytest.approx(2.0), ( + "a stale root must not silence a recreated conversation's spend" + ) + + +class _MutateOnTreeLoad: + """Store proxy that commits *hazard* the first time the tree is scanned. + + The builder reads the conversation, then scans the spawn tree. Anything + committed in between makes the earlier read stale — which is the window + the freshness refresh exists to close, whoever performed that read. + """ + + def __init__(self, inner: object, hazard: Callable[[], None]) -> None: + self._inner = inner + self._hazard = hazard + self._fired = False + + def list_conversations(self, *args: object, **kwargs: object) -> object: + if not self._fired: + self._fired = True + self._hazard() + return self._inner.list_conversations(*args, **kwargs) # type: ignore[attr-defined] + + def __getattr__(self, name: str) -> object: + return getattr(self._inner, name) + + +@pytest.mark.parametrize("preloaded", [True, False], ids=["preload", "no-preload"]) +@pytest.mark.parametrize("hazard", ["switch", "delete"]) +def test_mid_build_change_fails_closed_on_every_provenance( + conversation_store: SqlAlchemyConversationStore, + preloaded: bool, + hazard: str, +) -> None: + """ + A change landing between the conversation read and the tree scan must + fail closed — regardless of who performed that read. + + Parametrized over provenance on purpose. An earlier revision gated the + refresh on ``conversation is not None``, so the caller-preload path was + closed and the builder's own read had the identical window wide open: a + switch enforced the old agent's spec, a deletion seeded empty usage and + authorized a $0 budget. Provenance is a row here, not a separate test, + so a third way of acquiring the row is covered by construction. + """ + import asyncio + + from omnigent.errors import OmnigentError + + agent_a = "2" * 32 + conv = conversation_store.create_conversation(title=f"mid-build-{hazard}", agent_id=agent_a) + conversation_store.set_session_usage(conv.id, {"total_cost_usd": 5.0}) + row = conversation_store.get_conversation(conv.id) + + if hazard == "switch": + + def _apply() -> None: + conversation_store.switch_conversation_agent( + conv.id, + new_agent_id=uuid.uuid4().hex, + new_agent_name="other", + new_agent_bundle_location="other/bundle", + new_agent_description=None, + copy_model_settings=False, + carry_history_into_native=False, + presentation_labels={}, + previous_builtin_id=None, + ) + + expected = "no longer resolves to agent" + else: + + def _apply() -> None: + asyncio.run(conversation_store.delete_conversation(conv.id)) + + expected = "disappeared" + + store = _MutateOnTreeLoad(conversation_store, _apply) + with pytest.raises(OmnigentError, match=expected): + build_policy_engine( + spec=AgentSpec(spec_version=1, name="x"), + conversation_id=conv.id, + conversation_store=store, # type: ignore[arg-type] + conversation=row if preloaded else None, + expected_agent_id=agent_a, + ) + + +def test_archived_descendant_spend_counts_toward_the_displayed_total( + conversation_store: SqlAlchemyConversationStore, +) -> None: + """ + Including archived rows in the tree changes DISPLAY as well as gating. + + The change was made so archiving cannot reset a cost gate, but + ``load_session_usage`` is also the display path for the session badge + (``session.usage`` SSE) and the usage report, so an archived + descendant's spend now shows in the total a user sees. That is the + deliberate choice — the alternative, separate tree semantics for + display and enforcement, lets the badge disagree with the gate that + blocks the user — and it belongs with the change that caused it rather + than in a later PR. + """ + from omnigent.runtime.policies.builder import load_session_usage + + parent = conversation_store.create_conversation() + child = conversation_store.create_conversation( + kind="sub_agent", parent_conversation_id=parent.id, title=_sub_agent_title() + ) + conversation_store.set_session_usage(parent.id, {"total_cost_usd": 1.0}) + conversation_store.set_session_usage(child.id, {"total_cost_usd": 2.0}) + conversation_store.update_conversation(child.id, archived=True) + + displayed = load_session_usage(parent.id, conversation_store) + assert displayed["total_cost_usd"] == pytest.approx(3.0), ( + "archived descendant spend must reach the displayed subtree total" + ) + + # The same number the badge and the enforcement seed derive from, so the + # two cannot diverge. + engine = build_policy_engine( + spec=AgentSpec(spec_version=1, name="parent"), + conversation_id=parent.id, + conversation_store=conversation_store, + ) + assert engine.usage["total_cost_usd"] == pytest.approx(3.0) + + +def test_engine_seeds_the_new_tree_when_a_child_is_recreated_elsewhere( + conversation_store: SqlAlchemyConversationStore, +) -> None: + """ + Everything the build derives must come from the verified tree. + + The preload only *suggests* a root. Deriving the tree root from the + pre-refresh row while taking rows from a corrected tree mixed two + epochs: a child deleted and recreated beneath another root, with a + caller still holding the old row, seeded the OLD tree's spend — the + refresh landed on the row and left the root and the tree pointing at + the previous tree. (Root-inherited session *policies* take the same + refreshed root and are exposed to the same risk, but only spend is + asserted below — this test does not exercise policy inheritance.) + """ + import asyncio + + old_root = conversation_store.create_conversation(title="old-root") + child = conversation_store.create_conversation( + kind="sub_agent", parent_conversation_id=old_root.id, title=_sub_agent_title() + ) + conversation_store.set_session_usage(old_root.id, {"total_cost_usd": 10.0}) + stale_row = conversation_store.get_conversation(child.id) + + asyncio.run(conversation_store.delete_conversation(child.id)) + new_root = conversation_store.create_conversation(title="new-root") + conversation_store.set_session_usage(new_root.id, {"total_cost_usd": 3.0}) + conversation_store.create_conversation( + conversation_id=child.id, + kind="sub_agent", + parent_conversation_id=new_root.id, + title=_sub_agent_title(), + ) + + engine = build_policy_engine( + spec=AgentSpec(spec_version=1, name="x"), + conversation_id=child.id, + conversation_store=conversation_store, + conversation=stale_row, + ) + + # The gate seeds from the whole tree the child is in NOW. + assert engine.usage["total_cost_usd"] == pytest.approx(3.0), ( + "the engine seeded a tree the conversation no longer belongs to" + ) + + +@pytest.mark.parametrize("preload", [True, False], ids=["preloaded", "no-preload"]) +@pytest.mark.parametrize("mutation", ["switch", "delete"]) +def test_engine_refuses_a_tree_assembled_across_a_change( + conversation_store: SqlAlchemyConversationStore, + preload: bool, + mutation: str, +) -> None: + """ + A paged tree does not share one read instant, so it cannot vouch for + itself. + + Rows on page one are read before page two. The refresh takes the + evaluated row from the tree, so for a paged tree that row may predate + the last page — "the tree read is newer" is true of the tree, not of any + row in it. A switch OR a delete landing during paging must be caught + the same way, whether or not the caller preloaded the row — all four + preload/no-preload x switch/delete combinations previously proceeded. + """ + import asyncio + + from omnigent.errors import OmnigentError + from omnigent.runtime.policies import builder as builder_mod + + root = conversation_store.create_conversation(title="paged-root") + child = conversation_store.create_conversation( + kind="sub_agent", parent_conversation_id=root.id, title=_sub_agent_title() + ) + row = conversation_store.get_conversation(child.id) + + class _MutateDuringPaging: + """Switches or deletes the evaluated conversation between tree pages.""" + + def __init__(self, inner: SqlAlchemyConversationStore) -> None: + self._inner = inner + self._pages = 0 + self._fired = False + + def list_conversations(self, *args: object, **kwargs: object): + page = self._inner.list_conversations(*args, **kwargs) # type: ignore[arg-type] + self._pages += 1 + # Fire once the page holding the evaluated row has been read — + # which page that is depends on listing order, so keying on "page + # one" made the interleave dialect-dependent. + if not self._fired and any(c.id == child.id for c in page.data): + self._fired = True + if mutation == "switch": + self._inner.switch_conversation_agent( + child.id, + new_agent_id=uuid.uuid4().hex, + new_agent_name="other", + new_agent_bundle_location="other/bundle", + new_agent_description=None, + copy_model_settings=False, + carry_history_into_native=False, + presentation_labels={}, + previous_builtin_id=None, + ) + else: + asyncio.run(self._inner.delete_conversation(child.id)) + return page + + def __getattr__(self, name: str): + return getattr(self._inner, name) + + # One row per page, so the tree genuinely pages. + original_page_size = builder_mod._SUBTREE_USAGE_PAGE_SIZE + builder_mod._SUBTREE_USAGE_PAGE_SIZE = 1 + try: + with pytest.raises(OmnigentError, match="moved while its spawn tree"): + build_policy_engine( + spec=AgentSpec(spec_version=1, name="x"), + conversation_id=child.id, + conversation_store=_MutateDuringPaging(conversation_store), # type: ignore[arg-type] + conversation=row if preload else None, + ) + finally: + builder_mod._SUBTREE_USAGE_PAGE_SIZE = original_page_size + + +@pytest.mark.parametrize( + ("shape", "links"), + [ + ("cycle", {"a": "b", "b": "a"}), + ("parent outside the tree", {"a": "missing"}), + ], +) +def test_ancestor_walk_discards_an_untrustworthy_chain(shape: str, links: dict[str, str]) -> None: + """ + A chain that cannot be walked to the root yields nothing, not a prefix. + + The docstring promised empty for a cyclic or broken chain while the loop + returned the part it had already walked — and the caller publishes each + returned id a cost event, so `A → B → A` notified B off a cycle and + `C → missing-D` notified an id that is not in the tree at all. + """ + from omnigent.runtime.policies.builder import ancestor_ids_from_tree + + tree = [ + Conversation( + id=node, + created_at=0, + updated_at=0, + root_conversation_id="a", + parent_conversation_id=parent, + ) + for node, parent in links.items() + ] + + assert ancestor_ids_from_tree(tree, "a") == [], shape diff --git a/tests/runtime/policies/test_session_cost_ask_routing.py b/tests/runtime/policies/test_session_cost_ask_routing.py index 581c2d2d4f..ca01e30c29 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() 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/helpers.py b/tests/server/helpers.py index 2ccdd0d1c7..e96c1456db 100644 --- a/tests/server/helpers.py +++ b/tests/server/helpers.py @@ -8,7 +8,8 @@ import re import tarfile import threading -from collections.abc import Callable +from collections.abc import Callable, Iterator +from contextlib import contextmanager from dataclasses import dataclass from typing import Any, ClassVar @@ -937,3 +938,30 @@ async def post( if body.get("type") == "cost_approval_popup": self.popup_seen.set() return httpx.Response(202, request=httpx.Request("POST", f"http://runner{url}")) + + +@contextmanager +def count_conversation_selects(engine: Any) -> Iterator[list[str]]: + """Collect the conversations-table SELECTs a request emits. + + Shared by the snapshot and usage-report read-count regressions, which + live in the files owning those endpoints. + """ + from sqlalchemy import event + + seen: list[str] = [] + + def _on_exec(conn, cursor, statement, parameters, context, executemany): + if ( + statement.lstrip().startswith("SELECT") + and "FROM conversations" in statement + and "conversation_items" not in statement + and "conversation_labels" not in statement + ): + seen.append(statement) + + event.listen(engine, "before_cursor_execute", _on_exec) + try: + yield seen + finally: + event.remove(engine, "before_cursor_execute", _on_exec) diff --git a/tests/server/integration/test_sessions_endpoints.py b/tests/server/integration/test_sessions_endpoints.py index 09b0a140b6..4439d860b4 100644 --- a/tests/server/integration/test_sessions_endpoints.py +++ b/tests/server/integration/test_sessions_endpoints.py @@ -4918,6 +4918,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/server/integration/test_sessions_policy_evaluate.py b/tests/server/integration/test_sessions_policy_evaluate.py index 2d4f44f457..8253718fee 100644 --- a/tests/server/integration/test_sessions_policy_evaluate.py +++ b/tests/server/integration/test_sessions_policy_evaluate.py @@ -24,10 +24,13 @@ from __future__ import annotations import asyncio +from pathlib import Path from typing import Any import httpx import pytest +import pytest_asyncio +from fastapi import FastAPI from omnigent.runtime import get_caps, session_stream from omnigent.runtime.caps import RuntimeCaps @@ -1162,3 +1165,170 @@ async def test_llm_response_allow_when_no_matching_policy( assert body["result"] in ("POLICY_ACTION_ALLOW", "POLICY_ACTION_UNSPECIFIED"), ( f"Expected ALLOW or UNSPECIFIED for no-policy session, got {body['result']}." ) + + +_EVALUATE_USER = "alice@example.com" + + +@pytest.fixture() +def auth_app(runtime_init: None, db_uri: str, tmp_path: Path) -> FastAPI: + """App with permissions + header auth, for route-level SQL budgeting.""" + from omnigent.runtime.agent_cache import AgentCache + from omnigent.server.app import create_app + from omnigent.server.auth import UnifiedAuthProvider + from omnigent.stores.agent_store.sqlalchemy_store import SqlAlchemyAgentStore + from omnigent.stores.artifact_store.local import LocalArtifactStore + from omnigent.stores.comment_store.sqlalchemy_store import SqlAlchemyCommentStore + from omnigent.stores.file_store.sqlalchemy_store import SqlAlchemyFileStore + from omnigent.stores.permission_store.sqlalchemy_store import SqlAlchemyPermissionStore + + artifact_store = LocalArtifactStore(str(tmp_path / "artifacts")) + return create_app( + agent_store=SqlAlchemyAgentStore(db_uri), + file_store=SqlAlchemyFileStore(db_uri), + conversation_store=SqlAlchemyConversationStore(db_uri), + artifact_store=artifact_store, + agent_cache=AgentCache(artifact_store=artifact_store, cache_dir=tmp_path / "cache"), + comment_store=SqlAlchemyCommentStore(db_uri), + permission_store=SqlAlchemyPermissionStore(db_uri), + auth_provider=UnifiedAuthProvider(source="header"), + ) + + +@pytest_asyncio.fixture() +async def auth_client(auth_app: FastAPI, mock_llm: Any, tmp_path: Path): + """Async client against the auth-enabled app.""" + from omnigent.runtime import set_harness_process_manager + from omnigent.runtime.harnesses.process_manager import HarnessProcessManager + + pm = HarnessProcessManager(tmp_parent=tmp_path / "harness_pm") + await pm.start() + set_harness_process_manager(pm) + transport = httpx.ASGITransport(app=auth_app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as c: + yield c + mock_llm.release_all() + set_harness_process_manager(None) + await pm.shutdown() + + +async def _seed_authenticated_session(auth_client: httpx.AsyncClient, db_uri: str) -> str: + """Create a real agent-bound session and grant it to the test user. + + Uses the real create path (a genuine uploaded bundle) so the agent + cache can load the spec — a fabricated agent row with a fake bundle + location cannot be loaded and the route would bail before the work + this budget measures. + """ + import json as _json + + from omnigent.server.auth import LEVEL_OWNER + from omnigent.stores.permission_store.sqlalchemy_store import SqlAlchemyPermissionStore + from tests.server.helpers import build_agent_bundle + + perms = SqlAlchemyPermissionStore(db_uri) + perms.ensure_user(_EVALUATE_USER) + headers = {"X-Forwarded-Email": _EVALUATE_USER} + resp = await auth_client.post( + "/v1/sessions", + data={"metadata": _json.dumps({"title": "budget"})}, + files={ + "bundle": ( + "agent.tar.gz", + # The agent MUST declare a policy: with no policies the route + # short-circuits via ``any_policies_apply`` before building an + # engine, so a budget measured without one would not cover the + # preload path at all (it silently didn't, first time round). + build_agent_bundle( + name="budget-agent", + guardrails={ + "policies": { + "allow_all": { + "type": "function", + "on": ["tool_call"], + "function": "tests.runtime.policies.conftest._always_allow", + } + } + }, + ), + "application/gzip", + ) + }, + headers=headers, + ) + assert resp.status_code == 201, resp.text + session_id = resp.json()["session_id"] + perms.grant(_EVALUATE_USER, session_id, LEVEL_OWNER) + return session_id + + +# Measured SQL-statement budget for the AUTHENTICATED evaluate route — the +# blocking PreToolUse hook. Counted as executed statements, not store-method +# calls: a store-call oracle cannot see that one call fans out to three +# statements (conversation + metadata + labels), which is exactly how an +# earlier "6 queries" claim survived being wrong by five. +# +# Measured (both dialects) for the seed below: one owner grant, one agent +# declaring one tool_call policy, one conversation, no declared labels, no +# children. Composition: ACL resolution 3, the handler's conversation load +# 3, session-policy lookup, agent row, spawn-tree scan 3. +# +# This is the number the PR must quote — an earlier description claimed 6. +# It is also the oracle for the preload optimisation: dropping the +# ``conversation=`` argument makes the builder re-read the conversation and +# pushes this up, which no behavioural assertion can see because both +# variants return the same verdict. +_EVALUATE_ROUTE_SQL_BUDGET = 11 + + +@pytest.mark.asyncio +async def test_authenticated_evaluate_route_sql_budget( + auth_client: httpx.AsyncClient, + db_uri: str, +) -> None: + """ + Pin the route's SQL count so the description cannot drift from the code, + and so removing the preload is detectable. + + This is the oracle for the preload optimisation: dropping the + ``conversation=`` argument in ``_build_engine`` makes the builder read + the conversation again, which changes this count. A behavioural + assertion alone cannot see that, because both variants return the same + verdict. + """ + from sqlalchemy import event as sa_event + + from omnigent.db.utils import _engine_cache + + session_id = await _seed_authenticated_session(auth_client, db_uri) + headers = {"X-Forwarded-Email": _EVALUATE_USER} + payload = { + "event": { + "type": "PHASE_TOOL_CALL", + "data": {"name": "Bash", "arguments": {"command": "ls"}}, + } + } + # Warm the spec / policy caches so the count is the steady-state one. + await auth_client.post( + f"/v1/sessions/{session_id}/policies/evaluate", json=payload, headers=headers + ) + + engine = _engine_cache[db_uri] + statements: list[str] = [] + + def _on_exec(conn, cursor, statement, parameters, context, executemany): + if not statement.lstrip().upper().startswith("PRAGMA"): + statements.append(statement) + + sa_event.listen(engine, "before_cursor_execute", _on_exec) + try: + resp = await auth_client.post( + f"/v1/sessions/{session_id}/policies/evaluate", json=payload, headers=headers + ) + finally: + sa_event.remove(engine, "before_cursor_execute", _on_exec) + + assert resp.status_code == 200, resp.text + assert len(statements) == _EVALUATE_ROUTE_SQL_BUDGET, [ + s.split("\n")[0][:80] for s in statements + ] diff --git a/tests/server/routes/test_event_path_load_once.py b/tests/server/routes/test_event_path_load_once.py new file mode 100644 index 0000000000..752879d145 --- /dev/null +++ b/tests/server/routes/test_event_path_load_once.py @@ -0,0 +1,328 @@ +"""Regression tests: the event hot path loads the conversation once. + +POST /v1/sessions/{id}/events fires per streamed chunk during active +turns (~98% of server query volume), so redundant conversation/ACL +reloads multiply hard. These pin two things at once, because a +query-count ceiling on its own would stay green if the work it is +counting disappeared entirely: + +1. the per-table query counts for an event, and +2. that the event's real effect still happens — the status forward + reaches the bound runner, the usage recompute publishes the right + subtree total. + +The routing read is deliberately NOT threaded from the handler's row: +runner resolution takes its own fresh single-column binding read so a +concurrent rebind is picked up. The rebind semantics themselves are +pinned in ``tests/runner/test_routing.py``, because the app builds its +own ``RunnerRouter`` over a real tunnel registry with no injection seam — +what IS testable here, and asserted below, is that every event issues its +own binding read rather than reusing a cached resolution. +""" + +from __future__ import annotations + +import re +from collections import Counter +from contextlib import contextmanager +from typing import Any + +import httpx +import pytest +from sqlalchemy import event + +from omnigent.db.utils import _engine_cache +from omnigent.runtime import session_stream, set_runner_client, set_runner_router +from omnigent.stores.conversation_store.sqlalchemy_store import ( + SqlAlchemyConversationStore, +) + +# Attribute every statement to its table so the counter cannot miss the +# metadata / labels / UPDATE traffic a conversations-only filter ignored. +_TABLE_RE = re.compile( + r"\bFROM\s+([a-z_][a-z0-9_]*)|\bINSERT\s+INTO\s+([a-z_][a-z0-9_]*)|\bUPDATE\s+([a-z_][a-z0-9_]*)", + re.IGNORECASE, +) + + +class _QueryCounts(Counter): + """Per-table statement counts, retaining the statements themselves.""" + + statements: list[str] + + +@contextmanager +def _count_queries_by_table(engine: Any): + """Count every non-PRAGMA statement, keyed by the table it touches.""" + counts: _QueryCounts = _QueryCounts() + counts.statements = [] + + def _on_exec(conn, cursor, statement, parameters, context, executemany): + if statement.lstrip().upper().startswith("PRAGMA"): + return + counts.statements.append(statement) + match = _TABLE_RE.search(statement) + table = next((g for g in match.groups() if g), "other") if match else "other" + counts[table.lower()] += 1 + + event.listen(engine, "before_cursor_execute", _on_exec) + try: + yield counts + finally: + event.remove(engine, "before_cursor_execute", _on_exec) + + +@contextmanager +def _record_published(session_id: str): + """Capture events published to *session_id*'s stream.""" + captured: list[dict[str, Any]] = [] + real_publish = session_stream.publish + + def _spy(conv_id: str, payload: dict[str, Any]) -> None: + if conv_id == session_id: + captured.append(payload) + real_publish(conv_id, payload) + + session_stream.publish = _spy # type: ignore[assignment] + try: + yield captured + finally: + session_stream.publish = real_publish # type: ignore[assignment] + + +class _RecordingRunnerClient: + """Minimal runner client that records the events posted to it.""" + + def __init__(self) -> None: + self.posts: list[tuple[str, Any]] = [] + + async def post( + self, + url: str, + *, + json: Any = None, + timeout: float | None = None, + ) -> httpx.Response: + del timeout + self.posts.append((url, json)) + return httpx.Response(200, json={}, request=httpx.Request("POST", url)) + + +class _ExplodingRunnerClient: + """Process-wide fallback client that must never be used. + + Delivery falling back to this client is what let the routed-binding + assertion pass for four review rounds. + """ + + async def post(self, url: str, **kwargs: Any) -> httpx.Response: + raise AssertionError( + f"event was delivered through the fallback client, not the bound runner: POST {url}" + ) + + +@pytest.fixture() +def runner_router_reset(): + """Install/remove runner globals around a test.""" + yield + set_runner_router(None) + set_runner_client(None) + + +@pytest.mark.asyncio +async def test_status_event_forwards_and_reads_each_table_once( + app: Any, + client: httpx.AsyncClient, + db_uri: str, + monkeypatch: pytest.MonkeyPatch, + runner_router_reset: None, +) -> None: + """ + A status event reaches the runner it is BOUND to, and reads each + conversation table once — not once per routing lookup on top of the + handler's own read. + + The event is delivered through a real :class:`RunnerRouter` over a real + registry, so the binding read under test actually selects the client + that receives the POST. Earlier versions of this test bound + ``runner_one`` without registering it: routing reported it offline, + delivery silently fell back to the process-wide client, and the + assertion passed while observing nothing about routing. The fallback + client is installed here as a trap that fails if it is ever used. + """ + from tests.server.helpers import register_test_runner + + store = SqlAlchemyConversationStore(db_uri) + conv = store.create_conversation(title="event-hot-path") + store.set_runner_id(conv.id, "runner_one") + + # Bind AND register: the app's own router resolves the binding through + # the read under test and then looks the runner up in the app's + # registry. Binding without registering is what made routing report the + # runner offline and delivery fall through to the global client. + register_test_runner(app, "runner_one") + runner = _RecordingRunnerClient() + routed_ids: list[str] = [] + + def _client_for_runner(runner_id: str) -> Any: + routed_ids.append(runner_id) + return runner + + # Real routing, stubbed transport: everything up to and including the + # binding lookup is production code; only the tunnel socket is replaced. + monkeypatch.setattr(app.state.runner_router, "_client_for_runner", _client_for_runner) + set_runner_client(_ExplodingRunnerClient()) # type: ignore[arg-type] + + payload = {"type": "external_session_status", "data": {"status": "idle"}} + await client.post(f"/v1/sessions/{conv.id}/events", json=payload) # warm + + engine = _engine_cache[db_uri] + baseline = len(runner.posts) + with _count_queries_by_table(engine) as counts: + resp = await client.post(f"/v1/sessions/{conv.id}/events", json=payload) + + assert resp.status_code == 202, resp.text + # Delivered to the BOUND runner, not merely delivered: the routed client + # is the one that received it, and the global fallback would have raised. + assert routed_ids and set(routed_ids) == {"runner_one"}, routed_ids + delivered = runner.posts[baseline:] + assert [body["type"] for _url, body in delivered] == ["external_session_status"], delivered + + # The handler's single conversation load (conversations + metadata + + # labels) plus routing's one narrow binding read — and nothing else. + # Before this change routing re-loaded the whole row, so metadata, + # conversations and labels were each read twice. + assert counts["conversation_labels"] == 1, dict(counts) + full_row_reads = [ + st + for st in counts.statements + if st.lstrip().upper().startswith("SELECT") + and "FROM conversations" in st + and "conversations.title" in st + ] + assert len(full_row_reads) == 1, full_row_reads + # Routing's read is the NARROW binding lookup: it projects the id and the + # runner binding only, never the row's other columns. + binding_reads = [ + st for st in counts.statements if "OUTER JOIN" in st.upper() and "runner_id" in st + ] + assert len(binding_reads) == 1, binding_reads + assert "conversations.title" not in binding_reads[0], binding_reads + + +@pytest.mark.asyncio +async def test_usage_event_publishes_subtree_total_with_exact_reads( + client: httpx.AsyncClient, + db_uri: str, +) -> None: + """ + The usage event recomputes the subtree total from the handler's row + (no root re-resolution) and publishes the correct value — a count + ceiling alone would pass if the recompute were deleted. + """ + store = SqlAlchemyConversationStore(db_uri) + conv = store.create_conversation(title="usage-hot-path") + child = store.create_conversation( + kind="sub_agent", parent_conversation_id=conv.id, title="usage:child" + ) + store.set_session_usage(child.id, {"total_cost_usd": 0.25}) + + payload = {"type": "external_session_usage", "data": {"cumulative_cost_usd": 0.5}} + await client.post(f"/v1/sessions/{conv.id}/events", json=payload) # warm + + engine = _engine_cache[db_uri] + with _record_published(conv.id) as published: + with _count_queries_by_table(engine) as counts: + resp = await client.post(f"/v1/sessions/{conv.id}/events", json=payload) + + assert resp.status_code == 202, resp.text + usage_events = [e for e in published if e.get("type") == "session.usage"] + assert usage_events, f"no session.usage published: {published}" + # Parent's own 0.5 plus the child's 0.25 — proves the subtree recompute + # ran and used the whole tree, not just this row. + assert usage_events[-1]["total_cost_usd"] == pytest.approx(0.75) + + # Per-table totals, not a ceiling on one table: a ceiling hid extra + # metadata and label statements entirely, and "<= 3" also passed with the + # anti-replay read removed. conversations = handler row, the anti-replay + # merge read (which must stay fresh) and the tree page scan; no root + # re-resolution in any of them. + assert dict(counts) == { + "conversations": 4, + "omnigent_conversation_metadata": 5, + "conversation_labels": 4, + }, dict(counts) + # The fourth conversations read is the ancestor fan-out loading its own + # tree at publication time. Reusing the tree summed above saved it and was + # wrong: two siblings then published from their own request-start + # snapshots, delivering a newer total followed by an older one. + + +@pytest.mark.asyncio +async def test_ancestor_publish_reads_the_tree_after_a_sibling_lands( + client: httpx.AsyncClient, + db_uri: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """ + A parent's badge must not go stale because it was computed from a tree + read at the START of the request that happened to publish it. + + Each usage event re-publishes every ancestor's subtree total. Taking that + total from the tree loaded at the start of the current request made the + figure a fact about request-arrival time, so a sibling's spend landing + mid-request was published as if it had not happened — two siblings + reporting around each other delivered a newer total and then an older one, + and the parent was left showing the smaller number. + + This pins that specific regression, not a universal "never moves + backwards" guarantee: a sibling's spend landing AFTER the ancestor + fan-out's own tree read but before publish can still be missed, exactly + as on unmodified `main` — that window is pre-existing and unrelated to + the request-start snapshot this PR removed. + + The interleave is constructed rather than hoped for: a sibling's spend is + committed immediately after this request's first tree load, which is + exactly the window. The fan-out reads its own tree afterwards, so the + published total includes it. + """ + from omnigent.server.routes._sessions import orchestration as orch + + store = SqlAlchemyConversationStore(db_uri) + parent = store.create_conversation(title="fanout-parent") + reporter = store.create_conversation( + kind="sub_agent", parent_conversation_id=parent.id, title="usage:reporter" + ) + sibling = store.create_conversation( + kind="sub_agent", parent_conversation_id=parent.id, title="usage:sibling" + ) + + payload = {"type": "external_session_usage", "data": {"cumulative_cost_usd": 1.0}} + await client.post(f"/v1/sessions/{reporter.id}/events", json=payload) # warm + + real_load = orch.load_session_tree + landed = {"done": False} + + def _load_then_land_sibling(*args: Any, **kwargs: Any): + tree = real_load(*args, **kwargs) + if not landed["done"]: + landed["done"] = True + # The sibling's spend commits here: after this request read a tree, + # before anything is published from it. + store.set_session_usage(sibling.id, {"total_cost_usd": 2.0}) + return tree + + monkeypatch.setattr(orch, "load_session_tree", _load_then_land_sibling) + + with _record_published(parent.id) as parent_events: + resp = await client.post(f"/v1/sessions/{reporter.id}/events", json=payload) + assert resp.status_code == 202, resp.text + assert landed["done"], "the interleave never happened; the test proves nothing" + + totals = [e["total_cost_usd"] for e in parent_events if e.get("type") == "session.usage"] + assert totals, f"the parent received no usage event: {parent_events}" + # reporter 1.0 + sibling 2.0, read after the sibling landed. + assert totals[-1] == pytest.approx(3.0), ( + f"the parent's published total {totals} was computed from a tree read " + f"before the sibling's spend landed" + ) diff --git a/tests/server/routes/test_sessions_mcp_proxy_policy_retry.py b/tests/server/routes/test_sessions_mcp_proxy_policy_retry.py index 67c9e10180..dbbca6614b 100644 --- a/tests/server/routes/test_sessions_mcp_proxy_policy_retry.py +++ b/tests/server/routes/test_sessions_mcp_proxy_policy_retry.py @@ -76,6 +76,29 @@ def get_conversation(self, session_id: str) -> Conversation | None: return None +def _engine_factory_expecting_preload(engine: object): + """Return a `_build_policy_engine_from_spec` double that requires the row. + + Every production caller of that wrapper has already loaded the + conversation to resolve the spec, and passes it so the builder can skip + its own read and confirm the agent binding. A double that merely + *accepts* ``conversation=None`` asserts nothing: dropping the preload at + the call site leaves every test here green, which is exactly what + happened. This one fails if the row does not arrive, so the contract is + observed by all seven tests rather than described by none. + """ + + def _factory(spec, session_id, conversation_store, conversation=None): + assert conversation is not None, ( + "the handler must pass the conversation row it already read; " + "without it the builder re-reads and loses the agent-binding check" + ) + assert conversation.id == session_id, (conversation.id, session_id) + return engine + + return _factory + + @dataclass class _StubAgentStore: """ @@ -209,7 +232,7 @@ async def test_forged_retry_with_deny_policy_is_rejected( monkeypatch.setattr( sessions_mod, "_build_policy_engine_from_spec", - lambda spec, session_id, conversation_store: deny_engine, + _engine_factory_expecting_preload(deny_engine), ) response = await _handle_mcp_tools_call( @@ -266,7 +289,7 @@ async def test_forged_retry_with_ask_policy_rejects_unknown_elicitation( monkeypatch.setattr( sessions_mod, "_build_policy_engine_from_spec", - lambda spec, session_id, conversation_store: ask_engine, + _engine_factory_expecting_preload(ask_engine), ) forged_eid = "elicit_FORGED_never_issued" @@ -318,7 +341,7 @@ async def test_retry_with_allow_policy_falls_through( monkeypatch.setattr( sessions_mod, "_build_policy_engine_from_spec", - lambda spec, session_id, conversation_store: allow_engine, + _engine_factory_expecting_preload(allow_engine), ) response = await _handle_mcp_tools_call( @@ -374,7 +397,7 @@ async def test_retry_session_mismatch_still_rejected( monkeypatch.setattr( sessions_mod, "_build_policy_engine_from_spec", - lambda spec, session_id, conversation_store: deny_engine, + _engine_factory_expecting_preload(deny_engine), ) params = _forged_retry_params() @@ -432,7 +455,7 @@ async def test_legitimate_retry_with_pending_entry_proceeds( monkeypatch.setattr( sessions_mod, "_build_policy_engine_from_spec", - lambda spec, session_id, conversation_store: ask_engine, + _engine_factory_expecting_preload(ask_engine), ) eid = "elicit_LEGITIMATE_server_issued" @@ -533,8 +556,8 @@ async def test_non_mcp_entry_popped_by_events_handler_on_accept( monkeypatch.setattr( sessions_mod, "_build_policy_engine_from_spec", - lambda spec, session_id, conversation_store: _FixedPolicyEngine( - result=PolicyResult(action=PolicyAction.ALLOW, reason=None) + _engine_factory_expecting_preload( + _FixedPolicyEngine(result=PolicyResult(action=PolicyAction.ALLOW, reason=None)) ), ) @@ -612,7 +635,7 @@ async def test_mcp_proxy_runner_supplied_actor_reaches_policy_engine( monkeypatch.setattr( sessions_mod, "_build_policy_engine_from_spec", - lambda spec, session_id, conversation_store: engine, + _engine_factory_expecting_preload(engine), ) params = {"name": "sys_os_shell", "arguments": {"command": "echo hi"}} diff --git a/tests/server/routes/test_sessions_runner_relay.py b/tests/server/routes/test_sessions_runner_relay.py index cc14fdcc31..46bf715d32 100644 --- a/tests/server/routes/test_sessions_runner_relay.py +++ b/tests/server/routes/test_sessions_runner_relay.py @@ -833,3 +833,251 @@ async def test_relay_running_edge_clears_stale_intentional_stop_marker() -> None sessions_module._runner_relay_tasks.clear() sessions_module._session_status_cache.pop(session_id, None) session_stream.close(session_id) + + +@pytest.mark.asyncio +async def test_relay_completion_with_usage_rolls_up_subtree_cost(db_uri: str) -> None: + """ + A relay turn reporting usage must price it, persist it, and publish the + SUBTREE total resolved from the TREE ROOT — not from the reporting + conversation. + + Deliberately runs the relay on a SUB-AGENT with the spend on a DESCENDANT, + so the published figure can only be right if the roll-up sums the tree + rather than the reporting conversation. + + What actually fails here, and why: rooting the tree at the reporter is + REPAIRED downstream — the tree loader verifies a supplied root and + resolves it again when the conversation is not in the tree it names — so + the cost stays correct and that mutation is caught by the PR owning that + contract, not here. This test owns the shared read: collapsing one + conversation read back into two is caught by the read count, which a + behavioural assertion cannot see. + """ + from contextlib import contextmanager + + from sqlalchemy import event as sa_event + + from omnigent.db.utils import _engine_cache + from omnigent.server.routes import sessions as sessions_module + + @contextmanager + def _count_conv_selects(engine): + seen: list[str] = [] + + def _on(conn, cursor, statement, params, context, many): + if ( + statement.lstrip().upper().startswith("SELECT") + and "FROM conversations" in statement + and "conversation_item" not in statement + and "conversation_label" not in statement + # PostgreSQL takes the usage row lock as its own id-only + # ``SELECT … FOR UPDATE``; SQLite gets the same serialisation + # from BEGIN IMMEDIATE and emits nothing. It is a lock, not a + # row read, so counting it would make this oracle assert a + # different number per dialect for no behavioural reason. + and "FOR UPDATE" not in statement.upper() + ): + seen.append(statement) + + sa_event.listen(engine, "before_cursor_execute", _on) + try: + yield seen + finally: + sa_event.remove(engine, "before_cursor_execute", _on) + + sessions_module._runner_relay_tasks.clear() + store = SqlAlchemyConversationStore(db_uri) + root = store.create_conversation() + reporter = store.create_conversation( + kind="sub_agent", parent_conversation_id=root.id, title="relay:reporter" + ) + descendant = store.create_conversation( + kind="sub_agent", parent_conversation_id=reporter.id, title="relay:grandchild" + ) + # Spend on a DESCENDANT of the reporter: inside the reporter's subtree, + # but only visible if the tree is loaded from the real root — its + # ``root_conversation_id`` is the root's, so a tree rooted at the + # reporter matches no rows at all. + store.set_session_usage(descendant.id, {"total_cost_usd": 0.25}) + session_id = reporter.id + + response_id = "resp_relay_usage_1" + turn_events: list[dict[str, Any]] = [ + { + "type": "response.in_progress", + "response": {"id": response_id, "model": "databricks-claude-sonnet-4-6"}, + }, + {"type": "response.output_text.delta", "delta": "hi"}, + { + "type": "response.completed", + "response": { + "id": response_id, + "model": "databricks-claude-sonnet-4-6", + "usage": { + "input_tokens": 1000, + "output_tokens": 500, + "total_tokens": 1500, + }, + }, + }, + ] + release = asyncio.Event() + fake_runner = _ScriptedRunnerClient(release, turn_events) + + collector = None + try: + handle = await sessions_module._ensure_runner_relay_ready( + session_id, + "runner_relay_usage", + fake_runner, # type: ignore[arg-type] + conversation_store=store, + ) + assert handle is not None + collector = await start_session_stream_collector(session_id) + + with _count_conv_selects(_engine_cache[db_uri]) as conv_selects: + release.set() + usage_events: list[dict[str, Any]] = [] + seen: list[str] = [] + while not seen or seen[-1] != "response.completed": + event = await collector.next_event() + seen.append(event["type"]) + if event["type"] == "session.usage": + usage_events.append(event) + + # The turn's own tokens landed on the reporting conversation. + reporter_usage = dict(store.get_conversation(session_id).session_usage) + assert reporter_usage.get("total_tokens") == 1500, reporter_usage + + assert usage_events, f"no session.usage published; saw {seen}" + published = usage_events[-1].get("total_cost_usd") + # The reporter's own turn is unpriced (test model is not in the + # pricing catalog), so the ONLY way a total appears is the tree-root + # scan reaching the descendant. Rooting the tree at the reporter + # matches no rows, so nothing would be published at all. + assert published == pytest.approx(0.25), (published, reporter_usage) + + # The invariant this PR owns, stated as "how many times is THIS + # session's row read", not as a total across every conversations + # SELECT the turn makes. A total also counts the tree scan and the + # ancestor walk, which belong to the PR below this one — so reverting + # that PR failed this oracle for a reason unrelated to the change + # under test. This count is identical with or without it. + # + # Two, and both are named: the relay's shared row (pricing + roll-up) + # and the relay-status path's own read, which stays separate on + # purpose because it must observe the newest labels. A full revert of + # the shared-row fix makes this four, not three: pricing, the + # subtree roll-up and the ancestor walk each go back to resolving + # their own row independently once the sharing is gone, plus the one + # unrelated relay-status read that was always separate. + own_row_reads = [ + q + for q in conv_selects + if "conversations.title" in q and "root_conversation_id = " not in q + ] + assert len(own_row_reads) == 2, [q.split("\n")[0][:70] for q in own_row_reads] + finally: + if collector is not None: + await collector.stop() + for task_handle in list(sessions_module._runner_relay_tasks.values()): + task_handle.task.cancel() + sessions_module._runner_relay_tasks.clear() + + +@pytest.mark.asyncio +async def test_relay_completion_stops_when_the_session_is_gone(db_uri: str) -> None: + """ + A completion whose session vanished must not bill a recreated row. + + The handler reads the conversation once and uses it for both the pricing + lookup and the subtree roll-up. Passing an absent row onward meant + ``None`` read as "re-read the row" to the accumulator and as "skip" to the + roll-up, so a session deleted and recreated under the same id had the + turn's usage persisted to the new row while the event reporting it was + suppressed — billed silently. + + The window is inside the handler, so it is constructed there: the shared + read observes the deletion, and the row exists again by the time a + re-reading accumulator would look. One read now decides for the whole + completion. + """ + from omnigent.server.routes import sessions as sessions_module + + sessions_module._runner_relay_tasks.clear() + store = SqlAlchemyConversationStore(db_uri) + doomed = store.create_conversation(title="relay:doomed") + session_id = doomed.id + + class _RecreateBetweenReads: + """Reports the deletion, then recreates the id behind the reader.""" + + def __init__(self, inner: SqlAlchemyConversationStore) -> None: + self._inner = inner + self.recreated = False + + def get_conversation(self, conversation_id: str): + row = self._inner.get_conversation(conversation_id) + if row is None and conversation_id == session_id and not self.recreated: + self.recreated = True + self._inner.create_conversation( + conversation_id=conversation_id, title="relay:recreated" + ) + return row + + def __getattr__(self, name: str): + return getattr(self._inner, name) + + proxy = _RecreateBetweenReads(store) + + response_id = "resp_relay_gone_1" + turn_events: list[dict[str, Any]] = [ + { + "type": "response.in_progress", + "response": {"id": response_id, "model": "databricks-claude-sonnet-4-6"}, + }, + { + "type": "response.completed", + "response": { + "id": response_id, + "model": "databricks-claude-sonnet-4-6", + "usage": {"input_tokens": 1000, "output_tokens": 500, "total_tokens": 1500}, + }, + }, + ] + release = asyncio.Event() + fake_runner = _ScriptedRunnerClient(release, turn_events) + + collector = None + try: + handle = await sessions_module._ensure_runner_relay_ready( + session_id, + "runner_relay_gone", + fake_runner, # type: ignore[arg-type] + conversation_store=proxy, # type: ignore[arg-type] + ) + assert handle is not None + collector = await start_session_stream_collector(session_id) + + await store.delete_conversation(session_id) + + release.set() + seen: list[str] = [] + while not seen or seen[-1] != "response.completed": + event = await collector.next_event() + seen.append(event["type"]) + + assert proxy.recreated, "the interleave never happened; the test proves nothing" + recreated = store.get_conversation(session_id) + assert recreated is not None + assert not dict(recreated.session_usage), ( + f"the vanished session's turn was billed to the recreated row: " + f"{dict(recreated.session_usage)}" + ) + finally: + if collector is not None: + await collector.stop() + for task_handle in list(sessions_module._runner_relay_tasks.values()): + task_handle.task.cancel() + sessions_module._runner_relay_tasks.clear() diff --git a/tests/server/routes/test_sessions_snapshot.py b/tests/server/routes/test_sessions_snapshot.py index ac4a095f48..a094bbcb8e 100644 --- a/tests/server/routes/test_sessions_snapshot.py +++ b/tests/server/routes/test_sessions_snapshot.py @@ -6,8 +6,10 @@ from dataclasses import dataclass, field from typing import Any +import httpx import pytest +from omnigent.db.utils import _engine_cache from omnigent.entities import Conversation, ConversationItem, MessageData, PagedList from omnigent.server.routes import sessions as _sessions_mod from omnigent.server.routes.sessions import ( @@ -19,6 +21,7 @@ _truncate_label, ) from omnigent.spec.types import AgentSpec, ExecutorSpec +from tests.server.helpers import count_conversation_selects async def _drain_runner_skills(session_id: str) -> None: @@ -108,11 +111,13 @@ def list_conversations( after: str | None = None, kind: str | None = "default", root_conversation_id: str | None = None, + include_archived: bool = False, ) -> PagedList[Conversation]: """Return the spawn tree sharing ``root_conversation_id``. ``load_session_usage`` walks the tree via this method to sum a - parent's subtree usage. With an explicit graph, return every + parent's subtree usage, and passes ``include_archived=True`` — + archived conversations still hold spend. With an explicit graph, return every conversation sharing the root; otherwise synthesize the single childless conversation the legacy tests expect. """ @@ -1961,3 +1966,74 @@ def set_labels(self, session_id: str, updates: dict[str, str]) -> None: captured["d6e1678fb446a1cf5a892e0df60aaba3"]["omnigent.last_task_error_code"] == "runner_error" ) + + +@pytest.mark.asyncio +async def test_snapshot_subtree_usage_reuses_conversation_row( + client: httpx.AsyncClient, + db_uri: str, +) -> None: + """ + The snapshot's subtree-usage recompute derives the tree root from the + row in hand instead of re-reading the conversation. Remaining SELECTs: + the handler's own row, the tree page scan, and the two deliberate + child-indicator / liveness id scans. + """ + from tests.server.helpers import create_test_session + + snap = await create_test_session(client, title="snapshot-load-once") + sid = snap["id"] + engine = _engine_cache[db_uri] + await client.get(f"/v1/sessions/{sid}") # warm + + with count_conversation_selects(engine) as selects: + resp = await client.get(f"/v1/sessions/{sid}") + + assert resp.status_code == 200 + # EXACT, not a ceiling: a ceiling let the test pass with the supplied-root + # argument removed (measured 5 without it, 4 with). The four are the + # handler's own row, the child-indicator id scan, the liveness id scan, + # and the subtree tree page — no per-session root re-resolution. + assert len(selects) == 4, [str(q)[:90] for q in selects] + + +@pytest.mark.asyncio +async def test_archived_descendant_spend_counts_toward_displayed_total( + client: httpx.AsyncClient, + db_uri: str, +) -> None: + """ + Archived spend is still spend, and it shows in the displayed total. + + This is a deliberate product decision, not an accident of the + cost-enforcement fix that motivated it. Including archived rows in the + spawn-tree load was required to stop an archive-after-preload seeding a + $0 budget and ALLOWing over-budget tool calls. The same tree feeds the + DISPLAY roll-up, so an archived sub-agent's spend now appears in the + parent's total as well. + + The alternative — two tree loads with different archived semantics, one + for enforcement and one for display — would let the badge disagree with + the gate, which is worse than the change being visible. Pinned here so + the display behaviour cannot be altered silently in either direction. + """ + from omnigent.stores.conversation_store.sqlalchemy_store import ( + SqlAlchemyConversationStore, + ) + from tests.server.helpers import create_test_session + + snap = await create_test_session(client, title="archived-display") + sid = snap["id"] + store = SqlAlchemyConversationStore(db_uri) + child = store.create_conversation( + kind="sub_agent", parent_conversation_id=sid, title="archived:child" + ) + store.set_session_usage(sid, {"total_cost_usd": 1.0}) + store.set_session_usage(child.id, {"total_cost_usd": 2.0}) + store.update_conversation(child.id, archived=True) + + resp = await client.get(f"/v1/sessions/{sid}") + assert resp.status_code == 200, resp.text + assert resp.json()["total_cost_usd"] == pytest.approx(3.0), ( + "archived descendant spend must remain in the displayed subtree total" + ) diff --git a/tests/server/routes/test_usage_report.py b/tests/server/routes/test_usage_report.py index cc026c7ea9..6bf6caee50 100644 --- a/tests/server/routes/test_usage_report.py +++ b/tests/server/routes/test_usage_report.py @@ -2,8 +2,10 @@ from __future__ import annotations +import httpx import pytest +from omnigent.db.utils import _engine_cache from omnigent.server.auth import RESERVED_USER_LOCAL from omnigent.server.routes.usage import ( _build_usage_report, @@ -14,6 +16,7 @@ from omnigent.stores.conversation_store.sqlalchemy_store import ( SqlAlchemyConversationStore, ) +from tests.server.helpers import count_conversation_selects _DAY = 86_400 # Agent ids are stored as 16-byte uuids, so tests use a valid 32-char hex id. @@ -199,3 +202,45 @@ def test_sum_daily_cost_range(db_uri: str) -> None: assert store.sum_daily_cost("alice", "0000-00-00") == 7.0 # all-time assert store.sum_daily_cost("alice", "2026-08-01") == 0.0 # nothing in range assert store.sum_daily_cost("nobody", "0000-00-00") == 0.0 + + +@pytest.mark.asyncio +async def test_usage_report_reuses_listed_rows_for_subtree_totals( + client: httpx.AsyncClient, + db_uri: str, +) -> None: + """ + ``GET /v1/usage`` holds each listed conversation, so the per-session + subtree recompute must not re-read it: one conversations SELECT for + the page plus one tree scan per session, and the totals still include + child spend (a count-only assertion would pass with the recompute + deleted). + """ + from tests.server.helpers import create_test_session + + snap = await create_test_session(client, title="usage-report") + sid = snap["id"] + store = SqlAlchemyConversationStore(db_uri) + child = store.create_conversation( + kind="sub_agent", parent_conversation_id=sid, title="usage:child" + ) + store.set_session_usage(sid, {"total_cost_usd": 0.10}) + store.set_session_usage(child.id, {"total_cost_usd": 0.05}) + + await client.get("/v1/usage") # warm caches + + engine = _engine_cache[db_uri] + with count_conversation_selects(engine) as selects: + resp = await client.get("/v1/usage") + + assert resp.status_code == 200, resp.text + body = resp.json() + rows = {s["id"]: s for s in body.get("sessions", [])} + assert sid in rows, body + # Parent 0.10 + child 0.05: proves the subtree recompute ran. + assert rows[sid]["cost_usd"] == pytest.approx(0.15) + # The listing itself + the tree scan(s); no per-session row re-read. + # EXACT: the page listing plus one tree scan for the listed session — + # no per-session conversation re-read. A ceiling here would pass with + # the supplied-root argument removed. + assert len(selects) == 2, [str(q)[:90] for q in selects] diff --git a/tests/stores/test_conversation_store.py b/tests/stores/test_conversation_store.py index ad27389e3d..187f929f44 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 ( @@ -5264,3 +5264,149 @@ 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() 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" + ) + + +def test_get_runner_ids_distinguishes_missing_from_unbound( + conversation_store: SqlAlchemyConversationStore, +) -> None: + """ + The two half-written states are distinct, and both must report as such. + + The AP row is the existence authority and the binding lives on the + Omnigent metadata row, written in a separate transaction — so either row + can outlive the other: + + - AP row gone, metadata left behind (the documented deletion tradeoff): + the conversation no longer exists, so it must be OMITTED. Reporting + its old runner routes events to a session that is gone instead of + 404-ing. + - AP row present, metadata absent (the same shape during creation, in + the opposite order): the conversation exists but is unbound, so it + must map to ``None`` — which the route turns into CONFLICT. Omitting + it produces a false NOT_FOUND. + + A metadata-rooted lookup answers the first correctly and the second + wrongly, which is why both directions are asserted from one table. + """ + from sqlalchemy import delete as sa_delete + + from omnigent.db.db_models import SqlConversation as _Conv + from omnigent.db.db_models import SqlConversationMetadata as _Meta + + bound = conversation_store.create_conversation(title="bound") + conversation_store.set_runner_id(bound.id, "runner_live") + orphaned_metadata = conversation_store.create_conversation(title="ap-row-gone") + conversation_store.set_runner_id(orphaned_metadata.id, "runner_old") + missing_metadata = conversation_store.create_conversation(title="metadata-gone") + + with conversation_store._conv_session() as session: + session.execute(sa_delete(_Conv).where(_Conv.id == orphaned_metadata.id)) + with conversation_store._session() as session: + session.execute(sa_delete(_Meta).where(_Meta.id == missing_metadata.id)) + + ids = [bound.id, orphaned_metadata.id, missing_metadata.id] + assert conversation_store.get_runner_ids(ids) == { + bound.id: "runner_live", + # present but unbound -> None, never omitted + missing_metadata.id: None, + # AP row gone -> omitted entirely, never its stale runner + } + # The full read agrees about existence in both directions. + assert conversation_store.get_conversation(orphaned_metadata.id) is None + assert conversation_store.get_conversation(missing_metadata.id) is not None diff --git a/tests/stores/test_conversation_store_split_db.py b/tests/stores/test_conversation_store_split_db.py index fb56e459fa..0086b9ca9a 100644 --- a/tests/stores/test_conversation_store_split_db.py +++ b/tests/stores/test_conversation_store_split_db.py @@ -596,3 +596,36 @@ def test_delete_conversation_keeps_template_agent( asyncio.run(store.delete_conversation(conv.id)) assert _col(omnigent_db, "agents", "id") == ["191cbf904e3223e9e00ac9a1abfe79a5"] + + +def test_get_runner_ids_distinguishes_missing_from_unbound_split_db( + store: SqlAlchemyConversationStore, +) -> None: + """Split-DB must draw the same two-state distinction as one bind. + + With the tables on separate binds there is no join to root the lookup + in, so existence comes from the AP bind and bindings are overlaid onto + it. Skipping that seeding drops an existing-but-unbound conversation and + reports a false NOT_FOUND — the same defect as a metadata-rooted lookup, + reached a different way. + """ + from sqlalchemy import delete as sa_delete + + from omnigent.db.db_models import SqlConversation as _Conv + from omnigent.db.db_models import SqlConversationMetadata as _Meta + + bound = store.create_conversation(title="bound") + store.set_runner_id(bound.id, "runner_live") + orphaned_metadata = store.create_conversation(title="ap-row-gone") + store.set_runner_id(orphaned_metadata.id, "runner_old") + missing_metadata = store.create_conversation(title="metadata-gone") + + with store._conv_session() as session: + session.execute(sa_delete(_Conv).where(_Conv.id == orphaned_metadata.id)) + with store._session() as session: + session.execute(sa_delete(_Meta).where(_Meta.id == missing_metadata.id)) + + assert store.get_runner_ids([bound.id, orphaned_metadata.id, missing_metadata.id]) == { + bound.id: "runner_live", + missing_metadata.id: None, + }