SupplyAgent/tests/supply_agent/test_find_agent_v2_run_timeout.py

from __future__ import annotations

import asyncio
from collections.abc import Generator
from contextlib import contextmanager
from datetime import datetime, timedelta

import pytest
from sqlalchemy import create_engine, select
from sqlalchemy.orm import Session, sessionmaker

from api.services import find_agent_v2 as api_service
from find_agent_v2 import service as find_agent_service
from find_agent_v2.agent import FindAgentV2
from find_agent_v2.models import (
    FindAgentV2Candidate,
    FindAgentV2Evidence,
    FindAgentV2Run,
    FindAgentV2Search,
)
from find_agent_v2.observability import NullObserver
from find_agent_v2.service import FindAgentV2Service, RUN_TIMEOUT_REASON


def test_service_progress_and_work_queues_are_not_limited_to_100(monkeypatch) -> None:
    engine = create_engine("sqlite+pysqlite:///:memory:")
    FindAgentV2Run.__table__.create(engine)
    FindAgentV2Search.__table__.create(engine)
    FindAgentV2Candidate.__table__.create(engine)
    FindAgentV2Evidence.__table__.create(engine)
    factory = sessionmaker(bind=engine, autoflush=False, autocommit=False)

    @contextmanager
    def session_scope() -> Generator[Session, None, None]:
        session = factory()
        try:
            yield session
            session.commit()
        finally:
            session.close()

    monkeypatch.setattr(find_agent_service, "get_session", session_scope)
    with factory.begin() as session:
        session.add(FindAgentV2Run(
            id=1,
            run_id="large-run",
            demand_word="批量测试",
            input_json='{"user_input":"test"}',
            rule_config_json="{}",
            status="running",
            candidate_count=137,
        ))
        candidates = []
        for candidate_id in range(1, 138):
            if candidate_id <= 50:
                detail_status, portrait_status, bucket = "pending", "pending", "pending_evaluation"
            elif candidate_id <= 90:
                detail_status, portrait_status, bucket = "success", "success", "pending_evaluation"
            elif candidate_id <= 120:
                detail_status, portrait_status, bucket = "failed", "success", "pending_evaluation"
            elif candidate_id <= 130:
                detail_status, portrait_status, bucket = "success", "success", "primary"
            else:
                detail_status, portrait_status, bucket = "success", "success", "rejected"
            candidates.append(FindAgentV2Candidate(
                id=candidate_id,
                run_id="large-run",
                aweme_id=f"video-{candidate_id}",
                detail_status=detail_status,
                portrait_status=portrait_status,
                decision_bucket=bucket,
            ))
        session.add_all(candidates)

    service = FindAgentV2Service()
    progress = service.get_candidate_progress("large-run")

    assert progress == {
        "total_count": 137,
        "pending_count": 120,
        "primary_count": 10,
        "rejected_count": 7,
        "detail_pending_count": 50,
        "detail_success_count": 40,
        "detail_failed_count": 30,
        "portrait_pending_count": 50,
        "portrait_success_count": 70,
        "portrait_failed_count": 0,
        "evidence_completed_count": 70,
        "evidence_success_count": 40,
    }
    assert service.snapshot("large-run").candidate_count == 137
    assert service.count_pending_candidates("large-run") == 120
    assert len(service.get_full_state("large-run")["candidates"]) == 100
    assert len(service.list_pending_evidence_ids("large-run", "detail", limit=16)) == 16
    assert len(service.list_ready_evaluation_ids("large-run", limit=100)) == 70
    report = service.get_report_state("large-run")
    assert len(report["primary_candidates"]) == 10
    assert report["summary"]["total_count"] == 137


