omnigent-ai--omnigent/tests/runtime/policies/test_engine_session_state.py

"""
Tests for PolicyEngine session_state — reading and writing per-turn
mutable state from function policy callables.

Verifies:
- Engine exposes current ``session_state`` via ``event["session_state"]``.
- ``PolicyResult.state_updates`` are shallow-merged into the hot cache.
- State accumulates within a single engine instance (visible on the next evaluation).
- Multiple policies in one pass accumulate updates in YAML order
  (last-write-wins per key).
- State updates are applied on DENY (not discarded on short-circuit).
- State updates are withheld on ASK (not applied until caller approves).
- Engine loaded with ``initial_session_state`` seeds the hot cache.
"""

from __future__ import annotations

from typing import Any

import pytest

from omnigent.llms.context_window import ModelPricing
from omnigent.policies.function import FunctionPolicy
from omnigent.policies.types import EvaluationContext
from omnigent.runtime.policies.engine import PolicyEngine
from omnigent.spec.types import (
    FunctionPolicySpec,
    FunctionRef,
    Phase,
    PhaseSelector,
    PolicyAction,
    StateUpdateAction,
)
from omnigent.stores.conversation_store.sqlalchemy_store import (
    SqlAlchemyConversationStore,
)


def _state_writing_policy(
    name: str,
    state_updates: dict[str, Any],
    *,
    action: str = "ALLOW",
) -> FunctionPolicy:
    """
    Build a :class:`FunctionPolicy` that returns fixed *state_updates*.

    :param name: Policy name, e.g. ``"counter_policy"``.
    :param state_updates: Key/value pairs to return as
        :attr:`PolicyResult.state_updates`.
    :param action: Decision string passed to the result, e.g. ``"ALLOW"``.
    :returns: A :class:`FunctionPolicy` ready for engine use.
    """

    def _evaluate(_event: dict[str, Any]) -> dict[str, Any]:
        return {"result": action, "state_updates": state_updates}

    spec = FunctionPolicySpec(
        name=name,
        on=[PhaseSelector(phase=Phase.REQUEST)],
        function=FunctionRef(path="test.not.used"),
    )
    return FunctionPolicy(spec, _evaluate)


def _state_capturing_policy(bucket: dict[str, Any]) -> FunctionPolicy:
    """
    Build a :class:`FunctionPolicy` that records ``event["session_state"]``
    into *bucket* for assertion.

    :param bucket: Dict to write the captured state snapshot into under
        key ``"session_state"``.
    :returns: A capturing :class:`FunctionPolicy`.
    """

    def _evaluate(event: dict[str, Any]) -> dict[str, Any]:
        bucket["session_state"] = dict(event["session_state"])
        return {"result": "ALLOW"}

    spec = FunctionPolicySpec(
        name="capture",
        on=[PhaseSelector(phase=Phase.REQUEST)],
        function=FunctionRef(path="test.not.used"),
    )
    return FunctionPolicy(spec, _evaluate)


def _model_capturing_policy(bucket: dict[str, Any]) -> FunctionPolicy:
    """
    Build a :class:`FunctionPolicy` that records ``event["context"]["model"]``
    into *bucket* for assertion.

    :param bucket: Dict to write the captured model into under key
        ``"model"`` (present even when the value is ``None``).
    :returns: A capturing :class:`FunctionPolicy`.
    """

    def _evaluate(event: dict[str, Any]) -> dict[str, Any]:
        bucket["model"] = event["context"]["model"]
        return {"result": "ALLOW"}

    spec = FunctionPolicySpec(
        name="capture_model",
        on=[PhaseSelector(phase=Phase.TOOL_CALL)],
        function=FunctionRef(path="test.not.used"),
    )
    return FunctionPolicy(spec, _evaluate)


