ai-video-fullstack/backend/services/test_runs/orchestrator.py

"""In-process MVP scheduler for persisted batch test runs."""

from __future__ import annotations

import asyncio
from dataclasses import dataclass, field
from datetime import UTC, datetime

from loguru import logger
from sqlalchemy import select

from db.models import Assistant, ModelResource, TestCase, TestRun, TestRunCase
from db.session import SessionLocal
from models import RuntimeModelResource
from services.config_resolver import resolve_runtime_config
from services.runtime_variables import prepare_dynamic_config
from services.test_case_assets import load_test_asset_bytes
from services.test_runs.errors import TestExecutionError
from services.test_runs.evaluator import EvaluationEngine
from services.test_runs.prompt_improvement import (
    PROMPT_IMPROVEMENT_SNAPSHOT_KEY,
    build_prompt_improvement_snapshot,
)
from services.test_runs.text_runner import (
    RawTurnResult,
    TextPipelineRunner,
    TextStartupResult,
)
from test_schemas import TestCaseDefinition


@dataclass
class _RunControl:
    stop_event: asyncio.Event = field(default_factory=asyncio.Event)
    reason: str | None = None
    case_tasks: set[asyncio.Task[None]] = field(default_factory=set)


class TestRunOrchestrator:
    """Run cases concurrently while keeping PostgreSQL as the source of truth."""

    def __init__(self) -> None:
        self._run_tasks: dict[str, asyncio.Task[None]] = {}
        self._controls: dict[str, _RunControl] = {}

    def start(self, run_id: str) -> None:
        current = self._run_tasks.get(run_id)
        if current is not None and not current.done():
            return
        control = _RunControl()
        self._controls[run_id] = control
        task = asyncio.create_task(self._execute_run(run_id, control), name=f"test-run-{run_id}")
        self._run_tasks[run_id] = task
        task.add_done_callback(lambda _task: self._forget(run_id))

    async def cancel(self, run_id: str) -> None:
        control = self._controls.get(run_id)
        if control is None:
            return
        control.reason = "manual"
        control.stop_event.set()
        for task in tuple(control.case_tasks):
            task.cancel()

    async def shutdown(self) -> None:
        tasks = [task for task in self._run_tasks.values() if not task.done()]
        for task in tasks:
            task.cancel()
        if tasks:
            await asyncio.gather(*tasks, return_exceptions=True)
        self._run_tasks.clear()
        self._controls.clear()

    def _forget(self, run_id: str) -> None:
        self._run_tasks.pop(run_id, None)
        self._controls.pop(run_id, None)

    async def _execute_run(self, run_id: str, control: _RunControl) -> None:
        try:
            async with SessionLocal() as session:
                run = await session.get(TestRun, run_id)
                if run is None:
                    return
                run.status = "running"
                run.started_at = run.started_at or datetime.now(UTC)
                await session.commit()
                config = dict(run.config or {})
                try:
                    cfg = await resolve_runtime_config(session, str(run.assistant_id or ""))
                    # Validate the run-level values once. Each case/attempt gets
                    # a fresh prepared copy below so conversation-scoped system
                    # variables never leak between concurrent cases or retries.
                    prepare_dynamic_config(
                        cfg,
                        config.get("dynamicVariables") or {},
                        assistant_id=str(run.assistant_id or ""),
                    )
                    assistant = await session.get(
                        Assistant,
                        str(run.assistant_id or ""),
                    )
                    config[PROMPT_IMPROVEMENT_SNAPSHOT_KEY] = (
                        build_prompt_improvement_snapshot(
                            cfg,
                            source_prompt=assistant.prompt if assistant else None,
                        )
                    )
                    run.config = config
                    await session.commit()
                except Exception as exc:
                    await self._fail_run_setup(run_id, str(exc))
                    return
                evaluator_resource_id = str(
                    config.get("evaluatorModelResourceId") or ""
                )
                evaluator_row = await session.get(
                    ModelResource,
                    evaluator_resource_id,
                )
                if (
                    evaluator_row is None
                    or evaluator_row.capability != "LLM"
                    or not evaluator_row.enabled
                ):
                    await self._fail_run_setup(
                        run_id,
                        "评估模型不存在、已停用或不是 LLM 资源",
                    )
                    return
                evaluator_resource = RuntimeModelResource(
                    id=evaluator_row.id,
                    name=evaluator_row.name,
                    capability=evaluator_row.capability,
                    interface_type=evaluator_row.interface_type,
                    values=evaluator_row.values or {},
                    secrets=evaluator_row.secrets or {},
                    support_image_input=bool(evaluator_row.support_image_input),
                )
                rows = (
                    await session.execute(
                        select(TestRunCase)
                        .where(TestRunCase.run_id == run_id)
                        .order_by(TestRunCase.position)
                    )
                ).scalars().all()

            semaphore = asyncio.Semaphore(int(config.get("concurrency") or 3))
            tasks = {
                asyncio.create_task(
                    self._execute_case(
                        run_id,
                        row.id,
                        cfg,
                        str(run.assistant_id or ""),
                        evaluator_resource,
                        config,
                        semaphore,
                        control,
                    ),
                    name=f"test-run-case-{row.id}",
                )
                for row in rows
            }
            control.case_tasks = tasks
            if tasks:
                await asyncio.gather(*tasks, return_exceptions=True)
            await self._finish_run(run_id, control)
        except asyncio.CancelledError:
            await self._interrupt_run(run_id)
            raise
        except Exception as exc:  # noqa: BLE001 - background task must persist failure
            logger.exception(f"批量测试运行失败: run_id={run_id}: {exc}")
            await self._fail_run_setup(run_id, str(exc))

    async def _execute_case(
        self,
        run_id: str,
        run_case_id: str,
        cfg,
        assistant_id: str,
        evaluator_resource: RuntimeModelResource,
        config: dict,
        semaphore: asyncio.Semaphore,
        control: _RunControl,
    ) -> None:
        async with semaphore:
            if control.stop_event.is_set():
                await self._mark_skipped(run_case_id)
                return

            async with SessionLocal() as session:
                row = await session.get(TestRunCase, run_case_id)
                if row is None:
                    return
                snapshot = dict(row.case_snapshot or {})
                definition = TestCaseDefinition.model_validate(snapshot)
                max_attempts = row.max_attempts
                try:
                    image_assets = await load_test_asset_bytes(session, definition)
                except Exception as exc:
                    await self._mark_error(
                        run_case_id,
                        TestExecutionError(
                            code="TEST_IMAGE_READ_FAILED",
                            message=str(exc) or "测试图片无法读取",
                            stage="pipeline",
                            retryable=True,
                        ),
                    )
                    return

            for attempt in range(1, max_attempts + 1):
                if control.stop_event.is_set():
                    await self._mark_skipped(run_case_id)
                    return
                await self._mark_running(run_case_id, attempt)
                try:
                    timeout_seconds = int(
                        config.get("timeoutSecs")
                        or config.get("timeout_secs")
                        or 60
                    )
                    async with asyncio.timeout(timeout_seconds):
                        attempt_cfg = prepare_dynamic_config(
                            cfg.model_copy(deep=True),
                            config.get("dynamicVariables") or {},
                            assistant_id=assistant_id,
                        )

                        async def save_turn_progress(turn: RawTurnResult) -> None:
                            await self._mark_turn_progress(run_case_id, turn)

                        async def save_startup_progress(
                            startup: TextStartupResult,
                        ) -> None:
                            await self._mark_startup_progress(
                                run_case_id,
                                startup,
                            )

                        execution = await TextPipelineRunner().run(
                            attempt_cfg,
                            definition,
                            timeout_seconds=timeout_seconds,
                            image_assets=image_assets,
                            on_turn_complete=save_turn_progress,
                            on_startup_complete=save_startup_progress,
                        )
                        evaluated = await EvaluationEngine(evaluator_resource).evaluate(
                            definition,
                            execution,
                        )
                except asyncio.CancelledError:
                    await self._mark_skipped(run_case_id)
                    raise
                except TimeoutError:
                    error = TestExecutionError(
                        code="CASE_TIMEOUT",
                        message=f"测试用例执行与评估超过 {timeout_seconds} 秒",
                        stage="pipeline",
                        retryable=True,
                    )
                    if attempt < max_attempts and not control.stop_event.is_set():
                        continue
                    await self._mark_error(run_case_id, error)
                    if str(
                        config.get("errorStrategy")
                        or config.get("error_strategy")
                    ) == "stop_on_error":
                        await self._request_stop(run_id, control, "execution_error")
                    return
                except TestExecutionError as exc:
                    if exc.retryable and attempt < max_attempts and not control.stop_event.is_set():
                        continue
                    await self._mark_error(run_case_id, exc)
                    if str(config.get("errorStrategy") or config.get("error_strategy")) == "stop_on_error":
                        await self._request_stop(run_id, control, "execution_error")
                    return
                except Exception as exc:  # noqa: BLE001 - normalize unknown runtime errors
                    error = TestExecutionError(
                        code="UNEXPECTED_EXECUTION_ERROR",
                        message=str(exc) or type(exc).__name__,
                        stage="pipeline",
                        retryable=False,
                    )
                    await self._mark_error(run_case_id, error)
                    if str(config.get("errorStrategy") or config.get("error_strategy")) == "stop_on_error":
                        await self._request_stop(run_id, control, "execution_error")
                    return

                status = "pass" if evaluated.passed else "fail"
                await self._mark_evaluated(
                    run_case_id,
                    status,
                    evaluated,
                    execution.startup,
                )
                if status == "fail" and str(
                    config.get("failureStrategy") or config.get("failure_strategy")
                ) == "stop_on_fail":
                    await self._request_stop(run_id, control, "assertion_failure")
                return

    async def _request_stop(
        self,
        run_id: str,
        control: _RunControl,
        reason: str,
    ) -> None:
        if control.reason is None:
            control.reason = reason
        control.stop_event.set()
        async with SessionLocal() as session:
            run = await session.get(TestRun, run_id)
            if run is not None and run.stop_reason is None:
                run.stop_reason = reason
                await session.commit()

    async def _mark_running(self, run_case_id: str, attempt: int) -> None:
        async with SessionLocal() as session:
            row = await session.get(TestRunCase, run_case_id)
            if row is None:
                return
            row.status = "running"
            row.attempt_count = attempt
            row.started_at = row.started_at or datetime.now(UTC)
            row.finished_at = None
            row.execution_error = None
            row.result = {
                "turns": [],
                "startup": {
                    "completed": False,
                    "greeting": "",
                    "assistantOutputs": [],
                    "autoConfirmedMessages": [],
                    "runtimeRecords": [],
                },
                "overallCriteria": [],
            }
            await session.commit()

    async def _mark_startup_progress(
        self,
        run_case_id: str,
        startup: TextStartupResult,
    ) -> None:
        async with SessionLocal() as session:
            row = await session.get(TestRunCase, run_case_id)
            if row is None or row.status != "running":
                return
            result = dict(row.result or {})
            result["startup"] = {
                "completed": True,
                "greeting": startup.greeting,
                "assistantOutputs": list(startup.assistant_outputs),
                "autoConfirmedMessages": list(
                    startup.auto_confirmed_messages
                ),
                "runtimeRecords": list(startup.runtime_records),
            }
            row.result = result
            await session.commit()

    async def _mark_turn_progress(
        self,
        run_case_id: str,
        turn: RawTurnResult,
    ) -> None:
        """Persist completed turns so polling clients can show live progress."""
        async with SessionLocal() as session:
            row = await session.get(TestRunCase, run_case_id)
            if row is None or row.status != "running":
                return
            result = dict(row.result or {})
            turns = list(result.get("turns") or [])
            payload = {
                "id": turn.id,
                "index": turn.index,
                "userInput": turn.user_input,
                "image": (
                    turn.image.model_dump(mode="json", by_alias=True)
                    if turn.image is not None
                    else None
                ),
                "assistantReply": turn.assistant_reply,
                "toolCalls": [item.as_result() for item in turn.tool_calls],
                "runtimeRecords": list(turn.runtime_records),
                "evaluations": [],
            }
            turns = [item for item in turns if item.get("id") != turn.id]
            turns.append(payload)
            turns.sort(key=lambda item: int(item.get("index") or 0))
            row.result = {
                "turns": turns,
                "startup": dict(result.get("startup") or {}),
                "overallCriteria": list(result.get("overallCriteria") or []),
            }
            await session.commit()

    async def _mark_evaluated(
        self,
        run_case_id: str,
        status: str,
        evaluated,
        startup: TextStartupResult,
    ) -> None:
        async with SessionLocal() as session:
            row = await session.get(TestRunCase, run_case_id)
            if row is None:
                return
            row.status = status
            row.result = {
                "turns": [
                    item.model_dump(mode="json", by_alias=True)
                    for item in evaluated.turns
                ],
                "startup": {
                    "completed": True,
                    "greeting": startup.greeting,
                    "assistantOutputs": list(startup.assistant_outputs),
                    "autoConfirmedMessages": list(
                        startup.auto_confirmed_messages
                    ),
                    "runtimeRecords": list(startup.runtime_records),
                },
                "overallCriteria": [
                    item.model_dump(mode="json", by_alias=True)
                    for item in evaluated.overall_criteria
                ],
            }
            row.execution_error = None
            row.finished_at = datetime.now(UTC)
            if row.test_case_id:
                test_case = await session.get(TestCase, row.test_case_id)
                if test_case is not None:
                    test_case.last_result = status
            await session.commit()

    async def _mark_error(
        self,
        run_case_id: str,
        error: TestExecutionError,
    ) -> None:
        async with SessionLocal() as session:
            row = await session.get(TestRunCase, run_case_id)
            if row is None:
                return
            row.status = "error"
            row.execution_error = error.as_dict()
            row.finished_at = datetime.now(UTC)
            if row.test_case_id:
                test_case = await session.get(TestCase, row.test_case_id)
                if test_case is not None:
                    test_case.last_result = "fail"
            await session.commit()

    async def _mark_skipped(self, run_case_id: str) -> None:
        async with SessionLocal() as session:
            row = await session.get(TestRunCase, run_case_id)
            if row is None or row.status in {"pass", "fail", "error", "skipped"}:
                return
            row.status = "skipped"
            row.finished_at = datetime.now(UTC)
            await session.commit()

    async def _finish_run(self, run_id: str, control: _RunControl) -> None:
        async with SessionLocal() as session:
            run = await session.get(TestRun, run_id)
            if run is None:
                return
            waiting = (
                await session.execute(
                    select(TestRunCase).where(
                        TestRunCase.run_id == run_id,
                        TestRunCase.status.in_(["waiting", "running"]),
                    )
                )
            ).scalars().all()
            for row in waiting:
                row.status = "skipped"
                row.finished_at = datetime.now(UTC)
            manual = bool(run.cancel_requested or control.reason == "manual")
            run.status = "cancelled" if manual else "completed"
            run.stop_reason = "manual" if manual else control.reason or run.stop_reason
            run.finished_at = datetime.now(UTC)
            await session.commit()

    async def _fail_run_setup(self, run_id: str, message: str) -> None:
        error = TestExecutionError(
            code="RUN_SETUP_FAILED",
            message=message or "批量测试初始化失败",
            stage="pipeline",
            retryable=False,
        )
        async with SessionLocal() as session:
            run = await session.get(TestRun, run_id)
            if run is None:
                return
            rows = (
                await session.execute(
                    select(TestRunCase).where(
                        TestRunCase.run_id == run_id,
                        TestRunCase.status.in_(["waiting", "running"]),
                    )
                )
            ).scalars().all()
            now = datetime.now(UTC)
            for row in rows:
                row.status = "error"
                row.execution_error = error.as_dict()
                row.finished_at = now
            run.status = "completed"
            run.stop_reason = "execution_error"
            run.finished_at = now
            await session.commit()

    async def _interrupt_run(self, run_id: str) -> None:
        await _mark_run_interrupted(run_id, "后端服务停止，运行已中断")


async def _mark_run_interrupted(run_id: str, message: str) -> None:
    async with SessionLocal() as session:
        run = await session.get(TestRun, run_id)
        if run is None or run.status not in {"queued", "running"}:
            return
        rows = (
            await session.execute(
                select(TestRunCase).where(
                    TestRunCase.run_id == run_id,
                    TestRunCase.status.in_(["waiting", "running"]),
                )
            )
        ).scalars().all()
        now = datetime.now(UTC)
        for row in rows:
            if row.status == "running":
                row.status = "error"
                row.execution_error = {
                    "code": "SERVER_RESTARTED",
                    "message": message,
                    "stage": "pipeline",
                    "retryable": True,
                }
            else:
                row.status = "skipped"
            row.finished_at = now
        run.status = "completed"
        run.stop_reason = "execution_error"
        run.finished_at = now
        await session.commit()


async def recover_interrupted_test_runs() -> None:
    async with SessionLocal() as session:
        run_ids = (
            await session.execute(
                select(TestRun.id).where(TestRun.status.in_(["queued", "running"]))
            )
        ).scalars().all()
    for run_id in run_ids:
        await _mark_run_interrupted(str(run_id), "后端服务重启，运行已中断")


test_run_orchestrator = TestRunOrchestrator()