admin-git/DealDocumentScreening
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix issues hardening by plan. M1: .scratch/pipeline-hardening/issues/001

Browse source 
- 011
This commit is contained in:
febux 2026-09-05 23:18:19 +03:00
parent 17054c1e99
commit e5fd285328
37 changed files with 1176 additions and 112 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
.env.example
View file

@ -122,6 +122,8 @@ API_HOST=0.0.0.0
API_PORT=8000
API_METRICS_PORT=9100
B2B_DEFAULT_RATE_LIMIT_RPS=3 # per API key; mirrors Ollama Pro concurrency
# Bearer token protecting /metrics. Leave empty to keep the endpoint open (default).
METRICS_BEARER_TOKEN=
CORS_ORIGINS= # comma-separated, future web SPA
# SSE streaming of analysis status (GET /api/v1/reports/{id}/events):
@ -150,6 +152,10 @@ PASSWORD_RESET_TTL_MINUTES=60
PASSWORD_MIN_LENGTH=8
WEB_APP_BASE_URL=http://localhost:5173 # SPA base — used to build reset + magic links
# --- Auth rate limiting ---
AUTH_RATE_LIMIT_IP_RPS=5 # per-IP bucket for public auth endpoints
AUTH_RATE_LIMIT_EMAIL_RPS=2 # per-recipient bucket for mail-sending auth endpoints
# --- Passkeys (WebAuthn) ---
PASSKEY_ENABLED=true # toggle /api/v1/auth/passkeys/* routes
PASSKEY_RP_ID=localhost # effective domain of the webUI (must match the browser URL host)

3
docs/ARCHITECTURE.md
View file

@ -1142,6 +1142,7 @@ None of 1–5 requires touching `core/` application code — only compose/infra.
| `API_PORT` | `8000` | |
| `API_METRICS_PORT` | `9100` | |
| `B2B_DEFAULT_RATE_LIMIT_RPS` | `3` | per API key; mirrors Ollama Pro concurrency, overridable per `api_keys.rate_limit_rps` |
| `METRICS_BEARER_TOKEN` | (empty) | Bearer token protecting `/metrics`; empty leaves the endpoint open |
| `CORS_ORIGINS` | (empty, future web) | |
### Auth (JWT + Telegram identity verification + webUI + admin panel)
@ -1157,6 +1158,8 @@ None of 1–5 requires touching `core/` application code — only compose/infra.
| `WEB_APP_BASE_URL` | `http://localhost:5173` | base URL of the future web SPA; used for password-reset links |
| `PASSWORD_RESET_TTL_MINUTES` | `60` | reset-link validity |
| `PASSWORD_MIN_LENGTH` | `8` | enforced at register, reset, and in the `/admin` create form |
| `AUTH_RATE_LIMIT_IP_RPS` | `5` | per-IP token-bucket rate for public auth endpoints (register, login, forgot/reset password, magic-link) |
| `AUTH_RATE_LIMIT_EMAIL_RPS` | `2` | per-recipient token-bucket rate for mail-sending auth endpoints (forgot-password, magic-link request) |
| `WEB_ADMIN_ENABLED` | `true` | toggle the `/admin/*` server-rendered management UI |
| `ADMIN_REQUIRED_ROLE` | `admin` | `users.role` value required to enter `/admin` (must match the DB CHECK constraint) |

9
src/contract_check/api/admin/router.py
View file

