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

"""
Tests for the V0 event dict that FunctionPolicy callables receive.

With the V0 Service Policy interface, callables receive an ``event``
dict (and optionally a ``config`` dict) instead of the old
``(ctx, context)`` pair. Verifies:

- Event dict carries the expected V0 shape (type, target, data,
  context).
- The engine's label hot cache is exposed read-only to callables via
  ``event["context"]["labels"]`` (the advisor cost-plan guard gates on
  it), and mutating the exposed copy never corrupts engine state.
- Engine labels still accumulate correctly across evaluations
  (verified via ``engine.labels``).
- Policies see each other's set_labels writes via the engine's
  ``apply_label_writes`` (sequential across evaluations).
"""

from __future__ import annotations

from typing import Any

import pytest

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,
)
from omnigent.stores.conversation_store.sqlalchemy_store import (
    SqlAlchemyConversationStore,
)
from tests.runtime.policies.conftest import make_fixed_policy


def _capturing_policy(bucket: dict[str, Any]) -> FunctionPolicy:
    """Build a FunctionPolicy that records the V0 event it
    receives into *bucket*. Used to inspect what the engine
    passed at evaluate time."""

    def _evaluate(event: dict[str, Any]) -> dict[str, Any]:
        bucket["event"] = dict(event)  # copy to capture snapshot
        return {"result": "ALLOW"}

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


def _build(
    store: SqlAlchemyConversationStore,
    policies: list,
    *,
    initial_labels: dict[str, str] | None = None,
) -> PolicyEngine:
    """Build engine + fresh conversation."""
    conv = store.create_conversation()
    return PolicyEngine(
        policies=policies,
        label_defs={},
        ask_timeout=30,
        conversation_id=conv.id,
        initial_labels=initial_labels or {},
        conversation_store=store,
    )


# ── V0 event shape ──


@pytest.mark.asyncio
async def test_event_carries_v0_shape(
    conversation_store: SqlAlchemyConversationStore,
) -> None:
    """FunctionPolicy callable receives a V0-shaped event dict
    with type, target, data, and context keys."""
    bucket: dict[str, Any] = {}
    policy = _capturing_policy(bucket)
    engine = _build(
        conversation_store,
        [policy],
        initial_labels={"integrity": "1"},
    )
    await engine.evaluate(EvaluationContext(phase=Phase.REQUEST, content="hello"))
    event = bucket["event"]
    assert event["type"] == "request"
    assert event["target"] is None  # REQUEST phase has no tool_name
    assert event["data"] == "hello"
    assert "actor" in event["context"]


@pytest.mark.asyncio
async def test_event_context_usage_carries_total_cost_usd(
    conversation_store: SqlAlchemyConversationStore,
) -> None:
    """``event["context"]["usage"]`` carries ``total_cost_usd``.

    A cost-budget policy reads
    ``event["context"]["usage"]["total_cost_usd"]``; this pins that the
    engine seeds it and the function adapter forwards it, so the
    ``UsageContext.total_cost_usd`` schema field reflects what actually
    flows at runtime. If the adapter dropped the key (or the engine
    default omitted it), the policy would see no cost and never fire —
    the assertion below would KeyError / mismatch.
    """
    bucket: dict[str, Any] = {}
    policy = _capturing_policy(bucket)
    conv = conversation_store.create_conversation()
    engine = PolicyEngine(
        policies=[policy],
        label_defs={},
        ask_timeout=30,
        conversation_id=conv.id,
        initial_labels={},
        initial_usage={
            "input_tokens": 1000,
            "output_tokens": 200,
            "total_tokens": 1200,
            "total_cost_usd": 0.42,
        },
        conversation_store=conversation_store,
    )
    await engine.evaluate(EvaluationContext(phase=Phase.REQUEST, content="hi"))
    usage = bucket["event"]["context"]["usage"]
    # 0.42 is the seeded session cost — proves the priced figure reaches
    # the policy, not just the token counts.
    assert usage["total_cost_usd"] == 0.42
    assert usage["input_tokens"] == 1000


