feat: reserve chat quota transactionally in SQLite

Author: jiwenxu025-boop

scripts/bazi_engine/api.py

@@ -731,12 +731,17 @@ async def chat_api(request: Request):
         return JSONResponse({"error": "AI 功能未启用"}, status_code=503)
     from .chat import (
         FREE_DAILY_LIMIT,
+        _use_sqlite_runtime_store,
         build_messages,
         call_deepseek_stream,
         check_free_quota,
         consume_code,
         consume_free_quota,
         filter_sensitive,
+        release_quota_reservation,
+        reserve_activation_code,
+        reserve_free_quota,
+        settle_quota_reservation,
         validate_code,
     )
 
@@ -758,8 +763,27 @@ async def reject_gen():
             yield "data: [DONE]\n\n"
         return StreamingResponse(reject_gen(), media_type="text/event-stream")
 
-    # 2. 权限检查
-    if activation_code:
+    # 2. 权限检查。SQLite 模式先事务预占，首个 token 到达后才结算。
+    sqlite_runtime = _use_sqlite_runtime_store()
+    reservation_id = None
+    if sqlite_runtime and activation_code:
+        reservation = reserve_activation_code(activation_code)
+        if reservation is None:
+            _valid, _remaining, msg = validate_code(activation_code)
+            async def invalid_gen():
+                yield f"data: {json.dumps({'token': msg})}\n\n"
+                yield "data: [DONE]\n\n"
+            return StreamingResponse(invalid_gen(), media_type="text/event-stream")
+        reservation_id = reservation.reservation_id
+    elif sqlite_runtime:
+        reservation = reserve_free_quota(client_ip)
+        if reservation is None:
+            async def quota_gen():
+                yield f"data: {json.dumps({'token': f'今日免费追问次数（{FREE_DAILY_LIMIT}次）已用完。点击"解锁追问"获取激活码。'})}\n\n"
+                yield "data: [DONE]\n\n"
+            return StreamingResponse(quota_gen(), media_type="text/event-stream")
+        reservation_id = reservation.reservation_id
+    elif activation_code:
         valid, _remaining, msg = validate_code(activation_code)
         if not valid:
             async def invalid_gen():
@@ -780,14 +804,20 @@ async def quota_gen():
     # 4. 流式响应（扣费在首个token到达后执行，避免错误响应也扣费）
     async def stream_chat():
         consumed = False
-        async for chunk in call_deepseek_stream(messages):
-            if not consumed and chunk.startswith('data: {"token"'):
-                consumed = True
-                if activation_code:
-                    consume_code(activation_code)
-                else:
-                    consume_free_quota(client_ip)
-            yield chunk
+        try:
+            async for chunk in call_deepseek_stream(messages):
+                if not consumed and chunk.startswith('data: {"token"'):
+                    consumed = True
+                    if reservation_id:
+                        settle_quota_reservation(reservation_id)
+                    elif activation_code:
+                        consume_code(activation_code)
+                    else:
+                        consume_free_quota(client_ip)
+                yield chunk
+        finally:
+            if reservation_id and not consumed:
+                release_quota_reservation(reservation_id)
 
     return _limited_stream_response(stream_chat())
 

scripts/bazi_engine/chat.py

@@ -14,12 +14,16 @@
 
 from ._deepseek_config import DEEPSEEK_API_URL, DEEPSEEK_KEY, DEEPSEEK_MODEL
 from ._http import shared_async_client
+from .runtime_store import QuotaReservation, RuntimeStore
 
 # ═══════════════════════════════════════════════════════════════
 # 配置
 # ═══════════════════════════════════════════════════════════════
 
 _ACTIVATION_FILE = Path(__file__).resolve().parent / "activation_codes.json"
+_RUNTIME_DB = Path(
+    os.getenv("BAZI_RUNTIME_DB_PATH", Path(__file__).resolve().parents[2] / "data" / "runtime.sqlite3")
+)
 _RUNTIME_DATA_LOCK = threading.RLock()
 
 # ═══════════════════════════════════════════════════════════════
@@ -383,8 +387,33 @@ async def call_deepseek_stream(messages: list[dict]) -> AsyncGenerator[str]:
 _default_codes = dict(_DEMO_CODES) if _USE_DEMO_CODES else {}
 
 
