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):