"""Idempotency-Key support for state-changing endpoints. Usage (opt-in per-endpoint): from fastapi import Header from app.core.idempotency import IdempotencyService @router.post("/assets/transfer", response_model=TransferResponse) async def transfer( payload: TransferRequest, idempotency_key: str | None = Header(None, alias="Idempotency-Key"), db: AsyncSession = Depends(get_db), user: AuthUser = Depends(get_current_user), ): idem = IdempotencyService(db, scope="asset.transfer", key=idempotency_key) cached = await idem.lookup() if cached is not None: return cached # ... do the actual work ... transfer_row = await do_transfer(...) result = TransferResponse(...) await idem.record(result, resource_id=transfer_row.id) return result Semantics (RFC idempotency-key-header aligned): * First call with a fresh key → reserves a ``pending`` row, returns ``None`` so the handler proceeds normally. On success the handler calls ``record()`` which flips the row to ``completed`` and stores the response. * Replay with the same key + scope after success → ``lookup()`` returns the cached response dict; the handler short-circuits and returns it. * Replay while an earlier call is still running → ``lookup()`` raises ``ConflictError`` ("request still in progress"). The client should retry later with the same key. * Replay after a failed earlier call → the stale ``failed`` row is removed and a fresh ``pending`` one is reserved; the handler runs again. * No key in the header → ``lookup()`` returns ``None``, ``record()`` is a no-op. Idempotency is strictly opt-in by the client. Keys are namespaced by ``scope`` (e.g. ``"asset.transfer"``) so the same client key can be reused across endpoints without collision. A background job (``core_idempotency_cleanup`` scheduled handler) deletes expired rows hourly. TTL defaults to 24 hours. """ from __future__ import annotations import json import logging import uuid from datetime import UTC, datetime, timedelta from typing import Any from fastapi.encoders import jsonable_encoder from sqlalchemy import DateTime, Integer, String, Text, UniqueConstraint, select from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import Mapped, mapped_column from sqlalchemy.sql import func from app.core.database import Base from app.core.exceptions import ( IdempotencyInProgressError, IdempotencyPayloadMismatchError, ) from app.core.sql_types import GUID as Uuid DEFAULT_TTL_HOURS = 24 # How long a ``pending`` reservation may sit before a retry is allowed to take # it over. This is the backstop for the ONE failure mode the explicit # ``mark_failed`` path cannot cover: the process dying mid-handler (OOM, a # deploy replacing the container, a hard kill). Without it such a key answers # 409 for the full 24h TTL — the user cannot retry their own write, and no # amount of care in the handler helps because no code ran. # # 10 minutes, not 10 seconds: taking over a reservation whose first call is # genuinely still running would double-execute the exact operations this # guard protects. No HTTP handler here outlives the proxy timeouts (60-120s) # by an order of magnitude, so a 10-minute-old ``pending`` row means the # request is gone, not slow. IN_PROGRESS_TAKEOVER_MINUTES = 10 MAX_RESPONSE_BYTES = 256 * 1024 # 256 KiB — bigger responses aren't cached STATE_PENDING = "pending" STATE_COMPLETED = "completed" STATE_FAILED = "failed" class IdempotencyRecord(Base): """One ``Idempotency-Key`` reservation, scoped by endpoint family. A ``(key, scope)`` pair maps to exactly one operation. The response is stored as JSON so replays return the exact bytes the first call produced, including any generated UUIDs or timestamps — this is what makes the replay "idempotent from the client's point of view". Rows are never updated in place after ``completed``; a failed row gets deleted and a fresh reservation is made (see service code). """ __tablename__ = "core_idempotency_records" __table_args__ = ( UniqueConstraint("key", "scope", name="uq_idempotency_key_scope"), ) __mapper_args__ = {"eager_defaults": True} id: Mapped[uuid.UUID] = mapped_column(Uuid, primary_key=True, default=uuid.uuid4) key: Mapped[str] = mapped_column(String(128), nullable=False) # 128 (matches ``key``) so any ``f"{audit_action}:{uuid}"`` scope # fits — see migration 268. ``leaver_blocked.force_bypass_grace_advance: # {uuid}`` is 78 chars and 500'd on TiDB (strict mode, 1406) when this # was 64. scope: Mapped[str] = mapped_column(String(128), nullable=False) state: Mapped[str] = mapped_column(String(16), nullable=False, default=STATE_PENDING) # SHA-256 hex of the raw request body the key was first used with (see # migration 432). Binds the key to ONE payload: a reused key with a # different body must 422, not replay the first call's success. Nullable # — rows written before the column existed (and hand-wired call sites # that don't pass a fingerprint) skip the comparison. request_fingerprint: Mapped[str | None] = mapped_column( String(64), nullable=True, ) resource_id: Mapped[uuid.UUID | None] = mapped_column(Uuid, nullable=True) response_status: Mapped[int | None] = mapped_column(Integer, nullable=True) response_body_json: Mapped[str | None] = mapped_column(Text, nullable=True) expires_at: Mapped[datetime] = mapped_column( DateTime(timezone=True), nullable=False, ) created_at: Mapped[datetime] = mapped_column( DateTime(timezone=True), nullable=False, server_default=func.now(), ) updated_at: Mapped[datetime] = mapped_column( DateTime(timezone=True), nullable=False, server_default=func.now(), onupdate=func.now(), ) class IdempotencyService: """Per-request helper bound to a (scope, key) pair. ``key`` may be ``None`` (client didn't opt in); in that case ``lookup()`` always returns ``None`` and ``record()`` is a no-op. """ def __init__( self, db: AsyncSession, *, scope: str, key: str | None, ttl_hours: int = DEFAULT_TTL_HOURS, takeover_minutes: int | None = IN_PROGRESS_TAKEOVER_MINUTES, request_fingerprint: str | None = None, ) -> None: self.db = db self.scope = scope self.key = (key or "").strip() or None self.ttl_hours = ttl_hours # SHA-256 hex of the raw request body. ``None`` (hand-wired call # sites) disables the payload-binding check for this request; a # stored NULL (pre-432 row) disables it for that row. Both sides # must be present for a mismatch to be an error, so deploys and # non-fingerprinting callers stay seamless. self.request_fingerprint = request_fingerprint # ``None`` disables abandoned-reservation takeover for this scope. # Correct for operations whose handler can legitimately outlive the # window — a provider fan-out is not bounded by the proxy timeout, # because uvicorn does not cancel the coroutine when the client # disconnects. Taking such a reservation over runs the fan-out a # second time, concurrently with the first. self.takeover_minutes = takeover_minutes self._reserved_row: IdempotencyRecord | None = None @property def reserved_row_id(self) -> uuid.UUID | None: """Id of the reservation THIS instance created, if any. The deferred release needs it to avoid settling a successor's row. Reads ``__dict__`` rather than the mapped attribute ON PURPOSE: this is called from an exception path where a failed flush has already put the session in pending-rollback and expired the identity map. Ordinary attribute access would then emit a refresh SELECT, raise ``PendingRollbackError``, mask the handler's real exception, and skip the deferred release. ``__dict__`` never triggers loader machinery — absent id (nothing flushed yet, or the state was expired) reads as ``None``. """ row = self._reserved_row if row is None: return None row_id = row.__dict__.get("id") return row_id if isinstance(row_id, uuid.UUID) else None def _check_fingerprint(self, existing: IdempotencyRecord) -> None: """422 when the key is being reused with a DIFFERENT request body. Only when both sides carry a fingerprint: ``None`` on either side means "not fingerprinted" (a hand-wired call site, or a row written before migration 432) and the comparison is skipped. """ if ( self.request_fingerprint is not None and existing.request_fingerprint is not None and existing.request_fingerprint != self.request_fingerprint ): raise IdempotencyPayloadMismatchError() def _fresh_reservation(self, now: datetime) -> IdempotencyRecord: return IdempotencyRecord( key=self.key, scope=self.scope, state=STATE_PENDING, request_fingerprint=self.request_fingerprint, expires_at=now + timedelta(hours=self.ttl_hours), ) async def lookup(self) -> dict[str, Any] | None: """Return cached response dict if the operation already succeeded. Raises ``ConflictError`` if the operation is still in progress. Returns ``None`` when the caller should proceed with the work (either no key was provided, or a fresh reservation was made). """ if self.key is None: return None # Row-lock to serialize concurrent requests with the same key. existing = await self.db.scalar( select(IdempotencyRecord) .where( IdempotencyRecord.key == self.key, IdempotencyRecord.scope == self.scope, ) .with_for_update(), ) now = datetime.now(UTC) if existing is None: # ``with_for_update`` above cannot lock a row that does not # exist yet, so two concurrent first-time requests with the # same key can both fall through here, both INSERT, and the # ``uq_idempotency_key_scope`` constraint then 500s the # loser with an IntegrityError instead of the intended # ConflictError. Attempt the reservation inside a SAVEPOINT: # on a unique-violation the duplicate-key INSERT blocked # until the winner committed, so the winning row is now # visible — re-read it and fall through to the existing-row # handling below (which raises ConflictError for the # in-progress PENDING winner). Post-review fix 2026-06-14. self._reserved_row = self._fresh_reservation(now) try: # ``add`` MUST be inside the savepoint: on a unique # violation the SAVEPOINT rollback then cleanly expunges # the failed-INSERT object, so a later autoflush won't # retry it and re-raise. (Adding before begin_nested # leaves it pending in ``session.new`` after rollback.) async with self.db.begin_nested(): self.db.add(self._reserved_row) await self.db.flush() return None except IntegrityError: self._reserved_row = None existing = await self.db.scalar( select(IdempotencyRecord) .where( IdempotencyRecord.key == self.key, IdempotencyRecord.scope == self.scope, ) .with_for_update(), ) if existing is None: # The winner rolled its own reservation back between # our IntegrityError and this re-read — nothing to # serialize against. Surface as in-progress; the # caller's retry path re-attempts cleanly. raise IdempotencyInProgressError() from None # fall through to existing-row handling below # Treat expired rows as absent — reclaim the slot. MySQL/TiDB # DATETIME strips tzinfo on round-trip, so a row loaded fresh in # a NEW request comes back naive while ``now`` is aware — the # bare comparison raised TypeError and 500'd every cross-request # replay (same-request lookups hit the identity map's aware # in-memory object, which is why tests never caught it). existing_expires = existing.expires_at if existing_expires.tzinfo is None: existing_expires = existing_expires.replace(tzinfo=UTC) if existing_expires <= now: await self.db.delete(existing) await self.db.flush() self._reserved_row = self._fresh_reservation(now) self.db.add(self._reserved_row) await self.db.flush() return None if existing.state == STATE_COMPLETED: # A cached success may only answer the SAME request. Replaying # it for a different body would report "success" for a write # that never executed. self._check_fingerprint(existing) if existing.response_body_json is None: return {} try: return json.loads(existing.response_body_json) except json.JSONDecodeError: # Corrupt cached payload — safer to drop and replay than to # serve garbage. Rare, but survive-able. await self.db.delete(existing) await self.db.flush() self._reserved_row = self._fresh_reservation(now) self.db.add(self._reserved_row) await self.db.flush() return None if existing.state == STATE_PENDING: created = existing.created_at if created is not None and created.tzinfo is None: created = created.replace(tzinfo=UTC) abandoned = ( self.takeover_minutes is not None and created is not None and created <= now - timedelta(minutes=self.takeover_minutes) ) if not abandoned: # A different body under a live reservation is a client bug, # not a retry — 422 before the in-progress 409, because # "retry later" is a lie when the retry can never match. self._check_fingerprint(existing) raise IdempotencyInProgressError() # Abandoned reservation (see IN_PROGRESS_TAKEOVER_MINUTES) — # fall through to the same reclaim path a ``failed`` row takes. # STATE_FAILED, or an abandoned STATE_PENDING — let the caller retry # under the same key. The handler re-runs with THIS request's body, # so the fresh reservation carries this request's fingerprint. await self.db.delete(existing) await self.db.flush() self._reserved_row = self._fresh_reservation(now) self.db.add(self._reserved_row) await self.db.flush() return None async def record( self, response: Any, *, resource_id: uuid.UUID | str | None = None, status_code: int = 200, ) -> None: """Flip the reserved row to ``completed`` and store the response. ``response`` may be a Pydantic model, a dict, or any JSON-able value. Oversized responses (>256 KiB) are stored with a marker instead of the full body — the replay will return ``{}`` so the client is expected to detect completion some other way (e.g. by polling the resource). """ if self.key is None or self._reserved_row is None: return payload: str | None try: if hasattr(response, "model_dump"): data = response.model_dump(mode="json") # type: ignore[assignment] else: # ``jsonable_encoder`` — NOT ``json.dumps(default=str)``. # ``default=str`` renders a datetime via ``str(dt)``, which is # space-separated ("2026-07-27 20:02:24+00:00"), while the LIVE # response goes through FastAPI's encoder and is ISO-8601 # ("2026-07-27T20:02:24+00:00"). So a replay used to hand the # client a different wire format than the first call — and # ``new Date("2026-07-27 20:02:24+00:00")`` is **Invalid Date** # on WebKit, i.e. every timestamp in a replayed response # rendered as garbage inside the Feishu iOS webview. Found # 2026-07-28 by an end-to-end replay test that asserted the two # bodies are equal. data = jsonable_encoder(response) payload = json.dumps(data) if len(payload) > MAX_RESPONSE_BYTES: payload = json.dumps( {"_truncated": True, "reason": "response exceeds cache size"}, ) except (TypeError, ValueError): payload = json.dumps({"_unserializable": True}) self._reserved_row.state = STATE_COMPLETED self._reserved_row.resource_id = ( uuid.UUID(str(resource_id)) if resource_id else None ) self._reserved_row.response_status = status_code self._reserved_row.response_body_json = payload await self.db.flush() async def mark_failed(self) -> None: """Explicitly mark the reservation failed (letting future retries go). Normally called from an exception handler to avoid the reservation blocking legitimate client retries. If not called, the row stays ``pending`` until TTL expiry and retries raise ``ConflictError`` until then — that's usually the wrong behaviour. """ if self.key is None or self._reserved_row is None: return self._reserved_row.state = STATE_FAILED await self.db.flush() async def release_reservation( scope: str, key: str | None, row_id: uuid.UUID | None = None, ) -> None: """Mark a ``pending`` reservation failed in its OWN transaction. ``IdempotencyService.mark_failed`` only flushes — it rides on the request session, exactly like the reservation itself. That is right when the handler and the reservation live or die together, and wrong for the majority of the decorated handlers, which ``commit()`` internally: handler commits (reservation row is now durable, state='pending') → a line AFTER the commit raises (serialisation, a lookup, …) → mark_failed() flushes state='failed' → get_db's exception path calls session.rollback() → the 'failed' flush is undone; the row stays 'pending' The client then gets 409 on every retry of a write that never completed. This runs the state change on a fresh session with its own commit, so it survives the request session's rollback. Deliberately quiet: an idempotency reservation that fails to release is a degraded retry experience (409 until the takeover window), never a correctness problem, and this runs on a path where an exception is already propagating — raising here would replace the real error with a bookkeeping one. """ if not key: return from app.core.database import async_session_factory try: async with async_session_factory() as session: criteria = [ IdempotencyRecord.key == key, IdempotencyRecord.scope == scope, IdempotencyRecord.state == STATE_PENDING, ] if row_id is not None: # Match the row THIS request reserved, not merely "a pending # row with these coordinates". # # The release is deferred by a second, and a client retry # arrives sooner than that (the mobile retry backoff starts # at 500 ms). Without this the release lands on the # SUCCESSOR's reservation: request B is mid-flight, its row # gets marked `failed`, and a third caller with the same key # is then waved through to run the guarded operation # concurrently with B — the exact outcome the guard exists # to prevent, caused by the code meant to protect it. criteria.append(IdempotencyRecord.id == row_id) row = await session.scalar( select(IdempotencyRecord).where(*criteria) ) if row is None: # The request session rolled the reservation back (nothing # to release), another path already settled it, or — with # ``row_id`` — the pending row now belongs to a different, # live attempt. All three mean "leave it alone". return row.state = STATE_FAILED await session.commit() except Exception: # pragma: no cover - defensive logging.getLogger(__name__).warning( "idempotency: could not release reservation scope=%s", scope, exc_info=True, ) async def cleanup_expired(db: AsyncSession, *, batch_size: int = 500) -> int: """Delete expired idempotency rows in bounded batches. Returns total deleted. Called from the ``core.idempotency_cleanup`` scheduled job. Batched (not a single unbounded ``DELETE``): a table that has silently accumulated for months can hold a huge expired set, and on TiDB one giant DELETE can exceed the transaction size limit ("transaction too large") and fail EVERY run — leaving the table unbounded, the worse failure mode. Each batch selects up to ``batch_size`` expired ids and deletes them in its own committed transaction, mirroring the other bounded sweeps in the codebase. """ from sqlalchemy import delete, select now = datetime.now(UTC) total = 0 while True: ids = ( await db.execute( select(IdempotencyRecord.id) .where(IdempotencyRecord.expires_at <= now) .order_by(IdempotencyRecord.expires_at) .limit(batch_size) ) ).scalars().all() if not ids: break result = await db.execute( delete(IdempotencyRecord).where(IdempotencyRecord.id.in_(ids)), ) await db.commit() total += result.rowcount or 0 if len(ids) < batch_size: break return total __all__ = [ "DEFAULT_TTL_HOURS", "IdempotencyRecord", "IdempotencyService", "STATE_COMPLETED", "STATE_FAILED", "STATE_PENDING", "cleanup_expired", ]