Merge branch 'main' of https://github.com/onyx-dot-app/onyx into bugfix/dependency-updates

Author: Richard Kuo (Onyx)

Diff for: backend/ee/onyx/main.py

@@ -1,3 +1,6 @@
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+
 from fastapi import FastAPI
 from httpx_oauth.clients.google import GoogleOAuth2
 from httpx_oauth.clients.openid import BASE_SCOPES
@@ -44,21 +47,36 @@
 from onyx.main import get_application as get_application_base
 from onyx.main import include_auth_router_with_prefix
 from onyx.main import include_router_with_global_prefix_prepended
+from onyx.main import lifespan as lifespan_base
 from onyx.utils.logger import setup_logger
 from onyx.utils.variable_functionality import global_version
 from shared_configs.configs import MULTI_TENANT
 
 logger = setup_logger()
 
 
+@asynccontextmanager
+async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
+    """Small wrapper around the lifespan of the MIT application.
+    Basically just calls the base lifespan, and then adds EE-only
+    steps after."""
+
+    async with lifespan_base(app):
+        # seed the Onyx environment with LLMs, Assistants, etc. based on an optional
+        # environment variable. Used to automate deployment for multiple environments.
+        seed_db()
+
+        yield
+
+
 def get_application() -> FastAPI:
     # Anything that happens at import time is not guaranteed to be running ee-version
     # Anything after the server startup will be running ee version
     global_version.set_ee()
 
     test_encryption()
 
-    application = get_application_base()
+    application = get_application_base(lifespan_override=lifespan)
 
     if MULTI_TENANT:
         add_tenant_id_middleware(application, logger)
@@ -166,10 +184,6 @@ def get_application() -> FastAPI:
     # Ensure all routes have auth enabled or are explicitly marked as public
     check_ee_router_auth(application)
 
-    # seed the Onyx environment with LLMs, Assistants, etc. based on an optional
-    # environment variable. Used to automate deployment for multiple environments.
-    seed_db()
-
     # for debugging discovered routes
     # for route in application.router.routes:
     #     print(f"Path: {route.path}, Methods: {route.methods}")

Diff for: backend/onyx/background/celery/tasks/beat_schedule.py

@@ -179,14 +179,24 @@ def make_cloud_generator_task(task: dict[str, Any]) -> dict[str, Any]:
     },
     {
         "name": f"{ONYX_CLOUD_CELERY_TASK_PREFIX}_check-available-tenants",
-        "task": OnyxCeleryTask.CHECK_AVAILABLE_TENANTS,
+        "task": OnyxCeleryTask.CLOUD_CHECK_AVAILABLE_TENANTS,
         "schedule": timedelta(minutes=10),
         "options": {
             "queue": OnyxCeleryQueues.MONITORING,
             "priority": OnyxCeleryPriority.HIGH,
             "expires": BEAT_EXPIRES_DEFAULT,
         },
     },
+    {
+        "name": f"{ONYX_CLOUD_CELERY_TASK_PREFIX}_monitor-celery-pidbox",
+        "task": OnyxCeleryTask.CLOUD_MONITOR_CELERY_PIDBOX,
+        "schedule": timedelta(hours=4),
+        "options": {
+            "queue": OnyxCeleryQueues.MONITORING,
+            "priority": OnyxCeleryPriority.HIGH,
+            "expires": BEAT_EXPIRES_DEFAULT,
+        },
+    },
 ]
 
 # tasks that only run self hosted

Diff for: backend/onyx/background/celery/tasks/indexing/tasks.py

@@ -44,6 +44,7 @@
 from onyx.configs.constants import CELERY_INDEXING_LOCK_TIMEOUT
 from onyx.configs.constants import CELERY_INDEXING_WATCHDOG_CONNECTOR_TIMEOUT
 from onyx.configs.constants import CELERY_TASK_WAIT_FOR_FENCE_TIMEOUT
+from onyx.configs.constants import OnyxCeleryPriority
 from onyx.configs.constants import OnyxCeleryQueues
 from onyx.configs.constants import OnyxCeleryTask
 from onyx.configs.constants import OnyxRedisConstants