def _build_engine(
    store: SqlAlchemyConversationStore,
    policies: list[FunctionPolicy],
    *,
    initial_session_state: dict[str, Any] | None = None,
    initial_model: str | None = None,
) -> PolicyEngine:
    """
    Build a :class:`PolicyEngine` with a fresh conversation.

    :param store: Backing store used to create the conversation and
        handle label writes.
    :param policies: Ordered list of policies to run.
    :param initial_session_state: Seed state for the engine's hot cache,
        e.g. ``{"call_count": 3}``. ``None`` means start empty.
    :param initial_model: Seed model for the engine's model context,
        e.g. ``"opus"``. ``None`` means the model is undeterminable.
    :returns: A ready :class:`PolicyEngine`.
    """
    conv = store.create_conversation()
    return PolicyEngine(
        policies=policies,
        label_defs={},
        ask_timeout=30,
        conversation_id=conv.id,
        initial_labels={},
        initial_session_state=initial_session_state or {},
        initial_model=initial_model,
        conversation_store=store,
    )


# ── Reading session_state from the event ──────────────────────────────────────


@pytest.mark.asyncio
async def test_event_carries_session_state_key(
    conversation_store: SqlAlchemyConversationStore,
) -> None:
    """
    Function policy callables receive ``event["session_state"]`` as a
    dict. Defaults to an empty dict when no state has been written yet.

    What breaks if this fails: policies that read ``event["session_state"]``
    get a KeyError or ``None`` instead of the expected dict, breaking every
    stateful policy callable.
    """
    bucket: dict[str, Any] = {}
    engine = _build_engine(conversation_store, [_state_capturing_policy(bucket)])
    await engine.evaluate(EvaluationContext(phase=Phase.REQUEST, content="hello"))
    assert bucket["session_state"] == {}


@pytest.mark.asyncio
async def test_event_session_state_reflects_initial_seed(
    conversation_store: SqlAlchemyConversationStore,
) -> None:
    """
    When the engine is seeded with ``initial_session_state``, that state
    is visible in ``event["session_state"]`` on the first evaluation.

    What breaks if this fails: policies resuming a conversation lose the
    state accumulated in prior turns — every turn starts from scratch.
    """
    bucket: dict[str, Any] = {}
    engine = _build_engine(
        conversation_store,
        [_state_capturing_policy(bucket)],
        initial_session_state={"call_count": 3, "last_tool": "read_file"},
    )
    await engine.evaluate(EvaluationContext(phase=Phase.REQUEST, content="hi"))
    assert bucket["session_state"] == {"call_count": 3, "last_tool": "read_file"}


# ── Injecting the active model ──────────────────────────────────────────────────


@pytest.mark.asyncio
async def test_event_context_carries_injected_model(
    conversation_store: SqlAlchemyConversationStore,
) -> None:
    """
    The engine injects ``initial_model`` into ``event["context"]["model"]``.

    What breaks if this fails: model-aware policies (e.g. the cost gate's
    force-downgrade branch) can never read the active model, so they
    cannot tell an expensive model from a cheap one and the downgrade
    gate is dead.
    """
    bucket: dict[str, Any] = {}
    engine = _build_engine(
        conversation_store,
        [_model_capturing_policy(bucket)],
        initial_model="databricks-claude-opus-4-8",
    )
    await engine.evaluate(
        EvaluationContext(
            phase=Phase.TOOL_CALL,
            content={"name": "sys_os_shell", "arguments": {}},
            tool_name="sys_os_shell",
        )
    )
    assert bucket["model"] == "databricks-claude-opus-4-8"


@pytest.mark.asyncio
async def test_event_context_model_is_none_when_unseeded(
    conversation_store: SqlAlchemyConversationStore,
) -> None:
    """
    With no ``initial_model``, ``event["context"]["model"]`` is ``None``.

    What breaks if this fails: an undeterminable model would surface as a
    truthy MagicMock / missing key instead of ``None``, breaking the
    cost gate's "fail closed on unknown model" branch.
    """
    bucket: dict[str, Any] = {}
    engine = _build_engine(conversation_store, [_model_capturing_policy(bucket)])
    await engine.evaluate(
        EvaluationContext(
            phase=Phase.TOOL_CALL,
            content={"name": "sys_os_shell", "arguments": {}},
            tool_name="sys_os_shell",
        )
    )
    assert bucket["model"] is None


