diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 8e74b80..d72ffbc 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -90,7 +90,7 @@ jobs: # Real-DB integration tests -- the offline suite above stubs Postgres, # so this is the only job that exercises the actual schema end to end. - name: Test (real-Postgres integration) - run: pytest -q tests/test_auth_integration.py tests/test_feedback_integration.py tests/test_actions_integration.py tests/test_llm_credentials_integration.py + run: pytest -q tests/test_auth_integration.py tests/test_feedback_integration.py tests/test_actions_integration.py tests/test_llm_credentials_integration.py tests/test_extraction_concurrency_integration.py web: name: Web (lint · type · test · build) diff --git a/apps/api/app/db/models/llm_credential.py b/apps/api/app/db/models/llm_credential.py index 194c0e0..bb7672c 100644 --- a/apps/api/app/db/models/llm_credential.py +++ b/apps/api/app/db/models/llm_credential.py @@ -80,9 +80,10 @@ class UserLlmCredential(Base): model: Mapped[str] = mapped_column(Text, nullable=False) # Stamped by a successful /test only, via the revision-conditioned UPDATE. last_verified_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True)) - # Incremented on every PUT; the /test success stamp conditions on it so a - # credential replaced mid-test can never be marked verified by the old - # credential's test call. + # Incremented when a PUT changes the credential's contents (key/provider/ + # base_url/model) -- not on flag-only edits; the /test success stamp + # conditions on it so a credential replaced mid-test can never be marked + # verified by the old credential's test call. revision: Mapped[int] = mapped_column(Integer, nullable=False, server_default="1") # Separate from having a credential at all: extraction runs on a settled # verdict subset, classification runs on every ingested message, so a diff --git a/apps/api/app/routes/llm_settings.py b/apps/api/app/routes/llm_settings.py index c7dca15..0f0d06a 100644 --- a/apps/api/app/routes/llm_settings.py +++ b/apps/api/app/routes/llm_settings.py @@ -58,6 +58,7 @@ LlmTestResultOut, LlmUsageOut, ) +from app.services.nlp import persistence from app.services.nlp.extractor import test_credential from app.services.nlp.providers import ( PROVIDER_PRESETS, @@ -121,8 +122,13 @@ def _lock_guard_user(db: Session, user_id: UUID) -> None: REPLACES the old single-row `with_for_update()` PUT used to take on the credential row itself -- no route below takes a separate credential-row lock of its own; this is THE lock. + + Delegates to `persistence.lock_guard_user` -- the same implementation + `record_extraction`'s D2 credential-identity check takes, so there's one + SQL statement for this lock, not two copies that could drift (plan: + 2026-08-19-extraction-cost-hardening D2). """ - db.execute(select(AppUser.id).where(AppUser.id == user_id).with_for_update()) + persistence.lock_guard_user(db, user_id) def _credential_rows(db: Session, user_id: UUID) -> list[UserLlmCredential]: @@ -725,12 +731,15 @@ async def put_llm_settings( def _apply_update(target: UserLlmCredential) -> None: # Read BEFORE mutating -- both flag resolvers and `material_changed` # need `target`'s state as it stood before this write. A flag-only - # edit (no new key, same provider/base_url/model) - # must NOT clear last_verified_at -- that'd be a confusing "your - # verified key just went unverified" regression for a save that - # didn't touch the credential itself. + # edit, OR a re-send of the exact same key with the same + # provider/base_url/model, must NOT clear last_verified_at, bump + # revision, or reset extraction attempts -- none of those changed + # what credential is stored, so treating it as material would reopen + # rows that are still capped under the SAME dead credential. + # `target.api_key` is an EncryptedText column that reads back as + # plaintext, so this `!=` is a correct plaintext comparison. material_changed = ( - api_key_provided + (api_key_provided and effective_api_key != target.api_key) or target.provider != provider or target.base_url != base_url or target.model != model @@ -745,9 +754,17 @@ def _apply_update(target: UserLlmCredential) -> None: target.base_url = base_url target.api_key = effective_api_key target.model = model - target.revision += 1 if material_changed: + target.revision += 1 target.last_verified_at = None + # D2 recovery: a material rewrite is exactly the fix a + # credential-class cap-jump was waiting for -- reset this + # caller's terminally-capped rows so they're claimable again + # under whatever now-different credential this write leaves + # active. A flag-only edit changes nothing about the credential + # itself, so it deliberately does NOT reset (the row is still + # capped under the SAME dead credential). + persistence.reset_failed_extraction_attempts(db, current_user.id) target.updated_at = now if row is None: @@ -771,6 +788,11 @@ def _apply_update(target: UserLlmCredential) -> None: ) row.updated_at = now db.add(row) + # D2 recovery: this caller's first-ever credential always lands + # active -- same "the credential just changed" reasoning as a + # material PUT above, for whatever failed rows predate having any + # credential at all (e.g. rows capped under the server fallback). + persistence.reset_failed_extraction_attempts(db, current_user.id) try: db.commit() except IntegrityError: @@ -904,6 +926,12 @@ def delete_llm_settings( Guard-locked (D6) like every other structural mutation, so a concurrent create/activate for this user can't interleave with the wipe. + + Deliberately does NOT reset D2's terminally-capped ``failed`` rows + (unlike PUT/create/activate) -- there's no credential left to retry + under, so resetting here would just leave them capped-at-zero with + nothing to claim them until a NEW credential shows up, at which point + that credential's own mutation resets them anyway. """ _lock_guard_user(db, current_user.id) rows = _credential_rows(db, current_user.id) @@ -1005,6 +1033,12 @@ async def create_llm_credential( ) row.updated_at = now db.add(row) + if not rows: + # D2 recovery: landing active because it's the caller's first + # credential is the same "the active credential just changed" event + # as a material PUT -- reset. An inactive spare (rows non-empty) + # changes nothing about which credential is active, so it does not. + persistence.reset_failed_extraction_attempts(db, current_user.id) db.commit() return _credential_summary(row) @@ -1053,6 +1087,11 @@ def activate_llm_credential( old_active.is_active = False db.flush() target.is_active = True + # D2 recovery: a genuine switch changes which credential is active, + # so any of the caller's rows capped under the old one deserve a + # fresh claimable life. A call that targets the already-active + # credential (this branch doesn't run) changes nothing. + persistence.reset_failed_extraction_attempts(db, current_user.id) db.commit() return _credential_summary(target) @@ -1067,6 +1106,10 @@ def delete_llm_credential( (switch to another with ``/activate`` first, or use the singular ``DELETE /settings/llm`` kill switch to remove everything at once). 404s if ``credential_id`` isn't the caller's. + + Always removes an INACTIVE spare, so the active credential (whatever + D2's cap-jumped rows are capped under) never changes here -- no + `persistence.reset_failed_extraction_attempts` call, same reasoning as the kill switch. """ _lock_guard_user(db, current_user.id) rows = _credential_rows(db, current_user.id) diff --git a/apps/api/app/services/nlp/classifier.py b/apps/api/app/services/nlp/classifier.py index 6c83587..74a632d 100644 --- a/apps/api/app/services/nlp/classifier.py +++ b/apps/api/app/services/nlp/classifier.py @@ -351,44 +351,46 @@ def _classify_attempt( # gone -- kept as one constant anyway since there's still only one place # that needs it. _CLASSIFICATION_PROMPT = ( - "You classify an email into exactly ONE label, from the RECIPIENT's point " - "of view. Decide what the recipient must actually DO.\n\n" - "Labels:\n" - "- needs_reply: the recipient is personally expected to WRITE BACK. A real " - "person asks them a question, requests info, or awaits their response.\n" - "- action_required: the recipient must personally complete a concrete " - "off-email task with a real consequence or deadline -- pay an invoice, sign " - "a document, submit a form, RSVP to a real invitation, reset a password they " - "must change. NOT a reply, and NOT optional.\n" - "- fyi: informational / automated / transactional mail to read for awareness. " - "Receipts, order & shipping updates, notifications, statements, newsletters " - "you subscribed to, calendar notices, app/system alerts. This is the DEFAULT " - "when no genuine personal task or reply is required.\n" - "- promotional: marketing, sales, offers, deals, or bulk mail trying to get " - "you to buy or click. Has a commercial/advertising intent.\n" - "- security_alert: account or login security -- verification codes, new " - "sign-ins, suspicious activity, password/2FA notices.\n" - "- spam: junk, scams, or phishing.\n\n" - "CRITICAL boundary rules (this is where mistakes happen):\n" - "1. Marketing CTAs are NOT action_required. 'Shop now', 'click here', " - "'limited time', 'upgrade today', auto-renewal notices -> promotional or fyi.\n" - "2. A soft/optional/automated 'confirm', 'renew', 'click', 'verify' is NOT " - "action_required. Only use action_required when the recipient genuinely has " - "to do the task or face a consequence.\n" - "3. When unsure between action_required and fyi, choose fyi. Most automated " - "and bulk email is fyi.\n" - "4. needs_reply requires a real human awaiting YOUR written response, not an " - "automated 'do not reply' message.\n\n" - "Examples:\n" - "- 'Your Amazon order has shipped, arriving Tuesday' -> fyi\n" - "- 'Your monthly statement is ready to view' -> fyi\n" - "- 'Invoice #429 is due Friday, pay to avoid late fees' -> action_required\n" - "- 'Please sign the attached contract by EOD' -> action_required\n" - "- 'Hey, can you send me the report when you get a chance?' -> needs_reply\n" - "- '50% off this weekend only -- shop now!' -> promotional\n" - "- 'New sign-in to your account from a new device' -> security_alert\n\n" - "Return JSON only with keys: label, confidence (0-1), rationale." - ) + "You classify an email into exactly ONE label, from the RECIPIENT's " + "point of view -- what must they actually DO?\n\n" + "Labels:\n" + "- needs_reply: recipient must WRITE BACK -- a real person asked a " + "question, requested info, or awaits a reply.\n" + "- action_required: a concrete off-email task with a real deadline or " + "consequence -- pay, sign, submit a form, RSVP, reset a required " + "password. Not a reply, not optional.\n" + "- fyi: informational/automated/transactional mail for awareness " + "(receipts, shipping updates, statements, newsletters, system " + "alerts). DEFAULT when no genuine task or reply is required.\n" + "- promotional: marketing/sales/offers/bulk mail with commercial " + "intent.\n" + "- security_alert: account/login security -- codes, new sign-ins, " + "suspicious activity, password/2FA notices.\n" + "- spam: junk, scams, or phishing.\n\n" + "Boundary rules:\n" + "1. Marketing CTAs ('shop now', 'click here', 'limited time', " + "auto-renewal) -> promotional or fyi, not action_required.\n" + "2. A soft/optional/automated 'confirm'/'renew'/'click'/'verify' is not " + "action_required -- only when the recipient must act or face a real " + "consequence.\n" + "3. Unsure between action_required and fyi? Choose fyi -- most " + "automated/bulk mail is fyi.\n" + "4. needs_reply needs a real human awaiting YOUR written reply, not an " + "automated do-not-reply message.\n\n" + "Examples:\n" + "- '50% off this weekend only -- shop now!' -> promotional, not " + "action_required.\n" + "- 'Your monthly statement is ready to view' -> fyi, not " + "action_required.\n" + "- 'Automated notice: your account was updated. Do not reply.' -> " + "fyi, not needs_reply.\n\n" + "Return JSON only with keys: label, confidence (0-1), rationale." +) + +# D9 (2026-08-19-extraction-cost-hardening): prod p50 body is ~2,450 chars, +# so the median call is untouched by this cap -- it only trims the long tail, +# where the label signal is already in the first few thousand chars anyway. +_CLASSIFICATION_TEXT_MAX_LEN = 3000 # The BYOK classification call's own token budget -- a label/confidence/ # rationale reply is far smaller than extraction's structured JSON, so this @@ -547,7 +549,7 @@ def _classify_llm_user( result = call_chat_completion( credential, prompt=_CLASSIFICATION_PROMPT, - user_content=f"Email:\n{text[:6000]}", + user_content=f"Email:\n{text[:_CLASSIFICATION_TEXT_MAX_LEN]}", max_tokens=_CLASSIFICATION_MAX_TOKENS, timeout=_CLASSIFICATION_TIMEOUT_S, policy=policy, diff --git a/apps/api/app/services/nlp/extraction_run.py b/apps/api/app/services/nlp/extraction_run.py index f40ad65..4d077b3 100644 --- a/apps/api/app/services/nlp/extraction_run.py +++ b/apps/api/app/services/nlp/extraction_run.py @@ -32,6 +32,7 @@ from app.services.nlp.extractor import ACTION_LABELS, NoAction, extract_action_with_usage from app.services.nlp.llm_client import WORKER_RETRIES from app.services.nlp.persistence import ( + CREDENTIAL_CLASS_FAILURE_CATEGORIES, MAX_ATTEMPTS, PENDING_LEASE, claim_action_item, @@ -50,23 +51,32 @@ # The extraction-sweep twin of backfill.py's _CONSECUTIVE_NO_VERDICT_LIMIT # (plan: phase 3 of the LLM-failure work -- Phase 2 added the classification -# side of this, this closes the same gap for run_extraction_sweep). Same -# value, same reasoning: after this many CONSECUTIVE "failed" buckets in a -# row, stop issuing more extraction calls for the rest of this sweep. Not -# 1 -- a single odd message shouldn't abort an otherwise-healthy sweep -- -# and the counter resets to zero on any non-"failed" bucket, so only a -# genuine losing streak trips it. Each attempt this stop prevents can be a -# real BYOK-billed call very likely to fail the same way as the last few. +# side of this, this closes the same gap for run_extraction_sweep). After +# this many CONSECUTIVE credential-class failures (plan: +# 2026-08-19-extraction-cost-hardening D3 -- +# CREDENTIAL_CLASS_FAILURE_CATEGORIES: the SAME credential would fail every +# retry identically), stop issuing more extraction calls for the rest of +# this sweep. Not 1 -- a single odd message shouldn't abort an otherwise- +# healthy sweep. Each attempt this stop prevents can be a real BYOK-billed +# call very likely to fail the same way as the last few. # -# This streak is fed ONLY by a genuine LLM-result failure -- the "failed" -# bucket _claim_extract_record returns itself, which always means an actual -# extraction attempt reached (or tried to reach) a provider and came back -# empty (see ExtractionAttempt's docstring: `result is None` iff the call or -# its parse failed). It is deliberately NOT fed by a containment/DB failure -# (run_extraction_sweep's own fan-out `except Exception` below) -- those -# carry no failure_category and say nothing about the LLM provider's health, -# so counting them here would misdiagnose a Postgres hiccup as "the user's -# provider is rate-limited" (Codex review, phase 3 should-fix). +# D3 refines this from the ORIGINAL "three consecutive failed BUCKETS" +# breaker: that version missed the prod retry storm entirely, because the +# storm's calls surfaced under the "skipped" bucket (a genuine attempt whose +# record got fenced out by a competing reclaim), which sailed straight past +# a bucket == "failed" check. The streak now reads the per-attempt failure +# CATEGORY instead, threaded out of `_claim_extract_record` as this +# function's third return value -- a credential-class category feeds the +# streak regardless of which bucket it surfaced under, so a run of +# skipped-but-credential-dead attempts trips it exactly like a run of +# failed-bucket ones would. Resets to zero on anything else: a success, a +# non-credential-class category (rate limits recover on their own), or a +# clean skip with no category at all. It is deliberately NOT fed (and is +# explicitly reset) by a containment/DB failure (run_extraction_sweep's own +# fan-out `except Exception` below) -- those carry no failure_category and +# say nothing about the LLM provider's health, so counting them here would +# misdiagnose a Postgres hiccup as "the user's provider is dead" (Codex +# review, phase 3 should-fix; D3 carries the same reasoning forward). _CONSECUTIVE_FAILURE_LIMIT = 3 @@ -89,8 +99,16 @@ def _empty_counts() -> dict: # e.g. {"http_429": 8}, populated by run_extraction_sweep only; other # callers of _empty_counts() (run_extraction_for_message, the disabled # shape) legitimately leave it empty, same as an old result predating - # this field. - return {"processed": 0, **{bucket: 0 for bucket in _RESULT_BUCKETS}, "failure_categories": {}} + # this field. `aborted` (plan: 2026-08-19-extraction-cost-hardening D3) + # is the same kind of additive field: how many candidates the breaker + # skipped by aborting the rest of a sweep -- always 0 outside + # run_extraction_sweep's own credential-class trip. + return { + "processed": 0, + **{bucket: 0 for bucket in _RESULT_BUCKETS}, + "failure_categories": {}, + "aborted": 0, + } def _disabled_result() -> dict: @@ -121,12 +139,25 @@ def terminalize_expired_pending(db: Session, *, user_id: UUID | None = None) -> db.commit() -def _call_context_from_resolved(credential, source: str | None) -> CallContext: - """Map `ResolvedExtraction.source` onto `CallContext.payer` -- 'fallback' - is the operator's own key, and only a 'user' source should ever get - billed onto the account whose usage we're recording (plan §1, §3). +def _call_context_from_resolved(resolved) -> CallContext: + """Build a `CallContext` from a `ResolvedExtraction`: maps `.source` onto + `CallContext.payer` -- 'fallback' is the operator's own key, and only a + 'user' source should ever get billed onto the account whose usage we're + recording (plan §1, §3) -- and carries the resolved row's `credential_id`/ + `revision` straight through (plan: 2026-08-19-extraction-cost-hardening + D2), so `record_extraction`'s cap-jump can tell whether a rotation raced + the attempt that used this context. `getattr` with a `None` default, + not a direct attribute read -- the real `ResolvedExtraction` always + carries both, but plenty of tests stand in a bare + ``SimpleNamespace(credential=..., source=...)`` for it, and `None` is + the exact fallback-credential value anyway (no per-user row to rotate). """ - return CallContext(credential=credential, payer="user" if source == "user" else "operator") + return CallContext( + credential=resolved.credential, + payer="user" if resolved.source == "user" else "operator", + credential_id=getattr(resolved, "credential_id", None), + revision=getattr(resolved, "revision", None), + ) def _fence_claim_to_failed(message_id: UUID, claim_token: UUID) -> None: @@ -199,14 +230,23 @@ def _claim_extract_record( acc: UsageAccumulator | None = None, failure_categories: dict[str, int] | None = None, deadline: float | None = None, -) -> tuple[str, UUID | None]: +) -> tuple[str, UUID | None, str | None]: """One claim -> extract -> record cycle for ``message_id``. - Returns ``(bucket, user_id)`` -- ``bucket`` is one of - extracted/no_action/ineligible/failed/skipped/unavailable, and + Returns ``(bucket, user_id, failure_category)`` -- ``bucket`` is one of + extracted/no_action/ineligible/failed/skipped/unavailable/deadline; ``user_id`` is the message's owner when known (``None`` if the message or its thread has vanished), so callers can attach it to a task result - without a second query. + without a second query; ``failure_category`` (plan: + 2026-08-19-extraction-cost-hardening D3) is the attempt's own + ``ExtractionAttempt.failure_category`` when a real wire call was made and + failed, ``None`` otherwise -- including on every SUCCESS bucket, and on + ``"skipped"`` when no attempt was even made (a lost claim race). This is + what lets ``run_extraction_sweep``'s breaker read the credential-class + signal regardless of which bucket a genuine attempt's failure surfaced + under (a fenced-out record from a competing reclaim still reports its + real category here, even though its bucket is ``"skipped"``, not + ``"failed"``). ``call_context`` is the already-resolved credential AND who pays for it, for a sweep's user (``run_extraction_sweep`` resolves it once per run, @@ -219,16 +259,15 @@ def _claim_extract_record( given, otherwise a fresh one built here once the user id is known -- that single call is a run of one. - ``failure_categories`` is an out-param, not a return value: when given, a - call whose attempt actually failed (``attempt.failure_category`` set) - increments its bucket in place. Threading it through the return tuple - instead would mean every existing caller (and the many tests unpacking - ``bucket, user_id = ...``) has to change shape for a count only - ``run_extraction_sweep`` needs -- ``None`` (the default) just skips it. + ``failure_categories`` is an out-param, not (only) the return value: when + given, a call whose attempt actually failed (``attempt.failure_category`` + set) increments its bucket in place -- this is the sweep's own aggregate + breakdown across the whole run, distinct from the third return value's + per-call signal that feeds the streak. """ message = db.get(MailMessage, message_id) if message is None: - return "skipped", None + return "skipped", None, None # Frozen lock order: MailThread -> ActionItem. This FOR UPDATE also # serializes against a concurrent set_thread_done -- whichever commits @@ -238,7 +277,7 @@ def _claim_extract_record( select(MailThread).where(MailThread.id == message.thread_id).with_for_update() ).scalar_one_or_none() if thread is None: - return "skipped", None + return "skipped", None, None thread_done = thread.done_at is not None subject = thread.subject @@ -254,8 +293,8 @@ def _claim_extract_record( # Nothing was claimed, so a plain commit (not a rollback) is # enough to release the thread lock we just took. db.commit() - return "unavailable", user_id - call_context = _call_context_from_resolved(resolved.credential, resolved.source) + return "unavailable", user_id, None + call_context = _call_context_from_resolved(resolved) if acc is None: acc = UsageAccumulator(user_id) @@ -267,13 +306,18 @@ def _claim_extract_record( user_id=user_id, thread_done=thread_done, force=force, + model_version=f"{call_context.credential.provider}:{call_context.credential.model}", ) # Release the thread lock before the (possibly slow) LLM call -- no DB # transaction is checked out while extract_action runs. db.commit() if claim_token is None: - return "skipped", user_id + return "skipped", user_id, None + # Set once the wire call actually returns, below -- stays None for the + # deadline-exhausted path (nothing was sent) so every return before that + # point reports "no attempt, no category" honestly. + category: str | None = None try: # Final verify pass: the sweep's pre-claim deadline check can't see # time spent WAITING on the claim's own locks -- re-check here, @@ -299,10 +343,9 @@ def _claim_extract_record( policy=WORKER_RETRIES, ) result = attempt.result - if failure_categories is not None and attempt.failure_category is not None: - failure_categories[attempt.failure_category] = ( - failure_categories.get(attempt.failure_category, 0) + 1 - ) + category = attempt.failure_category + if failure_categories is not None and category is not None: + failure_categories[category] = failure_categories.get(category, 0) + 1 # Lock the classification row BEFORE reading its label and hold it # through the record commit: a concurrent reclassify-back upserts @@ -344,6 +387,11 @@ def _claim_extract_record( thread_done=current_done_at is not None, label_still_actionable=label_still_actionable, already_replied=already_replied, + model_version=f"{call_context.credential.provider}:{call_context.credential.model}", + user_id=user_id, + failure_category=category, + credential_id=call_context.credential_id, + credential_revision=call_context.revision, ) # Only a "user" payer gets recorded -- the operator's fallback key @@ -383,7 +431,7 @@ def _claim_extract_record( # reclaimable -- but report the stop distinctly so the sweep can # end with status "timed_out" instead of "ok" (verify pass 3). _rollback_discard_and_fence(db, acc, message.id, claim_token, exc) - return "deadline", user_id + return "deadline", user_id, None except SoftTimeLimitExceeded as exc: # Celery's soft time limit is a plain Exception, so the generic # `except Exception` below WOULD otherwise catch it -- and @@ -397,25 +445,35 @@ def _claim_extract_record( _rollback_discard_and_fence(db, acc, message.id, claim_token, exc) raise except Exception as exc: + # `category` may already be set here -- the wire call can succeed + # and THEN something else in this try block (the classification + # lock, the usage flush) throws. Reporting it is correct: a real + # credential-class signal from the provider happened, even though + # this particular row settles via the containment fence rather than + # a normal record. _rollback_discard_and_fence(db, acc, message.id, claim_token, exc) - return "failed", user_id + return "failed", user_id, category if not recorded: # Our lease expired and someone else reclaimed the row before this - # write landed -- a stale, harmless no-op, not a bug. + # write landed -- a stale, harmless no-op, not a bug. `category` + # still reports honestly here: the wire call genuinely happened and + # (if it failed) genuinely failed -- this is the D1 stale-lease/ + # competing-reclaim combination the plan's own repro test drives: + # bucket "skipped" with a real failure_category attached. logger.warning( "action extraction record fenced out (stale claim_token) for message %s", message.id, ) - return "skipped", user_id + return "skipped", user_id, category if not label_still_actionable: - return "ineligible", user_id + return "ineligible", user_id, category if isinstance(result, NoAction): - return "no_action", user_id + return "no_action", user_id, category if result is None: - return "failed", user_id - return "extracted", user_id + return "failed", user_id, category + return "extracted", user_id, category def run_extraction_for_message(db: Session, message_id: UUID) -> dict: @@ -433,7 +491,7 @@ def run_extraction_for_message(db: Session, message_id: UUID) -> dict: if not extraction_feature_enabled(): return _disabled_result() - bucket, user_id = _claim_extract_record(db, message_id) + bucket, user_id, _category = _claim_extract_record(db, message_id) if bucket == "unavailable": return _disabled_result() @@ -447,7 +505,13 @@ def run_extraction_for_message(db: Session, message_id: UUID) -> dict: def _message_driven_candidates( - db: Session, user_id: UUID, *, since_days: int, limit: int, force: bool + db: Session, + user_id: UUID, + *, + since_days: int, + limit: int, + force: bool, + model_version: str | None = None, ) -> list[UUID]: """Never-claimed or currently-claimable messages: the user's messages whose classification label is actionable, thread not done, within @@ -455,6 +519,14 @@ def _message_driven_candidates( (``force``-widened) -- otherwise a settled/live-pending message would burn a selection slot for nothing. + ``model_version`` (plan: 2026-08-19-extraction-cost-hardening D5) is the + CURRENT ``f"{provider}:{model}"``, threaded straight to + ``claimable_predicate`` -- only meaningful when ``force`` is set, since + that's the only case a settled ``extracted``/``no_action`` row can ever + become claimable again (comparing against a stale/None model_version + when ``force`` is off is harmless: that branch of the predicate never + fires without ``force`` regardless). + Messages with a NULL ``sent_at`` are possible by design (see ``backfill.py``'s ``latest_message_ordering``) -- coalescing to ``created_at`` for both the cutoff filter and the ordering keeps them @@ -473,7 +545,10 @@ def _message_driven_candidates( MailThread.done_at.is_(None), Classification.label.in_(ACTION_LABELS), sent_at >= cutoff, - or_(ActionItem.id.is_(None), claimable_predicate(force=force)), + or_( + ActionItem.id.is_(None), + claimable_predicate(force=force, model_version=model_version), + ), ) .order_by(sent_at.desc().nullslast()) .limit(limit) @@ -587,8 +662,13 @@ def run_extraction_sweep( if resolved.credential is None: return _disabled_result() - call_context = _call_context_from_resolved(resolved.credential, resolved.source) + call_context = _call_context_from_resolved(resolved) acc = UsageAccumulator(user_id) + # D5: what "the current model" means for this run, for + # claimable_predicate(force=True)'s IS DISTINCT FROM comparison -- the + # sweep's OWN candidate query is the only claimable_predicate call site + # that ever runs under force, so this is threaded exactly once, here. + current_model_version = f"{call_context.credential.provider}:{call_context.credential.model}" terminalize_expired_pending(db, user_id=user_id) @@ -596,32 +676,35 @@ def run_extraction_sweep( message_ids = _recovery_candidates(db, user_id, limit=limit) else: message_ids = _message_driven_candidates( - db, user_id, since_days=since_days, limit=limit, force=force + db, + user_id, + since_days=since_days, + limit=limit, + force=force, + model_version=current_model_version, ) counts = _empty_counts() # _CONSECUTIVE_FAILURE_LIMIT's early stop (plan: phase 3 of the - # LLM-failure work, mirroring Phase 2's run_backfill treatment) -- reset - # on any bucket other than "failed", tripped after 3 in a row. Fed ONLY - # by _claim_extract_record's own "failed" return (a genuine LLM-result - # failure) -- see this module's _CONSECUTIVE_FAILURE_LIMIT comment for - # why a containment/DB failure (the `except Exception` branch below) - # must never touch this counter. - consecutive_failures = 0 - # Distinct from consecutive_failures: which early-stop tripped, if any, + # LLM-failure work, mirroring Phase 2's run_backfill treatment; refined + # by D3 -- see this module's _CONSECUTIVE_FAILURE_LIMIT comment for the + # full "why category, not bucket" reasoning). Reset on anything that + # isn't a credential-class category, tripped after 3 in a row. + consecutive_credential_class_failures = 0 + # Distinct from the streak counter: which early-stop tripped, if any, # decides the reported status. "llm_unavailable" is a diagnosis about - # the provider and must only ever come from a genuine LLM failure + # the provider and must only ever come from a genuine credential-class # streak; "timed_out" is a neutral, no-diagnosis stop for the deadline # backstop and for Celery's own soft time limit (mirrors # extraction_recovery_tick's existing status name for the same signal). stop_reason: str | None = None - for message_id in message_ids: + for index, message_id in enumerate(message_ids): if deadline is not None and time.monotonic() >= deadline: stop_reason = "timed_out" break counts["processed"] += 1 try: - bucket, _user_id = _claim_extract_record( + bucket, _user_id, category = _claim_extract_record( db, message_id, force=force, call_context=call_context, acc=acc, failure_categories=counts["failure_categories"], deadline=deadline, ) @@ -646,15 +729,18 @@ def run_extraction_sweep( # This is a CONTAINMENT failure -- a DB error, or anything # _claim_extract_record itself couldn't recover from -- not an # LLM result: it carries no failure_category and must NOT feed - # consecutive_failures/llm_unavailable below, or a single - # Postgres hiccup would misreport as "the user's LLM provider - # is rate-limited, try again later" (Codex review, phase 3 - # should-fix). + # the credential-class streak below, or a single Postgres + # hiccup would misreport as "the user's LLM provider is dead" + # (Codex review, phase 3 should-fix). Explicitly reset rather + # than left alone: an unrelated DB blip between two genuine + # provider failures must not silently bridge them into one + # continuous trip either. db.rollback() logger.exception( "action extraction sweep failed for message %s", message_id ) counts["failed"] += 1 + consecutive_credential_class_failures = 0 continue if bucket == "deadline": @@ -669,17 +755,25 @@ def run_extraction_sweep( break counts[bucket] += 1 - if bucket == "failed": - consecutive_failures += 1 - if consecutive_failures >= _CONSECUTIVE_FAILURE_LIMIT: + if category in CREDENTIAL_CLASS_FAILURE_CATEGORIES: + consecutive_credential_class_failures += 1 + if consecutive_credential_class_failures >= _CONSECUTIVE_FAILURE_LIMIT: # Every remaining candidate is very likely to fail the same # way -- stop spending BYOK-billed attempts on it. The # counts collected so far are still exactly right; the loop - # just never reaches the rest of message_ids. + # just never reaches the rest of message_ids -- report how + # many of them this abort skipped. + logger.warning( + "action extraction sweep for user %s aborting after %d " + "consecutive %r failures -- credential looks dead for " + "this run", + user_id, _CONSECUTIVE_FAILURE_LIMIT, category, + ) stop_reason = "llm_unavailable" + counts["aborted"] = len(message_ids) - (index + 1) break else: - consecutive_failures = 0 + consecutive_credential_class_failures = 0 # Reporting "ok" for a run that quit early would be the same lie this # whole plan exists to remove: the counts would be truthful about what was # attempted and silent about the candidates never reached. Mirrors diff --git a/apps/api/app/services/nlp/extractor.py b/apps/api/app/services/nlp/extractor.py index fb1fd28..8e0d213 100644 --- a/apps/api/app/services/nlp/extractor.py +++ b/apps/api/app/services/nlp/extractor.py @@ -111,13 +111,29 @@ class ExtractionAttempt: def _build_message_text( - subject: str | None, sender: str | None, snippet: str | None, body_text: str | None + subject: str | None, + sender: str | None, + snippet: str | None, + body_text: str | None, + received_at: datetime | None, ) -> str: - """Assemble the email text handed to the model alongside the prompt.""" + """ + Assemble the email text handed to the model alongside the prompt. + + Carries the `Received:` header line (D7: prompt-caching hygiene) -- + the instruction block in `_EXTRACTION_PROMPT` is now byte-identical + across every call, so only this per-message text varies. + """ + received_at_str = ( + received_at.astimezone(timezone.utc).isoformat() + if received_at is not None + else "unknown" + ) return "\n".join( [ f"From: {sender or '(unknown)'}", f"Subject: {subject or '(no subject)'}", + f"Received: {received_at_str}", "", snippet or "", body_text or "", @@ -125,46 +141,39 @@ def _build_message_text( ).strip() -def _build_prompt(received_at: datetime | None) -> str: - # received_at anchors relative deadlines ("by Friday", "end of week") -- - # without it the model would have to guess "today", and guessing a date - # it can't see is exactly the invented-deadline failure mode this - # feature exists to avoid. - received_at_str = ( - received_at.astimezone(timezone.utc).isoformat() - if received_at is not None - else "unknown" - ) - return ( - "You read one email already flagged as needing a reply or a " - "concrete action from the RECIPIENT, and extract the single " - "underlying obligation as structured JSON.\n\n" - f"The message's received_at (ISO 8601 UTC) is: {received_at_str}. " - "Resolve any relative deadline ('by Friday', 'end of week', 'in 3 " - "days') against this timestamp -- never against today's date.\n\n" - "Rules:\n" - "- has_action: false if there is actually no concrete task -- a " - "marketing CTA, an optional/automated 'confirm'/'renew'/'verify', " - "or informational text with nothing for the recipient to DO. Same " - "boundary discipline as classification: when unsure, false.\n" - "- kind: one of reply, payment, signature, form, rsvp, deadline, " - "other -- pick the closest match.\n" - "- title: imperative and concrete (e.g. 'Pay invoice #429'), at " - f"most {_TITLE_MAX_LEN} characters.\n" - "- due_at: the resolved deadline as an ISO 8601 UTC timestamp, or " - "null if no deadline is stated or clearly implied. NEVER invent a " - "date that isn't in the text.\n" - "- due_is_date_only: true if only a calendar date was given with " - "no time of day -- in that case set due_at to 23:59:59 UTC of that " - "date.\n" - "- due_raw: the deadline phrase as written in the email (e.g. 'by " - "end of week'), or null.\n" - "- amount / currency: a stated monetary amount and its currency " - "code, or null for both. NEVER invent an amount.\n" - "- confidence: 0-1, how confident you are in this extraction.\n\n" - "Return JSON only with keys: has_action, kind, title, due_at, " - "due_is_date_only, due_raw, amount, currency, confidence." - ) +# D7 (2026-08-19-extraction-cost-hardening): a plain module constant, not a +# function -- received_at used to make this a per-call string, but it now +# lives in the message text (_build_message_text) instead, so the +# instruction block is the same bytes on every call and can be cached. +# received_at still anchors relative deadlines ("by Friday", "end of +# week") -- without it the model would have to guess "today", and guessing +# a date it can't see is exactly the invented-deadline failure mode this +# feature exists to avoid. +_EXTRACTION_PROMPT = ( + "You read one email already flagged as needing a reply or a concrete " + "action from the RECIPIENT and extract the single underlying " + "obligation as JSON. Resolve relative deadlines ('by Friday', 'end of " + "week', 'in 3 days') against the email's Received: line, never " + "today's date.\n\n" + "Rules:\n" + "- has_action: false for a marketing CTA, an optional/automated " + "'confirm'/'renew'/'verify', or text with nothing for the recipient " + "to DO. Unsure -> false.\n" + "- kind: reply, payment, signature, form, rsvp, deadline, or other -- " + "closest match.\n" + "- title: imperative, concrete (e.g. 'Pay invoice #429'), max " + f"{_TITLE_MAX_LEN} chars.\n" + "- due_at: resolved deadline as ISO 8601 UTC, or null if none stated " + "or clearly implied. Never invent a date.\n" + "- due_is_date_only: true if only a date was given, no time -- then " + "set due_at to 23:59:59 UTC of that date.\n" + "- due_raw: the deadline phrase as written, or null.\n" + "- amount / currency: a stated amount and its currency code, or null " + "for both. Never invent an amount.\n" + "- confidence: 0-1.\n\n" + "Return JSON only with keys: has_action, kind, title, due_at, " + "due_is_date_only, due_raw, amount, currency, confidence." +) def _parse_due_at(due_at_raw: object, due_is_date_only: object) -> tuple[datetime | None, str | None]: @@ -317,11 +326,10 @@ def _extract_attempt( still records exactly one attempt's worth of usage either way (plan: phase 3 of the LLM-failure work). """ - message_text = _build_message_text(subject, sender, snippet, body_text) - prompt = _build_prompt(received_at) + message_text = _build_message_text(subject, sender, snippet, body_text, received_at) try: - call_result = _call_llm(credential, prompt, message_text, policy) + call_result = _call_llm(credential, _EXTRACTION_PROMPT, message_text, policy) except ExtractionCallError as exc: logger.warning( "Action extraction failed for provider %s: %s", credential.provider, exc.category @@ -469,10 +477,9 @@ def test_credential( thread a policy through from, and getting one transient retry for free matches every other inline call site's treatment. """ - prompt = _build_prompt(None) start = time.monotonic() try: - call_result = _call_llm(credential, prompt, _TEST_CREDENTIAL_MESSAGE, policy) + call_result = _call_llm(credential, _EXTRACTION_PROMPT, _TEST_CREDENTIAL_MESSAGE, policy) except ExtractionCallError as exc: return False, exc.category, int((time.monotonic() - start) * 1000) diff --git a/apps/api/app/services/nlp/persistence.py b/apps/api/app/services/nlp/persistence.py index 9064fb1..b9d86ea 100644 --- a/apps/api/app/services/nlp/persistence.py +++ b/apps/api/app/services/nlp/persistence.py @@ -20,11 +20,11 @@ from datetime import datetime, timedelta, timezone from uuid import UUID, uuid4 -from sqlalchemy import and_, case, or_, update +from sqlalchemy import and_, case, or_, select, update from sqlalchemy.dialects.postgresql import insert from sqlalchemy.orm import Session -from app.db.models import ActionItem, Classification +from app.db.models import ActionItem, AppUser, Classification, UserLlmCredential from app.services.nlp.extractor import ExtractedAction, NoAction # A claim's lease: past this age, a fresh worker may steal a still-"pending" @@ -33,6 +33,14 @@ PENDING_LEASE = timedelta(minutes=30) MAX_ATTEMPTS = 3 +# Categories where every retry with the SAME credential fails identically -- +# a bad key (401), a forbidden model (403), or a retired/renamed model +# (404). `blocked_by_policy` is deliberately excluded: it reflects mutable +# deployment/DNS policy, not the credential itself, and policy recovery +# must not require a credential mutation to revive rows (plan: +# 2026-08-19-extraction-cost-hardening D2). +CREDENTIAL_CLASS_FAILURE_CATEGORIES = frozenset({"http_401", "http_403", "http_404"}) + # Marks a label set by a human in the console rather than a model -- lets # every write path (this module, the route, backfill) agree on the string # without importing it from a route. `upsert_classification`'s write-time @@ -41,9 +49,10 @@ OPERATOR_MODEL_VERSION = "user-override" # Bounded snapshot cap for classification_feedback.input_text: the first -# FEEDBACK_TEXT_MAX chars of build_classification_text(...)'s output. Matches -# the LLM paths' own 6,000-char cap and sits ~4x the local model's 256-token -# window, so it's generous enough to cover what either classifier actually +# FEEDBACK_TEXT_MAX chars of build_classification_text(...)'s output. Sized +# independently of either classifier's own cap (classifier.py's LLM path +# truncates at 3,000 chars as of D9; the local encoder's window is smaller +# still) -- 6,000 stays generous enough to cover what either one actually # saw without storing unbounded email bodies forever. FEEDBACK_TEXT_MAX = 6000 @@ -105,16 +114,32 @@ def upsert_classification( return "written" if result.rowcount else "protected" -def claimable_predicate(*, force: bool = False): +def claimable_predicate(*, force: bool = False, model_version: str | None = None): """WHERE-clause fragment: is an EXISTING ``action_item`` row a valid claim target right now? Settled rows (``extracted`` / ``no_action``) are terminal; only ``ineligible`` (a reclassify-back must always be able to re-extract), ``failed`` under the attempt cap, and an expired ``pending`` lease under - the cap remain retryable. ``force`` widens this to any non-pending row -- - an operator backfill can re-run a settled extraction, but force must - never steal a claim a live worker currently holds. + the cap remain retryable. + + ``force`` (plan: 2026-08-19-extraction-cost-hardening D5) widens this, + but no longer to "any non-pending row" -- re-running a settled + extraction for free is only worth paying for when the model that + produced it has since changed: + + - ``failed`` rows are claimable regardless of the attempt cap (retrying + failures is force's whole job). + - ``extracted``/``no_action`` rows are claimable ONLY when the row's + stored ``model_version`` differs from the CURRENT ``model_version`` + (``f"{provider}:{model}"``, threaded from the caller's resolved + credential), via SQL ``IS DISTINCT FROM`` -- plain ``!=`` is + NULL-poisoned and would silently exclude every historical row with no + stored model_version, the opposite of what a one-time re-extraction + needs. Same model -> not claimable even under force; the paid result + stands. + - a live ``pending`` claim is still NEVER stolen, forced or not -- it's + simply absent from every disjunct below. """ lease_cutoff = datetime.now(timezone.utc) - PENDING_LEASE base = or_( @@ -126,9 +151,116 @@ def claimable_predicate(*, force: bool = False): ActionItem.attempts < MAX_ATTEMPTS, ), ) - if force: - return or_(base, ActionItem.outcome != "pending") - return base + if not force: + return base + return or_( + base, + ActionItem.outcome == "failed", + and_( + ActionItem.outcome.in_(("extracted", "no_action")), + ActionItem.model_version.is_distinct_from(model_version), + ), + ) + + +def reset_failed_extraction_attempts(db: Session, user_id: UUID) -> None: + """D2 recovery: zero ``attempts`` on the caller's ``outcome='failed'`` + ``action_item`` rows, under the SAME guard-locked transaction as the + credential mutation that's about to commit. + + A credential-class failure (D2's cap-jump in `record_extraction` below) leaves + a row terminally capped at ``MAX_ATTEMPTS`` under the credential that + just failed it -- the whole point is to stop burning calls on a + credential known to be dead. But once the user actually FIXES that + credential (a material PUT, an activate, or a first-create landing + active), those rows must get a fresh claimable life, or D2's storm fix + would trade "storms forever" for "silently stuck forever" instead. + + Callers gate this on what actually changed: a material rewrite or an + activate that switches which credential is active, yes; a flag-only PUT, + an inactive spare create, a single spare delete, or the kill switch, no + -- none of those change anything about the credential a failed row was + capped under, so resetting there would just re-spend calls against the + exact same dead credential. + """ + db.execute( + update(ActionItem) + .where(ActionItem.user_id == user_id, ActionItem.outcome == "failed") + .values(attempts=0) + ) + + +def lock_guard_user(db: Session, user_id: UUID) -> None: + """The single per-user serialization anchor (PR #64's D6 guard lock, + ``llm_settings._lock_guard_user``): a ``SELECT ... FOR UPDATE`` on + ``app_user`` that every STRUCTURAL credential mutation takes before + reading or writing this user's credential rows. + + Defined here rather than in the routes module so this extraction-record + path can take the SAME lock without services reaching up into routes -- + ``llm_settings.py`` delegates to this implementation instead of keeping + its own copy. Used by ``record_extraction``'s D2 credential-identity + check: taking this lock before re-reading the active credential is what + stops a PUT/activate from committing between that read and the fenced + cap-jump write it guards (the TOCTOU Codex's pass 3 review caught). + """ + db.execute(select(AppUser.id).where(AppUser.id == user_id).with_for_update()) + + +def _credential_class_cap_jump_attempts( + db: Session, + *, + user_id: UUID | None, + failure_category: str | None, + credential_id: UUID | None, + credential_revision: int | None, +) -> int | None: + """The attempts value a credential-class failure should jump the row + to, or ``None`` to leave ``attempts`` untouched (plan D2). + + ``None`` short-circuits whenever there's nothing to jump: a + non-credential-class category (retryable or `blocked_by_policy`), or no + ``user_id`` to check identity against (callers that don't care about D2 + at all -- most of this module's own test suite -- simply never pass + ``failure_category``, so this stays a no-op for them, unchanged from + before D2 existed). + + Identity check, ATOMIC with the caller's write: takes ``lock_guard_user`` + THEN re-reads the user's currently active credential, all in the SAME + transaction ``record_extraction``'s fenced UPDATE commits -- so a + PUT/activate can no longer land between the read and the write. Match + (attempt's credential id+revision == the active row's) -> ``MAX_ATTEMPTS``, + terminalizing the row under the credential that actually produced this + failure. Mismatch -- a rotation raced this call -- -> ``0``, so the claim + this attempt already spent doesn't strand the row capped under a + credential nobody's using anymore; the row gets a fresh claimable life + under whichever credential is active now (fenced on the still-owned + claim token, same as any other `record_extraction` write). + + A ``None`` ``credential_id`` (the operator/fallback key, never a + per-user row a PUT/activate can rotate) has nothing to compare identity + against, so the cap always applies straight -- no rotation is possible + without a row to rotate. + """ + if failure_category not in CREDENTIAL_CLASS_FAILURE_CATEGORIES: + return None + if user_id is None: + return None + lock_guard_user(db, user_id) + if credential_id is None: + return MAX_ATTEMPTS + current = db.execute( + select(UserLlmCredential.id, UserLlmCredential.revision).where( + UserLlmCredential.user_id == user_id, UserLlmCredential.is_active + ) + ).first() + if ( + current is not None + and current.id == credential_id + and current.revision == credential_revision + ): + return MAX_ATTEMPTS + return 0 def _done_stamp_columns(thread_done: bool, now: datetime) -> dict: @@ -155,6 +287,7 @@ def claim_action_item( user_id: UUID, thread_done: bool, force: bool = False, + model_version: str | None = None, ) -> UUID | None: """Atomically claim ``message_id``'s ``action_item`` row for one extraction attempt. @@ -166,6 +299,17 @@ def claim_action_item( between claim and record must never leave an open row on a done thread, where it would surface as an obligation after an auto-reopen. A forced reclaim of a dismissed row must preserve ``dismissed`` byte-for-byte. + + ``model_version`` (plan D5) is threaded straight to this claim's own + ``claimable_predicate(force=force, model_version=model_version)`` -- + without it, the claim's WHERE clause would compare against ``NULL`` + (``IS DISTINCT FROM NULL``, true for any non-null stored value), a much + WEAKER gate than the one the caller's candidate SELECT already applied + to pick this row, and every settled row with SOME stored model_version + would claim regardless of whether it actually differs from the current + one. The candidate SELECT and this claim must agree on what "current" + means, or a message the SELECT rejected as same-model could still claim + here. """ now = datetime.now(timezone.utc) token = uuid4() @@ -190,7 +334,25 @@ def claim_action_item( .values(**insert_values) .on_conflict_do_nothing(constraint="uq_action_item_message") ) - if db.execute(insert_stmt).rowcount: + # D4 (plan: 2026-08-19-extraction-cost-hardening, D1's actual root + # cause): preserve_rowcount matters here exactly like it does in + # upsert_classification -- without it, psycopg3 reports -1 (truthy!) + # for an INSERT ... ON CONFLICT DO NOTHING statement even when the + # conflict path fired and ZERO rows were actually inserted. Two + # concurrent first-ever claims on the SAME brand-new message (the + # owner's two mail connections firing concurrent sweeps) would both + # read a truthy rowcount here and BOTH believe they'd won the claim -- + # the loser gets handed a claim_token that doesn't correspond to any + # row in the DB, spends a real wire call on it, and then fences out at + # record time (0 rows updated) as bucket "skipped" carrying a real + # failure_category. That's the exact prod storm signature: the + # existing consecutive-FAILURE breaker never sees it (it's a "skipped" + # bucket, not "failed"), so nothing bounds the loser's connection to a + # handful of calls the way D3 now does for a genuine failure streak. + # Repro: test_extraction_concurrency_integration.py's + # test_two_concurrent_claims_on_a_brand_new_row_self_limits fails + # without this line and passes with it. + if db.execute(insert_stmt.execution_options(preserve_rowcount=True)).rowcount: return token # Someone already has a row for this message -- try to (re)claim it. @@ -203,7 +365,10 @@ def claim_action_item( } claim_stmt = ( update(ActionItem) - .where(ActionItem.message_id == message_id, claimable_predicate(force=force)) + .where( + ActionItem.message_id == message_id, + claimable_predicate(force=force, model_version=model_version), + ) .values(**update_values) ) if db.execute(claim_stmt).rowcount: @@ -237,6 +402,11 @@ def record_extraction( thread_done: bool, label_still_actionable: bool, already_replied: bool = False, + model_version: str | None = None, + user_id: UUID | None = None, + failure_category: str | None = None, + credential_id: UUID | None = None, + credential_revision: int | None = None, ) -> bool: """Fenced write of an extraction attempt's outcome. @@ -252,7 +422,33 @@ def record_extraction( result decides: ``ExtractedAction`` -> ``extracted``, ``NoAction`` -> ``no_action`` (extraction fields cleared -- a forced re-extraction that finds nothing must not leave stale fields from a prior attempt), ``None`` - -> ``failed`` (fields untouched, transient failure). + -> ``failed`` (fields untouched, transient failure -- except for a + credential-class category, see below). + + ``model_version`` (plan D6) is the attempt's ``f"{provider}:{model}"``, + written on BOTH settled branches: ``extracted`` keeps preferring + ``result.model_version`` (the extractor always sets it to this same + value, so this is not a behavior change -- see + ``test_record_extraction_extracted_model_version_matches_call_context``), + and ``no_action`` -- which carries no model version of its own -- now + stamps the parameter instead of hard-writing ``None``. This is what lets + D5's ``claimable_predicate(force=True)`` tell "this settled row came from + the CURRENT model" from "the model has since changed" via + ``IS DISTINCT FROM``. + + ``failure_category``/``user_id``/``credential_id``/``credential_revision`` + (plan D2) only matter on the ``result is None`` (failed) branch, and only + when ``failure_category`` is one of ``CREDENTIAL_CLASS_FAILURE_CATEGORIES`` + -- every retry with the SAME credential would fail identically, so this + jumps ``attempts`` straight to ``MAX_ATTEMPTS`` (terminal) instead of + leaving it for the caller to burn two more calls on the same dead + credential. The identity check is atomic with this write (see + ``_credential_class_cap_jump_attempts``): a mismatch -- caught mid-flight + by a credential rotation -- writes ``attempts=0`` instead, so the claim + this attempt already spent doesn't strand the row capped under a + credential nobody's using anymore. Callers that don't pass + ``failure_category`` (most of this module's own test suite) get the + exact same ``failed``-with-fields-untouched behavior as before D2. Independent of which branch above fires: this record also resolves the row's existing ``"open"`` status to ``"done"`` when either ``thread_done`` @@ -295,10 +491,19 @@ def record_extraction( amount=None, currency=None, source_confidence=None, - model_version=None, + model_version=model_version, ) else: values["outcome"] = "failed" + cap_jump = _credential_class_cap_jump_attempts( + db, + user_id=user_id, + failure_category=failure_category, + credential_id=credential_id, + credential_revision=credential_revision, + ) + if cap_jump is not None: + values["attempts"] = cap_jump resolves_already_answered = ( already_replied and isinstance(result, ExtractedAction) and result.kind == "reply" diff --git a/apps/api/app/services/nlp/providers.py b/apps/api/app/services/nlp/providers.py index 7a0ffa9..e24e082 100644 --- a/apps/api/app/services/nlp/providers.py +++ b/apps/api/app/services/nlp/providers.py @@ -68,10 +68,21 @@ class CallContext: """Who a resolved credential bills, for the usage-recording call sites Wave 2b wires up. Kept here, next to `LlmCredential`, rather than in `llm_client` -- that module owns transport and security only, never - `user_id` or billing attribution (see plan §3).""" + `user_id` or billing attribution (see plan §3). + + `credential_id`/`revision` (plan: 2026-08-19-extraction-cost-hardening + D2) are the resolved row's identity at RESOLVE time -- `None` for a + fallback/operator credential, which has no per-user row to rotate. + `extraction_run` threads these through to `record_extraction`'s + credential-class cap-jump check, so a PUT/activate that rotates the + credential mid-sweep can be told apart from the one that actually + produced a given failure. + """ credential: LlmCredential payer: Literal["user", "operator"] + credential_id: UUID | None = None + revision: int | None = None PROVIDER_PRESETS: dict[str, str] = { diff --git a/apps/api/tests/test_actions_integration.py b/apps/api/tests/test_actions_integration.py index 1958ce8..615ec65 100644 --- a/apps/api/tests/test_actions_integration.py +++ b/apps/api/tests/test_actions_integration.py @@ -243,7 +243,7 @@ def test_extraction_run_stamps_action_item_user_id_from_the_locked_thread( provider="openai", base_url="https://example.invalid/v1", api_key="k", model="m" ) - bucket, derived_user_id = extraction_run._claim_extract_record( + bucket, derived_user_id, _category = extraction_run._claim_extract_record( db_session, message.id, call_context=CallContext(credential=stub_credential, payer="operator"), diff --git a/apps/api/tests/test_classifier.py b/apps/api/tests/test_classifier.py index 3ce39b3..c99183b 100644 --- a/apps/api/tests/test_classifier.py +++ b/apps/api/tests/test_classifier.py @@ -507,6 +507,39 @@ def fake_call(cred, *, prompt, user_content, max_tokens, timeout, policy=None): assert captured["timeout"] == classifier._CLASSIFICATION_TIMEOUT_S +@pytest.mark.parametrize( + "content_len,expected_text_len", + [(2999, 2999), (3000, 3000), (3001, 3000)], +) +def test_classify_user_content_truncates_at_classification_text_max_len( + monkeypatch, content_len, expected_text_len +): + """D9 (extraction-cost-hardening): classification's own cap is 3,000 + chars, tighter than extraction's 6,000 -- pins the exact truncation + boundary rather than just trusting the slice syntax.""" + from app.services.nlp import classifier + + monkeypatch.setattr(classifier.settings, "classifier_backend", "llm") + + credential = LlmCredential( + provider="openai", base_url="https://api.openai.com/v1", api_key="user-key", model="gpt-4o-mini" + ) + captured = {} + + def fake_call(cred, *, prompt, user_content, max_tokens, timeout, policy=None): + captured["user_content"] = user_content + content = json.dumps({"label": "fyi", "confidence": 0.5, "rationale": "r"}) + return LlmCallResult(content=content, usage=None) + + monkeypatch.setattr(classifier, "call_chat_completion", fake_call) + + routing = ClassificationRouting(mode="user", credential=credential) + classify("x" * content_len, routing=routing) + + body = captured["user_content"].removeprefix("Email:\n") + assert len(body) == expected_text_len + + def test_classify_with_usage_threads_the_caller_supplied_policy_to_the_wire_call(monkeypatch): """The entry-point-assignment contract (plan: phase 3 of the LLM-failure work): whatever RetryPolicy the caller passes to classify_with_usage() diff --git a/apps/api/tests/test_extraction_concurrency_integration.py b/apps/api/tests/test_extraction_concurrency_integration.py new file mode 100644 index 0000000..7807263 --- /dev/null +++ b/apps/api/tests/test_extraction_concurrency_integration.py @@ -0,0 +1,468 @@ +"""Real-Postgres coverage for the extraction retry-storm investigation +(plan: docs/plans/2026-08-19-extraction-cost-hardening-plan.md, D1/D2/D4). + +Same idiom as test_llm_credentials_integration.py: opt-in via +TEST_DATABASE_URL, schema built with `Base.metadata.create_all` against a +`*_test`-only database, real row-lock (and, here, real DBAPI rowcount) +behavior proven against actual Postgres -- test_extraction_run.py's own +module docstring says exactly why this tier exists: no fake `_FakeDB` can +ever catch a real driver's rowcount-reporting quirk, because a fake's +rowcount is always whatever the test canned it to be. This file IS that +live QA, promoted to the CI-gated real-Postgres tier. + +D1's actual finding, established here (NOT the two-connection race the +plan's own candidate-path framing led with -- that turned out to be a +symptom, not the mechanism): `claim_action_item`'s fresh-row INSERT +(`ON CONFLICT DO NOTHING`) never set `execution_options(preserve_rowcount= +True)`. Under psycopg3, an INSERT ... ON CONFLICT statement that takes the +DO-NOTHING path (the conflict fired, zero rows actually inserted) reports +`rowcount == -1` -- truthy -- without that execution option, exactly the +same trap `upsert_classification` already documents and guards against a +few functions up in this same module. `claim_action_item` treated any +truthy rowcount as "my insert won" and returned a claim_token, so EVERY +reclaim of an EXISTING action_item row -- no concurrency required at all, +proven single-threaded below -- silently failed: the INSERT always +conflicts (the row already exists), always reports a phantom "success", and +`claim_action_item` returns BEFORE it ever reaches the real UPDATE that +would have reclaimed the row and bumped `attempts`. The caller believes it +holds a valid claim, spends a real wire call, and then `record_extraction`'s +fenced UPDATE affects zero rows (the row's real `claim_token` is `None`, +its real `outcome` is `'failed'`, neither matches the phantom token) -- +bucket `"skipped"`, with the attempt's genuine `failure_category` still +attached. That's the EXACT prod triple: `skipped` bucket + populated +`failure_categories` + `attempts` that never advance past their first-ever +value, because the write that would have advanced them never happened. + +`test_reclaim_of_an_existing_failed_row_is_never_silently_dropped` proves +this directly, one session, zero threads. The remaining two tests show it +also governs what happens under REAL concurrency (two connections' sweeps +racing the same candidates, per the operator's setup) -- and that D2 bounds +whatever damage a single genuine failure can still do once D4's fix is in +place. +""" + +from __future__ import annotations + +import os +import threading +from datetime import datetime, timedelta, timezone +from urllib.parse import urlparse +from uuid import uuid4 + +import pytest +from sqlalchemy import create_engine, text, update +from sqlalchemy.orm import sessionmaker + +from app.db.base import Base +from app.db.models import ( + ActionItem, + AppUser, + Classification, + MailMessage, + MailThread, + ProviderAccount, + UserLlmCredential, +) +from app.services.nlp import extraction_run, persistence +from app.services.nlp.extractor import ExtractionAttempt +from app.services.nlp.providers import CallContext, LlmCredential + +TEST_DATABASE_URL = os.environ.get("TEST_DATABASE_URL") + +pytestmark = pytest.mark.skipif( + not TEST_DATABASE_URL, + reason="set TEST_DATABASE_URL to run the real-Postgres integration tier", +) + +_JOIN_TIMEOUT = 10.0 +_EVENT_TIMEOUT = 5.0 + + +@pytest.fixture(scope="session") +def _engine(): + """Build the schema once per test run against a real Postgres. + + Refuses anything but a `*_test` database up front -- every test in this + module truncates app_user, and a typo'd env var pointing at dev or prod + must fail loudly here, not quietly wipe real data. + """ + db_name = urlparse(TEST_DATABASE_URL).path.lstrip("/") + if not db_name.endswith("_test"): + raise RuntimeError( + f"TEST_DATABASE_URL points at {db_name!r}, not a *_test database -- " + "refusing to run tests that TRUNCATE tables against it" + ) + engine = create_engine(TEST_DATABASE_URL) + with engine.begin() as conn: + conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector")) + Base.metadata.create_all(engine) + try: + yield engine + finally: + engine.dispose() + + +@pytest.fixture +def session_factory(_engine): + """A fresh `sessionmaker` per test, with a clean `app_user` table. + TRUNCATE ... CASCADE takes provider_account/mail_thread/mail_message/ + classification/action_item/user_llm_credential with it, so each test + starts from empty. Every worker thread below opens its OWN session from + this factory -- never sharing a session across threads. + """ + factory = sessionmaker(bind=_engine, autoflush=False, autocommit=False) + setup = factory() + setup.execute(text("TRUNCATE app_user CASCADE")) + setup.commit() + setup.close() + return factory + + +def _seed(factory): + """One user, one active BYOK credential, one thread, one actionable + message with no existing action_item row (a never-claimed candidate -- + the exact shape the prod storm's fresh-mail candidates start as). + Returns (user_id, credential_id, credential_revision, message_id). + """ + db = factory() + try: + user = AppUser(email=f"{uuid4()}@example.com") + db.add(user) + db.flush() + + credential = UserLlmCredential( + user_id=user.id, + name="default", + is_active=True, + provider="openai", + base_url="https://api.openai.com/v1", + api_key="sk-test-key-1234", + model="gpt-4o-mini", + revision=1, + ) + db.add(credential) + db.flush() + + account = ProviderAccount( + user_id=user.id, provider="gmail", external_user_id="op@example.com", + access_token="access", + ) + db.add(account) + db.flush() + + thread = MailThread( + user_id=user.id, provider_account_id=account.id, provider="gmail", + provider_thread_id=str(uuid4()), subject="Invoice due", + ) + db.add(thread) + db.flush() + + message = MailMessage( + thread_id=thread.id, provider_message_id=str(uuid4()), + snippet="hi", body_text="pay by friday", + ) + db.add(message) + db.flush() + db.add( + Classification( + message_id=message.id, label="needs_reply", confidence=0.9, + model_version="test", + ) + ) + db.commit() + return user.id, credential.id, credential.revision, message.id + finally: + db.close() + + +def _call_context(credential_id, revision, *, model="gpt-4o-mini") -> CallContext: + return CallContext( + credential=LlmCredential( + provider="openai", base_url="https://api.openai.com/v1", + api_key="sk-test-key-1234", model=model, + ), + payer="user", + credential_id=credential_id, + revision=revision, + ) + + +def _row_state(factory, message_id): + db = factory() + try: + row = db.query(ActionItem).filter(ActionItem.message_id == message_id).one() + return { + "outcome": row.outcome, + "attempts": row.attempts, + "claim_token": row.claim_token, + } + finally: + db.close() + + +def test_reclaim_of_an_existing_failed_row_is_never_silently_dropped(session_factory): + """D1's cornerstone repro -- and it needs no concurrency at all. + + One session, sequential, no threads: claim a brand-new row (a genuine + INSERT), settle it to `failed` (a normal 404), then claim it again -- + the ordinary "next sweep retries a failed row" path every extraction + sweep relies on. Before the D4 fix (missing `preserve_rowcount=True` + on `claim_action_item`'s INSERT), the second `claim_action_item` call + returned a real-looking `claim_token` while the row underneath was + COMPLETELY UNTOUCHED -- still `outcome='failed'`, `attempts=1`, + `claim_token=None`. Confirmed directly against this DB before the fix + landed (git-stashed persistence.py, same script this test encodes): a + second claim's `rowcount` on the conflicting INSERT came back truthy + (psycopg3's -1-for-ON-CONFLICT-DO-NOTHING quirk), so `claim_action_item` + returned early and never reached the UPDATE that would have actually + reclaimed the row. + + With the fix, this claim genuinely reclaims: `outcome='pending'`, + `attempts=2`, `claim_token` matching what was just returned. That's the + whole D1/D4 story -- every retry of every failed extraction row was + silently a no-op, which is exactly "the attempts-cap machinery... is + demonstrably NOT terminalizing these" from the plan's own prod evidence. + """ + user_id, credential_id, revision, message_id = _seed(session_factory) + + db = session_factory() + try: + token1 = persistence.claim_action_item( + db, message_id=message_id, thread_id=db.get(MailMessage, message_id).thread_id, + user_id=user_id, thread_done=False, + ) + db.commit() + finally: + db.close() + + db = session_factory() + try: + recorded = persistence.record_extraction( + db, message_id=message_id, claim_token=token1, result=None, + thread_done=False, label_still_actionable=True, + ) + db.commit() + finally: + db.close() + assert recorded is True + + state = _row_state(session_factory, message_id) + assert state == {"outcome": "failed", "attempts": 1, "claim_token": None} + + # The retry every sweep after this one relies on: reclaim the same, + # now-`failed`, row. + db = session_factory() + try: + token2 = persistence.claim_action_item( + db, message_id=message_id, thread_id=db.get(MailMessage, message_id).thread_id, + user_id=user_id, thread_done=False, + ) + db.commit() + finally: + db.close() + + assert token2 is not None + assert token2 != token1 + + state = _row_state(session_factory, message_id) + # This is the assertion that fails without the D4 fix: state stays + # {"failed", 1, None} -- claim_action_item lied about reclaiming it. + assert state == {"outcome": "pending", "attempts": 2, "claim_token": token2} + + +def test_two_concurrent_claims_on_a_brand_new_row_self_limits(monkeypatch, session_factory): + """The concurrent variant of the SAME D4 bug: two REAL sessions, in two + REAL threads, racing `_claim_extract_record` on the SAME never-claimed + message at essentially the same instant (a `threading.Barrier`, not a + scripted pause). A genuine simultaneous INSERT race is a real Postgres + conflict too -- the loser's insert blocks until the winner commits, then + resolves as a conflict -- so this is the SAME missing-`preserve_rowcount` + defect, just triggered by two connections instead of one session's + retry. Without the fix, the loser's conflicting INSERT ALSO reports a + phantom truthy rowcount and returns a bogus token, wasting a real wire + call and fencing out at record time -- bucket "skipped" with a real + category, never "skipped" with `None`. With the fix, the loser's insert + correctly resolves to a genuine claim failure BEFORE any wire call -- + bucket "skipped", category `None`, exactly matching what Codex's + original source trace predicted for a clean race (plan changelog v2). + """ + user_id, credential_id, revision, message_id = _seed(session_factory) + + def fake_extract(*, credential, **kwargs): + return ExtractionAttempt( + result=None, provider_call_succeeded=True, usage=None, + llm_attempted=True, fallback_used=False, failure_category="http_404", + ) + + monkeypatch.setattr(extraction_run, "extract_action_with_usage", fake_extract) + + barrier = threading.Barrier(2) + results = {} + errors = {} + + def worker(name): + db = session_factory() + try: + barrier.wait(timeout=_EVENT_TIMEOUT) + categories: dict = {} + bucket, uid, category = extraction_run._claim_extract_record( + db, message_id, + call_context=_call_context(credential_id, revision), + failure_categories=categories, + ) + results[name] = (bucket, category) + except BaseException as exc: + errors[name] = exc + finally: + db.close() + + threads = [threading.Thread(target=worker, args=(name,)) for name in ("A", "B")] + for t in threads: + t.start() + for t in threads: + t.join(timeout=_JOIN_TIMEOUT) + + # Catch a hung/still-running thread and a worker exception BEFORE + # touching `results` -- a bare KeyError from a thread that silently + # died is a much worse failure message than surfacing what it raised. + for t in threads: + assert not t.is_alive(), f"worker thread {t.name} did not finish in time" + assert not errors, f"worker thread(s) raised: {errors}" + + outcomes = [results["A"], results["B"]] + winners = [r for r in outcomes if r[0] == "failed"] + losers = [r for r in outcomes if r[0] == "skipped"] + + assert len(winners) == 1, f"expected exactly one winner, got {outcomes}" + assert len(losers) == 1, f"expected exactly one loser, got {outcomes}" + assert winners[0][1] == "http_404" # the winner's own genuine attempt + assert losers[0][1] is None # the loser never reached the wire at all + + # And the row itself: one genuine attempt, D2's cap-jump already applied + # (a first-ever claim's credential-class failure caps straight to + # MAX_ATTEMPTS since there's nothing to have raced a rotation against). + state = _row_state(session_factory, message_id) + assert state["outcome"] == "failed" + assert state["attempts"] == persistence.MAX_ATTEMPTS + + +def test_stale_lease_competing_reclaim_bounded_by_d2(monkeypatch, session_factory): + """A SEPARATE, secondary race from the D4 bug above -- included because + the plan's own D1 candidate list named it, and because it's a genuine + interleaving worth bounding regardless of D4: thread A claims the row + (a genuine first-ever claim, attempts=1), then blocks INSIDE its own + "wire call" (a stubbed `extract_action_with_usage`, gated on a real + `threading.Event` -- not a scripted mock). While A is blocked, the test + backdates the row's `last_attempted_at` past `PENDING_LEASE` directly in + Postgres, simulating exactly what an orphaned/crashed worker's claim + looks like after 30 minutes: a live claim nobody is coming back to + finish. Thread B then reclaims it for real (Postgres's own claimable_ + predicate, not a stub) -- a completely normal expired-pending reclaim, + attempts -> 2 -- and completes its OWN failed attempt, recording + normally. Only THEN is A released to finish its stale attempt and try to + record with its now-superseded claim_token. + + This ALSO produces the skipped-bucket-with-a-real-category combination + on A's own return, but via a genuinely different mechanism than D4's bug + (a truly expired lease, not a misreported rowcount) -- it needs an + already-orphaned claim to exist first (a crashed/killed worker), which + is a MUCH rarer precondition than "a second reclaim of any failed row," + the D4 bug's trigger. This test's real job is proving D2 bounds it + either way: B's own record (a genuine credential-class failure, same + active credential the whole time -- no rotation) cap-jumps the row + straight to MAX_ATTEMPTS. The row is terminal after this ONE extra + wasted call (B's), not a storm that repeats every 30 minutes. + """ + user_id, credential_id, revision, message_id = _seed(session_factory) + + release_a = threading.Event() + a_blocked = threading.Event() + + def fake_extract(*, credential, **kwargs): + if threading.current_thread().name == "A": + a_blocked.set() + if not release_a.wait(timeout=_EVENT_TIMEOUT): + raise AssertionError("main thread never released A") + return ExtractionAttempt( + result=None, provider_call_succeeded=True, usage=None, + llm_attempted=True, fallback_used=False, failure_category="http_404", + ) + + monkeypatch.setattr(extraction_run, "extract_action_with_usage", fake_extract) + + results = {} + errors = {} + + def run_a(): + db = session_factory() + try: + bucket, uid, category = extraction_run._claim_extract_record( + db, message_id, + call_context=_call_context(credential_id, revision), + failure_categories={}, + ) + results["A"] = (bucket, category) + except BaseException as exc: + errors["A"] = exc + finally: + db.close() + + thread_a = threading.Thread(target=run_a, name="A") + thread_a.start() + assert a_blocked.wait(timeout=_EVENT_TIMEOUT), "A never reached its wire call" + + # A's claim has committed (the claim transaction closes before the wire + # call runs) and is sitting there genuinely 'pending' -- back-date it + # past the lease, exactly like a claim a crashed worker never returned + # to finish would look after 30 real minutes. + backdate_db = session_factory() + try: + backdate_db.execute( + update(ActionItem) + .where(ActionItem.message_id == message_id) + .values( + last_attempted_at=datetime.now(timezone.utc) + - persistence.PENDING_LEASE + - timedelta(minutes=1) + ) + ) + backdate_db.commit() + finally: + backdate_db.close() + + # B reclaims for real -- a completely ordinary expired-pending reclaim, + # no stubbing of the claim/record path at all. + db_b = session_factory() + try: + bucket_b, uid_b, category_b = extraction_run._claim_extract_record( + db_b, message_id, + call_context=_call_context(credential_id, revision), + failure_categories={}, + ) + finally: + db_b.close() + + # Now let A's stale attempt finish and try to record. + release_a.set() + thread_a.join(timeout=_JOIN_TIMEOUT) + + assert not thread_a.is_alive(), "thread A did not finish in time" + assert not errors, f"worker thread(s) raised: {errors}" + + # B: a normal, successful record of a genuine failure. + assert bucket_b == "failed" + assert category_b == "http_404" + + # A: the D1 triple -- skipped bucket, real category, fenced out by B's + # reclaim. + assert "A" in results, "thread A never completed" + bucket_a, category_a = results["A"] + assert bucket_a == "skipped" + assert category_a == "http_404" + + # The row itself: D2's cap-jump already terminalized it off B's write -- + # bounded to the one extra call B just spent, not a repeating storm. + state = _row_state(session_factory, message_id) + assert state["outcome"] == "failed" + assert state["attempts"] == persistence.MAX_ATTEMPTS + assert state["claim_token"] is None diff --git a/apps/api/tests/test_extraction_run.py b/apps/api/tests/test_extraction_run.py index c358398..26dd9a0 100644 --- a/apps/api/tests/test_extraction_run.py +++ b/apps/api/tests/test_extraction_run.py @@ -93,6 +93,9 @@ def scalars(self): def all(self): return self._rows + def first(self): + return self._rows[0] if self._rows else None + class _FakeDB: """Records every statement persistence.py hands to execute() and answers @@ -283,8 +286,21 @@ def test_claimable_predicate_sql_shape(): assert "outcome = 'pending' AND action_item.last_attempted_at <" in default_sql assert "!=" not in default_sql - forced_sql = _compiled_sql(persistence.claimable_predicate(force=True)) - assert "outcome != 'pending'" in forced_sql + # D5: force no longer widens to "any non-pending row" -- `failed` is + # claimable regardless of the attempt cap, and a settled + # extracted/no_action row is claimable only when its stored + # model_version differs from the CURRENT one (IS DISTINCT FROM, not a + # NULL-poisoned !=). + forced_sql = _compiled_sql( + persistence.claimable_predicate(force=True, model_version="gemini:gemini-x") + ) + assert "outcome = 'failed'" in forced_sql + assert "action_item.attempts < 3" in forced_sql # still present for the UNforced disjunct + assert "outcome != 'pending'" not in forced_sql # a live pending row is never stolen + assert ( + "action_item.outcome IN ('extracted', 'no_action') " + "AND action_item.model_version IS DISTINCT FROM 'gemini:gemini-x'" + ) in forced_sql # --------------------------------------------------------------------------- @@ -358,13 +374,19 @@ def test_claim_action_item_force_never_steals_live_pending(): db = _FakeDB([_FakeResult(rowcount=0), _FakeResult(rowcount=1)]) persistence.claim_action_item( db, message_id=uuid4(), thread_id=uuid4(), user_id=uuid4(), - thread_done=False, force=True, + thread_done=False, force=True, model_version="gemini:gemini-x", ) where_sql = _compiled_sql(db.statements[1]) - # force widens to "any non-pending row" -- a live (unexpired) pending row - # is deliberately absent from every disjunct, forced or not. - assert "outcome != 'pending'" in where_sql + # D5: force no longer widens to "any non-pending row" -- a live + # (unexpired) pending row is deliberately absent from every disjunct, + # forced or not, and neither is a settled row on the SAME model. + assert "outcome != 'pending'" not in where_sql assert "outcome = 'pending'" in where_sql # still present: the expired-lease branch + assert "outcome = 'failed'" in where_sql + assert ( + "action_item.outcome IN ('extracted', 'no_action') " + "AND action_item.model_version IS DISTINCT FROM 'gemini:gemini-x'" + ) in where_sql def test_claim_action_item_reclaim_on_done_thread_preserves_dismissed(): @@ -479,6 +501,191 @@ def test_record_extraction_label_mismatch_wins_as_ineligible(): assert "kind" not in params +def test_record_extraction_extracted_model_version_matches_call_context(): + # D6: the extracted branch keeps preferring result.model_version -- this + # asserts they're the SAME value a real caller produces (both derived + # from the same call_context.credential), rather than silently letting + # the two diverge. + db = _FakeDB([_FakeResult(rowcount=1)]) + call_context_model_version = "gemini:gemini-x" + + persistence.record_extraction( + db, message_id=uuid4(), claim_token=uuid4(), + result=_extracted_action(model_version=call_context_model_version), + thread_done=False, label_still_actionable=True, + model_version=call_context_model_version, + ) + + params = _compiled_params(db.statements[0]) + assert params["model_version"] == call_context_model_version + + +def test_record_extraction_no_action_writes_passed_model_version(): + # D6: no_action no longer hard-writes None -- it stamps the caller's + # model_version, same as the extracted branch. + db = _FakeDB([_FakeResult(rowcount=1)]) + + persistence.record_extraction( + db, message_id=uuid4(), claim_token=uuid4(), + result=NoAction(), thread_done=False, label_still_actionable=True, + model_version="gemini:gemini-x", + ) + + params = _compiled_params(db.statements[0]) + assert params["model_version"] == "gemini:gemini-x" + + +# --------------------------------------------------------------------------- +# persistence.record_extraction -- D2 credential-class cap-jump +# --------------------------------------------------------------------------- + + +def test_record_extraction_credential_class_failure_caps_attempts_on_identity_match(): + credential_id = uuid4() + user_id = uuid4() + # [guard-lock (ignored), active-credential re-read, the fenced UPDATE] + db = _FakeDB([ + _FakeResult(), + _FakeResult(rows=[SimpleNamespace(id=credential_id, revision=2)]), + _FakeResult(rowcount=1), + ]) + + persistence.record_extraction( + db, message_id=uuid4(), claim_token=uuid4(), + result=None, thread_done=False, label_still_actionable=True, + user_id=user_id, failure_category="http_404", + credential_id=credential_id, credential_revision=2, + ) + + params = _compiled_params(db.statements[-1]) + assert params["outcome"] == "failed" + assert params["attempts"] == persistence.MAX_ATTEMPTS + + +def test_record_extraction_credential_class_failure_resets_attempts_on_identity_mismatch(): + # Rotation race: the row's user has since switched to a DIFFERENT + # credential (different id, or same id but a bumped revision) by the + # time this record transaction re-reads it under the guard lock -- the + # mismatch resets attempts to 0 instead of stranding the row capped + # under a credential nobody's using anymore. + user_id = uuid4() + attempted_credential_id = uuid4() + now_active_credential_id = uuid4() + db = _FakeDB([ + _FakeResult(), + _FakeResult(rows=[SimpleNamespace(id=now_active_credential_id, revision=1)]), + _FakeResult(rowcount=1), + ]) + + persistence.record_extraction( + db, message_id=uuid4(), claim_token=uuid4(), + result=None, thread_done=False, label_still_actionable=True, + user_id=user_id, failure_category="http_401", + credential_id=attempted_credential_id, credential_revision=1, + ) + + params = _compiled_params(db.statements[-1]) + assert params["attempts"] == 0 + + +def test_record_extraction_credential_class_failure_same_id_different_revision_is_a_mismatch(): + user_id = uuid4() + credential_id = uuid4() + db = _FakeDB([ + _FakeResult(), + _FakeResult(rows=[SimpleNamespace(id=credential_id, revision=5)]), + _FakeResult(rowcount=1), + ]) + + persistence.record_extraction( + db, message_id=uuid4(), claim_token=uuid4(), + result=None, thread_done=False, label_still_actionable=True, + user_id=user_id, failure_category="http_403", + credential_id=credential_id, credential_revision=1, # stale revision + ) + + params = _compiled_params(db.statements[-1]) + assert params["attempts"] == 0 + + +def test_record_extraction_credential_class_failure_no_active_row_is_a_mismatch(): + # The credential was deleted entirely (kill switch) between the attempt + # and this record -- no active row at all is treated the same as any + # other identity mismatch: fresh claimable life, not a strand. + user_id = uuid4() + db = _FakeDB([ + _FakeResult(), + _FakeResult(rows=[]), + _FakeResult(rowcount=1), + ]) + + persistence.record_extraction( + db, message_id=uuid4(), claim_token=uuid4(), + result=None, thread_done=False, label_still_actionable=True, + user_id=user_id, failure_category="http_404", + credential_id=uuid4(), credential_revision=1, + ) + + params = _compiled_params(db.statements[-1]) + assert params["attempts"] == 0 + + +def test_record_extraction_credential_class_failure_fallback_credential_caps_straight(): + # No per-user row exists for the operator/fallback key (credential_id is + # None) -- there's nothing a PUT/activate could have rotated, so the cap + # applies directly, no active-credential re-read needed. + user_id = uuid4() + db = _FakeDB([_FakeResult(), _FakeResult(rowcount=1)]) + + persistence.record_extraction( + db, message_id=uuid4(), claim_token=uuid4(), + result=None, thread_done=False, label_still_actionable=True, + user_id=user_id, failure_category="http_404", + credential_id=None, credential_revision=None, + ) + + params = _compiled_params(db.statements[-1]) + assert params["attempts"] == persistence.MAX_ATTEMPTS + # Only two statements: the guard lock and the fenced UPDATE -- no + # active-credential re-read was needed at all. + assert len(db.statements) == 2 + + +def test_record_extraction_non_credential_class_failure_leaves_attempts_untouched(): + # http_429 (rate-limited) is transient and retries on its own -- no + # guard lock, no cap-jump, same as before D2 existed. + db = _FakeDB([_FakeResult(rowcount=1)]) + + persistence.record_extraction( + db, message_id=uuid4(), claim_token=uuid4(), + result=None, thread_done=False, label_still_actionable=True, + user_id=uuid4(), failure_category="http_429", + credential_id=uuid4(), credential_revision=1, + ) + + assert len(db.statements) == 1 # no guard lock, no credential re-read + params = _compiled_params(db.statements[0]) + assert "attempts" not in params + + +def test_record_extraction_credential_class_failure_without_user_id_is_a_noop(): + # A caller that doesn't pass user_id (most of this module's own callers + # in this test file) gets the exact same failed-fields-untouched + # behavior D2 shipped on top of -- no identity check is even possible + # without a user to check it for. + db = _FakeDB([_FakeResult(rowcount=1)]) + + persistence.record_extraction( + db, message_id=uuid4(), claim_token=uuid4(), + result=None, thread_done=False, label_still_actionable=True, + failure_category="http_404", + ) + + assert len(db.statements) == 1 + params = _compiled_params(db.statements[0]) + assert "attempts" not in params + + def test_record_extraction_stale_token_affects_zero_rows(): db = _FakeDB([_FakeResult(rowcount=0)]) @@ -522,7 +729,7 @@ def test_record_extraction_thread_done_stamps_status_on_every_branch( def test_claim_extract_record_missing_message_is_skipped(): session = _FakeSession(message=None) - bucket, user_id = extraction_run._claim_extract_record(session, uuid4()) + bucket, user_id, *_ = extraction_run._claim_extract_record(session, uuid4()) assert (bucket, user_id) == ("skipped", None) assert session.executed == [] @@ -532,7 +739,7 @@ def test_claim_extract_record_missing_thread_is_skipped(): message = _make_message() session = _FakeSession(message=message, thread=None) - bucket, user_id = extraction_run._claim_extract_record(session, message.id) + bucket, user_id, *_ = extraction_run._claim_extract_record(session, message.id) assert (bucket, user_id) == ("skipped", None) @@ -547,7 +754,7 @@ def test_claim_extract_record_not_claimable_skips_before_extraction(monkeypatch) lambda **k: pytest.fail("extraction must not run without a claim"), ) - bucket, user_id = extraction_run._claim_extract_record(session, message.id, call_context=_CTX) + bucket, user_id, *_ = extraction_run._claim_extract_record(session, message.id, call_context=_CTX) assert bucket == "skipped" assert user_id == thread.user_id @@ -625,7 +832,7 @@ def fake_record(db, **kwargs): monkeypatch.setattr(extraction_run, "record_extraction", fake_record) - bucket, _ = extraction_run._claim_extract_record(session, message.id, call_context=_CTX) + bucket, _, *_ = extraction_run._claim_extract_record(session, message.id, call_context=_CTX) classification_stmt = next( s for s in session.executed if s.column_descriptions[0]["name"] == "Classification" @@ -742,7 +949,7 @@ def test_claim_extract_record_stale_record_is_skipped(monkeypatch): ) monkeypatch.setattr(extraction_run, "record_extraction", lambda *a, **k: False) - bucket, _ = extraction_run._claim_extract_record(session, message.id, call_context=_CTX) + bucket, _, *_ = extraction_run._claim_extract_record(session, message.id, call_context=_CTX) assert bucket == "skipped" @@ -770,7 +977,7 @@ def test_claim_extract_record_outcome_buckets(monkeypatch, label, result, expect ) monkeypatch.setattr(extraction_run, "record_extraction", lambda *a, **k: True) - bucket, user_id = extraction_run._claim_extract_record(session, message.id, call_context=_CTX) + bucket, user_id, *_ = extraction_run._claim_extract_record(session, message.id, call_context=_CTX) assert bucket == expected assert user_id == thread.user_id @@ -794,7 +1001,7 @@ def test_claim_extract_record_failure_categories_out_param_records_attempt_categ ) categories: dict[str, int] = {} - bucket, _ = extraction_run._claim_extract_record( + bucket, _, *_ = extraction_run._claim_extract_record( session, message.id, call_context=_CTX, failure_categories=categories ) @@ -818,7 +1025,7 @@ def test_claim_extract_record_failure_categories_none_is_a_silent_no_op(monkeypa lambda **k: _extraction_attempt(None, provider_call_succeeded=False, failure_category="http_429"), ) - bucket, _ = extraction_run._claim_extract_record(session, message.id, call_context=_CTX) + bucket, _, *_ = extraction_run._claim_extract_record(session, message.id, call_context=_CTX) assert bucket == "failed" @@ -837,7 +1044,7 @@ def test_claim_extract_record_success_never_touches_failure_categories(monkeypat ) categories: dict[str, int] = {} - bucket, _ = extraction_run._claim_extract_record( + bucket, _, *_ = extraction_run._claim_extract_record( session, message.id, call_context=_CTX, failure_categories=categories ) @@ -877,7 +1084,7 @@ def _boom(**kwargs): fresh_session = MagicMock() monkeypatch.setattr(extraction_run, "SessionLocal", lambda: nullcontext(fresh_session)) - bucket, user_id = extraction_run._claim_extract_record(session, message.id, call_context=_CTX) + bucket, user_id, *_ = extraction_run._claim_extract_record(session, message.id, call_context=_CTX) assert bucket == "failed" assert user_id == thread.user_id @@ -1071,7 +1278,7 @@ def test_claim_extract_record_failing_usage_flush_does_not_block_business_commit fake_acc = _FakeAcc(thread.user_id, flush_raises=SQLAlchemyError("boom")) - bucket, _ = extraction_run._claim_extract_record( + bucket, _, *_ = extraction_run._claim_extract_record( session, message.id, call_context=_USER_CTX, acc=fake_acc ) @@ -1094,7 +1301,7 @@ def _boom(**kwargs): monkeypatch.setattr(extraction_run, "SessionLocal", lambda: nullcontext(fresh_session)) fake_acc = _FakeAcc(thread.user_id) - bucket, _ = extraction_run._claim_extract_record( + bucket, _, *_ = extraction_run._claim_extract_record( session, message.id, call_context=_USER_CTX, acc=fake_acc ) @@ -1128,7 +1335,8 @@ def test_run_extraction_for_message_delegates_and_includes_user_id(monkeypatch): _enable_extraction(monkeypatch) user_id = uuid4() monkeypatch.setattr( - extraction_run, "_claim_extract_record", lambda db, mid, force=False: ("extracted", user_id) + extraction_run, "_claim_extract_record", + lambda db, mid, force=False: ("extracted", user_id, None), ) result = extraction_run.run_extraction_for_message(MagicMock(), uuid4()) @@ -1284,6 +1492,34 @@ def test_run_extraction_sweep_recovery_uses_row_driven_candidates(monkeypatch): extraction_run.run_extraction_sweep(MagicMock(), uuid4(), recovery=True) +def test_run_extraction_sweep_breaker_trips_during_recovery_too(monkeypatch): + """D3's streak logic isn't message-driven-only -- a recovery=True sweep + (the beat tick's row-driven pass) hits the exact same loop, so three + consecutive credential-class failures abort it identically.""" + _enable_extraction(monkeypatch) + monkeypatch.setattr( + extraction_run, "resolve_extraction_credential", + lambda db, uid: SimpleNamespace(credential=_CRED, source="fallback"), + ) + monkeypatch.setattr(extraction_run, "terminalize_expired_pending", lambda *a, **k: None) + message_ids = [uuid4(), uuid4(), uuid4(), uuid4()] + monkeypatch.setattr(extraction_run, "_recovery_candidates", lambda *a, **k: message_ids) + + calls = [] + + def fake_claim(db, mid, force=False, call_context=None, acc=None, failure_categories=None, deadline=None): + calls.append(mid) + return "failed", uuid4(), "http_404" + + monkeypatch.setattr(extraction_run, "_claim_extract_record", fake_claim) + + result = extraction_run.run_extraction_sweep(MagicMock(), uuid4(), recovery=True) + + assert calls == message_ids[:3] + assert result["status"] == "llm_unavailable" + assert result["aborted"] == 1 + + def test_run_extraction_sweep_counts_a_bucket_per_message(monkeypatch): _enable_extraction(monkeypatch) monkeypatch.setattr( @@ -1297,7 +1533,7 @@ def test_run_extraction_sweep_counts_a_bucket_per_message(monkeypatch): monkeypatch.setattr( extraction_run, "_claim_extract_record", lambda db, mid, force=False, call_context=None, acc=None, failure_categories=None, deadline=None: ( - next(buckets), uuid4() + next(buckets), uuid4(), None ), ) @@ -1370,7 +1606,7 @@ def test_run_extraction_sweep_isolates_one_bad_message_and_continues(monkeypatch def fake_claim(db, mid, force=False, call_context=None, acc=None, failure_categories=None, deadline=None): if mid == bad_message: raise RuntimeError("boom") - return "extracted", uuid4() + return "extracted", uuid4(), None monkeypatch.setattr(extraction_run, "_claim_extract_record", fake_claim) db = MagicMock() @@ -1395,8 +1631,10 @@ def fake_claim(db, mid, force=False, call_context=None, acc=None, failure_catego def test_run_extraction_sweep_stops_after_three_consecutive_failures_with_partial_counts( monkeypatch, ): - """Four candidates, every one a "failed" bucket -- the sweep stops after - the third and never even attempts the fourth.""" + """Four candidates, every one a "failed" bucket with a credential-class + category (D3: the streak reads the category, not the bucket) -- the + sweep stops after the third and never even attempts the fourth, and + reports how many it skipped via the additive `aborted` count.""" _enable_extraction(monkeypatch) monkeypatch.setattr( extraction_run, "resolve_extraction_credential", @@ -1410,7 +1648,7 @@ def test_run_extraction_sweep_stops_after_three_consecutive_failures_with_partia def fake_claim(db, mid, force=False, call_context=None, acc=None, failure_categories=None, deadline=None): calls.append(mid) - return "failed", uuid4() + return "failed", uuid4(), "http_404" monkeypatch.setattr(extraction_run, "_claim_extract_record", fake_claim) @@ -1423,14 +1661,16 @@ def fake_claim(db, mid, force=False, call_context=None, acc=None, failure_catego assert result["status"] == "llm_unavailable" assert result["processed"] == 3 assert result["failed"] == 3 + assert result["aborted"] == 1 # the fourth candidate, never even attempted def test_run_extraction_sweep_isolated_failures_among_successes_never_stop_and_counter_resets( monkeypatch, ): """Pattern: failed, extracted, failed, failed, extracted -- 3 failures - total but the longest CONSECUTIVE run is 2, under the threshold, so all - 5 candidates get attempted.""" + total but the longest CONSECUTIVE run of credential-class CATEGORIES + (D3: the streak reads the category, not the bucket) is 2, under the + threshold, so all 5 candidates get attempted.""" _enable_extraction(monkeypatch) monkeypatch.setattr( extraction_run, "resolve_extraction_credential", @@ -1441,9 +1681,10 @@ def test_run_extraction_sweep_isolated_failures_among_successes_never_stop_and_c monkeypatch.setattr(extraction_run, "_message_driven_candidates", lambda *a, **k: message_ids) buckets = iter(["failed", "extracted", "failed", "failed", "extracted"]) + categories = iter(["http_404", None, "http_404", "http_404", None]) def fake_claim(db, mid, force=False, call_context=None, acc=None, failure_categories=None, deadline=None): - return next(buckets), uuid4() + return next(buckets), uuid4(), next(categories) monkeypatch.setattr(extraction_run, "_claim_extract_record", fake_claim) @@ -1455,6 +1696,66 @@ def fake_claim(db, mid, force=False, call_context=None, acc=None, failure_catego assert result["extracted"] == 2 +def test_run_extraction_sweep_three_consecutive_rate_limits_do_not_trip_the_breaker(monkeypatch): + """D3: http_429 is retryable and recovers on its own -- three consecutive + "failed" buckets with a non-credential-class category must NOT trip the + breaker, even though the OLD bucket-only breaker would have stopped + here. All 4 candidates get attempted, status "ok".""" + _enable_extraction(monkeypatch) + monkeypatch.setattr( + extraction_run, "resolve_extraction_credential", + lambda db, uid: SimpleNamespace(credential=_CRED, source="fallback"), + ) + monkeypatch.setattr(extraction_run, "terminalize_expired_pending", lambda *a, **k: None) + message_ids = [uuid4(), uuid4(), uuid4(), uuid4()] + monkeypatch.setattr(extraction_run, "_message_driven_candidates", lambda *a, **k: message_ids) + + calls = [] + + def fake_claim(db, mid, force=False, call_context=None, acc=None, failure_categories=None, deadline=None): + calls.append(mid) + return "failed", uuid4(), "http_429" + + monkeypatch.setattr(extraction_run, "_claim_extract_record", fake_claim) + + result = extraction_run.run_extraction_sweep(MagicMock(), uuid4()) + + assert calls == message_ids # never stopped early + assert result["status"] == "ok" + assert result["processed"] == 4 + assert result["failed"] == 4 + assert result["aborted"] == 0 + + +def test_run_extraction_sweep_credential_class_streak_resets_on_a_success_between(monkeypatch): + """D3: config, config, SUCCESS, config, config -- two credential-class + failures, a success that resets the streak, then two more. The longest + consecutive run is 2, under the threshold, so the breaker never trips + even though there are 4 credential-class failures total.""" + _enable_extraction(monkeypatch) + monkeypatch.setattr( + extraction_run, "resolve_extraction_credential", + lambda db, uid: SimpleNamespace(credential=_CRED, source="fallback"), + ) + monkeypatch.setattr(extraction_run, "terminalize_expired_pending", lambda *a, **k: None) + message_ids = [uuid4() for _ in range(5)] + monkeypatch.setattr(extraction_run, "_message_driven_candidates", lambda *a, **k: message_ids) + + buckets = iter(["failed", "failed", "extracted", "failed", "failed"]) + categories = iter(["http_403", "http_403", None, "http_403", "http_403"]) + + def fake_claim(db, mid, force=False, call_context=None, acc=None, failure_categories=None, deadline=None): + return next(buckets), uuid4(), next(categories) + + monkeypatch.setattr(extraction_run, "_claim_extract_record", fake_claim) + + result = extraction_run.run_extraction_sweep(MagicMock(), uuid4()) + + assert result["processed"] == 5 + assert result["status"] == "ok" + assert result["aborted"] == 0 + + def test_run_extraction_sweep_containment_failure_does_not_trigger_llm_unavailable(monkeypatch): """SHOULD-FIX (Codex review, phase 3): a message that blows up _claim_extract_record entirely (a DB error, or anything even its own @@ -1491,6 +1792,37 @@ def fake_claim(db, mid, force=False, call_context=None, acc=None, failure_catego assert db.rollback.call_count == 4 +def test_run_extraction_sweep_containment_failure_resets_the_credential_class_streak(monkeypatch): + """D3: two credential-class failures, an unrelated containment error, + then two more -- 4 real provider failures total but never 3 IN A ROW, + because the containment error explicitly resets the streak instead of + silently bridging the two runs into one continuous trip.""" + _enable_extraction(monkeypatch) + monkeypatch.setattr( + extraction_run, "resolve_extraction_credential", + lambda db, uid: SimpleNamespace(credential=_CRED, source="fallback"), + ) + monkeypatch.setattr(extraction_run, "terminalize_expired_pending", lambda *a, **k: None) + message_ids = [uuid4() for _ in range(5)] + monkeypatch.setattr(extraction_run, "_message_driven_candidates", lambda *a, **k: message_ids) + + calls = [] + + def fake_claim(db, mid, force=False, call_context=None, acc=None, failure_categories=None, deadline=None): + calls.append(mid) + if len(calls) == 3: + raise RuntimeError("boom") + return "failed", uuid4(), "http_404" + + monkeypatch.setattr(extraction_run, "_claim_extract_record", fake_claim) + + result = extraction_run.run_extraction_sweep(MagicMock(), uuid4()) + + assert calls == message_ids # never stopped early + assert result["status"] == "ok" + assert result["aborted"] == 0 + + def test_run_extraction_sweep_soft_time_limit_propagates_and_stops_the_sweep(monkeypatch): """BLOCKER (Codex review, phase 3): SoftTimeLimitExceeded is a plain Exception, so it must be special-cased ahead of the generic containment @@ -1514,7 +1846,7 @@ def fake_claim(db, mid, force=False, call_context=None, acc=None, failure_catego calls.append(mid) if len(calls) == 2: raise SoftTimeLimitExceeded - return "extracted", uuid4() + return "extracted", uuid4(), None monkeypatch.setattr(extraction_run, "_claim_extract_record", fake_claim) db = MagicMock() @@ -1564,12 +1896,12 @@ def test_run_extraction_sweep_inner_deadline_on_last_candidate_reports_timed_out monkeypatch.setattr(extraction_run, "terminalize_expired_pending", lambda db, **k: None) monkeypatch.setattr( extraction_run, "resolve_extraction_credential", - lambda db, user_id: SimpleNamespace(credential=object(), source="user"), + lambda db, user_id: SimpleNamespace(credential=_CRED, source="user"), ) monkeypatch.setattr(extraction_run, "_message_driven_candidates", lambda *a, **k: message_ids) def fake_claim(db, mid, force=False, call_context=None, acc=None, failure_categories=None, deadline=None): - return "deadline", uuid4() + return "deadline", uuid4(), None monkeypatch.setattr(extraction_run, "_claim_extract_record", fake_claim) @@ -1604,7 +1936,7 @@ def test_run_extraction_sweep_deadline_exceeded_mid_sweep_reports_partial_counts def fake_claim(db, mid, force=False, call_context=None, acc=None, failure_categories=None, deadline=None): calls.append(mid) - return "extracted", uuid4() + return "extracted", uuid4(), None monkeypatch.setattr(extraction_run, "_claim_extract_record", fake_claim) @@ -1613,6 +1945,10 @@ def fake_claim(db, mid, force=False, call_context=None, acc=None, failure_catego assert calls == message_ids[:2] # the third candidate's deadline check stopped it first assert result["status"] == "timed_out" assert result["processed"] == 2 + # Proves the two claimed candidates actually ran the success path, + # not a swallowed unpack error from a stand-in returning the wrong shape. + assert result["extracted"] == 2 + assert result["failed"] == 0 def test_run_extraction_sweep_end_to_end_terminalizes_expired_at_cap_row(monkeypatch): @@ -1688,12 +2024,12 @@ def test_run_extraction_sweep_resolves_credential_once_and_threads_it(monkeypatc def fake_claim(db, mid, force=False, call_context=None, acc=None, failure_categories=None, deadline=None): captured_call_contexts.append(call_context) captured_accs.append(acc) - return "extracted", uuid4() + return "extracted", uuid4(), None monkeypatch.setattr(extraction_run, "_claim_extract_record", fake_claim) user_id = uuid4() - extraction_run.run_extraction_sweep(MagicMock(), user_id) + result = extraction_run.run_extraction_sweep(MagicMock(), user_id) assert resolve_calls == [user_id] # exactly one resolution for the whole run assert [ctx.credential for ctx in captured_call_contexts] == [ @@ -1701,6 +2037,10 @@ def fake_claim(db, mid, force=False, call_context=None, acc=None, failure_catego ] assert [ctx.payer for ctx in captured_call_contexts] == ["operator", "operator"] assert captured_accs[0] is captured_accs[1] # one shared accumulator, not one per message + # Proves both claims actually ran the success path, not a swallowed + # unpack error from a stand-in returning the wrong shape. + assert result["extracted"] == 2 + assert result["failed"] == 0 # --------------------------------------------------------------------------- @@ -1739,11 +2079,20 @@ def test_message_driven_candidates_force_widens_existing_row_filter(): db = _FakeDB([_FakeResult()]) db.execute = lambda stmt: (db.statements.append(stmt), _FakeResult())[1] extraction_run._message_driven_candidates( - db, uuid4(), since_days=30, limit=10, force=True + db, uuid4(), since_days=30, limit=10, force=True, model_version="gemini:gemini-x" ) sql = _compiled_sql(db.statements[0]) - assert "action_item.outcome != 'pending'" in sql + # D5: force claims failed rows unconditionally, and extracted/no_action + # rows only when their stored model_version differs from the current one + # -- no longer "any non-pending row" (see claimable_predicate's own + # tests for the full SQL-shape assertion). + assert "action_item.outcome != 'pending'" not in sql + assert "action_item.outcome = 'failed'" in sql + assert ( + "action_item.outcome IN ('extracted', 'no_action') " + "AND action_item.model_version IS DISTINCT FROM 'gemini:gemini-x'" + ) in sql def test_message_driven_candidates_coalesces_null_sent_at_for_cutoff_and_order(): diff --git a/apps/api/tests/test_extractor.py b/apps/api/tests/test_extractor.py index df9bc8d..af40329 100644 --- a/apps/api/tests/test_extractor.py +++ b/apps/api/tests/test_extractor.py @@ -180,12 +180,16 @@ def test_extract_action_date_only_resolves_to_end_of_day_utc(monkeypatch): assert result.due_at == datetime(2024, 3, 15, 23, 59, 59, tzinfo=timezone.utc) -def test_extract_action_prompt_includes_received_at(monkeypatch): +def test_extract_action_user_content_includes_received_at(monkeypatch): + """D7 (extraction-cost-hardening): received_at moved out of the prompt + and into the per-message user_content (a Received: header line) so the + instruction block stays byte-identical across calls.""" calls = [] monkeypatch.setattr(extractor, "call_chat_completion", _stub_returns(json.dumps(VALID_PAYLOAD), calls=calls)) received_at = datetime(2024, 3, 10, 12, 0, 0, tzinfo=timezone.utc) extract_action(**_default_kwargs(received_at=received_at)) - assert received_at.isoformat() in calls[0]["prompt"] + assert received_at.isoformat() in calls[0]["user_content"] + assert received_at.isoformat() not in calls[0]["prompt"] def test_extract_action_model_version_attribution(monkeypatch): diff --git a/apps/api/tests/test_llm_settings.py b/apps/api/tests/test_llm_settings.py index d0d7403..82d0ed2 100644 --- a/apps/api/tests/test_llm_settings.py +++ b/apps/api/tests/test_llm_settings.py @@ -148,10 +148,23 @@ def __init__(self, rows: dict[str, object] | None = None): # test pin that the guard ran (and ran FIRST) without caring about # the AppUser row's actual contents, which this fake never models. self.guard_locks: list[str] = [] + # user_id hexes `persistence.reset_failed_extraction_attempts` targeted, in call order + # -- lets a D2 recovery test assert whether a mutation reset this + # caller's failed rows without this fake needing to model + # `action_item` rows at all. + self.reset_calls: list[str] = [] def execute(self, stmt): self.statements.append(stmt) if isinstance(stmt, Update): + compiled = _compiled(stmt) + if "action_item" in compiled: + match = re.search(r"action_item\.user_id = '([0-9a-f]{32})'", compiled) + if match: + self.reset_calls.append(match.group(1)) + result = MagicMock() + result.rowcount = 0 + return result result = MagicMock() matched = self._match_update(stmt) if matched is not None: @@ -1593,7 +1606,9 @@ def test_put_classification_fallback_local_flag_only_edit_leaves_last_verified_a assert body["classification_fallback_local"] is True assert body["last_verified_at"] is not None assert row.last_verified_at == verified_at - assert row.revision == 4 + # A flag-only edit doesn't touch the credential's contents, so revision + # stays put too -- see test_put_flag_only_edit_leaves_revision_alone. + assert row.revision == 3 # --------------------------------------------------------------------------- @@ -1715,7 +1730,9 @@ def test_put_flag_only_edit_leaves_last_verified_at_untouched(user): body = resp.json() assert body["last_verified_at"] is not None assert row.last_verified_at == verified_at - assert row.revision == 4 + # A flag-only edit doesn't touch the credential's contents, so revision + # stays put too -- see test_put_flag_only_edit_leaves_revision_alone. + assert row.revision == 3 assert row.classification_byok is True @@ -1766,7 +1783,11 @@ def test_put_model_change_clears_last_verified_at_even_when_api_key_is_absent(us assert row.api_key == "sk-stored-key-1234" -def test_put_revision_increments_on_a_flag_only_edit_same_as_any_other_write(user): +def test_put_flag_only_edit_leaves_revision_alone(user): + # Revision tracks the credential's CONTENTS (key/provider/base_url/ + # model), not "was this row written to" -- a flag-only PUT racing an + # in-flight credential-class failure must not fence that failure's D2 + # identity check (credential_id + revision) into a bogus mismatch. row = _make_row(user_id=user.id, provider="openai", model="gpt-4o-mini", revision=5) db = _CredentialDB({user.id.hex: row}) _override(user, db) @@ -1779,7 +1800,7 @@ def test_put_revision_increments_on_a_flag_only_edit_same_as_any_other_write(use }, ) assert resp.status_code == 200 - assert row.revision == 6 + assert row.revision == 5 # --------------------------------------------------------------------------- @@ -2723,4 +2744,144 @@ def test_classifier_mix_two_user_isolation_each_gets_their_own_scoped_query(): assert resp_b.json() == {"classifier_mix": [{"kind": "user_key", "count": 7}]} compiled_b = _compiled(db_b.execute.call_args_list[0].args[0]) assert f"mail_thread.user_id = '{user_b.id.hex}'" in compiled_b - assert user_a.id.hex not in compiled_b + + +# --------------------------------------------------------------------------- +# D2 recovery: credential mutations reset a caller's terminally-capped +# `failed` action_item rows under the guard lock (plan: +# 2026-08-19-extraction-cost-hardening D2). The matrix the plan names: +# first-create, inactive-create, material PUT, flag-only PUT, activate, kill +# switch. `db.reset_calls` (via `_CredentialDB`'s `action_item` UPDATE +# interception) is `[user_id.hex]` once per reset, `[]` when a mutation +# deliberately doesn't reset. +# --------------------------------------------------------------------------- + + +def test_put_material_change_resets_extraction_attempts(user): + row = _make_row(user_id=user.id, provider="openai", model="gpt-4o-mini", revision=1) + db = _CredentialDB({user.id.hex: row}) + _override(user, db) + resp = TestClient(app).put( + "/api/v1/settings/llm", + json={"provider": "openai", "model": "a-different-model"}, + ) + assert resp.status_code == 200 + assert db.reset_calls == [user.id.hex] + + +def test_put_flag_only_edit_does_not_reset_extraction_attempts(user): + row = _make_row(user_id=user.id, provider="openai", model="gpt-4o-mini", revision=1) + db = _CredentialDB({user.id.hex: row}) + _override(user, db) + resp = TestClient(app).put( + "/api/v1/settings/llm", + json={"provider": "openai", "model": "gpt-4o-mini", "classification_byok": True}, + ) + assert resp.status_code == 200 + assert db.reset_calls == [] + + +def test_put_resend_of_the_identical_stored_key_does_not_reset_extraction_attempts(user): + # A caller who re-sends the SAME key their client already had cached + # (e.g. a form re-submit) hasn't changed the credential at all -- + # `api_key_provided` alone used to be enough to call this "material," + # which wrongly reopened rows still capped under this very credential. + verified_at = datetime.now(timezone.utc) + row = _make_row( + user_id=user.id, + provider="openai", + model="gpt-4o-mini", + api_key="sk-stored-key-1234", + revision=5, + last_verified_at=verified_at, + ) + db = _CredentialDB({user.id.hex: row}) + _override(user, db) + resp = TestClient(app).put( + "/api/v1/settings/llm", + json={ + "provider": "openai", + "model": "gpt-4o-mini", + "api_key": "sk-stored-key-1234", + }, + ) + assert resp.status_code == 200 + assert row.revision == 5 + assert row.last_verified_at == verified_at + assert db.reset_calls == [] + + +def test_put_first_create_resets_extraction_attempts(user): + # A brand-new user's first PUT always creates an active row -- same + # "the active credential just changed" event as a material rewrite. + db = _CredentialDB() + _override(user, db) + resp = TestClient(app).put( + "/api/v1/settings/llm", + json={"provider": "openai", "api_key": "sk-abcdefgh", "model": "gpt-4o-mini"}, + ) + assert resp.status_code == 200 + assert db.reset_calls == [user.id.hex] + + +def test_create_credential_first_one_resets_extraction_attempts(user): + db = _CredentialDB() + _override(user, db) + resp = TestClient(app).post( + "/api/v1/settings/llm/credentials", + json={"name": "work", "provider": "openai", "api_key": "sk-abcdefgh", "model": "m"}, + ) + assert resp.status_code == 201 + assert db.reset_calls == [user.id.hex] + + +def test_create_credential_inactive_spare_does_not_reset_extraction_attempts(user): + existing = _make_row(user_id=user.id, name="default", is_active=True) + db = _CredentialDB({user.id.hex: existing}) + _override(user, db) + resp = TestClient(app).post( + "/api/v1/settings/llm/credentials", + json={"name": "backup", "provider": "gemini", "api_key": "sk-abcdefgh", "model": "m"}, + ) + assert resp.status_code == 201 + assert db.reset_calls == [] + + +def test_activate_credential_switch_resets_extraction_attempts(user): + old_active = _make_row(user_id=user.id, name="default", is_active=True) + target = _make_row(user_id=user.id, name="work", is_active=False) + db = _CredentialDB({user.id.hex: [old_active, target]}) + _override(user, db) + resp = TestClient(app).post(f"/api/v1/settings/llm/credentials/{target.id}/activate") + assert resp.status_code == 200 + assert db.reset_calls == [user.id.hex] + + +def test_activate_credential_already_active_does_not_reset_extraction_attempts(user): + target = _make_row(user_id=user.id, name="default", is_active=True) + db = _CredentialDB({user.id.hex: target}) + _override(user, db) + resp = TestClient(app).post(f"/api/v1/settings/llm/credentials/{target.id}/activate") + assert resp.status_code == 200 + assert db.reset_calls == [] + + +def test_delete_kill_switch_does_not_reset_extraction_attempts(user): + # No credential is left to retry under -- resetting here would just + # strand the rows capped-at-zero until a NEW credential shows up. + row = _make_row(user_id=user.id) + db = _CredentialDB({user.id.hex: row}) + _override(user, db) + resp = TestClient(app).delete("/api/v1/settings/llm") + assert resp.status_code == 204 + assert db.reset_calls == [] + + +def test_delete_credential_by_id_does_not_reset_extraction_attempts(user): + active = _make_row(user_id=user.id, name="default", is_active=True) + spare = _make_row(user_id=user.id, name="spare", is_active=False) + db = _CredentialDB({user.id.hex: [active, spare]}) + _override(user, db) + resp = TestClient(app).delete(f"/api/v1/settings/llm/credentials/{spare.id}") + assert resp.status_code == 200 + assert db.reset_calls == []