diff --git a/.kiro/specs/managed-kb-migration/tasks.md b/.kiro/specs/managed-kb-migration/tasks.md index 1ac5419b..e9365560 100644 --- a/.kiro/specs/managed-kb-migration/tasks.md +++ b/.kiro/specs/managed-kb-migration/tasks.md @@ -836,4 +836,27 @@ All three flags — managed-default, migration, and reconciler arming — ship * both needed manual repair. - Overlaps task 14.4 (one-click retry) and the report-only reconciler, which already knows how to join Bedrock's view against ours. + - **Backend built (report-only), branch `feat/kb-deadletter-reconcile`.** + `kb_migration/document_reconciler.py` is the missing second writer of `DOC#` + status: once a day it finds rows stuck non-terminal (`uploading`/`chunking`/ + `embedding`) past a 60-minute grace gate, asks Bedrock the ground truth per + document, and — the §5.37 case — drives a stranded-but-**retrievable** document + to `complete`. It reuses the consumer's own probes (`document_status`, the + `equals`-on-`document_id` retrievability search, its status-set constants, and + `set_document_terminal`) so §5.37/§5.38/§5.39 live in one place, not four. A + `FAILED` document is driven to `failed`; a `NOT_FOUND` one (dead-lettered before + ingest) is **re-ingested** from the S3 bytes — the scheduled form of 14.4's + one-click retry. Modelled on `reconciler.py`: **ships disarmed** + (`MANAGED_KB_DOC_RECONCILER_ARMED`, empty ⇒ off), per-run action limit applies + in both modes so the report is trustworthy, grace gate is a pure function of the + row's own `updatedAt` and fails closed. `terminal`/`deleting` rows are never + candidates (a soft-deleted doc must not be resurrected). Guards in + `tests/lambdas/test_kb_document_reconciler.py`, mutation-verified (neutering the + retrievability gate fails `test_indexed_but_not_retrievable_is_left_short_of_complete`; + widening `NON_TERMINAL_STATUSES` fails `test_terminal_and_deleting_rows_are_never_candidates`). + - **Remaining (deploy-gated follow-up, not in this PR):** wire the reconciler's + own Lambda + EventBridge schedule + IAM in `kb-migration-construct.ts`, then set + and eventually flip `MANAGED_KB_DOC_RECONCILER_ARMED`. The flag is exempted in + `test_kb_migration_env_contract.py`'s `OPTIONAL_OVERRIDES` until that wiring + lands. - _HANDOFF §5.37 · Requirements: 21.2_ diff --git a/backend/src/apis/app_api/kb_migration/document_reconciler.py b/backend/src/apis/app_api/kb_migration/document_reconciler.py new file mode 100644 index 00000000..ec7f5c74 --- /dev/null +++ b/backend/src/apis/app_api/kb_migration/document_reconciler.py @@ -0,0 +1,755 @@ +"""Dead-letter reconciler for stranded managed-KB documents. + +Task 16.5 (HANDOFF §5.37). The companion to :mod:`reconciler`, which reconciles +whole knowledge bases against ``ListKnowledgeBases``; this one reconciles +individual ``DOC#`` rows against Bedrock's *document* view. + +The gap it closes +----------------- +On the managed path a document's ``DOC#`` status is written by exactly one writer: +the ingestion consumer (``kb_migration/ingestion_consumer.py``). That consumer +polls Bedrock until the document is genuinely retrievable and only then writes +``complete``. It is the right design — but it is the *only* writer, and it runs +inside a Lambda whose asynchronous retry is capped at **2** attempts (a hard +service limit). When an event exhausts those retries and dead-letters, the row is +left in a non-terminal state (``uploading`` / ``chunking`` / ``embedding``) with +**nothing left to revisit it** — even though Bedrock frequently finished indexing +the document seconds after the final attempt was dead-lettered, so the content is +sitting in the knowledge base fully retrievable. + +That combination is invisible and permanent, because the retrieval status filter +(``rag_service._filter_vectors_by_document_status``) serves **only** ``complete`` +documents. A stranded row's chunks are dropped from every query: the user was told +their upload worked, the content really is in the knowledge base, and the +assistant will never cite it. Two such documents occurred in dev and both needed a +manual DynamoDB edit. + +This module is the missing second writer. Once a day it finds ``DOC#`` rows stuck +non-terminal past a grace period, asks Bedrock the ground truth for each, and — the +§5.37 case — drives a stranded-but-retrievable document to ``complete``. It also +handles the two neighbouring outcomes the same probe reveals: a document Bedrock +reports ``FAILED`` is driven to ``failed`` (it was never going to recover), and a +document Bedrock has never heard of (``NOT_FOUND`` — dead-lettered *before* the +ingest was ever accepted) is **re-ingested** from the bytes still in S3. That +re-ingest is the one-click retry of task 14.4 arriving on a schedule instead of a +button. + +Ground truth comes from the consumer, not a second copy +------------------------------------------------------- +Every decision here reuses the consumer's own probes — :func:`document_status` +(``GetKnowledgeBaseDocuments``), the ``equals``-on-``document_id`` retrievability +search, its status-set constants, and its terminal-write function. That reuse is +deliberate: those functions carry three hard-won lessons (§5.37 the poll budget, +§5.38 that an unfiltered retrievability search finds the wrong document, §5.39 +that the live service returns statuses the SDK enum omits). A reconciler that +re-derived any of them would be a fourth place for the same bug to live. The one +thing this module does *not* reuse is the consumer's *polling* — it takes a single +retrievability reading per document rather than waiting, because it is sweeping a +fleet, not shepherding one upload. + +Report-only, and armed separately +---------------------------------- +Modelled on :mod:`reconciler`. It ships **disarmed**: it logs exactly what it +would have done and writes nothing, so its judgement can be checked against real +data before it is trusted to correct records. Arming is one flag, +:data:`FLAG_DOC_RECONCILER_ARMED`, and an **empty string reads as off** — an unset +GitHub Actions variable expands to ``""``. The per-run action limit +(:func:`max_actions_per_run`) applies in **both** modes, so a report never claims +more corrections than an armed run would actually make. + +The grace gate is a pure function of the row's own ``updatedAt``, never of +discovery time — the same discipline as the KB reconciler's ``createdAt`` gate. A +row whose ``updatedAt`` cannot be read is left alone: without proof that a document +has been stuck *longer than a legitimate in-flight ingestion could take*, touching +it risks racing an upload that is still, correctly, being worked on. + +Import boundary +--------------- +Module-level imports are stdlib plus the stdlib-only ``ingestion_consumer`` / +``reconciler`` / ``kb_backend.records`` siblings; ``boto3`` and the heavy +``ManagedKbBackend`` are function-local. DynamoDB is reached through the raw table +resource, matching every other module in this package. +""" + +from __future__ import annotations + +import logging +import os +from dataclasses import dataclass, field +from datetime import datetime, timedelta, timezone +from typing import Any, Callable, Dict, Iterator, List, Optional + +from apis.app_api.kb_migration import ingestion_consumer as ic +from apis.shared.kb_backend.metrics import emit_count + +logger = logging.getLogger() +logger.setLevel(logging.INFO) + +# ── Flags ──────────────────────────────────────────────────────────────────── +# +# The arming flag. Absent, empty, or anything not in the truthy set means the +# reconciler reports and corrects nothing. Same allow-list as the KB reconciler +# and the dispatcher: the failure being designed around is a value that is present +# but empty (``bool("")`` is off by luck, ``bool("false")`` is not). +FLAG_DOC_RECONCILER_ARMED = "MANAGED_KB_DOC_RECONCILER_ARMED" + +_TRUTHY = frozenset({"1", "true", "yes", "on", "enabled"}) + +# ── Document statuses ───────────────────────────────────────────────────────── +# +# The non-terminal ``DOC#`` states a stranded document can be parked in. Taken to +# match ``apis/app_api/documents/ingestion/status.py``'s ``DocumentStatus`` literal +# minus its terminal members. ``deleting`` is deliberately NOT here: a +# soft-deleted document is being removed on purpose and must never be resurrected +# to ``complete``. +NON_TERMINAL_STATUSES = frozenset({"uploading", "chunking", "embedding"}) + +# ── Tunables, resolved at call time ────────────────────────────────────────── +# +# Read inside the functions that use them rather than bound as default arguments: +# a default argument is evaluated once at import, so a test overriding it silently +# gets the production value instead. Same reason the KB reconciler does this. + +#: A document younger than this is not yet evidence of a dead-letter — it may still +#: be legitimately in flight. The ingestion consumer waits up to +#: ``INDEXED_POLL_TIMEOUT_SECONDS`` (600 s) inside one invocation, and an event +#: gets 1 + 2 deliveries spread over a few minutes before it dead-letters, so a +#: genuinely-working document can still be non-terminal for ~15 minutes. 60 minutes +#: is comfortably past that, which is the correct direction: the cost of waiting one +#: more daily pass is nil, and the cost of racing an in-flight upload is marking it +#: from under the consumer. +STUCK_MIN_AGE_MINUTES = 60.0 + +#: Bounds the corrective work of a single run — mark-completes, re-ingests and +#: fails together — so a bug in the join, or a sudden flood of stranded rows, +#: costs at most this many actions before someone reads the report. Applied in +#: report-only mode too, so the report is trustworthy. +MAX_ACTIONS_PER_RUN = 25 + +#: Hard ceiling above which the env override is ignored. A larger sweep should +#: require repeated observed runs, not a variable edit. +MAX_ACTIONS_CEILING = 100 + +#: Bounds the join itself. A reconciler that walked an unbounded table would time +#: out mid-pass and produce a partial report indistinguishable from a complete one. +MAX_RECORDS_PER_RUN = 5000 + +# ── Action kinds ────────────────────────────────────────────────────────────── +ACTION_MARK_COMPLETE = "mark_complete" +ACTION_RE_INGEST = "re_ingest" +ACTION_MARK_FAILED = "mark_failed" + +# ── Metrics ────────────────────────────────────────────────────────────────── +METRIC_STRANDED_FOUND = "KbStrandedDocumentsFound" +METRIC_MARKED_COMPLETE = "KbStrandedDocumentsCompleted" +METRIC_RE_INGESTED = "KbStrandedDocumentsReingested" +METRIC_MARKED_FAILED = "KbStrandedDocumentsFailed" +METRIC_LIMIT_REACHED = "KbDocumentReconcilerLimitReached" + + +@dataclass +class PlannedAction: + """One correction the reconciler intends, and the evidence for it.""" + + assistant_id: str + document_id: str + kind: str + current_status: str + bedrock_status: str + performed: bool = False + error: Optional[str] = None + + def to_dict(self) -> Dict[str, Any]: + return { + "assistantId": self.assistant_id, + "documentId": self.document_id, + "kind": self.kind, + "currentStatus": self.current_status, + "bedrockStatus": self.bedrock_status, + "performed": self.performed, + "error": self.error, + } + + +@dataclass +class DocumentReconcileReport: + """What one run found and what it did (or, disarmed, would have done). + + ``armed`` lives on the report, not only in the logs, so a stored artifact is + self-describing: an operator reading last night's output should not have to go + and check what the flag was set to at the time. + """ + + armed: bool = False + managed_records: int = 0 + documents_scanned: int = 0 + stranded: int = 0 + planned_actions: List[PlannedAction] = field(default_factory=list) + skipped_too_young: List[str] = field(default_factory=list) + skipped_in_flight: List[str] = field(default_factory=list) + skipped_not_retrievable: List[str] = field(default_factory=list) + limit_reached: bool = False + + @property + def actions_performed(self) -> int: + return sum(1 for action in self.planned_actions if action.performed) + + def _count(self, kind: str) -> int: + return sum(1 for action in self.planned_actions if action.kind == kind) + + def to_dict(self) -> Dict[str, Any]: + return { + "armed": self.armed, + "mode": "armed" if self.armed else "report-only", + "managedRecords": self.managed_records, + "documentsScanned": self.documents_scanned, + "stranded": self.stranded, + "plannedActions": [action.to_dict() for action in self.planned_actions], + "plannedByKind": { + ACTION_MARK_COMPLETE: self._count(ACTION_MARK_COMPLETE), + ACTION_RE_INGEST: self._count(ACTION_RE_INGEST), + ACTION_MARK_FAILED: self._count(ACTION_MARK_FAILED), + }, + "actionsPerformed": self.actions_performed, + "skippedTooYoung": self.skipped_too_young, + "skippedInFlight": self.skipped_in_flight, + "skippedNotRetrievable": self.skipped_not_retrievable, + "limitReached": self.limit_reached, + } + + +# ── Flag and tunable readers ───────────────────────────────────────────────── +def doc_reconciler_armed() -> bool: + """Whether the reconciler may write. Defaults to **off**. + + An empty string is off: an unset repository or environment variable expands to + ``""`` in GitHub Actions, and stamping a truthiness test on the raw value is + the exact bug the KB reconciler was bitten by. + """ + raw = os.environ.get(FLAG_DOC_RECONCILER_ARMED) + if not raw: + return False + return raw.strip().lower() in _TRUTHY + + +def _env_float(name: str, default: float) -> float: + raw = os.environ.get(name) + if not raw: + return default + try: + return float(raw) + except ValueError: + logger.warning(f"{name}={raw!r} is not a number; falling back to {default}") + return default + + +def _env_int(name: str, default: int) -> int: + raw = os.environ.get(name) + if not raw: + return default + try: + return int(raw) + except ValueError: + logger.warning(f"{name}={raw!r} is not an integer; falling back to {default}") + return default + + +def stuck_min_age_minutes() -> float: + return _env_float("MANAGED_KB_DOC_STUCK_MIN_AGE_MINUTES", STUCK_MIN_AGE_MINUTES) + + +def max_actions_per_run() -> int: + """The per-run action bound, clamped so the environment cannot lift it. + + The env var may lower the limit but not raise it past + :data:`MAX_ACTIONS_CEILING`. A bound any variable can set to a million is not a + bound; this one caps how much a single bad run can churn before its report is + read. + """ + requested = _env_int("MANAGED_KB_DOC_RECONCILER_MAX_ACTIONS", MAX_ACTIONS_PER_RUN) + if requested > MAX_ACTIONS_CEILING: + logger.warning( + f"MANAGED_KB_DOC_RECONCILER_MAX_ACTIONS={requested} exceeds the ceiling " + f"of {MAX_ACTIONS_CEILING}; clamping. Run the reconciler repeatedly " + f"rather than raising this." + ) + return MAX_ACTIONS_CEILING + return max(requested, 0) + + +def max_records_per_run() -> int: + return _env_int("MANAGED_KB_DOC_RECONCILER_MAX_RECORDS", MAX_RECORDS_PER_RUN) + + +# ── DynamoDB plumbing ──────────────────────────────────────────────────────── +def _table(): + import boto3 + + return boto3.resource("dynamodb").Table(os.environ["DYNAMODB_ASSISTANTS_TABLE_NAME"]) + + +def _now() -> datetime: + return datetime.now(timezone.utc) + + +def iter_document_records(assistant_id: str) -> Iterator[Dict[str, Any]]: + """Every ``DOC#`` row for one assistant, paging the query to exhaustion. + + Paged for the same reason the reconciler pages its scans: a truncated read + would make a stranded document on a later page look absent, and the run would + silently skip it. + """ + from boto3.dynamodb.conditions import Key + + table = _table() + kwargs: Dict[str, Any] = { + "KeyConditionExpression": Key("PK").eq(f"AST#{assistant_id}") + & Key("SK").begins_with("DOC#"), + } + while True: + response = table.query(**kwargs) + for item in response.get("Items") or []: + yield item + start = response.get("LastEvaluatedKey") + if not start: + return + kwargs["ExclusiveStartKey"] = start + + +# ── Age gate (pure function of the row's own updatedAt) ────────────────────── +def document_age_minutes(updated_at: Any, now: Optional[datetime] = None) -> Optional[float]: + """Minutes since the row's ``updatedAt``, or ``None`` if it cannot be read. + + Reuses the KB reconciler's timestamp parser so the accepted shapes (aware + datetime, ISO string, epoch number) are identical across the feature. + """ + from apis.app_api.kb_migration.reconciler import parse_aws_timestamp + + stamped = parse_aws_timestamp(updated_at) + if stamped is None: + return None + return ((now or _now()) - stamped).total_seconds() / 60.0 + + +def document_is_stuck_long_enough( + updated_at: Any, + now: Optional[datetime] = None, + min_age_minutes: Optional[float] = None, +) -> bool: + """Whether a non-terminal row has been stuck long enough to reconcile. + + A missing or unparseable ``updatedAt`` returns ``False`` — fail-safe. Without + proof the row has been stuck longer than a legitimate ingestion can take, + correcting it risks racing the consumer that is still, correctly, working it. + The input is the row's own ``updatedAt``, never discovery time: the answer must + not depend on when this process happened to look. + """ + if min_age_minutes is None: + min_age_minutes = stuck_min_age_minutes() + from apis.app_api.kb_migration.reconciler import parse_aws_timestamp + + stamped = parse_aws_timestamp(updated_at) + if stamped is None: + return False + return (now or _now()) - stamped > timedelta(minutes=min_age_minutes) + + +# ── Retrievability: a single reading, not a poll ───────────────────────────── +def is_retrievable(backend: Any, kb_ref: str, document_id: str) -> bool: + """Whether a retrieval really returns ``document_id`` right now. + + ONE reading, not the consumer's poll: this sweeps a fleet, so waiting per + document would blow the run's time budget. The ``equals`` filter on + ``document_id`` is the load-bearing part and is the whole reason §5.38 exists + — an *unfiltered* search for a document id returns whatever the reranker + prefers (a document id is meaningless to an embedding model), so it confirms + the wrong document as retrievable and gets worse as a knowledge base grows. + Filtered, a non-empty result *is* proof and an empty one is a true negative. + + A probe that itself errors is treated as "not retrievable" for this pass — the + reconciler simply leaves the row for the next run rather than acting on a + failed reading. + """ + import asyncio + + document_filter = {"equals": {"key": "document_id", "value": document_id}} + try: + chunks = asyncio.run( + backend.search(kb_ref, document_id, 5, retrieval_filter=document_filter) + ) + except Exception as exc: # noqa: BLE001 - a probe failure is not a verdict + logger.warning(f"retrievability probe for {document_id} failed: {exc}") + return False + + # The filter already restricts the result set to this document; the per-chunk + # check is a belt-and-braces guard against a filter a future API change ignores. + for chunk in chunks or []: + metadata = getattr(chunk, "metadata", None) or {} + if metadata.get("document_id") == document_id: + return True + return False + + +# ── The run ────────────────────────────────────────────────────────────────── +def reconcile_documents( + client=None, + backend_factory: Optional[Callable[[str], Any]] = None, + armed: Optional[bool] = None, + now: Optional[datetime] = None, +) -> DocumentReconcileReport: + """One document-reconciliation pass. + + ``armed`` defaults to :func:`doc_reconciler_armed`, i.e. to the flag, i.e. to + off. It is an argument only so a test can exercise the armed path without + mutating process environment — never so a caller can conveniently turn writing + on. + + ``backend_factory`` builds a :class:`ManagedKbBackend` for an assistant; it is + injectable so tests can supply a stub that models Bedrock's document view. The + default constructs a real backend keyed on the assistant id (``App_KB_Id`` == + ``assistant_id`` in this phase) with the injected control-plane ``client``. + """ + from apis.shared.kb_backend.records import ENGINE_MANAGED, resolve_engine + + if armed is None: + armed = doc_reconciler_armed() + if now is None: + now = _now() + if backend_factory is None: + backend_factory = _default_backend_factory(client) + + report = DocumentReconcileReport(armed=armed) + record_limit = max_records_per_run() + action_limit = max_actions_per_run() + min_age = stuck_min_age_minutes() + + for record in _iter_managed_records(record_limit, report): + if resolve_engine(record) != ENGINE_MANAGED: + continue + if not record.get("awsKbId"): + # Managed but not provisioned yet: there is no knowledge base to probe, + # so a non-terminal document here is waiting on provisioning, not + # dead-lettered. + continue + + assistant_id = _assistant_id_of(record) + if not assistant_id: + continue + report.managed_records += 1 + backend = None # built lazily, only if this assistant has a stranded row + + for document in iter_document_records(assistant_id): + report.documents_scanned += 1 + status = str(document.get("status") or "") + if status not in NON_TERMINAL_STATUSES: + continue + + document_id = str(document.get("documentId") or "") + if not document_id: + logger.warning( + f"skipping malformed DOC# row {document.get('PK')}/" + f"{document.get('SK')}: no documentId" + ) + continue + + if not document_is_stuck_long_enough( + document.get("updatedAt"), now=now, min_age_minutes=min_age + ): + report.skipped_too_young.append(document_id) + continue + + report.stranded += 1 + + if backend is None: + backend = backend_factory(assistant_id) + + if report.limit_reached or len(report.planned_actions) >= action_limit: + report.limit_reached = True + emit_count(METRIC_LIMIT_REACHED) + logger.warning( + f"per-run action limit of {action_limit} reached; {document_id} " + f"and any further stranded documents are left for the next run" + ) + break + + _reconcile_one( + assistant_id, document, document_id, status, backend, armed, report + ) + + if report.limit_reached: + break + + if report.stranded: + emit_count(METRIC_STRANDED_FOUND, value=report.stranded) + + logger.info( + f"document reconcile complete: mode={'armed' if armed else 'report-only'} " + f"managedRecords={report.managed_records} " + f"documentsScanned={report.documents_scanned} stranded={report.stranded} " + f"planned={len(report.planned_actions)} performed={report.actions_performed} " + f"tooYoung={len(report.skipped_too_young)} " + f"inFlight={len(report.skipped_in_flight)} " + f"notRetrievable={len(report.skipped_not_retrievable)} " + f"limitReached={report.limit_reached}" + ) + return report + + +def _iter_managed_records( + record_limit: int, report: DocumentReconcileReport +) -> Iterator[Dict[str, Any]]: + """KB_Records, bounded, so a huge table cannot make the pass run forever. + + Reuses the KB reconciler's ``iter_kb_records`` (the ``SK begins_with KB#`` + scan) rather than a second implementation, so tombstones and key prefixes are + handled identically. + """ + from apis.app_api.kb_migration.reconciler import iter_kb_records + + seen = 0 + for record in iter_kb_records(): + if seen >= record_limit: + report.limit_reached = True + logger.warning( + f"stopping the KB_Record walk at {record_limit}; this run is partial" + ) + return + seen += 1 + yield record + + +def _reconcile_one( + assistant_id: str, + document: Dict[str, Any], + document_id: str, + status: str, + backend: Any, + armed: bool, + report: DocumentReconcileReport, +) -> None: + """Classify one stranded document against Bedrock and act (or report). + + The classification mirrors the ingestion consumer's own handling exactly, + because it uses the consumer's status sets — the point of a reconciler is to + reach the terminal state the dead-lettered invocation would have reached, not + to invent a new policy. + """ + bedrock_status, bedrock_updated_at = ic.document_status(backend, assistant_id, document_id) + + # Still indexing: not stranded, just slow. Leave it — a later pass (or a + # redelivery that beat the DLQ) will finish it. The grace gate makes this rare. + if bedrock_status in ic.DOC_STATUSES_IN_FLIGHT: + report.skipped_in_flight.append(document_id) + logger.info( + f"document {document_id} is {bedrock_status} in the knowledge base; " + f"still indexing, leaving it" + ) + return + + # Terminal failure on Bedrock's side: retrying cannot help, so drive the row to + # the terminal state the consumer would have written. + if bedrock_status in ic.DOC_STATUSES_FAILED: + _plan( + report, assistant_id, document_id, ACTION_MARK_FAILED, status, bedrock_status, + armed, + lambda: ic.set_document_terminal( + assistant_id, document_id, ic.STATUS_FAILED, + error=f"the knowledge base reports this document as {bedrock_status}", + ), + METRIC_MARKED_FAILED, + f"[report-only] WOULD mark {document_id} failed ({bedrock_status})", + ) + return + + # Indexed (fully or partially): if it is genuinely retrievable, this is the + # §5.37 case — a stranded-but-servable document — and it is driven to complete. + if bedrock_status in (ic.DOC_STATUS_INDEXED, *ic.DOC_STATUSES_PARTIAL): + if not is_retrievable(backend, assistant_id, document_id): + # Indexed but not yet queryable, or the probe failed. Do NOT claim + # complete: that is exactly the "upload worked but the assistant cannot + # see it" report the consumer exists to prevent. Leave it for next run. + report.skipped_not_retrievable.append(document_id) + logger.info( + f"document {document_id} is {bedrock_status} but not retrievable " + f"this pass; leaving it short of complete" + ) + return + indexed_at = bedrock_updated_at or ic._now_iso() + _plan( + report, assistant_id, document_id, ACTION_MARK_COMPLETE, status, bedrock_status, + armed, + lambda: ic.set_document_terminal( + assistant_id, document_id, ic.STATUS_COMPLETE, + indexed_at=indexed_at, retrievable_at=ic._now_iso(), + ), + METRIC_MARKED_COMPLETE, + f"[report-only] WOULD mark {document_id} complete ({bedrock_status}, " + f"confirmed retrievable)", + ) + return + + # NOT_FOUND: Bedrock never received this document. The event dead-lettered + # before the ingest was accepted, so the bytes in S3 were never submitted. + # Re-ingest them — the scheduled form of task 14.4's one-click retry. + if bedrock_status == ic.DOC_STATUS_NOT_FOUND: + source = _document_source(document, document_id) + if source is None: + logger.warning( + f"document {document_id} is NOT_FOUND but its row lacks an s3Key; " + f"cannot re-ingest, leaving it" + ) + report.skipped_not_retrievable.append(document_id) + return + _plan( + report, assistant_id, document_id, ACTION_RE_INGEST, status, bedrock_status, + armed, + lambda: _reingest(backend, assistant_id, source), + METRIC_RE_INGESTED, + f"[report-only] WOULD re-ingest {document_id} (NOT_FOUND in the " + f"knowledge base; bytes still in S3)", + ) + return + + # Any other value: the live service returns statuses the SDK enum omits + # (§5.39). Treat unknown as "still working" and leave it, logging the value so + # it can be classified rather than silently mishandled. + report.skipped_in_flight.append(document_id) + logger.warning( + f"document {document_id} reported unrecognised Bedrock status " + f"{bedrock_status!r}; treating it as in-flight and leaving it" + ) + + +def _plan( + report: DocumentReconcileReport, + assistant_id: str, + document_id: str, + kind: str, + current_status: str, + bedrock_status: str, + armed: bool, + perform: Callable[[], None], + metric: str, + report_only_message: str, +) -> None: + """Record a planned action and, if armed, perform it. + + The action is appended to the report whether or not it runs, so the report-only + artifact describes exactly what an armed run would do. When armed, a failure is + captured on the action rather than raised — one bad document must not end the + sweep — matching the KB reconciler's per-orphan error handling. + """ + action = PlannedAction( + assistant_id=assistant_id, + document_id=document_id, + kind=kind, + current_status=current_status, + bedrock_status=bedrock_status, + ) + report.planned_actions.append(action) + + if not armed: + logger.warning(f"{report_only_message}. Set {FLAG_DOC_RECONCILER_ARMED} to arm.") + return + + try: + perform() + action.performed = True + emit_count(metric) + logger.info(f"{kind} performed for document {document_id}") + except Exception as exc: # noqa: BLE001 - one bad document must not end the run + action.error = str(exc) + logger.error(f"{kind} failed for document {document_id}: {exc}", exc_info=True) + + +def _reingest(backend: Any, assistant_id: str, source: Any) -> None: + import asyncio + + asyncio.run(backend.ingest(assistant_id, source)) + + +def _document_source(document: Dict[str, Any], document_id: str) -> Optional[Any]: + """Reconstruct a ``DocumentSource`` for re-ingest from the ``DOC#`` row. + + The row carries ``s3Key`` and ``filename`` (from the ``Document`` model's + aliases), which is everything the managed backend needs to re-submit bytes + already in S3. Returns ``None`` when ``s3Key`` is absent — an old row that + predates the attribute cannot be re-ingested from here and is left for a + re-upload. + """ + from apis.shared.kb_backend.protocol import DocumentSource + + s3_key = document.get("s3Key") + if not s3_key: + return None + filename = document.get("filename") or "" + return DocumentSource(document_id=document_id, filename=filename, s3_key=str(s3_key)) + + +def _default_backend_factory(client) -> Callable[[str], Any]: + """Build a real :class:`ManagedKbBackend` per assistant, sharing one client. + + Function-local import: ``ManagedKbBackend`` pulls in the retrieval stack and + must not be a module-level import in this size-constrained Lambda. + """ + def factory(_assistant_id: str) -> Any: + from apis.shared.kb_backend.managed_backend import ManagedKbBackend + + return ManagedKbBackend(agent_client=client) + + return factory + + +def _assistant_id_of(record: Dict[str, Any]) -> str: + pk = str(record.get("PK") or "") + return pk[len("AST#") :] if pk.startswith("AST#") else "" + + +def lambda_handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]: + """Scheduled entry point. Returns the report so it lands in the invocation log. + + The invocation event is deliberately **not** consulted for arming, for the same + reason as the KB reconciler: an event payload is the one input an operator does + not review, so honouring an ``armed`` field in it would let any principal + holding ``lambda:InvokeFunction`` correct records while every reviewable setting + still said report-only. If the event disagrees with the flag, the flag wins and + the disagreement is logged. + """ + requested = (event or {}).get("armed") + if requested is not None: + logger.warning( + f"ignoring armed={requested!r} from the invocation event: arming is " + f"controlled only by {FLAG_DOC_RECONCILER_ARMED}" + ) + report = reconcile_documents() + return {"statusCode": 200, "report": report.to_dict()} + + +__all__ = [ + "ACTION_MARK_COMPLETE", + "ACTION_MARK_FAILED", + "ACTION_RE_INGEST", + "FLAG_DOC_RECONCILER_ARMED", + "MAX_ACTIONS_CEILING", + "MAX_ACTIONS_PER_RUN", + "MAX_RECORDS_PER_RUN", + "METRIC_LIMIT_REACHED", + "METRIC_MARKED_COMPLETE", + "METRIC_MARKED_FAILED", + "METRIC_RE_INGESTED", + "METRIC_STRANDED_FOUND", + "NON_TERMINAL_STATUSES", + "STUCK_MIN_AGE_MINUTES", + "DocumentReconcileReport", + "PlannedAction", + "doc_reconciler_armed", + "document_age_minutes", + "document_is_stuck_long_enough", + "is_retrievable", + "iter_document_records", + "lambda_handler", + "max_actions_per_run", + "max_records_per_run", + "reconcile_documents", + "stuck_min_age_minutes", +] diff --git a/backend/tests/lambdas/test_kb_document_reconciler.py b/backend/tests/lambdas/test_kb_document_reconciler.py new file mode 100644 index 00000000..b68ce02e --- /dev/null +++ b/backend/tests/lambdas/test_kb_document_reconciler.py @@ -0,0 +1,677 @@ +"""Dead-letter document reconciler — task 16.5, HANDOFF §5.37. + +The reconciler is the missing *second* writer of ``DOC#`` status. The ingestion +consumer is the only writer today, and when its event dead-letters (Lambda async +retry is capped at 2) a document Bedrock finished indexing is left parked +non-terminal forever — and the retrieval filter serves only ``complete``, so its +content is in the knowledge base and invisible to every query. + +Four assertions here are the reason the file exists, and each guards a mistake a +green suite would otherwise hide: + +**A stranded-but-retrievable document is driven to ``complete`` — but only when it +is genuinely retrievable.** The §5.37 fix. Marking it complete on ``INDEXED`` +alone, without confirming a filtered retrieval returns it, recreates the exact +"upload worked but the assistant cannot see it" report the consumer was built to +prevent. Both halves are asserted. + +**Report-only really is a no-op.** The shipped mode plans every action and writes +nothing; the arming flag treats an empty string as off. + +**The grace gate reads the row's own ``updatedAt``, never discovery time**, and +fails closed when it cannot be read — so an in-flight upload is never marked from +under the consumer. + +**Terminal and soft-deleted rows are untouchable.** ``complete``/``failed`` are +done; ``deleting`` is being removed on purpose and must never be resurrected. + +No test contacts AWS. DynamoDB is moto; the managed backend is a stub that models +Bedrock's document view, mirroring ``test_kb_ingestion_consumer``. +""" + +import types +from datetime import datetime, timedelta, timezone + +import boto3 +import pytest +from moto import mock_aws + +from apis.app_api.kb_migration import document_reconciler as dr + +REGION = "us-east-1" +TABLE = "test-doc-reconciler" +NOW = datetime(2026, 6, 1, 12, 0, 0, tzinfo=timezone.utc) + + +def _iso(moment): + return moment.strftime("%Y-%m-%dT%H:%M:%SZ") + + +# Ages relative to NOW. +OLD = _iso(NOW - timedelta(hours=2)) # comfortably past the 60-minute gate +YOUNG = _iso(NOW - timedelta(minutes=5)) # still plausibly in flight + + +@pytest.fixture() +def table(monkeypatch): + monkeypatch.setenv("AWS_DEFAULT_REGION", REGION) + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "testing") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "testing") + monkeypatch.setenv("AWS_SESSION_TOKEN", "testing") + monkeypatch.setenv("DYNAMODB_ASSISTANTS_TABLE_NAME", TABLE) + # Never inherited from the developer's shell: the reconciler is disarmed unless + # something says otherwise. + monkeypatch.delenv(dr.FLAG_DOC_RECONCILER_ARMED, raising=False) + + with mock_aws(): + boto3.client("dynamodb", region_name=REGION).create_table( + TableName=TABLE, + KeySchema=[ + {"AttributeName": "PK", "KeyType": "HASH"}, + {"AttributeName": "SK", "KeyType": "RANGE"}, + ], + AttributeDefinitions=[ + {"AttributeName": "PK", "AttributeType": "S"}, + {"AttributeName": "SK", "AttributeType": "S"}, + ], + BillingMode="PAY_PER_REQUEST", + ) + yield boto3.resource("dynamodb", region_name=REGION).Table(TABLE) + + +@pytest.fixture(autouse=True) +def no_metrics(monkeypatch): + monkeypatch.setattr(dr, "emit_count", lambda *a, **k: None) + + +# ── Fixtures for the table ──────────────────────────────────────────────────── +def _seed_kb(table, assistant_id, *, engine="managed", aws_kb_id="KB1"): + item = { + "PK": f"AST#{assistant_id}", + "SK": f"KB#{assistant_id}", + "appKbId": assistant_id, + } + if engine: + item["retrievalEngine"] = engine + if aws_kb_id: + item["awsKbId"] = aws_kb_id + item["awsDataSourceId"] = f"DS{aws_kb_id}" + table.put_item(Item=item) + + +def _seed_doc(table, assistant_id, document_id, *, status, updated_at=OLD, s3_key=None, filename=None): + item = { + "PK": f"AST#{assistant_id}", + "SK": f"DOC#{document_id}", + "documentId": document_id, + "status": status, + "updatedAt": updated_at, + } + if s3_key is not None: + item["s3Key"] = s3_key + if filename is not None: + item["filename"] = filename + table.put_item(Item=item) + + +def _doc(table, assistant_id, document_id): + return table.get_item( + Key={"PK": f"AST#{assistant_id}", "SK": f"DOC#{document_id}"} + ).get("Item") + + +# ── The stub backend: models Bedrock's document view ───────────────────────── +class _FakeBackend: + """Models the parts of ManagedKbBackend the reconciler leans on. + + ``statuses`` maps document_id -> the status Bedrock's ``GetKnowledgeBaseDocuments`` + reports (INDEXED, FAILED, IN_PROGRESS, TEXT_INDEXED, PARTIALLY_INDEXED, + NOT_FOUND, or anything unrecognised). ``retrievable`` is the set of document ids + a *filtered* retrieval returns — modelling that INDEXED does not imply + retrievable, and that the id-as-query search only works when filtered. + """ + + def __init__(self, statuses=None, retrievable=None, other_documents=("DOC-someone-else",)): + self._statuses = dict(statuses or {}) + self._retrievable = set(retrievable or []) + self._other_documents = list(other_documents) + self.search_filters = [] + self.ingested = [] + self._agent_client = _FakeAgent(self) + + # -- the private surface `ic.document_status` reuses ----------------------- + def _agent(self): + return self._agent_client + + def _locate(self, kb_ref): + return ("KB1", "DS1") + + # -- the protocol surface -------------------------------------------------- + async def ingest(self, kb_ref, source): + self.ingested.append(source.document_id) + + async def search(self, kb_ref, query, top_k=5, retrieval_filter=None): + """Honours an ``equals`` filter on ``document_id``; otherwise ranks badly. + + The unfiltered branch returns the *other* documents — what the real service + did (§5.38): an unfiltered search for a document id returns whatever the + reranker prefers. A reconciler that dropped the filter would confirm the + wrong document as retrievable. + """ + self.search_filters.append(retrieval_filter) + wanted = None + if retrieval_filter: + equals = retrieval_filter.get("equals") or {} + if equals.get("key") == "document_id": + wanted = equals.get("value") + + if wanted is not None: + doc_ids = [wanted] if wanted in self._retrievable else [] + else: + doc_ids = list(self._other_documents) + + return [types.SimpleNamespace(metadata={"document_id": d}) for d in doc_ids] + + +class _FakeAgent: + def __init__(self, owner): + self._owner = owner + + def get_knowledge_base_documents(self, **kwargs): + identifiers = kwargs.get("documentIdentifiers") or [{}] + doc_id = (identifiers[0].get("custom") or {}).get("id") + status = self._owner._statuses.get(doc_id, "NOT_FOUND") + if status == "NOT_FOUND": + return {"documentDetails": []} + return { + "documentDetails": [ + { + "status": status, + "identifier": {"dataSourceType": "CUSTOM", "custom": {"id": doc_id}}, + "updatedAt": datetime(2026, 5, 30, tzinfo=timezone.utc), + } + ] + } + + +def _run(table, backend, **kwargs): + """Run a pass with an injected backend factory returning ``backend``.""" + kwargs.setdefault("now", NOW) + kwargs.setdefault("backend_factory", lambda _assistant_id: backend) + return dr.reconcile_documents(**kwargs) + + +# ── The arming flag ────────────────────────────────────────────────────────── +class TestArmingFlag: + @pytest.mark.parametrize("value", ["", " ", "0", "false", "False", "off", "no", "disabled"]) + def test_falsy_and_empty_values_are_off(self, monkeypatch, value): + monkeypatch.setenv(dr.FLAG_DOC_RECONCILER_ARMED, value) + assert dr.doc_reconciler_armed() is False + + def test_unset_is_off(self, monkeypatch): + monkeypatch.delenv(dr.FLAG_DOC_RECONCILER_ARMED, raising=False) + assert dr.doc_reconciler_armed() is False + + @pytest.mark.parametrize("value", ["1", "true", "TRUE", "yes", "on", "enabled", " true "]) + def test_affirmative_values_arm(self, monkeypatch, value): + monkeypatch.setenv(dr.FLAG_DOC_RECONCILER_ARMED, value) + assert dr.doc_reconciler_armed() is True + + def test_reconcile_defaults_to_the_flag(self, table, monkeypatch): + monkeypatch.setenv(dr.FLAG_DOC_RECONCILER_ARMED, "") + _seed_kb(table, "ast-1") + _seed_doc(table, "ast-1", "doc-1", status="uploading") + backend = _FakeBackend(statuses={"doc-1": "INDEXED"}, retrievable={"doc-1"}) + + report = _run(table, backend) # note: no armed= override, so it reads the flag + + assert report.armed is False + assert report.to_dict()["mode"] == "report-only" + + +# ── The grace gate ─────────────────────────────────────────────────────────── +class TestGraceGate: + def test_a_recently_updated_document_is_left_alone(self, table): + """TRAP: acting on a young row races an ingestion still in flight.""" + _seed_kb(table, "ast-1") + _seed_doc(table, "ast-1", "doc-young", status="uploading", updated_at=YOUNG) + backend = _FakeBackend(statuses={"doc-young": "INDEXED"}, retrievable={"doc-young"}) + + report = _run(table, backend, armed=True) + + assert report.skipped_too_young == ["doc-young"] + assert report.planned_actions == [] + assert _doc(table, "ast-1", "doc-young")["status"] == "uploading" + + def test_a_long_stuck_document_is_reconciled(self, table): + _seed_kb(table, "ast-1") + _seed_doc(table, "ast-1", "doc-old", status="uploading", updated_at=OLD) + backend = _FakeBackend(statuses={"doc-old": "INDEXED"}, retrievable={"doc-old"}) + + report = _run(table, backend, armed=True) + + assert report.skipped_too_young == [] + assert report.stranded == 1 + + def test_missing_or_unparseable_updated_at_fails_closed(self): + assert dr.document_is_stuck_long_enough(None, now=NOW) is False + assert dr.document_is_stuck_long_enough("not-a-date", now=NOW) is False + + def test_the_gate_is_a_pure_function_of_updated_at(self): + old = NOW - timedelta(hours=2) + young = NOW - timedelta(minutes=5) + assert dr.document_is_stuck_long_enough(_iso(old), now=NOW) is True + assert dr.document_is_stuck_long_enough(_iso(young), now=NOW) is False + # Same answer regardless of when it is asked — the property a discovery-time + # clock does not have. + assert dr.document_is_stuck_long_enough(_iso(old), now=NOW + timedelta(days=9)) is True + + def test_min_age_is_read_at_call_time(self, monkeypatch): + stamped = NOW - timedelta(minutes=30) + assert dr.document_is_stuck_long_enough(_iso(stamped), now=NOW) is False + monkeypatch.setattr(dr, "STUCK_MIN_AGE_MINUTES", 10.0) + assert dr.document_is_stuck_long_enough(_iso(stamped), now=NOW) is True + + +# ── The §5.37 fix: stranded-but-retrievable → complete ─────────────────────── +class TestMarkComplete: + def test_a_stranded_retrievable_document_is_completed_when_armed(self, table): + _seed_kb(table, "ast-1") + _seed_doc(table, "ast-1", "doc-1", status="embedding", updated_at=OLD) + backend = _FakeBackend(statuses={"doc-1": "INDEXED"}, retrievable={"doc-1"}) + + report = _run(table, backend, armed=True) + + assert [a.kind for a in report.planned_actions] == [dr.ACTION_MARK_COMPLETE] + assert report.actions_performed == 1 + row = _doc(table, "ast-1", "doc-1") + assert row["status"] == "complete" + assert row["retrievableAt"] + assert row["indexedAt"] + + def test_indexed_but_not_retrievable_is_left_short_of_complete(self, table): + """MUTATION GUARD: dropping the retrievability check would complete this. + + Bedrock says INDEXED, but a filtered retrieval returns nothing — so the + content is not yet queryable. Marking it complete here is the exact bug the + consumer's retrievability poll exists to prevent. + """ + _seed_kb(table, "ast-1") + _seed_doc(table, "ast-1", "doc-1", status="uploading", updated_at=OLD) + backend = _FakeBackend(statuses={"doc-1": "INDEXED"}, retrievable=set()) + + report = _run(table, backend, armed=True) + + assert report.planned_actions == [] + assert report.skipped_not_retrievable == ["doc-1"] + assert _doc(table, "ast-1", "doc-1")["status"] == "uploading" + + def test_partially_indexed_and_retrievable_is_completed(self, table): + _seed_kb(table, "ast-1") + _seed_doc(table, "ast-1", "doc-1", status="chunking", updated_at=OLD) + backend = _FakeBackend(statuses={"doc-1": "PARTIALLY_INDEXED"}, retrievable={"doc-1"}) + + report = _run(table, backend, armed=True) + + assert [a.kind for a in report.planned_actions] == [dr.ACTION_MARK_COMPLETE] + assert _doc(table, "ast-1", "doc-1")["status"] == "complete" + + def test_report_only_plans_but_writes_nothing(self, table): + """MUTATION GUARD: report-only must plan the action and perform none of it.""" + _seed_kb(table, "ast-1") + _seed_doc(table, "ast-1", "doc-1", status="uploading", updated_at=OLD) + backend = _FakeBackend(statuses={"doc-1": "INDEXED"}, retrievable={"doc-1"}) + + report = _run(table, backend, armed=False) + + assert [a.kind for a in report.planned_actions] == [dr.ACTION_MARK_COMPLETE] + assert report.actions_performed == 0 + assert _doc(table, "ast-1", "doc-1")["status"] == "uploading" + assert backend.ingested == [] + + +# ── Bedrock FAILED → failed ────────────────────────────────────────────────── +class TestMarkFailed: + @pytest.mark.parametrize("bedrock_status", ["FAILED", "METADATA_UPDATE_FAILED"]) + def test_a_failed_document_is_marked_failed_when_armed(self, table, bedrock_status): + _seed_kb(table, "ast-1") + _seed_doc(table, "ast-1", "doc-1", status="embedding", updated_at=OLD) + backend = _FakeBackend(statuses={"doc-1": bedrock_status}) + + report = _run(table, backend, armed=True) + + assert [a.kind for a in report.planned_actions] == [dr.ACTION_MARK_FAILED] + row = _doc(table, "ast-1", "doc-1") + assert row["status"] == "failed" + assert row.get("ingestionError") + + def test_a_failed_document_is_not_confused_for_retrievable(self, table): + """A FAILED document must never be probed for retrievability and completed.""" + _seed_kb(table, "ast-1") + _seed_doc(table, "ast-1", "doc-1", status="uploading", updated_at=OLD) + backend = _FakeBackend(statuses={"doc-1": "FAILED"}, retrievable={"doc-1"}) + + report = _run(table, backend, armed=True) + + assert [a.kind for a in report.planned_actions] == [dr.ACTION_MARK_FAILED] + assert _doc(table, "ast-1", "doc-1")["status"] == "failed" + + +# ── Bedrock NOT_FOUND → re-ingest (task 14.4 overlap) ──────────────────────── +class TestReIngest: + def test_a_not_found_document_is_re_ingested_when_armed(self, table): + _seed_kb(table, "ast-1") + _seed_doc( + table, "ast-1", "doc-1", status="uploading", updated_at=OLD, + s3_key="assistants/ast-1/documents/doc-1/report.pdf", filename="report.pdf", + ) + backend = _FakeBackend(statuses={"doc-1": "NOT_FOUND"}) + + report = _run(table, backend, armed=True) + + assert [a.kind for a in report.planned_actions] == [dr.ACTION_RE_INGEST] + assert report.actions_performed == 1 + assert backend.ingested == ["doc-1"] + # Re-ingest re-fires the pipeline; the consumer drives it to complete, so + # the reconciler leaves the row non-terminal rather than claiming success. + assert _doc(table, "ast-1", "doc-1")["status"] == "uploading" + + def test_report_only_does_not_re_ingest(self, table): + _seed_kb(table, "ast-1") + _seed_doc( + table, "ast-1", "doc-1", status="uploading", updated_at=OLD, + s3_key="assistants/ast-1/documents/doc-1/report.pdf", filename="report.pdf", + ) + backend = _FakeBackend(statuses={"doc-1": "NOT_FOUND"}) + + report = _run(table, backend, armed=False) + + assert [a.kind for a in report.planned_actions] == [dr.ACTION_RE_INGEST] + assert backend.ingested == [] + + def test_not_found_without_an_s3_key_is_not_re_ingested(self, table): + """An old row with no s3Key cannot be re-ingested from here.""" + _seed_kb(table, "ast-1") + _seed_doc(table, "ast-1", "doc-1", status="uploading", updated_at=OLD) # no s3Key + backend = _FakeBackend(statuses={"doc-1": "NOT_FOUND"}) + + report = _run(table, backend, armed=True) + + assert report.planned_actions == [] + assert backend.ingested == [] + assert report.skipped_not_retrievable == ["doc-1"] + + +# ── In-flight and unknown statuses are left alone ───────────────────────────── +class TestLeftAlone: + @pytest.mark.parametrize("bedrock_status", ["STARTING", "PENDING", "IN_PROGRESS", "TEXT_INDEXED"]) + def test_in_flight_documents_are_left_alone(self, table, bedrock_status): + _seed_kb(table, "ast-1") + _seed_doc(table, "ast-1", "doc-1", status="uploading", updated_at=OLD) + backend = _FakeBackend(statuses={"doc-1": bedrock_status}) + + report = _run(table, backend, armed=True) + + assert report.planned_actions == [] + assert report.skipped_in_flight == ["doc-1"] + assert _doc(table, "ast-1", "doc-1")["status"] == "uploading" + + def test_an_unknown_bedrock_status_is_treated_as_in_flight(self, table): + """§5.39: the live service returns statuses the SDK enum omits. Unknown + must mean 'keep waiting', not 'give up'.""" + _seed_kb(table, "ast-1") + _seed_doc(table, "ast-1", "doc-1", status="uploading", updated_at=OLD) + backend = _FakeBackend(statuses={"doc-1": "SOME_NEW_STATUS_AWS_ADDED"}) + + report = _run(table, backend, armed=True) + + assert report.planned_actions == [] + assert report.skipped_in_flight == ["doc-1"] + + +# ── Terminal and soft-deleted rows are untouchable ─────────────────────────── +class TestTerminalRowsIgnored: + @pytest.mark.parametrize("status", ["complete", "failed", "deleting"]) + def test_terminal_and_deleting_rows_are_never_candidates(self, table, status): + """MUTATION GUARD: NON_TERMINAL_STATUSES must exclude these. + + ``deleting`` is the dangerous one — a soft-deleted document driven back to + ``complete`` would resurrect content the user removed. + """ + _seed_kb(table, "ast-1") + _seed_doc(table, "ast-1", "doc-1", status=status, updated_at=OLD) + # Bedrock would say INDEXED+retrievable, which WOULD complete a candidate. + backend = _FakeBackend(statuses={"doc-1": "INDEXED"}, retrievable={"doc-1"}) + + report = _run(table, backend, armed=True) + + assert report.stranded == 0 + assert report.planned_actions == [] + assert _doc(table, "ast-1", "doc-1")["status"] == status + + +# ── Engine and provisioning scoping ────────────────────────────────────────── +class TestScoping: + def test_legacy_assistant_documents_are_ignored(self, table): + """A record with no retrievalEngine is legacy; its DOC# status is owned by + the legacy pipeline, not this reconciler.""" + _seed_kb(table, "ast-legacy", engine=None, aws_kb_id=None) + _seed_doc(table, "ast-legacy", "doc-1", status="uploading", updated_at=OLD) + backend = _FakeBackend(statuses={"doc-1": "INDEXED"}, retrievable={"doc-1"}) + + report = _run(table, backend, armed=True) + + assert report.managed_records == 0 + assert report.documents_scanned == 0 + assert report.planned_actions == [] + + def test_an_unprovisioned_managed_record_is_skipped(self, table): + """Managed but no awsKbId: there is no knowledge base to probe yet.""" + _seed_kb(table, "ast-prov", engine="managed", aws_kb_id=None) + _seed_doc(table, "ast-prov", "doc-1", status="uploading", updated_at=OLD) + backend = _FakeBackend(statuses={"doc-1": "INDEXED"}, retrievable={"doc-1"}) + + report = _run(table, backend, armed=True) + + assert report.managed_records == 0 + assert report.planned_actions == [] + + +# ── Per-run action limit ───────────────────────────────────────────────────── +class TestPerRunActionLimit: + def _five_stranded(self, table): + _seed_kb(table, "ast-1") + statuses, retrievable = {}, set() + for i in range(5): + _seed_doc(table, "ast-1", f"doc-{i}", status="uploading", updated_at=OLD) + statuses[f"doc-{i}"] = "INDEXED" + retrievable.add(f"doc-{i}") + return _FakeBackend(statuses=statuses, retrievable=retrievable) + + def test_the_limit_caps_planned_actions_in_report_only_mode(self, table, monkeypatch): + monkeypatch.setattr(dr, "MAX_ACTIONS_PER_RUN", 2) + backend = self._five_stranded(table) + + report = _run(table, backend, armed=False) + + assert len(report.planned_actions) == 2 + assert report.limit_reached is True + + def test_the_limit_caps_actual_actions_when_armed(self, table, monkeypatch): + monkeypatch.setattr(dr, "MAX_ACTIONS_PER_RUN", 2) + backend = self._five_stranded(table) + + report = _run(table, backend, armed=True) + + assert report.actions_performed == 2 + assert report.limit_reached is True + completed = sum( + 1 for i in range(5) if (_doc(table, "ast-1", f"doc-{i}") or {})["status"] == "complete" + ) + assert completed == 2 + + def test_the_environment_can_lower_the_limit_but_not_lift_it(self, monkeypatch): + monkeypatch.setenv("MANAGED_KB_DOC_RECONCILER_MAX_ACTIONS", "3") + assert dr.max_actions_per_run() == 3 + monkeypatch.setenv("MANAGED_KB_DOC_RECONCILER_MAX_ACTIONS", "1000000") + assert dr.max_actions_per_run() == dr.MAX_ACTIONS_CEILING + + def test_a_negative_limit_does_not_become_unbounded(self, monkeypatch): + monkeypatch.setenv("MANAGED_KB_DOC_RECONCILER_MAX_ACTIONS", "-5") + assert dr.max_actions_per_run() == 0 + + def test_the_limit_is_read_at_call_time(self, monkeypatch): + assert dr.max_actions_per_run() == dr.MAX_ACTIONS_PER_RUN + monkeypatch.setattr(dr, "MAX_ACTIONS_PER_RUN", 3) + assert dr.max_actions_per_run() == 3 + + +# ── Retrievability probe (§5.38) ───────────────────────────────────────────── +class TestRetrievabilityProbe: + def test_it_filters_on_document_id_by_equals(self, table): + """MUTATION GUARD: an unfiltered search returns the wrong document (§5.38). + + The backend returns OTHER documents when unfiltered. ``is_retrievable`` must + pass the ``equals`` filter, so a document not in the retrievable set returns + False even though the search would otherwise return chunks. + """ + backend = _FakeBackend(retrievable=set()) # nothing is retrievable + + assert dr.is_retrievable(backend, "ast-1", "doc-1") is False + assert backend.search_filters == [{"equals": {"key": "document_id", "value": "doc-1"}}] + + def test_it_returns_true_only_for_its_own_document(self, table): + backend = _FakeBackend(retrievable={"doc-1"}) + assert dr.is_retrievable(backend, "ast-1", "doc-1") is True + assert dr.is_retrievable(backend, "ast-1", "doc-other") is False + + def test_a_probe_error_is_not_a_positive(self): + class Boom: + async def search(self, *a, **k): + raise RuntimeError("retrieve failed") + + assert dr.is_retrievable(Boom(), "ast-1", "doc-1") is False + + +# ── The lambda handler ─────────────────────────────────────────────────────── +class TestLambdaHandler: + @pytest.fixture() + def one_stranded(self, table, monkeypatch): + _seed_kb(table, "ast-1") + _seed_doc(table, "ast-1", "doc-1", status="uploading", updated_at=OLD) + backend = _FakeBackend(statuses={"doc-1": "INDEXED"}, retrievable={"doc-1"}) + # lambda_handler builds its own backend via the default factory, which we + # cannot inject through the event — so patch the factory builder. + monkeypatch.setattr(dr, "_default_backend_factory", lambda client: (lambda _a: backend)) + # And pin 'now' well ahead of the row's OLD stamp is unnecessary: OLD is + # relative to a fixed 2026 date, comfortably older than real wall-clock. + return backend + + def test_it_returns_the_serialized_report(self, table, monkeypatch): + monkeypatch.setattr( + dr, "reconcile_documents", lambda **kw: dr.DocumentReconcileReport(armed=False) + ) + result = dr.lambda_handler({}, None) + assert result["statusCode"] == 200 + assert result["report"]["mode"] == "report-only" + + @pytest.mark.parametrize("payload", [True, "true", 1, "1", "yes"]) + def test_the_event_cannot_arm_the_reconciler(self, table, one_stranded, payload): + result = dr.lambda_handler({"armed": payload}, None) + assert result["report"]["mode"] == "report-only" + assert result["report"]["actionsPerformed"] == 0 + assert one_stranded.ingested == [] + assert _doc(table, "ast-1", "doc-1")["status"] == "uploading" + # The finding is still reported — suppressing the write did not suppress it. + assert result["report"]["stranded"] == 1 + + def test_the_flag_is_what_arms_it(self, table, one_stranded, monkeypatch): + monkeypatch.setenv(dr.FLAG_DOC_RECONCILER_ARMED, "true") + result = dr.lambda_handler({"armed": False}, None) + assert result["report"]["mode"] == "armed" + assert result["report"]["actionsPerformed"] == 1 + assert _doc(table, "ast-1", "doc-1")["status"] == "complete" + + def test_an_ignored_arming_request_is_logged(self, table, one_stranded, caplog): + import logging + + with caplog.at_level(logging.WARNING): + dr.lambda_handler({"armed": True}, None) + assert any( + "ignoring armed" in r.message and dr.FLAG_DOC_RECONCILER_ARMED in r.message + for r in caplog.records + ) + + +# ── Mixed and degenerate cases ─────────────────────────────────────────────── +class TestMixedRun: + def test_all_outcomes_in_one_pass(self, table): + _seed_kb(table, "ast-1") + _seed_doc(table, "ast-1", "d-complete", status="uploading", updated_at=OLD) + _seed_doc(table, "ast-1", "d-failed", status="embedding", updated_at=OLD) + _seed_doc( + table, "ast-1", "d-reingest", status="uploading", updated_at=OLD, + s3_key="assistants/ast-1/documents/d-reingest/x.pdf", filename="x.pdf", + ) + _seed_doc(table, "ast-1", "d-inflight", status="chunking", updated_at=OLD) + _seed_doc(table, "ast-1", "d-young", status="uploading", updated_at=YOUNG) + _seed_doc(table, "ast-1", "d-done", status="complete", updated_at=OLD) + backend = _FakeBackend( + statuses={ + "d-complete": "INDEXED", + "d-failed": "FAILED", + "d-reingest": "NOT_FOUND", + "d-inflight": "IN_PROGRESS", + }, + retrievable={"d-complete"}, + ) + + report = _run(table, backend, armed=True) + + kinds = {a.document_id: a.kind for a in report.planned_actions} + assert kinds == { + "d-complete": dr.ACTION_MARK_COMPLETE, + "d-failed": dr.ACTION_MARK_FAILED, + "d-reingest": dr.ACTION_RE_INGEST, + } + assert report.skipped_in_flight == ["d-inflight"] + assert report.skipped_too_young == ["d-young"] + assert _doc(table, "ast-1", "d-complete")["status"] == "complete" + assert _doc(table, "ast-1", "d-failed")["status"] == "failed" + assert _doc(table, "ast-1", "d-done")["status"] == "complete" # untouched + + def test_an_empty_table_is_a_clean_no_op(self, table): + backend = _FakeBackend() + report = _run(table, backend, armed=True) + payload = report.to_dict() + assert payload["managedRecords"] == 0 + assert payload["stranded"] == 0 + assert payload["plannedActions"] == [] + assert payload["actionsPerformed"] == 0 + + def test_a_failing_write_does_not_end_the_run(self, table, monkeypatch): + """One bad document must not stop the reconciler reaching the others.""" + _seed_kb(table, "ast-1") + _seed_doc(table, "ast-1", "doc-a", status="uploading", updated_at=OLD) + _seed_doc(table, "ast-1", "doc-b", status="uploading", updated_at=OLD) + backend = _FakeBackend( + statuses={"doc-a": "INDEXED", "doc-b": "INDEXED"}, + retrievable={"doc-a", "doc-b"}, + ) + + calls = {"n": 0} + real = dr.ic.set_document_terminal + + def flaky(*args, **kwargs): + calls["n"] += 1 + if calls["n"] == 1: + raise RuntimeError("dynamo throttled") + return real(*args, **kwargs) + + monkeypatch.setattr(dr.ic, "set_document_terminal", flaky) + + report = _run(table, backend, armed=True) + + assert len(report.planned_actions) == 2 + assert report.actions_performed == 1 + assert any(a.error for a in report.planned_actions) diff --git a/backend/tests/supply_chain/test_kb_migration_env_contract.py b/backend/tests/supply_chain/test_kb_migration_env_contract.py index e535343b..4ade0b58 100644 --- a/backend/tests/supply_chain/test_kb_migration_env_contract.py +++ b/backend/tests/supply_chain/test_kb_migration_env_contract.py @@ -57,6 +57,14 @@ "MANAGED_KB_ORPHAN_MIN_AGE_HOURS", "MANAGED_KB_RECONCILER_MAX_DELETIONS", "MANAGED_KB_RECONCILER_MAX_SCANNED", + # Document-reconciler arming flag (task 16.5). Report-only by default and an + # empty string reads as off, so the module is correct with the variable unset. + # The document reconciler's own Lambda + schedule + arming are a deliberate + # deploy-gated follow-up (not wired in the PR that added document_reconciler.py), + # so the construct is not yet required to set it. When that Lambda is added, + # this exemption should go and the construct should set the flag (off), the way + # MANAGED_KB_RECONCILER_ARMED is spot-pinned below. + "MANAGED_KB_DOC_RECONCILER_ARMED", # Throttle for the last-retrieved write, defaulted in idleness.py. "KB_LAST_RETRIEVED_THROTTLE_HOURS", # Namespace override; metrics.py derives one from the project prefix.