def test_missing_evidence_response_marks_requested_candidate_failed(monkeypatch) -> None:
    engine = create_engine("sqlite+pysqlite:///:memory:")
    FindAgentV2Candidate.__table__.create(engine)
    FindAgentV2Evidence.__table__.create(engine)
    factory = sessionmaker(bind=engine, autoflush=False, autocommit=False)
    next_evidence_id = 1

    @contextmanager
    def session_scope() -> Generator[Session, None, None]:
        nonlocal next_evidence_id
        session = factory()
        try:
            yield session
            for instance in session.new:
                if isinstance(instance, FindAgentV2Evidence) and instance.id is None:
                    instance.id = next_evidence_id
                    next_evidence_id += 1
            session.commit()
        finally:
            session.close()

    monkeypatch.setattr(find_agent_service, "get_session", session_scope)
    with factory.begin() as session:
        session.add_all([
            FindAgentV2Candidate(
                id=candidate_id,
                run_id="partial-response",
                aweme_id=f"video-{candidate_id}",
                detail_status="pending",
                portrait_status="pending",
                decision_bucket="pending_evaluation",
            )
            for candidate_id in (1, 2)
        ])

    service = FindAgentV2Service()
    service.save_details(
        "partial-response",
        [{"content_id": "video-2", "title": "success"}],
        [],
        requested_aweme_ids=["video-1", "video-2"],
    )
    service.save_portraits(
        "partial-response",
        [{"aweme_id": "video-2", "age_normalization": {}}],
        requested_aweme_ids=["video-1", "video-2"],
    )

    with factory() as session:
        missing = session.get(FindAgentV2Candidate, 1)
        successful = session.get(FindAgentV2Candidate, 2)
        evidence = list(session.scalars(select(FindAgentV2Evidence).order_by(
            FindAgentV2Evidence.id,
        )))
    assert missing is not None and successful is not None
    assert (missing.detail_status, missing.portrait_status) == ("failed", "failed")
    assert (successful.detail_status, successful.portrait_status) == ("success", "success")
    assert [item.status for item in evidence].count("failed") == 2
    assert [item.status for item in evidence].count("success") == 2
    assert all(
        "上游" in str(item.error_message)
        for item in evidence
        if item.status == "failed"
    )


def test_admin_list_marks_runs_older_than_60_minutes_failed(monkeypatch) -> None:
    engine = create_engine("sqlite+pysqlite:///:memory:")
    FindAgentV2Run.__table__.create(engine)
    factory = sessionmaker(bind=engine, autoflush=False, autocommit=False)

    @contextmanager
    def session_scope() -> Generator[Session, None, None]:
        session = factory()
        try:
            yield session
            session.commit()
        finally:
            session.close()

    now = datetime(2026, 8, 14, 12, 0)
    with factory.begin() as session:
        session.add_all([
            FindAgentV2Run(
                id=1, run_id="overdue", demand_word="超时任务", input_json="{}",
                rule_config_json="{}", status="running",
                create_time=now - timedelta(minutes=61), update_time=now,
            ),
            FindAgentV2Run(
                id=2, run_id="active", demand_word="运行任务", input_json="{}",
                rule_config_json="{}", status="running",
                create_time=now - timedelta(minutes=59), update_time=now,
            ),
        ])

    expire = api_service._expire_overdue_runs
    monkeypatch.setattr(api_service, "get_session", session_scope)
    monkeypatch.setattr(
        api_service, "_expire_overdue_runs", lambda session: expire(session, now=now),
    )

    listed = api_service.list_runs()
    by_id = {item["run_id"]: item for item in listed["items"]}
    assert by_id["overdue"]["status"] == "failed"
    assert by_id["overdue"]["outcome_status"] == "failed"
    assert by_id["overdue"]["stop_reason"] == RUN_TIMEOUT_REASON
    assert by_id["active"]["status"] == "running"

    second_page = api_service.list_runs(page=2, page_size=1)
    assert second_page["page"] == 2
    assert second_page["page_size"] == 1
    assert second_page["total"] == 2
    assert second_page["total_pages"] == 2
    assert second_page["has_previous"] is True
    assert second_page["has_next"] is False
    assert len(second_page["items"]) == 1

    with factory() as session:
        stored = session.scalar(select(FindAgentV2Run).where(
            FindAgentV2Run.run_id == "overdue",
        ))
        assert stored is not None and stored.status == "failed"


@pytest.mark.asyncio
async def test_agent_hard_timeout_marks_run_failed(monkeypatch) -> None:
    class Service:
        failed_reason = ""

        def require_run(self, _run_id):
            return {
                "status": "failed" if self.failed_reason else "running",
                "outcome_status": "failed" if self.failed_reason else None,
                "demand_word": "测试", "current_round": 0, "valid_primary_count": 0,
            }

        def set_obagent_run_uid(self, *_args):
            return None

        def fail_run(self, _run_id, reason):
            self.failed_reason = reason

        def add_usage(self, *_args):
            return None

    class Runner:
        usage = {}

        def reset_usage(self):
            return None

    service = Service()
    agent = FindAgentV2(
        service=service, node_runner=Runner(), observer=NullObserver(),
        max_runtime_seconds=0.01,
    )

    async def slow_run(**_kwargs):
        await asyncio.sleep(0.1)

    monkeypatch.setattr(agent, "_arun_inner", slow_run)
    result = await agent.arun(run_id="timeout", user_input="task")

    assert result.status == "failed"
    assert result.stop_reason == RUN_TIMEOUT_REASON
    assert service.failed_reason == RUN_TIMEOUT_REASON