fix(kb): wait for the knowledge base to be ACTIVE before its data source

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",

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