@ -16,12 +16,12 @@ from src.contract_check.api.admin.auth import ADMIN_COOKIE, ADMIN_COOKIE_MAX_AGE
from src.contract_check.api.admin.billing import router as billing_router
from src.contract_check.api.admin.templating import templates
from src.contract_check.api.admin.users import router as users_router
from src.contract_check.api.deps import AsyncSessionDep
from src.contract_check.api.deps import AsyncSessionDep, RateLimiterDep, require_auth_rate_limit
from src.contract_check.core.auth import create_access_token
from src.contract_check.core.config import get_settings
from src.contract_check.core.db.repositories import UserRepository
from src.contract_check.core.logging import get_logger
from src.contract_check.core.security.passwords import verify_password
from src.contract_check.core.security.passwords import verify_password_async
log = get_logger(__name__)
@ -71,12 +71,15 @@ async def login_form(
@router.post("/admin/login", include_in_schema=False)
async def login(
request: Request,
session: AsyncSessionDep,
rate_limiter: RateLimiterDep,
email: Annotated[str, Form()],
password: Annotated[str, Form()],
) -> RedirectResponse:
"""Verify email/password + admin role, set the cookie, redirect to users."""
email_norm = email.lower().strip()
await require_auth_rate_limit(request, rate_limiter, email=email_norm)
user = await UserRepository(session).get_by_email(email_norm)
invalid = RedirectResponse(
url="/admin/login?reason=invalid", status_code=status.HTTP_303_SEE_OTHER
@ -84,7 +87,7 @@ async def login(
if user is None:
return invalid
if not user.password_hash or not verify_password(password, user.password_hash):
if not user.password_hash or not await verify_password_async(password, user.password_hash):
return invalid
if not user.is_active:
return RedirectResponse(

4
src/contract_check/api/app.py
View file

@ -166,7 +166,7 @@ async def _ensure_default_admin(settings: Any, redis_client: Any) -> None:
available during lifespan setup. Requires web auth (Redis) to be available.
"""
from src.contract_check.core.db.session import create_session_factory
from src.contract_check.core.security.passwords import hash_password
from src.contract_check.core.security.passwords import hash_password_async
logger = get_logger(__name__)
if redis_client is None:
@ -182,7 +182,7 @@ async def _ensure_default_admin(settings: Any, redis_client: Any) -> None:
session,
email=settings.admin_default_email.lower().strip(),
name="Administrator",
password_hash=hash_password(settings.admin_default_password),
password_hash=await hash_password_async(settings.admin_default_password),
)
await repo.set_role(user.id, settings.admin_required_role)
await session.commit()

57
src/contract_check/api/deps.py
View file

@ -6,7 +6,7 @@ from collections.abc import AsyncIterator
from typing import Annotated, Any
from uuid import UUID
from fastapi import Depends, Header, HTTPException, Query, Request
from fastapi import Depends, Header, HTTPException, Query, Request, status
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
@ -32,9 +32,9 @@ from src.contract_check.core.logging import get_logger
from src.contract_check.core.mq.publisher import Publisher
from src.contract_check.core.notifications.publisher import NotificationPublisher
from src.contract_check.core.passkeys import PasskeyChallengeStore
from src.contract_check.core.rate_limit import RateLimiter
from src.contract_check.core.rate_limit import RateLimiter, RateLimitResult
from src.contract_check.core.s3.port import Storage
from src.contract_check.core.security.passwords import hash_password
from src.contract_check.core.security.passwords import hash_password_async
from src.contract_check.core.tokens import hash_token
log = get_logger(__name__)
@ -182,11 +182,60 @@ async def bind_telegram_to_user(
async def set_user_password(session: AsyncSession, user_id: UUID, password: str) -> None:
"""Set/rotate a web password for a user (Telegram -> web UI access)."""
hashed = hash_password(password)
hashed = await hash_password_async(password)
await UserRepository(session).set_password(user_id, hashed)
await session.commit()
def get_rate_limiter(request: Request) -> RateLimiter:
"""Return the app-state rate limiter (Redis or in-memory fallback)."""
limiter: RateLimiter = request.app.state.rate_limiter
return limiter
RateLimiterDep = Annotated[RateLimiter, Depends(get_rate_limiter)]
def _auth_rate_limit_headers(result: RateLimitResult) -> dict[str, str]:
retry = int(result.retry_after_sec or 1)
return {"Retry-After": str(retry)}
async def require_auth_rate_limit(
request: Request,
rate_limiter: RateLimiter,
*,
email: str | None = None,
) -> None:
"""Throttle public auth endpoints by IP and, optionally, by recipient email.
Mail-sending endpoints should pass ``email`` so a single IP cannot bomb an
inbox; other endpoints pass only the IP bucket.
"""
settings = get_settings()
ip = request.client.host if request.client else "unknown"
ip_key = f"auth:ip:{ip}"
ip_result = await rate_limiter.allow(ip_key, settings.auth_rate_limit_ip_rps)
if not ip_result.allowed:
raise HTTPException(
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
detail="rate limit exceeded",
headers=_auth_rate_limit_headers(ip_result),
)
if email:
email_key = f"auth:email:{email.lower().strip()}"
email_result = await rate_limiter.allow(email_key, settings.auth_rate_limit_email_rps)
if not email_result.allowed:
raise HTTPException(
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
detail="rate limit exceeded",
headers=_auth_rate_limit_headers(email_result),
)
AuthRateLimiterDep = Annotated[RateLimiter, Depends(get_rate_limiter)]
async def fetch_user_by_email(session: AsyncSession, email: str) -> User | None:
"""Fetch a user by email (case-sensitive — normalize upstream). Returns None if not found."""
return await UserRepository(session).get_by_email(email)

10
src/contract_check/api/routes/auth/magic_links.py
View file

@ -4,13 +4,15 @@ import datetime as dt
import secrets
import uuid
from fastapi import APIRouter, HTTPException, status
from fastapi import APIRouter, HTTPException, Request, status
from src.contract_check.api.deps import (
AsyncSessionDep,
NotificationPublisherDep,
RateLimiterDep,
fetch_user_by_email,
fetch_user_by_id_full,
require_auth_rate_limit,
)
from src.contract_check.api.routes.auth.support import (
_build_magic_link,
@ -39,8 +41,10 @@ router = APIRouter(tags=["auth"])
status_code=status.HTTP_200_OK,
)
async def magic_link_request(
request: Request,
session: AsyncSessionDep,
publisher: NotificationPublisherDep,
rate_limiter: RateLimiterDep,
body: MagicLinkRequest,
) -> MagicLinkResponse:
"""Generate a one-time login token, store its hash, and enqueue a notification.
@ -51,6 +55,7 @@ async def magic_link_request(
"""
_require_magic_link_enabled()
email_normalized = body.email.lower().strip()
await require_auth_rate_limit(request, rate_limiter, email=email_normalized)
settings = get_settings()
generic = MagicLinkResponse(
success=True,
@ -112,7 +117,9 @@ async def magic_link_request(
status_code=status.HTTP_200_OK,
)
async def magic_link_verify(
request: Request,
session: AsyncSessionDep,
rate_limiter: RateLimiterDep,
body: MagicLinkVerifyRequest,
) -> SingleTokenAuthResponse:
"""Exchange a magic-link token for an access JWT.
@ -121,6 +128,7 @@ async def magic_link_verify(
once. Expired tokens are also cleared (they can never be replayed).
"""
_require_magic_link_enabled()
await require_auth_rate_limit(request, rate_limiter)
token_hash = _hash_reset_token(body.token)
users = UserRepository(session)

33
src/contract_check/api/routes/auth/password.py
View file

@ -4,16 +4,18 @@ import datetime as dt
import secrets
import uuid
from fastapi import APIRouter, HTTPException, status
from fastapi import APIRouter, HTTPException, Request, status
from src.contract_check.api.deps import (
AsyncSessionDep,
CurrentUserDep,
NotificationPublisherDep,
RateLimiterDep,
RefreshStoreDep,
create_email_user,
fetch_user_by_email,
fetch_user_by_id_full,
require_auth_rate_limit,
)
from src.contract_check.api.routes.auth.support import (
_build_reset_link,
@ -36,7 +38,7 @@ from src.contract_check.core.config import get_settings
from src.contract_check.core.db.repositories import UserRepository
from src.contract_check.core.logging import get_logger
from src.contract_check.core.mq.messages import NotificationMessage
from src.contract_check.core.security.passwords import hash_password, verify_password
from src.contract_check.core.security.passwords import hash_password_async, verify_password_async
log = get_logger(__name__)
router = APIRouter(tags=["auth"])
@ -48,12 +50,15 @@ router = APIRouter(tags=["auth"])
status_code=status.HTTP_201_CREATED,
)
async def register(
request: Request,
session: AsyncSessionDep,
refresh_store: RefreshStoreDep,
rate_limiter: RateLimiterDep,
body: RegisterRequest,
) -> TokenPairResponse:
"""Register a new email/password user and issue a JWT pair."""
_require_web_auth_enabled()
await require_auth_rate_limit(request, rate_limiter)
email_normalized = body.email.lower().strip()
existing = await fetch_user_by_email(session, email_normalized)
@ -61,7 +66,7 @@ async def register(
# Do not leak which emails are registered.
raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail="email already registered")
hashed = hash_password(body.password)
hashed = await hash_password_async(body.password)
user = await create_email_user(
session, email=email_normalized, name=body.name.strip(), password_hash=hashed
)
@ -76,18 +81,21 @@ async def register(
status_code=status.HTTP_200_OK,
)
async def login(
request: Request,
session: AsyncSessionDep,
refresh_store: RefreshStoreDep,
rate_limiter: RateLimiterDep,
body: LoginRequest,
) -> TokenPairResponse:
"""Email + password -> JWT pair."""
_require_web_auth_enabled()
email_normalized = body.email.lower().strip()
await require_auth_rate_limit(request, rate_limiter, email=email_normalized)
user = await fetch_user_by_email(session, email_normalized)
if user is None or not user.password_hash:
raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="invalid credentials")
if not verify_password(body.password, user.password_hash):
if not await verify_password_async(body.password, user.password_hash):
raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="invalid credentials")
if not user.is_active:
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="account disabled")
@ -124,8 +132,10 @@ async def logout(
status_code=status.HTTP_202_ACCEPTED,
)
async def forgot_password(
request: Request,
session: AsyncSessionDep,
publisher: NotificationPublisherDep,
rate_limiter: RateLimiterDep,
body: ForgotPasswordRequest,
) -> OkResponse:
"""Generate a reset token, store its hash + expiry, and enqueue a notification.
@ -135,6 +145,7 @@ async def forgot_password(
"""
_require_web_auth_enabled()
email_normalized = body.email.lower().strip()
await require_auth_rate_limit(request, rate_limiter, email=email_normalized)
user = await fetch_user_by_email(session, email_normalized)
if user is None:
@ -193,8 +204,10 @@ async def forgot_password(
status_code=status.HTTP_200_OK,
)
async def reset_password(
request: Request,
session: AsyncSessionDep,
refresh_store: RefreshStoreDep,
rate_limiter: RateLimiterDep,
body: ResetPasswordRequest,
) -> OkResponse:
"""Verify a reset token and set the new password.
@ -203,6 +216,7 @@ async def reset_password(
all active refresh tokens for the user (forcing re-login everywhere).
"""
_require_web_auth_enabled()
await require_auth_rate_limit(request, rate_limiter)
token_hash = _hash_reset_token(body.token)
users = UserRepository(session)
@ -215,7 +229,7 @@ async def reset_password(
if expires_at is None or expires_at < now:
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="reset token expired")
new_hash = hash_password(body.password)
new_hash = await hash_password_async(body.password)
await users.reset_password(user_id, password_hash=new_hash)
await session.commit()
@ -236,9 +250,11 @@ async def reset_password(
status_code=status.HTTP_200_OK,
)
async def change_password(
request: Request,
session: AsyncSessionDep,
refresh_store: RefreshStoreDep,
user: CurrentUserDep,
rate_limiter: RateLimiterDep,
body: ChangePasswordRequest,
) -> OkResponse:
"""Rotate the password of the authenticated user.
@ -247,14 +263,17 @@ async def change_password(
revoked, forcing re-login on other devices.
"""
_require_web_auth_enabled()
await require_auth_rate_limit(request, rate_limiter)
db_user = await fetch_user_by_id_full(session, user.user_id)
if db_user is None or not db_user.password_hash:
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="no password set")
if not verify_password(body.current_password, db_user.password_hash):
if not await verify_password_async(body.current_password, db_user.password_hash):
raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="invalid credentials")
await UserRepository(session).set_password(user.user_id, hash_password(body.new_password))
await UserRepository(session).set_password(
user.user_id, await hash_password_async(body.new_password)
)
await session.commit()
# Best-effort revoke of existing sessions; ignore Redis hiccups so the

13
src/contract_check/api/routes/b2b.py
View file

@ -22,6 +22,7 @@ from src.contract_check.api.deps import (
)
from src.contract_check.api.schemas import (
ApiKeyCreatedResponse,
ApiKeyMonthlyRequest,
ApiKeyResponse,
ApiKeyUsageDetailResponse,
B2BUsageResponse,
@ -119,7 +120,8 @@ async def get_usage(
"""Return current-month usage for the authenticating API key."""
repo = ApiKeyRepository(session)
key = await repo.get_by_id(auth.api_key_id)
assert key is not None
if key is None:
raise HTTPException(status_code=404, detail="api key not found")
_, monthly_used = await repo.get_monthly_quota_state(auth.api_key_id)
@ -276,7 +278,7 @@ async def get_key_usage(
if key is None or key.user_id != user.user_id:
raise HTTPException(status_code=404, detail="key not found")
monthly_result = await repo.get_usage_by_key(key_id)
monthly_result = await repo.get_monthly_usage(key_id)
return ApiKeyUsageDetailResponse(
key_id=str(key_id),
@ -286,7 +288,10 @@ async def get_key_usage(
monthly_used=key.monthly_used,
resets_at=key.resets_at.isoformat() if key.resets_at else None,
monthly_requests=[
{"month": req.created_at.isoformat() if req.created_at else None, "requests": 1}
for req in monthly_result
ApiKeyMonthlyRequest(
month=month.isoformat() if month else None,
requests=count,
)
for month, count in monthly_result
],
)

16
src/contract_check/api/routes/metrics.py
View file

@ -2,12 +2,24 @@
from __future__ import annotations
from fastapi import APIRouter, Response
from fastapi import APIRouter, Header, HTTPException, Response, status
from prometheus_client import CONTENT_TYPE_LATEST, generate_latest
from src.contract_check.core.config import get_settings
router = APIRouter(tags=["metrics"])
@router.get("/metrics")
async def metrics() -> Response:
async def metrics(authorization: str | None = Header(default=None)) -> Response:
token = get_settings().metrics_bearer_token
if token:
if not authorization or not authorization.lower().startswith("bearer "):
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED, detail="missing bearer token"
)
if authorization[7:].strip() != token:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED, detail="invalid bearer token"
)
return Response(generate_latest(), media_type=CONTENT_TYPE_LATEST)

2
src/contract_check/api/schemas/__init__.py
View file

@ -48,6 +48,7 @@ from src.contract_check.api.schemas.auth import (
)
from src.contract_check.api.schemas.b2b import (
ApiKeyCreatedResponse,
ApiKeyMonthlyRequest,
ApiKeyResponse,
ApiKeyUsageDetailResponse,
B2BUsageResponse,
@ -167,5 +168,6 @@ __all__ = [
"ApiKeyCreatedResponse",
"B2BUsageResponse",
"ApiKeyUsageDetailResponse",
"ApiKeyMonthlyRequest",
"RevokeApiKeyResponse",
]

3
src/contract_check/api/schemas/b2b.py
View file

@ -3,7 +3,6 @@
from __future__ import annotations
import uuid
from typing import Any
from pydantic import BaseModel
@ -51,7 +50,7 @@ class ApiKeyUsageDetailResponse(BaseModel):
monthly_quota: int | None
monthly_used: int
resets_at: str | None
monthly_requests: list[dict[str, Any]]
monthly_requests: list[ApiKeyMonthlyRequest]
class RevokeApiKeyResponse(BaseModel):

63
src/contract_check/api/services.py
View file

@ -6,7 +6,6 @@ import uuid
from typing import TYPE_CHECKING
from fastapi import HTTPException
from sqlalchemy import text
from src.contract_check.api.schemas import DocumentUploadResponse
from src.contract_check.core.billing.errors import NoCredits
@ -16,7 +15,11 @@ from src.contract_check.core.billing.quota import (
)
from src.contract_check.core.config import get_settings
from src.contract_check.core.credits import adjust_credits, reserve_credit
from src.contract_check.core.db.repositories import DocumentRepository, JobRepository
from src.contract_check.core.db.repositories import (
DocumentRepository,
JobRepository,
UserRepository,
)
from src.contract_check.core.db.repositories.credits import CreditsRepository
from src.contract_check.core.extraction.formats import SUPPORTED_SUFFIXES
from src.contract_check.core.logging import get_logger, new_correlation_id
@ -73,6 +76,12 @@ async def upload_and_enqueue(
detail=f"unsupported format {suffix!r}; supported: {sorted(SUPPORTED_SUFFIXES)}",
)
settings = get_settings()
# Reject Billing Holds before reading the body or writing storage.
if await UserRepository(session).get_billing_hold(user_id):
raise HTTPException(status_code=402, detail="billing hold")
credits_repo = CreditsRepository(session)
doc_repo = DocumentRepository(session)
job_repo = JobRepository(session)
@ -82,13 +91,12 @@ async def upload_and_enqueue(
s3_key = original_key(str(user_id), str(document_id), suffix)
content_type = file.content_type or _content_type_from_suffix(suffix)
try:
data = await file.read()
data = await _read_upload_with_limit(file, settings.max_upload_bytes)
if len(data) == 0:
raise HTTPException(status_code=400, detail="empty file")
try:
await storage.put(s3_key, data, content_type=content_type)
except HTTPException:
raise
except Exception as exc:
log.error("s3_upload_failed", document_id=str(document_id), error=str(exc))
raise HTTPException(status_code=500, detail="failed to store document") from exc
@ -111,19 +119,20 @@ async def upload_and_enqueue(
)
except Exception as exc:
log.error("db_enqueue_failed", document_id=str(document_id), error=str(exc))
# Best-effort cleanup so a DB failure does not orphan the stored object.
try:
await storage.delete(s3_key)
except Exception as cleanup_exc: # noqa: BLE001
log.warning(
"upload_cleanup_failed",
document_id=str(document_id),
s3_key=s3_key,
error=str(cleanup_exc),
)
raise HTTPException(status_code=500, detail="failed to enqueue document") from exc
# Reserve credit *after* the document row exists so the ledger can link it.
result = await session.execute(
text("SELECT billing_hold FROM users WHERE id = :u"),
{"u": user_id},
)
row = result.first()
if row is not None and row[0]:
raise HTTPException(status_code=402, detail="billing hold")
source = "credits"
if get_settings().plans_enabled:
if settings.plans_enabled:
try:
source = await reserve_document_slot(session, user_id, document_id)
except NoCredits:
@ -167,3 +176,25 @@ async def upload_and_enqueue(
correlation_id=str(correlation_id),
credits_left=credits_left,
)
async def _read_upload_with_limit(file: UploadFile, max_bytes: int) -> bytes:
"""Read ``file`` in chunks, raising 413 if ``max_bytes`` is exceeded.
The file is never fully buffered in memory beyond ``max_bytes + 1``.
"""
chunks: list[bytes] = []
total = 0
chunk_size = 64 * 1024
while True:
chunk = await file.read(chunk_size)
if not chunk:
break
total += len(chunk)
if total > max_bytes:
raise HTTPException(
status_code=413,
detail=f"file exceeds maximum size of {max_bytes} bytes",
)
chunks.append(chunk)
return b"".join(chunks)

15
src/contract_check/core/billing/fulfillment.py
View file

@ -17,7 +17,9 @@ from src.contract_check.core.db.repositories import (
SubscriptionsRepository,
)
from src.contract_check.core.db.repositories.invoices import InvoiceRow
from src.contract_check.core.db.repositories.subscriptions import SubscriptionError
from src.contract_check.core.logging import get_logger
from src.contract_check.core.metrics import billing_invoices_poisoned
if TYPE_CHECKING:
from sqlalchemy.ext.asyncio import AsyncSession
@ -135,14 +137,17 @@ async def _fulfill(
start=now,
end=now + dt.timedelta(days=30),
)
except Exception as exc:
log.error(
"subscription_fulfill_failed",
except SubscriptionError as exc:
log.warning(
"subscription_fulfill_poisoned",
invoice_id=str(invoice["id"]),
payment_id=payment_id,
error=str(exc),
user_id=str(invoice["user_id"]),
reason=str(exc),
)
raise
billing_invoices_poisoned.inc()
await invoices.set_cancelled(invoice["id"])
return
await invoices.set_succeeded(
invoice["id"],
paid_at=dt.datetime.now(tz=dt.UTC),

32
src/contract_check/core/billing/quota.py
View file

@ -13,6 +13,7 @@ from typing import TYPE_CHECKING
from sqlalchemy import text
from src.contract_check.core.billing.errors import NoCredits
from src.contract_check.core.db.enums import RefundPolicyLike
if TYPE_CHECKING:
from sqlalchemy.ext.asyncio import AsyncSession
@ -94,3 +95,34 @@ async def release_document_slot(session: AsyncSession, document_id: uuid.UUID) -
{"d": document_id},
)
return result.first() is not None
async def compensate_document_slot(
session: AsyncSession,
document_id: uuid.UUID,
failure_class: str,
policy: RefundPolicyLike,
) -> tuple[bool, bool]:
"""Compensate a consumed Document Slot exactly once.
Quota-paid documents get their quota reservation back and ``refunded`` set.
Credits-paid documents get one Credit back via the flag-guarded refund.
Idempotent: a second call is a no-op because ``documents.refunded`` is set
by both paths.
Returns ``(quota_released, credit_refunded)`` for observability/tests.
"""
from src.contract_check.core.credits import refund_credit
quota_released = await release_document_slot(session, document_id)
if quota_released:
# Quota-paid: set the shared compensation marker; no credit refund.
await session.execute(
text("UPDATE documents SET refunded = TRUE WHERE id = :d AND refunded = FALSE"),
{"d": document_id},
)
return True, False
# Credits-paid: flag-guarded refund (sets ``refunded`` itself).
credit_refunded = await refund_credit(session, document_id, failure_class, policy)
return False, credit_refunded

18
src/contract_check/core/config.py
View file

@ -56,6 +56,10 @@ class Settings(BaseSettings):
s3_server_side_encryption: bool = False
doc_retention_days: int = 7
text_retention_days: int = 30
max_upload_bytes: int = Field(
default=25 * 1024 * 1024,
description="Maximum uploaded file size in bytes (default 25 MiB).",
)
# --- billing ---
refund_policy: RefundPolicy = "all"
@ -109,6 +113,10 @@ class Settings(BaseSettings):
api_port: int = 8000
api_metrics_port: int = 9100
b2b_default_rate_limit_rps: int = 3
metrics_bearer_token: str = Field(
default="",
description="Bearer token protecting /metrics. Empty leaves the endpoint open.",
)
# SSE streaming of report status (GET /api/v1/reports/{id}/events).
sse_poll_interval_seconds: float = 1.0
@ -149,6 +157,16 @@ class Settings(BaseSettings):
# Minimum password length enforced at register / reset.
password_min_length: int = 8
# --- auth rate limiting ---
auth_rate_limit_ip_rps: int = Field(
default=5,
description="Per-IP token-bucket rate for public auth endpoints (login, register, etc.).",
)
auth_rate_limit_email_rps: int = Field(
default=2,
description="Per-recipient token-bucket rate for mail-sending auth endpoints.",
)
# --- passkeys (WebAuthn) ---
passkey_enabled: bool = Field(
default=True,

2
src/contract_check/core/db/enums.py
View file

@ -29,7 +29,7 @@ DOC_STATUSES: tuple[str, ...] = (
"failed",
"manual_review",
)
DOC_TERMINAL: tuple[str, ...] = ("done", "failed")
DOC_TERMINAL: tuple[str, ...] = ("done", "failed", "manual_review")
# jobs.status
JobStatus = Literal["pending", "running", "retrying", "dlq", "done"]

18
src/contract_check/core/db/repositories/api_keys.py
View file

@ -191,3 +191,21 @@ class ApiKeyRepository:
{"k": api_key_id},
)
return int(result.scalar_one())
async def get_monthly_usage(self, api_key_id: uuid.UUID) -> list[tuple[dt.datetime, int]]:
"""Return per-calendar-month request counts, newest first.
Each tuple is ``(month_start, request_count)``. Months with zero
requests are omitted.
"""
result = await self._session.execute(
text(
"SELECT DATE_TRUNC('month', created_at) AS month, count(*) "
"FROM api_key_requests "
"WHERE api_key_id = :k "
"GROUP BY month "
"ORDER BY month DESC"
),
{"k": api_key_id},
)
return [(row[0], int(row[1])) for row in result.all()]

37
src/contract_check/core/db/repositories/credits.py
View file

@ -76,30 +76,36 @@ class CreditsRepository:
failure_class: FailureClass | str,
policy: RefundPolicyLike,
) -> bool:
"""Refund one credit once, idempotent via ``documents.refunded``."""
"""Refund one credit once, idempotent via ``documents.refunded``.
The ``documents.refunded`` flag is flipped first; its row lock serializes
concurrent refund attempts so exactly one transaction wins and writes the
ledger event.
"""
if not should_refund(failure_class, policy):
return False
result = await self._session.execute(
# Flag-first: row lock on documents serializes concurrent refunds.
flag_result = await self._session.execute(
text(
"UPDATE users SET credits_left = credits_left + 1 "
"WHERE id = (SELECT user_id FROM documents "
" WHERE id = :d AND refunded = FALSE) "
"RETURNING credits_left"
"UPDATE documents SET refunded = TRUE "
"WHERE id = :d AND refunded = FALSE "
"RETURNING user_id"
),
{"d": document_id},
)
row = result.first()
row = flag_result.first()
if row is None:
return False
balance_after = int(row[0])
user_id = row[0]
# Need the user_id for the ledger; capture it from the same UPDATE.
user_result = await self._session.execute(
text("SELECT user_id FROM documents WHERE id = :d AND refunded = FALSE"),
{"d": document_id},
balance_result = await self._session.execute(
text(
"UPDATE users SET credits_left = credits_left + 1 WHERE id = :u RETURNING credits_left"
),
{"u": user_id},
)
user_id = user_result.scalar_one()
balance_after = int(balance_result.scalar_one())
await self._events.append(
user_id,
@ -108,11 +114,6 @@ class CreditsRepository:
balance_after=balance_after,
document_id=document_id,
)
await self._session.execute(
text("UPDATE documents SET refunded = TRUE WHERE id = :d"),
{"d": document_id},
)
return True
async def adjust(

12
src/contract_check/core/db/repositories/reports.py
View file

@ -37,6 +37,8 @@ class ReportRepository:
prompt_tokens: int,
eval_tokens: int,
latency_ms: int,
prescreen_meta: dict[str, Any] | None = None,
prescreen_result_id: uuid.UUID | None = None,
) -> uuid.UUID:
"""Idempotent insert for the analyze worker (ON CONFLICT document_id)."""
content_text = json.dumps(content_json, ensure_ascii=False)
@ -44,15 +46,17 @@ class ReportRepository:
text(
"INSERT INTO reports "
"(document_id, content_json, markdown, model_used, "
" prompt_tokens, eval_tokens, latency_ms) "
"VALUES (:d, CAST(:content AS jsonb), :md, :model, :pt, :et, :lat) "
" prompt_tokens, eval_tokens, latency_ms, prescreen_result_id, prescreen_meta) "
"VALUES (:d, CAST(:content AS jsonb), :md, :model, :pt, :et, :lat, :prid, :pm) "
"ON CONFLICT (document_id) DO UPDATE SET "
" content_json = EXCLUDED.content_json, "
" markdown = EXCLUDED.markdown, "
" model_used = EXCLUDED.model_used, "
" prompt_tokens = EXCLUDED.prompt_tokens, "
" eval_tokens = EXCLUDED.eval_tokens, "
" latency_ms = EXCLUDED.latency_ms "
" latency_ms = EXCLUDED.latency_ms, "
" prescreen_result_id = EXCLUDED.prescreen_result_id, "
" prescreen_meta = EXCLUDED.prescreen_meta "
"RETURNING id"
),
{
@ -63,6 +67,8 @@ class ReportRepository:
"pt": prompt_tokens,
"et": eval_tokens,
"lat": latency_ms,
"prid": prescreen_result_id,
"pm": json.dumps(prescreen_meta) if prescreen_meta is not None else None,
},
)
row = result.first()

20
src/contract_check/core/db/repositories/subscriptions.py
View file

@ -67,16 +67,30 @@ class SubscriptionsRepository:
start: dt.datetime,
end: dt.datetime,
) -> uuid.UUID:
"""Create an active subscription from a paid invoice (ticket 015)."""
"""Create an active subscription from a paid invoice (ticket 015).
Idempotent: if a subscription already exists for this origin invoice,
return its id. If the user has a different active/past_due subscription,
raise ``SubscriptionError`` so the caller can poison the invoice.
"""
existing = await self._session.execute(
text(
"SELECT 1 FROM subscriptions "
"SELECT id FROM subscriptions "
"WHERE user_id = :u AND status IN ('active','past_due') "
"LIMIT 1"
),
{"u": user_id},
)
if existing.first() is not None:
row = existing.first()
if row is not None:
current_id = row[0]
# If this invoice already granted the current subscription, no-op.
same = await self._session.execute(
text("SELECT 1 FROM subscriptions WHERE id = :id AND origin_invoice_id = :inv"),
{"id": current_id, "inv": origin_invoice_id},
)
if same.first() is not None:
return current_id
raise SubscriptionError("user already has an active subscription")
result = await self._session.execute(

4
src/contract_check/core/metrics.py
View file

@ -78,6 +78,10 @@ billing_invoices_reconciled = Counter(
"contract_check_billing_invoices_reconciled_total",
"Pending invoices reconciled by polling the provider.",
)
billing_invoices_poisoned = Counter(
"contract_check_billing_invoices_poisoned_total",
"Invoices cancelled because they cannot be fulfilled (duplicate subscription).",
)
billing_refunds_processed = Counter(
"contract_check_billing_refunds_processed_total",
"Refunds processed by the billing API.",

33
src/contract_check/core/mq/consumer.py
View file

@ -192,7 +192,18 @@ class Consumer[MsgT: RetryableMessage]:
failure_class = self.classify(exc)
error = f"{type(exc).__name__}: {exc}"
new_attempt = attempt + 1
try:
await self.on_failure(payload, failure_class, new_attempt, error)
except Exception as hook_exc: # noqa: BLE001
log.error(
"on_failure_hook_failed",
queue=self.queue,
correlation_id=cid,
error=f"{type(hook_exc).__name__}: {hook_exc}",
)
await self._to_retry(message, payload, attempt)
await message.ack()
return
if _is_terminal(exc):
log.warning(
@ -201,7 +212,18 @@ class Consumer[MsgT: RetryableMessage]:
correlation_id=cid,
failure_class=failure_class,
)
try:
await self.on_dlq(payload, failure_class, error)
except Exception as hook_exc: # noqa: BLE001
log.error(
"on_dlq_hook_failed",
queue=self.queue,
correlation_id=cid,
error=f"{type(hook_exc).__name__}: {hook_exc}",
)
await self._to_retry(message, payload, attempt)
await message.ack()
return
await self._to_dlq(message, failure_class, error, attempt=new_attempt)
return
@ -213,7 +235,18 @@ class Consumer[MsgT: RetryableMessage]:
attempt=new_attempt,
failure_class=failure_class,
)
try:
await self.on_dlq(payload, failure_class, error)
except Exception as hook_exc: # noqa: BLE001
log.error(
"on_dlq_hook_failed",
queue=self.queue,
correlation_id=cid,
error=f"{type(hook_exc).__name__}: {hook_exc}",
)
await self._to_retry(message, payload, attempt)
await message.ack()
return
await self._to_dlq(message, failure_class, error, attempt=new_attempt)
else:
await self._to_retry(message, payload, new_attempt)

7
src/contract_check/core/mq/messages.py
View file

@ -69,16 +69,19 @@ class AnalyzeRequested(PipelineMessage):
"""worker-prescreen → contracts.x[analyze] → worker-analyze.
Inherits the same shape as DocumentExtracted but carries additional
structured metadata discovered during prescreening.
structured metadata discovered during prescreening. ``filename`` is
optional so that a direct extract-to-analyze message (DocumentExtracted)
also validates against this superset model.
"""
extracted_s3_key: str
filename: str
filename: str | None = None
char_count: int
ocr_used: bool
is_structured: bool = False
has_tables: bool = False
prescreen_meta: dict | None = None
prescreen_result_id: uuid.UUID | None = None
class PrescreenCompleted(PipelineMessage):

15
src/contract_check/core/security/passwords.py
View file

@ -3,10 +3,15 @@
Pure crypto helpers — no FastAPI, no DB. Used by the api at register/login and
by password-reset. argon2-cffi is the OWASP-recommended PHC-format hasher; the
encoded hash embeds salt + params so the verifier auto-detects them on verify.
The synchronous helpers remain for non-async call sites; the `*_async`
variants run the CPU-heavy work off the event loop.
"""
from __future__ import annotations
import asyncio
from argon2 import PasswordHasher
from argon2.exceptions import InvalidHash, VerificationError, VerifyMismatchError
@ -53,3 +58,13 @@ def needs_rehash(hashed: str) -> bool:
return _hasher.check_needs_rehash(hashed)
except (InvalidHash, TypeError):
return False
async def hash_password_async(plain: str) -> str:
"""Event-loop-friendly wrapper for :func:`hash_password`."""
return await asyncio.to_thread(hash_password, plain)
async def verify_password_async(plain: str, hashed: str) -> bool:
"""Event-loop-friendly wrapper for :func:`verify_password`."""
return await asyncio.to_thread(verify_password, plain, hashed)

14
src/contract_check/worker_analyze/consumer.py
View file

@ -6,20 +6,19 @@ from src.contract_check.core.db.enums import FailureClass
from src.contract_check.core.db.session import create_session_factory
from src.contract_check.core.llm.port import LLMProvider
from src.contract_check.core.logging import get_logger
from src.contract_check.core.metrics import analyze_duration
from src.contract_check.core.mq.consumer import Consumer
from src.contract_check.core.mq.messages import DocumentExtracted
from src.contract_check.core.mq.messages import AnalyzeRequested
from src.contract_check.worker_analyze.handler import AnalyzeHandler
log = get_logger(__name__)
class AnalyzeConsumer(Consumer[DocumentExtracted]):
class AnalyzeConsumer(Consumer[AnalyzeRequested]):
"""Consumes `analyze.q`, runs the LLM analysis, persists the report."""
queue: str = "analyze.q"
routing_key: str = "analyze"
message_model = DocumentExtracted
message_model = AnalyzeRequested
def __init__(
self,
@ -44,13 +43,12 @@ class AnalyzeConsumer(Consumer[DocumentExtracted]):
def classify(self, exc: BaseException) -> FailureClass:
return self._handler.classify(exc)
@analyze_duration.time()
async def handle(self, payload: DocumentExtracted) -> None:
async def handle(self, payload: AnalyzeRequested) -> None:
await self._handler.handle(payload)
async def on_failure(
self,
payload: DocumentExtracted,
payload: AnalyzeRequested,
failure_class: FailureClass,
attempt: int,
error: str,
@ -58,7 +56,7 @@ class AnalyzeConsumer(Consumer[DocumentExtracted]):
await self._handler.on_failure(payload, failure_class, attempt, error)
async def on_dlq(
self, payload: DocumentExtracted, failure_class: FailureClass, error: str
self, payload: AnalyzeRequested, failure_class: FailureClass, error: str
) -> None:
await self._handler.on_terminal_failure(payload, failure_class, error)

18
src/contract_check/worker_analyze/handler.py
View file

@ -17,9 +17,8 @@ from typing import TYPE_CHECKING
from src.contract_check.core.analysis.analyzer import RenderMetrics, render_markdown
from src.contract_check.core.analysis.checklist import checklist_for_prompt
from src.contract_check.core.billing.quota import release_document_slot
from src.contract_check.core.billing.quota import compensate_document_slot
from src.contract_check.core.config import get_settings
from src.contract_check.core.credits import refund_credit
from src.contract_check.core.db.enums import DOC_TERMINAL, FailureClass
from src.contract_check.core.db.repositories import (
DocumentRepository,
@ -36,7 +35,7 @@ from src.contract_check.core.llm.factory import build_llm_provider
from src.contract_check.core.llm.port import LLMProvider
from src.contract_check.core.logging import get_logger
from src.contract_check.core.metrics import analyze_duration, llm_fell_back, llm_tokens, mq_failed
from src.contract_check.core.mq.messages import DocumentExtracted
from src.contract_check.core.mq.messages import AnalyzeRequested
from src.contract_check.core.s3.minio_storage import MinioStorage
if TYPE_CHECKING:
@ -80,7 +79,7 @@ class AnalyzeHandler:
) -> tuple[str | None, str | None]:
return await DocumentRepository(session).get_status_and_filename_for_update(document_id)
async def handle(self, payload: DocumentExtracted) -> None:
async def handle(self, payload: AnalyzeRequested) -> None:
async with self._session_factory() as session:
docs = DocumentRepository(session)
jobs = JobRepository(session)
@ -98,7 +97,7 @@ class AnalyzeHandler:
text_bytes = await self._storage.get(payload.extracted_s3_key)
contract_text = text_bytes.decode("utf-8")
source_name = filename or payload.extracted_s3_key
source_name = payload.filename or filename or payload.extracted_s3_key
provider = await self._provider_instance()
with analyze_duration.time():
@ -139,6 +138,8 @@ class AnalyzeHandler:
prompt_tokens=result.prompt_tokens,
eval_tokens=result.eval_tokens,
latency_ms=int(result.latency_sec * 1000),
prescreen_meta=payload.prescreen_meta,
prescreen_result_id=payload.prescreen_result_id,
)
await docs.update_status(payload.document_id, status="done", stage="done")
await jobs.mark_done(payload.document_id, "analyze")
@ -163,7 +164,7 @@ class AnalyzeHandler:
async def on_failure(
self,
payload: DocumentExtracted,
payload: AnalyzeRequested,
failure_class: FailureClass,
attempt: int,
error: str,
@ -187,7 +188,7 @@ class AnalyzeHandler:
await session.commit()
async def on_terminal_failure(
self, payload: DocumentExtracted, failure_class: FailureClass, error: str
self, payload: AnalyzeRequested, failure_class: FailureClass, error: str
) -> None:
mq_failed.labels(queue="analyze", failure_class=failure_class).inc()
async with self._session_factory() as session:
@ -197,8 +198,7 @@ class AnalyzeHandler:
payload.document_id, "analyze", failure_class=failure_class, error=error
)
await docs.mark_failed(payload.document_id, stage=failure_class)
await release_document_slot(session, payload.document_id)
await refund_credit(
await compensate_document_slot(
session,
payload.document_id,
failure_class,

26
src/contract_check/worker_billing/scheduler.py
View file

@ -16,6 +16,7 @@ state transitions that don't need the provider).
from __future__ import annotations
import datetime as dt
import hashlib
import uuid
from typing import TYPE_CHECKING
@ -45,6 +46,11 @@ RENEWAL_WINDOW_DAYS = 3
RECONCILE_AFTER_MINUTES = 15
PAST_DUE_GRACE_DAYS = 7
# Postgres advisory lock id for the billing scheduler tick. All replicas use
# the same key so only one scheduler runs a tick at a time.
SCHEDULER_LOCK_KEY = hashlib.sha256(b"billing_scheduler_tick").hexdigest()[:16]
SCHEDULER_LOCK_INT = int(SCHEDULER_LOCK_KEY, 16)
async def run_tick(
session: AsyncSession,
@ -52,12 +58,30 @@ async def run_tick(
now: dt.datetime,
settings: Settings,
) -> None:
"""Run one scheduler tick inside an open session/transaction."""
"""Run one scheduler tick inside an open session/transaction.
Takes a Postgres advisory lock; a second replica's tick skips when the
lock is held. The lock is released when the tick completes or raises.
"""
lock_result = await session.execute(
text("SELECT pg_try_advisory_lock(:key)"),
{"key": SCHEDULER_LOCK_INT},
)
if not lock_result.scalar_one():
log.info("scheduler_tick_lock_held_skipping")
return
try:
if provider is not None and settings.yookassa_enabled:
await _create_renewal_invoices(session, provider, now, settings)
await _expire_and_roll_subscriptions(session, now, settings)
if provider is not None:
await _reconcile_pending_invoices(session, provider, now, settings)
finally:
await session.execute(
text("SELECT pg_advisory_unlock(:key)"),
{"key": SCHEDULER_LOCK_INT},
)
async def _create_renewal_invoices(

6
src/contract_check/worker_extract/handler.py
View file

@ -16,9 +16,8 @@ import asyncio
import uuid
from typing import TYPE_CHECKING
from src.contract_check.core.billing.quota import release_document_slot
from src.contract_check.core.billing.quota import compensate_document_slot
from src.contract_check.core.config import get_settings
from src.contract_check.core.credits import refund_credit
from src.contract_check.core.db.enums import DOC_TERMINAL, FailureClass
from src.contract_check.core.db.repositories import DocumentRepository, JobRepository
from src.contract_check.core.extraction import (
@ -215,8 +214,7 @@ class ExtractHandler:
payload.document_id, "extract", failure_class=failure_class, error=error
)
await docs.mark_failed(payload.document_id, stage=failure_class)
await release_document_slot(session, payload.document_id)
await refund_credit(
await compensate_document_slot(
session,
payload.document_id,
failure_class,

7
src/contract_check/worker_prescreen/handler.py
View file

@ -19,9 +19,8 @@ import uuid
from datetime import UTC, date, datetime
from typing import TYPE_CHECKING
from src.contract_check.core.billing.quota import release_document_slot
from src.contract_check.core.billing.quota import compensate_document_slot
from src.contract_check.core.config import get_settings
from src.contract_check.core.credits import refund_credit
from src.contract_check.core.db.enums import DOC_TERMINAL, FailureClass
from src.contract_check.core.db.repositories import (
DocumentRepository,
@ -440,6 +439,7 @@ class PrescreenHandler:
is_structured=payload.is_structured,
has_tables=payload.has_tables,
prescreen_meta=completed.model_dump(),
prescreen_result_id=prescreen_result_id,
attempt=payload.attempt,
)
await publisher.publish(next_msg, routing_key=self._publish_routing_key)
@ -606,8 +606,7 @@ class PrescreenHandler:
error=error,
)
await docs.mark_failed(payload.document_id, stage=failure_class)
await release_document_slot(session, payload.document_id)
await refund_credit(
await compensate_document_slot(
session,
payload.document_id,
failure_class,

8
tests/integration/test_analyze_worker.py
View file

@ -22,7 +22,7 @@ from sqlalchemy import text
from contract_check.core.db.session import create_session_factory
from contract_check.core.llm.ollama_cloud import LLMQuotaError, OllamaCloudProvider
from contract_check.core.mq.messages import DocumentExtracted
from contract_check.core.mq.messages import AnalyzeRequested
from contract_check.core.s3 import extracted_key
from contract_check.core.s3.minio_storage import MinioStorage
from contract_check.worker_analyze.handler import AnalyzeHandler
@ -156,7 +156,7 @@ async def test_analyze_worker_saves_report_and_marks_done(
return_value=httpx.Response(200, json=_chat_body(json.dumps(_VALID_FINDINGS)))
)
await AnalyzeHandler(session_factory=sess, provider=provider).handle(
DocumentExtracted(
AnalyzeRequested(
correlation_id=correlation_id,
document_id=document_id,
user_id=user_id,
@ -214,7 +214,7 @@ async def test_analyze_worker_terminal_failure_refunds_and_dlqs(
mock.post(URL).mock(return_value=httpx.Response(429, json={"error": "quota"}))
with pytest.raises(LLMQuotaError):
await handler.handle(
DocumentExtracted(
AnalyzeRequested(
correlation_id=correlation_id,
document_id=document_id,
user_id=user_id,
@ -225,7 +225,7 @@ async def test_analyze_worker_terminal_failure_refunds_and_dlqs(
)
await handler.on_terminal_failure(
DocumentExtracted(
AnalyzeRequested(
correlation_id=correlation_id,
document_id=document_id,
user_id=user_id,

123
tests/integration/test_auth_flow.py
View file

@ -12,6 +12,7 @@ from urllib.parse import urlencode
import httpx
import pytest
from asgi_lifespan import LifespanManager
from sqlalchemy import text
pytestmark = pytest.mark.integration
@ -166,3 +167,125 @@ async def test_auth_me_introspects_jwt(
body = introspect.json()
assert body["telegram_id"] == telegram_id
assert body["type"] == "access"
async def _strict_rate_limit_client() -> httpx.AsyncClient:
"""Build a fresh ASGI client with very low auth rate limits for deterministic tests.
Uses an isolated Redis DB so concurrent/sequential tests do not share buckets.
"""
import os
import redis.asyncio as redis
from contract_check.api.app import create_app
from contract_check.core.config import get_settings
isolated_redis_url = "redis://localhost:17379/15"
r = redis.from_url(isolated_redis_url)
await r.flushdb()
await r.aclose()
os.environ["AUTH_RATE_LIMIT_IP_RPS"] = "1"
os.environ["AUTH_RATE_LIMIT_EMAIL_RPS"] = "1"
os.environ["REDIS_URL"] = isolated_redis_url
get_settings.cache_clear()
app = create_app()
client = httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test")
async with LifespanManager(app):
yield client
await client.aclose()
os.environ.pop("AUTH_RATE_LIMIT_IP_RPS", None)
os.environ.pop("AUTH_RATE_LIMIT_EMAIL_RPS", None)
os.environ.pop("REDIS_URL", None)
get_settings.cache_clear()
async def test_auth_register_rate_limited_by_ip() -> None:
import uuid
async for client in _strict_rate_limit_client():
email_ok = f"ratelimit-ok-{uuid.uuid4().hex[:8]}@example.com"
r = await client.post(
"/api/v1/auth/register",
json={
"email": email_ok,
"password": "strong-pass-123",
"name": "Rate",
},
)
assert r.status_code == 201
email_blocked = f"ratelimit-blocked-{uuid.uuid4().hex[:8]}@example.com"
r = await client.post(
"/api/v1/auth/register",
json={
"email": email_blocked,
"password": "strong-pass-123",
"name": "Rate",
},
)
assert r.status_code == 429
assert "Retry-After" in r.headers
async def test_auth_login_rate_limited_by_email() -> None:
import uuid
async for client in _strict_rate_limit_client():
email = f"rate-email-{uuid.uuid4().hex[:8]}@example.com"
payload = {"email": email, "password": "wrong"}
r = await client.post("/api/v1/auth/login", json=payload)
assert r.status_code == 401
r = await client.post("/api/v1/auth/login", json=payload)
assert r.status_code == 429
assert "Retry-After" in r.headers
async def test_auth_forgot_password_rate_limited_by_email() -> None:
import uuid
async for client in _strict_rate_limit_client():
email = f"rate-forgot-{uuid.uuid4().hex[:8]}@example.com"
payload = {"email": email}
r = await client.post("/api/v1/auth/forgot-password", json=payload)
assert r.status_code == 202
r = await client.post("/api/v1/auth/forgot-password", json=payload)
assert r.status_code == 429
assert "Retry-After" in r.headers
async def test_metrics_open_when_token_unset(client: httpx.AsyncClient) -> None:
r = await client.get("/metrics")
assert r.status_code == 200
assert "python_info" in r.text or "# TYPE" in r.text
async def test_metrics_gated_when_token_set() -> None:
import os
from asgi_lifespan import LifespanManager
from contract_check.api.app import create_app
from contract_check.core.config import get_settings
os.environ["METRICS_BEARER_TOKEN"] = "super-secret-metrics-token"
get_settings.cache_clear()
app = create_app()
client = httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test")
async with LifespanManager(app):
r = await client.get("/metrics")
assert r.status_code == 401
r = await client.get("/metrics", headers={"Authorization": "Bearer wrong"})
assert r.status_code == 401
r = await client.get(
"/metrics", headers={"Authorization": "Bearer super-secret-metrics-token"}
)
assert r.status_code == 200
await client.aclose()
os.environ.pop("METRICS_BEARER_TOKEN", None)
get_settings.cache_clear()

54
tests/integration/test_b2b_api.py
View file

@ -178,6 +178,60 @@ async def test_b2b_key_management_requires_user_jwt(
assert response.status_code == 401
async def test_b2b_key_usage_aggregated_by_month(
api_key_client: tuple[httpx.AsyncClient, str, int],
infra: dict[str, str],
db_session, # noqa: ANN001
) -> None:
from sqlalchemy import text
client, api_key, telegram_id = api_key_client
token = await user_token(client, infra, telegram_id)
list_resp = await client.get(
"/api/v1/b2b/keys",
headers={"Authorization": f"Bearer {token}"},
)
assert list_resp.status_code == 200
key_id = uuid.UUID(list_resp.json()[0]["id"])
# Seed two requests last month and one this month, each linked to a document.
result = await db_session.execute(
text("SELECT id FROM users WHERE telegram_id = :t"), {"t": telegram_id}
)
user_id = result.scalar_one()
doc_ids = [uuid.uuid4() for _ in range(3)]
await db_session.execute(
text(
"INSERT INTO documents (id, user_id, s3_key, filename, mime, bytes, status) "
"VALUES (:id1, :u, 's3/1', '1.pdf', 'application/pdf', 1, 'done'), "
"(:id2, :u, 's3/2', '2.pdf', 'application/pdf', 1, 'done'), "
"(:id3, :u, 's3/3', '3.pdf', 'application/pdf', 1, 'done')"
),
{"id1": doc_ids[0], "id2": doc_ids[1], "id3": doc_ids[2], "u": user_id},
)
await db_session.execute(
text(
"INSERT INTO api_key_requests (api_key_id, document_id, created_at) "
"VALUES (:k, :d1, now() - interval '1 month'), "
"(:k, :d2, now() - interval '1 month'), "
"(:k, :d3, now())"
),
{"k": key_id, "d1": doc_ids[0], "d2": doc_ids[1], "d3": doc_ids[2]},
)
await db_session.commit()
response = await client.get(
f"/api/v1/b2b/keys/{key_id}/usage",
headers={"Authorization": f"Bearer {token}"},
)
assert response.status_code == 200
body = response.json()
assert len(body["monthly_requests"]) == 2
counts = sorted(entry["requests"] for entry in body["monthly_requests"])
assert counts == [1, 2]
async def test_b2b_revoke_key_blocks_usage(
api_key_client: tuple[httpx.AsyncClient, str, int],
infra: dict[str, str],

40
tests/integration/test_prescreen_worker.py
View file

@ -281,3 +281,43 @@ async def test_prescreen_llm_failure_records_error_not_fatal(
status, stage = doc.one()
assert status == "manual_review"
assert stage == "manual_review"
async def test_prescreen_redelivery_for_manual_review_is_no_op(
infra: dict[str, str],
) -> None:
sess = create_session_factory()
document_id = uuid.uuid4()
user_id, ext_key = await _seed_document(
infra, telegram_id=777_555_667, document_id=document_id, contract_text=LOW_CONF_TEXT
)
handler = PrescreenHandler(
session_factory=sess,
publish_routing_key=_DEAD_ROUTING_KEY,
extractor=HybridMetaExtractor(
provider=StubProvider(RuntimeError("quota exceeded")),
fallback_enabled=True,
fallback_threshold=0.99,
heuristic=HeuristicExtractor(),
),
)
await handler.handle(_requested(document_id, user_id, ext_key, LOW_CONF_TEXT))
# Force-update a field so we can detect a re-run.
async with sess() as session:
await session.execute(
text("UPDATE prescreen_results SET confidence_score = 0.123 WHERE document_id = :d"),
{"d": document_id},
)
await session.commit()
# Redeliver the same prescreen message: it should be a no-op.
await handler.handle(_requested(document_id, user_id, ext_key, LOW_CONF_TEXT))
async with sess() as session:
result = await session.execute(
text("SELECT confidence_score FROM prescreen_results WHERE document_id = :d"),
{"d": document_id},
)
assert float(result.scalar_one()) == pytest.approx(0.123)

66
tests/integration/test_upload_pipeline.py
View file

@ -18,7 +18,9 @@ from typing import TYPE_CHECKING
import aio_pika
import httpx
import pytest
from asgi_lifespan import LifespanManager
from contract_check.api.app import create_app
from tests.integration.conftest import user_token
if TYPE_CHECKING:
@ -129,3 +131,67 @@ async def test_upload_document_reserves_credit_and_enqueues(
assert found, "expected pipeline message not found on extract.q or analyze.q"
finally:
await connection.close()
async def test_upload_rejected_before_storage_on_billing_hold(
client: httpx.AsyncClient,
infra: dict[str, str],
pdf_bytes: bytes,
db_session: AsyncSession,
) -> None:
from sqlalchemy import text
telegram_id = 123456790
await db_session.execute(
text(
"INSERT INTO users (telegram_id, credits_left, billing_hold) VALUES (:t, 5, TRUE) "
"ON CONFLICT (telegram_id) DO UPDATE SET credits_left = 5, billing_hold = TRUE"
),
{"t": telegram_id},
)
await db_session.commit()
token = await user_token(client, infra, telegram_id)
response = await client.post(
"/api/v1/documents",
headers={"Authorization": f"Bearer {token}"},
files={"file": ("contract.pdf", pdf_bytes, "application/pdf")},
)
assert response.status_code == 402
assert "billing hold" in response.text.lower()
async def test_upload_rejects_oversized_file(
infra: dict[str, str],
db_session: AsyncSession,
monkeypatch: pytest.MonkeyPatch,
) -> None:
from sqlalchemy import text
telegram_id = 123456791
await db_session.execute(
text(
"INSERT INTO users (telegram_id, credits_left) VALUES (:t, 5) "
"ON CONFLICT (telegram_id) DO UPDATE SET credits_left = 5"
),
{"t": telegram_id},
)
await db_session.commit()
monkeypatch.setenv("MAX_UPLOAD_BYTES", "10")
from contract_check.core.config import get_settings
get_settings.cache_clear()
app = create_app()
async with httpx.AsyncClient(
transport=httpx.ASGITransport(app=app), base_url="http://test"
) as client:
async with LifespanManager(app):
token = await user_token(client, infra, telegram_id)
response = await client.post(
"/api/v1/documents",
headers={"Authorization": f"Bearer {token}"},
files={"file": ("big.pdf", b"x" * 100, "application/pdf")},
)
assert response.status_code == 413

142
tests/unit/test_compensation.py Normal file
View file

@ -0,0 +1,142 @@
"""Unit tests for Document Slot compensation (ticket 001)."""
from __future__ import annotations
import uuid
import pytest
from sqlalchemy import text
from contract_check.core.billing.quota import compensate_document_slot, reserve_document_slot
pytestmark = pytest.mark.unit
async def _seed_user(session, *, telegram_id: int, credits: int = 0) -> uuid.UUID:
result = await session.execute(
text("INSERT INTO users (telegram_id, credits_left) VALUES (:t, :c) RETURNING id"),
{"t": telegram_id, "c": credits},
)
await session.flush()
return result.scalar_one()
async def _seed_plan_and_subscription(session, user_id: uuid.UUID) -> uuid.UUID:
result = await session.execute(
text(
"INSERT INTO plans (code, name, price_kopecks, monthly_quota, is_active, sort) "
"VALUES ('basic', 'Basic', 9900, 5, TRUE, 1) RETURNING code"
)
)
plan_code = result.scalar_one()
result = await session.execute(
text(
"INSERT INTO subscriptions "
"(id, user_id, plan_code, status, current_period_start, current_period_end) "
"VALUES (gen_random_uuid(), :u, :p, 'active', now() - interval '1 day', now() + interval '30 days') "
"RETURNING id"
),
{"u": user_id, "p": plan_code},
)
await session.flush()
return result.scalar_one()
async def _seed_document(session, user_id: uuid.UUID) -> uuid.UUID:
result = await session.execute(
text(
"INSERT INTO documents (user_id, s3_key, filename, mime, bytes, status) "
"VALUES (:u, 's3/key', 'file.pdf', 'application/pdf', 100, 'queued') "
"RETURNING id"
),
{"u": user_id},
)
await session.flush()
return result.scalar_one()
async def test_compensate_quota_paid_restores_quota_not_credits(db_session) -> None:
user_id = await _seed_user(db_session, telegram_id=1001, credits=3)
sub_id = await _seed_plan_and_subscription(db_session, user_id)
document_id = await _seed_document(db_session, user_id)
source = await reserve_document_slot(db_session, user_id, document_id)
assert source == "quota"
quota_released, credit_refunded = await compensate_document_slot(
db_session, document_id, "llm_quota", "all"
)
assert quota_released is True
assert credit_refunded is False
row = await db_session.execute(
text("SELECT refunded FROM documents WHERE id = :d"), {"d": document_id}
)
assert row.scalar_one() is True
credits = await db_session.execute(
text("SELECT credits_left FROM users WHERE id = :u"), {"u": user_id}
)
assert credits.scalar_one() == 3
quota = await db_session.execute(
text("SELECT count(*) FROM quota_usage WHERE subscription_id = :s"),
{"s": sub_id},
)
assert quota.scalar_one() == 0
async def test_compensate_credits_paid_refunds_once(db_session) -> None:
user_id = await _seed_user(db_session, telegram_id=1002, credits=3)
document_id = await _seed_document(db_session, user_id)
source = await reserve_document_slot(db_session, user_id, document_id)
assert source == "credits"
quota_released, credit_refunded = await compensate_document_slot(
db_session, document_id, "llm_quota", "all"
)
assert quota_released is False
assert credit_refunded is True
row = await db_session.execute(
text(
"SELECT credits_left, refunded FROM users u JOIN documents d ON d.user_id = u.id "
"WHERE d.id = :d"
),
{"d": document_id},
)
credits_left, refunded = row.one()
assert credits_left == 3
assert refunded is True
# Second execution must be a no-op.
quota_released2, credit_refunded2 = await compensate_document_slot(
db_session, document_id, "llm_quota", "all"
)
assert quota_released2 is False
assert credit_refunded2 is False
credits2 = await db_session.execute(
text("SELECT credits_left FROM users WHERE id = :u"), {"u": user_id}
)
assert credits2.scalar_one() == 3
async def test_compensate_infra_only_policy_skips_extraction_failed(db_session) -> None:
user_id = await _seed_user(db_session, telegram_id=1003, credits=3)
document_id = await _seed_document(db_session, user_id)
source = await reserve_document_slot(db_session, user_id, document_id)
assert source == "credits"
quota_released, credit_refunded = await compensate_document_slot(
db_session, document_id, "extraction_failed", "infra_only"
)
assert quota_released is False
assert credit_refunded is False
row = await db_session.execute(
text("SELECT refunded FROM documents WHERE id = :d"), {"d": document_id}
)
assert row.scalar_one() is False

152
tests/unit/test_consumer_hook_guard.py Normal file
View file

@ -0,0 +1,152 @@
"""Unit tests for the RabbitMQ consumer base hook crash guard (ticket 007)."""
from __future__ import annotations
import uuid
from typing import Any
import pytest
from aio_pika import Message
from contract_check.core.db.enums import FailureClass
from contract_check.core.mq.consumer import Consumer
from contract_check.core.mq.messages import PipelineMessage
from contract_check.core.mq.topology import H_ATTEMPT
pytestmark = pytest.mark.unit
class _FakeChannel:
def __init__(self) -> None:
self.published: list[tuple[Message, str]] = []
self.acks: list[Message] = []
async def get_exchange(self, name: str, ensure: bool = True) -> _FakeExchange:
return _FakeExchange(self)
@property
def default_exchange(self) -> _FakeExchange:
return _FakeExchange(self)
class _FakeExchange:
def __init__(self, channel: _FakeChannel) -> None:
self._channel = channel
async def publish(self, message: Message, routing_key: str) -> None:
self._channel.published.append((message, routing_key))
class _FakeIncomingMessage:
def __init__(self, body: bytes, headers: dict[str, Any] | None = None) -> None:
self.body = body
self.headers = headers or {}
self.correlation_id = str(uuid.uuid4())
self.content_type = "application/json"
self._ack = False
async def ack(self) -> None:
self._ack = True
class _TestMessage(PipelineMessage):
pass
class _CrashHookConsumer(Consumer[_TestMessage]):
queue = "test.q"
routing_key = "test"
message_model = _TestMessage
retry_exchange = "retry.x"
def __init__(self) -> None:
super().__init__(
"amqp://test",
origin="test",
prefetch=1,
max_attempts=3,
retry_base_ms=1000,
)
self.channel = _FakeChannel()
self._channel = self.channel
self.failure_calls: list[tuple[Any, ...]] = []
self.dlq_calls: list[tuple[Any, ...]] = []
self.raise_on_failure = True
self.raise_on_dlq = True
async def handle(self, payload: _TestMessage) -> None:
raise RuntimeError("handler failed")
def classify(self, exc: BaseException) -> FailureClass:
return "infra"
async def on_failure(
self,
payload: _TestMessage,
failure_class: FailureClass,
attempt: int,
error: str,
) -> None:
self.failure_calls.append((payload, failure_class, attempt, error))
if self.raise_on_failure:
raise RuntimeError("on_failure crashed")
async def on_dlq(self, payload: _TestMessage, failure_class: FailureClass, error: str) -> None:
self.dlq_calls.append((payload, failure_class, error))
if self.raise_on_dlq:
raise RuntimeError("on_dlq crashed")
async def test_on_failure_crash_republishes_to_retry_with_unchanged_attempt() -> None:
consumer = _CrashHookConsumer()
consumer.raise_on_failure = True
consumer.raise_on_dlq = False
payload = _TestMessage(
correlation_id=uuid.uuid4(),
document_id=uuid.uuid4(),
user_id=uuid.uuid4(),
attempt=2,
)
msg = _FakeIncomingMessage(
payload.model_dump_json().encode("utf-8"),
headers={H_ATTEMPT: 2},
)
await consumer._on_message(msg)
assert len(consumer.channel.published) == 1
published, routing_key = consumer.channel.published[0]
assert routing_key == "retry.test"
assert published.headers["x-attempt"] == 2 # unchanged
assert msg._ack is True
async def test_on_dlq_crash_republishes_to_retry_with_unchanged_attempt() -> None:
from contract_check.core.errors import TerminalError
class _TerminalConsumer(_CrashHookConsumer):
async def handle(self, payload: _TestMessage) -> None:
raise TerminalError("terminal")
consumer = _TerminalConsumer()
consumer.raise_on_failure = False
consumer.raise_on_dlq = True
payload = _TestMessage(
correlation_id=uuid.uuid4(),
document_id=uuid.uuid4(),
user_id=uuid.uuid4(),
attempt=1,
)
msg = _FakeIncomingMessage(
payload.model_dump_json().encode("utf-8"),
headers={H_ATTEMPT: 1},
)
await consumer._on_message(msg)
assert len(consumer.channel.published) == 1
published, routing_key = consumer.channel.published[0]
assert routing_key == "retry.test"
assert published.headers["x-attempt"] == 1 # unchanged
assert msg._ack is True
assert len(consumer.dlq_calls) == 1

182
tests/unit/test_worker_billing.py
View file

@ -6,6 +6,7 @@ state transitions without touching the network.
from __future__ import annotations
import asyncio
import datetime as dt
import uuid
from dataclasses import dataclass
@ -73,6 +74,9 @@ class FakeProvider(PaymentProvider):
async def factory():
import os
os.environ["DATABASE_URL"] = (
"postgresql+asyncpg://contract_check:contract_check@localhost:15432/contract_check"
)
os.environ["PLANS_ENABLED"] = "true"
os.environ["YOOKASSA_ENABLED"] = "true"
os.environ["YOOKASSA_SHOP_ID"] = "123456"
@ -81,6 +85,7 @@ async def factory():
get_settings.cache_clear()
f = create_session_factory()
yield f
os.environ.pop("DATABASE_URL", None)
os.environ.pop("PLANS_ENABLED", None)
os.environ.pop("YOOKASSA_ENABLED", None)
os.environ.pop("YOOKASSA_SHOP_ID", None)
@ -276,3 +281,180 @@ async def test_scheduler_reconciles_pending_invoice(factory) -> None:
)
).scalar_one()
assert status == "succeeded"
async def test_scheduler_tick_advisory_lock_prevents_double_run(factory) -> None:
provider = FakeProvider(payments={})
async with factory() as session:
user_id = await _seed_user(session)
sub_id = await _seed_subscription(
session, user_id, auto_renew=True, period_end_offset_days=2
)
await session.commit()
now = dt.datetime.now(tz=dt.UTC)
settings = get_settings()
# First tick holds the advisory lock until it commits.
task1 = asyncio.create_task(run_tick(session, provider, now, settings))
# Give task1 time to acquire the lock.
await asyncio.sleep(0.1)
# Second tick on a separate connection should skip.
async with factory() as session2:
await run_tick(session2, provider, now, settings)
await task1
count = (
await session.execute(
text("SELECT count(*) FROM invoices WHERE subscription_id = :s"),
{"s": sub_id},
)
).scalar_one()
assert int(count) == 1
async def test_subscription_fulfillment_idempotent_by_invoice(factory) -> None:
from contract_check.core.billing.fulfillment import apply_payment_status
provider = FakeProvider(payments={})
async with factory() as session:
user_id = await _seed_user(session, telegram_id=123_456_700)
# Ensure no leftover subscription for this user from an interrupted run.
await session.execute(text("DELETE FROM subscriptions WHERE user_id = :u"), {"u": user_id})
plan_code = "pro"
await session.execute(
text(
"INSERT INTO plans (code, name, price_kopecks, monthly_quota, is_active, sort) "
"VALUES (:code, :name, 149000, 10, TRUE, 1) "
"ON CONFLICT (code) DO NOTHING"
),
{"code": plan_code, "name": "Pro"},
)
inv_id = uuid.uuid4()
await session.execute(
text(
"INSERT INTO invoices (id, user_id, amount, status, kind, external_id) "
"VALUES (:id, :u, 149000, 'pending', 'subscription', 'pay-sub')"
),
{"id": inv_id, "u": user_id},
)
provider.payments["pay-sub"] = {
"amount_kopecks": 149000,
"status": "succeeded",
"metadata": {"invoice_id": str(inv_id), "kind": "subscription", "plan_code": plan_code},
}
await session.commit()
await apply_payment_status(session, provider, "pay-sub")
await apply_payment_status(session, provider, "pay-sub")
subs = await session.execute(
text("SELECT count(*) FROM subscriptions WHERE user_id = :u"), {"u": user_id}
)
assert int(subs.scalar_one()) == 1
inv_status = await session.execute(
text("SELECT status FROM invoices WHERE id = :id"), {"id": inv_id}
)
assert inv_status.scalar_one() == "succeeded"
async def test_subscription_fulfillment_poison_on_conflicting_subscription(factory) -> None:
from contract_check.core.billing.fulfillment import apply_payment_status
provider = FakeProvider(payments={})
async with factory() as session:
user_id = await _seed_user(session, telegram_id=123_456_701)
await session.execute(text("DELETE FROM subscriptions WHERE user_id = :u"), {"u": user_id})
plan_code = "pro"
await session.execute(
text(
"INSERT INTO plans (code, name, price_kopecks, monthly_quota, is_active, sort) "
"VALUES (:code, :name, 149000, 10, TRUE, 1) "
"ON CONFLICT (code) DO NOTHING"
),
{"code": plan_code, "name": "Pro"},
)
# Pre-existing active subscription from a different invoice.
other_inv_id = uuid.uuid4()
await session.execute(
text(
"INSERT INTO invoices (id, user_id, amount, status, kind) "
"VALUES (:id, :u, 149000, 'succeeded', 'subscription')"
),
{"id": other_inv_id, "u": user_id},
)
await session.execute(
text(
"INSERT INTO subscriptions "
"(id, user_id, plan_code, status, current_period_start, current_period_end, origin_invoice_id) "
"VALUES (gen_random_uuid(), :u, :plan, 'active', now(), now() + interval '30 days', :inv)"
),
{"u": user_id, "plan": plan_code, "inv": other_inv_id},
)
inv_id = uuid.uuid4()
await session.execute(
text(
"INSERT INTO invoices (id, user_id, amount, status, kind, external_id) "
"VALUES (:id, :u, 149000, 'pending', 'subscription', 'pay-sub2')"
),
{"id": inv_id, "u": user_id},
)
provider.payments["pay-sub2"] = {
"amount_kopecks": 149000,
"status": "succeeded",
"metadata": {"invoice_id": str(inv_id), "kind": "subscription", "plan_code": plan_code},
}
await session.commit()
await apply_payment_status(session, provider, "pay-sub2")
inv_status = await session.execute(
text("SELECT status FROM invoices WHERE id = :id"), {"id": inv_id}
)
assert inv_status.scalar_one() == "cancelled"
async def test_credit_refund_is_single_fire_under_concurrency(factory) -> None:
from contract_check.core.credits import refund_credit
async with factory() as session:
user_id = await _seed_user(session, telegram_id=123_456_702)
await session.execute(
text("UPDATE users SET credits_left = 0 WHERE id = :u"), {"u": user_id}
)
doc_id = uuid.uuid4()
await session.execute(
text(
"INSERT INTO documents "
"(id, user_id, s3_key, filename, mime, bytes, status, stage, refunded) "
"VALUES (:id, :u, 's3', 'f.pdf', 'application/pdf', 100, 'failed', 'failed', FALSE)"
),
{"id": doc_id, "u": user_id},
)
await session.commit()
async def _refund() -> bool:
async with factory() as s:
result = await refund_credit(s, doc_id, "llm_quota", "all")
await s.commit()
return result
results = await asyncio.gather(_refund(), _refund(), _refund())
assert sum(1 for r in results if r) == 1
async with factory() as session:
balance = await session.execute(
text("SELECT credits_left FROM users WHERE id = :u"), {"u": user_id}
)
assert balance.scalar_one() == 1
events = await session.execute(
text(
"SELECT count(*) FROM credit_events "
"WHERE user_id = :u AND kind = 'refund_auto' AND document_id = :d"
),
{"u": user_id, "d": doc_id},
)
assert events.scalar_one() == 1