fix(connectors): enable Create once required checks pass, without re-running checks (#15883)

.agents/feature-map/components/connectors.md

@@ -462,7 +462,10 @@ within the blocking budget blocks the pairing. Indexing and perm-sync attempts k
 the legacy validation. Full runs execute as the `RUN_CAPABILITY_CHECKS` Celery task
 on the `capability_checks` queue. A beat task, `CHECK_FOR_STALE_CAPABILITY_RUNS`,
 retires dead runs. Draft runs on an unsaved form (`/admin/connector-checks/runs`)
-reuse their result at creation when the form is unchanged. A newer run with the
+reuse their result at creation when the form is unchanged. A check the draft
+run is still running is not started again: the draft task marks its pending
+results in flight, creation waits for them within its blocking budget, and the
+background run of unfinished checks waits for them too. A newer run with the
 same draft key supersedes an older one, and the cancel endpoint does the same
 without a successor: the task stops before its next check. The plan endpoint
 (`/admin/connector-checks/plan`) lists the same checks without a credential or a

backend/onyx/background/celery/tasks/capability_checks/tasks.py

@@ -18,18 +18,26 @@ from uuid import UUID
 from celery import Task, shared_task
 
 from onyx.background.celery.apps.app_base import task_logger
-from onyx.configs.constants import OnyxCeleryTask
+from onyx.configs.constants import DocumentSource, OnyxCeleryTask
 from onyx.connectors.capability_checks.draft_runs import (
     DRAFT_CHECK_TIMEOUT_SECONDS,
+    DRAFT_RUN_TTL_SECONDS,
     DraftCheckStateKind,
     DraftRunPairScope,
     DraftRunStatus,
     apply_check_result,
     cache_draft_result,
+    cached_check_result,
+    clear_draft_result_in_flight,
+    form_draft_result_keys,
+    is_draft_result_in_flight,
     is_superseded,
     load_draft_run,
+    mark_draft_results_in_flight,
     save_draft_run,
+    wait_for_draft_result,
 )
+from onyx.connectors.capability_checks.form_state import validate_form_state
 from onyx.connectors.capability_checks.models import CapabilityCheckResult
 from onyx.connectors.capability_checks.registry import get_capability_checks
 from onyx.connectors.capability_checks.runner import (
@@ -41,6 +49,7 @@ from onyx.connectors.capability_checks.runner import (
 )
 from onyx.connectors.config_hash import compute_connector_config_hash
 from onyx.connectors.models import InputType
+from onyx.connectors.registry import CONNECTOR_CLASS_MAP
 from onyx.db.connector import fetch_connector_by_id
 from onyx.db.connector_credential_pair import get_connector_credential_pair
 from onyx.db.credential_capability import (
@@ -55,6 +64,50 @@ from onyx.db.credentials import (
 )
 from onyx.db.engine.sql_engine import get_session_with_current_tenant
 from onyx.db.enums import AccessType, CapabilityCheckTrigger
+from onyx.db.models import Credential
+
+
+def _await_draft_results(
+    check_ids: list[str],
+    *,
+    credential: Credential,
+    source: DocumentSource,
+    access_type: AccessType | None,
+    connector_specific_config: dict[str, Any] | None,
+) -> tuple[list[CapabilityCheckResult], list[str]]:
+    """The results a create-form draft run of the same form is still to write
+    for ``check_ids``, awaited up to each check's draft timeout, and the ids
+    of the checks that still need to run."""
+    checks = get_capability_checks(source)
+    keys = form_draft_result_keys(
+        [check for check in checks if check.check_id in check_ids],
+        credential_id=credential.id,
+        credential_updated_at=credential.time_updated,
+        source=source,
+        access_type=access_type,
+        form_values=validate_form_state(
+            CONNECTOR_CLASS_MAP[source].config_class, connector_specific_config or {}
+        ).values,
+        cc_pair_id=None,
+    )
+    awaited: list[CapabilityCheckResult] = []
+    to_run: list[str] = []
+    for check_id in check_ids:
+        key = keys.get(check_id)
+        cached = (
+            wait_for_draft_result(key, DRAFT_CHECK_TIMEOUT_SECONDS)
+            if key is not None and is_draft_result_in_flight(key)
+            else None
+        )
+        if cached is None:
+            to_run.append(check_id)
+            continue
+        awaited.extend(
+            cached_check_result(check, cached)
+            for check in checks
+            if check.check_id == check_id
+        )
+    return awaited, to_run
 
 
 @shared_task(  # ty: ignore[invalid-argument-type]
@@ -132,6 +185,18 @@ def run_capability_checks_task(
                 )
                 if cc_pair is not None and parsed_access_type is None:
                     parsed_access_type = cc_pair.access_type
+        # A run of the checks creation left unfinished first waits for the
+        # ones a draft run of the same form is still running.
+        awaited_results: list[CapabilityCheckResult] = []
+        run_check_ids = check_ids
+        if check_ids is not None:
+            awaited_results, run_check_ids = _await_draft_results(
+                check_ids,
+                credential=credential,
+                source=source,
+                access_type=parsed_access_type,
+                connector_specific_config=config,
+            )
         report = generate_capability_report(
             credential,
             source=source,
@@ -140,16 +205,17 @@ def run_capability_checks_task(
             input_type=input_type,
             trigger=parsed_trigger,
             access_type=parsed_access_type,
-            check_ids=frozenset(check_ids) if check_ids is not None else None,
+            check_ids=frozenset(run_check_ids) if run_check_ids is not None else None,
         )
-        if prior_results:
-            report = merge_capability_results(
-                report,
-                [
-                    CapabilityCheckResult.model_validate(result)
-                    for result in prior_results
-                ],
-            )
+        merged_prior: list[CapabilityCheckResult] = [
+            *awaited_results,
+            *(
+                CapabilityCheckResult.model_validate(result)
+                for result in prior_results or []
+            ),
+        ]
+        if merged_prior:
+            report = merge_capability_results(report, merged_prior)
         with get_session_with_current_tenant() as db_session:
             completed_row = upsert_completed_capability_report(
                 db_session,
@@ -216,6 +282,9 @@ def run_draft_capability_checks_task(
     pending = [
         check for check in snapshot.checks if check.state == DraftCheckStateKind.PENDING
     ]
+    pending_result_keys: list[str] = list(
+        dict.fromkeys(run.result_cache_keys[check.check_id] for check in pending)
+    )
     timeout_by_check_id = {
         check.check_id: min(
             effective_check_timeout_seconds(check), DRAFT_CHECK_TIMEOUT_SECONDS
@@ -246,7 +315,9 @@ def run_draft_capability_checks_task(
             and check.capability == result.capability
         )
         apply_check_result(check_state, result)
-        cache_draft_result(run.result_cache_keys[result.check_id], result)
+        result_key = run.result_cache_keys[result.check_id]
+        cache_draft_result(result_key, result)
+        clear_draft_result_in_flight(result_key)
         if is_superseded(run):
             raise _DraftRunSupersededError()
         mark_next_running()
@@ -265,6 +336,9 @@ def run_draft_capability_checks_task(
             snapshot.status = DraftRunStatus.FAILED_TO_RUN
             save_draft_run(run)
             return
+        # A validation of the same form waits for these results instead of
+        # running the checks again.
+        mark_draft_results_in_flight(pending_result_keys, DRAFT_RUN_TTL_SECONDS)
         mark_next_running()
         save_draft_run(run)
         pair_scope: DraftRunPairScope | None = run.pair_scope
@@ -292,6 +366,10 @@ def run_draft_capability_checks_task(
         snapshot.status = DraftRunStatus.FAILED_TO_RUN
         save_draft_run(run)
         raise
+    finally:
+        # A check the run did not finish is not coming.
+        for result_key in pending_result_keys:
+            clear_draft_result_in_flight(result_key)
 
 
 @shared_task(  # ty: ignore[invalid-argument-type]

backend/onyx/connectors/capability_checks/creation.py

@@ -25,8 +25,10 @@ from onyx.connectors.capability_checks.draft_runs import (
     CachedDraftResult,
     DraftCheckStateKind,
     cached_check_result,
-    draft_result_cache_key,
+    form_draft_result_keys,
     get_cached_draft_result,
+    is_draft_result_in_flight,
+    wait_for_draft_result,
 )
 from onyx.connectors.capability_checks.form_state import validate_form_state
 from onyx.connectors.capability_checks.models import (
@@ -66,6 +68,29 @@ CREATION_BLOCKING_BUDGET_SECONDS = 3.0
 _REJECTED_MESSAGE = "Did not finish before the connector was rejected."
 
 
+def _draft_result_keys(
+    checks: list[CapabilityCheck[Any]],
+    *,
+    credential: Credential,
+    source: DocumentSource,
+    access_type: AccessType,
+    connector_specific_config: dict[str, Any],
+    cc_pair_id: int | None,
+) -> dict[str, str]:
+    """check_id to the draft result key a draft run of this form uses."""
+    return form_draft_result_keys(
+        checks,
+        credential_id=credential.id,
+        credential_updated_at=credential.time_updated,
+        source=source,
+        access_type=access_type,
+        form_values=validate_form_state(
+            CONNECTOR_CLASS_MAP[source].config_class, connector_specific_config
+        ).values,
+        cc_pair_id=cc_pair_id,
+    )
+
+
 def _cached_draft_results(
     checks: list[CapabilityCheck[Any]],
     *,
@@ -78,24 +103,19 @@ def _cached_draft_results(
     """check_id to its cached draft result, for the checks a draft run already
     ran with this credential, access type and form values. With ``cc_pair_id``,
     only that pair's dry runs count; without it, only create-form runs."""
-    form_values = validate_form_state(
-        CONNECTOR_CLASS_MAP[source].config_class, connector_specific_config
-    ).values
+    keys = _draft_result_keys(
+        checks,
+        credential=credential,
+        source=source,
+        access_type=access_type,
+        connector_specific_config=connector_specific_config,
+        cc_pair_id=cc_pair_id,
+    )
     cached_by_check_id: dict[str, CachedDraftResult] = {}
-    for check in checks:
-        cached = get_cached_draft_result(
-            draft_result_cache_key(
-                credential_id=credential.id,
-                credential_updated_at=credential.time_updated,
-                source=source,
-                access_type=access_type,
-                check=check,
-                form_values=form_values,
-                cc_pair_id=cc_pair_id,
-            )
-        )
+    for check_id, key in keys.items():
+        cached = get_cached_draft_result(key)
         if cached is not None:
-            cached_by_check_id[check.check_id] = cached
+            cached_by_check_id[check_id] = cached
     return cached_by_check_id
 
 
@@ -197,9 +217,27 @@ def _run_check_group(
     ).check_results
 
 
+def _await_draft_check(
+    checks: list[CapabilityCheck[Any]], *, check_id: str, result_key: str
+) -> list[CapabilityCheckResult] | None:
+    """The result a draft run of the same form is still to write for the
+    check, with a row per capability it is mirrored on; ``None`` when it does
+    not land within the blocking budget."""
+    cached = wait_for_draft_result(result_key, CREATION_BLOCKING_BUDGET_SECONDS)
+    if cached is None:
+        return None
+    return [
+        cached_check_result(check, cached)
+        for check in checks
+        if check.check_id == check_id
+    ]
+
+
 def _run_checks_within_budget(
     check_ids: list[str],
     *,
+    awaited: dict[str, str],
+    checks: list[CapabilityCheck[Any]],
     credential_id: int,
     source: DocumentSource,
     connector_specific_config: dict[str, Any],
@@ -217,12 +255,15 @@ def _run_checks_within_budget(
     discarded. The threads read a credential loaded on a session that is
     already closed, so they share no session with the request.
     """
-    if not check_ids:
+    if not check_ids and not awaited:
         return [], frozenset()
     with get_session_with_current_tenant() as db_session:
         credential = fetch_credential_by_id(credential_id, db_session)
         if credential is None:
             raise RuntimeError(f"Credential {credential_id} was deleted.")
+    # A check a draft run of this form is still running is awaited, not run
+    # again: its result lands in the draft cache.
+    awaited_ids: list[str] = list(awaited)
     # TODO(evan-onyx): concurrent creation checks can add load on rate-limited
     # sources; they may need a concurrency cap.
     groups = run_functions_tuples_in_parallel(
@@ -241,13 +282,25 @@ def _run_checks_within_budget(
                 (),
             )
             for check_id in check_ids
+        ]
+        + [
+            (
+                partial(
+                    _await_draft_check,
+                    checks,
+                    check_id=check_id,
+                    result_key=awaited[check_id],
+                ),
+                (),
+            )
+            for check_id in awaited_ids
         ],
         timeout=CREATION_BLOCKING_BUDGET_SECONDS,
         timeout_callback=lambda *_: None,
     )
     finished: list[CapabilityCheckResult] = []
     unfinished: set[str] = set()
-    for check_id, group in zip(check_ids, groups, strict=True):
+    for check_id, group in zip([*check_ids, *awaited_ids], groups, strict=True):
         if group is None:
             unfinished.add(check_id)
         else:
@@ -400,12 +453,28 @@ def run_named_checks_within_budget(
         connector_specific_config=connector_specific_config,
         cc_pair_id=cc_pair_id,
     )
+    to_check: list[str] = list(
+        dict.fromkeys(
+            check.check_id for check in checks if check.check_id not in reused
+        )
+    )
+    keys = _draft_result_keys(
+        checks,
+        credential=credential,
+        source=source,
+        access_type=access_type,
+        connector_specific_config=connector_specific_config,
+        cc_pair_id=cc_pair_id,
+    )
+    awaited: dict[str, str] = {
+        check_id: keys[check_id]
+        for check_id in to_check
+        if is_draft_result_in_flight(keys[check_id])
+    }
     finished, unfinished = _run_checks_within_budget(
-        list(
-            dict.fromkeys(
-                check.check_id for check in checks if check.check_id not in reused
-            )
-        ),
+        [check_id for check_id in to_check if check_id not in awaited],
+        awaited=awaited,
+        checks=checks,
         credential_id=credential.id,
         source=source,
         connector_specific_config=connector_specific_config,

backend/onyx/connectors/capability_checks/draft_runs.py

@@ -12,6 +12,8 @@ from other pairs.
 """
 
 import math
+import time
+from collections.abc import Iterable
 from datetime import datetime, timedelta, timezone
 from enum import Enum
 from typing import Any
@@ -56,6 +58,9 @@ DRAFT_RESULT_CACHE_TTL_SECONDS = 10 * 60
 _DRAFT_RUN_KEY_PREFIX = "capability_check_draft_run"
 _DRAFT_LATEST_RUN_KEY_PREFIX = "capability_check_draft_latest_run"
 _DRAFT_RESULT_KEY_PREFIX = "capability_check_draft_result"
+# Set on a result key while a draft run is still to write it.
+_IN_FLIGHT_SUFFIX = ":in_flight"
+_DRAFT_RESULT_POLL_SECONDS = 0.2
 _DRAFT_START_LOCK_PREFIX = "capability_check_draft_start"
 _CC_PAIR_KEY_PART = "cc_pair:"
 # The config-hash part of the result cache key for checks that never read the
@@ -299,6 +304,71 @@ def cache_draft_result(key: str, result: CapabilityCheckResult) -> None:
     )
 
 
+def form_draft_result_keys(
+    checks: Iterable[CapabilityCheck[Any]],
+    *,
+    credential_id: int,
+    credential_updated_at: datetime,
+    source: DocumentSource,
+    access_type: AccessType | None,
+    form_values: dict[str, Any] | None,
+    cc_pair_id: int | None,
+) -> dict[str, str]:
+    """check_id to its result cache key, for a form's draft results."""
+    return {
+        check.check_id: draft_result_cache_key(
+            credential_id=credential_id,
+            credential_updated_at=credential_updated_at,
+            source=source,
+            access_type=access_type,
+            check=check,
+            form_values=form_values,
+            cc_pair_id=cc_pair_id,
+        )
+        for check in checks
+    }
+
+
+def _in_flight_key(result_key: str) -> str:
+    return f"{result_key}{_IN_FLIGHT_SUFFIX}"
+
+
+def mark_draft_results_in_flight(result_keys: Iterable[str], ttl_seconds: int) -> None:
+    """Marks results that a draft run is still to write, so a validation of
+    the same form waits for them instead of running the checks again."""
+    cache = get_cache_backend()
+    for result_key in result_keys:
+        cache.set(_in_flight_key(result_key), "1", ex=ttl_seconds)
+
+
+def clear_draft_result_in_flight(result_key: str) -> None:
+    get_cache_backend().delete(_in_flight_key(result_key))
+
+
+def is_draft_result_in_flight(result_key: str) -> bool:
+    return get_cache_backend().get(_in_flight_key(result_key)) is not None
+
+
+def wait_for_draft_result(
+    result_key: str, timeout_seconds: float
+) -> CachedDraftResult | None:
+    """The result a draft run writes under ``result_key``, once it lands.
+    ``None`` when it does not land within ``timeout_seconds``, or the run
+    stops without it (an INDETERMINATE result is not cached)."""
+    deadline = time.monotonic() + timeout_seconds
+    while True:
+        cached = get_cached_draft_result(result_key)
+        if cached is not None:
+            return cached
+        if not is_draft_result_in_flight(result_key):
+            # The run may have written the result just before clearing.
+            return get_cached_draft_result(result_key)
+        remaining = deadline - time.monotonic()
+        if remaining <= 0:
+            return None
+        time.sleep(min(_DRAFT_RESULT_POLL_SECONDS, remaining))
+
+
 def cached_check_result(
     check: CapabilityCheck[Any], cached: CachedDraftResult
 ) -> CapabilityCheckResult:

backend/tests/external_dependency_unit/connectors/capability_checks/test_named_checks_at_creation.py

@@ -414,3 +414,80 @@ def test_failed_enqueue_marks_the_run_failed_to_run(
     assert (
         _run_status(db_session, slack_pair) == CapabilityReportRunStatus.FAILED_TO_RUN
     )
+
+
+def _start_draft_run_in_background(
+    harness: _Harness, slack_pair: ConnectorCredentialPair
+) -> threading.Thread:
+    """Starts a create-form draft run of the slack form and runs its task in a
+    thread, as a worker would."""
+    start_draft_capability_check_run(
+        user_id=uuid4(),
+        credential=slack_pair.credential,
+        source=DocumentSource.SLACK,
+        config_class=SlackConnectorConfig,
+        access_type=AccessType.PUBLIC,
+        draft_key=uuid4().hex,
+        form_values={"channels": ["a"]},
+    )
+    task_kwargs = harness.send_task.call_args.kwargs["kwargs"]
+    harness.send_task.reset_mock()
+    worker = threading.Thread(
+        target=run_draft_capability_checks_task, kwargs=task_kwargs, daemon=True
+    )
+    worker.start()
+    return worker
+
+
+def _wait_until_started(harness: _Harness, check_id: str) -> None:
+    deadline = time.monotonic() + 10
+    while check_id not in harness.runs:
+        assert time.monotonic() < deadline, f"{check_id} never started"
+        time.sleep(0.05)
+
+
+@pytest.mark.usefixtures("tenant_context")
+def test_creation_awaits_a_check_the_draft_run_is_still_running(
+    db_session: Session, harness: _Harness, slack_pair: ConnectorCredentialPair
+) -> None:
+    release = harness.make_slow(_CHANNELS)
+    worker = _start_draft_run_in_background(harness, slack_pair)
+    _wait_until_started(harness, _CHANNELS)
+    # The draft finishes the check while creation waits for it.
+    threading.Timer(0.3, release.set).start()
+
+    assert _validate(db_session, slack_pair, {"channels": ["a"]}) is True
+    worker.join(timeout=10)
+
+    # Creation used the draft's result instead of running the check again.
+    assert harness.runs.count(_CHANNELS) == 1
+    harness.send_task.assert_not_called()
+    assert _stored_statuses(db_session, slack_pair)[_CHANNELS] == "passed"
+
+
+@pytest.mark.usefixtures("tenant_context")
+def test_background_run_awaits_a_draft_check_that_outlasts_the_budget(
+    db_session: Session, harness: _Harness, slack_pair: ConnectorCredentialPair
+) -> None:
+    release = harness.make_slow(_CHANNELS)
+    worker = _start_draft_run_in_background(harness, slack_pair)
+    _wait_until_started(harness, _CHANNELS)
+
+    # The draft is still running the check when the budget ends.
+    assert _validate(db_session, slack_pair, {"channels": ["a"]}) is True
+    assert _run_status(db_session, slack_pair) == CapabilityReportRunStatus.RUNNING
+    task_kwargs = harness.send_task.call_args.kwargs["kwargs"]
+    assert task_kwargs["check_ids"] == [_CHANNELS]
+
+    # The background run waits for the draft's result instead of starting
+    # another copy.
+    threading.Timer(0.3, release.set).start()
+    run_capability_checks_task(**task_kwargs)
+    worker.join(timeout=10)
+
+    assert harness.runs.count(_CHANNELS) == 1
+    assert _stored_statuses(db_session, slack_pair) == {
+        _TOKEN: "passed",
+        _CHANNELS: "passed",
+        _PERM_SYNC: "skipped",
+    }

web/src/lib/connectors/checks/hooks.test.tsx

@@ -209,6 +209,40 @@ it("lets Create run without a run when no required check applies", async () => {
   await waitFor(() => expect(result.current.checks.createReady).toBe(true));
 });
 
+it("lets Create run once required checks pass, while optional ones still run", async () => {
+  startMock.mockResolvedValue(
+    run(
+      [
+        check({ ...BINDING, state: "passed" }),
+        check({ ...CONTENT, state: "passed" }),
+        check({ check_id: "discovery", required: false, state: "running" }),
+      ],
+      "running"
+    )
+  );
+  const { result } = renderChecks();
+  act(() => result.current.checks.begin());
+
+  await waitFor(() => expect(result.current.checks.createReady).toBe(true));
+});
+
+it("holds Create while a required check still runs", async () => {
+  startMock.mockResolvedValue(
+    run(
+      [
+        check({ ...BINDING, state: "passed" }),
+        check({ ...CONTENT, state: "running" }),
+      ],
+      "running"
+    )
+  );
+  const { result } = renderChecks();
+  act(() => result.current.checks.begin());
+
+  await waitFor(() => expect(result.current.checks.formUnlocked).toBe(true));
+  expect(result.current.checks.createReady).toBe(false);
+});
+
 it("ignores non-required failures and unverified checks for Create", async () => {
   startMock.mockResolvedValue(
     run([

web/src/lib/connectors/checks/hooks.ts

@@ -393,11 +393,15 @@ export function useConnectorChecks({
       )
     : plan !== undefined &&
       !applicable.some((check) => check.validates_binding);
+  // A current run of this form that started fine. It need not be finished:
+  // `createReady` waits only on required checks (see `blocks`), so a slow
+  // optional check such as discovery does not hold Create back.
   const ranThisForm: boolean =
     !!session &&
     !session.requesting &&
     session.ranWithForm === formKey &&
-    snapshot?.status === "completed";
+    snapshot !== null &&
+    snapshot.status !== "failed_to_run";
   const createReady: boolean =
     formUnlocked &&
     (!applicable.some((check) => check.required) ||