feat(mt): Writer Hooks for tenant work gating (2/3) (#10246)

backend/ee/onyx/background/celery/tasks/doc_permission_syncing/tasks.py

@@ -80,6 +80,7 @@ from onyx.redis.redis_connector_doc_perm_sync import RedisConnectorPermissionSyn
 from onyx.redis.redis_pool import get_redis_client
 from onyx.redis.redis_pool import get_redis_replica_client
 from onyx.redis.redis_pool import redis_lock_dump
+from onyx.redis.redis_tenant_work_gating import maybe_mark_tenant_active
 from onyx.server.runtime.onyx_runtime import OnyxRuntime
 from onyx.server.utils import make_short_id
 from onyx.utils.logger import doc_permission_sync_ctx
@@ -208,6 +209,11 @@ def check_for_doc_permissions_sync(self: Task, *, tenant_id: str) -> bool | None
                 if _is_external_doc_permissions_sync_due(cc_pair):
                     cc_pair_ids_to_sync.append(cc_pair.id)
 
+        # Tenant-work-gating hook: refresh this tenant's active-set membership
+        # whenever doc-permission sync has any due cc_pairs to dispatch.
+        if cc_pair_ids_to_sync:
+            maybe_mark_tenant_active(tenant_id)
+
         lock_beat.reacquire()
         for cc_pair_id in cc_pair_ids_to_sync:
             payload_id = try_creating_permissions_sync_task(

backend/ee/onyx/background/celery/tasks/external_group_syncing/tasks.py

@@ -69,6 +69,7 @@ from onyx.redis.redis_connector_ext_group_sync import (
 )
 from onyx.redis.redis_pool import get_redis_client
 from onyx.redis.redis_pool import get_redis_replica_client
+from onyx.redis.redis_tenant_work_gating import maybe_mark_tenant_active
 from onyx.server.runtime.onyx_runtime import OnyxRuntime
 from onyx.server.utils import make_short_id
 from onyx.utils.logger import format_error_for_logging
@@ -202,6 +203,11 @@ def check_for_external_group_sync(self: Task, *, tenant_id: str) -> bool | None:
                 if _is_external_group_sync_due(cc_pair):
                     cc_pair_ids_to_sync.append(cc_pair.id)
 
+        # Tenant-work-gating hook: refresh this tenant's active-set membership
+        # whenever external-group sync has any due cc_pairs to dispatch.
+        if cc_pair_ids_to_sync:
+            maybe_mark_tenant_active(tenant_id)
+
         lock_beat.reacquire()
         for cc_pair_id in cc_pair_ids_to_sync:
             payload_id = try_creating_external_group_sync_task(

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

@@ -59,6 +59,7 @@ from onyx.redis.redis_connector_delete import RedisConnectorDelete
 from onyx.redis.redis_connector_delete import RedisConnectorDeletePayload
 from onyx.redis.redis_pool import get_redis_client
 from onyx.redis.redis_pool import get_redis_replica_client
+from onyx.redis.redis_tenant_work_gating import maybe_mark_tenant_active
 from onyx.server.metrics.deletion_metrics import inc_deletion_blocked
 from onyx.server.metrics.deletion_metrics import inc_deletion_completed
 from onyx.server.metrics.deletion_metrics import inc_deletion_fence_reset
@@ -172,6 +173,11 @@ def check_for_connector_deletion_task(self: Task, *, tenant_id: str) -> bool | N
             for cc_pair in cc_pairs:
                 cc_pair_ids.append(cc_pair.id)
 
+        # Tenant-work-gating hook: any cc_pair means deletion could have
+        # cleanup work to do for this tenant on some cycle.
+        if cc_pair_ids:
+            maybe_mark_tenant_active(tenant_id)
+
         # try running cleanup on the cc_pair_ids
         for cc_pair_id in cc_pair_ids:
             with get_session_with_current_tenant() as db_session:

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

@@ -108,6 +108,7 @@ from onyx.redis.redis_pool import get_redis_client
 from onyx.redis.redis_pool import get_redis_replica_client
 from onyx.redis.redis_pool import redis_lock_dump
 from onyx.redis.redis_pool import SCAN_ITER_COUNT_DEFAULT
+from onyx.redis.redis_tenant_work_gating import maybe_mark_tenant_active
 from onyx.redis.redis_utils import is_fence
 from onyx.server.metrics.connector_health_metrics import on_connector_error_state_change
 from onyx.server.metrics.connector_health_metrics import on_connector_indexing_success
@@ -896,6 +897,11 @@ def check_for_indexing(self: Task, *, tenant_id: str) -> int | None:
 
                 secondary_cc_pair_ids = standard_cc_pair_ids
 
+        # Tenant-work-gating hook: refresh this tenant's active-set membership
+        # whenever indexing actually has work to dispatch.
+        if primary_cc_pair_ids or secondary_cc_pair_ids:
+            maybe_mark_tenant_active(tenant_id)
+
         # Flag CC pairs in repeated error state for primary/current search settings
         with get_session_with_current_tenant() as db_session:
             for cc_pair_id in primary_cc_pair_ids:

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

@@ -72,6 +72,7 @@ from onyx.redis.redis_hierarchy import get_source_node_id_from_cache
 from onyx.redis.redis_hierarchy import HierarchyNodeCacheEntry
 from onyx.redis.redis_pool import get_redis_client
 from onyx.redis.redis_pool import get_redis_replica_client
+from onyx.redis.redis_tenant_work_gating import maybe_mark_tenant_active
 from onyx.server.metrics.pruning_metrics import observe_pruning_diff_duration
 from onyx.server.runtime.onyx_runtime import OnyxRuntime
 from onyx.server.utils import make_short_id
@@ -228,6 +229,11 @@ def check_for_pruning(self: Task, *, tenant_id: str) -> bool | None:
                 for cc_pair_entry in cc_pairs:
                     cc_pair_ids.append(cc_pair_entry.id)
 
+            # Tenant-work-gating hook: any cc_pair means pruning could have
+            # work to do for this tenant on some cycle.
+            if cc_pair_ids:
+                maybe_mark_tenant_active(tenant_id)
+
             for cc_pair_id in cc_pair_ids:
                 lock_beat.reacquire()
                 with get_session_with_current_tenant() as db_session:

backend/onyx/background/celery/tasks/vespa/document_sync.py

@@ -15,6 +15,7 @@ from onyx.configs.constants import OnyxCeleryTask
 from onyx.configs.constants import OnyxRedisConstants
 from onyx.db.document import construct_document_id_select_by_needs_sync
 from onyx.db.document import count_documents_by_needs_sync
+from onyx.redis.redis_tenant_work_gating import maybe_mark_tenant_active
 from onyx.utils.logger import setup_logger
 
 # Redis keys for document sync tracking
@@ -150,6 +151,10 @@ def try_generate_stale_document_sync_tasks(
         logger.info("No stale documents found. Skipping sync tasks generation.")
         return None
 
+    # Tenant-work-gating hook: refresh this tenant's active-set membership
+    # whenever vespa sync actually has stale docs to dispatch.
+    maybe_mark_tenant_active(tenant_id)
+
     logger.info(
         f"Stale documents found (at least {stale_doc_count}). Generating sync tasks in one batch."
     )

backend/onyx/redis/redis_tenant_work_gating.py

@@ -39,10 +39,9 @@ def mark_tenant_active(tenant_id: str) -> None:
     current timestamp as the score). Best-effort — a Redis failure is logged
     and swallowed so it never breaks a writer path.
 
-    Call sites:
-    - Top of each gated beat-task consumer when its "is there work?" query
-      returns a non-empty result.
-    - cc_pair create lifecycle hook.
+    Raw write; does not check the feature flag. Writer call sites should
+    use `maybe_mark_tenant_active` instead so the feature flag gates the
+    ZADD.
     """
     if not MULTI_TENANT:
         return
@@ -55,6 +54,23 @@ def mark_tenant_active(tenant_id: str) -> None:
         logger.exception(f"mark_tenant_active failed: tenant_id={tenant_id}")
 
 
+def maybe_mark_tenant_active(tenant_id: str) -> None:
+    """Convenience wrapper for writer call sites: records the tenant only
+    when the feature flag is on. Fully defensive — never raises, so a Redis
+    outage or flag-read failure can't abort the calling task."""
+    try:
+        # Local import to avoid a module-load cycle: OnyxRuntime imports
+        # onyx.redis.redis_pool, so a top-level import here would wedge on
+        # certain startup paths.
+        from onyx.server.runtime.onyx_runtime import OnyxRuntime
+
+        if not OnyxRuntime.get_tenant_work_gating_enabled():
+            return
+        mark_tenant_active(tenant_id)
+    except Exception:
+        logger.exception(f"maybe_mark_tenant_active failed: tenant_id={tenant_id}")
+
+
 def get_active_tenants(ttl_seconds: int) -> set[str] | None:
     """Return tenants whose last-seen timestamp is within `ttl_seconds` of
     now.

backend/onyx/server/documents/cc_pair.py

@@ -60,6 +60,7 @@ from onyx.db.permission_sync_attempt import (
 from onyx.redis.redis_connector import RedisConnector
 from onyx.redis.redis_connector_utils import get_deletion_attempt_snapshot
 from onyx.redis.redis_pool import get_redis_client
+from onyx.redis.redis_tenant_work_gating import maybe_mark_tenant_active
 from onyx.server.documents.models import CCPairFullInfo
 from onyx.server.documents.models import CCPropertyUpdateRequest
 from onyx.server.documents.models import CCStatusUpdateRequest
@@ -581,6 +582,10 @@ def associate_credential_to_connector(
             processing_mode=metadata.processing_mode,
         )
 
+        # Tenant-work-gating lifecycle hook: keep new-tenant latency to
+        # seconds instead of one full-fanout interval.
+        maybe_mark_tenant_active(tenant_id)
+
         # trigger indexing immediately
         client_app.send_task(
             OnyxCeleryTask.CHECK_FOR_INDEXING,

backend/onyx/server/documents/connector.py

@@ -125,6 +125,7 @@ from onyx.file_store.file_store import FileStore
 from onyx.file_store.file_store import get_default_file_store
 from onyx.key_value_store.interface import KvKeyNotFoundError
 from onyx.redis.redis_pool import get_redis_client
+from onyx.redis.redis_tenant_work_gating import maybe_mark_tenant_active
 from onyx.server.documents.models import AuthStatus
 from onyx.server.documents.models import AuthUrl
 from onyx.server.documents.models import ConnectorBase
@@ -1640,6 +1641,10 @@ def create_connector_with_mock_credential(
             groups=connector_data.groups,
         )
 
+        # Tenant-work-gating lifecycle hook: keep new-tenant latency to
+        # seconds instead of one full-fanout interval.
+        maybe_mark_tenant_active(tenant_id)
+
         # trigger indexing immediately
         client_app.send_task(
             OnyxCeleryTask.CHECK_FOR_INDEXING,

backend/onyx/server/features/build/sandbox/tasks/tasks.py

@@ -19,6 +19,7 @@ from onyx.configs.constants import OnyxRedisLocks
 from onyx.db.engine.sql_engine import get_session_with_current_tenant
 from onyx.db.enums import SandboxStatus
 from onyx.redis.redis_pool import get_redis_client
+from onyx.redis.redis_tenant_work_gating import maybe_mark_tenant_active
 from onyx.server.features.build.configs import SANDBOX_BACKEND
 from onyx.server.features.build.configs import SANDBOX_IDLE_TIMEOUT_SECONDS
 from onyx.server.features.build.configs import SandboxBackend
@@ -110,6 +111,10 @@ def cleanup_idle_sandboxes_task(self: Task, *, tenant_id: str) -> None:  # noqa:
                 task_logger.debug("No idle sandboxes found")
                 return
 
+            # Tenant-work-gating hook: refresh this tenant's active-set
+            # membership whenever sandbox cleanup has work to do.
+            maybe_mark_tenant_active(tenant_id)
+
             task_logger.info(
                 f"Found {len(idle_sandboxes)} idle sandboxes to put to sleep"
             )

backend/tests/external_dependency_unit/tenant_work_gating/test_tenant_work_gating.py

@@ -19,6 +19,7 @@ from onyx.redis.redis_tenant_work_gating import _SET_KEY
 from onyx.redis.redis_tenant_work_gating import cleanup_expired
 from onyx.redis.redis_tenant_work_gating import get_active_tenants
 from onyx.redis.redis_tenant_work_gating import mark_tenant_active
+from onyx.redis.redis_tenant_work_gating import maybe_mark_tenant_active
 
 
 @pytest.fixture(autouse=True)
@@ -157,3 +158,26 @@ def test_rendered_key_is_cloud_prefixed() -> None:
     raw = RedisPool().get_raw_client()
     assert raw.zscore("cloud:active_tenants", "tenant_a") is not None
     assert raw.zscore("active_tenants", "tenant_a") is None
+
+
+def test_maybe_mark_is_noop_when_gating_disabled() -> None:
+    """Writer-side API: when the feature flag is off, the call must not
+    write to Redis so deploys are inert."""
+    with patch(
+        "onyx.server.runtime.onyx_runtime.OnyxRuntime.get_tenant_work_gating_enabled",
+        return_value=False,
+    ):
+        maybe_mark_tenant_active("tenant_a")
+
+    assert get_active_tenants(ttl_seconds=60) == set()
+
+
+def test_maybe_mark_writes_when_gating_enabled() -> None:
+    """Writer-side API: when the feature flag is on, the call must write."""
+    with patch(
+        "onyx.server.runtime.onyx_runtime.OnyxRuntime.get_tenant_work_gating_enabled",
+        return_value=True,
+    ):
+        maybe_mark_tenant_active("tenant_a")
+
+    assert get_active_tenants(ttl_seconds=60) == {"tenant_a"}