From 46b112e98f7c28d1b9e3fb7a75be0d29451fd157 Mon Sep 17 00:00:00 2001 From: derrickfink Date: Thu, 27 Aug 2026 21:48:08 -0600 Subject: [PATCH] fix(kb): wait for the knowledge base to be ACTIVE before its data source MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The first knowledge base that got past IAM was created and then abandoned: ConflictException: The Knowledge Base is not in a valid status. Wait for the knowledge base to reach a valid status and try again. `CreateKnowledgeBase` returns while the knowledge base is still `CREATING`. This module's own header records 47-124 s to `ACTIVE` (n=7) — the knowledge was there, just not in the code path, which called `CreateDataSource` immediately. `ConflictException` is deliberately absent from RETRYABLE_ERROR_CODES, because a genuine conflict should fail fast, and `_call`'s backoff tops out near 60 s anyway — short of the measured upper bound. Retrying the dependent call is the wrong shape regardless: it spends attempts on a precondition instead of waiting for it. So the wait is explicit and bounded at 300 s, roughly 2.5x the observed worst case and well inside the worker's 15-minute timeout. A terminal status (FAILED/DELETING) fails on the first poll rather than waiting out the budget. Consequence worth recording: the create succeeded, the data source did not, and `attach_aws_ids` never ran — so the knowledge base existed in AWS with no record pointing at it. Two pre-existing decisions contained it. The `clientToken` is derived from `app_kb_id`, so the next attempt *adopts* that knowledge base rather than creating a second one; and tags are written at create, so the reconciler can find it either way. The failure message now says so, because an operator reading it needs to know a retry is safe. WHY NO TEST CAUGHT THIS `FakeBedrockAgent.create_knowledge_base` returned `status: "ACTIVE"`, which the real API never does. Every provisioning test therefore skipped straight past the window where this breaks. The fake now returns `CREATING` and answers `get_knowledge_base` from a configurable status sequence, so the wait is exercised rather than assumed. Five new tests cover it, including that the data source is never attempted while CREATING and that the budget is read at call time. Mutation verified: deleting the wait fails the named tests. Tests: 2,321 passed across tests/shared, tests/lambdas and tests/property. --- .../apis/shared/kb_backend/provisioning.py | 119 +++++++++++++++++- .../tests/shared/test_managed_kb_backend.py | 107 +++++++++++++++- 2 files changed, 223 insertions(+), 3 deletions(-) diff --git a/backend/src/apis/shared/kb_backend/provisioning.py b/backend/src/apis/shared/kb_backend/provisioning.py index 4bdda009..e9b2fba2 100644 --- a/backend/src/apis/shared/kb_backend/provisioning.py +++ b/backend/src/apis/shared/kb_backend/provisioning.py @@ -447,8 +447,90 @@ async def _call( # ── The saga ───────────────────────────────────────────────────────────────── -def _complete(item: Mapping[str, Any]) -> bool: - return bool(item.get("awsKbId")) and bool(item.get("awsDataSourceId")) +#: Statuses a knowledge base can hold while still on its way to usable. +KB_PENDING_STATUSES = ("CREATING", "UPDATING") + +#: The status the data-source create requires. +KB_ACTIVE_STATUS = "ACTIVE" + +#: Terminal-bad statuses. Waiting on these would burn the whole budget to reach +#: the same conclusion the first poll already supports. +KB_FAILED_STATUSES = ("FAILED", "DELETING", "DELETE_UNSUCCESSFUL") + +#: Ceiling on the wait. The measured range to ACTIVE is 47-124 s (n=7), so this +#: is roughly 2.5x the observed worst case — comfortably inside the worker's +#: 15-minute Lambda timeout, and short enough that a genuinely stuck creation is +#: reported within one migration step rather than silently holding a lease. +KB_ACTIVE_WAIT_SECONDS = 300.0 + +#: Poll interval. Not a tuning knob worth an env var: the operation being waited +#: on takes tens of seconds, so anything finer just adds API calls. +KB_ACTIVE_POLL_SECONDS = 5.0 + + +class KnowledgeBaseNotReady(Exception): + """A knowledge base did not reach ``ACTIVE`` within the wait budget.""" + + +async def _wait_for_knowledge_base_active( + client: Any, + aws_kb_id: str, + *, + what: str, + sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, + budget_seconds: Optional[float] = None, + interval_seconds: Optional[float] = None, +) -> str: + """Block until ``aws_kb_id`` is ``ACTIVE``, or raise. + + ``CreateKnowledgeBase`` returns while the knowledge base is still + ``CREATING``; anything that touches it before ``ACTIVE`` is refused with a + ``ConflictException`` telling you to wait. Retrying the *dependent* call is + the wrong shape — it burns attempts on a precondition rather than waiting for + it — so the precondition is waited on directly. + + Budget and interval are resolved at call time, never bound as default + arguments: a module-level default is captured at import and silently ignores a + test's override, which has already cost this feature a 33-second test. + """ + budget = KB_ACTIVE_WAIT_SECONDS if budget_seconds is None else budget_seconds + interval = KB_ACTIVE_POLL_SECONDS if interval_seconds is None else interval_seconds + + waited = 0.0 + last_status = "UNKNOWN" + while True: + described = await asyncio.to_thread( + lambda: client.get_knowledge_base(knowledgeBaseId=aws_kb_id) + ) + last_status = str( + (described.get("knowledgeBase") or {}).get("status") or "UNKNOWN" + ) + if last_status == KB_ACTIVE_STATUS: + if waited: + logger.info( + f"kb {aws_kb_id} reached {KB_ACTIVE_STATUS} after {waited:.0f}s; " + f"proceeding to {what}" + ) + return last_status + if last_status in KB_FAILED_STATUSES: + # Failing here rather than waiting out the budget: the status is + # terminal, so the only thing more waiting buys is a later report. + raise KnowledgeBaseNotReady( + f"kb {aws_kb_id} is {last_status}, which will never reach " + f"{KB_ACTIVE_STATUS}; refusing {what}" + ) + if waited >= budget: + 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." + ) + await sleep(interval) + waited += interval + + +def _complete(item: Mapping[str, Any]) -> bool: return bool(item.get("awsKbId")) and bool(item.get("awsDataSourceId")) def _resource_name(app_kb_id: str, project_prefix: Optional[str] = None) -> str: @@ -478,6 +560,11 @@ async def provision_managed_kb( environment: Optional[str] = None, max_attempts: int = MAX_PROVISION_ATTEMPTS, sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, + # How long to wait for the knowledge base to reach ACTIVE before creating its + # data source. Optional-None rather than a bound module constant so a test can + # actually override them (see _wait_for_knowledge_base_active). + budget_seconds: Optional[float] = None, + interval_seconds: Optional[float] = None, ) -> ProvisionedKnowledgeBase: """Provision, or adopt, the managed knowledge base for ``app_kb_id``. @@ -585,6 +672,33 @@ async def provision_managed_kb( ) aws_kb_id = response["knowledgeBase"]["knowledgeBaseId"] + # 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 + # (n=7). CreateDataSource against a CREATING knowledge base is refused: + # + # ConflictException: The Knowledge Base is not in a valid status. + # Wait for the knowledge base to reach a valid status and try again. + # + # `ConflictException` is deliberately absent from RETRYABLE_ERROR_CODES — a + # 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, + what="CreateDataSource", + sleep=sleep, + budget_seconds=budget_seconds, + interval_seconds=interval_seconds, + ) + aws_data_source_id = (existing or {}).get("awsDataSourceId") if not aws_data_source_id: ds_response = await _call( @@ -646,6 +760,7 @@ async def provision_managed_kb( "EMBEDDING_MODEL_ID", "EMBEDDING_MODEL_TYPE", "IMAGE_EXTRACTION_STATUS", + "KnowledgeBaseNotReady", "KNOWLEDGE_BASE_TYPE", "ProvisionedKnowledgeBase", "ProvisioningError", diff --git a/backend/tests/shared/test_managed_kb_backend.py b/backend/tests/shared/test_managed_kb_backend.py index db6ac341..30c8b69c 100644 --- a/backend/tests/shared/test_managed_kb_backend.py +++ b/backend/tests/shared/test_managed_kb_backend.py @@ -74,9 +74,12 @@ def __init__( *, on_create=None, create_failures: Optional[List[Exception]] = None, + status_sequence: Optional[List[str]] = None, ) -> 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._status_sequence = list(status_sequence or ["ACTIVE"]) self.ingest_calls: List[Dict[str, Any]] = [] self.delete_calls: List[Dict[str, Any]] = [] self.start_ingestion_job_calls: List[Dict[str, Any]] = [] @@ -105,7 +108,26 @@ def create_knowledge_base(self, **kwargs): if token not in self._by_token: self._counter += 1 self._by_token[token] = f"KB{self._counter:08d}" - return {"knowledgeBase": {"knowledgeBaseId": self._by_token[token], "status": "ACTIVE"}} + # 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"}} + + def get_knowledge_base(self, **kwargs): + """Status poll. Yields each queued status once, then settles on the last. + + Default is a single ACTIVE so tests that do not care about the wait are + unaffected; `status_sequence` lets one drive CREATING -> ACTIVE or a + terminal failure. + """ + self.thread_idents.append(threading.get_ident()) + self.get_kb_calls.append(kwargs) + if len(self._status_sequence) > 1: + status = self._status_sequence.pop(0) + else: + status = self._status_sequence[0] + return {"knowledgeBase": {"knowledgeBaseId": kwargs["knowledgeBaseId"], "status": status}} def create_data_source(self, **kwargs): self.thread_idents.append(threading.get_ident()) @@ -747,6 +769,89 @@ async def test_a_crash_after_the_data_source_still_converges(self, table): # =========================================================================== +class TestWaitsForActiveBeforeTheDataSource: + """The defect that orphaned the first real knowledge base in dev. + + `CreateKnowledgeBase` returns while the knowledge base is `CREATING` — this + module's header records 47-124 s to `ACTIVE` (n=7) — and `CreateDataSource` + against a creating knowledge base is refused: + + ConflictException: The Knowledge Base is not in a valid status. + + The create succeeded, the data source did not, `attach_aws_ids` never ran, and + the knowledge base was left in AWS with nothing pointing at it. + + Nothing caught it because the fake returned `ACTIVE` from `create_knowledge_base`, + which the real API never does. The fake now returns `CREATING`, so these tests + exercise the wait rather than skipping past it. + """ + + @pytest.mark.asyncio + async def test_the_data_source_is_created_only_after_active(self, table): + client = FakeBedrockAgent(status_sequence=["CREATING", "CREATING", "ACTIVE"]) + await _provision(client) + + assert client.create_kb_calls, "no knowledge base was created" + assert client.create_ds_calls, "no data source was created" + # Polled until ACTIVE rather than charging ahead. + assert len(client.get_kb_calls) == 3, ( + f"expected three status polls, saw {len(client.get_kb_calls)}" + ) + + @pytest.mark.asyncio + async def test_no_data_source_while_the_knowledge_base_is_creating(self, table): + """The precondition is waited on, not retried through.""" + client = FakeBedrockAgent(status_sequence=["CREATING"]) + with pytest.raises(p.KnowledgeBaseNotReady, match="still CREATING"): + await _provision(client, budget_seconds=10.0, interval_seconds=5.0) + + assert client.create_kb_calls, "the knowledge base should still be created" + assert client.create_ds_calls == [], ( + "CreateDataSource must not be attempted against a CREATING knowledge " + "base — that is the ConflictException this wait exists to prevent" + ) + + @pytest.mark.asyncio + async def test_a_terminal_status_fails_immediately(self, table): + """FAILED will never become ACTIVE, so waiting only delays the report.""" + client = FakeBedrockAgent(status_sequence=["FAILED"]) + with pytest.raises(p.KnowledgeBaseNotReady, match="FAILED"): + await _provision(client, budget_seconds=300.0, interval_seconds=5.0) + + assert len(client.get_kb_calls) == 1, ( + "a terminal status should be acted on after one poll, not waited out" + ) + assert client.create_ds_calls == [] + + @pytest.mark.asyncio + async def test_the_budget_is_read_at_call_time(self, table): + """Bound as a default argument the budget would be unpatchable. + + Asserted by giving two calls different budgets on the same import. + """ + slow = FakeBedrockAgent(status_sequence=["CREATING"]) + with pytest.raises(p.KnowledgeBaseNotReady): + await _provision(slow, budget_seconds=5.0, interval_seconds=5.0) + few = len(slow.get_kb_calls) + + slower = FakeBedrockAgent(status_sequence=["CREATING"]) + with pytest.raises(p.KnowledgeBaseNotReady): + await _provision(slower, budget_seconds=25.0, interval_seconds=5.0) + + assert len(slower.get_kb_calls) > few, ( + "a larger budget polled no more than a smaller one, so the value is " + "not being read at call time" + ) + + @pytest.mark.asyncio + async def test_the_failure_message_says_the_retry_is_safe(self, table): + """An operator reading this needs to know a retry will not duplicate.""" + 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) + + class TestOffEventLoop: @pytest.mark.asyncio async def test_create_knowledge_base_runs_in_a_worker_thread(self, table):