ForcePilot/backend/test/test_agent_run_service.py

from __future__ import annotations

from contextlib import asynccontextmanager
from types import SimpleNamespace

import pytest
from sqlalchemy.exc import IntegrityError

import yuxi.services.agent_run_service as agent_run_service


@pytest.mark.asyncio
async def test_stream_agent_run_events_emits_error_and_close_on_db_error(monkeypatch: pytest.MonkeyPatch):
    @asynccontextmanager
    async def fake_session_ctx():
        yield object()

    class BrokenRepo:
        def __init__(self, db):
            self.db = db

        async def get_run_for_user(self, run_id: str, user_id: str):
            del run_id, user_id
            raise RuntimeError("db down")

    monkeypatch.setattr(agent_run_service.pg_manager, "get_async_session_context", fake_session_ctx)
    monkeypatch.setattr(agent_run_service, "AgentRunRepository", BrokenRepo)

    chunks = []
    async for chunk in agent_run_service.stream_agent_run_events(
        run_id="run-1",
        after_seq="0",
        current_user_id="1",
    ):
        chunks.append(chunk)

    assert len(chunks) == 2
    assert chunks[0].startswith("event: error")
    assert '"reason": "db_error"' in chunks[0]
    assert chunks[1].startswith("event: close")


@pytest.mark.asyncio
async def test_stream_agent_run_events_reads_redis_and_close_terminal(monkeypatch: pytest.MonkeyPatch):
    @asynccontextmanager
    async def fake_session_ctx():
        yield object()

    class Repo:
        def __init__(self, db):
            self.db = db

        async def get_run_for_user(self, run_id: str, user_id: str):
            del run_id, user_id
            return SimpleNamespace(status="completed")

    calls = {"count": 0}

    async def fake_list_events(run_id: str, *, after_seq: str, limit: int):
        del run_id, after_seq, limit
        calls["count"] += 1
        if calls["count"] == 1:
            return [
                {
                    "seq": "1700000000000-0",
                    "event_type": "loading",
                    "payload": {"items": [{"status": "loading", "response": "你"}]},
                    "ts": 1700000000000,
                }
            ]
        return []

    async def fake_last_seq(run_id: str):
        del run_id
        return "1700000000000-0"

    monkeypatch.setattr(agent_run_service.pg_manager, "get_async_session_context", fake_session_ctx)
    monkeypatch.setattr(agent_run_service, "AgentRunRepository", Repo)
    monkeypatch.setattr(agent_run_service, "list_run_stream_events", fake_list_events)
    monkeypatch.setattr(agent_run_service, "get_last_run_stream_seq", fake_last_seq)
    monkeypatch.setattr(agent_run_service, "SSE_POLL_INTERVAL_SECONDS", 0)

    chunks = []
    async for chunk in agent_run_service.stream_agent_run_events(
        run_id="run-1",
        after_seq="0",
        current_user_id="1",
    ):
        chunks.append(chunk)

    assert any(item.startswith("event: loading") for item in chunks)
    assert chunks[-1].startswith("event: close")
    assert '"last_seq": "1700000000000-0"' in chunks[-1]


@pytest.mark.asyncio
async def test_create_agent_run_commits_before_enqueue(monkeypatch: pytest.MonkeyPatch):
    class FakeDB:
        def __init__(self):
            self.order: list[str] = []
            self.committed = False

        async def commit(self):
            self.order.append("commit")
            self.committed = True

        async def rollback(self):
            raise AssertionError("rollback should not be called")

    db = FakeDB()
    created_run = SimpleNamespace(
        id="run-1",
        thread_id="thread-1",
        status="pending",
        request_id="req-1",
        user_id="1",
    )

    class Repo:
        def __init__(self, db_session):
            self.db = db_session

        async def get_run_by_request_id(self, request_id: str):
            del request_id
            return None

        async def create_run(self, **kwargs):
            assert kwargs["request_id"] == "req-1"
            return created_run

    class ConvRepo:
        def __init__(self, db_session):
            self.db = db_session

        async def get_conversation_by_thread_id(self, thread_id: str):
            del thread_id
            return SimpleNamespace(user_id="1", status="active")

    class Queue:
        async def enqueue_job(self, job_name: str, run_id: str, _job_id: str):
            assert job_name == "process_agent_run"
            assert run_id == "run-1"
            assert _job_id == "run:run-1"
            db.order.append("enqueue")
            assert db.committed is True

    async def fake_get_arq_pool():
        return Queue()

    monkeypatch.setattr(agent_run_service.agent_manager, "get_agent", lambda agent_id: object())
    monkeypatch.setattr(agent_run_service, "ConversationRepository", ConvRepo)
    monkeypatch.setattr(agent_run_service, "AgentRunRepository", Repo)
    monkeypatch.setattr(agent_run_service, "get_arq_pool", fake_get_arq_pool)

    result = await agent_run_service.create_agent_run_view(
        agent_id="ChatbotAgent",
        query="hello",
        config={"thread_id": "thread-1", "request_id": "req-1"},
        image_content=None,
        current_user_id="1",
        db=db,
    )

    assert db.order == ["commit", "enqueue"]
    assert result["run_id"] == "run-1"
    assert result["request_id"] == "req-1"


