omnigent-ai--omnigent/tests/server/integration/test_policy_ask_lifecycle_e2e.py

"""
End-to-end integration test for the ASK policy approve/refuse lifecycle.

Exercises the full user journey:

1. Create a session.
2. Attach an ASK policy via ``POST /v1/sessions/{session_id}/policies``
   using the registered ``ask_on_os_tools`` handler.
3. Trigger policy evaluation via ``POST /v1/sessions/{id}/policies/evaluate``
   with a Bash tool call.
4. Observe the parked elicitation in the session snapshot.
5. Resolve with accept (approve) or decline (refuse) via
   ``POST /v1/sessions/{id}/elicitations/{eid}/resolve``.
6. Assert the evaluate endpoint returns ``POLICY_ACTION_ALLOW`` (approve)
   or ``POLICY_ACTION_DENY`` (refuse).

Uses the shared ``runtime_init`` and ``mock_llm`` fixtures from
``tests/server/conftest.py``, plus a per-module policy-enabled app
fixture that wires in a :class:`SqlAlchemyPolicyStore` so the session
policy CRUD routes are mounted and the evaluate endpoint picks up
session-attached policies.
"""

from __future__ import annotations

import asyncio
from collections.abc import AsyncIterator
from pathlib import Path
from typing import Any

import httpx
import pytest
import pytest_asyncio
from fastapi import FastAPI

from omnigent.runtime import pending_elicitations, session_stream
from omnigent.runtime.agent_cache import AgentCache
from omnigent.server.app import create_app
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.conversation_store.sqlalchemy_store import (
    SqlAlchemyConversationStore,
)
from omnigent.stores.file_store.sqlalchemy_store import SqlAlchemyFileStore
from omnigent.stores.policy_store.sqlalchemy_store import SqlAlchemyPolicyStore
from tests.server.conftest import ControllableMockClient
from tests.server.helpers import create_test_agent

pytestmark = pytest.mark.asyncio


# ── Fixtures ────────────────────────────────────────────────


@pytest.fixture()
def policy_app(runtime_init: None, db_uri: str, tmp_path: Path) -> FastAPI:
    """
    FastAPI app with a policy store wired in.

    The standard ``app`` fixture from ``conftest.py`` omits the policy
    store, so the session policy CRUD routes are not mounted. This
    fixture adds one so ``POST /v1/sessions/{id}/policies`` and the
    evaluate endpoint's ``get_policy_store()`` both see session-attached
    policies. The runtime-global ``_policy_store`` is wired separately
    in the ``client`` fixture via monkeypatch.

    :param runtime_init: Fixture that initializes the runtime with a
        mock LLM.
    :param db_uri: Test database URI.
    :param tmp_path: Pytest temporary directory fixture.
    """
    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),
        policy_store=SqlAlchemyPolicyStore(db_uri),
    )


