From 7542d907622bcb0cc07faae070bcf01d795f6024 Mon Sep 17 00:00:00 2001 From: derrickfink Date: Mon, 31 Aug 2026 11:29:16 -0600 Subject: [PATCH 1/3] fix(kb): record the knowledge base id before anything else can fail MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The retry after the ACTIVE fix did not reach the data source. It could not: ConflictException: KnowledgeBase with name dev-boisestateai-v2-kb-ast-1a90784a7f18 already exists. `awsKbId` was only ever written by `attach_aws_ids`, which runs after BOTH AWS creates succeed. So a failure between them left a record with no identifier, and every later attempt re-entered the create path and was refused, permanently, because the name was taken. The record was unrecoverable by any number of retries. The `clientToken` does not cover this, and this commit corrects that claim wherever it appears — including in a message I added two commits ago. AWS idempotency tokens expire within minutes; a retry an hour later is a genuinely new request that collides on the unique name. Two changes: * `records.attach_knowledge_base_id` persists `awsKbId` the moment the create returns, guarded on `attribute_not_exists(awsKbId)` so a straggler cannot overwrite a newer attempt's identifier. Everything after that point is resumable. * Provisioning adopts by name on a name-collision `ConflictException`. This is what recovers records already stuck in that state, including the one in dev — without it the only recourse is deleting the knowledge base by hand. Adoption is safe because the name is derived from `app_kb_id`, so a collision can only be this knowledge base's own earlier attempt. The status conflict and the name conflict share the `ConflictException` code, so they are told apart by message. WHY TWO TESTS CERTIFIED THE BUG `test_the_record_survives_as_a_discoverable_retry_anchor` asserted `"awsKbId" not in anchor` — the missing write, encoded as a requirement — and `test_the_retry_does_not_create_a_second_knowledge_base` asserted the retry re-issues the create and relied on the fake deduplicating it. `FakeBedrockAgent` modelled `clientToken` dedup as permanent and did not model name uniqueness at all, so both passed against a design that could not recover. The fake now enforces name uniqueness, treats tokens as expired by default, and answers `list_knowledge_bases`. Both tests are rewritten to assert what actually makes the window survivable, and a third covers the stuck-record state directly. Mutations verified caught: dropping the immediate persist, and dropping adopt-by-name. Tests: 2,465 passed across shared, lambdas, property, architecture and supply_chain. --- .../apis/shared/kb_backend/provisioning.py | 140 ++++++++++++++---- backend/src/apis/shared/kb_backend/records.py | 43 ++++++ .../tests/shared/test_managed_kb_backend.py | 133 ++++++++++++++--- 3 files changed, 273 insertions(+), 43 deletions(-) diff --git a/backend/src/apis/shared/kb_backend/provisioning.py b/backend/src/apis/shared/kb_backend/provisioning.py index e9b2fba2..a18178d9 100644 --- a/backend/src/apis/shared/kb_backend/provisioning.py +++ b/backend/src/apis/shared/kb_backend/provisioning.py @@ -468,6 +468,55 @@ 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: + 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 +572,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 @@ -655,22 +704,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 +780,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..d5ed9973 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, diff --git a/backend/tests/shared/test_managed_kb_backend.py b/backend/tests/shared/test_managed_kb_backend.py index 30c8b69c..f51e2745 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. @@ -685,10 +734,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 +776,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 +797,46 @@ 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_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 +946,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: From 6420f148024311f595246667e7d3eb308a19a9a6 Mon Sep 17 00:00:00 2001 From: derrickfink Date: Mon, 31 Aug 2026 12:19:45 -0600 Subject: [PATCH 2/3] fix(kb): drop the embedding pin, defer verify, and complete a migration in dev MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A migration now runs shadow -> verify -> promote -> retain and serves from the managed backend. Three further defects, all found by driving the state machine locally against dev instead of through a deploy cycle. THE EMBEDDING PIN AND MANAGED RERANKING ARE MUTUALLY EXCLUSIVE (Req 8.5 amended) Req 8.5 pinned titan-embed-text-v2:0 via `embeddingModelType: CUSTOM`; Req 11.2 requires `rerankingModelType: MANAGED`. AWS rejects the combination, and §13 had measured the two separately, never together. Measured all four: CUSTOM + MANAGED -> ValidationException CUSTOM + NONE -> ok, scores 1.00/0.982/0.952 (flat) default + MANAGED -> ok, scores 0.413/0.199 (separated) default + NONE -> ok The pin loses, for a better reason than "it buys little": it protected a failure mode that cannot occur here. On S3 Vectors *we* embed the question, so the query model must match the index. Managed retrieval sends text and managed ingestion sends text — we never produce a vector, so Bedrock embeds both sides and consistency is its invariant. §13 measured the pin as worth nothing (9/9 identical) against reranking being "what makes a small context cap defensible". `embeddingModelId`/`embeddingDimensions` are no longer recorded either: with no pin, nothing here knows what Bedrock chose, and a field naming Titan on a knowledge base embedded with something else is worse than an absent one. VERIFY FAILED A GOOD MIGRATION FOR BEING ASKED TOO EARLY The canary returned nothing because the freshly-ingested document was not yet queryable, and that was terminal. Measured ~45 s from ingest to retrievable on a fresh knowledge base, against the docstring's 0.75-1.03 s (a warm figure). Now defers via `records.defer_verify`, bounded at MAX_VERIFY_ATTEMPTS, so latency reads as latency and only a corpus that never answers fails. ADOPTION TOOK A KNOWLEDGE BASE THAT WAS BEING DELETED Found while recreating one locally: the delete had not finished, adopt-by-name took the DELETING knowledge base, and the ACTIVE wait then refused it. Adoption now skips terminal statuses — the name is about to free up, so a fresh create is right. ALSO A test that only passed while MANAGED_KB_SERVICE_ROLE_ARN was absent now deletes it explicitly; the variable is needed in backend/src/.env for the local driver. `scripts/local-dev/run-kb-migration.py` drives the state machine in-process against dev. Three of the last five defects would have been minutes each with it. HANDOFF.md is rewritten: it previously said "Nothing deployed", which has been untrue since 1.16.0 shipped the feature to production behind flags. Tests: 6,817 backend passed (5 pre-existing Strands failures), 626 infra. --- .kiro/specs/managed-kb-migration/HANDOFF.md | 211 +++++++++++++++--- .../managed-kb-migration/requirements.md | 38 +++- .../src/apis/app_api/kb_migration/worker.py | 42 +++- .../apis/shared/kb_backend/provisioning.py | 67 ++++-- backend/src/apis/shared/kb_backend/records.py | 40 ++++ .../tests/lambdas/test_kb_migration_worker.py | 51 ++++- .../tests/shared/test_managed_kb_backend.py | 75 +++++-- scripts/local-dev/run-kb-migration.py | 185 +++++++++++++++ 8 files changed, 632 insertions(+), 77 deletions(-) create mode 100644 scripts/local-dev/run-kb-migration.py 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 a18178d9..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} @@ -510,8 +529,22 @@ async def _find_knowledge_base_by_name(client, name: str) -> Optional[str]: 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: - return summary.get("knowledgeBaseId") + 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 @@ -673,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, diff --git a/backend/src/apis/shared/kb_backend/records.py b/backend/src/apis/shared/kb_backend/records.py index d5ed9973..5a5af07e 100644 --- a/backend/src/apis/shared/kb_backend/records.py +++ b/backend/src/apis/shared/kb_backend/records.py @@ -582,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 f51e2745..f82ec3be 100644 --- a/backend/tests/shared/test_managed_kb_backend.py +++ b/backend/tests/shared/test_managed_kb_backend.py @@ -434,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). - - A drift in any of these three values is not a migration but a rebuild, so - the numbers are asserted rather than merely present. + 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". + + 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()[ @@ -662,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) @@ -805,6 +831,29 @@ async def test_the_retry_resumes_instead_of_re_creating(self, table): 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. 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()) From fdf15d21f2388d4fd313355352f99246b9654b61 Mon Sep 17 00:00:00 2001 From: derrickfink Date: Mon, 31 Aug 2026 14:21:10 -0600 Subject: [PATCH 3/3] fix(kb): grant bedrock:StartIngestionJob, which authorizes direct ingestion MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A document uploaded to a promoted knowledge base in dev went to `failed` with: AccessDeniedException ... IngestKnowledgeBaseDocuments ... not authorized to perform: bedrock:StartIngestionJob on resource: knowledge-base/M8WQZVQJ8X AWS authorizes `IngestKnowledgeBaseDocuments` under the adjacent action name `bedrock:StartIngestionJob`; both appear in one statement in AWS's direct-ingestion prerequisites. The grant carried only the name matching the API call, so it reviewed as complete, deployed clean, and failed on the first real upload — the same shape as the missing `bedrock:TagResource`. The worker had the identical gap. It receives the same grant and calls the same API, so the first migration driven by the deployed dispatcher would have failed identically. It stayed invisible because every migration so far was driven by scripts/local-dev/run-kb-migration.py under an SSO identity broader than either Lambda role. The action is easy to mistake for a mistake: Requirement 9.2 forbids *calling* StartIngestionJob (0.1 RPS account-wide, one document per ten seconds) and nothing does. Holding it is authorization, not invocation. A docblock and a separately-named test carry that reason so the obvious cleanup fails a test that explains itself. `bedrock:ListKnowledgeBaseDocuments` is in AWS's example policy and left out on purpose: no code path calls it. Guards: three tests, including one asserting both the worker and the ingestion-consumer roles carry the action. Mutation-tested — removing the action fails exactly four tests, each named for the reason. --- .../managed-kb/managed-kb-role-construct.ts | 33 ++++++++++++++++ infrastructure/test/kb-migration.test.ts | 20 ++++++++++ infrastructure/test/managed-kb.test.ts | 38 +++++++++++++++++++ 3 files changed, 91 insertions(+) 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',