diff --git a/.agents/feature-map/components/connectors.md b/.agents/feature-map/components/connectors.md index 4919b1976f..50ee56fd9a 100644 --- a/.agents/feature-map/components/connectors.md +++ b/.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 diff --git a/backend/onyx/background/celery/tasks/capability_checks/tasks.py b/backend/onyx/background/celery/tasks/capability_checks/tasks.py index 132f6914a3..c6d0eee4bf 100644 --- a/backend/onyx/background/celery/tasks/capability_checks/tasks.py +++ b/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] diff --git a/backend/onyx/connectors/capability_checks/creation.py b/backend/onyx/connectors/capability_checks/creation.py index 6650fcbe7a..f9421dc326 100644 --- a/backend/onyx/connectors/capability_checks/creation.py +++ b/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, diff --git a/backend/onyx/connectors/capability_checks/draft_runs.py b/backend/onyx/connectors/capability_checks/draft_runs.py index ca93f16cf5..f3b67aa307 100644 --- a/backend/onyx/connectors/capability_checks/draft_runs.py +++ b/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: diff --git a/backend/tests/external_dependency_unit/connectors/capability_checks/test_named_checks_at_creation.py b/backend/tests/external_dependency_unit/connectors/capability_checks/test_named_checks_at_creation.py index 98ce6ddf53..6001c69bf8 100644 --- a/backend/tests/external_dependency_unit/connectors/capability_checks/test_named_checks_at_creation.py +++ b/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", + } diff --git a/web/src/lib/connectors/checks/hooks.test.tsx b/web/src/lib/connectors/checks/hooks.test.tsx index c43bdc5450..c4f80f3718 100644 --- a/web/src/lib/connectors/checks/hooks.test.tsx +++ b/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([ diff --git a/web/src/lib/connectors/checks/hooks.ts b/web/src/lib/connectors/checks/hooks.ts index 267ca60d90..573c3682d2 100644 --- a/web/src/lib/connectors/checks/hooks.ts +++ b/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) ||