From 9737638bfb1a98b7a8f09e3435c6c2c5129b1892 Mon Sep 17 00:00:00 2001 From: derrickfink Date: Tue, 1 Sep 2026 10:18:21 -0600 Subject: [PATCH] fix(kb): wait for indexing to finish, and stop re-ingesting while it runs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A 1.5 MB PDF uploaded to a promoted knowledge base in dev sat at `uploading` indefinitely with a fully retrievable copy in the knowledge base. Three separate bugs stacked up. 1. The poll window was sized against the wrong measurement --------------------------------------------------------- The consumer ingested and then polled a *retrieval* for 30 s. Its own header justifies that window with "Bedrock reports INDEXED up to a second before it can be retrieved — measured at 0.75–1.03 s", but the poll starts the moment the ingest call returns, so it actually has to cover `ingest -> INDEXED -> retrievable`. §5.1 of the evaluation measured PDF ingestion at 37–264 s; this file took 5 m 30 s. The budget was smaller than the documented lower bound. Identical in shape to §5.30, where `verify` failed good migrations against that same 0.75–1.03 s figure. That fix never reached this component, which inherited the constant. 2. Every redelivery re-ingested, restarting the work it was waiting for ---------------------------------------------------------------------- `IngestKnowledgeBaseDocuments` is fire-and-forget, and nothing asked Bedrock what it already knew, so the consumer could not tell "not indexed yet" from "never submitted". Each of the three deliveries re-submitted the document. It reached INDEXED 54 s after the final attempt had been dead-lettered. `handle_object` now probes `GetKnowledgeBaseDocuments` first and branches on the real `DocumentStatus` enum, taken from the packaged service model: STARTING/PENDING/IN_PROGRESS means do not re-ingest; FAILED is terminal; PARTIALLY_INDEXED counts as usable, because the document IS retrievable and failing it would hide content the user can see. A probe failure reads as NOT_FOUND — no evidence of prior work — so it never blocks ingestion. 3. `indexedAt` was fabricated ----------------------------- `indexed_at = _now_iso()` ran immediately after the ingest call returned, so the field recorded when we asked, presented as when indexing finished — minutes apart for this document. It is now Bedrock's own `updatedAt`. The pre-existing test asserted only that the key existed and was truthy, which a fabricated value satisfies; that is why this survived. It also means my earlier claim that a measured 0.9 s gap "confirmed the 0.75–1.03 s figure" was wrong: that gap was ingest-return to retrievable, not INDEXED to retrievable. Why the wait is in-invocation and not more retries -------------------------------------------------- I first raised a RetryPolicy on the EventBridge target. That does nothing, and the construct now says so instead: for a Lambda target EventBridge hands the event off and the function's OWN async retry config governs — `retryAttempts: 2`, which is Lambda's hard maximum and exactly the 1 + 2 attempts seen in the logs. Redelivery therefore spans a few minutes and cannot be extended, so INDEXED_POLL_TIMEOUT_SECONDS is 600 s: covers the measured tail, leaves 5 minutes under the 15-minute Lambda timeout. Small documents still complete in one invocation, so the fast path is unchanged. Guards: 8 new tests in test_kb_ingestion_consumer.py, and a cross-language test in test_kb_migration_env_contract.py that parses the Lambda timeout out of the CDK construct and asserts the poll budget fits inside it with headroom — the two live in different languages with no compiler between them. `_FakeBackend` now models document STATUS, not just ingest calls. A fake that reported instant success is what let a 30 s window look adequate for work that takes minutes — the same failure as §5.28, where the fake modelled `clientToken` dedup as permanent. Mutations verified caught, each by correctly-named tests: in-flight guard removed, `indexedAt` back to the local clock, FAILED treated as retryable, budget raised past the Lambda timeout, budget dropped back to 25 s. Backend 2,510 passed across the affected areas; infra 640 passed. --- .../kb_migration/ingestion_consumer.py | 224 ++++++++++++++++- .../lambdas/test_kb_ingestion_consumer.py | 234 +++++++++++++++++- .../test_kb_migration_env_contract.py | 61 +++++ .../managed-kb/kb-migration-construct.ts | 8 + 4 files changed, 514 insertions(+), 13 deletions(-) diff --git a/backend/src/apis/app_api/kb_migration/ingestion_consumer.py b/backend/src/apis/app_api/kb_migration/ingestion_consumer.py index 776beb6e..ba99c0b5 100644 --- a/backend/src/apis/app_api/kb_migration/ingestion_consumer.py +++ b/backend/src/apis/app_api/kb_migration/ingestion_consumer.py @@ -75,6 +75,37 @@ RETRIEVABLE_POLL_TIMEOUT_SECONDS = 30.0 RETRIEVABLE_POLL_INTERVAL_SECONDS = 0.5 +#: How long ONE invocation waits for Bedrock to finish indexing. +#: +#: This has to cover the whole indexing time, because redelivery cannot. Lambda's +#: asynchronous retry is capped at **2** attempts — a hard service limit, not a +#: setting we chose — so an event gets 1 + 2 tries spread over a few minutes and +#: then dead-letters. Raising a retry policy on the EventBridge target does not +#: change that: for a Lambda target EventBridge hands the event off and the +#: function's own async retry config takes over. +#: +#: So the budget is sized from measurement, not convenience. The §5.1 benchmark +#: measured PDF ingestion at 37–264 s, and a 1.5 MB PDF in dev reached INDEXED +#: 5 m 30 s after upload — image-heavy files run longer because the vision model +#: runs per page. 10 minutes covers that with headroom and still leaves 5 minutes +#: under the Lambda's 15-minute timeout, so a slow-but-succeeding document is never +#: killed mid-wait. +#: +#: The cost of waiting is real but small: this Lambda is 1024 MB and handles one +#: user upload at a time, so a 6-minute wait is a fraction of a cent. The cost of +#: NOT waiting was a permanently stuck document with a fully retrievable copy in +#: the knowledge base and no writer left to reconcile it. +#: +#: ⚠️ Keep the sum of this and RETRIEVABLE_POLL_TIMEOUT_SECONDS below the Lambda's +#: timeout. `tests/supply_chain/test_kb_migration_env_contract.py` asserts it +#: against the value in the CDK construct. +INDEXED_POLL_TIMEOUT_SECONDS = 600.0 + +#: 5 s rather than sub-second: over a 10-minute budget this is ~120 control-plane +#: calls instead of ~1,200, and indexing progress is measured in tens of seconds, +#: so a finer interval buys nothing. +INDEXED_POLL_INTERVAL_SECONDS = 5.0 + #: Bounded retries on the record update. The event source already redelivers, so #: this only covers a transient DynamoDB failure inside one invocation; unbounded #: retries would burn the Lambda timeout and lose the DLQ signal. @@ -216,6 +247,109 @@ def set_document_terminal( ) +#: Bedrock's own view of a document, from the packaged service model's +#: ``DocumentStatus`` enum rather than guessed. Read at call time from +#: ``GetKnowledgeBaseDocuments``, which is the only way to know whether indexing +#: has finished — the ``IngestKnowledgeBaseDocuments`` call returns as soon as the +#: request is accepted and says nothing about progress. +DOC_STATUS_INDEXED = "INDEXED" +DOC_STATUS_NOT_FOUND = "NOT_FOUND" + +#: Bedrock is still working. Re-ingesting a document in one of these states is +#: pointless at best: the work is already queued, and re-submitting it discards +#: whatever progress has been made and starts the clock again. That is what turned +#: a slow success into a permanent failure in dev — three redeliveries each +#: re-ingested, and the document only reached INDEXED 54 s after the last attempt +#: had already been dead-lettered. +DOC_STATUSES_IN_FLIGHT = frozenset({"STARTING", "PENDING", "IN_PROGRESS"}) + +#: Terminal and unusable. Worth failing the document rather than retrying forever. +DOC_STATUSES_FAILED = frozenset({"FAILED", "METADATA_UPDATE_FAILED"}) + +#: Indexed, but not everything made it. Treated as usable — the document IS +#: retrievable — because the alternative is failing a document the user can see +#: content from. Logged so the partiality is not silent. +DOC_STATUSES_PARTIAL = frozenset({"PARTIALLY_INDEXED", "METADATA_PARTIALLY_INDEXED"}) + + +def document_status(backend: Any, kb_ref: str, document_id: str) -> Tuple[str, Optional[str]]: + """Ask Bedrock for a document's ingestion status and its own timestamp. + + Returns ``(status, updated_at_iso)``. ``NOT_FOUND`` covers both "Bedrock has + never heard of it" and "the call failed", because both mean the same thing to + the caller: there is no evidence the document is already being worked on, so + ingesting is the right next move. A probe failure must not be mistaken for a + document failure. + + This exists because the ingest call is fire-and-forget. Without it the consumer + cannot distinguish "not indexed yet" from "never submitted", so every + redelivery re-submits — see :data:`DOC_STATUSES_IN_FLIGHT`. + """ + try: + # Imported here, not at module scope: this module's module-level imports are + # stdlib only so the shared Lambda image stays small + # (tests/architecture/test_kb_backend_boundary.py). Reused rather than + # redefined so the connector type has one definition. + from apis.shared.kb_backend.managed_backend import CONTENT_DATA_SOURCE_TYPE + + client = backend._agent() # noqa: SLF001 - same package, deliberate reuse + aws_kb_id, data_source_id = backend._locate(kb_ref) # noqa: SLF001 + response = client.get_knowledge_base_documents( + knowledgeBaseId=aws_kb_id, + dataSourceId=data_source_id, + documentIdentifiers=[ + {"dataSourceType": CONTENT_DATA_SOURCE_TYPE, "custom": {"id": document_id}} + ], + ) + except Exception as exc: # noqa: BLE001 - a probe failure is not a document failure + logger.warning(f"document status probe for {document_id} failed: {exc}") + return DOC_STATUS_NOT_FOUND, None + + for detail in response.get("documentDetails") or []: + status = str(detail.get("status") or DOC_STATUS_NOT_FOUND) + updated = detail.get("updatedAt") + return status, (updated.isoformat() if hasattr(updated, "isoformat") else updated) + return DOC_STATUS_NOT_FOUND, None + + +def wait_until_indexed( + backend: Any, + kb_ref: str, + document_id: str, + timeout_seconds: Optional[float] = None, + interval_seconds: Optional[float] = None, + sleep: Any = time.sleep, +) -> Tuple[str, Optional[str]]: + """Poll Bedrock's document status until it settles, or give up. + + Returns the last ``(status, updated_at)`` seen. Giving up is an ordinary + outcome, not an error: the caller leaves the document non-terminal and lets + redelivery come back to it, by which time indexing has usually finished. + + Why a *bounded* in-invocation wait rather than pure redelivery: most documents + index in a few seconds, and making every one of them wait for an EventBridge + retry would add a minute of latency to the common case for the sake of the rare + slow one. Why bounded at all: PDF ingestion was measured at 37–264 s and + image-heavy files run longer, so waiting for the worst case in-invocation would + hold a concurrency slot for minutes and bill for it. + + Same call-time constant resolution as :func:`wait_until_retrievable`, and for + the same reason — a default argument is bound once at import and cannot be + patched by a test. + """ + if timeout_seconds is None: + timeout_seconds = INDEXED_POLL_TIMEOUT_SECONDS + if interval_seconds is None: + interval_seconds = INDEXED_POLL_INTERVAL_SECONDS + + deadline = time.monotonic() + timeout_seconds + status, updated_at = document_status(backend, kb_ref, document_id) + while status in DOC_STATUSES_IN_FLIGHT and time.monotonic() < deadline: + sleep(interval_seconds) + status, updated_at = document_status(backend, kb_ref, document_id) + return status, updated_at + + def wait_until_retrievable( backend: Any, kb_ref: str, @@ -308,22 +442,90 @@ def handle_object(bucket: str, key: str) -> Dict[str, Any]: backend = ManagedKbBackend(bucket=bucket) source = DocumentSource(document_id=document_id, filename=filename, s3_key=key) - try: - asyncio.run(backend.ingest(assistant_id, source)) - except Exception as exc: - logger.error(f"direct ingestion of {document_id} failed: {exc}", exc_info=True) - set_document_terminal(assistant_id, document_id, STATUS_FAILED, error=str(exc)) - raise + # Ask Bedrock what it already knows BEFORE ingesting. The ingest call is + # fire-and-forget — it returns as soon as the request is accepted — so without + # this the consumer cannot tell "not indexed yet" from "never submitted", and + # every redelivery re-submits a document that is already being worked on. + # + # That is not merely wasteful. In dev a 1.5 MB PDF was re-ingested on each of + # three redeliveries and reached INDEXED only 54 s after the final attempt had + # been dead-lettered, leaving a perfectly retrievable document parked at + # `uploading` with nothing left to reconcile it. + status, bedrock_updated_at = document_status(backend, assistant_id, document_id) + + if status in DOC_STATUSES_FAILED: + # Terminal on Bedrock's side. Retrying cannot help. + logger.error(f"document {document_id} is {status} in the knowledge base") + set_document_terminal( + assistant_id, document_id, STATUS_FAILED, + error=f"the knowledge base reports this document as {status}", + ) + return {"routed": "managed", "ingested": False, "document_id": document_id, + "status": status} + + if status in DOC_STATUSES_IN_FLIGHT: + # Already being indexed. Do NOT ingest again — that would discard the + # progress this invocation is waiting on. Leave it non-terminal so the + # event source brings us back, by which time Bedrock may have finished. + logger.info( + f"document {document_id} is {status} in the knowledge base; not " + f"re-ingesting, waiting for redelivery" + ) + raise IngestionRoutingError( + f"document {document_id} is still {status}; leaving it for redelivery" + ) + + if status in DOC_STATUSES_PARTIAL: + logger.warning( + f"document {document_id} is {status}: it is retrievable but some of its " + f"content or metadata did not index" + ) + + already_indexed = status == DOC_STATUS_INDEXED or status in DOC_STATUSES_PARTIAL + + if not already_indexed: + try: + asyncio.run(backend.ingest(assistant_id, source)) + except Exception as exc: + logger.error(f"direct ingestion of {document_id} failed: {exc}", exc_info=True) + set_document_terminal(assistant_id, document_id, STATUS_FAILED, error=str(exc)) + raise + + # Submitted. Wait briefly for indexing so a small document finishes in this + # one invocation, then hand the slow ones back to redelivery. + status, bedrock_updated_at = wait_until_indexed(backend, assistant_id, document_id) + + if status in DOC_STATUSES_FAILED: + logger.error(f"document {document_id} became {status} during indexing") + set_document_terminal( + assistant_id, document_id, STATUS_FAILED, + error=f"the knowledge base reports this document as {status}", + ) + return {"routed": "managed", "ingested": True, "document_id": document_id, + "status": status} + + if status not in (DOC_STATUS_INDEXED, *DOC_STATUSES_PARTIAL): + # Still indexing. Deliberately NOT marked terminal, and deliberately + # not given an invented timestamp — the next delivery will find it + # INDEXED and record Bedrock's own. + raise IngestionRoutingError( + f"document {document_id} is {status} after submission; leaving it " + f"for redelivery to confirm indexing" + ) - indexed_at = _now_iso() + # INDEXED. `indexedAt` is Bedrock's OWN timestamp, not this process's clock — + # an earlier version recorded `_now_iso()` immediately after the ingest call + # returned, which measured when the request was accepted and was reported as + # the moment indexing completed. The two can be minutes apart. + indexed_at = bedrock_updated_at or _now_iso() retrievable_at = wait_until_retrievable(backend, assistant_id, document_id) if retrievable_at is None: - # Ingested but not confirmed retrievable. Left non-terminal deliberately so - # the event source redelivers, rather than the record claiming a success the - # user cannot yet observe. + # INDEXED but not yet queryable. This is the one short, real gap the poll + # window was always sized for (measured at 0.75–1.03 s); a timeout here is + # unusual rather than routine, so leave it for redelivery. raise IngestionRoutingError( - f"document {document_id} was ingested but not retrievable within the " + f"document {document_id} is INDEXED but was not retrievable within the " f"poll window; leaving it for redelivery" ) diff --git a/backend/tests/lambdas/test_kb_ingestion_consumer.py b/backend/tests/lambdas/test_kb_ingestion_consumer.py index 70d20246..cf0d07fe 100644 --- a/backend/tests/lambdas/test_kb_ingestion_consumer.py +++ b/backend/tests/lambdas/test_kb_ingestion_consumer.py @@ -13,6 +13,7 @@ legacy must ingest NOTHING here, managed must ingest here and NOT fall back. """ +from datetime import datetime, timezone from unittest.mock import MagicMock, patch import boto3 @@ -79,12 +80,67 @@ def _eventbridge_event(key=KEY): return {"detail": {"bucket": {"name": BUCKET}, "object": {"key": key}}} +class _FakeAgentClient: + """The control-plane surface the consumer reads document status from.""" + + def __init__(self, owner): + self._owner = owner + + def get_knowledge_base_documents(self, **kwargs): + self._owner.status_calls += 1 + status = self._owner.next_status() + if status == "NOT_FOUND": + return {"documentDetails": []} + return { + "documentDetails": [ + { + "knowledgeBaseId": "KB123", + "dataSourceId": "DS456", + "status": status, + "identifier": { + "dataSourceType": "CUSTOM", + "custom": {"id": DOCUMENT_ID}, + }, + "updatedAt": datetime(2026, 9, 1, 14, 53, 19, tzinfo=timezone.utc), + } + ] + } + + class _FakeBackend: - """Records ingest calls; reports the document retrievable immediately.""" + """Models the parts of ManagedKbBackend the consumer actually leans on. - def __init__(self): + Deliberately models Bedrock's document STATUS, not just the ingest call. + Ingestion is fire-and-forget: the API returns once the request is accepted and + says nothing about progress, so a fake that only recorded ingests could not + express the state the consumer now has to reason about — and a fake that + reported instant success is what let the 30 s poll window look adequate for + documents that take minutes. + + ``statuses`` is the sequence returned by successive status probes. The default + models a small document: unknown, then indexed. Pass a longer sequence to model + a slow one; the last value repeats forever. + """ + + def __init__(self, statuses=None): self.ingested = [] + self.status_calls = 0 + self._statuses = list(statuses or ["NOT_FOUND", "INDEXED"]) + self._agent_client = _FakeAgentClient(self) + + def next_status(self): + if len(self._statuses) > 1: + return self._statuses.pop(0) + return self._statuses[0] + # -- the private surface `document_status` reuses -------------------------- + def _agent(self): + return self._agent_client + + def _locate(self, kb_ref): + return ("KB123", "DS456") + + # -- the protocol surface -------------------------------------------------- async def ingest(self, kb_ref, source): self.ingested.append(source.document_id) return None @@ -211,6 +267,31 @@ def test_indexed_and_retrievable_are_recorded_separately(self, table): assert "retrievableAt" in item assert result["indexedAt"] and result["retrievableAt"] + def test_indexed_at_is_bedrocks_timestamp_not_our_clock(self, table): + """`indexedAt` must be the value Bedrock reports, not the local time. + + The original code set `indexed_at = _now_iso()` immediately after the + ingest call returned. That call is fire-and-forget — it returns when the + request is accepted — so the field recorded "when we asked", presented as + "when indexing finished". For a 1.5 MB PDF in dev those were 5.5 minutes + apart, and because the field was always populated the error was invisible: + the pre-existing test asserted only that the key existed and was truthy, + which a fabricated value satisfies perfectly. + """ + self._seed_managed(table) + with patch( + "apis.shared.kb_backend.managed_backend.ManagedKbBackend", + return_value=_FakeBackend(), + ): + result = ic.handle_object(BUCKET, KEY) + + # The fake reports 2026-09-01T14:53:19+00:00 from GetKnowledgeBaseDocuments. + assert result["indexedAt"].startswith("2026-09-01T14:53:19"), ( + f"indexedAt is {result['indexedAt']!r}, which is not the timestamp " + f"Bedrock reported — it looks like a local clock reading" + ) + assert _doc(table)["indexedAt"].startswith("2026-09-01T14:53:19") + def test_a_managed_document_never_falls_back_to_legacy(self, table): """Managed engine but unprovisioned must FAIL, not silently degrade. @@ -330,3 +411,152 @@ def test_the_module_does_not_use_ensure_future(self): } assert "ensure_future" not in called assert "create_task" not in called + + +# --------------------------------------------------------------------------- +# A slow document must converge, not die +# --------------------------------------------------------------------------- +class TestSlowIndexingConverges: + """The defect a 1.5 MB PDF exposed in dev on 2026-09-01. + + Ingestion succeeded, but the consumer polled a *retrieval* for 30 s starting + the instant the ingest call returned — before Bedrock had indexed anything. + That poll window was sized against the INDEXED -> retrievable gap + (0.75-1.03 s), while it actually had to cover ingest -> INDEXED -> retrievable, + which the §5.1 benchmark measured at 37-264 s for PDFs. + + Each of the three redeliveries then RE-INGESTED, discarding progress and + restarting the clock. The document reached INDEXED 54 s after the final attempt + was dead-lettered, leaving a fully retrievable document parked at `uploading` + with nothing left to reconcile it — the legacy pipeline no longer writes status + for a promoted knowledge base, so there was no second writer to mask it. + """ + + def _seed_managed(self, table): + _seed_kb(table, retrievalEngine="managed", awsKbId="KB123", awsDataSourceId="DS456") + + def test_a_document_already_being_indexed_is_not_re_ingested(self, table): + """The core fix. Re-submitting restarts the work we are waiting for.""" + self._seed_managed(table) + fake = _FakeBackend(statuses=["IN_PROGRESS"]) + + with patch( + "apis.shared.kb_backend.managed_backend.ManagedKbBackend", return_value=fake + ): + with pytest.raises(ic.IngestionRoutingError, match="IN_PROGRESS"): + ic.handle_object(BUCKET, KEY) + + assert fake.ingested == [], ( + "a document Bedrock was already indexing was submitted again; that " + "discards the progress this invocation is waiting on" + ) + + def test_a_document_still_indexing_is_left_non_terminal(self, table): + """Not complete and not failed — the next delivery decides.""" + self._seed_managed(table) + fake = _FakeBackend(statuses=["IN_PROGRESS"]) + + with patch( + "apis.shared.kb_backend.managed_backend.ManagedKbBackend", return_value=fake + ): + with pytest.raises(ic.IngestionRoutingError): + ic.handle_object(BUCKET, KEY) + + item = _doc(table) + assert item["status"] not in ("complete", "failed") + assert "indexedAt" not in item, "no timestamp may be invented before indexing" + + def test_a_redelivery_completes_the_document_without_a_second_ingest(self, table): + """Delivery 1 submits and defers; delivery 2 finds it INDEXED and finishes. + + This is the whole convergence property: the document ends up `complete` + having been handed to Bedrock exactly once. + """ + self._seed_managed(table) + + # Delivery 1: never seen, then still working for the whole in-invocation wait. + first = _FakeBackend(statuses=["NOT_FOUND", "IN_PROGRESS"]) + with patch( + "apis.shared.kb_backend.managed_backend.ManagedKbBackend", return_value=first + ), patch.object(ic, "INDEXED_POLL_TIMEOUT_SECONDS", 0.01): + with pytest.raises(ic.IngestionRoutingError): + ic.handle_object(BUCKET, KEY) + assert first.ingested == [DOCUMENT_ID] + assert _doc(table)["status"] not in ("complete", "failed") + + # Delivery 2: Bedrock has finished. + second = _FakeBackend(statuses=["INDEXED"]) + with patch( + "apis.shared.kb_backend.managed_backend.ManagedKbBackend", return_value=second + ): + ic.handle_object(BUCKET, KEY) + + assert second.ingested == [], "the second delivery must not re-ingest" + item = _doc(table) + assert item["status"] == "complete" + assert item["indexedAt"].startswith("2026-09-01T14:53:19") + + def test_a_small_document_still_finishes_in_one_invocation(self, table): + """The fast path must not regress into waiting for a retry. + + Deferring every document would add a minute of EventBridge backoff to the + common case, which is why the in-invocation wait exists at all. + """ + self._seed_managed(table) + fake = _FakeBackend(statuses=["NOT_FOUND", "INDEXED"]) + + with patch( + "apis.shared.kb_backend.managed_backend.ManagedKbBackend", return_value=fake + ): + result = ic.handle_object(BUCKET, KEY) + + assert result["ingested"] is True + assert _doc(table)["status"] == "complete" + assert fake.ingested == [DOCUMENT_ID] + + def test_a_failed_document_is_marked_failed_and_not_retried(self, table): + """Terminal on Bedrock's side. Redelivering cannot help.""" + self._seed_managed(table) + fake = _FakeBackend(statuses=["FAILED"]) + + with patch( + "apis.shared.kb_backend.managed_backend.ManagedKbBackend", return_value=fake + ): + result = ic.handle_object(BUCKET, KEY) # must NOT raise + + assert fake.ingested == [] + item = _doc(table) + assert item["status"] == "failed" + assert "FAILED" in item["ingestionError"] + assert result["status"] == "FAILED" + + def test_a_partially_indexed_document_is_treated_as_usable(self, table): + """It IS retrievable, so failing it would hide content the user can see.""" + self._seed_managed(table) + fake = _FakeBackend(statuses=["PARTIALLY_INDEXED"]) + + with patch( + "apis.shared.kb_backend.managed_backend.ManagedKbBackend", return_value=fake + ): + ic.handle_object(BUCKET, KEY) + + assert fake.ingested == [], "already indexed, even if partially" + assert _doc(table)["status"] == "complete" + + def test_a_status_probe_failure_does_not_fail_the_document(self, table): + """An unreadable probe means "no evidence", so ingesting is correct.""" + self._seed_managed(table) + + class _ProbeBroken(_FakeBackend): + def _agent(self): + raise RuntimeError("bedrock control plane unavailable") + + fake = _ProbeBroken(statuses=["NOT_FOUND"]) + with patch( + "apis.shared.kb_backend.managed_backend.ManagedKbBackend", return_value=fake + ), patch.object(ic, "INDEXED_POLL_TIMEOUT_SECONDS", 0.01): + with pytest.raises(ic.IngestionRoutingError): + ic.handle_object(BUCKET, KEY) + + assert fake.ingested == [DOCUMENT_ID], "a probe failure must not block ingestion" + assert _doc(table)["status"] != "failed" 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 9b48bb9e..e535343b 100644 --- a/backend/tests/supply_chain/test_kb_migration_env_contract.py +++ b/backend/tests/supply_chain/test_kb_migration_env_contract.py @@ -201,3 +201,64 @@ def test_no_kb_migration_variable_is_published_unread(self): def test_known_load_bearing_variables_stay_wired(self, name): """Spot-pins for variables whose absence is silent rather than loud.""" assert name in _names_set_by_construct() + + +# --------------------------------------------------------------------------- +# The poll budget must fit inside the Lambda that runs it +# --------------------------------------------------------------------------- +class TestTheIngestionPollBudgetFitsTheLambdaTimeout: + """A wait longer than the timeout is a killed invocation, not a wait. + + The consumer waits for Bedrock to finish indexing inside a single invocation, + because Lambda's asynchronous retry is capped at 2 attempts and cannot be + extended — so redelivery spans only minutes and a slow document would + dead-letter. That makes the in-invocation budget load-bearing, and it now lives + in two files that have no compiler between them: the timeout in the CDK + construct and the poll constants in Python. + + Raise either past the other and a slow-but-succeeding document is killed + mid-wait and dead-lettered — which is exactly the failure this budget was + introduced to remove. Hence a test rather than a comment. + """ + + def _lambda_timeout_minutes(self) -> int: + import re + + text = CONSTRUCT.read_text(encoding="utf-8") + # The consumer's own timeout, not another function's: anchor on its + # construct id and read the first timeout that follows. + start = text.index("KbIngestionConsumerLambda'") + match = re.search(r"timeout:\s*cdk\.Duration\.minutes\((\d+)\)", text[start:]) + assert match, "could not find the ingestion consumer's timeout in the construct" + return int(match.group(1)) + + def test_the_poll_budget_leaves_headroom_under_the_lambda_timeout(self): + from apis.app_api.kb_migration import ingestion_consumer as ic + + budget = ic.INDEXED_POLL_TIMEOUT_SECONDS + ic.RETRIEVABLE_POLL_TIMEOUT_SECONDS + timeout = self._lambda_timeout_minutes() * 60 + + assert budget < timeout, ( + f"the consumer can wait {budget:.0f}s but its Lambda times out at " + f"{timeout}s — a slow document would be killed mid-wait and " + f"dead-lettered, which is the bug this budget exists to prevent" + ) + # Headroom for the ingest call, the S3 read and cold start. + assert timeout - budget >= 120, ( + f"only {timeout - budget:.0f}s of headroom between the poll budget and " + f"the Lambda timeout; leave at least 120s for the ingest call itself" + ) + + def test_the_budget_covers_the_measured_indexing_tail(self): + """264 s was the slowest PDF in the §5.1 benchmark; dev saw 5 m 30 s. + + Pinned as a literal rather than compared to a constant, because asserting a + constant against itself proves nothing. This number is a property of + Bedrock's indexing behaviour, not a knob. + """ + from apis.app_api.kb_migration import ingestion_consumer as ic + + assert ic.INDEXED_POLL_TIMEOUT_SECONDS >= 330, ( + "the budget no longer covers the 5 m 30 s indexing time observed in dev " + "for a 1.5 MB PDF" + ) diff --git a/infrastructure/lib/constructs/managed-kb/kb-migration-construct.ts b/infrastructure/lib/constructs/managed-kb/kb-migration-construct.ts index e5841e9f..9419b550 100644 --- a/infrastructure/lib/constructs/managed-kb/kb-migration-construct.ts +++ b/infrastructure/lib/constructs/managed-kb/kb-migration-construct.ts @@ -538,6 +538,14 @@ export class KbMigrationConstruct extends Construct { }, }, }); + // NO RetryPolicy here, deliberately. It would look like the fix for slow + // indexing and would do nothing: for a Lambda target EventBridge hands the + // event off asynchronously, and from that point the function's OWN async + // retry config governs — `retryAttempts: 2` above, which is Lambda's hard + // maximum. That is the 1 + 2 attempts observed in dev before a document was + // dead-lettered. Redelivery therefore spans only a few minutes and cannot be + // extended, which is why the consumer waits for indexing WITHIN one + // invocation (INDEXED_POLL_TIMEOUT_SECONDS) rather than relying on retries. this.documentsEventRule.addTarget( new targets.LambdaFunction(this.ingestionConsumerLambda, { deadLetterQueue: this.ingestionConsumerDlq,