"""Token 预算守卫：累计消耗超阈值时优雅收尾。"""

from app.application.agent.extensions import (
    AgentRunContext,
    ExtensionRegistry,
    TurnSnapshot,
)
from app.application.agent.python_skills import PythonSkillDefinition

# 累计 total_tokens 软上限。达到后 TurnPolicy 要求停止，让 Agent 用当前内容自然
# 收尾——区别于 max_iterations 的硬截断（truncated=True，回复可能只说了一半）。
_BUDGET_TOTAL_TOKENS = 20000


async def _stop_when_budget_exhausted(
    turn: TurnSnapshot, _run: AgentRunContext
) -> bool:
    return turn.usage.total_tokens >= _BUDGET_TOTAL_TOKENS


def _setup(registry: ExtensionRegistry) -> None:
    registry.add_turn_policy(_stop_when_budget_exhausted)


SKILL = PythonSkillDefinition(
    name="budget_guard",
    description="累计 token 消耗超预算时优雅停止 Agent 循环，避免成本失控。",
    triggers=("预算", "budget", "长任务"),
    instructions=(
        "本轮受 token 预算约束。若已接近上限，请用当前已有信息给出最终回复，"
        "不要发起新的工具调用或长篇推理。"
    ),
    setup=_setup,
)