@pytest.mark.asyncio
async def test_caller_supplied_model_wins_over_engine_resolved(
    conversation_store: SqlAlchemyConversationStore,
) -> None:
    """
    A model already on the context is preferred over ``initial_model``.

    This is the race-free path for codex: the native hook reads the live
    model from ``config.toml`` at gate time and stamps it on the context.
    It must win over the engine's build-time resolution (which can lag
    behind a ``/model`` switch). What breaks if this fails: a user who
    downgrades via ``/model`` stays blocked because the gate keeps seeing
    the stale expensive model the engine resolved at build time.
    """
    bucket: dict[str, Any] = {}
    engine = _build_engine(
        conversation_store,
        [_model_capturing_policy(bucket)],
        initial_model="gpt-5.5",  # stale/expensive value resolved at build time
    )
    await engine.evaluate(
        EvaluationContext(
            phase=Phase.TOOL_CALL,
            content={"name": "sys_os_shell", "arguments": {}},
            tool_name="sys_os_shell",
            model="gpt-5.4",  # live value the hook read from config.toml
        )
    )
    # The hook-supplied gpt-5.4 reaches the policy, NOT the engine's gpt-5.5.
    assert bucket["model"] == "gpt-5.4"


# ── Writing state_updates ──────────────────────────────────────────────────────


@pytest.mark.asyncio
async def test_state_updates_merge_into_hot_cache(
    conversation_store: SqlAlchemyConversationStore,
) -> None:
    """
    A policy returning ``state_updates`` causes the engine's hot cache
    to reflect the merged state after evaluation.

    What breaks if this fails: ``engine.session_state`` never changes
    regardless of what policies return — state accumulation is silently
    broken.
    """
    policy = _state_writing_policy("counter", {"call_count": 1})
    engine = _build_engine(conversation_store, [policy])
    assert engine.session_state == {}
    await engine.evaluate(EvaluationContext(phase=Phase.REQUEST, content="x"))
    assert engine.session_state == {"call_count": 1}


@pytest.mark.asyncio
async def test_state_updates_shallow_merge_preserves_unmentioned_keys(
    conversation_store: SqlAlchemyConversationStore,
) -> None:
    """
    state_updates are a shallow merge: keys not mentioned in the update
    are left untouched in the hot cache.

    What breaks if this fails: a policy updating ``call_count`` silently
    wipes out ``last_tool`` and any other previously-set keys.
    """
    policy = _state_writing_policy("updater", {"call_count": 5})
    engine = _build_engine(
        conversation_store,
        [policy],
        initial_session_state={"last_tool": "write_file", "call_count": 4},
    )
    await engine.evaluate(EvaluationContext(phase=Phase.REQUEST, content="x"))
    assert engine.session_state == {"last_tool": "write_file", "call_count": 5}


@pytest.mark.asyncio
async def test_multiple_policies_accumulate_state_updates_last_write_wins(
    conversation_store: SqlAlchemyConversationStore,
) -> None:
    """
    When multiple policies in one evaluation pass return state_updates
    for the same key, the last policy in YAML order wins.

    What breaks if this fails: the first policy always wins, or updates
    are silently dropped, violating the documented merge semantics.
    """
    policy_a = _state_writing_policy("policy_a", {"x": "from_a", "y": "from_a"})
    policy_b = _state_writing_policy("policy_b", {"x": "from_b"})
    engine = _build_engine(conversation_store, [policy_a, policy_b])
    await engine.evaluate(EvaluationContext(phase=Phase.REQUEST, content="x"))
    # policy_b wins on "x"; "y" from policy_a is untouched.
    assert engine.session_state == {"x": "from_b", "y": "from_a"}


