Vatxzz commited on
Commit
55f5f8f
·
1 Parent(s): f1ba44b

feat: usage table + daily quota accounting

Browse files
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