diff --git a/dev/benchmarks/omnigent/journeys.py b/dev/benchmarks/omnigent/journeys.py index afc0bdd758..9c47a945a9 100644 --- a/dev/benchmarks/omnigent/journeys.py +++ b/dev/benchmarks/omnigent/journeys.py @@ -679,6 +679,86 @@ async def _measure_read_runner_file(env: BenchEnvironment, ctx: JourneyContext) await env.read_runner_file(session_id, _RUNNER_FILE_PATH) +# ── policy evaluate ────────────────────────────────────────── + +_POLICY_EVALUATE_PAYLOAD = { + "event": { + "type": "PHASE_TOOL_CALL", + "data": {"name": "Bash", "arguments": {"command": "ls"}}, + } +} + + +async def _setup_policy_evaluate_session(env: BenchEnvironment) -> str: + """Create an agent-bound session with a declared policy and warm the caches. + + The agent must declare at least one policy so ``any_policies_apply`` is + true and the full engine build (single tree scan + preloaded conversation) + runs on every evaluate call. A zero-policy spec short-circuits before the + build, which would measure the wrong path. + + Two warm calls are made before returning so the agent-spec and + session-policy caches are populated; the measured iteration then reflects + steady-state overhead, not cold-cache cost. + """ + assert env.client is not None + import io + import tarfile + + import yaml + + config: dict[str, object] = { + "spec_version": 1, + "name": "bench-policy-agent", + "guardrails": { + "policies": { + "allow_all": { + "type": "function", + "on": ["tool_call"], + "function": "tests.runtime.policies.conftest._always_allow", + } + } + }, + } + buf = io.BytesIO() + with tarfile.open(fileobj=buf, mode="w:gz") as tar: + payload = yaml.safe_dump(config).encode() + info = tarfile.TarInfo("config.yaml") + info.size = len(payload) + tar.addfile(info, io.BytesIO(payload)) + bundle = buf.getvalue() + + resp = await env.client.post( + "/v1/agents", + files={"bundle": ("agent.tar.gz", bundle, "application/gzip")}, + ) + resp.raise_for_status() + agent_id = resp.json()["id"] + + session_resp = await env.client.post("/v1/sessions", json={"agent_id": agent_id}) + session_resp.raise_for_status() + session_id = session_resp.json()["id"] + + # Warm the spec + policy caches — the measured iteration is steady-state. + for _ in range(2): + await env.client.post( + f"/v1/sessions/{session_id}/policies/evaluate", + json=_POLICY_EVALUATE_PAYLOAD, + ) + + return session_id + + +async def _measure_policy_evaluate(env: BenchEnvironment, ctx: JourneyContext) -> None: + session_id = cast(str, ctx) # _setup_policy_evaluate_session + assert env.client is not None + resp = await env.client.post( + f"/v1/sessions/{session_id}/policies/evaluate", + json=_POLICY_EVALUATE_PAYLOAD, + ) + resp.raise_for_status() + + # ── registry ───────────────────────────────────────────────── ALL_JOURNEYS: dict[str, Journey] = { @@ -754,6 +834,15 @@ async def _measure_read_runner_file(env: BenchEnvironment, ctx: JourneyContext) concurrency_safe=True, description="POST /v1/sessions/{id}/comments — create a review comment.", ), + Journey( + name="policy_evaluate", + kind="latency", + measure=_measure_policy_evaluate, + setup=_setup_policy_evaluate_session, + concurrency_safe=True, + description="POST /v1/sessions/{id}/policies/evaluate — PreToolUse hook " + "(single tree scan, preloaded conversation row, caches warm).", + ), # Runner (full-turn) journeys — with_runner=True, openai-agents, mock LLM. Journey( name="session_cold_start", diff --git a/omnigent/runtime/policies/builder.py b/omnigent/runtime/policies/builder.py index 12320e7d6f..eed4ae9621 100644 --- a/omnigent/runtime/policies/builder.py +++ b/omnigent/runtime/policies/builder.py @@ -18,6 +18,7 @@ from __future__ import annotations import logging +from dataclasses import dataclass from typing import Any import cachetools @@ -160,28 +161,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, @@ -286,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, @@ -302,11 +283,9 @@ def build_policy_engine( call through, they just always ALLOW. 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. + exists yet in ``conversation_labels``, seeds missing keys via + :meth:`ConversationStore.set_labels`. 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* @@ -326,6 +305,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. @@ -360,21 +358,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 @@ -386,22 +400,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, @@ -411,15 +526,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 = ( @@ -427,7 +551,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 @@ -684,6 +815,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. @@ -700,10 +832,14 @@ 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) + if existing is None: + existing = _load_existing_labels(conversation_id, conversation_store) to_seed = { key: ldef.initial for key, ldef in label_defs.items() @@ -763,36 +899,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). @@ -843,6 +949,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 @@ -866,6 +974,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 @@ -879,10 +995,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. @@ -992,6 +1221,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: @@ -1000,6 +1236,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, @@ -1037,6 +1305,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/server/routes/_sessions/helpers.py b/omnigent/server/routes/_sessions/helpers.py index 7fe79cddba..46e801f72f 100644 --- a/omnigent/server/routes/_sessions/helpers.py +++ b/omnigent/server/routes/_sessions/helpers.py @@ -6667,22 +6667,33 @@ 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 ) - return cast( - PolicyEngine, - build_policy_engine( - spec=spec, - conversation_id=session_id, - conversation_store=conversation_store, - default_policies=caps.default_policies, - policy_store=get_policy_store(), - server_llm=caps.llm, - host_connection=host_connection, - ), + return build_policy_engine( + 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, + host_connection=host_connection, ) @@ -6734,22 +6745,27 @@ async def _apply_pending_policy_ask_writes( return # Non-MCP relay path: pop and apply writes here since no retry # will arrive. - _pending_policy_ask_writes.pop(elicitation_id, None) # Resolve the agent spec + build the engine off the event loop: the # lookup, cold-cache bundle fetch, and engine construction are all # blocking DB/IO. spec = await asyncio.to_thread(_load_agent_spec_for_session, conv, agent_store) if spec is None: + _pending_policy_ask_writes.pop(elicitation_id, 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 ) + # Pop only after the engine build succeeds: a raise here (e.g. a + # concurrent agent rebind) would otherwise lose the approved writes + # with no retry possible. + _pending_policy_ask_writes.pop(elicitation_id, None) # The label/state writes hit the DB synchronously too — keep them # off the loop. if pending.set_labels: await asyncio.to_thread(engine.apply_label_writes, pending.set_labels) if pending.state_updates: - await asyncio.to_thread(engine.apply_state_updates, pending.state_updates) + with contextlib.suppress(ConversationNotFoundError): + await asyncio.to_thread(engine.apply_state_updates, pending.state_updates) def _build_actor(user_id: str | None) -> dict[str, str] | None: @@ -6856,11 +6872,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): @@ -7145,7 +7163,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 da6c280f87..98cf5449ba 100644 --- a/omnigent/server/routes/_sessions/orchestration.py +++ b/omnigent/server/routes/_sessions/orchestration.py @@ -83,6 +83,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 @@ -191,7 +194,6 @@ # primitives) live in _sessions.helpers. from omnigent.server.routes._sessions.helpers import ( SessionLiveness, - _ancestor_session_ids, _await_settled_managed_launch, _build_new_item, _build_policy_engine_from_spec, @@ -868,6 +870,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. @@ -880,18 +883,35 @@ def _publish_subtree_cost_to_ancestors( display side.) For each ancestor of *session_id*, recompute its subtree priced cost and publish a ``session.usage`` event carrying it. + One tree load serves the whole walk. Every ancestor of this session is in + the same tree, so both the chain and each ancestor's sum are derived from + one set of freshly-read rows — previously each ancestor paged the tree + again, and the chain came from a conversation row that may have been read + before a concurrent delete/recreate moved this session elsewhere. + Sync (does store reads + SSE fan-out); call via :func:`asyncio.to_thread`, mirroring the elicitation ancestor-publish helpers. ``session_stream.publish`` is safe to call from a worker thread. - :param conv_store: Store used to discover ancestors and sum each - ancestor's subtree usage. + :param conv_store: Store used to load the tree. :param session_id: The child session whose usage just changed, e.g. ``"conv_child123"``. + :param conv: The child's already-loaded conversation row, when the caller + holds one. Only its ``root_conversation_id`` is used, and only as a + hint: :func:`load_session_tree` verifies the tree it names actually + contains this session and resolves the root itself when it does not. + The parameter belongs to this function, not to whichever caller first + needed it, so that no caller can be removed and leave a signature + behind that its remaining callers already depend on. :returns: None. """ - for ancestor_id in _ancestor_session_ids(conv_store, session_id): - ancestor_usage = load_session_usage(ancestor_id, conv_store) + tree = load_session_tree( + session_id, + conv_store, + conv.root_conversation_id if conv is not None else None, + ) + for ancestor_id in ancestor_ids_from_tree(tree, session_id): + ancestor_usage = _sum_subtree_usage(tree, ancestor_id) subtree_cost = _priced_cost_for_display(ancestor_usage) usage_by_model = _usage_by_model_for_display(ancestor_usage) if subtree_cost is None and usage_by_model is None: @@ -1915,7 +1935,8 @@ async def _hold_native_ask_gate_impl( if result.set_labels: engine.apply_label_writes(result.set_labels) if result.state_updates: - engine.apply_state_updates(result.state_updates) + with contextlib.suppress(ConversationNotFoundError): + engine.apply_state_updates(result.state_updates) return approved @@ -6073,7 +6094,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: @@ -6230,7 +6251,7 @@ async def _evaluate_input_policy( request_content = {"user_content": user_text, "attachments": attachments} 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, @@ -8120,7 +8141,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: @@ -8200,7 +8221,8 @@ async def _handle_mcp_tools_call( if _pending.set_labels: await asyncio.to_thread(engine.apply_label_writes, _pending.set_labels) if _pending.state_updates: - await asyncio.to_thread(engine.apply_state_updates, _pending.state_updates) + with contextlib.suppress(ConversationNotFoundError): + await asyncio.to_thread(engine.apply_state_updates, _pending.state_updates) else: # ALLOW — policy no longer requires approval (e.g. label # state changed between the original ASK and this retry). diff --git a/omnigent/server/routes/sessions/routes_hooks.py b/omnigent/server/routes/sessions/routes_hooks.py index ab75c1f20d..a85fd958b9 100644 --- a/omnigent/server/routes/sessions/routes_hooks.py +++ b/omnigent/server/routes/sessions/routes_hooks.py @@ -15,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.runner.routing import RunnerRouter from omnigent.runtime import ( @@ -700,7 +701,14 @@ async def evaluate_policy( code=ErrorCode.INVALID_INPUT, ) - 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.", @@ -760,7 +768,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. @@ -769,6 +777,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``. """ @@ -776,20 +788,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/tests/runtime/policies/test_builder.py b/tests/runtime/policies/test_builder.py index d03b73f8d0..004cd6769b 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,829 @@ 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_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_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/server/integration/test_sessions_policy_evaluate.py b/tests/server/integration/test_sessions_policy_evaluate.py index 6842b38e68..6096f0ec71 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 @@ -1359,3 +1362,170 @@ async def test_policy_evaluate_gates_every_first_party_tool_result_shape( assert resp.json()["result"] == "POLICY_ACTION_DENY", ( f"{why}: the tool-scoped policy did not see a tool name — {resp.text[:160]}" ) + + +_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_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_snapshot.py b/tests/server/routes/test_sessions_snapshot.py index 50dccec243..a1143c6f0e 100644 --- a/tests/server/routes/test_sessions_snapshot.py +++ b/tests/server/routes/test_sessions_snapshot.py @@ -108,11 +108,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. """