@pytest_asyncio.fixture()
async def client(
    policy_app: FastAPI,
    mock_llm: ControllableMockClient,
    tmp_path: Path,
    db_uri: str,
    monkeypatch: pytest.MonkeyPatch,
) -> AsyncIterator[httpx.AsyncClient]:
    """
    Async HTTP client wired to the policy-enabled app.

    Mirrors the shared ``client`` fixture but targets the
    ``policy_app`` so session policy CRUD routes are available.
    Also patches the runtime's ``_policy_store`` global so the
    evaluate endpoint picks up session-attached policies.

    :param policy_app: The policy-enabled FastAPI app.
    :param mock_llm: Controllable mock LLM (released on teardown).
    :param tmp_path: Pytest temporary directory fixture.
    :param db_uri: Test database URI.
    :param monkeypatch: Pytest monkeypatch fixture.
    """
    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)

    # Wire the policy store into the runtime global so
    # ``get_policy_store()`` returns it during evaluate.
    from omnigent.runtime import _globals

    monkeypatch.setattr(_globals, "_policy_store", SqlAlchemyPolicyStore(db_uri))

    transport = httpx.ASGITransport(app=policy_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()


# ── Helpers ─────────────────────────────────────────────────


async def _create_session(client: httpx.AsyncClient, agent_id: str) -> str:
    """
    Create a session bound to an agent.

    :param client: Test HTTP client.
    :param agent_id: Agent to bind.
    :returns: New session id.
    """
    resp = await client.post("/v1/sessions", json={"agent_id": agent_id})
    assert resp.status_code == 201, f"create failed: {resp.status_code} {resp.text}"
    return resp.json()["id"]


async def _attach_ask_policy(
    client: httpx.AsyncClient,
    session_id: str,
) -> str:
    """
    Attach the registered ``ask_on_os_tools`` ASK policy to a session.

    This builtin policy ASKs for approval on any file or shell tool
    call (Bash, Read, Write, Edit, Glob, Grep). Using a registered
    handler avoids the policy registry allowlist rejection.

    :param client: Test HTTP client.
    :param session_id: Session to attach the policy to.
    :returns: The created policy id.
    """
    resp = await client.post(
        f"/v1/sessions/{session_id}/policies",
        json={
            "name": "test_ask_policy",
            "type": "python",
            "handler": "omnigent.policies.builtins.safety.ask_on_os_tools",
        },
    )
    assert resp.status_code == 200, f"attach policy failed: {resp.status_code} {resp.text}"
    body = resp.json()
    assert body["name"] == "test_ask_policy"
    return body["id"]


def _tool_call_request(
    tool_name: str = "Bash",
    arguments: dict[str, Any] | None = None,
) -> dict[str, Any]:
    """
    Build a ``PHASE_TOOL_CALL`` policy-evaluate request.

    :param tool_name: Tool name, e.g. ``"Bash"``.
    :param arguments: Tool arguments dict.
    :returns: JSON body for ``POST /v1/sessions/{id}/policies/evaluate``.
    """
    return {
        "event": {
            "type": "PHASE_TOOL_CALL",
            "target": "",
            "data": {
                "name": tool_name,
                "arguments": arguments or {},
            },
            "context": {},
        },
    }


async def _drain_elicitation_id(
    session_id: str,
    *,
    timeout_s: float = 5.0,
) -> str:
    """
    Block on the session SSE stream until a
    ``response.elicitation_request`` arrives; return its id.

    :param session_id: Session to subscribe to.
    :param timeout_s: Max seconds to wait before failing the test.
    :returns: The published ``elicitation_id``.
    """
    async with asyncio.timeout(timeout_s):
        async for event in session_stream.subscribe(session_id):
            if event.get("type") == "response.elicitation_request":
                eid = event.get("elicitation_id")
                assert isinstance(eid, str) and eid, f"missing id: {event!r}"
                return eid
    raise AssertionError("subscribe loop ended without an elicitation event")


# ── Tests ───────────────────────────────────────────────────


async def test_ask_policy_approve_flow(
    client: httpx.AsyncClient,
) -> None:
    """
    Attach ASK policy, evaluate, approve → ALLOW.

    Full journey: create session → attach ``ask_on_os_tools`` policy
    → trigger evaluate with a Bash tool call → observe pending
    elicitation in the session snapshot → resolve with accept →
    evaluate returns ``POLICY_ACTION_ALLOW``. Proves the session-
    attached policy fires, parks a real server-side Future, and the
    URL-based resolve wakes it with the correct verdict.
    """
    agent = await create_test_agent(client, "test-ask-approve")
    session_id = await _create_session(client, agent["id"])
    await _attach_ask_policy(client, session_id)

    drain = asyncio.create_task(_drain_elicitation_id(session_id))
    evaluate = None
    try:
        await asyncio.sleep(0.05)

        # The evaluate POST parks until the verdict arrives.
        evaluate = asyncio.create_task(
            client.post(
                f"/v1/sessions/{session_id}/policies/evaluate",
                json=_tool_call_request("Bash"),
            )
        )

        # Learn the elicitation id from the stream.
        elicitation_id = await drain

        # Verify the session snapshot shows a pending elicitation.
        snapshot = await client.get(f"/v1/sessions/{session_id}")
        assert snapshot.status_code == 200, snapshot.text
        pending = snapshot.json().get("pending_elicitations", [])
        pending_ids = [p["elicitation_id"] for p in pending]
        assert elicitation_id in pending_ids, (
            f"elicitation {elicitation_id} not in snapshot pending list: {pending_ids}"
        )

        # Approve.
        verdict = await client.post(
            f"/v1/sessions/{session_id}/elicitations/{elicitation_id}/resolve",
            json={"action": "accept"},
        )
        assert verdict.status_code == 202, verdict.text

        # The parked evaluate call should now return ALLOW.
        resp = await evaluate
        assert resp.status_code == 200, resp.text
        assert resp.json()["result"] == "POLICY_ACTION_ALLOW"
    finally:
        for task in [drain, evaluate]:
            if task is not None and not task.done():
                task.cancel()
                await asyncio.gather(task, return_exceptions=True)
        pending_elicitations.reset_for_tests()


async def test_ask_policy_refuse_flow(
    client: httpx.AsyncClient,
) -> None:
    """
    Attach ASK policy, evaluate, refuse → DENY.

    Same setup as the approve flow but resolves with ``decline``.
    The evaluate endpoint must collapse the ASK to
    ``POLICY_ACTION_DENY`` — fail-closed. Proves the session-attached
    ASK policy's refuse path terminates correctly and the DENY
    sentinel propagates.
    """
    agent = await create_test_agent(client, "test-ask-refuse")
    session_id = await _create_session(client, agent["id"])
    await _attach_ask_policy(client, session_id)

    drain = asyncio.create_task(_drain_elicitation_id(session_id))
    evaluate = None
    try:
        await asyncio.sleep(0.05)

        evaluate = asyncio.create_task(
            client.post(
                f"/v1/sessions/{session_id}/policies/evaluate",
                json=_tool_call_request("Bash"),
            )
        )

        elicitation_id = await drain

        # Refuse.
        verdict = await client.post(
            f"/v1/sessions/{session_id}/elicitations/{elicitation_id}/resolve",
            json={"action": "decline"},
        )
        assert verdict.status_code == 202, verdict.text

        # The parked evaluate call should now return DENY.
        resp = await evaluate
        assert resp.status_code == 200, resp.text
        body = resp.json()
        assert body["result"] == "POLICY_ACTION_DENY"
    finally:
        for task in [drain, evaluate]:
            if task is not None and not task.done():
                task.cancel()
                await asyncio.gather(task, return_exceptions=True)
        pending_elicitations.reset_for_tests()