feat: usage table + daily quota accounting
Browse files- backend/app/config.py +6 -0
- backend/app/quota.py +61 -0
- backend/app/store/db.py +42 -0
- backend/tests/test_quota.py +21 -0
backend/app/config.py
CHANGED
|
@@ -74,6 +74,12 @@ class Settings(BaseSettings):
|
|
| 74 |
reddit_client_id: str = ""
|
| 75 |
reddit_client_secret: str = ""
|
| 76 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 77 |
# worker / queue
|
| 78 |
max_attempts: int = 3
|
| 79 |
worker_poll_seconds: float = 1.0
|
|
|
|
| 74 |
reddit_client_id: str = ""
|
| 75 |
reddit_client_secret: str = ""
|
| 76 |
|
| 77 |
+
# quotas (per UTC day)
|
| 78 |
+
quota_cards_per_day: int = 10
|
| 79 |
+
quota_chat_per_day: int = 30
|
| 80 |
+
quota_connections_refresh_per_day: int = 3
|
| 81 |
+
quota_ip_cards_per_day: int = 30
|
| 82 |
+
|
| 83 |
# worker / queue
|
| 84 |
max_attempts: int = 3
|
| 85 |
worker_poll_seconds: float = 1.0
|
backend/app/quota.py
ADDED
|
@@ -0,0 +1,61 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Per-user daily quotas. Chat-style routes raise 429; card creation degrades
|
| 2 |
+
instead (see cards.py) so a save never hard-fails."""
|
| 3 |
+
|
| 4 |
+
from __future__ import annotations
|
| 5 |
+
|
| 6 |
+
from datetime import datetime, time, timedelta, timezone
|
| 7 |
+
|
| 8 |
+
from fastapi import Depends, HTTPException, Request
|
| 9 |
+
|
| 10 |
+
from app.auth import get_owner
|
| 11 |
+
from app.config import get_settings
|
| 12 |
+
from app.store import db
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
def _resets_at() -> str:
|
| 16 |
+
"""ISO timestamp of the next UTC midnight (quota reset)."""
|
| 17 |
+
now = datetime.now(timezone.utc)
|
| 18 |
+
tomorrow = datetime.combine(now.date() + timedelta(days=1), time.min, tzinfo=timezone.utc)
|
| 19 |
+
return tomorrow.isoformat()
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
def spend(kind: str, limit_attr: str):
|
| 23 |
+
"""Dependency factory: spend one unit of `kind` or raise 429."""
|
| 24 |
+
|
| 25 |
+
async def _dep(owner_id: str = Depends(get_owner)) -> str:
|
| 26 |
+
limit = getattr(get_settings(), limit_attr)
|
| 27 |
+
async with db.session() as s:
|
| 28 |
+
allowed, used = await db.spend_usage(
|
| 29 |
+
s, owner_id=owner_id, kind=kind, limit=limit
|
| 30 |
+
)
|
| 31 |
+
if not allowed:
|
| 32 |
+
raise HTTPException(
|
| 33 |
+
status_code=429,
|
| 34 |
+
detail={
|
| 35 |
+
"error": "quota", "kind": kind,
|
| 36 |
+
"used": used, "limit": limit, "resets_at": _resets_at(),
|
| 37 |
+
},
|
| 38 |
+
)
|
| 39 |
+
return owner_id
|
| 40 |
+
|
| 41 |
+
return _dep
|
| 42 |
+
|
| 43 |
+
|
| 44 |
+
async def card_budget(owner_id: str, request: Request) -> bool:
|
| 45 |
+
"""Card-creation budget. Enforces the per-IP cap (429) and returns whether
|
| 46 |
+
the owner still has AI budget today (False -> degrade, never fail)."""
|
| 47 |
+
settings = get_settings()
|
| 48 |
+
ip = request.client.host if request.client else "unknown"
|
| 49 |
+
async with db.session() as s:
|
| 50 |
+
ip_ok, _ = await db.spend_usage(
|
| 51 |
+
s, owner_id=f"ip:{ip}", kind="cards", limit=settings.quota_ip_cards_per_day
|
| 52 |
+
)
|
| 53 |
+
if not ip_ok:
|
| 54 |
+
raise HTTPException(status_code=429, detail={
|
| 55 |
+
"error": "quota", "kind": "ip", "limit": settings.quota_ip_cards_per_day,
|
| 56 |
+
"used": settings.quota_ip_cards_per_day, "resets_at": _resets_at(),
|
| 57 |
+
})
|
| 58 |
+
allowed, _ = await db.spend_usage(
|
| 59 |
+
s, owner_id=owner_id, kind="cards", limit=settings.quota_cards_per_day
|
| 60 |
+
)
|
| 61 |
+
return allowed
|
backend/app/store/db.py
CHANGED
|
@@ -275,6 +275,8 @@ class JobRow(Base):
|
|
| 275 |
state: Mapped[str] = mapped_column(String, default=JobState.QUEUED.value, index=True)
|
| 276 |
attempts: Mapped[int] = mapped_column(Integer, default=0)
|
| 277 |
last_error: Mapped[str | None] = mapped_column(Text, nullable=True)
|
|
|
|
|
|
|
| 278 |
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=_utcnow)
|
| 279 |
started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
| 280 |
finished_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
|
@@ -296,6 +298,39 @@ class ConnectionRow(Base):
|
|
| 296 |
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=_utcnow)
|
| 297 |
|
| 298 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 299 |
# --------------------------------------------------------------------------- #
|
| 300 |
# Engine / session lifecycle
|
| 301 |
# --------------------------------------------------------------------------- #
|
|
@@ -383,6 +418,13 @@ async def init_db() -> None:
|
|
| 383 |
"owner_id": "TEXT",
|
| 384 |
},
|
| 385 |
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 386 |
await _add_missing_columns(
|
| 387 |
conn,
|
| 388 |
"concepts",
|
|
|
|
| 275 |
state: Mapped[str] = mapped_column(String, default=JobState.QUEUED.value, index=True)
|
| 276 |
attempts: Mapped[int] = mapped_column(Integer, default=0)
|
| 277 |
last_error: Mapped[str | None] = mapped_column(Text, nullable=True)
|
| 278 |
+
# Past-quota card creation: worker skips AI structuring, paragraph fallback.
|
| 279 |
+
degraded: Mapped[bool] = mapped_column(Boolean, default=False)
|
| 280 |
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=_utcnow)
|
| 281 |
started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
| 282 |
finished_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
|
|
|
| 298 |
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=_utcnow)
|
| 299 |
|
| 300 |
|
| 301 |
+
class UsageRow(Base):
|
| 302 |
+
"""Daily metered usage. One row per (owner, UTC day, kind); owner_id also
|
| 303 |
+
stores "ip:<addr>" rows for the anonymous-farming IP cap."""
|
| 304 |
+
|
| 305 |
+
__tablename__ = "usage"
|
| 306 |
+
|
| 307 |
+
owner_id: Mapped[str] = mapped_column(String, primary_key=True)
|
| 308 |
+
day: Mapped[str] = mapped_column(String, primary_key=True) # "YYYY-MM-DD" UTC
|
| 309 |
+
kind: Mapped[str] = mapped_column(String, primary_key=True)
|
| 310 |
+
count: Mapped[int] = mapped_column(Integer, default=0)
|
| 311 |
+
|
| 312 |
+
|
| 313 |
+
def _today() -> str:
|
| 314 |
+
"""Current UTC day key."""
|
| 315 |
+
return datetime.now(timezone.utc).strftime("%Y-%m-%d")
|
| 316 |
+
|
| 317 |
+
|
| 318 |
+
async def spend_usage(
|
| 319 |
+
db_session: AsyncSession, *, owner_id: str, kind: str, limit: int
|
| 320 |
+
) -> tuple[bool, int]:
|
| 321 |
+
"""Increment today's counter unless at limit. Returns (allowed, used_after)."""
|
| 322 |
+
day = _today()
|
| 323 |
+
row = await db_session.get(UsageRow, (owner_id, day, kind))
|
| 324 |
+
if row is None:
|
| 325 |
+
row = UsageRow(owner_id=owner_id, day=day, kind=kind, count=0)
|
| 326 |
+
db_session.add(row)
|
| 327 |
+
if row.count >= limit:
|
| 328 |
+
return False, row.count
|
| 329 |
+
row.count += 1
|
| 330 |
+
await db_session.commit()
|
| 331 |
+
return True, row.count
|
| 332 |
+
|
| 333 |
+
|
| 334 |
# --------------------------------------------------------------------------- #
|
| 335 |
# Engine / session lifecycle
|
| 336 |
# --------------------------------------------------------------------------- #
|
|
|
|
| 418 |
"owner_id": "TEXT",
|
| 419 |
},
|
| 420 |
)
|
| 421 |
+
await _add_missing_columns(
|
| 422 |
+
conn,
|
| 423 |
+
"jobs",
|
| 424 |
+
{
|
| 425 |
+
"degraded": "BOOLEAN DEFAULT 0", # quota-degraded card creation
|
| 426 |
+
},
|
| 427 |
+
)
|
| 428 |
await _add_missing_columns(
|
| 429 |
conn,
|
| 430 |
"concepts",
|
backend/tests/test_quota.py
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Quota accounting: daily counters, limits, per-owner/per-kind isolation."""
|
| 2 |
+
|
| 3 |
+
from app.store import db
|
| 4 |
+
|
| 5 |
+
|
| 6 |
+
async def test_spend_usage_counts_and_caps(database) -> None:
|
| 7 |
+
async with db.session() as s:
|
| 8 |
+
for i in range(3):
|
| 9 |
+
allowed, used = await db.spend_usage(s, owner_id="u1", kind="chat", limit=3)
|
| 10 |
+
assert allowed and used == i + 1
|
| 11 |
+
allowed, used = await db.spend_usage(s, owner_id="u1", kind="chat", limit=3)
|
| 12 |
+
assert not allowed and used == 3
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
async def test_usage_is_per_owner_and_per_kind(database) -> None:
|
| 16 |
+
async with db.session() as s:
|
| 17 |
+
await db.spend_usage(s, owner_id="u1", kind="chat", limit=3)
|
| 18 |
+
allowed, used = await db.spend_usage(s, owner_id="u2", kind="chat", limit=3)
|
| 19 |
+
assert allowed and used == 1
|
| 20 |
+
allowed, used = await db.spend_usage(s, owner_id="u1", kind="cards", limit=3)
|
| 21 |
+
assert allowed and used == 1
|