diff --git a/.kiro/specs/managed-kb-migration/HANDOFF.md b/.kiro/specs/managed-kb-migration/HANDOFF.md index ca16cfbc..7a583b04 100644 --- a/.kiro/specs/managed-kb-migration/HANDOFF.md +++ b/.kiro/specs/managed-kb-migration/HANDOFF.md @@ -1,9 +1,28 @@ # Managed KB Migration — Handoff -**Last updated:** 2026-08-26 (groups 11–13, group 14 backend half, tag contract, **14.3 upgrade UX + enrolment surface**) · **Branch:** `feature/kb-migration` · **Nothing deployed** +**Last updated:** 2026-08-31 · **Shipped to production in 1.16.0, inert behind flags** · +**A migration has now completed end to end in dev** -Working state for this feature so a fresh session can pick it up without re-deriving -anything. Read this, then `tasks.md`. +Working state for this feature so a fresh session can pick it up without +re-deriving anything. Read this, then `tasks.md`. + +--- + +## 0. Read this first + +Three things invalidate earlier versions of this document: + +1. **It is deployed.** The feature shipped to production in release 1.16.0 and the + platform deploy succeeded on 2026-08-28, so `GSI7`, the Bedrock service role and + the four Lambdas exist in **both** dev and prod. Earlier revisions of this file + said "Nothing deployed"; that is no longer true. +2. **Five defects were found only by running it**, in a row, each one step further + along than the last. Every one reviewed clean and deployed clean. They are + §5 items 24–28 and they are the most useful part of this document. +3. **Iterate locally.** `scripts/local-dev/run-kb-migration.py` drives the whole + state machine in-process against dev with your SSO credentials. Three of the + five defects would have been minutes of work instead of a merge → image build → + deploy → 15-minute-tick cycle each. Use it. --- @@ -11,11 +30,29 @@ anything. Read this, then `tasks.md`. | | | |---|---| -| Spec | Complete, audited 3× to clean. 25 requirements, 201 criteria, 0 dangling refs | -| Implementation | Groups **1–13** done, plus group 14 except 14.5. 14.4's one-click document retry is deferred. 4 subtasks left: 14.4's retry, 14.5, and group 15 | -| Tests | 617 infra (jest) · **6,603** backend (pytest, 6 m 20 s) · **1,886** frontend (vitest, 7 s) · 5 pre-existing unrelated failures | -| Deployed | **Nothing.** No `cdk deploy`, no AWS mutation, at any point | -| Feature flags | `migrationEnabled` **on in development**, off in production. `newDefault` and `reconcilerArmed` off in both (explicit `false`, set as GitHub Environment variables) | +| Spec | Complete. Requirement **8.5 was amended by measurement on 2026-08-31** — see §5.28 | +| Implementation | Groups 1–14 except 14.5. A migration has completed `shadow → verify → promote → retain` in dev and serves from the managed backend | +| Tests | 626 infra (jest) · ~6,780 backend (pytest) · 1,936 frontend (vitest) · 5 pre-existing unrelated Strands failures | +| Deployed | **dev and prod.** Flags off in prod; `migrationEnabled` on in dev | +| Open PRs | **#898** (persist + adopt-by-name) — validated locally, still needs merging | +| Uncommitted | the 8.5 amendment, the reranking/embedding change, the verify-defer, the local driver | + +### Flag state (GitHub Environment variables) + +| Flag | development | production | +|---|---|---| +| `CDK_MANAGED_KB_MIGRATION_ENABLED` | `true` | `false` | +| `CDK_MANAGED_KB_NEW_DEFAULT` | `false` | `false` | +| `CDK_MANAGED_KB_RECONCILER_ARMED` | `false` | `false` | +| `CDK_TAG_ENVIRONMENT` | `dev` | `prod` | + +`newDefault` has **no reader anywhere in `backend/src`** — "new knowledge bases are +created managed" is design §14.7 steps 5–8, a follow-up spec. Setting it does +nothing, which is worth knowing before someone flips it expecting an effect. + +⚠️ **Production carries every defect fixed after 1.16.0 shipped.** It cannot fire, +because nothing enrols while `migrationEnabled` is false. Do not turn that flag on +in prod until #898 and the uncommitted work have landed and shipped. ### Commits (16 on the branch, all pushed) @@ -91,6 +128,30 @@ There is **no eslint config** in the repo despite the steering docs mentioning ESLint; `npx eslint` fails with "couldn't find an eslint.config.*". Type-check with `tsc --noEmit` and build with `ng build` instead. +### Driving a migration locally (do this before deploying anything) + +```bash +cd backend +uv run python ../scripts/local-dev/run-kb-migration.py --show +uv run python ../scripts/local-dev/run-kb-migration.py --break-lease +``` + +Runs the worker's steps in-process against dev with your SSO credentials, so the +whole state machine iterates in seconds. `--break-lease` clears the 15-minute lease +between steps and defers `dueAt` 20 minutes out so the deployed dispatcher does not +race you. + +It needs five variables in `backend/src/.env` that the app-api task definition does +not carry — copy them from the deployed worker Lambda: +`MANAGED_KB_SERVICE_ROLE_ARN`, `MANAGED_KB_TAG_VALUE_PREFIX`, +`MANAGED_KB_TAG_VALUE_ENVIRONMENT`, `MANAGED_KB_METRIC_NAMESPACE`, +`KB_MIGRATION_RETAIN_DAYS`. + +**What it does not prove:** the worker Lambda's IAM role (your SSO identity is +broader), the CDK environment wiring, or the image contents. Those are deploy-time +concerns — check them by deploying. Getting the logic right here first is the point, +and two of the five defects below were IAM/wiring and could only surface that way. + ### Running the upgrade UI locally ```bash @@ -367,36 +428,120 @@ saying why that number is a property of AWS rather than a knob. --- +### The five that only running it revealed + +These came out in sequence over 2026-08-27 to 08-31, each one step further into the +saga than the last. Every one reviewed clean, deployed clean, and did nothing or +failed on first real use. If you read only one section of this file, read this one. + +24. **The dispatcher could not find the worker.** The construct set + `MANAGED_KB_WORKER_FUNCTION_NAME`; `dispatcher.py` reads + `KB_MIGRATION_WORKER_FUNCTION_NAME` — the convention its siblings use. Every + tick raised `RuntimeError: ... is not set`, on a fifteen-minute schedule, in + silence unless someone read the logs. `kb-sync` does not have this bug for + exactly one reason: `kb-sync.test.ts` asserts the variable exists. + A second mismatch found by the same sweep: the construct published + `MANAGED_KB_RETENTION_WINDOW_DAYS`, which nothing reads, while + `worker._retain_days()` reads `KB_MIGRATION_RETAIN_DAYS` — so Req 15.11's + configured window was silently replaced by the code's 30-day floor. **Guard:** + `tests/supply_chain/test_kb_migration_env_contract.py` now asserts every + `os.environ` name the handlers read is set by the construct, and that the + construct publishes nothing unread. + +25. **`bedrock:TagResource` was not granted.** `CreateKnowledgeBase` is called + *with* tags and AWS authorises the tagging as a separate action, so the grant + reviewed as complete and failed the moment a real knowledge base was created. + `bedrock:ListTagsForResource` was missing for the same reason, with a quieter + failure: the reconciler fails closed on a tag read, so every knowledge base + would look untagged and the orphan sweep would report a clean account forever. + +26. **`CreateDataSource` ran against a `CREATING` knowledge base.** + `CreateKnowledgeBase` returns before the knowledge base is usable — this + module's own header records 47–124 s to `ACTIVE` — and the code called + `CreateDataSource` immediately. `ConflictException` is deliberately not + retryable and `_call`'s backoff tops out near 60 s, so the wait had to be + explicit. **Why no test caught it:** `FakeBedrockAgent.create_knowledge_base` + returned `status: "ACTIVE"`, which the real API never does. + +27. **The knowledge base id was never recorded until both creates succeeded.** + `attach_aws_ids` needs both identifiers, so a failure between them left a + record with no `awsKbId` — and every later attempt re-entered the create path + and was refused, permanently, because the *name* was taken. **The + `clientToken` does not save this**: AWS idempotency tokens expire within + minutes. Earlier revisions of this document and of the module header claimed + otherwise; they were wrong. Fixed by `records.attach_knowledge_base_id` + (persist immediately) plus adopt-by-name for records already stuck. + **Why no test caught it:** two tests *certified* the bug — + `test_the_record_survives_as_a_discoverable_retry_anchor` asserted + `"awsKbId" not in anchor`, and its sibling relied on the fake modelling + `clientToken` dedup as **permanent** while not modelling name uniqueness at + all. The fake now enforces name uniqueness and treats tokens as expired by + default. + +28. **The embedding pin and managed reranking are mutually exclusive.** Req 8.5 + pinned `titan-embed-text-v2:0` via `embeddingModelType: CUSTOM`; Req 11.2 + requires `rerankingModelType: MANAGED`. AWS rejects the combination, and the + §13 evaluation had measured the two **separately, never together**. Req 8.5 is + now amended: the pin protected a failure mode that cannot occur in managed mode + (we never embed the query — Bedrock embeds both sides), and the evaluation + measured the pin as worth nothing while reranking measurably separates scores. + Confirmed in dev: pinned + `NONE` gives flat 1.00/0.982/0.952; unpinned + + `MANAGED` gives 0.413/0.199. + +29. **`verify` failed a good migration for being asked too early.** The canary + retrieval returned nothing because the freshly-ingested document was not yet + queryable, and that was treated as terminal. Measured: **~45 s** from ingest to + retrievable on a fresh knowledge base, against the docstring's "0.75–1.03 s" + (a warm-knowledge-base figure). `verify` now defers via + `records.defer_verify`, bounded at `MAX_VERIFY_ATTEMPTS`. Adoption also learned + to skip knowledge bases in `DELETING`, found the same way: a local recreate + adopted one mid-delete. + +--- + ## 6. Remaining work -| Group | Subtasks | Notes | -|---|---|---| -| **14** Surfaces | 1½ | **14.5** admin surface (filter by engine, stored bytes + document counts, bulk migrate, per-KB retry) — not started. **14.4** is surfaced but its one-click document retry is deferred; see the deferral below. 14.0–14.3, 14.6, 14.7 are **done**. | -| **15** Pre-promotion verification | 3 | The gate before any real traffic moves. | +### Do these first -### Known deferrals (correct, not oversights) +| | | +|---|---| +| **Merge #898** | persist-the-id + adopt-by-name. Validated locally; it is what unstuck the dev record | +| **Commit the uncommitted** | Req 8.5 amendment, the embedding/reranking change, the verify-defer, the adoption `DELETING` filter, `scripts/local-dev/run-kb-migration.py` | +| **Then deploy and click through in dev** | the local run proves the logic; the deploy proves the IAM and the wiring | -- **One-click document reprocess (Req 21.2).** Ingestion is S3-event-triggered - (`documents/ingestion/handler.py`) and there is **no reprocess endpoint** — the - only document writes are upload-url, import, upload-failed and delete. A retry - control therefore needs new backend that re-fires the pipeline against bytes - already in S3, which is a change to a live ingestion path. Deliberately not - improvised. The card directs the user to re-upload via "Add files", a retry path - that works today. **Close by building the endpoint or by amending Req 21.2 to - accept re-upload** — do not leave it ambiguous. -- **`backend/Dockerfile.kb-migration`** does not exist yet, on purpose. The real image - needs five artefacts that do not exist: the handler modules, their - `requirements.txt`, a case in `scripts/build/build-one.sh`, `backend.yml` jobs, and - entries in the **hand-maintained** lists in - `backend/tests/supply_chain/test_dockerfile_pinning.py` and - `test_lambda_image_imports.py`. Per platform-as-bootstrap, CDK ships the bootstrap - stub and the **workflow** ships the real image. -- **Reconciler EventBridge wiring** (Reqs 14.1, 14.7) — `infrastructure/`, platform - group. Backend code never deploys before the IAM and resources it requires. -- **Group 7's snapshot reservation now has its caller** (`run_shadow`), reserving - the whole corpus before anything is provisioned. +### Open, in rough order ---- +| Group | Notes | +|---|---| +| **14.4** one-click document retry | Req 21.2. Ingestion is S3-event-triggered and there is no reprocess endpoint, so this needs new backend against a live pipeline. The card currently directs the user to re-upload, which works today. Close it by building the endpoint **or** by amending Req 21.2 to accept re-upload | +| **14.5** admin surface | not started. Filter by engine, stored bytes, document counts, bulk migrate, per-KB retry | +| **15.1** packaged-SDK probe | the *static* half is done and passing (`boto3==1.43.68` carries `MANAGED`, the embedding members, `FLOAT32`, all four document ops, no `AWS_DATA_PATH`). The live half has now effectively been done by hand — a real create → ingest → retrieve → promote succeeded in dev | +| **15.2** ingestion-concurrency probe | unanswered. Do not size a wide fleet migration before it | +| **15.3** full matrix | run it once the above land | + +### Known unknown, picked up mid-flight + +**`document_id` and `relevance` come back empty from the facade.** After promotion, +`search_assistant_knowledgebase_with_formatting` returned two real chunks with +correct content and correct `distance` ordering (−0.4226 then −0.1616 — negation of +relevance, so ascending distance is descending relevance, as designed). But +`document_id` was `None` and `relevance` was `None` on the facade output. + +`document_id` is the **join key for the document-status filter**, which fails +closed — so chunks with no `document_id` should have been dropped and were not. +Either the filter is not being applied on the managed path, or the id is being lost +between `managed_backend._to_chunk` and the facade. Worth resolving before any real +traffic moves; it was found with ~5 minutes of context left rather than chased. + +Reproduce with `scripts/local-dev/run-kb-migration.py ast-1a90784a7f18 --show` +followed by a facade query — note `resolver` has no `get_backend`; find the real +accessor. + +### Known deferrals (correct, not oversights) + +- **Reconciler EventBridge wiring** (Reqs 14.1, 14.7) — the rule exists and is + enabled; `reconcilerArmed` stays off so it reports rather than deletes. +- **Group 7's snapshot reservation** has its caller (`run_shadow`). ## 7. File map diff --git a/.kiro/specs/managed-kb-migration/requirements.md b/.kiro/specs/managed-kb-migration/requirements.md index a543fd7a..7d4b9b4e 100644 --- a/.kiro/specs/managed-kb-migration/requirements.md +++ b/.kiro/specs/managed-kb-migration/requirements.md @@ -317,8 +317,42 @@ capability we are paying for. real connector type in `managedKnowledgeBaseConnectorConfiguration.connectorParameters`. 4. THE system SHALL use connector type `CUSTOM`. -5. THE system SHALL set `embeddingModelType: CUSTOM` pinned to - `amazon.titan-embed-text-v2:0` at `FLOAT32` (the service-model enum value; lowercase is rejected) and 1024 dimensions. +5. THE system SHALL NOT send an embedding pin — no `embeddingModelType`, no + `embeddingModelArn`, no `embeddingModelConfiguration` — and SHALL let Bedrock + choose and manage the embedding model for a managed knowledge base. + + > **Amended 2026-08-31, by measurement.** This criterion previously required + > `embeddingModelType: CUSTOM` pinned to `amazon.titan-embed-text-v2:0` at + > `FLOAT32` and 1024 dimensions. That was carried over from the Legacy_Backend + > without re-deriving it, and it is wrong here for two independent reasons. + > + > **It protected a failure mode that cannot occur.** On S3 Vectors *we* embed + > the user's question (`s3vectors_backend` → `apis.shared.embeddings`), so the + > query model must match the model that indexed the documents or the similarity + > search compares vectors from different spaces. Managed retrieval sends + > `retrievalQuery={"text": …}` and managed ingestion sends `inlineContent` + > text — we never produce a vector. Bedrock embeds both sides itself, so + > consistency is the service's invariant and not ours to get wrong by omission. + > + > **AWS refuses the pin together with managed reranking.** Measured against + > dev, all four combinations: + > + > | Embedding | `rerankingModelType` | Result | + > |---|---|---| + > | `CUSTOM` (pinned) | `MANAGED` | `ValidationException` | + > | `CUSTOM` (pinned) | `NONE` | ok — scores 1.00 / 0.982 / 0.952 (flat) | + > | default (managed) | `MANAGED` | ok — scores 0.413 / 0.199 (separated) | + > | default (managed) | `NONE` | ok | + > + > The pin and Requirement 11.2's reranking are therefore mutually exclusive. + > §13 measured the pin as worth nothing ("identical cold-ingest time and + > identical answer quality to the built-in embedding — 9/9 either way") and + > reranking as worth a great deal ("the reranker is what makes a small context + > cap defensible"). Keeping reranking is the side with evidence behind it. + > + > Requirement 8.8's immutability still holds and now matters more: the choice + > cannot be revisited per knowledge base after creation. + 6. THE system SHALL enable `mediaExtractionConfiguration.imageExtractionConfiguration.imageExtractionStatus = ENABLED` on the data source. diff --git a/backend/src/apis/app_api/kb_migration/worker.py b/backend/src/apis/app_api/kb_migration/worker.py index 89165eea..49da7230 100644 --- a/backend/src/apis/app_api/kb_migration/worker.py +++ b/backend/src/apis/app_api/kb_migration/worker.py @@ -75,6 +75,15 @@ #: worker's knowledge base is picked up again the same hour. Requirement 15.13. LEASE_MINUTES = 15 +#: How long to wait before re-asking whether the managed corpus is queryable. +#: Measured ~45 s for a first ingest into a fresh knowledge base, so this re-asks +#: a little either side of that rather than guessing a single number. +VERIFY_RETRY_SECONDS = 60 + +#: Bound on those deferrals. Ten minutes of "not queryable" is no longer latency, +#: it is a corpus that will never answer — and the owner deserves to be told. +MAX_VERIFY_ATTEMPTS = 10 + #: Catch-up passes before the worker gives up waiting for quiet. A knowledge base #: whose owner is actively uploading may never converge; stopping is correct — #: the record stays in ``shadow``, the dispatcher brings it back, and the corpus @@ -578,9 +587,36 @@ async def run_verify( canary_text = _canary_query(complete) chunks = await backend.search(app_kb_id, canary_text, 5) if not chunks: - raise VerificationFailed( - f"canary retrieval on kb {app_kb_id} returned nothing; the managed " - f"corpus is not queryable yet" + # NOT a verification failure — a verification that has not happened yet. + # A first ingest into a fresh knowledge base took ~45 s to become + # retrievable when measured against dev; the docstring's 0.75-1.03 s was a + # warm-knowledge-base figure. Failing here marked a perfectly good + # migration `failed` for being asked too early, and the owner then saw a + # retry button for a problem that would have resolved itself. + attempts = await asyncio.to_thread( + r.defer_verify, + assistant_id, + app_kb_id, + generation, + _iso(_now() + timedelta(seconds=VERIFY_RETRY_SECONDS)), + ) + if attempts > MAX_VERIFY_ATTEMPTS: + raise VerificationFailed( + f"canary retrieval on kb {app_kb_id} still returned nothing after " + f"{attempts} attempts over ~" + f"{attempts * VERIFY_RETRY_SECONDS // 60} minutes; the corpus " + f"never became queryable" + ) + logger.info( + f"kb {app_kb_id}: corpus not queryable yet (attempt {attempts}); " + f"deferring verify by {VERIFY_RETRY_SECONDS}s" + ) + return StepResult( + assistant_id=assistant_id, + app_kb_id=app_kb_id, + from_state=r.VERIFY, + to_state=r.VERIFY, + detail=f"corpus not queryable yet; deferred (attempt {attempts})", ) retrieved_ids = {chunk.document_id for chunk in chunks if chunk.document_id} diff --git a/backend/src/apis/shared/kb_backend/provisioning.py b/backend/src/apis/shared/kb_backend/provisioning.py index e9b2fba2..1557df33 100644 --- a/backend/src/apis/shared/kb_backend/provisioning.py +++ b/backend/src/apis/shared/kb_backend/provisioning.py @@ -298,20 +298,39 @@ def knowledge_base_payload( ``storageConfiguration`` is absent by construction — there is no key to accidentally set to ``None``, because a managed knowledge base has no vector store and sending one is rejected. + + NO EMBEDDING PIN. Requirement 8.5 originally pinned + ``amazon.titan-embed-text-v2:0`` via ``embeddingModelType: CUSTOM``, and that + was carried over from the legacy path without re-deriving it. It does not + apply here, for two independent reasons: + + 1. **The pin protected a failure mode that cannot occur in managed mode.** On + S3 Vectors *we* embed the user's question (``s3vectors_backend`` calls + ``apis.shared.embeddings``), so the query model must match the model that + indexed the documents or the similarity search compares vectors from + different spaces. Managed retrieval sends ``retrievalQuery={"text": ...}`` + and ingestion sends ``inlineContent`` text — we never produce a vector, so + Bedrock embeds both sides itself and consistency is the service's + invariant, not ours to get wrong. + 2. **AWS refuses the pin together with managed reranking.** Measured, both + ways:: + + CUSTOM + rerankingModelType=MANAGED -> ValidationException + CUSTOM + NONE -> ok, scores 1.0/0.982/0.952 (flat) + default + MANAGED -> ok + default + NONE -> ok + + The evaluation recorded the pin as worth nothing measurable ("identical + answer quality — 9/9 either way") and reranking as worth a lot ("the + reranker is what makes a small context cap defensible"). Given they are + mutually exclusive, keeping reranking is the side with evidence behind it. + + The choice is immutable per knowledge base, which is why it is argued here + rather than left as a default someone might flip casually. """ validate_client_token(client_token) - managed: Dict[str, Any] = { - # Requirement 8.5: pinned, and immutable from here on. - "embeddingModelType": EMBEDDING_MODEL_TYPE, - "embeddingModelArn": embedding_model_arn(region), - "embeddingModelConfiguration": { - "bedrockEmbeddingModelConfiguration": { - "dimensions": EMBEDDING_DIMENSIONS, - "embeddingDataType": EMBEDDING_DATA_TYPE, - } - }, - } + managed: Dict[str, Any] = {} if kms_key_arn: # Requirement 20.5, only where customer-managed encryption is required. managed["serverSideEncryptionConfiguration"] = {"kmsKeyArn": kms_key_arn} @@ -468,6 +487,69 @@ async def _call( KB_ACTIVE_POLL_SECONDS = 5.0 +#: Fragments that identify a "name already taken" conflict, as opposed to the +#: status conflict the ACTIVE wait handles. Matched on the message because AWS +#: uses one error code (``ConflictException``) for both. +_NAME_CONFLICT_FRAGMENTS = ("already exists",) + + +def _is_name_conflict(exc: BaseException) -> bool: + """Whether ``exc`` is AWS refusing a duplicate knowledge base *name*. + + Distinguished from the status conflict by message, because both arrive as + ``ConflictException``. Getting this wrong in either direction is bad: treating + a status conflict as a name conflict would adopt while still CREATING, and + treating a name conflict as fatal leaves a record that can never be retried. + """ + message = str(exc).lower() + if not any(f in message for f in _NAME_CONFLICT_FRAGMENTS): + return False + response = getattr(exc, "response", None) + if isinstance(response, Mapping): + code = response.get("Error", {}).get("Code") + # Only trust the message when the code agrees it is a conflict. + return code in ("ConflictException", "ValidationException") + return True + + +async def _find_knowledge_base_by_name(client, name: str) -> Optional[str]: + """Return the id of the knowledge base called ``name``, or ``None``. + + Names are unique per account and derived from ``app_kb_id``, so a match is + always this knowledge base's own earlier attempt — never someone else's. + + Paginated fully rather than reading the first page: a truncated scan that + missed the match would fall through to "the name is taken but I cannot find + it", turning a recoverable state into a permanent failure. + """ + next_token: Optional[str] = None + while True: + kwargs: Dict[str, Any] = {"maxResults": 100} + if next_token: + kwargs["nextToken"] = next_token + page = await asyncio.to_thread(lambda: client.list_knowledge_bases(**kwargs)) + for summary in page.get("knowledgeBaseSummaries") or []: + if summary.get("name") != name: + continue + status = str(summary.get("status") or "") + if status in KB_FAILED_STATUSES: + # Adopting a knowledge base that is on its way out guarantees a + # failure one step later, at the ACTIVE wait. Skip it: the name is + # about to free up, so a fresh create is the right move and the + # next attempt will make it. Observed while recreating a knowledge + # base locally — the delete had not finished, adoption took the + # DELETING one, and the run failed on a resource nobody wanted. + logger.info( + f"ignoring knowledge base {summary.get('knowledgeBaseId')} named " + f"{name}: it is {status} and cannot be adopted" + ) + continue + return summary.get("knowledgeBaseId") + next_token = page.get("nextToken") + if not next_token: + return None + + class KnowledgeBaseNotReady(Exception): """A knowledge base did not reach ``ACTIVE`` within the wait budget.""" @@ -523,8 +605,8 @@ async def _wait_for_knowledge_base_active( raise KnowledgeBaseNotReady( f"kb {aws_kb_id} was still {last_status} after {waited:.0f}s " f"(budget {budget:.0f}s); refusing {what}. The migration is " - f"resumable: the clientToken is deterministic, so the next attempt " - f"adopts this knowledge base rather than creating another." + f"resumable: the id is already recorded on the KB_Record, so the " + f"next attempt resumes from it rather than creating another." ) await sleep(interval) waited += interval @@ -624,8 +706,12 @@ async def provision_managed_kb( owner_user_id=owner_user_id, provisioning_state=r.PROVISIONING, client_token=kb_token, - embedding_model_id=EMBEDDING_MODEL_ID, - embedding_dimensions=EMBEDDING_DIMENSIONS, + # Not recorded: with no pin, Bedrock chooses the embedding model and + # nothing here would know which. A field naming Titan on a knowledge + # base Bedrock embedded with something else is worse than an absent + # one — nothing reads these for retrieval, so they can only mislead. + embedding_model_id=None, + embedding_dimensions=0, image_extraction=True, parser_config={ "imageExtractionStatus": IMAGE_EXTRACTION_STATUS, @@ -655,22 +741,70 @@ async def provision_managed_kb( aws_kb_id = (existing or {}).get("awsKbId") if not aws_kb_id: - response = await _call( - client.create_knowledge_base, - knowledge_base_payload( - name=name, - role_arn=role_arn, - client_token=kb_token, - description=f"Managed knowledge base for {app_kb_id}", - tags=build_tags(app_kb_id, owner_user_id, project_prefix, environment), - region=region, - kms_key_arn=kms_key_arn, - ), - what="CreateKnowledgeBase", - max_attempts=max_attempts, - sleep=sleep, - ) - aws_kb_id = response["knowledgeBase"]["knowledgeBaseId"] + try: + response = await _call( + client.create_knowledge_base, + knowledge_base_payload( + name=name, + role_arn=role_arn, + client_token=kb_token, + description=f"Managed knowledge base for {app_kb_id}", + tags=build_tags(app_kb_id, owner_user_id, project_prefix, environment), + region=region, + kms_key_arn=kms_key_arn, + ), + what="CreateKnowledgeBase", + max_attempts=max_attempts, + sleep=sleep, + ) + aws_kb_id = response["knowledgeBase"]["knowledgeBaseId"] + except Exception as exc: + # ADOPT-BY-NAME. A knowledge base already carrying this name means a + # previous attempt created one and did not get to record its id, so + # this record has no `awsKbId` to resume from and the create can never + # succeed again — the name is taken, permanently. + # + # The `clientToken` does NOT cover this, contrary to what the module + # header implies. AWS idempotency tokens expire in minutes; a retry an + # hour later is a new request that collides on the unique name. This is + # exactly how the first real migration became unretryable. + # + # Adopting is safe because the name is derived from `app_kb_id`, so a + # collision can only be *this* knowledge base's own earlier attempt. + if not _is_name_conflict(exc): + raise + adopted = await _find_knowledge_base_by_name(client, name) + if not adopted: + # The name is taken but nothing matching is visible — do not guess. + raise + logger.warning( + f"kb {app_kb_id}: adopting existing knowledge base {adopted}; its " + f"name was already taken, which means an earlier attempt created it " + f"without recording the id" + ) + aws_kb_id = adopted + emit_count(METRIC_PROVISION_ADOPTED, dimensions={"appKbId": app_kb_id}) + + # Persist the identifier NOW, before anything else can fail. Everything + # after this point is resumable; before it, a failure loses a paying + # resource and blocks every future attempt on the name. + try: + await asyncio.to_thread( + r.attach_knowledge_base_id, + assistant_id, + app_kb_id, + aws_kb_id, + _now_iso(), + ) + except r.TransitionLost: + # Another worker recorded an id first. Theirs wins; ours is either the + # same knowledge base (adopted by name) or a duplicate that the + # reconciler will find by tag. + logger.info( + f"kb {app_kb_id}: another worker recorded awsKbId first; deferring" + ) + refreshed = await asyncio.to_thread(r.get_kb_record, assistant_id, app_kb_id) + aws_kb_id = (refreshed or {}).get("awsKbId") or aws_kb_id # CreateKnowledgeBase returns as soon as the knowledge base is CREATING, not # when it is usable — this module's own header records 47–124 s to ACTIVE @@ -683,13 +817,6 @@ async def provision_managed_kb( # genuine conflict must fail fast — and `_call`'s backoff tops out around 60 s # anyway, short of the measured upper bound. So the wait is explicit rather # than a widened retry set. - # - # This is also why the failure orphaned a knowledge base on first run: the - # create succeeded, the data source did not, and `attach_aws_ids` never ran, so - # nothing recorded the id. The deterministic `clientToken` means a retry adopts - # that knowledge base rather than creating a second one, and the tags written - # at create make it discoverable by the reconciler — but the orphan existed at - # all only because of this missing wait. await _wait_for_knowledge_base_active( client, aws_kb_id, diff --git a/backend/src/apis/shared/kb_backend/records.py b/backend/src/apis/shared/kb_backend/records.py index 34256679..5a5af07e 100644 --- a/backend/src/apis/shared/kb_backend/records.py +++ b/backend/src/apis/shared/kb_backend/records.py @@ -306,6 +306,49 @@ def create_provisioning( return item +def attach_knowledge_base_id( + assistant_id: str, + app_kb_id: str, + aws_kb_id: str, + now_iso: str, +) -> None: + """Record ``awsKbId`` the moment the knowledge base exists in AWS. + + Deliberately separate from :func:`attach_aws_ids`, which needs both + identifiers and flips the record to ``active``. This one runs *between* the + two AWS creates, and exists because the gap between them is where a paying + resource can be lost. + + Without it: ``CreateKnowledgeBase`` succeeds, ``CreateDataSource`` fails, and + nothing has recorded the identifier — so the retry re-enters the create path + and AWS refuses it, permanently, because the *name* is already taken: + + ConflictException: KnowledgeBase with name ... already exists. + + The ``clientToken`` does not save this. AWS idempotency tokens expire within + minutes, so a retry hours or days later is a genuinely new request that + collides on the unique name. A record stuck this way can never be retried + successfully — which is exactly what happened to the first real migration. + + Guarded on ``attribute_not_exists(awsKbId)`` so a late-returning create from + an abandoned attempt cannot overwrite the identifier a newer one recorded, + and on still being ``provisioning`` so it cannot resurrect a torn-down record. + """ + _conditional( + _table().update_item, + Key={"PK": kb_pk(assistant_id), "SK": kb_sk(app_kb_id)}, + UpdateExpression="SET awsKbId = :kb, updatedAt = :now", + ConditionExpression=( + "attribute_not_exists(awsKbId) AND provisioningState = :provisioning" + ), + ExpressionAttributeValues={ + ":kb": aws_kb_id, + ":now": now_iso, + ":provisioning": PROVISIONING, + }, + ) + + def attach_aws_ids( assistant_id: str, app_kb_id: str, @@ -539,6 +582,46 @@ def set_migration_state( ) +def defer_verify( + assistant_id: str, + app_kb_id: str, + generation: int, + due_at: str, +) -> int: + """Push ``verify`` out and count the attempt. Returns the new attempt count. + + "The corpus is not queryable yet" is not a verification failure — it is a + verification that has not happened. Treating it as terminal marked a migration + `failed` for the crime of being asked too early: the document was ingested + correctly and became retrievable ~45 s later. + + The module docstring's "INDEXED precedes retrievable by 0.75-1.03 s" was + measured on a warm knowledge base; a first ingest into a fresh one is far + slower. Rather than encode either number, this defers and re-asks, bounded by + the caller so it cannot defer forever. + + Guarded on still being ``verify`` at this generation, so a deferral cannot + resurrect a migration that has since been promoted, failed or rolled back. + """ + response = _conditional( + _table().update_item, + Key={"PK": kb_pk(assistant_id), "SK": kb_sk(app_kb_id)}, + UpdateExpression="SET GSI7_SK = :due ADD verifyAttempts :one", + ConditionExpression=( + "migrationGeneration = :gen AND migrationState = :verify" + ), + ExpressionAttributeValues={ + ":due": due_at, + ":one": Decimal(1), + ":gen": Decimal(generation), + ":verify": VERIFY, + }, + ReturnValues="UPDATED_NEW", + ) + attempts = (response or {}).get("Attributes", {}).get("verifyAttempts") + return int(attempts) if attempts is not None else 1 + + def acquire_lease( assistant_id: str, app_kb_id: str, diff --git a/backend/tests/lambdas/test_kb_migration_worker.py b/backend/tests/lambdas/test_kb_migration_worker.py index 87965147..aa355316 100644 --- a/backend/tests/lambdas/test_kb_migration_worker.py +++ b/backend/tests/lambdas/test_kb_migration_worker.py @@ -625,22 +625,51 @@ def test_a_document_with_no_hash_still_contributes_a_changing_value(self): assert worker.manifest_entry(item) == "d9:2026-08-01T00:00:00Z" @pytest.mark.asyncio - async def test_verify_requires_a_canary_retrieval_to_return_something(self): - """Requirement 15.7. Bedrock reporting a document INDEXED precedes it being - retrievable by 0.75-1.03 s, and a knowledge base can hold documents while - returning nothing, so "we ingested everything" and "retrieval works" are - separate claims.""" + async def test_an_unqueryable_corpus_defers_instead_of_failing(self): + """Requirement 15.7, corrected by measurement. + + "Not queryable yet" is a verification that has not happened, not one that + failed. The docstring's 0.75-1.03 s was measured on a warm knowledge base; + a first ingest into a fresh one took ~45 s in dev, and treating that as + terminal marked a good migration `failed` and showed its owner a retry + button for a problem that resolves itself. + """ + backend = StubBackend(chunks=[]) + deferred = [] + + def _defer(_assistant, _kb, _generation, due_at): + deferred.append(due_at) + return len(deferred) + + with patch.dict("os.environ", BASE_ENV, clear=True), patch.object( + worker, "list_document_items", return_value=[_doc("d1")] + ), patch.object(r, "defer_verify", _defer): + result = await worker.run_verify( + ASSISTANT_ID, ASSISTANT_ID, _kb_record(r.VERIFY), backend + ) + + assert deferred, "did not defer; an early canary would fail the migration" + assert result.to_state == r.VERIFY, "left verify on a deferral" + assert "not queryable" in result.detail + + @pytest.mark.asyncio + async def test_deferring_forever_eventually_fails(self): + """Bounded, because "not queryable" past some point is not latency. + + Matched on the attempt count rather than "not queryable", so this cannot + be satisfied by the deferral path it is meant to sit past. + """ backend = StubBackend(chunks=[]) with patch.dict("os.environ", BASE_ENV, clear=True), patch.object( worker, "list_document_items", return_value=[_doc("d1")] + ), patch.object( + r, "defer_verify", lambda *a, **k: worker.MAX_VERIFY_ATTEMPTS + 1 ): - # Matched on "not queryable", not on "canary": both failure messages - # mention the canary, so the looser pattern passed even with the - # empty-result check removed — the *other* check raised and the test - # could not tell the difference. - with pytest.raises(worker.VerificationFailed, match="not queryable"): - await worker.run_verify(ASSISTANT_ID, ASSISTANT_ID, _kb_record(r.VERIFY), backend) + with pytest.raises(worker.VerificationFailed, match="attempts over"): + await worker.run_verify( + ASSISTANT_ID, ASSISTANT_ID, _kb_record(r.VERIFY), backend + ) @pytest.mark.asyncio async def test_verify_rejects_a_canary_that_returns_foreign_documents(self): diff --git a/backend/tests/shared/test_managed_kb_backend.py b/backend/tests/shared/test_managed_kb_backend.py index 30c8b69c..f82ec3be 100644 --- a/backend/tests/shared/test_managed_kb_backend.py +++ b/backend/tests/shared/test_managed_kb_backend.py @@ -57,6 +57,16 @@ # --------------------------------------------------------------------------- +def _conflict(message: str) -> Exception: + """A ClientError shaped like the real ConflictException.""" + from botocore.exceptions import ClientError + + return ClientError( + {"Error": {"Code": "ConflictException", "Message": message}}, + "CreateKnowledgeBase", + ) + + class FakeBedrockAgent: """A ``bedrock-agent`` control-plane stub with real idempotency semantics. @@ -75,10 +85,15 @@ def __init__( on_create=None, create_failures: Optional[List[Exception]] = None, status_sequence: Optional[List[str]] = None, + token_still_valid: bool = False, ) -> None: self.create_kb_calls: List[Dict[str, Any]] = [] self.create_ds_calls: List[Dict[str, Any]] = [] self.get_kb_calls: List[Dict[str, Any]] = [] + self.list_kb_calls: List[Dict[str, Any]] = [] + self._by_name: Dict[str, str] = {} + self._names_by_id: Dict[str, str] = {} + self.token_still_valid = token_still_valid self._status_sequence = list(status_sequence or ["ACTIVE"]) self.ingest_calls: List[Dict[str, Any]] = [] self.delete_calls: List[Dict[str, Any]] = [] @@ -105,14 +120,48 @@ def create_knowledge_base(self, **kwargs): raise self._create_failures.pop(0) token = kwargs["clientToken"] - if token not in self._by_token: - self._counter += 1 - self._by_token[token] = f"KB{self._counter:08d}" + name = kwargs["name"] + + # NAME UNIQUENESS, which is what AWS actually enforces and what this fake + # used to ignore. `clientToken` deduplication is real but *expires* within + # minutes, so a retry an hour later is a new request that collides on the + # name. Modelling the token as permanent is why two tests certified a + # design that could not recover: the first real migration created a + # knowledge base, failed before recording its id, and every retry + # thereafter was refused with "already exists". + # + # `token_still_valid` selects which side of that expiry is being modelled. + # It defaults to False — the realistic case for any retry that is not + # within the same few minutes. + if token in self._by_token and self.token_still_valid: + return { + "knowledgeBase": { + "knowledgeBaseId": self._by_token[token], + "status": "CREATING", + } + } + if name in self._by_name: + raise _conflict(f"KnowledgeBase with name {name} already exists.") + + self._counter += 1 + kb_id = f"KB{self._counter:08d}" + self._by_token[token] = kb_id + self._by_name[name] = kb_id + self._names_by_id[kb_id] = name # CREATING, not ACTIVE — what the real API returns. The fake previously # claimed ACTIVE here, which is why nothing caught the provisioner calling - # CreateDataSource against a knowledge base that was still creating and - # getting a ConflictException in dev. - return {"knowledgeBase": {"knowledgeBaseId": self._by_token[token], "status": "CREATING"}} + # CreateDataSource against a knowledge base that was still creating. + return {"knowledgeBase": {"knowledgeBaseId": kb_id, "status": "CREATING"}} + + def list_knowledge_bases(self, **kwargs): + """Summaries, so adopt-by-name can find a knowledge base it did not record.""" + self.list_kb_calls.append(kwargs) + return { + "knowledgeBaseSummaries": [ + {"knowledgeBaseId": kb_id, "name": name, "status": "ACTIVE"} + for kb_id, name in self._names_by_id.items() + ] + } def get_knowledge_base(self, **kwargs): """Status poll. Yields each queued status once, then settles on the last. @@ -385,22 +434,42 @@ def test_storage_configuration_is_omitted_entirely(self): def test_role_arn_is_passed(self): assert self._payload()["roleArn"] == ROLE_ARN - def test_embedding_is_pinned_to_titan_v2_float32_1024(self): - """Requirement 8.5, and immutable from here on (8.8). + def test_no_embedding_pin_is_sent(self): + """Requirement 8.5, as amended. + + The pin was carried over from the legacy path without re-deriving it, and + it does not apply here: + + * On S3 Vectors *we* embed the user's question, so the query model must + match the model that indexed the documents. Managed retrieval sends + text and managed ingestion sends text — we never produce a vector, so + Bedrock embeds both sides and consistency is its invariant. + * AWS rejects the pin together with ``rerankingModelType: MANAGED``, and + the evaluation measured the pin as worth nothing ("identical answer + quality — 9/9") against reranking being "what makes a small context cap + defensible". - A drift in any of these three values is not a migration but a rebuild, so - the numbers are asserted rather than merely present. + Asserted as *absence*, because sending any of these keys is what breaks + reranking — and the failure is a ValidationException at query time, long + after the immutable choice was made. """ managed = self._payload()["knowledgeBaseConfiguration"][ "managedKnowledgeBaseConfiguration" ] - assert managed["embeddingModelType"] == "CUSTOM" - assert managed["embeddingModelArn"].endswith("amazon.titan-embed-text-v2:0") - bedrock_config = managed["embeddingModelConfiguration"][ - "bedrockEmbeddingModelConfiguration" - ] - assert bedrock_config["dimensions"] == 1024 - assert bedrock_config["embeddingDataType"] == "FLOAT32" + assert "embeddingModelType" not in managed + assert "embeddingModelArn" not in managed + assert "embeddingModelConfiguration" not in managed + + def test_reranking_stays_managed(self): + """The half of the tradeoff that has evidence behind it (Req 11.2). + + Kept next to the pin test on purpose: these two are mutually exclusive in + AWS, so anyone reinstating the pin should see this failing beside it. + """ + from apis.shared.kb_backend.managed_backend import retrieval_configuration + + managed = retrieval_configuration()["managedSearchConfiguration"] + assert managed["rerankingModelType"] == "MANAGED" def test_kms_key_is_only_sent_when_supplied(self): assert "serverSideEncryptionConfiguration" not in self._payload()[ @@ -613,7 +682,13 @@ def _lose(assistant_id, record): ) @pytest.mark.asyncio - async def test_provisioning_requires_a_service_role(self, table): + async def test_provisioning_requires_a_service_role(self, table, monkeypatch): + # `role_arn=None` falls back to MANAGED_KB_SERVICE_ROLE_ARN, so this only + # asserted what it meant while that variable happened to be absent from the + # environment. It is now present in `backend/src/.env` for the local + # migration driver — which `load_dotenv(override=True)` reads — so the + # absence has to be made explicit rather than assumed. + monkeypatch.delenv("MANAGED_KB_SERVICE_ROLE_ARN", raising=False) with pytest.raises(p.ProvisioningError, match="service role"): await _provision(FakeBedrockAgent(), role_arn=None) @@ -685,10 +760,24 @@ class TestCrashBetweenCreateAndRecordUpdate: The window that record-first ordering exists to make survivable: the AWS knowledge base exists, the record does not yet name it. + + ⚠️ These tests previously certified the opposite of what they now assert, and + that is how the first real migration became permanently unretryable. They + asserted `"awsKbId" not in anchor` — the missing write, enshrined as a + requirement — and leaned on `clientToken` deduplication to make the retry + safe, which the fake modelled as **permanent**. Real AWS idempotency tokens + expire within minutes, so the retry was a new request that collided on the + unique name and was refused forever: + + ConflictException: KnowledgeBase with name ... already exists. + + The identifier is now persisted the moment it exists, which is what actually + makes the window survivable. The token is a nice-to-have inside the expiry + window; it is not the mechanism. """ @pytest.mark.asyncio - async def test_the_record_survives_as_a_discoverable_retry_anchor(self, table): + async def test_the_identifier_is_recorded_before_anything_else_can_fail(self, table): client = FakeBedrockAgent() crashed = [] @@ -713,15 +802,15 @@ def _crash(*_args, **_kwargs): "orphan nothing can find (Requirement 7.8)" ) assert anchor["provisioningState"] == r.PROVISIONING - assert "awsKbId" not in anchor - assert anchor["clientToken"], ( - "the anchor carries no clientToken, so a retry cannot be deduplicated " - "and would create a second knowledge base" + assert anchor.get("awsKbId") == "KB00000001", ( + "the identifier was not recorded, so a later retry cannot resume from " + "it and will be refused because the name is already taken" ) @pytest.mark.asyncio - async def test_the_retry_does_not_create_a_second_knowledge_base(self, table): - """The whole point: one knowledge base across a crash and a retry.""" + async def test_the_retry_resumes_instead_of_re_creating(self, table): + """One knowledge base across a crash and a retry — by resuming, not by + re-issuing a create and hoping AWS deduplicates it.""" client = FakeBedrockAgent() original = r.attach_aws_ids @@ -734,12 +823,69 @@ async def test_the_retry_does_not_create_a_second_knowledge_base(self, table): result = await _provision(client) # the retry - assert len(client.create_kb_calls) == 2, "the retry did not re-issue the create" - assert client.distinct_knowledge_base_ids == {"KB00000001"}, ( - "the retry created a SECOND knowledge base: the persisted clientToken " - "was not reused, so AWS did not deduplicate" + assert len(client.create_kb_calls) == 1, ( + "the retry re-issued CreateKnowledgeBase. With an expired idempotency " + "token that is refused outright — the name is taken — so resuming from " + "the recorded id is the only thing that works" ) + assert client.distinct_knowledge_base_ids == {"KB00000001"} assert result.aws_kb_id == "KB00000001" + + @pytest.mark.asyncio + async def test_adoption_ignores_a_knowledge_base_being_deleted(self, table): + """Adopting a dying knowledge base guarantees a failure one step later. + + Seen while recreating one locally: the delete had not finished, adoption + took the DELETING knowledge base, and the ACTIVE wait then refused it. The + name is about to free up, so skipping is correct — the next attempt creates + fresh. + """ + from apis.shared.kb_backend import provisioning as prov + + class _Dying: + def list_knowledge_bases(self, **_kwargs): + return { + "knowledgeBaseSummaries": [ + {"knowledgeBaseId": "KBDYING", "name": "wanted", + "status": "DELETING"}, + ] + } + + found = await prov._find_knowledge_base_by_name(_Dying(), "wanted") + assert found is None, "adopted a knowledge base that is being deleted" + + @pytest.mark.asyncio + async def test_a_lost_identifier_is_recovered_by_adopting_the_name(self, table): + """The state the first real migration was actually stuck in. + + The knowledge base exists in AWS, the record has no `awsKbId` (it was + written before this fix), and the token has long expired. Without + adopt-by-name the create is refused forever and the migration can never + succeed — the operator's only recourse would be deleting the knowledge + base by hand. + """ + client = FakeBedrockAgent() + await _provision(client) # creates KB00000001 and records it + + # Simulate the pre-fix record: identifier dropped, still provisioning. + table.update_item( + Key={"PK": r.kb_pk(ASSISTANT_ID), "SK": r.kb_sk(APP_KB_ID)}, + UpdateExpression="REMOVE awsKbId, awsDataSourceId SET provisioningState = :p", + ExpressionAttributeValues={":p": r.PROVISIONING}, + ) + + result = await _provision(client) + + assert result.aws_kb_id == "KB00000001", "did not adopt the existing name" + assert client.distinct_knowledge_base_ids == {"KB00000001"}, ( + "a second knowledge base was created; the name collision should have " + "been resolved by adoption, not by another create" + ) + assert client.list_kb_calls, "adoption did not look the name up" + assert _record(table).get("awsKbId") == "KB00000001", ( + "the adopted identifier was not persisted, so the next attempt would " + "have to adopt all over again" + ) assert _record(table)["provisioningState"] == r.ACTIVE @pytest.mark.asyncio @@ -849,7 +995,7 @@ async def test_the_failure_message_says_the_retry_is_safe(self, table): client = FakeBedrockAgent(status_sequence=["CREATING"]) with pytest.raises(p.KnowledgeBaseNotReady) as excinfo: await _provision(client, budget_seconds=5.0, interval_seconds=5.0) - assert "adopts this knowledge base" in str(excinfo.value) + assert "already recorded on the KB_Record" in str(excinfo.value) class TestOffEventLoop: 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 047a5fc0..393e811c 100644 --- a/infrastructure/lib/constructs/managed-kb/managed-kb-role-construct.ts +++ b/infrastructure/lib/constructs/managed-kb/managed-kb-role-construct.ts @@ -167,6 +167,35 @@ export function grantManagedKbProvisioning( * from CRUD so an ingestion-only caller can never delete a knowledge * base. * + * WHY `bedrock:StartIngestionJob` IS HERE WHEN NOTHING CALLS IT. + * It looks wrong, and removing it is the obvious "cleanup". It is not. + * `IngestKnowledgeBaseDocuments` is *authorized* under the adjacent + * action name `bedrock:StartIngestionJob` — AWS's own direct-ingestion + * prerequisites list both in one statement + * (bedrock/latest/userguide/kb-direct-ingestion-prereq.html). So the + * API this platform calls and the IAM action that permits it have + * different names, and granting only the matching name fails at the + * first real upload with: + * + * AccessDeniedException ... not authorized to perform: + * bedrock:StartIngestionJob on resource: knowledge-base/XXXXXXXXXX + * + * This is NOT a contradiction of Requirement 9.2, which forbids + * *calling* `StartIngestionJob` (0.1 RPS account-wide, one document + * every ten seconds). `managed_backend.py` never calls it and must + * never start; the grant is about authorization only. + * + * Same failure shape as the missing `bedrock:TagResource` — an + * adjacent action AWS checks separately, so the grant reviews as + * complete and dies on first real use. Found exactly that way: a + * document uploaded to a promoted knowledge base in dev went to + * `failed` with the message above, while the local driver's broader + * SSO identity had been masking it. + * + * `bedrock:ListKnowledgeBaseDocuments` is in AWS's example policy and + * deliberately omitted: no code path calls it, and the docs permit + * omitting actions. A future caller fails loudly rather than silently. + * * Also intentionally unattached for now — wired in task 2.1 alongside * the migration Lambdas, via * `ManagedKbRoleConstruct.grantDirectIngestion()`. @@ -177,6 +206,10 @@ export function grantManagedKbDirectIngestion(config: AppConfig, role: iam.IRole effect: iam.Effect.ALLOW, actions: [ 'bedrock:IngestKnowledgeBaseDocuments', + // The IAM action that actually authorizes the call above. See the + // docblock: this is authorization, not an invocation of the 0.1 RPS + // ingestion-job API that Requirement 9.2 forbids. + 'bedrock:StartIngestionJob', 'bedrock:DeleteKnowledgeBaseDocuments', 'bedrock:GetKnowledgeBaseDocuments', ], diff --git a/infrastructure/test/kb-migration.test.ts b/infrastructure/test/kb-migration.test.ts index f7b5ac4e..f40f8289 100644 --- a/infrastructure/test/kb-migration.test.ts +++ b/infrastructure/test/kb-migration.test.ts @@ -859,6 +859,26 @@ describe('KbMigrationConstruct — IAM', () => { expect(holders.some((id) => /IngestionConsumerLambdaServiceRole/.test(id))).toBe(true); }); + it('gives both ingesting roles bedrock:StartIngestionJob, not just the Ingest action', () => { + // Regression for the defect that failed a real upload in dev. AWS + // authorizes `IngestKnowledgeBaseDocuments` under the adjacent action + // name `bedrock:StartIngestionJob`, so a grant carrying only the + // matching name deploys and reviews clean, then returns + // AccessDeniedException on the first document. + // + // Both roles are checked because both call + // `ingest_knowledge_base_documents`: the ingestion consumer surfaced it, + // and the worker had the identical gap — invisible until now only + // because every migration so far was driven locally under a broader SSO + // identity than the Lambda role. + const statements = allStatements(t).filter((s) => s.Sid === 'ManagedKbDirectIngestion'); + expect(statements).toHaveLength(2); + for (const s of statements) { + expect(s.Action).toContain('bedrock:StartIngestionJob'); + expect(s.Action).toContain('bedrock:IngestKnowledgeBaseDocuments'); + } + }); + it('gives the dispatcher namespace-conditioned metrics and nothing Bedrock-shaped', () => { const statements = allStatements(t).filter((s) => s.Sid === 'ManagedKbDispatchMetrics'); expect(statements).toHaveLength(1); diff --git a/infrastructure/test/managed-kb.test.ts b/infrastructure/test/managed-kb.test.ts index 159abbb5..48d0de67 100644 --- a/infrastructure/test/managed-kb.test.ts +++ b/infrastructure/test/managed-kb.test.ts @@ -361,12 +361,50 @@ describe('ManagedKbRoleConstruct — caller grants', () => { const s = statementBySid(t, 'ManagedKbDirectIngestion'); expect(s.Action).toEqual([ 'bedrock:IngestKnowledgeBaseDocuments', + 'bedrock:StartIngestionJob', 'bedrock:DeleteKnowledgeBaseDocuments', 'bedrock:GetKnowledgeBaseDocuments', ]); expect(s.Resource).toBe(KB_ARN_WILDCARD); }); + it('grants bedrock:StartIngestionJob, which is what authorizes IngestKnowledgeBaseDocuments', () => { + // Regression. AWS authorizes `IngestKnowledgeBaseDocuments` under the + // adjacent action name `bedrock:StartIngestionJob` — both appear in one + // statement in AWS's direct-ingestion prerequisites. Granting only the + // name that matches the API call deploys clean, reviews clean, and then + // fails every real upload: + // + // AccessDeniedException ... not authorized to perform: + // bedrock:StartIngestionJob on resource: knowledge-base/XXXXXXXXXX + // + // That is what happened in dev: a document added to a promoted knowledge + // base went straight to `failed`. It had been masked because the local + // driver runs under a broader SSO identity than the Lambda role. + // + // Asserted on its own, not just inside the array above, so the reason + // survives a future reordering or trimming of that list — and so the + // "obvious cleanup" of an action nothing calls fails a test that says why. + const s = statementBySid(t, 'ManagedKbDirectIngestion'); + expect(s.Action).toContain('bedrock:StartIngestionJob'); + + // Both halves are required together; neither alone is sufficient. + expect(s.Action).toContain('bedrock:IngestKnowledgeBaseDocuments'); + }); + + it('gives every holder of the ingestion grant the StartIngestionJob authorization', () => { + // This suite attaches the grant to one fake role; the real pairing of + // worker + ingestion consumer is asserted in kb-migration.test.ts. What + // matters here is that whoever holds the statement holds both halves, + // since neither action alone permits an upload. + const holders = policiesWithSid(t, 'ManagedKbDirectIngestion'); + expect(holders).toHaveLength(1); + for (const holder of holders) { + expect(holder.json).toContain('bedrock:StartIngestionJob'); + expect(holder.json).toContain('bedrock:IngestKnowledgeBaseDocuments'); + } + }); + it('grants PutMetricData on the same non-reserved namespace to every calling identity', () => { for (const sid of [ 'ManagedKbProvisionMetrics', diff --git a/scripts/local-dev/run-kb-migration.py b/scripts/local-dev/run-kb-migration.py new file mode 100644 index 00000000..2cd24d12 --- /dev/null +++ b/scripts/local-dev/run-kb-migration.py @@ -0,0 +1,185 @@ +#!/usr/bin/env python3 +"""Drive a managed-KB migration locally against the dev environment. + +Does what the dispatcher + worker Lambdas do, but in-process with your SSO +credentials, so the shadow -> verify -> promote -> retain path can be iterated in +seconds instead of a merge, an image build, a deploy and a 15-minute tick. + + cd backend + uv run python ../scripts/local-dev/run-kb-migration.py ast-1a90784a7f18 + uv run python ../scripts/local-dev/run-kb-migration.py ast-... --once + uv run python ../scripts/local-dev/run-kb-migration.py ast-... --show + +WHAT THIS DOES AND DOES NOT PROVE + +It talks to the real Bedrock, DynamoDB and S3 in dev, so it exercises the actual +API contracts, the real state machine and real documents. What it does NOT +exercise is the worker Lambda's IAM role — your SSO identity is broader — nor the +CDK environment wiring, nor the image contents. Those are deploy-time concerns and +are checked by deploying. Getting the logic right here first is the point. + +Writes to real dev records. Point it at a knowledge base you are willing to churn. +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import os +import sys +from pathlib import Path + +BACKEND_SRC = Path(__file__).resolve().parents[2] / "backend" / "src" +sys.path.insert(0, str(BACKEND_SRC)) + +from dotenv import load_dotenv # noqa: E402 + +load_dotenv(BACKEND_SRC / ".env", override=True) + +REQUIRED = ( + "DYNAMODB_ASSISTANTS_TABLE_NAME", + "S3_ASSISTANTS_DOCUMENTS_BUCKET_NAME", + "MANAGED_KB_SERVICE_ROLE_ARN", +) + + +def _preflight() -> None: + missing = [v for v in REQUIRED if not os.environ.get(v)] + if missing: + sys.exit( + f"missing env: {missing}\nCopy them from the deployed worker Lambda into " + f"backend/src/.env." + ) + + +def _show(assistant_id: str) -> None: + from apis.shared.kb_backend import records as r + + record = r.get_kb_record(assistant_id, assistant_id) + if not record: + print(" no KB record") + return + interesting = [ + "migrationState", "migrationGeneration", "provisioningState", "awsKbId", + "awsDataSourceId", "retrievalEngine", "migrationProgress", "migrationError", + "totalBytes", "retainUntil", "promotedAt", "GSI7_PK", + ] + for key in interesting: + if key in record: + value = record[key] + if isinstance(value, dict): + value = {k: str(v) for k, v in value.items()} + text = str(value) + print(f" {key:22} = {text[:110]}") + engine = "managed" if record.get("retrievalEngine") == "managed" else "legacy (absent)" + print(f" {'-> serving from':22} = {engine}") + + +async def _drive(assistant_id: str, once: bool, max_steps: int, break_lease: bool) -> int: + from apis.app_api.kb_migration import worker + from apis.shared.kb_backend import records as r + + for step in range(1, max_steps + 1): + record = r.get_kb_record(assistant_id, assistant_id) + if not record: + print(" no KB record — enrol from the UI first, or it was torn down") + return 1 + state = record.get("migrationState") + print(f"\n── step {step}: state={state} gen={record.get('migrationGeneration')}") + + if state in (r.RETAIN, r.MIGRATION_FAILED): + print(f" terminal: {state}") + if state == r.MIGRATION_FAILED: + print(f" error: {str(record.get('migrationError'))[:400]}") + _show(assistant_id) + return 0 if state == r.RETAIN else 2 + + if break_lease and record.get("migrationLeaseUntil"): + # Simulates lease expiry so consecutive local steps do not have to wait + # out the 15-minute window. Safe ONLY because --break-lease is paired + # with deferring dueAt, which keeps the deployed dispatcher from + # claiming the same record while we hold it. Never run this against a + # record a real worker may be mid-step on. + _table(r).update_item( + Key={"PK": r.kb_pk(assistant_id), "SK": r.kb_sk(assistant_id)}, + UpdateExpression="REMOVE migrationLeaseUntil", + ) + print(" (lease cleared for local stepping)") + + try: + result = await worker.run_step(assistant_id, assistant_id) + print(f" result: {json.dumps(result.as_log_fields(), default=str)[:300]}") + except Exception as exc: # noqa: BLE001 — this is a diagnostic driver + print(f" RAISED {type(exc).__name__}: {str(exc)[:400]}") + _show(assistant_id) + return 3 + + if once: + _show(assistant_id) + return 0 + + print(f"\n stopped after {max_steps} steps without reaching a terminal state") + _show(assistant_id) + return 4 + + +def _table(records_module): + import boto3 + + return boto3.resource("dynamodb").Table( + os.environ["DYNAMODB_ASSISTANTS_TABLE_NAME"] + ) + + +def _defer(assistant_id: str, minutes: int) -> None: + """Push dueAt out so the deployed dispatcher leaves this record alone. + + The dispatcher sweeps on `GSI7_SK <= now`, so a future dueAt makes the record + invisible to it without removing it from the index — which means nothing is + lost if this driver dies part-way. + """ + from datetime import datetime, timedelta, timezone + + from apis.shared.kb_backend import records as r + + due = (datetime.now(timezone.utc) + timedelta(minutes=minutes)).isoformat() + due = due.replace("+00:00", "Z") + _table(r).update_item( + Key={"PK": r.kb_pk(assistant_id), "SK": r.kb_sk(assistant_id)}, + UpdateExpression="SET GSI7_SK = :due", + ConditionExpression="attribute_exists(GSI7_SK)", + ExpressionAttributeValues={":due": due}, + ) + print(f" deferred dueAt to {due} so the deployed dispatcher skips it") + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("assistant_id") + parser.add_argument("--once", action="store_true", help="run a single step") + parser.add_argument("--show", action="store_true", help="print state and exit") + parser.add_argument("--max-steps", type=int, default=12) + parser.add_argument("--break-lease", action="store_true", + help="clear the lease between steps (implies --defer)") + parser.add_argument("--defer", type=int, default=0, metavar="MIN", + help="push dueAt out so the deployed dispatcher skips it") + args = parser.parse_args() + + _preflight() + print(f"table={os.environ['DYNAMODB_ASSISTANTS_TABLE_NAME']}") + print(f"kb={args.assistant_id}") + + if args.show: + _show(args.assistant_id) + return 0 + minutes = args.defer or (20 if args.break_lease else 0) + if minutes: + _defer(args.assistant_id, minutes) + return asyncio.run( + _drive(args.assistant_id, args.once, args.max_steps, args.break_lease) + ) + + +if __name__ == "__main__": + sys.exit(main())