wooo/awoooi
1
1
Fork 0
Code Issues 2 Pull Requests 1 Actions 2 Packages Projects Releases Wiki Activity

feat(ai): add guarded Claude fallback route

Browse Source
This commit is contained in:
ogt
2026-07-15 15:57:24 +08:00
parent f57845ba5e
commit 5c5eb48091
56 changed files with 4942 additions and 760 deletions
Download Patch File Download Diff File Expand all files Collapse all files

330
apps/api/tests/integration/test_gemini_cost_guard_real_redis.py
View File

@@ -21,16 +21,54 @@ import redis.asyncio as redis_async
from src.core import config as config_module
from src.services.ai_rate_limiter import (
DAILY_COST_MICRO_USD_KEY,
DAILY_COST_RESERVED_MICRO_USD_KEY,
DAILY_REQ_KEY,
DAILY_TOKEN_KEY,
DAILY_TOKEN_RESERVED_KEY,
GENERATION_RECEIPT_KEY,
GENERATION_RUN_IDEMPOTENCY_KEY,
REDIS_KEY_PREFIX,
TOTAL_COST_KEY,
TOTAL_COST_RESERVED_MICRO_USD_KEY,
AIRateLimiter,
)
def _configure_paid_provider_limits(
monkeypatch: pytest.MonkeyPatch,
provider: str,
*,
rpm: int = 20,
daily_cost_usd: float = 5.0,
total_cost_usd: float = 5.0,
alert_threshold_usd: float = 4.0,
) -> None:
prefix = provider.upper()
monkeypatch.setattr(config_module.settings, f"{prefix}_RPM_LIMIT", rpm)
monkeypatch.setattr(config_module.settings, f"{prefix}_DAILY_QUOTA", 100)
monkeypatch.setattr(
config_module.settings,
f"{prefix}_DAILY_TOKEN_LIMIT",
100_000,
)
monkeypatch.setattr(
config_module.settings,
f"{prefix}_DAILY_COST_LIMIT_USD",
daily_cost_usd,
)
monkeypatch.setattr(
config_module.settings,
f"{prefix}_TOTAL_COST_LIMIT_USD",
total_cost_usd,
)
monkeypatch.setattr(
config_module.settings,
f"{prefix}_COST_ALERT_THRESHOLD_USD",
alert_threshold_usd,
)
def _free_local_port() -> int:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
sock.bind(("127.0.0.1", 0))
@@ -179,20 +217,24 @@ def ephemeral_redis_url(tmp_path_factory: pytest.TempPathFactory) -> Iterator[st
@pytest.mark.integration
@pytest.mark.asyncio
@pytest.mark.parametrize(
("provider", "model"),
[
("gemini", "gemini-2.5-flash-lite"),
("claude", "claude-sonnet-5"),
],
)
async def test_lua_reservation_is_atomic_under_real_redis_concurrency(
ephemeral_redis_url: str,
monkeypatch: pytest.MonkeyPatch,
provider: str,
model: str,
) -> None:
client = redis_async.from_url(ephemeral_redis_url, decode_responses=False)
await client.ping()
await client.flushdb()
monkeypatch.setattr(config_module.settings, "GEMINI_RPM_LIMIT", 5)
monkeypatch.setattr(config_module.settings, "GEMINI_DAILY_QUOTA", 100)
monkeypatch.setattr(config_module.settings, "GEMINI_DAILY_TOKEN_LIMIT", 100_000)
monkeypatch.setattr(config_module.settings, "GEMINI_DAILY_COST_LIMIT_USD", 5.0)
monkeypatch.setattr(config_module.settings, "GEMINI_TOTAL_COST_LIMIT_USD", 5.0)
monkeypatch.setattr(config_module.settings, "GEMINI_COST_ALERT_THRESHOLD_USD", 4.0)
_configure_paid_provider_limits(monkeypatch, provider, rpm=5)
limiter = AIRateLimiter()
limiter._redis = client
@@ -202,9 +244,9 @@ async def test_lua_reservation_is_atomic_under_real_redis_concurrency(
reservations = await asyncio.gather(
*(
limiter.reserve_generation(
"gemini",
provider,
f"isolated concurrency verifier {index}",
model="gemini-2.5-flash-lite",
model=model,
max_output_tokens=1,
context={
"trace_id": f"trace-{index}",
@@ -225,7 +267,7 @@ async def test_lua_reservation_is_atomic_under_real_redis_concurrency(
for reservation in reservations:
receipt = await client.hgetall(
GENERATION_RECEIPT_KEY.format(
provider="gemini",
provider=provider,
receipt_id=reservation.receipt_id,
)
)
@@ -238,7 +280,7 @@ async def test_lua_reservation_is_atomic_under_real_redis_concurrency(
assert receipt[b"reserved_cost_micro_usd"] == b"0"
request_count = await client.get(
DAILY_REQ_KEY.format(provider="gemini", date="2026-07-14")
DAILY_REQ_KEY.format(provider=provider, date="2026-07-14")
)
assert int(request_count or 0) == 5
@@ -256,7 +298,7 @@ async def test_lua_reservation_is_atomic_under_real_redis_concurrency(
assert int(
await client.get(
DAILY_TOKEN_RESERVED_KEY.format(
provider="gemini",
provider=provider,
date="2026-07-14",
)
)
@@ -265,7 +307,7 @@ async def test_lua_reservation_is_atomic_under_real_redis_concurrency(
assert int(
await client.get(
DAILY_COST_RESERVED_MICRO_USD_KEY.format(
provider="gemini",
provider=provider,
date="2026-07-14",
)
)
@@ -273,7 +315,7 @@ async def test_lua_reservation_is_atomic_under_real_redis_concurrency(
) == 0
assert int(
await client.get(
TOTAL_COST_RESERVED_MICRO_USD_KEY.format(provider="gemini")
TOTAL_COST_RESERVED_MICRO_USD_KEY.format(provider=provider)
)
or 0
) == 0
@@ -284,21 +326,29 @@ async def test_lua_reservation_is_atomic_under_real_redis_concurrency(
@pytest.mark.integration
@pytest.mark.asyncio
@pytest.mark.parametrize(
("provider", "model"),
[
("gemini", "gemini-2.5-flash-lite"),
("claude", "claude-sonnet-5"),
],
)
async def test_same_run_is_idempotent_and_failure_sets_work_item_cooldown(
ephemeral_redis_url: str,
monkeypatch: pytest.MonkeyPatch,
provider: str,
model: str,
) -> None:
client = redis_async.from_url(ephemeral_redis_url, decode_responses=False)
await client.ping()
await client.flushdb()
monkeypatch.setattr(config_module.settings, "GEMINI_RPM_LIMIT", 20)
monkeypatch.setattr(config_module.settings, "GEMINI_DAILY_QUOTA", 100)
monkeypatch.setattr(config_module.settings, "GEMINI_DAILY_TOKEN_LIMIT", 100_000)
monkeypatch.setattr(config_module.settings, "GEMINI_DAILY_COST_LIMIT_USD", 5.0)
monkeypatch.setattr(config_module.settings, "GEMINI_TOTAL_COST_LIMIT_USD", 5.0)
monkeypatch.setattr(config_module.settings, "GEMINI_COST_ALERT_THRESHOLD_USD", 4.0)
monkeypatch.setattr(config_module.settings, "GEMINI_FAILURE_COOLDOWN_SECONDS", 60)
_configure_paid_provider_limits(monkeypatch, provider)
monkeypatch.setattr(
config_module.settings,
f"{provider.upper()}_FAILURE_COOLDOWN_SECONDS",
60,
)
limiter = AIRateLimiter()
limiter._redis = client
@@ -313,8 +363,9 @@ async def test_same_run_is_idempotent_and_failure_sets_work_item_cooldown(
reservations = await asyncio.gather(
*(
limiter.reserve_generation(
"gemini",
provider,
"same-run concurrency verifier",
model=model,
max_output_tokens=1,
context={**identity, "trace_id": f"trace-same-run-{index}"},
)
@@ -329,7 +380,7 @@ async def test_same_run_is_idempotent_and_failure_sets_work_item_cooldown(
}
assert len({reservation.receipt_id for reservation in reservations}) == 1
assert int(
await client.get(DAILY_REQ_KEY.format(provider="gemini", date="2026-07-14"))
await client.get(DAILY_REQ_KEY.format(provider=provider, date="2026-07-14"))
or 0
) == 1
@@ -339,8 +390,9 @@ async def test_same_run_is_idempotent_and_failure_sets_work_item_cooldown(
error_code="isolated_failure",
)
cooldown_block = await limiter.reserve_generation(
"gemini",
provider,
"retry after failure",
model=model,
max_output_tokens=1,
context={
"trace_id": "trace-retry",
@@ -355,6 +407,238 @@ async def test_same_run_is_idempotent_and_failure_sets_work_item_cooldown(
await client.aclose()
@pytest.mark.integration
@pytest.mark.asyncio
async def test_claude_daily_and_total_cost_caps_include_pending_reservations(
ephemeral_redis_url: str,
monkeypatch: pytest.MonkeyPatch,
) -> None:
client = redis_async.from_url(ephemeral_redis_url, decode_responses=False)
await client.ping()
await client.flushdb()
# One-byte prompt + one output token at the conservative Sonnet 5 ceiling
# is exactly 18 micro-USD: (1 * $3 + 1 * $15) / 1M.
boundary_usd = 18 / 1_000_000
_configure_paid_provider_limits(
monkeypatch,
"claude",
daily_cost_usd=boundary_usd,
total_cost_usd=5.0,
alert_threshold_usd=4.0,
)
limiter = AIRateLimiter()
limiter._redis = client
limiter._get_today = lambda: "2026-07-14" # type: ignore[method-assign]
limiter._seconds_until_tomorrow = lambda: 3600 # type: ignore[method-assign]
first = await limiter.reserve_generation(
"claude",
"x",
model="claude-sonnet-5",
max_output_tokens=1,
context={
"trace_id": "trace-claude-daily-1",
"run_id": "run-claude-daily-1",
"work_item_id": "work-claude-cost-boundary",
},
)
daily_block = await limiter.reserve_generation(
"claude",
"x",
model="claude-sonnet-5",
max_output_tokens=1,
context={
"trace_id": "trace-claude-daily-2",
"run_id": "run-claude-daily-2",
"work_item_id": "work-claude-cost-boundary-2",
},
)
assert first.allowed is True
assert first.estimated_cost_usd == boundary_usd
assert daily_block.allowed is False
assert daily_block.reason == "daily_cost_limit"
assert int(
await client.get(
DAILY_COST_RESERVED_MICRO_USD_KEY.format(
provider="claude",
date="2026-07-14",
)
)
or 0
) == 18
assert await limiter.finalize_generation(
first,
success=True,
prompt_tokens=1,
completion_tokens=1,
)
assert int(
await client.get(
DAILY_COST_MICRO_USD_KEY.format(
provider="claude",
date="2026-07-14",
)
)
or 0
) == 18
await client.flushdb()
_configure_paid_provider_limits(
monkeypatch,
"claude",
daily_cost_usd=5.0,
total_cost_usd=1.0,
alert_threshold_usd=1.0,
)
await client.set(TOTAL_COST_KEY.format(provider="claude"), "0.999982")
# Keep this verifier hermetic: the warning path remains acknowledged in the
# isolated Redis DB, so the boundary test cannot contact Telegram.
await client.set(f"{REDIS_KEY_PREFIX}cost_warning_sent:claude", "1")
total_first = await limiter.reserve_generation(
"claude",
"x",
model="claude-sonnet-5",
max_output_tokens=1,
context={
"trace_id": "trace-claude-total-1",
"run_id": "run-claude-total-1",
"work_item_id": "work-claude-total-boundary",
},
)
total_block = await limiter.reserve_generation(
"claude",
"x",
model="claude-sonnet-5",
max_output_tokens=1,
context={
"trace_id": "trace-claude-total-2",
"run_id": "run-claude-total-2",
"work_item_id": "work-claude-total-boundary-2",
},
)
assert total_first.allowed is True
assert total_block.allowed is False
assert total_block.reason == "total_cost_limit"
assert int(
await client.get(
TOTAL_COST_RESERVED_MICRO_USD_KEY.format(provider="claude")
)
or 0
) == 18
assert await limiter.finalize_generation(
total_first,
success=True,
prompt_tokens=1,
completion_tokens=1,
)
assert float(await client.get(TOTAL_COST_KEY.format(provider="claude")) or 0) == 1.0
await client.flushdb()
await client.aclose()
@pytest.mark.integration
@pytest.mark.asyncio
async def test_claude_finalize_is_atomic_and_idempotent_under_real_redis_race(
ephemeral_redis_url: str,
monkeypatch: pytest.MonkeyPatch,
) -> None:
client = redis_async.from_url(ephemeral_redis_url, decode_responses=False)
await client.ping()
await client.flushdb()
_configure_paid_provider_limits(monkeypatch, "claude")
limiter = AIRateLimiter()
limiter._redis = client
limiter._get_today = lambda: "2026-07-14" # type: ignore[method-assign]
limiter._seconds_until_tomorrow = lambda: 3600 # type: ignore[method-assign]
reservation = await limiter.reserve_generation(
"claude",
"finalize race",
model="claude-sonnet-5",
max_output_tokens=16,
context={
"trace_id": "trace-claude-finalize-race",
"run_id": "run-claude-finalize-race",
"work_item_id": "work-claude-finalize-race",
},
)
assert reservation.allowed is True
finalized = await asyncio.gather(
*(
limiter.finalize_generation(
reservation,
success=True,
prompt_tokens=7,
completion_tokens=3,
)
for _ in range(20)
)
)
assert finalized == [True] * 20
receipt_key = GENERATION_RECEIPT_KEY.format(
provider="claude",
receipt_id=reservation.receipt_id,
)
finalized_receipt = await client.hgetall(receipt_key)
assert finalized_receipt[b"status"] == b"succeeded"
assert finalized_receipt[b"actual_prompt_tokens"] == b"7"
assert finalized_receipt[b"actual_completion_tokens"] == b"3"
assert int(
await client.get(
DAILY_TOKEN_KEY.format(provider="claude", date="2026-07-14")
)
or 0
) == 10
assert int(
await client.get(
DAILY_COST_MICRO_USD_KEY.format(
provider="claude",
date="2026-07-14",
)
)
or 0
) == 66
assert int(
await client.get(
TOTAL_COST_RESERVED_MICRO_USD_KEY.format(provider="claude")
)
or 0
) == 0
adversarial = await asyncio.gather(
*(
limiter.finalize_generation(
reservation,
success=False,
prompt_tokens=999,
completion_tokens=999,
error_code="adversarial_replay",
)
for _ in range(20)
)
)
assert adversarial == [True] * 20
assert await client.hgetall(receipt_key) == finalized_receipt
assert int(
await client.get(
DAILY_TOKEN_KEY.format(provider="claude", date="2026-07-14")
)
or 0
) == 10
await client.flushdb()
await client.aclose()
@pytest.mark.integration
@pytest.mark.asyncio
async def test_finalized_receipt_and_run_claim_survive_configured_ttl_and_replay(