muse-agent-example/src/muse/任务运行/预算管理.py

"""版本化额度与逐调用预算账本；预留、取消和结算都在短事务中完成。"""

from __future__ import annotations

import math
from collections.abc import Callable, Iterator
from contextlib import contextmanager
from dataclasses import asdict, dataclass, replace
from datetime import datetime, timedelta
from decimal import ROUND_HALF_UP, Decimal, InvalidOperation
from typing import Any, LiteralString
from zoneinfo import ZoneInfo

import psycopg
from psycopg import sql
from psycopg.rows import dict_row
from psycopg.types.json import Jsonb

from muse.任务运行.任务领取 import 校验领取
from muse.任务运行.存储 import 任务存储
from muse.任务运行.模型 import 内容哈希, 领取凭证
from muse.任务运行.角色策略 import 角色身份
from muse.共享.调用身份 import 用途
from muse.共享.错误 import Muse错误
from muse.基础设施.数据库.连接 import 数据库工厂

_金额精度 = Decimal("0.000001")


class 预算错误(Muse错误):
    错误码 = "MUSE_BUDGET"


class 预算不足(预算错误):
    错误码 = "MUSE_BUDGET_EXHAUSTED"


class 预算状态冲突(预算错误):
    错误码 = "MUSE_BUDGET_STATE"


def 金额(值: Decimal | str | int) -> Decimal:
    try:
        if isinstance(值, bool):
            raise ValueError
        结果 = Decimal(str(值))
        if not 结果.is_finite() or 结果 < 0:
            raise ValueError
        return 结果.quantize(_金额精度, rounding=ROUND_HALF_UP)
    except (InvalidOperation, ValueError):
        raise 预算错误("金额必须是明确的非负有限数") from None


def _时区时间(值: datetime) -> datetime:
    if 值.tzinfo is None or 值.utcoffset() is None:
        raise 预算错误("预算时间必须包含时区")
    return 值


@dataclass(frozen=True, slots=True)
class 额度窗口:
    起点: datetime
    终点: datetime


