codex-app-server/cloud/quota.py

"""成本与配额的唯一判定源。

收敛前同一口径写在三处，注释里明文写着“改口径时三处
要一起改”（server._quota_status、worker._quota_exceeded、
automations._automation_quota_exceeded），且三处已经
不一致——这是本轨的 P0：
- 无价 token 上限的计数口径：worker 与 automations 只计
  source='actual' 的行，服务端 _quota_unpriced_tokens 计
  全部 usd IS NULL 的行（含 estimated 估算行）；而
  cost.py 的 P1 决议、CostStore._report 的
  unpriced_actual_tokens、服务端自己的项目成本报表
  （actual-only）都站在 actual-only 一边。服务端是孤例。
  本轨以 actual-only 为 canonical：估算行不是真实花费，
  绝不能触发上限；服务端 _quota_unpriced_tokens 改为本
  模块口径（行为变化仅限含 estimated 无价行的用户：此
  前可能被估算行误拦，此后不再误拦；真实花费的拦阻不变）。
- 默认配额解析：server、worker、automations 各一份
  CLOUD_DEFAULT_QUOTA_USD_MICROS 解析，现统一为
  default_quota_usd_micros（默认值 $10/月 = 10_000_000
  微美元，非法值回落，无配额行按默认值、显式 NULL 行
  才是不限额——三处旧语义一致，收敛不改语义）。
- 比较判定：used >= quota（USD）先于 unpriced >= cap
  （无价），两者任一命中即阻断；该顺序与三处旧代码一致，
  现收为 decide() 一处。调用方只负责取数与映射错误码
  （服务端 402、worker interrupted、自动化 skipped_quota），
  不得在调用点另写一套 >= 比较。

本模块只依赖标准库，是叶子模块：调用方传一个 DBAPI
连接（? 占位符，Postgres 适配层同款翻译），不 import
cloud 其它模块，避免循环。无价上限的数值本身仍由
cost.unpriced_token_cap() 拥有（环境变量
CLOUD_QUOTA_UNPRICED_TOKEN_CAP），decide() 只接受
cap 作参数，不再另起一套上限解析。

调阈边界（写清，不越界）：
- 本模块只判定、不建议、不生效。没有任何函数会算
  “建议把上限调到 X”，也没有任何函数会写回配额行或
  改环境变量。阈值变化的唯一生效路径是运维改配置：
  改 user_quotas 行（立即对下次查询生效）或改环境变量
  后重启进程（进程内 env 按调用时读取）。未来如需
  “自动算建议”，必须是独立的只读函数，其输出绝不能
  被 decide() 或任何准入点自动采纳。
- 归属口径冻结：按任务创建者 created_by 算，不按项目、
  不按操作人；协调员子任务无 created_by 时跟父任务走
  （该回查仍在 worker 侧，属主解析不是本模块的职责）。
- 窗口口径冻结：UTC 自然月，与 cost_records.day 同口径
  （day LIKE 'YYYY-MM%'）。
"""

import os
from datetime import datetime, timezone

DEFAULT_QUOTA_USD_MICROS = 10000000


def default_quota_usd_micros():
    """默认月配额（微美元）：环境变量非法时回落 $10。"""
    try:
        return max(0, int(os.environ.get(
            "CLOUD_DEFAULT_QUOTA_USD_MICROS", "10000000")))
    except (TypeError, ValueError):
        return DEFAULT_QUOTA_USD_MICROS


def utc_month(now=None):
    """UTC 自然月 'YYYY-MM'（与 cost_records.day 同口径）。"""
    moment = datetime.now(timezone.utc) if now is None else now
    return moment.strftime("%Y-%m")


def _row_value(row, name):
    try:
        return row[name]
    except (KeyError, IndexError, TypeError):
        return None


def quota_for_owner(conn, owner):
    """属主的配额行：无行回默认值，显式 NULL 行回 None（不限额）。

    只读：无行时不在此写行（写行由服务端的
    _apply_default_quota / _backfill_default_quota 负责）。
    表不存在时先幂等建表（与三处旧代码同款）。
    """
    conn.execute(
        "CREATE TABLE IF NOT EXISTS user_quotas("
        "user_id TEXT PRIMARY KEY, quota_usd_micros BIGINT,"
        " updated_at DOUBLE PRECISION)")
    row = conn.execute(
        "SELECT quota_usd_micros FROM user_quotas"
        " WHERE user_id = ?",
        (owner,)).fetchone()
    if row is None:
        return default_quota_usd_micros()
    return _row_value(row, "quota_usd_micros")


def monthly_usage(conn, owner, month):
    """当月已记账用量：(used_usd_micros, unpriced_actual_tokens)。

    USD：cost_records 与 tasks 按 created_by 连接，SUM
    usd_micros（estimated 行 usd 为 NULL，SUM 自动跳过，
    与旧三处一致）。无价：usd IS NULL 且 source='actual'
    的行 input + output 之和——estimated 行一律排除（见
    模块头 P0）。调用方须保证 cost_records 表已存在
    （服务端/worker 走 CostStore 懒建，自动化走表存在探针）。
    """
    used_row = conn.execute(
        "SELECT COALESCE(SUM(c.usd_micros), 0)"
        " FROM cost_records c JOIN tasks t ON c.task_id = t.id"
        " WHERE t.created_by = ? AND c.day LIKE ?",
        (owner, month + "%")).fetchone()
    used = _row_value(used_row, 0)
    if used is None:
        # sqlite3.Row 按列名取不到时按位置取。
        try:
            used = used_row[0]
        except (IndexError, TypeError):
            used = 0
    unpriced_row = conn.execute(
        "SELECT COALESCE(SUM(c.input_tokens + c.output_tokens), 0)"
        " FROM cost_records c JOIN tasks t ON c.task_id = t.id"
        " WHERE t.created_by = ? AND c.day LIKE ?"
        " AND c.usd_micros IS NULL AND c.source = 'actual'",
        (owner, month + "%")).fetchone()
    unpriced = _row_value(unpriced_row, 0)
    if unpriced is None:
        try:
            unpriced = unpriced_row[0]
        except (IndexError, TypeError):
            unpriced = 0
    return int(used or 0), int(unpriced or 0)


def decide(quota, used_usd, unpriced_used, cap, month):
    """唯一比较判定：放行回 None，否则回决定 dict。

    quota 为 None（显式不限额）时永远放行，即使无价
    已超——与服务端、worker 旧行为一致（不限额用户不
    受无价上限约束）。USD 门先于无价门；命中条件均为
    >=（用满即拦，不是超过才拦）。
    """
    if quota is None:
        return None
    if used_usd >= quota:
        return {
            "kind": "usd",
            "quota_usd_micros": quota,
            "used_usd_micros": used_usd,
            "month": month,
            "message": ("quota_exceeded: 本月配额已用完（已用 %d / "
                        "配额 %d 微美元，%s）" % (used_usd, quota, month)),
        }
    if unpriced_used >= cap:
        return {
            "kind": "unpriced",
            "unpriced_token_cap": cap,
            "unpriced_tokens_used": unpriced_used,
            "month": month,
            "message": ("quota_exceeded: 本月无价模型 token 已用完"
                        "（已用 %d / 上限 %d，%s）"
                        % (unpriced_used, cap, month)),
        }
    return None