""" ``build_policy_engine`` — construct a :class:`PolicyEngine` for a workflow. 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. Phase 2 scope: zero-policy and declared-policy paths both work; concrete Policy subclasses land in Phases 3+, and this builder will start instantiating them as those phases ship. """ from __future__ import annotations import logging from typing import Any import cachetools from omnigent.entities import Conversation from omnigent.entities import Policy as StoredPolicy from omnigent.errors import ErrorCode, OmnigentError from omnigent.llms.context_window import fetch_model_pricing from omnigent.policies.base import Policy from omnigent.policies.function import resolve_function_policy from omnigent.policies.schema import ( SESSION_COST_ASK_APPROVED_STATE_KEY, SESSION_COST_UNPRICED_APPROVED_KEY, ) from omnigent.policies.types import PolicyLLMClient from omnigent.runtime.credentials.databricks import resolve_databricks_workspace from omnigent.runtime.policies.engine import PolicyEngine from omnigent.spec.types import ( DEFAULT_ASK_TIMEOUT, AgentSpec, FunctionPolicySpec, FunctionRef, LabelDef, LLMConfig, PolicySpec, ) from omnigent.stores.conversation_store import ConversationStore from omnigent.stores.policy_store import PolicyStore _logger = logging.getLogger(__name__) # Dotted path of the per-user daily cost-budget factory. The engine is # seeded with the session owner's daily-cost rollup ONLY when a policy # set includes this handler — otherwise the owner + daily-cost lookups # are skipped entirely, so sessions/deployments that don't use it pay # nothing extra per evaluation. _USER_DAILY_COST_POLICY_PATH = "omnigent.policies.builtins.cost.user_daily_cost_budget" # Dotted path of the per-subagent cost-budget factory. The engine is # seeded with the subtree-scoped usage ONLY when a policy set includes # this handler — otherwise the subtree usage lookup is skipped. _SUBAGENT_COST_POLICY_PATH = "omnigent.policies.builtins.cost.subagent_cost_budget" # Hardcoded policy that always ASKs before sys_add_policy executes. # Injected unconditionally into every engine so agents cannot add # policies without user approval. _ASK_ON_ADD_POLICY_SPEC = FunctionPolicySpec( name="__ask_on_add_policy", on=None, function=FunctionRef( path="omnigent.policies.builtins.safety.ask_on_add_policy", arguments=None, ), ) # Bounded cache of ``conversation_id -> session owner``. The owner # (``LEVEL_OWNER`` grantee) is immutable for a session's lifetime, so # caching it avoids a ``session_permissions`` lookup on every # per-tool-call engine build. Only non-``None`` owners are cached (a # session is granted its owner atomically at creation, so ``None`` is a # transient single-user/pre-grant state, not worth caching). _SESSION_OWNER_CACHE: cachetools.LRUCache[str, str] = cachetools.LRUCache(maxsize=4096) # TTL cache of ``workspace_id -> list[PolicySpec]`` for DB-stored default # policies. Default policies are admin-managed and change infrequently, so # a short TTL (30 s) avoids one ``list_defaults()`` DB query per tool-call # evaluation while still propagating changes within half a minute. _DEFAULT_POLICY_SPECS_CACHE: cachetools.TTLCache[int, list[PolicySpec]] = cachetools.TTLCache( maxsize=256, ttl=30 ) # Invalidation-based LRU cache of ``(workspace_id, conversation_id) -> list[PolicySpec]`` # for session-scoped policies. Unlike defaults, session policies can be added # mid-session (via sys_add_policy), so a TTL would delay enforcement. Instead, # the cache is explicitly invalidated whenever a session policy is mutated via # the CRUD routes. Keyed by workspace to prevent cross-tenant leakage. # Bounded (LRU, 4096 entries) to match _SESSION_OWNER_CACHE and prevent unbounded # growth — LRU eviction handles sessions that end without any policy mutation. _SESSION_POLICY_SPECS_CACHE: cachetools.LRUCache[tuple[int, str], list[PolicySpec]] = ( cachetools.LRUCache(maxsize=4096) ) def _needs_user_daily_cost(specs: list[PolicySpec]) -> bool: """ Return whether any policy in *specs* is the per-user daily cost-budget. Drives the conditional injection: only when this returns ``True`` does :func:`build_policy_engine` resolve the owner and read the daily-cost rollup. :param specs: The merged policy specs for the engine. :returns: ``True`` when a :class:`FunctionPolicySpec` references the ``user_daily_cost_budget`` factory. """ return any( isinstance(s, FunctionPolicySpec) and s.function is not None and s.function.path == _USER_DAILY_COST_POLICY_PATH for s in specs ) def _needs_subtree_usage(specs: list[PolicySpec]) -> bool: """ Return whether any policy in *specs* is the per-subagent cost-budget. Drives the conditional injection: only when this returns ``True`` does :func:`build_policy_engine` compute the subtree usage seed. :param specs: The merged policy specs for the engine. :returns: ``True`` when a :class:`FunctionPolicySpec` references the ``subagent_cost_budget`` factory. """ return any( isinstance(s, FunctionPolicySpec) and s.function is not None and s.function.path == _SUBAGENT_COST_POLICY_PATH for s in specs ) def _normalize_usage_for_engine(usage: dict[str, float]) -> dict[str, float]: """ Normalize a usage dict for injection into the policy engine. Removes display-only fields (``by_model``) and converts the enforcement-cost field (``policy_cost_usd``) to the engine's canonical ``total_cost_usd`` key. Both operations are idempotent: if a field is absent, the operation is a no-op. :param usage: The usage dict to normalize (modified in-place). :returns: The normalized dict (same object, for chaining). """ usage.pop("by_model", None) policy_cost = usage.pop("policy_cost_usd", None) if policy_cost is not None: usage["total_cost_usd"] = policy_cost 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, ) -> str | None: """ Resolve a session's owner, caching the immutable result. :param conversation_id: The session, e.g. ``"conv_abc123"``. :param conversation_store: Store for the owner lookup. :returns: The owner user id, or ``None`` when the session has no owner grant (single-user mode). """ owner = _SESSION_OWNER_CACHE.get(conversation_id) if owner is not None: return owner owner = conversation_store.get_session_owner(conversation_id) if owner is not None: _SESSION_OWNER_CACHE[conversation_id] = owner return owner def _load_user_daily_cost( conversation_id: str, conversation_store: ConversationStore, ) -> dict[str, float | str]: """ Read the session owner's per-UTC-day cost rollup as the engine seed. Resolves the owner (cached) and reads ``{cost_usd, ask_approved_usd}`` for today (UTC), tagged with the owner's ``user_id`` so the budget policy can name whose spend tripped the gate. When the session has no owner grant (single-user mode), returns zeros (and no ``user_id``) so the per-user daily budget never trips — consistent with the write path, which also no-ops without an owner. :param conversation_id: The session, e.g. ``"conv_abc123"``. :param conversation_store: Store for the owner + daily-cost lookups. :returns: ``{"cost_usd": , "ask_approved_usd": , "user_id": }``; ``user_id`` omitted in single-user mode. """ from omnigent.db.utils import now_epoch, utc_day owner = _resolve_session_owner_cached(conversation_id, conversation_store) if owner is None: return {"cost_usd": 0.0, "ask_approved_usd": 0.0} state: dict[str, float | str] = dict( conversation_store.get_daily_cost_state(owner, utc_day(now_epoch())) ) state["user_id"] = owner return state def build_policy_engine( *, spec: AgentSpec, conversation_id: str, conversation_store: ConversationStore, connection_override: dict[str, str] | None = None, default_policies: list[PolicySpec] | None = None, policy_store: PolicyStore | None = None, server_llm: LLMConfig | None = None, host_connection: dict[str, str] | None = None, ) -> PolicyEngine: """ Construct the :class:`PolicyEngine` for one workflow. When ``spec.guardrails`` is ``None`` (no guardrails declared), *default_policies* is empty, and no session policies are stored, returns a no-op engine with empty policies and labels — the four enforcement sites still 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. Policy run order: session policies (from the CRUD API) first, then agent spec policies, then *default_policies* (server-wide admin policies). This lets user-configured session policies short-circuit on DENY before agent or admin policies run, and gives admin policies the last word on ALLOW/ASK decisions. For sub-agent conversations, session policies from the root (top-level) conversation are inherited and prepended before any child-specific session policies. This ensures guardrails set on the parent session (e.g. via ``sys_add_policy``) also govern spawned sub-agents. Policies with the same ``name`` on both root and child are deduplicated (child wins). :param spec: The parsed agent spec. :param conversation_id: The conversation this workflow is running on, e.g. ``"conv_abc123"``. :param conversation_store: The store used for label reads and writes. Held by the engine for the life of the workflow. :param connection_override: Fallback ``{"base_url", "api_key"}`` used by prompt policies whose spec declares no ``llm.connection``. Explicit policy / agent connections still win. :param default_policies: Server-wide policies appended after per-agent policies. Sourced from ``RuntimeCaps.default_policies`` (parsed from the server ``--config`` YAML at startup). ``None`` and ``[]`` both mean no server-wide policies. :param policy_store: Session-scoped policy store. When provided, enabled policies for ``conversation_id`` are loaded and inserted between agent and admin policies in the evaluation order. :param server_llm: Server-level LLM configuration from ``RuntimeCaps.llm``. When provided, a :class:`~omnigent.policies.types.PolicyLLMClient` is constructed and injected into every function policy's ``event["llm_client"]``. ``None`` means no server-level LLM — function policies see ``None``. :param host_connection: Per-request ``{"base_url", "api_key"}`` dict resolved from the caller's auth token (e.g. via :attr:`RuntimeCaps.policy_llm_connection_factory`). When provided, takes precedence over any connection derived from ``server_llm.connection`` / ``server_llm.profile``, so LLM calls are billed to the request caller rather than a static service credential. ``None`` falls back to the server-level connection. :returns: A :class:`PolicyEngine` ready for evaluation. """ 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 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 # Always require user approval before sys_add_policy executes. # Appended unconditionally so the guard is present even when # no other guardrails are declared (the noop-engine path below # is no longer reachable since all_policy_specs is never empty). all_policy_specs.append(_ASK_ON_ADD_POLICY_SPEC) label_defs = (guardrails.labels or {}) if guardrails else {} initial_labels = _seed_and_load_labels( conversation_id=conversation_id, label_defs=label_defs, conversation_store=conversation_store, ) initial_session_state = _load_session_state(conversation_id, conversation_store) # 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) for _root_key in ( SESSION_COST_ASK_APPROVED_STATE_KEY, SESSION_COST_UNPRICED_APPROVED_KEY, ): if _root_key in root_state: initial_session_state[_root_key] = root_state[_root_key] # 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 ) # 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 = ( _load_user_daily_cost(conversation_id, conversation_store) if _needs_user_daily_cost(all_policy_specs) else None ) initial_model = _resolve_session_model(conversation_id, conversation_store, spec) # 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 server_connection = _resolve_server_llm_connection(server_llm) # host_connection carries the per-request caller token (billed to # the caller). It takes precedence over the static server-level # connection so policy LLM calls are attributed to the right # identity. Falls back to server_connection when absent. policy_connection = host_connection or server_connection llm_client = _build_policy_llm_client(server_llm, policy_connection) # Fall back to the server's gateway connection for prompt-policy # classifiers (else they default to api.openai.com). effective_connection_override = connection_override or server_connection return PolicyEngine( policies=[ _instantiate_policy( s, agent_llm=spec.llm, connection_override=effective_connection_override, ) for s in all_policy_specs ], label_defs=label_defs, ask_timeout=guardrails.ask_timeout if guardrails else DEFAULT_ASK_TIMEOUT, conversation_id=conversation_id, initial_labels=initial_labels, initial_session_state=initial_session_state, initial_usage=initial_usage, initial_subtree_usage=initial_subtree_usage, initial_user_daily_cost=initial_user_daily_cost, token_pricing=token_pricing, initial_model=initial_model, conversation_store=conversation_store, root_conversation_id=root_conversation_id, llm_client=llm_client, ) def _resolve_server_llm_connection( server_llm: LLMConfig | None, ) -> dict[str, str] | None: """ Resolve the server-level LLM connection dict. Returns ``server_llm.connection`` directly when present; otherwise resolves ``server_llm.profile`` to a Databricks workspace connection. ``None`` when no server LLM is configured or it declares neither a connection nor a profile. :param server_llm: The server-level :class:`LLMConfig` from ``RuntimeCaps.llm``, or ``None``. :returns: A ``{"base_url", "api_key"}`` dict, or ``None``. :raises OSError: When ``profile`` is set but cannot be resolved. """ if server_llm is None: return None if server_llm.connection is not None: return server_llm.connection if server_llm.profile is not None: return _resolve_databricks_connection(server_llm.profile) return None def _build_policy_llm_client( server_llm: LLMConfig | None, connection: dict[str, str] | None, ) -> PolicyLLMClient | None: """ Construct a :class:`PolicyLLMClient` from server-level LLM config. Returns ``None`` when no server-level ``llm:`` config is present. The :class:`~omnigent.llms.client.Client` is instantiated lazily here (no constructor args — auth routes per-call via ``connection_params``). :param server_llm: The server-level :class:`LLMConfig` from ``RuntimeCaps.llm``. ``None`` when the server config has no ``llm:`` block. :param connection: The connection dict already resolved by :func:`_resolve_server_llm_connection` (shared with the classifier connection-override fallback). :returns: A :class:`PolicyLLMClient` wrapping the client with pre-bound model/connection/timeout, or ``None``. """ if server_llm is None: return None from omnigent.llms.client import Client # Models prefixed with ``databricks-`` (e.g. # ``databricks-claude-sonnet-4-6``) need the ``databricks/`` # provider prefix so the LLM adapter routes through # DatabricksAdapter (Chat Completions) rather than # OpenAIAdapter (Responses API). Without this, the request # hits ``/responses`` on the Databricks gateway → 400. model = server_llm.model if "/" not in model and model.startswith("databricks-"): model = f"databricks/{model}" return PolicyLLMClient( _client=Client(), _model=model, _connection=connection, _request_timeout=server_llm.request_timeout, ) def _resolve_databricks_connection(profile: str) -> dict[str, str]: """ Resolve a Databricks CLI profile to a connection dict. Uses :func:`~omnigent.runtime.credentials.databricks.resolve_databricks_workspace` to resolve the profile to workspace host + bearer token, then builds the ``{"base_url": ..., "api_key": ...}`` dict that the LLM adapter expects. :param profile: The Databricks CLI profile name, e.g. ``"my-workspace"``. :returns: A connection dict with ``base_url`` (workspace host + ``/serving-endpoints``) and ``api_key`` (bearer token), e.g. ``{"base_url": "https://host/serving-endpoints", "api_key": "dapi..."}``. :raises OSError: When the profile cannot be resolved. """ creds = resolve_databricks_workspace(profile) return { "base_url": creds.host + "/serving-endpoints", "api_key": creds.token, } def _instantiate_policy( spec: PolicySpec, *, agent_llm: LLMConfig | None, connection_override: dict[str, str] | None = None, ) -> Policy: """ Dispatch a :class:`PolicySpec` to the matching runtime :class:`Policy` subclass. :param spec: The declarative spec. :param agent_llm: The agent-level ``llm:`` config. Used as the default backend for :class:`PromptPolicy` when the policy didn't declare its own ``llm:`` override. Unused for Function policies. :param connection_override: Forwarded to the prompt classifier as a fallback when the policy / agent declare no connection. :returns: A :class:`Policy` subclass instance bound to the spec. :raises NotImplementedError: When ``spec`` is not a known :class:`PolicySpec` subclass — parser bug protection. """ if isinstance(spec, FunctionPolicySpec): return resolve_function_policy(spec) raise NotImplementedError( f"Policy type {type(spec).__name__} for {spec.name!r} is not " f"a known subclass of PolicySpec (FunctionPolicySpec).", ) def _build_noop_engine( *, conversation_id: str, conversation_store: ConversationStore, ) -> PolicyEngine: """ Build an engine for an agent with no guardrails declared. Kept as a named helper rather than inlined so the zero-policy path is grep-able ("why is every phase returning ALLOW?" → search for ``_build_noop_engine``). :param conversation_id: The conversation for the workflow. :param conversation_store: Writes from this engine still go through the store — useful if a later turn of the same conversation runs under an updated spec that does declare guardrails. :returns: An engine with zero policies and an empty label cache. """ # We still read the persisted labels and session_state (if any) # so an engine upgrade mid-conversation sees state its # predecessor wrote. existing = _load_existing_labels(conversation_id, conversation_store) initial_session_state = _load_session_state(conversation_id, conversation_store) # Inert here (no policies read usage), but kept identical to the live # engine's seed so the "engine usage == session-wide total" invariant # holds uniformly across both builders. initial_usage = _policy_usage_seed(conversation_id, conversation_store) return PolicyEngine( policies=[], label_defs={}, ask_timeout=DEFAULT_ASK_TIMEOUT, conversation_id=conversation_id, initial_labels=existing, initial_session_state=initial_session_state, initial_usage=initial_usage, conversation_store=conversation_store, ) def _seed_and_load_labels( *, conversation_id: str, label_defs: dict[str, LabelDef], conversation_store: ConversationStore, ) -> dict[str, str]: """ Seed declared initial values and return the current snapshot. Race-safe across concurrent workflows: only writes keys that are missing from the persisted state. If two workflows seed simultaneously, the dialect-specific UPSERT guarantees one writer wins per (conversation, key) pair and the other no-ops. :param conversation_id: The conversation to seed. :param label_defs: Per-key declarations from the spec. Keys with ``initial is None`` are skipped (those labels start unset until a policy writes them). :param conversation_store: Target for both the read and the seed UPSERT. :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. existing = _load_existing_labels(conversation_id, conversation_store) return existing def _load_existing_labels( conversation_id: str, conversation_store: ConversationStore, ) -> dict[str, str]: """ Load the current persisted label state. Empty dict when the conversation has no labels yet (or when the conversation itself does not exist yet — the caller is responsible for ordering conversation creation before engine build). :param conversation_id: Conversation to load. :param conversation_store: Store to read from. :returns: ``{key: value}`` map. Empty when nothing persisted. """ conv = conversation_store.get_conversation(conversation_id) if conv is None: return {} return dict(conv.labels) def _load_session_state( conversation_id: str, conversation_store: ConversationStore, ) -> dict[str, Any]: """ Load the current persisted session state. Empty dict when the conversation has no session state yet (or when the conversation itself does not exist yet — the caller is responsible for ordering conversation creation before engine build). :param conversation_id: Conversation to load, e.g. ``"conv_abc123"``. :param conversation_store: Store to read from. :returns: Session state dict. Empty when nothing persisted. """ conv = conversation_store.get_conversation(conversation_id) if conv is None: return {} 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). _SUBTREE_USAGE_PAGE_SIZE = 100 # Usage counters summed across a conversation subtree. Restricted to the # known numeric keys the PolicyEngine reads so an unexpected key in one # conversation's persisted usage can't leak into the aggregate. _SUMMABLE_USAGE_KEYS = ( "input_tokens", "output_tokens", "total_tokens", "cache_read_input_tokens", "cache_creation_input_tokens", "total_cost_usd", ) def _merge_by_model( aggregate: dict[str, dict[str, float]], per_conv: dict[str, Any], ) -> None: """ Deep-merge one conversation's ``by_model`` sub-dict into the subtree aggregate. Unions model keys and sums each numeric per-bucket value (the token counters and ``total_cost_usd``) within each model, so a parent's per-model view folds in sub-agents that ran a different model. Mutates ``aggregate`` in place. :param aggregate: The running subtree ``by_model`` map being built, keyed by raw harness model id, e.g. ``{"claude-sonnet-4-6": {"input_tokens": 1200}}``. :param per_conv: One conversation's ``session_usage["by_model"]`` dict. Non-dict model buckets (malformed persisted data) are skipped. """ for model, bucket in per_conv.items(): if not isinstance(bucket, dict): continue agg_bucket = aggregate.setdefault(model, {}) for key, value in bucket.items(): # Only sum genuine numerics; ``bool`` is an ``int`` subclass so # exclude it explicitly to avoid summing a stray flag as 1. if isinstance(value, (int, float)) and not isinstance(value, bool): agg_bucket[key] = agg_bucket.get(key, 0.0) + value def load_session_usage( conversation_id: str, conversation_store: ConversationStore, ) -> dict[str, Any]: """ Load cumulative session usage for a conversation **plus all of its sub-agent descendants** (the subtree total). A cost-ask policy on a parent must see what its sub-agents spent, but each conversation persists only its own usage, so this sums the conversation and every conversation it transitively spawned (all share one ``root_conversation_id``). Read-only; the per-conversation write path is server-side ``_accumulate_session_usage`` / ``_persist_native_cumulative_usage``, not this function. Public because the server's session snapshot / ``session.usage`` SSE use this per-node subtree total to DISPLAY a node's own cost (a sub-agent's badge shows only its subtree). Cost GATING does NOT use this per-node view — it seeds from the whole-tree total via :func:`_policy_usage_seed` (which calls this with the tree root), so a sub-agent gates against the full session spend rather than just its own subtree. :param conversation_id: Conversation to load, e.g. ``"conv_abc123"``. :param conversation_store: Store to read from. :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 ``policy_cost_usd`` (the ENFORCEMENT cost sum — see below; only keys present in at least one conversation appear). When any conversation in the subtree recorded a per-model breakdown, a nested ``by_model`` key maps each raw harness model id to its own summed token/cost buckets (folding in sub-agents that ran a different model). Empty when the conversation does not exist or no usage is recorded. Display callers read ``total_cost_usd``; the policy seed (:func:`_policy_usage_seed`) reads ``policy_cost_usd`` (both unaffected by ``by_model``). """ conv = conversation_store.get_conversation(conversation_id) if conv is None: return {} tree = _load_tree_conversations(conv.root_conversation_id, conversation_store) subtree_ids = _subtree_conversation_ids(tree, conversation_id) totals: dict[str, Any] = {} # Per-model breakdown summed across the subtree, parallel to the flat sums. by_model_totals: dict[str, dict[str, float]] = {} # Enforcement cost total, accumulated alongside the display sums so the # policy seed can pick it without a second tree pass. policy_cost_total = 0.0 any_policy_cost = False for tree_conv in tree: if tree_conv.id not in subtree_ids: continue session_usage = tree_conv.session_usage for key in _SUMMABLE_USAGE_KEYS: value = session_usage.get(key) if value is not None: totals[key] = totals.get(key, 0.0) + value # Per-model sub-dict (nested ``by_model`` key) is ignored by the flat # ``_SUMMABLE_USAGE_KEYS`` loop above; merge it separately so the flat # sum (used by policy gating) stays unchanged and backward-compatible. per_conv_by_model = session_usage.get("by_model") if isinstance(per_conv_by_model, dict): _merge_by_model(by_model_totals, per_conv_by_model) # Enforcement cost: prefer this conversation's ``policy_cost_usd`` # (claude-native's real-time figure incl. in-flight sub-agent spend), # else its displayed ``total_cost_usd`` (codex-native / relay don't # post the split). Kept separate from the ``total_cost_usd`` sum # above so the badge keeps the authoritative statusLine total. per_conv_policy_cost = session_usage.get("policy_cost_usd") if per_conv_policy_cost is None: per_conv_policy_cost = session_usage.get("total_cost_usd") if per_conv_policy_cost is not None: policy_cost_total += per_conv_policy_cost any_policy_cost = True if any_policy_cost: totals["policy_cost_usd"] = policy_cost_total if by_model_totals: totals["by_model"] = by_model_totals return totals def _policy_usage_seed( conversation_id: str, conversation_store: ConversationStore, ) -> dict[str, float]: """ SESSION-WIDE usage seed for the :class:`PolicyEngine`; cost = ENFORCEMENT total. Cost gating caps the **session** (the whole spawn tree), so this seeds from the tree-wide total — the spend rooted at ``root_conversation_id`` — not just the subtree rooted at the node being evaluated. A sub-agent gated on its own subtree would miss its parent's and siblings' spend, so the session could overshoot its budget while the orchestrator parent is parked (it makes no tool calls, so its own gate never fires). For the root conversation this equals the per-node subtree (its subtree IS the whole tree), so only sub-agents change behavior. The cost the gate reads (``total_cost_usd`` in the returned seed) is the ENFORCEMENT total — ``policy_cost_usd`` when present (claude-native's real-time figure that reflects in-flight sub-agent spend while the displayed statusLine ``S`` is frozen), falling back to the displayed ``total_cost_usd`` for harnesses that don't post the split (codex-native, relay). The ``policy_cost_usd`` key is then dropped so the engine's usage context carries only the standard counters. Display callers use :func:`load_session_usage` directly (per-node subtree, authoritative ``total_cost_usd`` = ``S``), which is why the cost-budget gate can read a higher in-flight / session-wide total than a node's badge shows mid-turn. :param conversation_id: Conversation to seed the engine for, e.g. ``"conv_child"`` (a sub-agent) or ``"conv_root"`` (the session root). :param conversation_store: Store to read the tree usage from. :returns: Whole-tree usage seed dict; when an enforcement cost exists its ``total_cost_usd`` is the enforcement total and no ``policy_cost_usd`` key remains. Empty when the conversation is absent or no usage is recorded. """ conv = conversation_store.get_conversation(conversation_id) if conv is None: return {} usage = load_session_usage(conv.root_conversation_id, conversation_store) return _normalize_usage_for_engine(usage) def _load_tree_conversations( root_conversation_id: str, conversation_store: ConversationStore, ) -> list[Conversation]: """ Page through every conversation in one spawn tree. Returns all conversations sharing ``root_conversation_id`` (the root plus every sub-agent, any ``kind``), paginating so a large tree is not silently truncated. The ``root_conversation_id`` column is indexed, so this is a bounded indexed scan per page. :param root_conversation_id: The tree's root conversation id (every conversation in a spawn tree shares it), e.g. ``"conv_abc123"``. :param conversation_store: Store to read from. :returns: All conversations in the tree, in store order. """ convs: list[Conversation] = [] after: str | None = None while True: page = conversation_store.list_conversations( limit=_SUBTREE_USAGE_PAGE_SIZE, after=after, # None disables the kind filter so sub_agent conversations # (not just "default") are included in the tree. kind=None, root_conversation_id=root_conversation_id, ) convs.extend(page.data) if not page.has_more or page.last_id is None: break after = page.last_id return convs def _subtree_conversation_ids( tree: list[Conversation], conversation_id: str, ) -> set[str]: """ Collect a conversation id plus all its transitive sub-agent descendants within a spawn tree. Walking the subtree (rather than summing the whole ``tree``) keeps the aggregate correct when the policy is evaluated on a mid-tree sub-agent: that node sees its own spend and its children's, but not its parent's or siblings'. :param tree: All conversations in the spawn tree (from :func:`_load_tree_conversations`); order-independent. :param conversation_id: The subtree root to walk from, e.g. ``"conv_abc123"``. :returns: Set of conversation ids in the subtree rooted at ``conversation_id`` (always includes ``conversation_id``). """ children_by_parent: dict[str, list[str]] = {} for tree_conv in tree: if tree_conv.parent_conversation_id is not None: children_by_parent.setdefault(tree_conv.parent_conversation_id, []).append( tree_conv.id ) subtree: set[str] = set() stack = [conversation_id] while stack: node = stack.pop() if node in subtree: continue subtree.add(node) stack.extend(children_by_parent.get(node, [])) return subtree def _load_default_policy_specs( policy_store: PolicyStore | None, ) -> list[PolicySpec]: """ Load enabled server-wide default policies from the store. These are policies created via ``POST /v1/policies`` (``session_id IS NULL``). They run after agent-spec policies and before YAML-based admin policies in the evaluation order. Results are cached per workspace for 30 s (see :data:`_DEFAULT_POLICY_SPECS_CACHE`) to avoid a ``list_defaults()`` DB round-trip on every tool-call evaluation. The cache is keyed by workspace id so multi-tenant deployments never share results across tenants. Call :func:`invalidate_default_policy_specs_cache` after any mutation to make changes visible before the TTL expires. :param policy_store: The policy store. ``None`` returns an empty list. :returns: List of :class:`FunctionPolicySpec` for enabled default policies, in ``created_at ASC`` order. :raises OmnigentError: If an enabled policy has an unsupported type. """ if policy_store is None: return [] from omnigent.db.db_models import current_workspace_id workspace_id = current_workspace_id() cached = _DEFAULT_POLICY_SPECS_CACHE.get(workspace_id) if cached is not None: return cached specs: list[PolicySpec] = [] for policy in policy_store.list_defaults(): if not policy.enabled: continue if policy.type != "python": # Skip unsupported types with a warning rather than raising. # A session-scoped policy of unsupported type fails loudly (blast # radius: one session); a default policy of unsupported type would # crash engine construction for every session server-wide. Log and # skip so a stale or manually-inserted row can't cause an outage. _logger.warning( "Skipping default policy %r (id=%r): unsupported type %r — " "only type='python' can be evaluated. Disable or delete this " "policy to suppress this warning.", policy.name, policy.id, policy.type, ) continue specs.append(_stored_policy_to_spec(policy)) _DEFAULT_POLICY_SPECS_CACHE[workspace_id] = specs return specs def invalidate_default_policy_specs_cache() -> None: """ Evict the current workspace's entry from the default-policy specs cache. Call this after any mutation (create, update, delete) of a default policy so the next :func:`build_policy_engine` call re-reads from the DB rather than serving a stale TTL entry. Scoped to the current workspace context via :func:`~omnigent.db.db_models.current_workspace_id`. """ from omnigent.db.db_models import current_workspace_id _DEFAULT_POLICY_SPECS_CACHE.pop(current_workspace_id(), None) def invalidate_session_policy_specs_cache(conversation_id: str) -> None: """ Evict a conversation's entry from the session policy specs cache. Call this after any mutation (create, update, delete) of a session policy so the next :func:`build_policy_engine` call re-reads from the DB. Scoped to the current workspace context. :param conversation_id: The session whose cache entry to evict, e.g. ``"conv_abc123"``. """ from omnigent.db.db_models import current_workspace_id _SESSION_POLICY_SPECS_CACHE.pop((current_workspace_id(), conversation_id), None) def _load_session_policy_specs( conversation_id: str, policy_store: PolicyStore | None, ) -> list[PolicySpec]: """ Load enabled session policies from the store and convert them to :class:`FunctionPolicySpec` instances. Results are cached per ``(workspace_id, conversation_id)`` and invalidated on any mutation via :func:`invalidate_session_policy_specs_cache`. There is no TTL — the cache entry is permanent until explicitly evicted, so session policy changes (including ``sys_add_policy``) take effect immediately on the next engine build. Only ``type="python"`` policies are instantiable today. An enabled policy of an unsupported type (e.g. ``type="url"``) raises :class:`OmnigentError` rather than being skipped, so a stored guardrail that never enforces fails loudly. :param conversation_id: The session whose policies to load, e.g. ``"conv_abc123"``. :param policy_store: The session-scoped policy store. ``None`` returns an empty list. :returns: List of :class:`FunctionPolicySpec` for enabled session policies, in ``created_at ASC`` order. :raises OmnigentError: If an enabled policy has an unsupported ``type`` (e.g. ``type="url"``). """ if policy_store is None: return [] from omnigent.db.db_models import current_workspace_id key = (current_workspace_id(), conversation_id) cached = _SESSION_POLICY_SPECS_CACHE.get(key) if cached is not None: return cached stored = policy_store.list_for_session(conversation_id) specs: list[PolicySpec] = [] for policy in stored: if not policy.enabled: continue specs.append(_stored_policy_to_spec(policy)) _SESSION_POLICY_SPECS_CACHE[key] = specs return specs def _stored_policy_to_spec(policy: StoredPolicy) -> PolicySpec: """ Convert a stored :class:`Policy` entity to a :class:`FunctionPolicySpec`. For ``type="python"``, creates a :class:`FunctionPolicySpec` with a :class:`FunctionRef` pointing at the stored handler path and optional factory params. Session policies fire on all phases (``on=None``) — the callable itself decides whether to act by inspecting ``event["type"]``. For ``type="url"``, raises :class:`OmnigentError` (URL policy evaluation is unimplemented). A stored policy that never enforces is a silent safety hole, so converting an unsupported type fails loudly rather than returning ``None``. :param policy: The stored session policy entity. :returns: A :class:`FunctionPolicySpec`. :raises OmnigentError: If the policy ``type`` cannot be evaluated yet (e.g. ``type="url"``). """ if policy.type == "python": return FunctionPolicySpec( name=policy.name, # Session policies self-select: on=None means the # engine skips phase filtering and always dispatches. on=None, function=FunctionRef( path=policy.handler, arguments=policy.factory_params, ), ) # Any non-"python" type (today only "url") cannot be evaluated yet. # Reject loudly and fail closed: a stored policy that silently never # enforces is worse than a visible failure the operator can act on. raise OmnigentError( f"Session policy {policy.name!r} (id {policy.id!r}) has unsupported " f"type {policy.type!r}; only type='python' policies can be evaluated " f"today. URL policy evaluation is a future extension. Remove or " f"disable this policy, since storing it does not enforce anything.", code=ErrorCode.INVALID_INPUT, ) __all__ = [ "build_policy_engine", "invalidate_default_policy_specs_cache", "invalidate_session_policy_specs_cache", ]