def 归属窗口(当前: datetime, 时区: str = "Asia/Shanghai") -> 额度窗口:
    """沿用每日 00/05/10/15/20 边界；末窗到午夜，避免默改既有额度合同。"""
    本地 = _时区时间(当前).astimezone(ZoneInfo(时区))
    起点 = 本地.replace(hour=(本地.hour // 5) * 5, minute=0, second=0, microsecond=0)
    终点 = (
        起点 + timedelta(hours=5) if 起点.hour < 20 else (起点 + timedelta(days=1)).replace(hour=0)
    )
    return 额度窗口(起点, 终点)


@dataclass(frozen=True, slots=True)
class 额度策略:
    账户ID: str
    版本: str
    窗口金额上限: Decimal
    窗口调用上限: int
    时区: str = "Asia/Shanghai"

    def __post_init__(self) -> None:
        object.__setattr__(self, "窗口金额上限", 金额(self.窗口金额上限))
        if not self.账户ID or not self.版本 or self.窗口金额上限 <= 0:
            raise 预算错误("额度策略需要身份、版本和正数金额上限")
        if type(self.窗口调用上限) is not int or self.窗口调用上限 < 1:
            raise 预算错误("窗口调用上限必须是正整数")
        ZoneInfo(self.时区)

    def 冻结(self) -> dict:
        return {**asdict(self), "窗口金额上限": str(self.窗口金额上限)}


@dataclass(frozen=True, slots=True)
class 角色预算:
    角色: str
    计划次数: int
    最多次数: int
    单次上限: Decimal

    def __post_init__(self) -> None:
        object.__setattr__(self, "单次上限", 金额(self.单次上限))
        if self.角色 not in 角色身份 or self.单次上限 <= 0:
            raise 预算错误("角色预算需要已登记角色和正数单次上限")
        if any(type(n) is not int or n < 1 for n in (self.计划次数, self.最多次数)):
            raise 预算错误("调用计划与安全容量必须为正整数")
        if self.计划次数 > self.最多次数:
            raise 预算错误("计划调用次数不能超过安全容量")


@dataclass(frozen=True, slots=True)
class 任务预算计划:
    总金额: Decimal
    角色: tuple[角色预算, ...]
    批准引用: str
    截止时间: datetime

    def __post_init__(self) -> None:
        object.__setattr__(self, "总金额", 金额(self.总金额))
        _时区时间(self.截止时间)
        if not self.批准引用 or not self.角色 or len({r.角色 for r in self.角色}) != len(self.角色):
            raise 预算错误("任务预算需要批准引用和不重复的完整角色计划")
        if self.总金额 < self.最坏预留:
            raise 预算不足("任务总预算不能覆盖计划调用的最坏预留")

    @property
    def 最坏预留(self) -> Decimal:
        return sum((r.单次上限 * r.计划次数 for r in self.角色), Decimal(0))

    def 冻结(self) -> dict:
        return {
            "总金额": str(self.总金额),
            "批准引用": self.批准引用,
            "截止时间": self.截止时间.isoformat(),
            "角色": [{**asdict(r), "单次上限": str(r.单次上限)} for r in self.角色],
        }


@dataclass(frozen=True, slots=True)
class 预算预留:
    调用ID: str
    任务ID: str
    尝试ID: str
    角色: str
    状态: str
    预留金额: Decimal
    实际金额: Decimal | None
    到期时间: datetime
    结算回执: str | None
    超预算: bool
    允许外发: bool = False


class 预算管理:
    def __init__(self, 数据库: 数据库工厂, 账户ID: str) -> None:
        self.数据库 = 数据库
        self.账户ID = 账户ID
        self._schema = "evaluation" if 数据库.用途 is 用途.评测 else "public"

    def _查询(self, 连: psycopg.Connection, 语句: LiteralString, 参数: tuple = ()) -> Any:
        return 连.cursor(row_factory=dict_row).execute(
            sql.SQL(语句).format(s=sql.Identifier(self._schema)), 参数
        )

    @contextmanager
    def _事务(self) -> Iterator[psycopg.Connection]:
        with self.数据库.连接() as 连, 连.transaction():
            yield 连

    def _锁账户(self, 连: psycopg.Connection) -> dict:
        行 = self._查询(
            连,
            "SELECT * FROM {s}.muse_budget_account "
            "WHERE run_purpose=%s AND account_id=%s FOR UPDATE",
            (self.数据库.用途.value, self.账户ID),
        ).fetchone()
        if 行 is None:
            raise 预算错误("额度账户尚未登记")
        return 行["policy"]

    def 登记策略(self, 策略: 额度策略) -> None:
        if 策略.账户ID != self.账户ID:
            raise 预算错误("额度策略与账户不匹配")
        内容 = 策略.冻结()
        with self._事务() as 连:
            self._查询(
                连,
                "INSERT INTO {s}.muse_budget_account "
                "(run_purpose,account_id,policy_version,policy,policy_hash) VALUES (%s,%s,%s,"
                "%s,%s) "
                "ON CONFLICT DO NOTHING",
                (self.数据库.用途.value, self.账户ID, 策略.版本, Jsonb(内容), 内容哈希(内容)),
            )
            if self._锁账户(连) != 内容:
                raise 预算状态冲突("账户已有冻结额度策略；不能就地改变在途预算")

    def 登记任务预算(self, 任务ID: str, 计划: 任务预算计划) -> None:
        with self._事务() as 连:
            self.登记任务预算于(连, 任务ID, 计划)

    def 登记任务预算于(self, 连, 任务ID: str, 计划: 任务预算计划) -> None:
        """沿用调用方事务，使任务、配置和预算同时可见或一起回滚。"""
        内容 = 计划.冻结()
        任务存储(连, self.数据库.用途).任务行(任务ID, 锁定=True)
        self._锁账户(连)
        self._查询(
            连,
            "INSERT INTO {s}.muse_task_budget (task_id,account_id,plan,plan_hash) "
            "VALUES (%s,%s,%s,%s) ON CONFLICT DO NOTHING",
            (任务ID, self.账户ID, Jsonb(内容), 内容哈希(内容)),
        )
        既有 = self._任务预算(连, 任务ID)
        if 既有["plan_hash"] != 内容哈希(内容):
            raise 预算状态冲突("任务预算已冻结，重复登记不能改变批准金额与计划")

    def _任务预算(self, 连: psycopg.Connection, 任务ID: str) -> dict:
        行 = self._查询(
            连,
            "SELECT * FROM {s}.muse_task_budget WHERE task_id=%s AND account_id=%s",
            (任务ID, self.账户ID),
        ).fetchone()
        if 行 is None:
            raise 预算错误("任务尚未登记本账户的预算计划")
        return 行

    def 核对任务预算(self, 任务ID: str, 角色: str) -> dict:
        """领取模型步骤前核对已批准的任务包络。

        这是只读预检，不预留金额、不延长截止时间，也不代替发送前的
        账户、窗口与逐调用原子复检。
        """

        if 角色 not in 角色身份:
            raise 预算错误("任务预算角色未登记")
        with self.数据库.连接(只读=True) as 连:
            登记 = self._任务预算(连, 任务ID)
            计划 = 登记["plan"]
            当前 = self._查询(连, "SELECT clock_timestamp() AS current_time").fetchone()[
                "current_time"
            ]
        if 登记["stopped"] or datetime.fromisoformat(计划["截止时间"]) <= 当前:
            raise 预算不足("任务预算已停止或超过总期限")
        角色计划 = next((项 for 项 in 计划["角色"] if 项["角色"] == 角色), None)
        if 角色计划 is None:
            raise 预算错误("当前角色未纳入已批准的任务调用计划")
        return {
            "账户ID": self.账户ID,
            "批准引用": 计划["批准引用"],
            "截止时间": 计划["截止时间"],
            "角色": 角色,
        }

    def 读取任务预算面(self, 任务ID: str) -> dict:
        with self.数据库.连接(只读=True) as 连:
            account = self._查询(
                连,
                "SELECT policy,policy_hash FROM {s}.muse_budget_account "
                "WHERE run_purpose=%s AND account_id=%s",
                (self.数据库.用途.value, self.账户ID),
            ).fetchone()
            row = self._查询(
                连,
                "SELECT plan,plan_hash,stopped FROM {s}.muse_task_budget "
                "WHERE task_id=%s AND account_id=%s",
                (任务ID, self.账户ID),
            ).fetchone()
        if account and 内容哈希(account["policy"]) != account["policy_hash"]:
            raise 预算错误("账户额度策略哈希不一致")
        if row and 内容哈希(row["plan"]) != row["plan_hash"]:
            raise 预算错误("任务预算哈希不一致")
        return {
            "account_id": self.账户ID,
            "account_policy": account["policy"] if account else None,
            "plan": row["plan"] if row else None,
            "stopped": row["stopped"] if row else False,
        }

    def 读取任务费用于(self, 连, 任务ID: str) -> dict:
        """查询完整尝试账本；失败调用不因没有业务产物而消失。"""
        任务存储(连, self.数据库.用途).任务行(任务ID)
        行集 = self._查询(
            连,
            "SELECT call_id,attempt_id,role_id,state,actual_amount,reserved_amount,over_budget "
            "FROM {s}.muse_budget_reservation WHERE run_purpose=%s AND account_id=%s "
            "AND task_id=%s ORDER BY call_id",
            (self.数据库.用途.value, self.账户ID, 任务ID),
        ).fetchall()
        known = sum(
            (r["actual_amount"] for r in 行集 if r["actual_amount"] is not None), Decimal(0)
        )
        unknown = any(r["state"] in {"in_flight", "unknown"} for r in 行集)
        pending = any(r["state"] == "reserved" for r in 行集)
        return {
            "known_total_usd": str(known),
            "total_usd": None if unknown or pending else str(known),
            "has_unknown": unknown,
            "has_pending": pending,
            "over_budget": any(r["over_budget"] for r in 行集),
            "calls": [
                {
                    **r,
                    "attempt_id": str(r["attempt_id"]),
                    "actual_amount": str(r["actual_amount"])
                    if r["actual_amount"] is not None
                    else None,
                    "reserved_amount": str(r["reserved_amount"]),
                }
                for r in 行集
            ],
        }

    def 回收过期(self) -> dict[str, int]:
        """未发送预留可释放；已发送过期只能转未知，不能自动抹掉成本。"""
        with self._事务() as 连:
            self._锁账户(连)
            条件 = (self.数据库.用途.value, self.账户ID)
            释放 = self._查询(
                连,
                "UPDATE {s}.muse_budget_reservation SET state='released',reason='expired' "
                "WHERE run_purpose=%s AND account_id=%s AND state='reserved' "
                "AND expires_at<=clock_timestamp()",
                条件,
            ).rowcount
            未知 = self._查询(
                连,
                "UPDATE {s}.muse_budget_reservation SET state='unknown',reason='expired' "
                "WHERE run_purpose=%s AND account_id=%s AND state='in_flight' "
                "AND expires_at<=clock_timestamp()",
                条件,
            ).rowcount
            return {"释放": 释放, "未知": 未知}

    def _调用行(self, 连: psycopg.Connection, 调用ID: str) -> dict | None:
        return self._查询(
            连,
            "SELECT * FROM {s}.muse_budget_reservation "
            "WHERE run_purpose=%s AND account_id=%s AND call_id=%s",
            (self.数据库.用途.value, self.账户ID, 调用ID),
        ).fetchone()

    def _所有调用(self, 连: psycopg.Connection) -> list[dict]:
        return self._查询(
            连,
            "SELECT * FROM {s}.muse_budget_reservation "
            "WHERE run_purpose=%s AND account_id=%s AND state<>'released'",
            (self.数据库.用途.value, self.账户ID),
        ).fetchall()

    def 预留(
        self,
        领取: 领取凭证,
        调用ID: str,
        角色: str,
        *,
        有效秒: float = 60,
        允许未发重绑: bool = False,
    ) -> 预算预留:
        if not 调用ID or not math.isfinite(有效秒) or 有效秒 <= 0:
            raise 预算错误("调用身份和正数预留期限必须提供")
        self.回收过期()
        with self._事务() as 连:
            校验领取(任务存储(连, self.数据库.用途), 领取)
            策略 = self._锁账户(连)
            请求哈希 = 内容哈希(
                {
                    "task": 领取.任务ID,
                    "attempt": 领取.尝试ID,
                    "role": 角色,
                    "account": self.账户ID,
                    "ttl": 有效秒,
                }
            )
            既有 = self._调用行(连, 调用ID)
            可重绑 = bool(
                允许未发重绑
                and 既有
                and 既有["state"] == "released"
                and 既有["sent_at"] is None
                and str(既有["task_id"]) == 领取.任务ID
                and 既有["role_id"] == 角色
            )
            if 既有 and not 可重绑:
                if 既有["request_hash"] != 请求哈希:
                    raise 预算状态冲突("同一调用身份不能重复绑定不同预算请求")
                return self._预留(既有)
            当前 = self._查询(连, "SELECT clock_timestamp() AS current_time").fetchone()[
                "current_time"
            ]
            窗口 = 归属窗口(当前, 策略["时区"])
            登记 = self._任务预算(连, 领取.任务ID)
            计划 = 登记["plan"]
            if 登记["stopped"] or datetime.fromisoformat(计划["截止时间"]) <= 当前:
                raise 预算不足("任务预算已停止或超过总期限")
            所有 = self._所有调用(连)
            if any(r["state"] == "unknown" for r in 所有):
                raise 预算不足("账户存在未知成本，必须先对账")
            本窗 = [r for r in 所有 if r["window_start"] == 窗口.起点]
            已占 = sum((self._占用(r) for r in 本窗), Decimal(0))
            角色表 = {r["角色"]: r for r in 计划["角色"]}
            if 角色 not in 角色表:
                raise 预算错误("当前角色未纳入已批准的任务调用计划")
            每次 = Decimal(角色表[角色]["单次上限"])
            if len(本窗) >= 策略["窗口调用上限"] or 已占 + 每次 > Decimal(策略["窗口金额上限"]):
                raise 预算不足(
                    "当前窗口的金额或调用次数不足", 上下文={"下个窗口": 窗口.终点.isoformat()}
                )
            本任务 = [r for r in 所有 if str(r["task_id"]) == 领取.任务ID]
            用量 = {role: sum(r["role_id"] == role for r in 本任务) for role in 角色表}
            目标 = 角色表[角色]
            if 用量[角色] >= min(目标["计划次数"], 目标["最多次数"]):
                raise 预算不足("当前角色的计划调用次数已耗尽")
            剩余 = sum(
                (
                    Decimal(r["单次上限"]) * max(r["计划次数"] - 用量[role], 0)
                    for role, r in 角色表.items()
                ),
                Decimal(0),
            )
            if sum((self._占用(r) for r in 本任务), Decimal(0)) + 剩余 > Decimal(计划["总金额"]):
                raise 预算不足("实际成本与剩余计划预留超过任务总预算")
            到期 = min(
                当前 + timedelta(seconds=有效秒),
                窗口.终点,
                datetime.fromisoformat(计划["截止时间"]),
            )
            if 可重绑:
                行 = self._查询(
                    连,
                    "UPDATE {s}.muse_budget_reservation SET attempt_id=%s,request_hash=%s,"
                    "state='reserved',reserved_amount=%s,window_start=%s,window_end=%s,"
                    "expires_at=%s,reason=NULL WHERE run_purpose=%s AND account_id=%s "
                    "AND call_id=%s AND state='released' AND sent_at IS NULL RETURNING *",
                    (
                        领取.尝试ID,
                        请求哈希,
                        每次,
                        窗口.起点,
                        窗口.终点,
                        到期,
                        self.数据库.用途.value,
                        self.账户ID,
                        调用ID,
                    ),
                ).fetchone()
                return self._预留(行)
            行 = self._查询(
                连,
                "INSERT INTO {s}.muse_budget_reservation "
                "(run_purpose,call_id,account_id,task_id,attempt_id,role_id,request_hash,state,"
                "reserved_amount,window_start,window_end,expires_at) "
                "VALUES (%s,%s,%s,%s,%s,%s,%s,'reserved',%s,%s,%s,%s) RETURNING *",
                (
                    self.数据库.用途.value,
                    调用ID,
                    self.账户ID,
                    领取.任务ID,
                    领取.尝试ID,
                    角色,
                    请求哈希,
                    每次,
                    窗口.起点,
                    窗口.终点,
                    到期,
                ),
            ).fetchone()
            return self._预留(行)

    def 标记已发送(
        self,
        领取: 领取凭证,
        调用ID: str,
        *,
        最长秒: float,
        登记尝试: bool = False,
        发送复检: Callable[[psycopg.Connection], None] | None = None,
    ) -> 预算预留:
        if not math.isfinite(最长秒) or 最长秒 <= 0:
            raise 预算错误("调用总期限必须大于零")
        self.回收过期()
        with self._事务() as 连:
            存储 = 任务存储(连, self.数据库.用途)
            _, 尝试 = 校验领取(存储, 领取)
            self._锁账户(连)
            行 = self._要求调用(连, 调用ID)
            if (str(行["task_id"]), str(行["attempt_id"])) != (领取.任务ID, 领取.尝试ID):
                raise 预算状态冲突("调用预留不属于当前尝试")
            if 行["state"] == "in_flight":
                return self._预留(行)
            # 用例授权检查和发送状态共用事务；回调不能做模型调用或另开写事务。
            if 发送复检 is not None:
                发送复检(连)
                校验领取(存储, 领取)
            # 锁等待也消耗租期与预算有效期，不能沿用复检前的时间。
            当前 = self._查询(连, "SELECT clock_timestamp() AS current_time").fetchone()[
                "current_time"
            ]
            if 行["state"] != "reserved" or 行["expires_at"] <= 当前:
                raise 预算状态冲突("只有有效且尚未发送的预留可以外发")
            登记 = self._任务预算(连, 领取.任务ID)
            if 登记["stopped"]:
                raise 预算不足("任务预算已经停止")
            if any(r["state"] == "unknown" for r in self._所有调用(连)):
                raise 预算不足("账户存在未知成本，必须先对账")
            到期 = min(
                当前 + timedelta(seconds=最长秒), datetime.fromisoformat(登记["plan"]["截止时间"])
            )
            if 到期 <= 当前:
                raise 预算不足("调用已经超过任务总期限")
            if 登记尝试:
                if 尝试["call_state"] not in {"not_sent", "saved"}:
                    raise 预算状态冲突("本次尝试的上一调用仍未完成证据对账")
                存储.登记调用(领取, 调用ID)
            self._查询(
                连,
                "UPDATE {s}.muse_budget_reservation SET state='in_flight',sent_at=%s,expires_at=%s "
                "WHERE run_purpose=%s AND call_id=%s",
                (当前, 到期, self.数据库.用途.value, 调用ID),
            )
            return replace(self._预留(self._要求调用(连, 调用ID)), 允许外发=True)

    def 结算(self, 调用ID: str, 实际金额: Decimal | None, *, 回执ID: str) -> 预算预留:
        if not 回执ID:
            raise 预算错误("结算需要实际回执身份")
        成本 = None if 实际金额 is None else 金额(实际金额)
        with self._事务() as 连:
            return self.结算于(连, 调用ID, 成本, 回执ID=回执ID)

    def 结算于(self, 连, 调用ID: str, 实际金额: Decimal | None, *, 回执ID: str) -> 预算预留:
        """调用事实与作者结算命令共用事务，不改变任务执行状态。"""
        if not 回执ID:
            raise 预算错误("结算需要实际回执身份")
        成本 = None if 实际金额 is None else 金额(实际金额)
        所属 = self._要求调用(连, 调用ID)
        任务存储(连, self.数据库.用途).任务行(str(所属["task_id"]), 锁定=True)
        self._锁账户(连)
        行 = self._要求调用(连, 调用ID)
        if 行["state"] == "settled":
            if 行["receipt_id"] != 回执ID or 行["actual_amount"] != 成本:
                raise 预算状态冲突("调用已经结算，不能改写成本或回执")
            return self._预留(行)
        if 行["state"] not in ("in_flight", "unknown"):
            raise 预算状态冲突("未发送的调用不能记作实际成本")
        超额 = 成本 is not None and 成本 > 行["reserved_amount"]
        self._查询(
            连,
            "UPDATE {s}.muse_budget_reservation SET state=%s,actual_amount=%s,"
            "receipt_id=%s,over_budget=%s WHERE run_purpose=%s AND call_id=%s",
            (
                "unknown" if 成本 is None else "settled",
                成本,
                回执ID,
                超额,
                self.数据库.用途.value,
                调用ID,
            ),
        )
        self._查询(
            连,
            "UPDATE {s}.muse_task SET accounting_revision=accounting_revision+1 WHERE task_id=%s",
            (行["task_id"],),
        )
        if 超额:
            self._查询(
                连,
                "UPDATE {s}.muse_task_budget SET stopped=true WHERE task_id=%s",
                (行["task_id"],),
            )
        return self._预留(self._要求调用(连, 调用ID))

    def 取消预留(self, 调用ID: str) -> 预算预留:
        with self._事务() as 连:
            self._锁账户(连)
            行 = self._要求调用(连, 调用ID)
            self._取消行(连, 行)
            return self._预留(self._要求调用(连, 调用ID))

    def 关闭任务预算(self, 任务ID: str) -> None:
        with self._事务() as 连:
            存储 = 任务存储(连, self.数据库.用途)
            存储.任务行(任务ID, 锁定=True)
            停止任务预算(存储, 任务ID)

    def _取消行(self, 连: psycopg.Connection, 行: dict) -> None:
        if 行["state"] in ("reserved", "in_flight"):
            self._查询(
                连,
                "UPDATE {s}.muse_budget_reservation SET state=%s,reason='cancelled' "
                "WHERE run_purpose=%s AND call_id=%s",
                (
                    "released" if 行["state"] == "reserved" else "unknown",
                    self.数据库.用途.value,
                    行["call_id"],
                ),
            )

    def 读取(self, 调用ID: str) -> 预算预留:
        with self._事务() as 连:
            return self.读取于(连, 调用ID)

    def 读取于(self, 连, 调用ID: str) -> 预算预留:
        """在调用方快照核验实际结算；不锁账户或改变状态。"""
        return self._预留(self._要求调用(连, 调用ID))

    def 窗口余额(self) -> dict:
        self.回收过期()
        with self._事务() as 连:
            策略 = self._锁账户(连)
            当前 = self._查询(连, "SELECT clock_timestamp() AS current_time").fetchone()[
                "current_time"
            ]
            窗口 = 归属窗口(当前, 策略["时区"])
            所有 = self._所有调用(连)
            本窗 = [r for r in 所有 if r["window_start"] == 窗口.起点]
            已知 = sum((r["actual_amount"] for r in 本窗 if r["state"] == "settled"), Decimal(0))
            在途 = sum((r["reserved_amount"] for r in 本窗 if r["state"] != "settled"), Decimal(0))
            return {
                "已知成本": 已知,
                "在途预留": 在途,
                "已占次数": len(本窗),
                "未知调用数": sum(r["state"] == "unknown" for r in 所有),
                "可用金额": max(Decimal(0), Decimal(策略["窗口金额上限"]) - 已知 - 在途),
                "窗口起点": 窗口.起点,
                "窗口终点": 窗口.终点,
            }

    def _要求调用(self, 连: psycopg.Connection, 调用ID: str) -> dict:
        行 = self._调用行(连, 调用ID)
        if 行 is None:
            raise 预算状态冲突("调用预算不存在或不属于本账户用途")
        return 行

    @staticmethod
    def _占用(行: dict) -> Decimal:
        return 行["actual_amount"] if 行["state"] == "settled" else 行["reserved_amount"]

    @staticmethod
    def _预留(行: dict) -> 预算预留:
        return 预算预留(
            行["call_id"],
            str(行["task_id"]),
            str(行["attempt_id"]),
            行["role_id"],
            行["state"],
            行["reserved_amount"],
            行["actual_amount"],
            行["expires_at"],
            行["receipt_id"],
            行["over_budget"],
        )


def 恢复任务预算(存储: 任务存储, 任务ID: str) -> None:
    """作者恢复在原批准范围内重新开放剩余额度，不重置已使用次数和费用。"""
    预算 = 存储.查询(
        "SELECT account_id,plan FROM {s}.muse_task_budget WHERE task_id=%s", (任务ID,)
    ).fetchone()
    if 预算 is None:
        return
    存储.查询(
        "SELECT account_id FROM {s}.muse_budget_account "
        "WHERE run_purpose=%s AND account_id=%s FOR UPDATE",
        (存储.用途.value, 预算["account_id"]),
    )
    当前 = 存储.查询("SELECT clock_timestamp() AS now").fetchone()["now"]
    if datetime.fromisoformat(预算["plan"]["截止时间"]) <= 当前:
        raise 预算不足("原任务预算已经截止")
    if 存储.查询(
        "SELECT call_id FROM {s}.muse_budget_reservation WHERE task_id=%s "
        "AND state IN ('unknown','in_flight') LIMIT 1",
        (任务ID,),
    ).fetchone():
        raise 预算不足("恢复前必须对账所有在途调用")
    存储.查询("UPDATE {s}.muse_task_budget SET stopped=false WHERE task_id=%s", (任务ID,))


def 停止任务预算(存储: 任务存储, 任务ID: str) -> None:
    """与任务终止共用短事务；按任务、账户顺序锁定，不丢失已发出的成本。"""
    预算 = 存储.查询(
        "SELECT account_id FROM {s}.muse_task_budget WHERE task_id=%s", (任务ID,)
    ).fetchone()
    if 预算 is None:
        return
    存储.查询(
        "SELECT account_id FROM {s}.muse_budget_account "
        "WHERE run_purpose=%s AND account_id=%s FOR UPDATE",
        (存储.用途.value, 预算["account_id"]),
    )
    存储.查询("UPDATE {s}.muse_task_budget SET stopped=true WHERE task_id=%s", (任务ID,))
    存储.查询(
        "UPDATE {s}.muse_budget_reservation SET state=CASE WHEN state='reserved' "
        "THEN 'released' ELSE 'unknown' END,reason='task_stopped' "
        "WHERE task_id=%s AND state IN ('reserved','in_flight')",
        (任务ID,),
    )


__all__ = [
    "额度策略",
    "额度窗口",
    "归属窗口",
    "角色预算",
    "任务预算计划",
    "预算预留",
    "预算管理",
    "预算错误",
    "预算不足",
    "预算状态冲突",
    "金额",
]