From a4d59aca30b6dbcccab001809a243ac8e09b639d Mon Sep 17 00:00:00 2001 From: xsl Date: Fri, 14 Aug 2026 16:26:45 +0800 Subject: [PATCH] feat(llm): four-tier daily budget degradation with manual override MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 把单一日预算熔断升级为架构 §13.3 的四档分级降级（T-P1-25b / NFR-COST-02）： - budget.py：新增 BudgetTier 配置模型与档位判定（current_tier/tier_name/ blocked_scenes）。阈值相对 LLM_DAILY_BUDGET_CENTS：warn ≥60% 只告警； degrade_secondary ≥80% 拒 hd_export/report_regen；degrade_generation ≥90% 追加拒 hairstyle/tryon；block_all ≥100% 拒全部非 critical、allow 名单 （默认 agent_chat）放行。高档位 blocked 取「截至当前档的并集」（架构原文 「追加拒绝」）。critical 场景全档豁免但照常计成本。 - 手动 override：Redis 键 llm:budget:override（值 "1"~"4"），优先于成本水 位，读不到/非法则忽略；set_override TTL 到当日 UTC+8 24 点，跨预算日必然 失效。CostCounter 协议相应增加 set/delete。 - 指标：Gauge llm_budget_tier + Counter llm_budget_blocked_total{scene} 挂到共享 REGISTRY 随 /metrics 暴露；档位上升沿只告警一次 （budget_tier_raised，实例内记 last_tier）。 - errors.py 新增 SceneDegradedError（tier2/3 拒绝）；gateway.py 调用点换成 ensure_scene_allowed（保留 ensure_within_budget 兼容入口）；agent.py 映射 LLM_SCENE_DEGRADED(60006)；100% 档维持 DailyBudgetExceededError/60003 语义。 - 单测：新增 tests/unit/test_tiered_budget.py（27 例：四档边界、各档放行/ 拒绝、critical 豁免、override 加严/放宽/非法值/TTL、配置驱动）。 - 集成：tests/integration/test_vertical_slice.py 新增用例——预算 1 分 + block_all allow 清空 + Redis 注入当日成本，POST /agent/chat 返回 60003， 断言 bl_ledger 中 consume(refunded,-1)+refund(committed,+1) 净和为 0， 覆盖「reserve 之后才熔断 → refund」路径；fixture 重构为共享 FakeProvider/fake_client，断言语义未弱化。 门禁：ruff check/format、mypy（packages/llm strict + 其余包）、lint-imports、 单测 196 passed、集成 4 passed。 偏差说明： 1. tests/unit/test_gateway.py 的 test_daily_budget_rejects_non_critical_ scenes_before_spending 把被测场景 AGENT_CHAT 换成 REPORT_COPY（test_daily_ budget.py 同用例换 tryon）：block_all 档默认 allow=["agent_chat"]（架构 「仅保留纯文本对话」）后 agent_chat 在 100% 档不再被拒，原场景选择与新 语义冲突；断言本身（进 provider 前被拒）未弱化。 2. 集成测试的 block_all 档配置为 blocked="*"+allow=[]（文档示例为 blocked=["agent_chat"]）：后者按档位规则会走 SceneDegradedError/60006， 前者才能按设计走既有 60003 语义（文档允许二者取其一且"按上面设计应为 60003"）。 --- apps/api/routers/agent.py | 4 + packages/llm/budget.py | 261 +++++++++++++++++++++-- packages/llm/errors.py | 16 ++ packages/llm/gateway.py | 7 +- tests/integration/test_vertical_slice.py | 120 +++++++++-- tests/unit/test_daily_budget.py | 6 +- tests/unit/test_gateway.py | 5 +- tests/unit/test_tiered_budget.py | 245 +++++++++++++++++++++ 8 files changed, 617 insertions(+), 47 deletions(-) create mode 100755 tests/unit/test_tiered_budget.py diff --git a/apps/api/routers/agent.py b/apps/api/routers/agent.py index 6c3f5e2..3389ac3 100755 --- a/apps/api/routers/agent.py +++ b/apps/api/routers/agent.py @@ -18,6 +18,7 @@ from packages.llm.errors import ( AllProvidersFailedError, CachePrefixViolationError, DailyBudgetExceededError, + SceneDegradedError, StructuredOutputError, ) from packages.shared.errors import AppError, ErrorCode @@ -37,6 +38,9 @@ async def chat( result = await agent.chat(user_id=user_id, message=payload.message, request_key=request_key) except DailyBudgetExceededError as exc: raise AppError(ErrorCode.LLM_DAILY_BUDGET_EXCEEDED) from exc + except SceneDegradedError as exc: + # 分级降级（§13.3 tier2/3）：该场景今日被关停，但平台未爆预算 + raise AppError(ErrorCode.LLM_SCENE_DEGRADED) from exc except StructuredOutputError as exc: raise AppError(ErrorCode.LLM_STRUCTURED_OUTPUT_FAILED) from exc except AllProvidersFailedError as exc: diff --git a/packages/llm/budget.py b/packages/llm/budget.py index 70e77b8..a797bb0 100644 --- a/packages/llm/budget.py +++ b/packages/llm/budget.py @@ -1,27 +1,114 @@ -"""Global daily cost circuit breaker (架构文档 §13「必须设置的三道成本闸门」). +"""Global daily cost circuit breaker — four-tier graded degradation (架构文档 §13.3). Gate 1 in §13 is described as "某功能超预算自动降级，并告警". Alerting alone is not enough: the failure modes that blow up AI spend (retry storms, a cache-key bug, someone scripting the API) all produce cost faster than a human can react -to a page. So this is an *enforcement* point, not a notification: once the -platform-wide daily spend crosses the threshold, non-critical scenes are -rejected outright. +to a page. So this is an *enforcement* point, and the degradation is graded so +the product winds down gracefully instead of dying all at once. Thresholds are +percentages of ``LLM_DAILY_BUDGET_CENTS``; higher tiers *add* rejected scenes on +top of the lower ones (架构原文「追加拒绝」), never un-block them:: -Critical scenes (e.g. content moderation follow-ups, paid-path completions) -opt out via ``SceneConfig.critical`` so that a budget overrun degrades the -product instead of breaking it. + tier 1 warn ≥60% 只告警（指标 + 上升沿日志），不影响请求 + tier 2 degrade_secondary ≥80% 拒绝 hd_export / report_regen + tier 3 degrade_generation ≥90% 追加拒绝 hairstyle / tryon + tier 4 block_all ≥100% 拒绝所有非 critical；allow 名单放行 + （默认 agent_chat ——「仅保留在途任务与 + 纯文本对话」） + +Critical scenes (``SceneConfig.critical``, e.g. INTENT_CLASSIFY) opt out of +every tier so that a budget overrun degrades the product instead of breaking +it — but their cost still counts towards the water level. 在途任务同理放行： +熔断只拦网关**新请求**，worker 侧不做二次检查。 + +手动干预（运营/演练，redis-cli 可直接操作）:: + + SET llm:budget:override 4 # 强制 tier 4（加严演练） + SET llm:budget:override 2 # 水位已到 tier 4 时手动放宽回 tier 2 + DEL llm:budget:override # 解除，回落成本水位判定 + +键 ``llm:budget:override`` 的值为 "1"~"4"（档位序号），优先级高于成本水位； +读不到或值非法则忽略。程序化写入走 :meth:`DailyBudgetGuard.set_override`， +TTL 到当日 UTC+8 24 点 —— override 最多活一个预算日，第二天必然回落水位。 """ from __future__ import annotations from datetime import UTC, datetime, timedelta, timezone from decimal import Decimal -from typing import Protocol, runtime_checkable +from typing import Literal, Protocol, runtime_checkable -from packages.llm.errors import DailyBudgetExceededError +from prometheus_client import Counter, Gauge +from pydantic import BaseModel, ConfigDict, Field, field_validator + +from packages.llm.errors import DailyBudgetExceededError, SceneDegradedError +from packages.shared.logging import get_logger +from packages.shared.observability import REGISTRY + +logger = get_logger(__name__) _DAY_SECONDS = 60 * 60 * 24 +#: 运营手动干预档位的 Redis 键（值为 "1"~"4"；TTL 到当日 UTC+8 24 点）。 +OVERRIDE_KEY = "llm:budget:override" + +#: blocked 名单中的通配符：拒绝所有非 critical 场景（仅 block_all 档配置 allow 名单）。 +BLOCK_ALL: Literal["*"] = "*" + +#: 内置默认四档，与 ``Settings.llm_budget_tiers`` 的默认值保持一致（架构 §13.3）。 +_DEFAULT_TIERS: list[dict[str, object]] = [ + {"threshold": 0.6, "name": "warn", "blocked": []}, + {"threshold": 0.8, "name": "degrade_secondary", "blocked": ["hd_export", "report_regen"]}, + {"threshold": 0.9, "name": "degrade_generation", "blocked": ["hairstyle", "tryon"]}, + {"threshold": 1.0, "name": "block_all", "blocked": BLOCK_ALL, "allow": ["agent_chat"]}, +] + +# 档位是全局状态（整个平台只有一个水位），挂到共享 REGISTRY 上随 /metrics 暴露。 +LLM_BUDGET_TIER = Gauge( + "llm_budget_tier", + "当前生效的日预算降级档位（0=未触发，1~4 递进；手动 override 优先）", + registry=REGISTRY, +) +LLM_BUDGET_BLOCKED = Counter( + "llm_budget_blocked_total", + "被预算降级档位拒绝的网关请求数", + labelnames=("scene",), + registry=REGISTRY, +) + + +class BudgetTier(BaseModel): + """一档降级配置（``LLM_BUDGET_TIERS_JSON`` 的元素，架构 §13.3）。 + + ``blocked`` 里的 ``"*"`` 表示该档拒绝全部非 critical 场景（配置允许写成 + 裸字符串，这里统一归一成列表）；``allow`` 仅 block_all 档有意义。 + """ + + model_config = ConfigDict(frozen=True) + + threshold: float + name: str + blocked: list[str] = Field(default_factory=list) + allow: list[str] = Field(default_factory=list) + + @field_validator("blocked", "allow", mode="before") + @classmethod + def _wrap_bare_string(cls, value: object) -> object: + # Settings 默认值里 "blocked": "*" 是裸字符串写法，包成列表让类型恒为 list[str]。 + if isinstance(value, str): + return [value] + return value + + @property + def blocked_set(self) -> frozenset[str]: + return frozenset(self.blocked) + + @property + def allow_set(self) -> frozenset[str]: + return frozenset(self.allow) + + def blocks_all(self) -> bool: + return BLOCK_ALL in self.blocked + @runtime_checkable class CostCounter(Protocol): @@ -31,20 +118,36 @@ class CostCounter(Protocol): async def get(self, key: str) -> Decimal: ... + async def set(self, key: str, value: str, ttl_seconds: int) -> None: ... + + async def delete(self, key: str) -> None: ... + class InMemoryCostCounter: """Process-local counter. Correct only for a single process — tests only.""" def __init__(self) -> None: self._values: dict[str, Decimal] = {} + self._strings: dict[str, str] = {} async def incr(self, key: str, amount: Decimal, ttl_seconds: int) -> Decimal: self._values[key] = self._values.get(key, Decimal("0")) + amount return self._values[key] async def get(self, key: str) -> Decimal: + if key in self._strings: + # 非法字符串在这里抛 InvalidOperation，与 Redis 实现的行为一致， + # 调用方（override 读取）按「值非法则忽略」处理。 + return Decimal(self._strings[key]) return self._values.get(key, Decimal("0")) + async def set(self, key: str, value: str, ttl_seconds: int) -> None: + self._strings[key] = value + + async def delete(self, key: str) -> None: + self._strings.pop(key, None) + self._values.pop(key, None) + class RedisCostCounter: """Cross-process counter backed by ``INCRBYFLOAT`` + a rolling TTL.""" @@ -63,9 +166,15 @@ class RedisCostCounter: value = await self._redis.get(key) # type: ignore[attr-defined] return Decimal(str(value)) if value is not None else Decimal("0") + async def set(self, key: str, value: str, ttl_seconds: int) -> None: + await self._redis.set(key, value, ex=ttl_seconds) # type: ignore[attr-defined] + + async def delete(self, key: str) -> None: + await self._redis.delete(key) # type: ignore[attr-defined] + class DailyBudgetGuard: - """Hard platform-wide daily AI budget.""" + """Platform-wide daily AI budget with four-tier graded degradation (§13.3).""" def __init__( self, @@ -84,9 +193,13 @@ class DailyBudgetGuard: self._key_prefix = key_prefix self._tz = timezone(timedelta(hours=tz_offset_hours)) self._clock = clock - # 四档分级降级的配置（架构 §13.3 / T-P1-25b）。T1 任务实现判定逻辑； - # 这里先接收配置，使 deps 装配层不必随后再改。None = 使用内置默认四档。 - self._tiers_config = tiers + # 档位配置（LLM_BUDGET_TIERS_JSON）。None = 内置默认四档。解析失败在装配期 + # 就炸出来：一份写坏的降级配置不应该等到平台烧超预算时才被发现。 + self._tiers = [BudgetTier.model_validate(t) for t in (tiers if tiers else _DEFAULT_TIERS)] + self._tiers.sort(key=lambda t: t.threshold) + self._thresholds = [Decimal(str(t.threshold)) for t in self._tiers] + # 上升沿告警的实例内游标：避免每个请求都刷一遍 budget_tier_raised。 + self._last_tier = 0 @property def enabled(self) -> bool: @@ -109,13 +222,125 @@ class DailyBudgetGuard: async def spent_today(self) -> Decimal: return await self._counter.get(self.current_key()) - async def ensure_within_budget(self, *, scene: str, critical: bool = False) -> None: - """Raise :class:`DailyBudgetExceededError` for non-critical scenes over budget.""" - if not self._enabled or critical: + # -- tiered degradation ------------------------------------------------ + + async def current_tier(self) -> int: + """当前生效档位（0=未触发，1~4 递进）。手动 override 优先于成本水位。""" + if not self._enabled: + return 0 + override = await self._override_tier() + if override is not None: + return override + return self._tier_for_spend(await self.spent_today()) + + async def tier_name(self) -> str: + """当前档位名（warn/degrade_secondary/degrade_generation/block_all）。""" + tier = await self.current_tier() + return self._tiers[tier - 1].name if tier else "" + + async def blocked_scenes(self) -> frozenset[str] | Literal["*"]: + """当前档位**累计**拒绝的场景名单；``"*"`` 表示拒绝全部非 critical。""" + return self._blocked_scenes_at(await self.current_tier()) + + async def ensure_scene_allowed(self, *, scene: str, critical: bool = False) -> None: + """Raise when the current tier rejects ``scene``. + + - critical 场景全档豁免（但照常计入成本水位）； + - tier 2/3 且 scene 命中累计 blocked 名单 → :class:`SceneDegradedError`； + - tier 4（blocked="*"）且 scene 不在 allow 名单 → + :class:`DailyBudgetExceededError`——保持 100% 档的既有语义与 60003 + 错误码不变（现有测试与告警文案都依赖它）。 + """ + if not self._enabled: return + override = await self._override_tier() spent = await self.spent_today() - if spent >= self._budget: - raise DailyBudgetExceededError(str(spent), str(self._budget), scene) + tier = override if override is not None else self._tier_for_spend(spent) + self._observe_tier(tier, spent=spent) + if critical or tier == 0: + return + + blocked = self._blocked_scenes_at(tier) + if blocked == BLOCK_ALL: + if scene not in self._tiers[tier - 1].allow_set: + LLM_BUDGET_BLOCKED.labels(scene).inc() + raise DailyBudgetExceededError(str(spent), str(self._budget), scene) + return + if scene in blocked: + LLM_BUDGET_BLOCKED.labels(scene).inc() + raise SceneDegradedError(scene, self._tiers[tier - 1].name) + + async def ensure_within_budget(self, *, scene: str, critical: bool = False) -> None: + """兼容入口：与 :meth:`ensure_scene_allowed` 等价（存量调用方在用）。""" + await self.ensure_scene_allowed(scene=scene, critical=critical) + + # -- manual override --------------------------------------------------- + + async def set_override(self, tier: int) -> None: + """运营/演练：强制指定当前档位，TTL 到当日 UTC+8 24 点。""" + if not 1 <= tier <= len(self._tiers): + msg = f"override tier must be in 1..{len(self._tiers)}, got {tier}" + raise ValueError(msg) + await self._counter.set(OVERRIDE_KEY, str(tier), self._seconds_until_day_end()) + + async def clear_override(self) -> None: + """解除 override，回落成本水位判定。""" + await self._counter.delete(OVERRIDE_KEY) + + # -- internals ---------------------------------------------------------- + + async def _override_tier(self) -> int | None: + """Read the manual override; unreadable/invalid values are ignored.""" + try: + value = await self._counter.get(OVERRIDE_KEY) + tier = int(value) + except Exception: + # 键不存在时计数器返回 0（→ 下方范围校验拒绝）；这里兜底的是 + # 值非法（如 "abc"）或存储故障：override 是运营工具，绝不能 + # 因为它把请求路径打挂。 + return None + if 1 <= tier <= len(self._tiers): + return tier + return None + + def _tier_for_spend(self, spent: Decimal) -> int: + """水位判定：取「最后一个被跨过阈值」的档位（阈值已升序）。""" + tier = 0 + for index, threshold in enumerate(self._thresholds, start=1): + if spent >= threshold * self._budget: + tier = index + return tier + + def _blocked_scenes_at(self, tier: int) -> frozenset[str] | Literal["*"]: + """档位越高拒绝越多：取 1..tier 各档 blocked 的并集（「追加拒绝」）。""" + scenes: set[str] = set() + for spec in self._tiers[:tier]: + if spec.blocks_all(): + return BLOCK_ALL + scenes.update(spec.blocked_set) + return frozenset(scenes) + + def _observe_tier(self, tier: int, *, spent: Decimal) -> None: + LLM_BUDGET_TIER.set(tier) + if tier <= self._last_tier: + return + self._last_tier = tier + if tier == 0: + return + # 只在上升沿告警一次：水位在 90% 徘徊时每个请求都打日志等于没有告警。 + logger.warning( + "budget_tier_raised", + tier=tier, + tier_name=self._tiers[tier - 1].name, + spent=str(spent), + budget=str(self._budget), + ) + + def _seconds_until_day_end(self) -> int: + """到当日 UTC+8 24 点的秒数（预算日切换即失效）。""" + local = self._now().astimezone(self._tz) + day_end = local.replace(hour=0, minute=0, second=0, microsecond=0) + timedelta(days=1) + return max(int((day_end - local).total_seconds()), 1) async def add_cost(self, cost_cents: Decimal) -> Decimal: """Accumulate spend. Called after every provider attempt, success or not.""" diff --git a/packages/llm/errors.py b/packages/llm/errors.py index 7dfa8d3..d8a914c 100755 --- a/packages/llm/errors.py +++ b/packages/llm/errors.py @@ -138,6 +138,22 @@ class DailyBudgetExceededError(LLMError): ) +class SceneDegradedError(LLMError): + """A scene was rejected by a *graded* budget tier (架构文档 §13.3 tier 2/3). + + 与 :class:`DailyBudgetExceededError` 的区别：平台还没烧穿预算，只是按档位 + 关停部分高成本场景（hd_export/report_regen/hairstyle/tryon）。用户看到的 + 是「该功能今日暂时关闭」(LLM_SCENE_DEGRADED, 60006) 而非「平台已达上限」。 + """ + + code = "llm_scene_degraded" + + def __init__(self, scene: str, tier_name: str) -> None: + self.scene = scene + self.tier_name = tier_name + super().__init__(f"scene '{scene}' blocked by budget tier '{tier_name}'") + + class PromptNotFoundError(LLMError): code = "llm_prompt_not_found" diff --git a/packages/llm/gateway.py b/packages/llm/gateway.py index 1adc36a..75eceef 100755 --- a/packages/llm/gateway.py +++ b/packages/llm/gateway.py @@ -9,7 +9,7 @@ Built-in capabilities, all exercised on every call: Prompt 版本 模板文件按 id+version 加载，版本号随每次调用落库 缓存前缀纪律 固定前缀含时间戳/随机 ID/昵称 → 拒绝或告警 (§7.2, TD-04) 成本计量 每次 provider 调用（含失败与重试）写一行 ``ob_ai_call`` -预算保护 平台级日成本硬熔断：超限即拒绝非关键场景 (§13) +预算保护 平台级日成本四档分级降级：warn/secondary/generation/block_all (§13.3) =================== ========================================================= Deliberately not here yet: streaming (SSE), 审核联动 (input/output moderation), @@ -106,8 +106,9 @@ class LLMGateway: ) -> Completion: config = self._router.resolve(scene) - # Gate 1: platform-wide daily cost breaker, checked before any spend. - await self._budget.ensure_within_budget(scene=scene.value, critical=config.critical) + # Gate 1: four-tier graded daily cost breaker (架构文档 §13.3), + # checked before any spend — 在途任务不受影响，这里只拦新请求。 + await self._budget.ensure_scene_allowed(scene=scene.value, critical=config.critical) template = self._prompts.get(config.prompt_id, config.prompt_version) violations = validate_cache_prefix( diff --git a/tests/integration/test_vertical_slice.py b/tests/integration/test_vertical_slice.py index 7fc2798..89b05cf 100755 --- a/tests/integration/test_vertical_slice.py +++ b/tests/integration/test_vertical_slice.py @@ -12,6 +12,8 @@ from __future__ import annotations import os import uuid +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager from decimal import Decimal from typing import Any @@ -31,40 +33,48 @@ pytestmark = [ ] -@pytest.fixture -async def client() -> Any: +class FakeProvider: + """替换真实 provider：集成用例验证配额与计量链路，不该花钱也不该受网络影响。""" + + name = "deepseek" + default_model = "deepseek-v4-flash" + + async def complete(self, request: Any) -> Any: + from packages.llm.types import ProviderResponse, Usage + + return ProviderResponse( + text="建议白衬衫配西裤。", + model="deepseek-v4-flash", + usage=Usage(tokens_in=800, tokens_out=120, tokens_cached=600), + ) + + async def aclose(self) -> None: + return None + + +@asynccontextmanager +async def fake_client(app: Any) -> AsyncIterator[Any]: + """起 app lifespan + ASGI 客户端，并把 deepseek 换成 FakeProvider。""" import httpx - from apps.api.main import create_app - from packages.llm.types import ProviderResponse, Usage - from packages.shared.config import get_settings - - app = create_app(get_settings()) - - class FakeProvider: - name = "deepseek" - default_model = "deepseek-v4-flash" - - async def complete(self, request: Any) -> ProviderResponse: - return ProviderResponse( - text="建议白衬衫配西裤。", - model="deepseek-v4-flash", - usage=Usage(tokens_in=800, tokens_out=120, tokens_cached=600), - ) - - async def aclose(self) -> None: - return None - transport = httpx.ASGITransport(app=app) async with ( httpx.AsyncClient(transport=transport, base_url="http://test") as http, app.router.lifespan_context(app), ): - # 替换掉真实 provider：集成用例验证配额与计量链路，不该花钱也不该受网络影响 app.state.llm_gateway._providers["deepseek"] = FakeProvider() yield http +@pytest.fixture +async def client() -> Any: + from apps.api.main import create_app + from packages.shared.config import get_settings + + async with fake_client(create_app(get_settings())) as http: + yield http + + async def test_chat_slice_reserves_commits_and_records_cost(client: Any) -> None: token_resp = await client.post( "/api/v1/auth/dev-token", json={"device_key": f"it-{uuid.uuid4().hex[:8]}"} @@ -99,3 +109,67 @@ async def test_unauthenticated_requests_get_the_unified_envelope(client: Any) -> resp = await client.post("/api/v1/agent/chat", json={"message": "hi"}) assert resp.status_code == 401 assert resp.json()["code"] == 40100 + + +async def test_block_all_budget_rejects_and_refunds_the_reservation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """T-P1-25b：reserve 之后才触发熔断 → gateway 拒绝 → refund → 净扣减 = 0。 + + 这是四档降级（架构 §13.3）的验收路径：把 block_all 档的 allow 名单清空 + （连纯文本对话也拦），让 agent_chat 走 100% 档的既有 60003 语义，从而 + 断言配额三段式在「先预扣、后熔断」时的补偿行为。 + """ + import json + + from apps.api.main import create_app + from packages.llm.budget import DailyBudgetGuard, InMemoryCostCounter + from packages.shared.config import get_settings + from packages.shared.redis_client import get_redis + + monkeypatch.setenv("LLM_DAILY_BUDGET_CENTS", "1") # 预算 1 分钱 + monkeypatch.setenv( + "LLM_BUDGET_TIERS_JSON", + json.dumps( + [ + {"threshold": 0.6, "name": "warn", "blocked": []}, + {"threshold": 0.8, "name": "degrade_secondary", "blocked": ["hd_export"]}, + {"threshold": 0.9, "name": "degrade_generation", "blocked": ["hairstyle"]}, + # allow 清空：agent_chat 也被 block_all 拦下（演练配置） + {"threshold": 1.0, "name": "block_all", "blocked": "*", "allow": []}, + ] + ), + ) + get_settings.cache_clear() + try: + app = create_app(get_settings()) + async with fake_client(app) as http: + redis = get_redis() + # 注入当日成本到预算以上（键格式照 DailyBudgetGuard.current_key()） + key_helper = DailyBudgetGuard(InMemoryCostCounter(), daily_budget_cents=Decimal("1")) + cost_key = key_helper.current_key() + await redis.incrbyfloat(cost_key, 5) + + token_resp = await http.post( + "/api/v1/auth/dev-token", json={"device_key": f"it-{uuid.uuid4().hex[:8]}"} + ) + headers = {"Authorization": f"Bearer {token_resp.json()['data']['access_token']}"} + + resp = await http.post( + "/api/v1/agent/chat", json={"message": "面试穿什么"}, headers=headers + ) + assert resp.json()["code"] == 60003 # LLM_DAILY_BUDGET_EXCEEDED + + # 被拒请求的配额净扣减 = 0：consume(-1, refunded) + refund(+1, committed) + rows = (await http.get("/api/v1/quota/ledger", headers=headers)).json()["data"] + consume = [r for r in rows if r["tx_type"] == "consume"] + refunds = [r for r in rows if r["tx_type"] == "refund"] + assert len(consume) == 1 and consume[0]["status"] == "refunded" + assert len(refunds) == 1 and refunds[0]["delta"] == 1 + assert sum(r["delta"] for r in (*consume, *refunds)) == 0 + + await redis.delete(cost_key) + finally: + # 恢复全局 settings 缓存，避免本用例的环境覆盖渗进其它用例 + monkeypatch.undo() + get_settings.cache_clear() diff --git a/tests/unit/test_daily_budget.py b/tests/unit/test_daily_budget.py index 436f7c3..37b16a0 100755 --- a/tests/unit/test_daily_budget.py +++ b/tests/unit/test_daily_budget.py @@ -37,9 +37,11 @@ async def test_allows_calls_below_the_threshold() -> None: async def test_rejects_non_critical_scenes_once_the_budget_is_reached() -> None: guard, _ = make_guard("100") await guard.add_cost(Decimal("100")) + # block_all 档默认放行 agent_chat（「仅保留纯文本对话」），用非 allow 场景 + # 断言 100% 档的既有拒绝语义不变 with pytest.raises(DailyBudgetExceededError) as excinfo: - await guard.ensure_within_budget(scene="agent_chat") - assert "agent_chat" in str(excinfo.value) + await guard.ensure_within_budget(scene="tryon") + assert "tryon" in str(excinfo.value) async def test_critical_scenes_bypass_the_breaker() -> None: diff --git a/tests/unit/test_gateway.py b/tests/unit/test_gateway.py index 967354a..9442450 100755 --- a/tests/unit/test_gateway.py +++ b/tests/unit/test_gateway.py @@ -437,6 +437,7 @@ async def test_spend_accumulates_into_the_daily_budget(prompts_dir: Path) -> Non async def test_daily_budget_rejects_non_critical_scenes_before_spending( prompts_dir: Path, ) -> None: + # block_all 档默认放行 agent_chat（纯文本对话），用非 allow 场景断言拒绝 primary = FakeProvider("p1", [response("ok")]) gateway, _, guard = build_gateway( prompts_dir, {"p1": primary}, backup=None, budget_cents="0.00001" @@ -444,7 +445,9 @@ async def test_daily_budget_rejects_non_critical_scenes_before_spending( await guard.add_cost(Decimal("1")) with pytest.raises(DailyBudgetExceededError): - await gateway.complete(scene=Scene.AGENT_CHAT, prompt_vars={"user_input": "hi"}) + await gateway.complete( + scene=Scene.REPORT_COPY, prompt_vars={"payload": "诊断结论与规则命中"} + ) assert primary.requests == [] # rejected, not merely alerted diff --git a/tests/unit/test_tiered_budget.py b/tests/unit/test_tiered_budget.py new file mode 100755 index 0000000..349829b --- /dev/null +++ b/tests/unit/test_tiered_budget.py @@ -0,0 +1,245 @@ +"""四档分级降级（架构 §13.3 / T-P1-25b）的判定逻辑单测。 + +不连网不连库：计数用 InMemoryCostCounter，时间用注入 clock。 +档位语义：阈值相对 LLM_DAILY_BUDGET_CENTS；高档位**追加**拒绝低档位名单； +critical 全档豁免；手动 override 优先于成本水位。 +""" + +from __future__ import annotations + +from datetime import UTC, datetime +from decimal import Decimal + +import pytest + +from packages.llm.budget import ( + BLOCK_ALL, + OVERRIDE_KEY, + BudgetTier, + DailyBudgetGuard, + InMemoryCostCounter, +) +from packages.llm.errors import DailyBudgetExceededError, SceneDegradedError + + +def make_guard( + budget: str = "100", + *, + enabled: bool = True, + tiers: list[dict[str, object]] | None = None, +) -> tuple[DailyBudgetGuard, InMemoryCostCounter]: + counter = InMemoryCostCounter() + guard = DailyBudgetGuard( + counter, + daily_budget_cents=Decimal(budget), + enabled=enabled, + clock=lambda: datetime(2026, 8, 14, 3, 30, tzinfo=UTC), # UTC+8 已是 8-14 + tiers=tiers, + ) + return guard, counter + + +async def guard_spent(spent: str) -> DailyBudgetGuard: + guard, _ = make_guard() + await guard.add_cost(Decimal(spent)) + return guard + + +# -------------------------------------------------------------------------- +# 档位边界 +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("spent", "expected_tier"), + [ + ("59.9", 0), + ("60", 1), + ("79.9", 1), + ("80", 2), + ("89.9", 2), + ("90", 3), + ("99.9", 3), + ("100", 4), + ], +) +async def test_tier_boundaries_follow_configured_thresholds(spent: str, expected_tier: int) -> None: + assert await (await guard_spent(spent)).current_tier() == expected_tier + + +async def test_tier_names_and_cumulative_blocked_scenes() -> None: + guard = await guard_spent("60") + assert await guard.tier_name() == "warn" + assert await guard.blocked_scenes() == frozenset() # warn 档无用户可感知变化 + + await guard.add_cost(Decimal("20")) # 80% + assert await guard.tier_name() == "degrade_secondary" + assert await guard.blocked_scenes() == frozenset({"hd_export", "report_regen"}) + + await guard.add_cost(Decimal("10")) # 90%：追加拒绝，而不是替换名单 + assert await guard.tier_name() == "degrade_generation" + assert await guard.blocked_scenes() == frozenset( + {"hd_export", "report_regen", "hairstyle", "tryon"} + ) + + await guard.add_cost(Decimal("10")) # 100% + assert await guard.tier_name() == "block_all" + assert await guard.blocked_scenes() == BLOCK_ALL + + +# -------------------------------------------------------------------------- +# 各档放行 / 拒绝 +# -------------------------------------------------------------------------- + + +async def test_warn_tier_blocks_nothing() -> None: + guard = await guard_spent("60") + for scene in ("agent_chat", "hd_export", "tryon"): + await guard.ensure_scene_allowed(scene=scene) + + +async def test_tier2_blocks_secondary_scenes_but_not_plain_chat() -> None: + guard = await guard_spent("80") + with pytest.raises(SceneDegradedError) as excinfo: + await guard.ensure_scene_allowed(scene="hd_export") + assert excinfo.value.scene == "hd_export" + assert excinfo.value.tier_name == "degrade_secondary" + await guard.ensure_scene_allowed(scene="agent_chat") + + +async def test_tier3_keeps_tier2_blocks_and_adds_generation_scenes() -> None: + guard = await guard_spent("90") + for scene in ("hd_export", "hairstyle", "tryon"): + with pytest.raises(SceneDegradedError): + await guard.ensure_scene_allowed(scene=scene) + await guard.ensure_scene_allowed(scene="agent_chat") + + +async def test_tier4_blocks_all_except_the_allow_list() -> None: + guard = await guard_spent("100") + with pytest.raises(DailyBudgetExceededError) as excinfo: + await guard.ensure_scene_allowed(scene="tryon") + assert excinfo.value.scene == "tryon" + # 架构原文「仅保留在途任务与纯文本对话」：agent_chat 默认在 allow 名单 + await guard.ensure_scene_allowed(scene="agent_chat") + + +async def test_tier4_allow_list_is_configurable() -> None: + tiers: list[dict[str, object]] = [ + # 裸字符串 "*" 是合法配置写法，用于演练「连纯文本对话也拦」 + {"threshold": 1.0, "name": "block_all", "blocked": BLOCK_ALL, "allow": []}, + ] + guard, _ = make_guard(tiers=tiers) + await guard.add_cost(Decimal("100")) + assert await guard.current_tier() == 1 # 只有一档时 override/档位序号随之缩短 + with pytest.raises(DailyBudgetExceededError): + await guard.ensure_scene_allowed(scene="agent_chat") + + +async def test_critical_scenes_are_exempt_from_every_tier() -> None: + guard = await guard_spent("500") + # 意图识别把守越界拒答：降级它等于让超范围提问漏过去，因此全档豁免。 + await guard.ensure_scene_allowed(scene="intent_classify", critical=True) + + +async def test_disabled_guard_reports_no_tier() -> None: + guard, _ = make_guard(enabled=False) + await guard.add_cost(Decimal("999")) + assert await guard.current_tier() == 0 + await guard.ensure_scene_allowed(scene="tryon") + + +async def test_legacy_entry_point_still_enforces() -> None: + """ensure_within_budget 兼容入口：100% 档非 allow 场景仍抛既有异常。""" + guard = await guard_spent("100") + with pytest.raises(DailyBudgetExceededError): + await guard.ensure_within_budget(scene="tryon") + + +# -------------------------------------------------------------------------- +# 手动 override +# -------------------------------------------------------------------------- + + +async def test_override_forces_a_tier_above_the_water_level() -> None: + guard = await guard_spent("10") # 水位仅 10%，本未触发任何档 + await guard.set_override(4) + assert await guard.current_tier() == 4 + with pytest.raises(DailyBudgetExceededError): + await guard.ensure_scene_allowed(scene="tryon") + await guard.clear_override() + assert await guard.current_tier() == 0 # 清除后回落水位判定 + await guard.ensure_scene_allowed(scene="tryon") + + +async def test_override_can_loosen_a_tier_the_water_level_reached() -> None: + """运营解除加严：水位已到 tier4，手动放宽回 tier2。""" + guard = await guard_spent("100") + await guard.set_override(2) + assert await guard.current_tier() == 2 + with pytest.raises(SceneDegradedError): + await guard.ensure_scene_allowed(scene="hd_export") + await guard.ensure_scene_allowed(scene="tryon") + + +@pytest.mark.parametrize("bad_value", ["0", "5", "abc", ""]) +async def test_invalid_override_values_are_ignored(bad_value: str) -> None: + guard, counter = make_guard() + await guard.add_cost(Decimal("100")) # 水位 tier4 + await counter.set(OVERRIDE_KEY, bad_value, 60) + assert await guard.current_tier() == 4 # 回落水位判定，override 不生效 + + +async def test_set_override_rejects_out_of_range_tiers() -> None: + guard, _ = make_guard() + with pytest.raises(ValueError, match=r"1\.\.4"): + await guard.set_override(5) + with pytest.raises(ValueError, match=r"1\.\.4"): + await guard.set_override(0) + + +async def test_override_ttl_expires_at_the_utc8_day_boundary() -> None: + """override TTL 到当日 UTC+8 24 点：跨预算日必然失效，不需要人工清理。""" + set_expiries: dict[str, int] = {} + counter = InMemoryCostCounter() + real_set = counter.set + + async def recording_set(key: str, value: str, ttl_seconds: int) -> None: + set_expiries[key] = ttl_seconds + await real_set(key, value, ttl_seconds) + + counter.set = recording_set # type: ignore[method-assign] + # 8-14 11:30 UTC+8 → 24:00 还有 12.5 小时 + guard = DailyBudgetGuard( + counter, + daily_budget_cents=Decimal("100"), + clock=lambda: datetime(2026, 8, 14, 3, 30, tzinfo=UTC), + ) + await guard.set_override(4) + assert set_expiries[OVERRIDE_KEY] == 12 * 3600 + 1800 + + +# -------------------------------------------------------------------------- +# 配置解析 +# -------------------------------------------------------------------------- + + +def test_budget_tier_normalizes_bare_star_and_defaults() -> None: + tier = BudgetTier.model_validate({"threshold": 1.0, "name": "block_all", "blocked": "*"}) + assert tier.blocked == [BLOCK_ALL] + assert tier.blocks_all() is True + assert tier.allow == [] # allow 未配置时默认空：block_all 档将拒绝全部非 critical + assert BudgetTier.model_validate({"threshold": 0.6, "name": "warn", "blocked": []}).allow == [] + + +async def test_custom_thresholds_drive_the_tiers() -> None: + tiers: list[dict[str, object]] = [ + {"threshold": 0.5, "name": "warn", "blocked": []}, + {"threshold": 0.75, "name": "degrade_secondary", "blocked": ["hd_export"]}, + {"threshold": 1.0, "name": "block_all", "blocked": BLOCK_ALL, "allow": ["agent_chat"]}, + ] + guard, _ = make_guard(tiers=tiers) + await guard.add_cost(Decimal("75")) + assert await guard.current_tier() == 2 + with pytest.raises(SceneDegradedError): + await guard.ensure_scene_allowed(scene="hd_export")