@pytest.mark.asyncio
async def test_state_updates_applied_on_deny(
    conversation_store: SqlAlchemyConversationStore,
) -> None:
    """
    state_updates from a DENYing policy are still applied — consistent
    with how set_labels are applied on DENY.

    What breaks if this fails: a policy that both denies AND records
    audit state loses its state write, breaking audit trails on
    blocked requests.
    """
    policy = _state_writing_policy("deny_and_log", {"blocked": True}, action="DENY")
    engine = _build_engine(conversation_store, [policy])
    result = await engine.evaluate(EvaluationContext(phase=Phase.REQUEST, content="x"))
    assert result.action == PolicyAction.DENY
    # State was still applied despite the DENY.
    assert engine.session_state == {"blocked": True}


@pytest.mark.asyncio
async def test_state_updates_withheld_on_ask(
    conversation_store: SqlAlchemyConversationStore,
) -> None:
    """
    state_updates from an ASKing policy are NOT applied to the hot cache —
    they are carried in the result for the caller to apply on approval only.

    What breaks if this fails: a policy that ASKs and is later denied would
    still apply its state writes, violating the §7.2 "no side effects from
    a denied ASK" invariant. The engine would accumulate state from every
    ASK, whether approved or denied.
    """
    policy = _state_writing_policy("ask_policy", {"flagged": True}, action="ASK")
    engine = _build_engine(conversation_store, [policy])
    result = await engine.evaluate(EvaluationContext(phase=Phase.REQUEST, content="x"))
    assert result.action == PolicyAction.ASK
    # State NOT applied to the hot cache — withheld pending approval.
    assert engine.session_state == {}
    # State is carried in the result for callers to apply on approval.
    assert result.state_updates is not None
    assert len(result.state_updates) == 1
    assert result.state_updates[0].key == "flagged"
    assert result.state_updates[0].action == StateUpdateAction.SET
    assert result.state_updates[0].value is True


@pytest.mark.asyncio
async def test_second_evaluation_sees_prior_state_updates(
    conversation_store: SqlAlchemyConversationStore,
) -> None:
    """
    A callable that writes state on one evaluation sees that state in
    ``event["session_state"]`` on the next evaluation of the same engine.

    What breaks if this fails: within-turn state accumulation is broken —
    each policy sees an empty dict regardless of what earlier policies wrote.
    """
    write_bucket: dict[str, Any] = {}

    def _write_then_capture(event: dict[str, Any]) -> dict[str, Any]:
        write_bucket["seen"] = dict(event["session_state"])
        return {"result": "ALLOW", "state_updates": {"step": "done"}}

    spec = FunctionPolicySpec(
        name="write_then_capture",
        on=[PhaseSelector(phase=Phase.REQUEST)],
        function=FunctionRef(path="test.not.used"),
    )
    policy = FunctionPolicy(spec, _write_then_capture)
    engine = _build_engine(conversation_store, [policy], initial_session_state={"step": "start"})

    # First evaluation — sees initial seed, writes "done".
    await engine.evaluate(EvaluationContext(phase=Phase.REQUEST, content="turn1"))
    assert write_bucket["seen"] == {"step": "start"}
    assert engine.session_state == {"step": "done"}

    # Second evaluation — sees the write from the first turn.
    await engine.evaluate(EvaluationContext(phase=Phase.REQUEST, content="turn2"))
    assert write_bucket["seen"] == {"step": "done"}


# ── Usage tracking ────────────────────────────────────────────────────────────


def _usage_capturing_policy(bucket: dict[str, Any]) -> FunctionPolicy:
    """
    Build a :class:`FunctionPolicy` that records ``event["context"]["usage"]``
    into *bucket* for assertion.

    :param bucket: Dict to write the captured usage snapshot into under
        key ``"usage"``.
    :returns: A capturing :class:`FunctionPolicy`.
    """

    def _evaluate(event: dict[str, Any]) -> dict[str, Any]:
        bucket["usage"] = dict(event["context"]["usage"])
        return {"result": "ALLOW"}

    spec = FunctionPolicySpec(
        name="usage_capture",
        on=[PhaseSelector(phase=Phase.REQUEST)],
        function=FunctionRef(path="test.not.used"),
    )
    return FunctionPolicy(spec, _evaluate)


