From e5fd2853284a912ee476d6c6f6086a4d1701416a Mon Sep 17 00:00:00 2001 From: febux Date: Sat, 5 Sep 2026 23:18:19 +0300 Subject: [PATCH] Fix issues hardening by plan. M1: .scratch/pipeline-hardening/issues/001 - 011 --- .env.example | 6 + docs/ARCHITECTURE.md | 3 + src/contract_check/api/admin/router.py | 9 +- src/contract_check/api/app.py | 4 +- src/contract_check/api/deps.py | 57 +++++- .../api/routes/auth/magic_links.py | 10 +- .../api/routes/auth/password.py | 33 +++- src/contract_check/api/routes/b2b.py | 13 +- src/contract_check/api/routes/metrics.py | 16 +- src/contract_check/api/schemas/__init__.py | 2 + src/contract_check/api/schemas/b2b.py | 3 +- src/contract_check/api/services.py | 65 +++++-- .../core/billing/fulfillment.py | 15 +- src/contract_check/core/billing/quota.py | 32 +++ src/contract_check/core/config.py | 18 ++ src/contract_check/core/db/enums.py | 2 +- .../core/db/repositories/api_keys.py | 18 ++ .../core/db/repositories/credits.py | 37 ++-- .../core/db/repositories/reports.py | 12 +- .../core/db/repositories/subscriptions.py | 20 +- src/contract_check/core/metrics.py | 4 + src/contract_check/core/mq/consumer.py | 39 +++- src/contract_check/core/mq/messages.py | 7 +- src/contract_check/core/security/passwords.py | 15 ++ src/contract_check/worker_analyze/consumer.py | 14 +- src/contract_check/worker_analyze/handler.py | 18 +- .../worker_billing/scheduler.py | 36 +++- src/contract_check/worker_extract/handler.py | 6 +- .../worker_prescreen/handler.py | 7 +- tests/integration/test_analyze_worker.py | 8 +- tests/integration/test_auth_flow.py | 123 ++++++++++++ tests/integration/test_b2b_api.py | 54 ++++++ tests/integration/test_prescreen_worker.py | 40 ++++ tests/integration/test_upload_pipeline.py | 66 +++++++ tests/unit/test_compensation.py | 142 ++++++++++++++ tests/unit/test_consumer_hook_guard.py | 152 +++++++++++++++ tests/unit/test_worker_billing.py | 182 ++++++++++++++++++ 37 files changed, 1176 insertions(+), 112 deletions(-) create mode 100644 tests/unit/test_compensation.py create mode 100644 tests/unit/test_consumer_hook_guard.py diff --git a/.env.example b/.env.example index b91d126..42c1b42 100644 --- a/.env.example +++ b/.env.example @@ -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) diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 2f9aa93..63ea2db 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -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) | diff --git a/src/contract_check/api/admin/router.py b/src/contract_check/api/admin/router.py index 76851ec..5cd1e09 100644 --- a/src/contract_check/api/admin/router.py +++ b/src/contract_check/api/admin/router.py @@ -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( diff --git a/src/contract_check/api/app.py b/src/contract_check/api/app.py index 363a731..806bf79 100644 --- a/src/contract_check/api/app.py +++ b/src/contract_check/api/app.py @@ -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() diff --git a/src/contract_check/api/deps.py b/src/contract_check/api/deps.py index 70afe21..bd4d9ee 100644 --- a/src/contract_check/api/deps.py +++ b/src/contract_check/api/deps.py @@ -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) diff --git a/src/contract_check/api/routes/auth/magic_links.py b/src/contract_check/api/routes/auth/magic_links.py index 542edb5..5a1d536 100644 --- a/src/contract_check/api/routes/auth/magic_links.py +++ b/src/contract_check/api/routes/auth/magic_links.py @@ -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) diff --git a/src/contract_check/api/routes/auth/password.py b/src/contract_check/api/routes/auth/password.py index 1ac4c27..da07525 100644 --- a/src/contract_check/api/routes/auth/password.py +++ b/src/contract_check/api/routes/auth/password.py @@ -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 diff --git a/src/contract_check/api/routes/b2b.py b/src/contract_check/api/routes/b2b.py index 9daaa55..5f7c903 100644 --- a/src/contract_check/api/routes/b2b.py +++ b/src/contract_check/api/routes/b2b.py @@ -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 ], ) diff --git a/src/contract_check/api/routes/metrics.py b/src/contract_check/api/routes/metrics.py index f1a33eb..7785954 100644 --- a/src/contract_check/api/routes/metrics.py +++ b/src/contract_check/api/routes/metrics.py @@ -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) diff --git a/src/contract_check/api/schemas/__init__.py b/src/contract_check/api/schemas/__init__.py index 17130ac..5a49f67 100644 --- a/src/contract_check/api/schemas/__init__.py +++ b/src/contract_check/api/schemas/__init__.py @@ -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", ] diff --git a/src/contract_check/api/schemas/b2b.py b/src/contract_check/api/schemas/b2b.py index 7e7d0c2..f4146fc 100644 --- a/src/contract_check/api/schemas/b2b.py +++ b/src/contract_check/api/schemas/b2b.py @@ -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): diff --git a/src/contract_check/api/services.py b/src/contract_check/api/services.py index 0fceae9..b483b71 100644 --- a/src/contract_check/api/services.py +++ b/src/contract_check/api/services.py @@ -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) + data = await _read_upload_with_limit(file, settings.max_upload_bytes) + if len(data) == 0: + raise HTTPException(status_code=400, detail="empty file") + try: - data = await file.read() - if len(data) == 0: - raise HTTPException(status_code=400, detail="empty file") 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) diff --git a/src/contract_check/core/billing/fulfillment.py b/src/contract_check/core/billing/fulfillment.py index a260f63..929c970 100644 --- a/src/contract_check/core/billing/fulfillment.py +++ b/src/contract_check/core/billing/fulfillment.py @@ -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), diff --git a/src/contract_check/core/billing/quota.py b/src/contract_check/core/billing/quota.py index 79fc509..574cb59 100644 --- a/src/contract_check/core/billing/quota.py +++ b/src/contract_check/core/billing/quota.py @@ -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 diff --git a/src/contract_check/core/config.py b/src/contract_check/core/config.py index 54e3023..d56d2eb 100644 --- a/src/contract_check/core/config.py +++ b/src/contract_check/core/config.py @@ -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, diff --git a/src/contract_check/core/db/enums.py b/src/contract_check/core/db/enums.py index 029ac2f..ddd04a3 100644 --- a/src/contract_check/core/db/enums.py +++ b/src/contract_check/core/db/enums.py @@ -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"] diff --git a/src/contract_check/core/db/repositories/api_keys.py b/src/contract_check/core/db/repositories/api_keys.py index ab3a78f..7efd252 100644 --- a/src/contract_check/core/db/repositories/api_keys.py +++ b/src/contract_check/core/db/repositories/api_keys.py @@ -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()] diff --git a/src/contract_check/core/db/repositories/credits.py b/src/contract_check/core/db/repositories/credits.py index 6aba554..e9d8e89 100644 --- a/src/contract_check/core/db/repositories/credits.py +++ b/src/contract_check/core/db/repositories/credits.py @@ -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( diff --git a/src/contract_check/core/db/repositories/reports.py b/src/contract_check/core/db/repositories/reports.py index 997c917..d3a0f77 100644 --- a/src/contract_check/core/db/repositories/reports.py +++ b/src/contract_check/core/db/repositories/reports.py @@ -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() diff --git a/src/contract_check/core/db/repositories/subscriptions.py b/src/contract_check/core/db/repositories/subscriptions.py index 7f095ca..cbfee8f 100644 --- a/src/contract_check/core/db/repositories/subscriptions.py +++ b/src/contract_check/core/db/repositories/subscriptions.py @@ -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( diff --git a/src/contract_check/core/metrics.py b/src/contract_check/core/metrics.py index 6111119..f88ddd6 100644 --- a/src/contract_check/core/metrics.py +++ b/src/contract_check/core/metrics.py @@ -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.", diff --git a/src/contract_check/core/mq/consumer.py b/src/contract_check/core/mq/consumer.py index 918ce18..1737703 100644 --- a/src/contract_check/core/mq/consumer.py +++ b/src/contract_check/core/mq/consumer.py @@ -192,7 +192,18 @@ class Consumer[MsgT: RetryableMessage]: failure_class = self.classify(exc) error = f"{type(exc).__name__}: {exc}" new_attempt = attempt + 1 - await self.on_failure(payload, failure_class, new_attempt, error) + 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, ) - await self.on_dlq(payload, failure_class, error) + 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, ) - await self.on_dlq(payload, failure_class, error) + 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) diff --git a/src/contract_check/core/mq/messages.py b/src/contract_check/core/mq/messages.py index 3cf554f..1626eb1 100644 --- a/src/contract_check/core/mq/messages.py +++ b/src/contract_check/core/mq/messages.py @@ -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): diff --git a/src/contract_check/core/security/passwords.py b/src/contract_check/core/security/passwords.py index 801562f..6ac834d 100644 --- a/src/contract_check/core/security/passwords.py +++ b/src/contract_check/core/security/passwords.py @@ -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) diff --git a/src/contract_check/worker_analyze/consumer.py b/src/contract_check/worker_analyze/consumer.py index b8ef621..80ae618 100644 --- a/src/contract_check/worker_analyze/consumer.py +++ b/src/contract_check/worker_analyze/consumer.py @@ -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) diff --git a/src/contract_check/worker_analyze/handler.py b/src/contract_check/worker_analyze/handler.py index 13f0913..ff1d052 100644 --- a/src/contract_check/worker_analyze/handler.py +++ b/src/contract_check/worker_analyze/handler.py @@ -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, diff --git a/src/contract_check/worker_billing/scheduler.py b/src/contract_check/worker_billing/scheduler.py index 21eb637..afd8700 100644 --- a/src/contract_check/worker_billing/scheduler.py +++ b/src/contract_check/worker_billing/scheduler.py @@ -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.""" - 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) + """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( diff --git a/src/contract_check/worker_extract/handler.py b/src/contract_check/worker_extract/handler.py index e2a53bb..9d63edd 100644 --- a/src/contract_check/worker_extract/handler.py +++ b/src/contract_check/worker_extract/handler.py @@ -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, diff --git a/src/contract_check/worker_prescreen/handler.py b/src/contract_check/worker_prescreen/handler.py index 428a3ba..ab9d5fe 100644 --- a/src/contract_check/worker_prescreen/handler.py +++ b/src/contract_check/worker_prescreen/handler.py @@ -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, diff --git a/tests/integration/test_analyze_worker.py b/tests/integration/test_analyze_worker.py index af8551d..6b3af81 100644 --- a/tests/integration/test_analyze_worker.py +++ b/tests/integration/test_analyze_worker.py @@ -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, diff --git a/tests/integration/test_auth_flow.py b/tests/integration/test_auth_flow.py index a69cc5a..fd03ec2 100644 --- a/tests/integration/test_auth_flow.py +++ b/tests/integration/test_auth_flow.py @@ -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() diff --git a/tests/integration/test_b2b_api.py b/tests/integration/test_b2b_api.py index 14153a3..8dfb0a2 100644 --- a/tests/integration/test_b2b_api.py +++ b/tests/integration/test_b2b_api.py @@ -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], diff --git a/tests/integration/test_prescreen_worker.py b/tests/integration/test_prescreen_worker.py index ad41e39..6b2f67d 100644 --- a/tests/integration/test_prescreen_worker.py +++ b/tests/integration/test_prescreen_worker.py @@ -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) diff --git a/tests/integration/test_upload_pipeline.py b/tests/integration/test_upload_pipeline.py index 9ee604f..43309af 100644 --- a/tests/integration/test_upload_pipeline.py +++ b/tests/integration/test_upload_pipeline.py @@ -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 diff --git a/tests/unit/test_compensation.py b/tests/unit/test_compensation.py new file mode 100644 index 0000000..4a45060 --- /dev/null +++ b/tests/unit/test_compensation.py @@ -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 diff --git a/tests/unit/test_consumer_hook_guard.py b/tests/unit/test_consumer_hook_guard.py new file mode 100644 index 0000000..bc9dee6 --- /dev/null +++ b/tests/unit/test_consumer_hook_guard.py @@ -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 diff --git a/tests/unit/test_worker_billing.py b/tests/unit/test_worker_billing.py index 95296eb..cba19e5 100644 --- a/tests/unit/test_worker_billing.py +++ b/tests/unit/test_worker_billing.py @@ -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