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/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/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/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/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 b75a18967..0c7c19db7 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 @@ -116,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"