From 8a35dbc656a0addabbae858399235032d26a5884 Mon Sep 17 00:00:00 2001 From: derrickfink Date: Mon, 31 Aug 2026 17:16:07 -0600 Subject: [PATCH 1/2] fix(kb): the legacy pipeline stands down for a promoted knowledge base MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Routing exclusivity was only ever enforced on one side. The managed ingestion consumer returns immediately for a legacy document, but the legacy pipeline had no engine gate at all, and its `s3:ObjectCreated` notification is still live alongside the consumer's EventBridge rule. So every document added to a promoted knowledge base was handled twice, and design.md §537's "a document is indexed on exactly one backend" was untrue. The duplicate vectors are the cheap half of the problem. The expensive half is that two writers owned one `status` field and the last one won by luck. Both outcomes were observed in dev: * `DOC-b5d5d8019f44` was marked `complete` by the legacy pipeline at +30 s while the managed knowledge base — which is what actually serves this assistant — could not answer for it until +95 s. Sixty-five seconds of "your document is ready" followed by an answer that does not mention it. That is precisely the failure ingestion_consumer.py's docstring says it polls to prevent, defeated by a second writer nobody had gated. * `DOC-d637491d6cb1`, an image-only 4-year flowchart, failed Docling outright (`Docling produced zero chunks`) and was marked `failed`, while Bedrock's image extraction indexed it successfully and served it. It only ended up reading `complete` because the managed consumer happened to finish second. Reverse the finishing order — purely a function of parse time — and a good, retrievable document reads `failed` permanently, with no retry endpoint to recover it (task 14.4 is still open). So `handler.py` now resolves the engine before writing any status and returns early for `managed`, doing no parse, no embed and no status write. Keyed on the ENGINE, not on the presence of a KB record and not on `migrationState`. During `shadow` and `verify` the legacy path is still authoritative and must keep working (Requirements 16.1, 16.6); only `promote` writes `retrievalEngine`. Two tests pin that, including one for a record in `shadow`. An unreadable record resolves to legacy on purpose, because the two errors are not symmetric: wrong towards legacy costs a duplicate index while the consumer still drives the document to a correct terminal state, whereas wrong towards skipping leaves an upload un-ingested with no error and no way out but re-uploading. Same convention as resolver.load_record. Delegates to `records.resolve_engine` rather than reading the attribute inline, so "absence means legacy" keeps exactly one definition — the same call the consumer makes for the mirror-image decision. That needs `apis/shared/kb_backend/` in the rag-ingestion image, which is cheap: the package's module-level imports are stdlib only, enforced by test_kb_backend_boundary.py. Added to the Dockerfile and to build-one.sh's SOURCE_DIRS so the content hash moves with the gate. No IAM or environment change: the Lambda already carries DYNAMODB_ASSISTANTS_TABLE_NAME and already holds dynamodb:GetItem (verified against the deployed function and role in dev). The bootstrap copy at infrastructure/bootstrap-assets/rag-ingestion/handler.py needs no gate — it is a 33-line no-op placeholder that indexes nothing. Guards: 8 tests in tests/ingestion/test_ingestion_engine_gate.py. Three mutations verified caught, each by correctly-named tests and each parsed before running so a syntax error could not masquerade as a detection: - gate removed -> the 3 "promoted KB is left alone" tests - fail-open inverted -> the unreadable-record test - Dockerfile COPY dropped -> test_lambda_image_imports[rag-ingestion] Backend suite 6,825 passed; the 5 failures are the pre-existing Strands SDK contract tests documented in HANDOFF §2. --- backend/Dockerfile.rag-ingestion | 6 + .../app_api/documents/ingestion/handler.py | 84 ++++++ .../ingestion/test_ingestion_engine_gate.py | 252 ++++++++++++++++++ scripts/build/build-one.sh | 5 + 4 files changed, 347 insertions(+) create mode 100644 backend/tests/ingestion/test_ingestion_engine_gate.py diff --git a/backend/Dockerfile.rag-ingestion b/backend/Dockerfile.rag-ingestion index 73dfd9a41..960c837ae 100644 --- a/backend/Dockerfile.rag-ingestion +++ b/backend/Dockerfile.rag-ingestion @@ -168,9 +168,15 @@ COPY backend/src/apis/app_api/documents/ingestion/ ${LAMBDA_TASK_ROOT} # tag notices changes. # embeddings/ — ingestion/embeddings/bedrock_embeddings.py re-exports it # timestamps.py — ingestion/status.py imports utc_now_iso +# kb_backend/ — handler.py reads `records.resolve_engine` to skip documents +# whose knowledge base has been promoted to the managed engine. +# Cheap to carry: the package's module-level imports are stdlib +# only (enforced by tests/architecture/test_kb_backend_boundary.py), +# so this adds files, not dependencies. COPY backend/src/apis/shared/__init__.py ${LAMBDA_TASK_ROOT}/apis/shared/__init__.py COPY backend/src/apis/shared/timestamps.py ${LAMBDA_TASK_ROOT}/apis/shared/timestamps.py COPY backend/src/apis/shared/embeddings/ ${LAMBDA_TASK_ROOT}/apis/shared/embeddings/ +COPY backend/src/apis/shared/kb_backend/ ${LAMBDA_TASK_ROOT}/apis/shared/kb_backend/ RUN touch ${LAMBDA_TASK_ROOT}/apis/__init__.py CMD [ "handler.lambda_handler" ] diff --git a/backend/src/apis/app_api/documents/ingestion/handler.py b/backend/src/apis/app_api/documents/ingestion/handler.py index 418a8ab66..3cff9ae99 100644 --- a/backend/src/apis/app_api/documents/ingestion/handler.py +++ b/backend/src/apis/app_api/documents/ingestion/handler.py @@ -106,12 +106,53 @@ def _detect_mime_type(content_type: Optional[str], filename: str) -> str: pass # dotenv not installed, running in Lambda +def _resolve_engine(assistant_id: str) -> str: + """Which engine serves this knowledge base: ``managed`` or ``s3vectors``. + + Delegates to ``records.resolve_engine`` rather than reading the attribute + here, so "absence means legacy" has exactly one definition in the codebase — + the same call the managed ingestion consumer makes for the mirror-image + decision. + + **Failures resolve to legacy, deliberately.** The two ways to be wrong are + not symmetric: + + * Wrong towards legacy on a promoted knowledge base: the document is indexed + twice and the two writers race on status. Wasteful and untidy, but the + document still ends up correct, because the managed consumer is also + running and it owns the terminal state. + * Wrong towards skipping on a legacy knowledge base: nothing indexes the + document at all. It sits un-ingested with no error, and the only path back + is a re-upload. + + The second is much worse, so an unreadable record runs the legacy pipeline. + This is the same convention as ``resolver.load_record``, which treats a + failed read and an absent record identically on the grounds that "an absent + opinion is the legacy opinion". + """ + from apis.shared.kb_backend.records import ENGINE_LEGACY, get_kb_record, resolve_engine + + try: + return resolve_engine(get_kb_record(assistant_id, assistant_id)) + except Exception as exc: # noqa: BLE001 — see the docstring: legacy is the safe default + logger.warning( + f"could not resolve the retrieval engine for assistant {assistant_id}; " + f"running the legacy pipeline, which is the safe default: {exc}" + ) + return ENGINE_LEGACY + + async def async_lambda_handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]: """ Async implementation of the Lambda handler """ from status import create_status_manager + # Function-local like every other import in this module: the Lambda image + # keeps its cold start down by not importing anything at module scope that a + # given invocation might not need. + from apis.shared.kb_backend.records import ENGINE_MANAGED + # Initialize status manager status_manager = create_status_manager() @@ -121,6 +162,49 @@ async def async_lambda_handler(event: Dict[str, Any], context: Any) -> Dict[str, # 1. Parse event and extract metadata event_data = _parse_s3_event(event) + # 1a. Routing exclusivity (design §537, Requirement 10.5). A promoted + # knowledge base is served by the managed backend, and the managed + # ingestion consumer is triggered by the same S3 upload through + # EventBridge. So for those documents this pipeline must do NOTHING: + # not parse, not embed, and above all not write document status. + # + # Both halves matter, and the status half is the one that bit us: + # + # * Indexing anyway spends Docling time and S3 Vectors storage on + # vectors that retrieval will never read, because `resolve_backend` + # sends a promoted record to the managed adapter with no fallback. + # * Writing status anyway means two writers own one field and the last + # one wins by luck. Measured in dev: a PDF was marked `complete` by + # this pipeline 65 s before the managed KB could answer for it, and + # an image-only PDF that Docling could not parse at all was marked + # `failed` while the managed KB was serving it perfectly. Reverse the + # finishing order and a good document reads `failed` forever. + # + # Deliberately keyed on the *engine*, not on "has a KB record" or on + # `migrationState`: during `shadow` and `verify` the legacy path is still + # authoritative and must keep working (Requirements 16.1, 16.6). Only + # `promote` writes `retrievalEngine`, so only a promoted knowledge base + # is skipped here. + engine = _resolve_engine(event_data["assistant_id"]) + if engine == ENGINE_MANAGED: + logger.info( + f"skipping document {event_data['document_id']} for assistant " + f"{event_data['assistant_id']}: it is served by the {engine!r} " + f"engine, so the managed ingestion consumer owns this document " + f"and its status" + ) + return { + "statusCode": 200, + "body": json.dumps( + { + "message": "Skipped: knowledge base is served by the managed engine", + "assistant_id": event_data["assistant_id"], + "document_id": event_data["document_id"], + "engine": engine, + } + ), + } + # 2. Update status to 'chunking' logger.info("Starting document processing") await status_manager.mark_chunking(assistant_id=event_data["assistant_id"], document_id=event_data["document_id"]) diff --git a/backend/tests/ingestion/test_ingestion_engine_gate.py b/backend/tests/ingestion/test_ingestion_engine_gate.py new file mode 100644 index 000000000..65715f2a5 --- /dev/null +++ b/backend/tests/ingestion/test_ingestion_engine_gate.py @@ -0,0 +1,252 @@ +"""The legacy ingestion pipeline must not touch a promoted knowledge base. + +Routing exclusivity (design §537, Requirement 10.5). Both pipelines are triggered +by the same S3 upload — the legacy `s3:ObjectCreated` notification and the managed +consumer's EventBridge rule — so exactly one of them has to stand down per +document. The consumer already returns immediately for a legacy document; these +tests cover the half that was missing, which is this pipeline standing down for a +*managed* one. + +WHY THE STATUS HALF MATTERS MORE THAN THE DUPLICATE VECTORS +Two writers owning one `status` field means the last writer wins by luck. Both +outcomes were observed in dev before this gate existed: + +* A PDF was marked `complete` by this pipeline 65 s before the managed knowledge + base could answer for it — "your document is ready", then an answer that does + not mention it. +* An image-only PDF that Docling could not parse at all (`Docling produced zero + chunks`) was marked `failed` while the managed knowledge base was serving it + correctly. That one only read `complete` in the end because the managed + consumer happened to finish second and overwrite it. Reverse the finishing + order — entirely a matter of document size and parse time — and a good, + retrievable document reads `failed` forever, with no way to retry it. + +So the assertions below are mostly about what is *not* written. +""" + +from __future__ import annotations + +import json +import sys +import types +from typing import Any, Dict, List, Optional + +import pytest + +from apis.app_api.documents.ingestion import handler as handler_module + +ASSISTANT_ID = "ast-gate-1" +DOCUMENT_ID = "DOC-gate-1" + + +class RecordingStatusManager: + """Captures every status transition the handler attempts.""" + + def __init__(self) -> None: + self.calls: List[str] = [] + + async def mark_chunking(self, **_: Any) -> None: + self.calls.append("chunking") + + async def mark_embedding(self, **_: Any) -> None: + self.calls.append("embedding") + + async def mark_complete(self, **_: Any) -> None: + self.calls.append("complete") + + async def mark_failed(self, **_: Any) -> None: + self.calls.append("failed") + + +@pytest.fixture +def status_manager(monkeypatch: pytest.MonkeyPatch) -> RecordingStatusManager: + """Stand in for the `status` module, which only resolves inside the image. + + `handler.py` does `from status import create_status_manager` as a bare + top-level import because the Dockerfile flattens `documents/ingestion/` onto + LAMBDA_TASK_ROOT. Injecting the module is how a test calls the handler + without reproducing that layout. + """ + recorder = RecordingStatusManager() + fake = types.ModuleType("status") + fake.create_status_manager = lambda: recorder # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "status", fake) + return recorder + + +@pytest.fixture +def no_real_pipeline(monkeypatch: pytest.MonkeyPatch) -> List[str]: + """Replace the Docling/embedding pipeline; record whether it was reached.""" + reached: List[str] = [] + + async def _fake_pipeline(**kwargs: Any) -> None: + reached.append(kwargs.get("document_id", "?")) + + monkeypatch.setattr(handler_module, "_process_document_pipeline", _fake_pipeline) + return reached + + +def _kb_record(engine: Optional[str]) -> Dict[str, Any]: + record: Dict[str, Any] = {"PK": f"AST#{ASSISTANT_ID}", "SK": f"KB#{ASSISTANT_ID}"} + if engine is not None: + record["retrievalEngine"] = engine + return record + + +def _set_record(monkeypatch: pytest.MonkeyPatch, record) -> None: + """Point `records.get_kb_record` at a fixed answer, or make it raise.""" + from apis.shared.kb_backend import records as r + + def _get(_assistant_id: str, _app_kb_id: str): + if isinstance(record, Exception): + raise record + return record + + monkeypatch.setattr(r, "get_kb_record", _get) + + +def _event() -> Dict[str, Any]: + return { + "Records": [ + { + "s3": { + "bucket": {"name": "docs-bucket"}, + "object": { + "key": ( + f"assistants/{ASSISTANT_ID}/documents/" + f"{DOCUMENT_ID}/flowchart.pdf" + ) + }, + } + } + ] + } + + +async def _invoke() -> Dict[str, Any]: + return await handler_module.async_lambda_handler(_event(), None) + + +class TestAPromotedKnowledgeBaseIsLeftAlone: + @pytest.mark.asyncio + async def test_no_document_status_is_written_at_all( + self, status_manager: RecordingStatusManager, no_real_pipeline, monkeypatch + ): + """The whole point. `complete` and `failed` both belong to the consumer.""" + _set_record(monkeypatch, _kb_record("managed")) + + response = await _invoke() + + assert status_manager.calls == [], ( + "the legacy pipeline wrote document status for a knowledge base served " + "by the managed engine; two writers on one field is how a good " + "document ends up reading 'failed'" + ) + assert response["statusCode"] == 200 + assert "Skipped" in json.loads(response["body"])["message"] + + @pytest.mark.asyncio + async def test_the_document_is_not_parsed_or_embedded( + self, status_manager, no_real_pipeline: List[str], monkeypatch + ): + """No Docling time and no S3 Vectors storage for vectors nothing reads.""" + _set_record(monkeypatch, _kb_record("managed")) + + await _invoke() + + assert no_real_pipeline == [] + + @pytest.mark.asyncio + async def test_it_returns_success_so_the_event_is_not_retried( + self, status_manager, no_real_pipeline, monkeypatch + ): + """A skip is a correct outcome, not a failure to redeliver.""" + _set_record(monkeypatch, _kb_record("managed")) + + response = await _invoke() + + assert response["statusCode"] == 200 + assert json.loads(response["body"])["engine"] == "managed" + + +class TestEveryOtherKnowledgeBaseStillRuns: + """The gate keys on the ENGINE, not on the presence of a record. + + During `shadow` and `verify` the legacy path is still authoritative and must + keep working (Requirements 16.1, 16.6). Only `promote` writes + `retrievalEngine`, so a record that exists but names no engine — which is + exactly a migration in flight — must not be skipped. + """ + + @pytest.mark.asyncio + async def test_a_record_with_no_engine_still_runs( + self, status_manager: RecordingStatusManager, no_real_pipeline: List[str], monkeypatch + ): + _set_record(monkeypatch, _kb_record(None)) + + await _invoke() + + assert status_manager.calls == ["chunking"] + assert no_real_pipeline == [DOCUMENT_ID] + + @pytest.mark.asyncio + async def test_a_migration_in_flight_still_runs( + self, status_manager: RecordingStatusManager, no_real_pipeline: List[str], monkeypatch + ): + """`shadow` is not `promoted`. Skipping here would strand live uploads.""" + record = _kb_record(None) + record["migrationState"] = "shadow" + _set_record(monkeypatch, record) + + await _invoke() + + assert status_manager.calls == ["chunking"] + assert no_real_pipeline == [DOCUMENT_ID] + + @pytest.mark.asyncio + async def test_no_record_at_all_still_runs( + self, status_manager: RecordingStatusManager, no_real_pipeline: List[str], monkeypatch + ): + """Every knowledge base that predates this feature. Absence ⇒ legacy.""" + _set_record(monkeypatch, None) + + await _invoke() + + assert status_manager.calls == ["chunking"] + assert no_real_pipeline == [DOCUMENT_ID] + + +class TestAnUnreadableRecordRunsTheLegacyPipeline: + """Fail towards legacy, deliberately — the two errors are not symmetric. + + Wrong towards legacy on a promoted knowledge base costs a duplicate index and + a status race, and the managed consumer still drives the document to a correct + terminal state. Wrong towards skipping on a legacy knowledge base means + nothing indexes the document at all: it sits un-ingested with no error, and + the only way out is a re-upload. The second is much worse. + """ + + @pytest.mark.asyncio + async def test_a_lookup_failure_does_not_skip_the_document( + self, status_manager: RecordingStatusManager, no_real_pipeline: List[str], monkeypatch + ): + _set_record(monkeypatch, RuntimeError("DynamoDB unavailable")) + + await _invoke() + + assert status_manager.calls == ["chunking"], ( + "an unreadable KB record skipped the document; a transient read " + "failure must not silently leave an upload un-ingested" + ) + assert no_real_pipeline == [DOCUMENT_ID] + + @pytest.mark.asyncio + async def test_a_lookup_failure_is_not_reported_as_a_document_failure( + self, status_manager: RecordingStatusManager, no_real_pipeline, monkeypatch + ): + """The user's document is fine; our read of an unrelated row was not.""" + _set_record(monkeypatch, RuntimeError("DynamoDB unavailable")) + + await _invoke() + + assert "failed" not in status_manager.calls diff --git a/scripts/build/build-one.sh b/scripts/build/build-one.sh index b75a18967..090d7a6c3 100755 --- a/scripts/build/build-one.sh +++ b/scripts/build/build-one.sh @@ -91,6 +91,11 @@ case "$SERVICE" in SOURCE_DIRS=( "backend/src/apis/app_api/documents/ingestion" "backend/src/apis/shared/embeddings" + # handler.py reads records.resolve_engine to skip documents whose + # knowledge base is promoted to the managed engine. Without this + # entry a change to that gate would not move the content hash, and + # the Lambda would keep running the previous image. + "backend/src/apis/shared/kb_backend" ) # shared/__init__.py and shared/timestamps.py are single files, # hashed as manifests. The requirements.lock lives inside the From e3398f309b039244cc7ba6dbf2721e6cde079e53 Mon Sep 17 00:00:00 2001 From: derrickfink Date: Mon, 31 Aug 2026 19:48:08 -0600 Subject: [PATCH 2/2] fix(kb): propagate document deletion to the managed knowledge base MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The mirror of the ingestion gate in the previous commit. Ingestion now routes by engine; deletion did not route at all. `cleanup_service` removed the legacy S3 Vectors copy and the `DOC#` row and never touched the managed knowledge base, so on a promoted knowledge base a deleted document stayed indexed forever. The managed delete path existed — `kb_backend/tombstones.py` — but its only callers are in the reconciler, which is report-only with `reconcilerArmed` off. Three consequences, none of which raised anything: * Storage was paid for indefinitely, at $5.00/GB-month against S3 Vectors' ~$0.15 — orphans on the expensive engine. * Retrieval quietly degraded. The status filter runs AFTER retrieval, so each orphan consumed a slot in `top_k` and was then dropped: a query could return five chunks and the model see two, with nothing logged. * The only thing preventing deleted content from being served was the fail-closed status filter, which has exactly one fail-open branch (`if not doc_ids: return vectors`). That line became load-bearing in a way it was never designed to be. `cleanup_document_resources` gains a third phase, engine-gated, conjoined into `all_succeeded` so a failure blocks the hard delete. The asymmetry with the ingestion gate is deliberate. There, an unreadable KB record resolves to legacy, because being wrong costs a duplicate index while the consumer still finishes the document. Here it must FAIL: the `DOC#` row is what the status filter joins against, so reporting success on a failed managed delete would remove the row *and* leave the content — the one combination that turns a storage leak into a disclosure. `ManagedKbNotProvisioned` is the exception, treated as terminal success: nothing was ever indexed, so there is nothing to remove and no reason to retry. IAM — the part with no code to give it away ------------------------------------------- Neither the app-api task role nor the kb-sync worker could delete from a managed knowledge base, so the code above would have deployed clean and failed on first use. New `grantManagedKbDocumentDeletion`, deliberately NOT `grantDirectIngestion`: these callers only ever remove documents, so a bug in the delete path cannot add content and a bug in the ingest path cannot remove it. It includes `bedrock:StartIngestionJob` and that again looks wrong. AWS groups the whole `KnowledgeBaseDocs` family with it in one statement, and this feature has already shipped a grant naming only the matching API and failing at runtime. Granting it costs nothing here, since `DeleteKnowledgeBaseDocuments` is itself the destructive verb. kb-sync also needed the image change: its worker calls `cleanup_document_resources` after soft-deleting a document whose upstream source has vanished, so `apis/shared/kb_backend/` is added to Dockerfile.kb-sync and to build-one.sh's kb-sync SOURCE_DIRS alongside the rag-ingestion entry. Pre-existing cleanup tests now declare their engine --------------------------------------------------- Ten blocks across tests/routes/test_cleanup_service.py and tests/property/test_pbt_cleanup_service.py patched `records.get_kb_record -> None`. They describe a legacy knowledge base — the state of every assistant predating this feature — and previously said nothing about the engine because there was no third phase to say it to. Left alone they reached a real DynamoDB read. Guards: 7 tests in tests/documents/test_managed_delete_propagation.py, plus grant-shape tests in managed-kb.test.ts and wiring assertions on the real roles in kb-sync.test.ts and platform-stack.test.ts — the latter because a grant proven on a fake role says nothing about the identity that runs the code, which is how two earlier defects on this feature shipped. Four mutations verified caught, each by correctly-named tests: - managed phase dropped from the conjunction -> test_a_failed_managed_delete_blocks_the_hard_delete - every KB treated as legacy -> 3 tests incl. test_the_document_is_deleted_from_the_managed_kb - unreadable record reports success -> test_an_unreadable_record_fails_rather_than_assuming_legacy - both grant attachments removed -> KbSyncConstruct + PlatformStack wiring tests Backend 6,832 passed; infra 640 passed. The 5 backend failures are the pre-existing Strands SDK contract tests documented in HANDOFF §2. --- backend/Dockerfile.kb-sync | 6 + .../documents/services/cleanup_service.py | 104 +++++++- .../test_managed_delete_propagation.py | 222 ++++++++++++++++++ .../property/test_pbt_cleanup_service.py | 27 +++ backend/tests/routes/test_cleanup_service.py | 63 +++++ .../constructs/app-api/app-api-iam-grants.ts | 18 +- .../constructs/kb-sync/kb-sync-construct.ts | 10 + .../managed-kb/managed-kb-role-construct.ts | 38 +++ infrastructure/test/kb-sync.test.ts | 31 +++ infrastructure/test/managed-kb.test.ts | 31 +++ infrastructure/test/platform-stack.test.ts | 37 +++ scripts/build/build-one.sh | 3 + 12 files changed, 588 insertions(+), 2 deletions(-) create mode 100644 backend/tests/documents/test_managed_delete_propagation.py diff --git a/backend/Dockerfile.kb-sync b/backend/Dockerfile.kb-sync index 046351a2d..3bfb4549d 100644 --- a/backend/Dockerfile.kb-sync +++ b/backend/Dockerfile.kb-sync @@ -37,6 +37,12 @@ COPY backend/src/apis/shared/dynamo_errors.py ${LAMBDA_TASK_ROOT}/apis/shared/dy COPY backend/src/apis/shared/sync_policies/ ${LAMBDA_TASK_ROOT}/apis/shared/sync_policies/ COPY backend/src/apis/shared/oauth/ ${LAMBDA_TASK_ROOT}/apis/shared/oauth/ COPY backend/src/apis/shared/embeddings/ ${LAMBDA_TASK_ROOT}/apis/shared/embeddings/ +# kb_backend/ — the worker reaches documents/services/cleanup_service.py, +# which reads records.resolve_engine and calls ManagedKbBackend to remove a +# deleted document from a promoted knowledge base. Module-level imports in +# that package are stdlib only (test_kb_backend_boundary.py), so this adds +# files, not dependencies. +COPY backend/src/apis/shared/kb_backend/ ${LAMBDA_TASK_ROOT}/apis/shared/kb_backend/ # assistants/ — document_service verifies assistant ownership via # get_assistant before soft-deleting; the worker's miss-eviction path # (_delete_missing_web_document) goes through it. diff --git a/backend/src/apis/app_api/documents/services/cleanup_service.py b/backend/src/apis/app_api/documents/services/cleanup_service.py index 047b62755..1cb085520 100644 --- a/backend/src/apis/app_api/documents/services/cleanup_service.py +++ b/backend/src/apis/app_api/documents/services/cleanup_service.py @@ -78,6 +78,16 @@ async def cleanup_document_resources( logger.error(f"Unexpected error in vector deletion for {document_id}: {e}", exc_info=True) vectors_deleted = False + try: + managed_deleted = await _delete_managed_documents_with_retries( + assistant_id, document_id, max_retries, base_delay + ) + except Exception as e: + logger.error( + f"Unexpected error in managed-KB deletion for {document_id}: {e}", exc_info=True + ) + managed_deleted = False + try: s3_deleted = await _delete_s3_with_retries( s3_key, max_retries, base_delay @@ -86,7 +96,7 @@ async def cleanup_document_resources( logger.error(f"Unexpected error in S3 deletion for {document_id}: {e}", exc_info=True) s3_deleted = False - all_succeeded = vectors_deleted and s3_deleted + all_succeeded = vectors_deleted and managed_deleted and s3_deleted if all_succeeded: try: @@ -239,6 +249,98 @@ async def _delete_vectors_with_retries( return False +async def _delete_managed_documents_with_retries( + assistant_id: str, + document_id: str, + max_retries: int, + base_delay: float, +) -> bool: + """Remove the document from the managed knowledge base, if there is one. + + The mirror of the legacy vector deletion above, and the reason it exists: + before this, deletion removed the S3 Vectors copy and the ``DOC#`` row but + never touched the managed knowledge base. On a promoted knowledge base the + content therefore stayed indexed forever. The corpus could only grow, at + $5.00/GB-month, and each orphan kept consuming a slot in every ``top_k`` + before the status filter dropped it — so answers quietly got thinner while + nothing errored. + + Returns ``True`` when there is nothing to do + ------------------------------------------- + A knowledge base that is not promoted has no managed copy, so "no managed + copy to delete" is success, not failure. Returning ``False`` there would + block ``hard_delete_document`` for every legacy document in the system. + + Why a failure here must block the hard delete + --------------------------------------------- + The caller only hard-deletes the ``DOC#`` row when every phase succeeds, and + that row is what the fail-closed status filter joins against. Leaving it in + place is what keeps a still-indexed managed chunk from being served while + this deletion is retried. Reporting success on a failed managed delete would + remove the row *and* leave the content — the one combination that turns a + storage leak into a disclosure. + + Deletion is idempotent + ---------------------- + ``DeleteKnowledgeBaseDocuments`` on an absent document is not an error, so a + retry after a partial failure is safe and needs no bookkeeping. + """ + from apis.shared.kb_backend.records import ENGINE_MANAGED, get_kb_record, resolve_engine + + try: + # App_KB_Id == assistant_id in this phase, so one value serves both. + engine = resolve_engine(get_kb_record(assistant_id, assistant_id)) + except Exception as e: + # Unlike the ingestion gate — where an unreadable record resolves to + # legacy so the upload still gets indexed — an unreadable record HERE + # must fail. If this knowledge base is in fact promoted, reporting + # success would hard-delete the row that keeps its chunks unserved. + logger.error( + f"could not resolve the retrieval engine for assistant {assistant_id} " + f"while deleting {document_id}; treating the managed deletion as failed " + f"so the document record survives for a retry: {e}" + ) + return False + + if engine != ENGINE_MANAGED: + return True + + from apis.shared.kb_backend.managed_backend import ManagedKbBackend, ManagedKbNotProvisioned + + backend = ManagedKbBackend() + for attempt in range(max_retries): + try: + await backend.delete_document(assistant_id, document_id) + logger.info( + f"deleted document {document_id} from the managed knowledge base " + f"for assistant {assistant_id}" + ) + return True + except ManagedKbNotProvisioned: + # Promoted but carrying no AWS identifiers is not a state a real + # migration produces — `promote` runs after provisioning. Nothing was + # ever indexed, so there is nothing to remove and no reason to retry. + logger.warning( + f"assistant {assistant_id} names the managed engine but has no AWS " + f"identifiers; nothing to delete for {document_id}" + ) + return True + except Exception as e: + delay = base_delay * (2 ** attempt) + random.uniform(0, 0.1) + logger.warning( + f"Managed-KB deletion attempt {attempt + 1}/{max_retries} failed for " + f"{document_id}: {e}, retrying in {delay:.2f}s" + ) + if attempt < max_retries - 1: + await asyncio.sleep(delay) + + logger.error( + f"Managed-KB deletion failed after {max_retries} attempts for {document_id}; " + f"the document record is being kept so the status filter continues to hide it" + ) + return False + + async def _delete_s3_with_retries( s3_key: str, max_retries: int, diff --git a/backend/tests/documents/test_managed_delete_propagation.py b/backend/tests/documents/test_managed_delete_propagation.py new file mode 100644 index 000000000..8adaacb6a --- /dev/null +++ b/backend/tests/documents/test_managed_delete_propagation.py @@ -0,0 +1,222 @@ +"""Deleting a document must remove it from the engine that actually serves it. + +Before this, `cleanup_service` deleted the legacy S3 Vectors copy and the `DOC#` +row and never touched the managed knowledge base. On a promoted knowledge base +the content therefore stayed indexed forever: + +* it kept being paid for, at $5.00/GB-month against S3 Vectors' ~$0.15; +* every orphan kept consuming a slot in `top_k`, because the status filter runs + *after* retrieval — so a query could return five chunks and the model see two, + with nothing logged and nothing raised; +* and the only thing preventing deleted content from being served was the + fail-closed status filter, which has exactly one fail-open branch. + +The most load-bearing assertion in this file is +`test_a_failed_managed_delete_keeps_the_document_record`: the `DOC#` row is what +the status filter joins against, so reporting success on a failed managed delete +would remove the row *and* leave the content — turning a storage leak into a +disclosure. +""" + +from __future__ import annotations + +from typing import Any, Dict, List, Optional + +import pytest + +from apis.app_api.documents.services import cleanup_service + +ASSISTANT_ID = "ast-del-1" +DOCUMENT_ID = "DOC-del-1" + + +class RecordingManagedBackend: + """Stands in for ManagedKbBackend, recording deletes and optionally failing.""" + + def __init__(self, raises: Optional[BaseException] = None) -> None: + self.deleted: List[str] = [] + self.attempts = 0 + self._raises = raises + + async def delete_document(self, kb_ref: str, document_id: str) -> None: + self.attempts += 1 + if self._raises is not None: + raise self._raises + self.deleted.append(document_id) + + +@pytest.fixture +def managed(monkeypatch: pytest.MonkeyPatch) -> RecordingManagedBackend: + backend = RecordingManagedBackend() + import apis.shared.kb_backend.managed_backend as mb + + monkeypatch.setattr(mb, "ManagedKbBackend", lambda *a, **k: backend) + return backend + + +def _set_engine(monkeypatch: pytest.MonkeyPatch, engine: Optional[str]) -> None: + from apis.shared.kb_backend import records as r + + record: Optional[Dict[str, Any]] + if engine is None: + record = None + else: + record = {"retrievalEngine": engine} if engine != "absent" else {} + + monkeypatch.setattr(r, "get_kb_record", lambda *_: record) + + +def _raise_on_lookup(monkeypatch: pytest.MonkeyPatch, exc: BaseException) -> None: + from apis.shared.kb_backend import records as r + + def _boom(*_: Any): + raise exc + + monkeypatch.setattr(r, "get_kb_record", _boom) + + +async def _delete() -> bool: + return await cleanup_service._delete_managed_documents_with_retries( + ASSISTANT_ID, DOCUMENT_ID, max_retries=3, base_delay=0.0 + ) + + +class TestAPromotedKnowledgeBaseHasTheDocumentRemoved: + @pytest.mark.asyncio + async def test_the_document_is_deleted_from_the_managed_kb( + self, managed: RecordingManagedBackend, monkeypatch + ): + _set_engine(monkeypatch, "managed") + + assert await _delete() is True + assert managed.deleted == [DOCUMENT_ID], ( + "a document deleted by its owner was left in the managed knowledge " + "base; the corpus can only grow and every orphan costs a top_k slot" + ) + + +class TestALegacyKnowledgeBaseIsUntouched: + """No managed copy to remove is success, not failure. + + Returning False here would block `hard_delete_document` for every legacy + document in the system — the deletion would never complete. + """ + + @pytest.mark.asyncio + async def test_an_absent_record_needs_no_managed_delete( + self, managed: RecordingManagedBackend, monkeypatch + ): + _set_engine(monkeypatch, None) + + assert await _delete() is True + assert managed.deleted == [] + + @pytest.mark.asyncio + async def test_a_record_with_no_engine_needs_no_managed_delete( + self, managed: RecordingManagedBackend, monkeypatch + ): + """A migration in flight: `shadow`/`verify` have not promoted anything.""" + _set_engine(monkeypatch, "absent") + + assert await _delete() is True + assert managed.deleted == [] + + +class TestFailuresKeepTheDocumentRecordAlive: + """The `DOC#` row is the safety net, so a failure must not report success.""" + + @pytest.mark.asyncio + async def test_a_failed_managed_delete_keeps_the_document_record( + self, monkeypatch + ): + backend = RecordingManagedBackend(raises=RuntimeError("AccessDenied")) + import apis.shared.kb_backend.managed_backend as mb + + monkeypatch.setattr(mb, "ManagedKbBackend", lambda *a, **k: backend) + _set_engine(monkeypatch, "managed") + + assert await _delete() is False, ( + "a failed managed delete reported success; the caller would then " + "hard-delete the DOC# row that the fail-closed status filter joins " + "against, leaving indexed content with nothing left to hide it" + ) + assert backend.attempts == 3, "every attempt should be retried" + + @pytest.mark.asyncio + async def test_an_unreadable_record_fails_rather_than_assuming_legacy( + self, managed: RecordingManagedBackend, monkeypatch + ): + """The opposite choice from the ingestion gate, on purpose. + + On ingest an unreadable record resolves to legacy, because the cost of + being wrong is a duplicate index while the consumer still finishes the + document. On delete the cost of being wrong is hard-deleting the row that + keeps a still-indexed chunk unserved, so it fails and retries instead. + """ + _raise_on_lookup(monkeypatch, RuntimeError("DynamoDB unavailable")) + + assert await _delete() is False + assert managed.deleted == [] + + @pytest.mark.asyncio + async def test_a_promoted_kb_with_no_aws_ids_is_not_retried_forever( + self, monkeypatch + ): + """Nothing was ever indexed, so there is nothing to remove.""" + from apis.shared.kb_backend.managed_backend import ManagedKbNotProvisioned + + backend = RecordingManagedBackend(raises=ManagedKbNotProvisioned("no awsKbId")) + import apis.shared.kb_backend.managed_backend as mb + + monkeypatch.setattr(mb, "ManagedKbBackend", lambda *a, **k: backend) + _set_engine(monkeypatch, "managed") + + assert await _delete() is True + assert backend.attempts == 1, "not provisioned is terminal, not transient" + + +class TestTheCleanupContractIncludesManagedDeletion: + @pytest.mark.asyncio + async def test_a_failed_managed_delete_blocks_the_hard_delete( + self, monkeypatch + ): + """End to end through `cleanup_document_resources`, not just the helper. + + The helper returning False is only useful if the caller conjoins it. This + is the assertion that would fail if someone computed `all_succeeded` + without the managed phase. + """ + hard_deleted: List[str] = [] + + async def _no_hard_delete(assistant_id: str, document_id: str) -> None: + hard_deleted.append(document_id) + + async def _ok(*_: Any, **__: Any) -> bool: + return True + + monkeypatch.setattr(cleanup_service, "_delete_vectors_with_retries", _ok) + monkeypatch.setattr(cleanup_service, "_delete_s3_with_retries", _ok) + + async def _managed_fails(*_: Any, **__: Any) -> bool: + return False + + monkeypatch.setattr( + cleanup_service, "_delete_managed_documents_with_retries", _managed_fails + ) + + import apis.app_api.documents.services.document_service as ds + + monkeypatch.setattr(ds, "hard_delete_document", _no_hard_delete) + + result = await cleanup_service.cleanup_document_resources( + document_id=DOCUMENT_ID, + assistant_id=ASSISTANT_ID, + s3_key=f"assistants/{ASSISTANT_ID}/documents/{DOCUMENT_ID}/f.pdf", + chunk_count=3, + ) + + assert result is False + assert hard_deleted == [], ( + "the DOC# row was hard-deleted even though the managed knowledge " + "base still holds the content" + ) diff --git a/backend/tests/property/test_pbt_cleanup_service.py b/backend/tests/property/test_pbt_cleanup_service.py index cca1f099b..92d429d81 100644 --- a/backend/tests/property/test_pbt_cleanup_service.py +++ b/backend/tests/property/test_pbt_cleanup_service.py @@ -104,6 +104,15 @@ def failing_s3_delete(**kwargs): side_effect=failing_fallback, ), patch("boto3.client", return_value=mock_s3_client), + # These cases are all about a LEGACY knowledge base — the state of + # every assistant that predates the managed-KB migration. Absence of + # a KB record is what "legacy" looks like, so the managed deletion + # phase resolves to a no-op. Stated explicitly rather than left to a + # real DynamoDB read. + patch( + "apis.shared.kb_backend.records.get_kb_record", + return_value=None, + ), patch( "apis.app_api.documents.services.cleanup_service.asyncio.sleep", side_effect=mock_sleep, @@ -208,6 +217,15 @@ async def failing_fallback(*args, **kwargs): side_effect=failing_fallback, ), patch("boto3.client", return_value=mock_s3_client), + # These cases are all about a LEGACY knowledge base — the state of + # every assistant that predates the managed-KB migration. Absence of + # a KB record is what "legacy" looks like, so the managed deletion + # phase resolves to a no-op. Stated explicitly rather than left to a + # real DynamoDB read. + patch( + "apis.shared.kb_backend.records.get_kb_record", + return_value=None, + ), patch( "apis.app_api.documents.services.cleanup_service.asyncio.sleep", new_callable=AsyncMock, @@ -286,6 +304,15 @@ async def succeeding_fallback(*args, **kwargs): side_effect=succeeding_fallback, ), patch("boto3.client", return_value=mock_s3_client), + # These cases are all about a LEGACY knowledge base — the state of + # every assistant that predates the managed-KB migration. Absence of + # a KB record is what "legacy" looks like, so the managed deletion + # phase resolves to a no-op. Stated explicitly rather than left to a + # real DynamoDB read. + patch( + "apis.shared.kb_backend.records.get_kb_record", + return_value=None, + ), patch( "apis.app_api.documents.services.document_service.hard_delete_document", mock_hard_delete, diff --git a/backend/tests/routes/test_cleanup_service.py b/backend/tests/routes/test_cleanup_service.py index b4c78fa8c..9b392f2c3 100644 --- a/backend/tests/routes/test_cleanup_service.py +++ b/backend/tests/routes/test_cleanup_service.py @@ -45,6 +45,15 @@ async def test_cleanup_returns_true_when_both_succeed(self): new_callable=AsyncMock, ), patch("boto3.client", return_value=mock_s3_client), + # These cases are all about a LEGACY knowledge base — the state of + # every assistant that predates the managed-KB migration. Absence of + # a KB record is what "legacy" looks like, so the managed deletion + # phase resolves to a no-op. Stated explicitly rather than left to a + # real DynamoDB read. + patch( + "apis.shared.kb_backend.records.get_kb_record", + return_value=None, + ), patch( "apis.app_api.documents.services.document_service.hard_delete_document", mock_hard_delete, @@ -83,6 +92,15 @@ async def test_cleanup_returns_false_when_vectors_fail(self): side_effect=Exception("vector fallback error"), ), patch("boto3.client", return_value=mock_s3_client), + # These cases are all about a LEGACY knowledge base — the state of + # every assistant that predates the managed-KB migration. Absence of + # a KB record is what "legacy" looks like, so the managed deletion + # phase resolves to a no-op. Stated explicitly rather than left to a + # real DynamoDB read. + patch( + "apis.shared.kb_backend.records.get_kb_record", + return_value=None, + ), patch( "apis.app_api.documents.services.cleanup_service.asyncio.sleep", new_callable=AsyncMock, @@ -123,6 +141,15 @@ async def test_cleanup_returns_false_when_s3_fails(self): new_callable=AsyncMock, ), patch("boto3.client", return_value=mock_s3_client), + # These cases are all about a LEGACY knowledge base — the state of + # every assistant that predates the managed-KB migration. Absence of + # a KB record is what "legacy" looks like, so the managed deletion + # phase resolves to a no-op. Stated explicitly rather than left to a + # real DynamoDB read. + patch( + "apis.shared.kb_backend.records.get_kb_record", + return_value=None, + ), patch( "apis.app_api.documents.services.cleanup_service.asyncio.sleep", new_callable=AsyncMock, @@ -165,6 +192,15 @@ async def test_cleanup_independent_phases(self): side_effect=Exception("vector fallback error"), ), patch("boto3.client", return_value=mock_s3_client), + # These cases are all about a LEGACY knowledge base — the state of + # every assistant that predates the managed-KB migration. Absence of + # a KB record is what "legacy" looks like, so the managed deletion + # phase resolves to a no-op. Stated explicitly rather than left to a + # real DynamoDB read. + patch( + "apis.shared.kb_backend.records.get_kb_record", + return_value=None, + ), patch( "apis.app_api.documents.services.cleanup_service.asyncio.sleep", new_callable=AsyncMock, @@ -214,6 +250,15 @@ async def fail_twice_then_succeed(*args, **kwargs): new_callable=AsyncMock, ), patch("boto3.client", return_value=mock_s3_client), + # These cases are all about a LEGACY knowledge base — the state of + # every assistant that predates the managed-KB migration. Absence of + # a KB record is what "legacy" looks like, so the managed deletion + # phase resolves to a no-op. Stated explicitly rather than left to a + # real DynamoDB read. + patch( + "apis.shared.kb_backend.records.get_kb_record", + return_value=None, + ), patch( "apis.app_api.documents.services.cleanup_service.asyncio.sleep", new_callable=AsyncMock, @@ -256,6 +301,15 @@ async def test_cleanup_calls_hard_delete_on_success(self): new_callable=AsyncMock, ), patch("boto3.client", return_value=mock_s3_client), + # These cases are all about a LEGACY knowledge base — the state of + # every assistant that predates the managed-KB migration. Absence of + # a KB record is what "legacy" looks like, so the managed deletion + # phase resolves to a no-op. Stated explicitly rather than left to a + # real DynamoDB read. + patch( + "apis.shared.kb_backend.records.get_kb_record", + return_value=None, + ), patch( "apis.app_api.documents.services.document_service.hard_delete_document", mock_hard_delete, @@ -294,6 +348,15 @@ async def test_cleanup_does_not_call_hard_delete_on_failure(self): side_effect=Exception("vector fallback error"), ), patch("boto3.client", return_value=mock_s3_client), + # These cases are all about a LEGACY knowledge base — the state of + # every assistant that predates the managed-KB migration. Absence of + # a KB record is what "legacy" looks like, so the managed deletion + # phase resolves to a no-op. Stated explicitly rather than left to a + # real DynamoDB read. + patch( + "apis.shared.kb_backend.records.get_kb_record", + return_value=None, + ), patch( "apis.app_api.documents.services.cleanup_service.asyncio.sleep", new_callable=AsyncMock, diff --git a/infrastructure/lib/constructs/app-api/app-api-iam-grants.ts b/infrastructure/lib/constructs/app-api/app-api-iam-grants.ts index a9bd4690c..f5ba769f0 100644 --- a/infrastructure/lib/constructs/app-api/app-api-iam-grants.ts +++ b/infrastructure/lib/constructs/app-api/app-api-iam-grants.ts @@ -23,7 +23,10 @@ import * as cdk from 'aws-cdk-lib'; import { Construct } from 'constructs'; import { AppConfig } from '../../config'; import { PlatformComputeRefs } from '../platform-compute-refs'; -import { grantManagedKbRetrieval } from '../managed-kb/managed-kb-role-construct'; +import { + grantManagedKbDocumentDeletion, + grantManagedKbRetrieval, +} from '../managed-kb/managed-kb-role-construct'; export interface AppApiIamGrantsProps { scope: Construct; @@ -657,6 +660,19 @@ export function grantAppApiPermissions(props: AppApiIamGrantsProps): void { // Lambdas alone. grantManagedKbRetrieval(config, taskRole); + // ── Managed knowledge bases (document deletion) ── + // `DELETE /assistants/{id}/documents/{doc}` reaches + // `cleanup_service._delete_managed_documents_with_retries`, which removes + // the document from the managed knowledge base when that assistant has + // been promoted. Without this grant the delete fails, the document row is + // deliberately kept so the fail-closed status filter keeps hiding the + // chunks, and the managed corpus grows forever at $5.00/GB-month. + // + // Deletion only — NOT `grantManagedKbDirectIngestion`. The App API must + // not be able to write a corpus; only the migration worker and the + // ingestion consumer do that. + grantManagedKbDocumentDeletion(config, taskRole); + // ── AgentCore WorkloadIdentity (OAuth vault token minting) ── // Grants the App API the data-plane actions used by /connectors/* // routes and shared/oauth/agentcore_identity.py: diff --git a/infrastructure/lib/constructs/kb-sync/kb-sync-construct.ts b/infrastructure/lib/constructs/kb-sync/kb-sync-construct.ts index b50a8c982..343e965f3 100644 --- a/infrastructure/lib/constructs/kb-sync/kb-sync-construct.ts +++ b/infrastructure/lib/constructs/kb-sync/kb-sync-construct.ts @@ -12,6 +12,7 @@ import * as path from 'path'; import { Construct } from 'constructs'; import { AppConfig } from '../../config'; +import { grantManagedKbDocumentDeletion } from '../managed-kb/managed-kb-role-construct'; export interface KbSyncConstructProps { config: AppConfig; @@ -144,6 +145,15 @@ export class KbSyncConstruct extends Construct { oauthProvidersTable.grantReadData(this.workerLambda); documentsBucket.grantPut(this.workerLambda); + // Worker: remove a vanished document from a promoted knowledge base. + // `kb_sync/worker.py` soft-deletes a document whose upstream source is + // gone and then calls `cleanup_service.cleanup_document_resources`, which + // deletes from the managed knowledge base when the assistant has been + // promoted. Same grant and same reasoning as the app-api task role: the + // sync worker removes documents, it never writes a corpus, so this is the + // deletion grant and not `grantDirectIngestion`. + grantManagedKbDocumentDeletion(config, this.workerLambda.role!); + // Worker: retrieve the policy creator's stored 3LO token from the // AgentCore Identity vault with no live user session. Mirrors app-api's // AgentCoreWorkloadIdentityAccess statement (app-api-iam-grants.ts) diff --git a/infrastructure/lib/constructs/managed-kb/managed-kb-role-construct.ts b/infrastructure/lib/constructs/managed-kb/managed-kb-role-construct.ts index 393e811c2..d7ddf05e7 100644 --- a/infrastructure/lib/constructs/managed-kb/managed-kb-role-construct.ts +++ b/infrastructure/lib/constructs/managed-kb/managed-kb-role-construct.ts @@ -218,6 +218,44 @@ export function grantManagedKbDirectIngestion(config: AppConfig, role: iam.IRole role.addToPrincipalPolicy(putMetricDataStatement(config, 'ManagedKbIngestMetrics')); } +/** + * Document-deletion grant: remove one document's content from a managed + * knowledge base. Held by the identities that service a user deleting a + * document — the app-api task role and the kb-sync worker — neither of + * which has any business ingesting. + * + * Separate from `grantManagedKbDirectIngestion` on purpose. That grant + * belongs to the migration worker and the ingestion consumer, which write + * corpora; these callers only ever remove. Splitting them means a bug in + * the delete path cannot add content and a bug in the ingest path cannot + * remove it. + * + * `bedrock:StartIngestionJob` IS REQUIRED HERE TOO, and again it looks + * wrong. AWS lists the whole `KnowledgeBaseDocs` family — ingest, get, + * list and delete — in one statement alongside `StartIngestionJob` in its + * direct-ingestion prerequisites, and we have already shipped one grant + * that named only the matching API and failed on first real use + * (`IngestKnowledgeBaseDocuments` authorized as `StartIngestionJob`). + * Rather than rediscover whether delete is authorized the same way from a + * production AccessDeniedException, it is granted. The cost of including + * it is nothing: it confers no ability this caller does not already need, + * because `DeleteKnowledgeBaseDocuments` is itself the destructive verb. + */ +export function grantManagedKbDocumentDeletion(config: AppConfig, role: iam.IRole): void { + role.addToPrincipalPolicy(new iam.PolicyStatement({ + sid: 'ManagedKbDocumentDeletion', + effect: iam.Effect.ALLOW, + actions: [ + 'bedrock:DeleteKnowledgeBaseDocuments', + // See the docblock: the IAM action AWS actually checks for the + // document-plane operations, not an invocation of the 0.1 RPS + // ingestion-job API that Requirement 9.2 forbids calling. + 'bedrock:StartIngestionJob', + ], + resources: [knowledgeBaseArnWildcard(config)], + })); +} + /** * Sharing grant: administer a knowledge base's *resource* policy * (Requirement 25.6). Separate from every other grant because it is the diff --git a/infrastructure/test/kb-sync.test.ts b/infrastructure/test/kb-sync.test.ts index b2f5f3222..437ef37bb 100644 --- a/infrastructure/test/kb-sync.test.ts +++ b/infrastructure/test/kb-sync.test.ts @@ -86,6 +86,37 @@ describe('KbSyncConstruct', () => { }); }); + it('worker may delete documents from a promoted managed knowledge base', () => { + // `kb_sync/worker.py` soft-deletes a document whose upstream source has + // vanished, then calls `cleanup_service.cleanup_document_resources`, which + // removes it from the managed knowledge base when the assistant has been + // promoted. Without the grant that delete fails, the DOC# row is kept on + // purpose so the fail-closed status filter keeps hiding the chunks, and the + // managed corpus grows forever at $5.00/GB-month. + // + // This assertion exists because the same class of gap has shipped twice on + // this feature: code that reads correctly with no IAM behind it. + t.hasResourceProperties('AWS::IAM::Policy', { + PolicyDocument: { + Statement: Match.arrayWith([ + Match.objectLike({ + Sid: 'ManagedKbDocumentDeletion', + Action: Match.arrayWith(['bedrock:DeleteKnowledgeBaseDocuments']), + }), + ]), + }, + }); + }); + + it('worker may not ingest into a managed knowledge base', () => { + // The sync worker removes documents; writing a corpus belongs to the + // migration worker and the ingestion consumer alone. + const policies = t.findResources('AWS::IAM::Policy'); + const json = JSON.stringify(policies); + expect(json).not.toContain('IngestKnowledgeBaseDocuments'); + expect(json).not.toContain('CreateKnowledgeBase'); + }); + it('custom metrics are namespace-conditioned', () => { t.hasResourceProperties('AWS::IAM::Policy', { PolicyDocument: { diff --git a/infrastructure/test/managed-kb.test.ts b/infrastructure/test/managed-kb.test.ts index 48d0de67b..c810b280d 100644 --- a/infrastructure/test/managed-kb.test.ts +++ b/infrastructure/test/managed-kb.test.ts @@ -12,6 +12,7 @@ import * as iam from 'aws-cdk-lib/aws-iam'; import { Match, Template } from 'aws-cdk-lib/assertions'; import { + grantManagedKbDocumentDeletion, grantManagedKbRetrieval, ManagedKbRoleConstruct, } from '../lib/constructs/managed-kb/managed-kb-role-construct'; @@ -66,6 +67,7 @@ function synthConstruct(): Template { const ingestor = new iam.Role(stack, 'FakeIngestor', { assumedBy: lambdaPrincipal }); const retriever = new iam.Role(stack, 'FakeRetriever', { assumedBy: lambdaPrincipal }); const sharer = new iam.Role(stack, 'FakeSharer', { assumedBy: lambdaPrincipal }); + const deleter = new iam.Role(stack, 'FakeDeleter', { assumedBy: lambdaPrincipal }); // Exercise the public methods (the surface task 2.1 calls) for // provisioning/ingestion and the free function for retrieval (the @@ -75,6 +77,7 @@ function synthConstruct(): Template { managedKb.grantResourcePolicyAdmin(sharer); managedKb.grantMetricsRead(provisioner); grantManagedKbRetrieval(config, retriever); + grantManagedKbDocumentDeletion(config, deleter); return Template.fromStack(stack); } @@ -405,6 +408,34 @@ describe('ManagedKbRoleConstruct — caller grants', () => { } }); + it('scopes document deletion to its own statement, separate from ingestion', () => { + // The app-api task role and the kb-sync worker delete documents from a + // promoted knowledge base; neither may write a corpus. Splitting the grants + // means a bug in the delete path cannot add content and a bug in the ingest + // path cannot remove it. + const s = statementBySid(t, 'ManagedKbDocumentDeletion'); + expect(s.Action).toEqual([ + 'bedrock:DeleteKnowledgeBaseDocuments', + 'bedrock:StartIngestionJob', + ]); + expect(s.Resource).toBe(KB_ARN_WILDCARD); + }); + + it('keeps the deleting caller away from ingestion and from knowledge-base CRUD', () => { + // Deletion of a *document* is not deletion of a knowledge base, and it is + // certainly not permission to write one. + const holders = policiesWithSid(t, 'ManagedKbDocumentDeletion'); + expect(holders).toHaveLength(1); + for (const forbidden of [ + 'IngestKnowledgeBaseDocuments', + 'CreateKnowledgeBase', + 'DeleteKnowledgeBase"', + 'iam:PassRole', + ]) { + expect(holders[0].json).not.toContain(forbidden); + } + }); + it('grants PutMetricData on the same non-reserved namespace to every calling identity', () => { for (const sid of [ 'ManagedKbProvisionMetrics', diff --git a/infrastructure/test/platform-stack.test.ts b/infrastructure/test/platform-stack.test.ts index 4cb43f8ed..adb5caa5d 100644 --- a/infrastructure/test/platform-stack.test.ts +++ b/infrastructure/test/platform-stack.test.ts @@ -355,4 +355,41 @@ describe('PlatformStack', () => { }); }); }); + describe('Managed knowledge base grants on the real compute roles', () => { + // The recurring failure on this feature is code that reads correctly with + // no IAM behind it — a grant on a fake role in a construct test proves the + // statement is well-formed, not that the identity which runs the code ever + // receives it. These assert the wiring, on the synthesized stack. + it('the app-api task role may delete documents from a managed KB', () => { + // `DELETE /assistants/{id}/documents/{doc}` reaches + // `cleanup_service._delete_managed_documents_with_retries`. Without this + // the delete fails, the DOC# row is kept so the fail-closed status filter + // keeps hiding the chunks, and the managed corpus grows forever. + const policies = { + ...template.findResources('AWS::IAM::Policy'), + ...template.findResources('AWS::IAM::ManagedPolicy'), + }; + const found = Object.values(policies).some((r) => { + const statements = + (r.Properties as { PolicyDocument?: { Statement?: Array<{ Sid?: string }> } }) + .PolicyDocument?.Statement ?? []; + return statements.some((st) => st.Sid === 'ManagedKbDocumentDeletion'); + }); + expect(found).toBe(true); + }); + + it('the app-api task role may retrieve from a managed KB', () => { + const policies = { + ...template.findResources('AWS::IAM::Policy'), + ...template.findResources('AWS::IAM::ManagedPolicy'), + }; + const found = Object.values(policies).some((r) => { + const statements = + (r.Properties as { PolicyDocument?: { Statement?: Array<{ Sid?: string }> } }) + .PolicyDocument?.Statement ?? []; + return statements.some((st) => st.Sid === 'ManagedKbRetrieve'); + }); + expect(found).toBe(true); + }); + }); }); diff --git a/scripts/build/build-one.sh b/scripts/build/build-one.sh index 090d7a6c3..0c7c19db7 100755 --- a/scripts/build/build-one.sh +++ b/scripts/build/build-one.sh @@ -121,6 +121,9 @@ case "$SERVICE" in # would ship stale code under an unchanged content-hash tag. SOURCE_DIRS=( "backend/src/apis/app_api/kb_sync" + # cleanup_service.py reads records.resolve_engine and calls + # ManagedKbBackend to delete from a promoted knowledge base. + "backend/src/apis/shared/kb_backend" "backend/src/apis/app_api/file_sources" "backend/src/apis/app_api/documents" "backend/src/apis/app_api/web_sources"