+def _use_sqlite_runtime_store() -> bool:
+    return os.getenv("BAZI_RUNTIME_STORE", "json").lower() == "sqlite"
+
+
+def _activation_seed_codes() -> dict:
+    codes = dict(_default_codes)
+    env_codes = os.getenv("ACTIVATION_CODES", "")
+    if env_codes:
+        with suppress(json.JSONDecodeError):
+            codes.update(json.loads(env_codes))
+    if os.getenv("BAZI_PUBLIC", "").lower() in ("1", "true", "yes"):
+        for code in _DEMO_CODES:
+            codes.pop(code, None)
+    return codes
+
+
+def _runtime_store() -> RuntimeStore:
+    store = RuntimeStore(_RUNTIME_DB)
+    store.seed_activation_codes(_activation_seed_codes())
+    return store
+
+
 def _load_codes() -> dict:
     """加载激活码（环境变量 + 本地文件合并，环境变量优先）"""
+    if _use_sqlite_runtime_store():
+        return _runtime_store().activation_codes()
+
     codes = dict(_default_codes)
 
     # 本地文件（首次自动创建，gitignore 保护）
@@ -442,6 +471,14 @@ def _atomic_json_write(path: Path, data: dict, **json_kwargs) -> None:
 
 def validate_code(code: str) -> tuple[bool, int, str]:
     """验证激活码。返回 (有效?, 剩余次数, 消息)"""
+    if _use_sqlite_runtime_store():
+        remaining = _runtime_store().activation_remaining(code)
+        if remaining is None:
+            return False, 0, "激活码无效"
+        if remaining <= 0:
+            return False, 0, "该激活码次数已用完"
+        return True, remaining, "有效"
+
     codes = _load_codes()
     entry = codes.get(code.strip().upper())
     if not entry:
@@ -454,6 +491,13 @@ def validate_code(code: str) -> tuple[bool, int, str]:
 
 def consume_code(code: str) -> tuple[bool, int]:
     """消耗一次激活码。返回 (成功?, 剩余次数)"""
+    if _use_sqlite_runtime_store():
+        reservation = reserve_activation_code(code)
+        if reservation is None:
+            return False, 0
+        _runtime_store().settle_reservation(reservation.reservation_id)
+        return True, reservation.remaining
+
     with _RUNTIME_DATA_LOCK:
         codes = _load_codes()
         entry = codes.get(code.strip().upper())
