"""四档分级降级（架构 §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")