@pytest.mark.asyncio
async def test_event_tool_call_carries_target(
    conversation_store: SqlAlchemyConversationStore,
) -> None:
    """On TOOL_CALL phase, event.target is the tool_name."""
    bucket: dict[str, Any] = {}

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

    spec = FunctionPolicySpec(
        name="capture",
        on=[PhaseSelector(phase=Phase.TOOL_CALL)],
        function=FunctionRef(path="test.not.used"),
    )
    policy = FunctionPolicy(spec, _evaluate)
    engine = _build(conversation_store, [policy])
    await engine.evaluate(
        EvaluationContext(
            phase=Phase.TOOL_CALL,
            content={"name": "web_search", "arguments": {"q": "test"}},
            tool_name="web_search",
        ),
    )
    event = bucket["event"]
    assert event["type"] == "tool_call"
    assert event["target"] == "web_search"
    assert event["data"] == {"name": "web_search", "arguments": {"q": "test"}}


@pytest.mark.asyncio
async def test_event_context_carries_labels_snapshot(
    conversation_store: SqlAlchemyConversationStore,
) -> None:
    """``event["context"]["labels"]`` carries the engine's label cache.

    The advisor cost-plan guard reads ``cost_control.plan`` from this
    field; if the engine stopped injecting labels (or the function
    adapter dropped the key), the guard would silently see no plan and
    never enforce — the value assertion below would fail.
    """
    bucket: dict[str, Any] = {}
    policy = _capturing_policy(bucket)
    engine = _build(
        conversation_store,
        [policy],
        initial_labels={"cost_control.plan": '{"v": 1}', "integrity": "1"},
    )
    await engine.evaluate(EvaluationContext(phase=Phase.REQUEST, content="hi"))
    # Exact-mapping assertion: both seeded labels arrive, nothing else.
    assert bucket["event"]["context"]["labels"] == {
        "cost_control.plan": '{"v": 1}',
        "integrity": "1",
    }
    # The exposed dict is a defensive copy — mutating it must not
    # corrupt the engine's hot cache (a shared reference would).
    bucket["event"]["context"]["labels"]["integrity"] = "0"
    assert engine.labels["integrity"] == "1"


# ── Engine labels still work correctly ──


@pytest.mark.asyncio
async def test_engine_labels_accumulate_across_evaluations(
    conversation_store: SqlAlchemyConversationStore,
) -> None:
    """Engine's hot cache reflects label writes from prior
    evaluations. A subsequent evaluation sees the effects of a
    prior policy's set_labels."""
    bucket_1: dict[str, Any] = {}
    bucket_2: dict[str, Any] = {}

    policy_1 = _capturing_policy(bucket_1)
    policy_2 = _capturing_policy(bucket_2)

    engine = _build(
        conversation_store,
        [policy_1],
        initial_labels={"integrity": "1"},
    )

    # First evaluation.
    await engine.evaluate(EvaluationContext(phase=Phase.REQUEST, content="x"))
    assert engine.labels == {"integrity": "1"}

    # Apply a write outside evaluate — bumps the hot cache.
    engine.apply_label_writes({"integrity": "0"})

    # Swap in policy_2 and re-evaluate.
    engine.policies = [policy_2]
    await engine.evaluate(EvaluationContext(phase=Phase.REQUEST, content="x"))
    # Engine's labels reflect the update.
    assert engine.labels == {"integrity": "0"}


@pytest.mark.asyncio
async def test_label_writer_then_reader_in_same_evaluation(
    conversation_store: SqlAlchemyConversationStore,
) -> None:
    """A fixed policy writes integrity=0; a later FunctionPolicy in
    the same evaluate() call sees the engine's initial label state
    (context is built once). After evaluate, engine reflects the
    accumulated writes."""
    bucket: dict[str, Any] = {}

    # Policy 1: FunctionPolicy writes integrity=0.
    writer = make_fixed_policy(
        name="writer",
        on=[PhaseSelector(phase=Phase.REQUEST)],
        action=PolicyAction.ALLOW,
        set_labels={"integrity": "0"},
    )

    # Policy 2: FunctionPolicy captures event.
    reader = _capturing_policy(bucket)

    engine = _build(
        conversation_store,
        [writer, reader],
        initial_labels={"integrity": "1"},
    )
    await engine.evaluate(EvaluationContext(phase=Phase.REQUEST, content="x"))

    # After evaluate, the cache reflects the accumulated writes.
    assert engine.labels == {"integrity": "0"}
    # The V0 event was received.
    assert bucket["event"]["type"] == "request"