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

"""Strict, deterministic tool simulations for text test runs."""

from __future__ import annotations

import asyncio
import json
from typing import Any

from models import RuntimeTool, RuntimeToolBinding
from services.runtime_variables import DynamicVariableStore
from services.tool_executor import ToolExecutionError, ToolExecutor
from test_schemas import ToolSimulation, ToolSimulationResponse, ToolSimulationStep


TOOL_SIMULATION_ERROR_CODES = {
    "UNMOCKED_TOOL_CALL",
    "TOOL_SIMULATION_EXHAUSTED",
    "TOOL_SIMULATION_SCOPE_MISMATCH",
}


class ToolSimulationError(ToolExecutionError):
    """A non-retryable test configuration error raised at tool-call time."""

    def __init__(self, code: str, message: str):
        self.code = code
        super().__init__(f"{code}: {message}")


class ToolSimulationResolver:
    """Resolve one case's fixed or ordered simulated tool responses."""

    def __init__(self, simulations: list[ToolSimulation]):
        self._by_tool_id = {item.tool_id: item for item in simulations}
        self._by_function = {item.function_name: item for item in simulations}
        self._positions = {item.tool_id: 0 for item in simulations}
        self._phase = "startup"
        self._turn_id: str | None = None
        self._failure: ToolSimulationError | None = None

    def enter_startup(self) -> None:
        self._phase = "startup"
        self._turn_id = None

    def enter_turn(self, turn_id: str) -> None:
        self._phase = "turn"
        self._turn_id = turn_id

    async def resolve(
        self,
        *,
        tool_id: str | None,
        function_name: str,
    ) -> dict[str, Any]:
        simulation = (
            self._by_tool_id.get(tool_id or "")
            or self._by_function.get(function_name)
        )
        if simulation is None:
            raise self._record_failure(
                "UNMOCKED_TOOL_CALL",
                f"工具 {function_name} 未配置工具模拟",
            )

        response: ToolSimulationResponse | ToolSimulationStep
        if simulation.mode == "fixed":
            response = simulation.response
        else:
            position = self._positions[simulation.tool_id]
            if position >= len(simulation.steps):
                raise self._record_failure(
                    "TOOL_SIMULATION_EXHAUSTED",
                    f"工具 {function_name} 的模拟响应已耗尽（第 {position + 1} 次调用）",
                )
            step = simulation.steps[position]
            if not self._scope_matches(step):
                raise self._record_failure(
                    "TOOL_SIMULATION_SCOPE_MISMATCH",
                    (
                        f"工具 {function_name} 的第 {position + 1} 条模拟步骤"
                        f"期望发生在{self._scope_label(step)}，"
                        f"实际发生在{self._current_scope_label()}"
                    ),
                )
            response = step
            self._positions[simulation.tool_id] = position + 1

        if response.delay_ms:
            await asyncio.sleep(response.delay_ms / 1000)
        return self._normalized_result(response)

    @property
    def failure(self) -> ToolSimulationError | None:
        return self._failure

    def raise_if_failed(self) -> None:
        """Re-raise errors swallowed by a higher-level client message stage."""

        if self._failure is not None:
            raise self._failure

    def _record_failure(self, code: str, message: str) -> ToolSimulationError:
        error = ToolSimulationError(code, message)
        if self._failure is None:
            self._failure = error
        return error

    def _scope_matches(self, step: ToolSimulationStep) -> bool:
        if step.scope.phase == "any":
            return True
        if step.scope.phase == "startup":
            return self._phase == "startup"
        return self._phase == "turn" and step.scope.turn_id == self._turn_id

    @staticmethod
    def _scope_label(step: ToolSimulationStep) -> str:
        if step.scope.phase == "any":
            return "任意阶段"
        if step.scope.phase == "startup":
            return "启动阶段"
        return f"轮次 {step.scope.turn_id}"

    def _current_scope_label(self) -> str:
        if self._phase == "startup":
            return "启动阶段"
        return f"轮次 {self._turn_id or 'unknown'}"

    @staticmethod
    def _normalized_result(response: ToolSimulationResponse) -> dict[str, Any]:
        payload = json.loads(response.body)
        status = "ok" if response.outcome == "success" else "error"
        if isinstance(payload, dict):
            return {**payload, "status": status}
        return {"status": status, "data": payload}


class MockToolExecutor(ToolExecutor):
    """Execute external runtime tools exclusively through one test resolver."""

    def __init__(
        self,
        store: DynamicVariableStore,
        resolver: ToolSimulationResolver,
    ):
        super().__init__(store)
        self._resolver = resolver

    async def execute(
        self,
        tool: RuntimeTool,
        arguments: dict[str, Any] | None = None,
        *,
        binding: RuntimeToolBinding | None = None,
        result_assignments: dict[str, str] | None = None,
    ) -> dict[str, Any]:
        result = await self._resolver.resolve(
            tool_id=tool.id,
            function_name=tool.function_name,
        )
        if result.get("status") != "ok":
            return {**result, "updated_variables": []}
        return self._apply_result_assignments(
            tool,
            result,
            result_assignments=(
                result_assignments
                if result_assignments is not None
                else (binding.result_assignments if binding is not None else {})
            ),
        )


class TextClientToolPort:
    """Return configured client-tool results without a browser transport."""

    def __init__(self, resolver: ToolSimulationResolver) -> None:
        self._resolver = resolver
        self.auto_confirmed_messages: list[dict[str, Any]] = []

    def supports(self, _function_name: str) -> bool:
        # Test runs advertise the virtual client, then fail strictly in call().
        # Returning False would silently skip an unconfigured client tool.
        return True

    async def call(
        self,
        function_name: str,
        arguments: dict[str, Any],
        **_kwargs: Any,
    ) -> dict[str, Any]:
        resolved = await self._resolver.resolve(
            tool_id=None,
            function_name=function_name,
        )
        status = str(resolved.get("status") or "error")
        payload = (
            resolved.get("data")
            if set(resolved).issubset({"status", "data"}) and "data" in resolved
            else {key: value for key, value in resolved.items() if key != "status"}
        )
        # ClientToolPort mirrors the browser transport envelope. The JSON
        # configured in the editor remains the client handler's root object.
        result: dict[str, Any] = {"status": status, "data": payload}
        if (
            status != "ok"
            and isinstance(payload, dict)
            and payload.get("message") is not None
        ):
            result["message"] = str(payload["message"])
        data = result.get("data")
        action = data.get("action") if isinstance(data, dict) else None
        if (
            function_name == "show_message"
            and result.get("status") == "ok"
            and action == "confirmed"
        ):
            self.auto_confirmed_messages.append(dict(arguments))
        return result