@@ -1234,8 +1235,9 @@ def connector_indexing_proxy_task(
 @shared_task(
     name=OnyxCeleryTask.CHECK_FOR_CHECKPOINT_CLEANUP,
     soft_time_limit=300,
+    bind=True,
 )
-def check_for_checkpoint_cleanup(*, tenant_id: str) -> None:
+def check_for_checkpoint_cleanup(self: Task, *, tenant_id: str) -> None:
     """Clean up old checkpoints that are older than 7 days."""
     locked = False
     redis_client = get_redis_client(tenant_id=tenant_id)
@@ -1256,14 +1258,15 @@ def check_for_checkpoint_cleanup(*, tenant_id: str) -> None:
                 task_logger.info(
                     f"Cleaning up checkpoint for index attempt {attempt.id}"
                 )
-                cleanup_checkpoint_task.apply_async(
+                self.app.send_task(
+                    OnyxCeleryTask.CLEANUP_CHECKPOINT,
                     kwargs={
                         "index_attempt_id": attempt.id,
                         "tenant_id": tenant_id,
                     },
                     queue=OnyxCeleryQueues.CHECKPOINT_CLEANUP,
+                    priority=OnyxCeleryPriority.MEDIUM,
                 )
-
     except Exception:
         task_logger.exception("Unexpected exception during checkpoint cleanup")
         return None
@@ -1287,5 +1290,17 @@ def cleanup_checkpoint_task(
     self: Task, *, index_attempt_id: int, tenant_id: str | None
 ) -> None:
     """Clean up a checkpoint for a given index attempt"""
-    with get_session_with_current_tenant() as db_session:
-        cleanup_checkpoint(db_session, index_attempt_id)
+
+    start = time.monotonic()
+
+    try:
+        with get_session_with_current_tenant() as db_session:
+            cleanup_checkpoint(db_session, index_attempt_id)
+    finally:
+        elapsed = time.monotonic() - start
+
+        task_logger.info(
+            f"cleanup_checkpoint_task completed: tenant_id={tenant_id} "
+            f"index_attempt_id={index_attempt_id} "
+            f"elapsed={elapsed:.2f}"
+        )

Diff for: backend/onyx/background/celery/tasks/monitoring/tasks.py

@@ -4,6 +4,7 @@
 from datetime import timedelta
 from itertools import islice
 from typing import Any
+from typing import cast
 from typing import Literal
 
 import psutil
@@ -1000,3 +1001,50 @@ def monitor_process_memory(self: Task, *, tenant_id: str) -> None:
 
     except Exception:
         task_logger.exception("Error in monitor_process_memory task")
+
+
+@shared_task(
+    name=OnyxCeleryTask.CLOUD_MONITOR_CELERY_PIDBOX, ignore_result=True, bind=True
+)
+def cloud_monitor_celery_pidbox(
+    self: Task,
+) -> None:
+    """
+    Celery can leave behind orphaned pidboxes from old workers that are idle and never cleaned up.
+    This task removes them based on idle time to avoid Redis clutter and overflowing the instance.
+    This is a real issue we've observed in production.
+
+    Note:
+    - Setting CELERY_ENABLE_REMOTE_CONTROL = False would prevent pidbox keys entirely,
+    but might also disable features like inspect, broadcast, and worker remote control.
+    Use with caution.
+    """
+
+    num_deleted = 0
+
+    MAX_PIDBOX_IDLE = 24 * 3600  # 1 day in seconds
+    r_celery: Redis = self.app.broker_connection().channel().client  # type: ignore
+    for key in r_celery.scan_iter("*.reply.celery.pidbox"):
+        key_bytes = cast(bytes, key)
+        key_str = key_bytes.decode("utf-8")
+        if key_str.startswith("_kombu"):
+            continue
+
+        idletime_raw = r_celery.object("idletime", key)
+        if idletime_raw is None:
+            continue
+
+        idletime = cast(int, idletime_raw)
+        if idletime < MAX_PIDBOX_IDLE:
+            continue
+
+        r_celery.delete(key)
+        task_logger.info(
+            f"Deleted idle pidbox: pidbox={key_str} "
+            f"idletime={idletime} "
+            f"max_idletime={MAX_PIDBOX_IDLE}"
+        )
+        num_deleted += 1
+
+    # Enable later in case we want some aggregate metrics
+    # task_logger.info(f"Deleted idle pidbox: pidbox={key_str}")

Diff for: backend/onyx/background/celery/tasks/tenant_provisioning/tasks.py

@@ -35,7 +35,7 @@
 
 
 @shared_task(
-    name=OnyxCeleryTask.CHECK_AVAILABLE_TENANTS,
+    name=OnyxCeleryTask.CLOUD_CHECK_AVAILABLE_TENANTS,
     queue=OnyxCeleryQueues.MONITORING,
     ignore_result=True,
     soft_time_limit=JOB_TIMEOUT,

Diff for: backend/onyx/configs/constants.py

@@ -398,7 +398,12 @@ class OnyxCeleryTask:
     CLOUD_MONITOR_CELERY_QUEUES = (
         f"{ONYX_CLOUD_CELERY_TASK_PREFIX}_monitor_celery_queues"
     )
-    CHECK_AVAILABLE_TENANTS = f"{ONYX_CLOUD_CELERY_TASK_PREFIX}_check_available_tenants"
+    CLOUD_CHECK_AVAILABLE_TENANTS = (
+        f"{ONYX_CLOUD_CELERY_TASK_PREFIX}_check_available_tenants"
+    )
+    CLOUD_MONITOR_CELERY_PIDBOX = (
+        f"{ONYX_CLOUD_CELERY_TASK_PREFIX}_monitor_celery_pidbox"
+    )
 
     # Tenant pre-provisioning
     PRE_PROVISION_TENANT = f"{ONYX_CLOUD_CELERY_TASK_PREFIX}_pre_provision_tenant"

Diff for: backend/onyx/main.py

@@ -21,6 +21,7 @@
 from sentry_sdk.integrations.fastapi import FastApiIntegration
 from sentry_sdk.integrations.starlette import StarletteIntegration
 from sqlalchemy.orm import Session
+from starlette.types import Lifespan
 
 from onyx import __version__
 from onyx.auth.schemas import UserCreate
@@ -275,8 +276,12 @@ def log_http_error(request: Request, exc: Exception) -> JSONResponse:
     )
 
 
-def get_application() -> FastAPI:
-    application = FastAPI(title="Onyx Backend", version=__version__, lifespan=lifespan)
+def get_application(lifespan_override: Lifespan | None = None) -> FastAPI:
+    application = FastAPI(
+        title="Onyx Backend",
+        version=__version__,
+        lifespan=lifespan_override or lifespan,
+    )
     if SENTRY_DSN:
         sentry_sdk.init(
             dsn=SENTRY_DSN,