@pytest.mark.asyncio
async def test_engine_starts_with_zero_usage(
    conversation_store: SqlAlchemyConversationStore,
) -> None:
    """
    Engine starts with all-zero usage counters when no initial_usage is
    provided.

    What breaks if this fails: the usage property would be missing keys or
    start with non-zero values, causing budget-enforcement policies to
    miscount from the first turn.
    """
    engine = _build_engine(conversation_store, [])
    assert engine.usage == {
        "input_tokens": 0,
        "output_tokens": 0,
        "total_tokens": 0,
        "cache_read_input_tokens": 0,
        "cache_creation_input_tokens": 0,
        "total_cost_usd": 0.0,
    }


@pytest.mark.asyncio
async def test_record_usage_accumulates(
    conversation_store: SqlAlchemyConversationStore,
) -> None:
    """
    After ``record_usage()`` calls, the engine's usage property reflects the
    cumulative values. Without pricing, total_cost_usd stays at 0.

    What breaks if this fails: usage counters would reset or not increment,
    making cumulative tracking useless for budget policies.
    """
    engine = _build_engine(conversation_store, [])
    engine.record_usage(input_tokens=100, output_tokens=50, total_tokens=150)
    assert engine.usage["input_tokens"] == 100
    assert engine.usage["output_tokens"] == 50
    assert engine.usage["total_tokens"] == 150
    assert engine.usage["cache_read_input_tokens"] == 0
    assert engine.usage["cache_creation_input_tokens"] == 0
    assert engine.usage["total_cost_usd"] == 0.0
    engine.record_usage(input_tokens=200, output_tokens=100, total_tokens=300)
    assert engine.usage["input_tokens"] == 300
    assert engine.usage["output_tokens"] == 150
    assert engine.usage["total_tokens"] == 450
    # Third call with cache tokens to verify accumulation.
    engine.record_usage(
        input_tokens=50,
        output_tokens=25,
        total_tokens=75,
        cache_read_input_tokens=3000,
        cache_creation_input_tokens=1000,
    )
    assert engine.usage["input_tokens"] == 350
    assert engine.usage["output_tokens"] == 175
    assert engine.usage["total_tokens"] == 525
    assert engine.usage["cache_read_input_tokens"] == 3000
    assert engine.usage["cache_creation_input_tokens"] == 1000


@pytest.mark.asyncio
async def test_record_usage_with_pricing_computes_cost(
    conversation_store: SqlAlchemyConversationStore,
) -> None:
    """
    When ``token_pricing`` (:class:`ModelPricing`) is provided,
    ``record_usage()`` computes ``total_cost_usd`` from the per-token
    rates using :func:`compute_llm_cost`.

    What breaks if this fails: budget-enforcement policies that read
    ``event["context"]["usage"]["total_cost_usd"]`` would always see 0,
    making cost-based gating impossible.
    """
    conv = conversation_store.create_conversation()
    # $3/M input, $15/M output (Claude Sonnet pricing), no cache rates.
    pricing = ModelPricing(
        input_per_token=3.0 / 1_000_000,
        output_per_token=15.0 / 1_000_000,
    )
    engine = PolicyEngine(
        policies=[],
        label_defs={},
        ask_timeout=30,
        conversation_id=conv.id,
        initial_labels={},
        token_pricing=pricing,
        conversation_store=conversation_store,
    )
    # 1000 input tokens @ $3/M = $0.003, 500 output tokens @ $15/M = $0.0075
    engine.record_usage(input_tokens=1000, output_tokens=500, total_tokens=1500)
    expected_cost = 1000 * 3.0 / 1_000_000 + 500 * 15.0 / 1_000_000
    assert engine.usage["total_cost_usd"] == pytest.approx(expected_cost)
    assert engine.usage["total_cost_usd"] == pytest.approx(0.0105)