@@ -498,6 +542,10 @@ def check_free_quota(client_id: str) -> tuple[bool, int]:
     IP 经哈希处理后存储，不保留原始 IP。
     注意: NAT/代理环境下多用户共用同一 IP，一个用户用完额度其他人也会被拦。
     """
+    if _use_sqlite_runtime_store():
+        remaining = _runtime_store().free_remaining(_hash_ip(client_id), time.strftime("%Y-%m-%d"), FREE_DAILY_LIMIT)
+        return remaining > 0, remaining
+
     with _RUNTIME_DATA_LOCK:
         today = time.strftime("%Y-%m-%d")
         data = _load_free_usage()
@@ -521,6 +569,13 @@ def check_free_quota(client_id: str) -> tuple[bool, int]:
 
 def consume_free_quota(client_id: str) -> int:
     """消耗一次免费额度。返回剩余次数"""
+    if _use_sqlite_runtime_store():
+        reservation = reserve_free_quota(client_id)
+        if reservation is None:
+            return 0
+        _runtime_store().settle_reservation(reservation.reservation_id)
+        return reservation.remaining
+
     with _RUNTIME_DATA_LOCK:
         today = time.strftime("%Y-%m-%d")
         data = _load_free_usage()
@@ -536,3 +591,27 @@ def consume_free_quota(client_id: str) -> int:
         data[key] = entry
         _save_free_usage(data)
         return FREE_DAILY_LIMIT - entry["count"]
+
+
+def reserve_activation_code(code: str) -> QuotaReservation | None:
+    """Atomically reserve a paid quota until the first response token is delivered."""
+    if not _use_sqlite_runtime_store():
+        return None
+    return _runtime_store().reserve_activation(code)
+
+
+def reserve_free_quota(client_id: str) -> QuotaReservation | None:
+    """Atomically reserve one daily free quota until the first response token is delivered."""
+    if not _use_sqlite_runtime_store():
+        return None
+    return _runtime_store().reserve_free(
+        _hash_ip(client_id), time.strftime("%Y-%m-%d"), FREE_DAILY_LIMIT,
+    )
+
+
+def settle_quota_reservation(reservation_id: str) -> bool:
+    return _runtime_store().settle_reservation(reservation_id)
+
+
+def release_quota_reservation(reservation_id: str) -> bool:
+    return _runtime_store().release_reservation(reservation_id)

scripts/bazi_engine/runtime_store.py

@@ -158,6 +158,17 @@ def activation_remaining(self, code: str) -> int | None:
             ).fetchone()
         return int(row["remaining"]) if row else None
 
+    def activation_codes(self) -> dict[str, dict]:
+        self.initialize()
+        with self.connect() as connection:
+            rows = connection.execute(
+                "SELECT code, remaining, note FROM activation_codes ORDER BY code"
+            ).fetchall()
+        return {
+            row["code"]: {"剩余": int(row["remaining"]), "备注": row["note"]}
+            for row in rows
+        }
+
     def free_remaining(self, client_hash: str, usage_date: str, limit: int) -> int:
         self.initialize()
         with self.connect() as connection:

scripts/tests/test_api.py

@@ -648,3 +648,62 @@ async def fake_stream(messages):
     assert "2026年 19岁 丙午流年，丙午大运" in system_prompt
     assert "2026年 19岁 丙午流年，甲辰大运" not in system_prompt
     assert system_prompt.rfind("【最终事实约束】") > system_prompt.rfind("classical-texts.md")
+
+
+def test_sqlite_chat_quota_releases_when_provider_returns_no_token(monkeypatch, tmp_path):
+    import bazi_engine.api as api_module
+    import bazi_engine.chat as chat_module
+    from bazi_engine.api import app
+    from bazi_engine.runtime_store import RuntimeStore
+
+    async def failed_stream(_messages):
+        yield "data: [ERROR] provider unavailable\n\n"
+
+    database_path = tmp_path / "runtime.sqlite3"
+    monkeypatch.setenv("BAZI_RUNTIME_STORE", "sqlite")
+    monkeypatch.setenv("ACTIVATION_CODES", '{"TEST": {"剩余": 1, "备注": "test"}}')
+    monkeypatch.setattr(api_module, "_AI_ENABLED", True)
+    monkeypatch.setattr(chat_module, "_RUNTIME_DB", database_path)
+    monkeypatch.setattr(chat_module, "build_messages", lambda *_args: [{"role": "user", "content": "test"}])
+    monkeypatch.setattr(chat_module, "call_deepseek_stream", failed_stream)
+
+    with TestClient(app).stream(
+        "POST",
+        "/api/chat",
+        json={"question": "测试", "chart_data": {}, "activation_code": "TEST"},
+    ) as response:
+        body = "".join(response.iter_text())
+
+    assert response.status_code == 200
+    assert "provider unavailable" in body
+    assert RuntimeStore(database_path).activation_remaining("TEST") == 1
+
+
+def test_sqlite_chat_quota_settles_after_first_token(monkeypatch, tmp_path):
+    import bazi_engine.api as api_module
+    import bazi_engine.chat as chat_module
+    from bazi_engine.api import app
+    from bazi_engine.runtime_store import RuntimeStore
+
+    async def successful_stream(_messages):
+        yield 'data: {"token":"ok"}\n\n'
+        yield "data: [DONE]\n\n"
+
+    database_path = tmp_path / "runtime.sqlite3"
+    monkeypatch.setenv("BAZI_RUNTIME_STORE", "sqlite")
+    monkeypatch.setenv("ACTIVATION_CODES", '{"TEST": {"剩余": 1, "备注": "test"}}')
+    monkeypatch.setattr(api_module, "_AI_ENABLED", True)
+    monkeypatch.setattr(chat_module, "_RUNTIME_DB", database_path)
+    monkeypatch.setattr(chat_module, "build_messages", lambda *_args: [{"role": "user", "content": "test"}])
+    monkeypatch.setattr(chat_module, "call_deepseek_stream", successful_stream)
+
+    with TestClient(app).stream(
+        "POST",
+        "/api/chat",
+        json={"question": "测试", "chart_data": {}, "activation_code": "TEST"},
+    ) as response:
+        body = "".join(response.iter_text())
+
+    assert response.status_code == 200
+    assert '"token":"ok"' in body or '"token": "ok"' in body
+    assert RuntimeStore(database_path).activation_remaining("TEST") == 0