@pytest.mark.asyncio
async def test_create_agent_run_handles_integrity_error_with_same_user_existing(monkeypatch: pytest.MonkeyPatch):
    class FakeDB:
        def __init__(self):
            self.commit_called = 0
            self.rollback_called = 0

        async def commit(self):
            self.commit_called += 1
            raise IntegrityError("insert", {"request_id": "req-1"}, Exception("duplicate"))

        async def rollback(self):
            self.rollback_called += 1

    db = FakeDB()
    existing_run = SimpleNamespace(
        id="run-existing",
        thread_id="thread-1",
        status="running",
        request_id="req-1",
        user_id="1",
    )
    state = {"lookup_count": 0}

    class Repo:
        def __init__(self, db_session):
            self.db = db_session

        async def get_run_by_request_id(self, request_id: str):
            del request_id
            state["lookup_count"] += 1
            if state["lookup_count"] == 1:
                return None
            return existing_run

        async def create_run(self, **kwargs):
            del kwargs
            return SimpleNamespace(
                id="run-new",
                thread_id="thread-1",
                status="pending",
                request_id="req-1",
                user_id="1",
            )

    class ConvRepo:
        def __init__(self, db_session):
            self.db = db_session

        async def get_conversation_by_thread_id(self, thread_id: str):
            del thread_id
            return SimpleNamespace(user_id="1", status="active")

    async def fake_get_arq_pool():
        raise AssertionError("should not enqueue on integrity fallback")

    monkeypatch.setattr(agent_run_service.agent_manager, "get_agent", lambda agent_id: object())
    monkeypatch.setattr(agent_run_service, "ConversationRepository", ConvRepo)
    monkeypatch.setattr(agent_run_service, "AgentRunRepository", Repo)
    monkeypatch.setattr(agent_run_service, "get_arq_pool", fake_get_arq_pool)

    result = await agent_run_service.create_agent_run_view(
        agent_id="ChatbotAgent",
        query="hello",
        config={"thread_id": "thread-1", "request_id": "req-1"},
        image_content=None,
        current_user_id="1",
        db=db,
    )

    assert db.commit_called == 1
    assert db.rollback_called == 1
    assert result["run_id"] == "run-existing"
    assert result["status"] == "running"


@pytest.mark.asyncio
async def test_create_agent_run_integrity_error_returns_409_for_other_user(monkeypatch: pytest.MonkeyPatch):
    class FakeDB:
        async def commit(self):
            raise IntegrityError("insert", {"request_id": "req-1"}, Exception("duplicate"))

        async def rollback(self):
            return None

    db = FakeDB()
    existing_run = SimpleNamespace(
        id="run-existing",
        thread_id="thread-1",
        status="pending",
        request_id="req-1",
        user_id="2",
    )
    state = {"lookup_count": 0}

    class Repo:
        def __init__(self, db_session):
            self.db = db_session

        async def get_run_by_request_id(self, request_id: str):
            del request_id
            state["lookup_count"] += 1
            if state["lookup_count"] == 1:
                return None
            return existing_run

        async def create_run(self, **kwargs):
            del kwargs
            return SimpleNamespace(
                id="run-new",
                thread_id="thread-1",
                status="pending",
                request_id="req-1",
                user_id="1",
            )

    class ConvRepo:
        def __init__(self, db_session):
            self.db = db_session

        async def get_conversation_by_thread_id(self, thread_id: str):
            del thread_id
            return SimpleNamespace(user_id="1", status="active")

    monkeypatch.setattr(agent_run_service.agent_manager, "get_agent", lambda agent_id: object())
    monkeypatch.setattr(agent_run_service, "ConversationRepository", ConvRepo)
    monkeypatch.setattr(agent_run_service, "AgentRunRepository", Repo)

    with pytest.raises(agent_run_service.HTTPException) as exc:
        await agent_run_service.create_agent_run_view(
            agent_id="ChatbotAgent",
            query="hello",
            config={"thread_id": "thread-1", "request_id": "req-1"},
            image_content=None,
            current_user_id="1",
            db=db,
        )

    assert exc.value.status_code == 409
    assert exc.value.detail == "request_id 冲突"