simbarag/tests/unit/test_agent_runtime.py

import asyncio
from types import SimpleNamespace
from unittest.mock import AsyncMock

import pytest
from langchain_core.messages import AIMessage
from quart import Quart

from blueprints.conversation import runtime
from metrics import AGENT_REQUESTS_TOTAL


@pytest.mark.parametrize(
    "channel", ["web", "email", "imessage", "whatsapp", "scheduled", "cron"]
)
@pytest.mark.parametrize("admin", [False, True])
async def test_channels_share_identity_memories_permissions_and_summary(
    monkeypatch, channel, admin
):
    app = Quart(__name__)
    user = SimpleNamespace(id="owner", is_admin=lambda: admin)
    app.main_agent = SimpleNamespace(
        ainvoke=AsyncMock(return_value={"messages": [AIMessage(content="answer")]}),
        aget_state=AsyncMock(return_value=SimpleNamespace(values={"messages": []})),
    )
    app.admin_agent = SimpleNamespace(
        ainvoke=AsyncMock(
            return_value={"messages": [AIMessage(content="admin answer")]}
        ),
        aget_state=AsyncMock(return_value=SimpleNamespace(values={"messages": []})),
    )
    summary = AIMessage(content="Saved constraints")
    app.conversation_summary_store = SimpleNamespace(
        update=AsyncMock(return_value=SimpleNamespace(as_message=lambda: summary))
    )
    prompt = AsyncMock(return_value="Remember the budget")
    monkeypatch.setattr(runtime, "build_system_prompt_with_memories", prompt)
    before = AGENT_REQUESTS_TOTAL.labels(channel=channel, result="success")._value.get()
    async with app.app_context():
        response = await runtime.invoke_agent_for_user(
            user, "What is my budget?", "thread", channel=channel
        )
    agent = app.admin_agent if admin else app.main_agent
    assert response["messages"][-1].content == ("admin answer" if admin else "answer")
    config = agent.ainvoke.call_args.kwargs["config"]["configurable"]
    assert config["user_id"] == "owner"
    assert config["is_admin"] is admin
    assert config["thread_id"] == "thread"
    assert config["channel"] == channel
    assert config["system_prompt"] == "Remember the budget"
    assert config["conversation_summary"] is summary
    assert config["request_id"]
    prompt.assert_awaited_once_with("owner", query="What is my budget?")
    assert (
        AGENT_REQUESTS_TOTAL.labels(channel=channel, result="success")._value.get()
        == before + 1
    )


async def test_anonymous_invocation_is_rejected():
    with pytest.raises(ValueError, match="authenticated user"):
        await runtime.invoke_agent_for_user(None, "hello", "thread")


@pytest.mark.parametrize(
    "error, outcome",
    [(RuntimeError("failed"), "error"), (asyncio.CancelledError(), "cancelled")],
)
async def test_shared_invocation_measures_failed_and_cancelled_generation(
    error, outcome
):
    agent = SimpleNamespace(ainvoke=AsyncMock(side_effect=error))
    turn = runtime.PreparedTurn(
        agent,
        {"messages": []},
        {
            "configurable": {
                "channel": "web",
                "request_id": "test",
                "thread_id": "t",
                "user_id": "u",
            }
        },
    )
    before = AGENT_REQUESTS_TOTAL.labels(channel="web", result=outcome)._value.get()
    with pytest.raises(type(error)):
        await runtime.invoke_prepared(turn)
    assert (
        AGENT_REQUESTS_TOTAL.labels(channel="web", result=outcome)._value.get()
        == before + 1
    )


async def test_concurrent_turns_on_one_thread_keep_every_exchange(
    monkeypatch, tmp_path
):
    import aiosqlite
    from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
    from langgraph.graph import END, START, MessagesState, StateGraph

    monkeypatch.setattr(
        runtime, "build_system_prompt_with_memories", AsyncMock(return_value="sys")
    )
    async with aiosqlite.connect(tmp_path / "checkpoints.db") as conn:
        saver = AsyncSqliteSaver(conn)
        await saver.setup()

        async def answer(state):
            # Yield so an unserialized second turn would read stale history.
            await asyncio.sleep(0.05)
            return {
                "messages": [
                    AIMessage(content="reply " + state["messages"][-1].content)
                ]
            }

        graph = StateGraph(MessagesState)
        graph.add_node("answer", answer)
        graph.add_edge(START, "answer")
        graph.add_edge("answer", END)
        agent = graph.compile(checkpointer=saver)

        app = Quart(__name__)
        app.main_agent = app.admin_agent = agent
        user = SimpleNamespace(id="owner", is_admin=lambda: False)
        async with app.app_context():
            await asyncio.gather(
                *[
                    runtime.invoke_agent_for_user(user, q, "same-thread")
                    for q in ("first", "second")
                ]
            )
            state = await agent.aget_state(
                {"configurable": {"thread_id": "same-thread"}}
            )
    contents = [m.content for m in state.values["messages"]]
    assert sorted(contents) == ["first", "reply first", "reply second", "second"]