@pytest.mark.asyncio
async def test_record_usage_with_cache_tokens_computes_cost(
    conversation_store: SqlAlchemyConversationStore,
) -> None:
    """
    When ``ModelPricing`` includes cache-read and cache-write rates,
    ``record_usage()`` prices cache tokens at their own rates instead
    of the plain input rate.

    What breaks if this fails: cache-read tokens (cheap) and cache-write
    tokens (expensive) would be priced at the plain input rate, over- or
    under-counting cost for Anthropic-style providers.
    """
    conv = conversation_store.create_conversation()
    # Anthropic-style pricing: $3/M input, $15/M output,
    # $0.30/M cache-read (0.1x), $3.75/M cache-write (1.25x).
    pricing = ModelPricing(
        input_per_token=3.0 / 1_000_000,
        output_per_token=15.0 / 1_000_000,
        cache_read_per_token=0.30 / 1_000_000,
        cache_write_per_token=3.75 / 1_000_000,
    )
    engine = PolicyEngine(
        policies=[],
        label_defs={},
        ask_timeout=30,
        conversation_id=conv.id,
        initial_labels={},
        token_pricing=pricing,
        conversation_store=conversation_store,
    )
    engine.record_usage(
        input_tokens=1000,
        output_tokens=500,
        total_tokens=1500,
        cache_read_input_tokens=5000,
        cache_creation_input_tokens=2000,
    )
    # 1000 * 3e-6 = 0.003  (non-cached input)
    # 500 * 15e-6 = 0.0075  (output)
    # 5000 * 0.3e-6 = 0.0015 (cache read)
    # 2000 * 3.75e-6 = 0.0075 (cache write)
    # total = 0.0195
    expected = (
        1000 * 3.0 / 1_000_000
        + 500 * 15.0 / 1_000_000
        + 5000 * 0.30 / 1_000_000
        + 2000 * 3.75 / 1_000_000
    )
    assert engine.usage["total_cost_usd"] == pytest.approx(expected)
    assert engine.usage["total_cost_usd"] == pytest.approx(0.0195)
    assert engine.usage["cache_read_input_tokens"] == 5000
    assert engine.usage["cache_creation_input_tokens"] == 2000


@pytest.mark.asyncio
async def test_event_context_carries_usage(
    conversation_store: SqlAlchemyConversationStore,
) -> None:
    """
    The ``event["context"]["usage"]`` dict carries the current cumulative
    token counts when a function policy is dispatched.

    What breaks if this fails: policy callables that read
    ``event["context"]["usage"]`` would get stale or empty usage data,
    breaking budget-enforcement or rate-limiting policies.
    """
    bucket: dict[str, Any] = {}
    policy = _usage_capturing_policy(bucket)
    engine = _build_engine(conversation_store, [policy])
    engine.record_usage(input_tokens=500, output_tokens=200, total_tokens=700)
    await engine.evaluate(EvaluationContext(phase=Phase.REQUEST, content="hi"))
    assert bucket["usage"]["input_tokens"] == 500
    assert bucket["usage"]["output_tokens"] == 200
    assert bucket["usage"]["total_tokens"] == 700


@pytest.mark.asyncio
async def test_record_usage_persists_to_store(
    conversation_store: SqlAlchemyConversationStore,
) -> None:
    """
    ``record_usage()`` writes the cumulative totals to the conversation's
    ``session_usage`` column so they survive across engine lifetimes.

    What breaks if this fails: usage counters reset to zero on every new
    turn because the persisted state is never written.
    """
    conv = conversation_store.create_conversation()
    engine = PolicyEngine(
        policies=[],
        label_defs={},
        ask_timeout=30,
        conversation_id=conv.id,
        initial_labels={},
        conversation_store=conversation_store,
    )
    engine.record_usage(
        input_tokens=100,
        output_tokens=50,
        total_tokens=150,
        cache_read_input_tokens=400,
        cache_creation_input_tokens=200,
    )
    reloaded = conversation_store.get_conversation(conv.id)
    assert reloaded is not None
    assert reloaded.session_usage["input_tokens"] == 100
    assert reloaded.session_usage["output_tokens"] == 50
    assert reloaded.session_usage["total_tokens"] == 150
    assert reloaded.session_usage["cache_read_input_tokens"] == 400
    assert reloaded.session_usage["cache_creation_input_tokens"] == 200
    assert reloaded.session_usage["total_cost_usd"] == 0.0