SupplyAgent/agents/find_agent/completion_guard.py

"""find_agent 基于持久化运行状态的结束守卫。"""
from __future__ import annotations

from collections.abc import Sequence
from typing import TYPE_CHECKING

from supply_agent.types import CompletionGuard, Message
from supply_infra.services.video_discovery_service import (
    get_video_discovery_service,
)

if TYPE_CHECKING:
    from supply_agent.agent.core import Agent

_TERMINAL_STATUSES = {"finished", "failed"}


def create_find_completion_guard(run_id: str) -> CompletionGuard:
    """创建只允许已进入 finished / failed 终态的 find_agent 结束守卫。"""
    normalized_run_id = str(run_id or "").strip()[:64]

    def guard(_response: Message, _messages: Sequence[Message]) -> str | None:
        if not normalized_run_id:
            return (
                "当前用户消息中缺少 run_id，无法确认运行状态。"
                "请不要直接输出最终结果。"
            )

        run = get_video_discovery_service().lookup_run(normalized_run_id)
        if run is None:
            return (
                f"run_id={normalized_run_id} 不存在，无法确认运行状态。"
                "请不要直接输出最终结果。"
            )

        status = str(run.get("status") or "").strip()
        if status in _TERMINAL_STATUSES:
            return None

        return (
            f"当前 run.status 仍为 {status or 'unknown'}。"
            "请继续处理；正常完成后先调用 "
            "update_video_discovery_run_status 将状态更新为 finished，"
            "无法完成时将状态更新为 failed，再输出最终结果。"
        )

    return guard


def configure_find_agent_completion_guard(agent: Agent, run_id: str) -> None:
    """在 find_agent 循环创建前按需注入状态结束守卫。"""
    if getattr(agent, "completion_guard", None) is not None:
        return
    agent.completion_guard = create_find_completion_guard(run_id)


__all__ = [
    "configure_find_agent_completion_guard",
    "create_find_completion_guard",
]