From 74f72df7b31352b26d4eae3bd2ba2827507ee2fd Mon Sep 17 00:00:00 2001 From: ali <117142933+muhammad-ali-e@users.noreply.github.com> Date: Thu, 23 Jul 2026 20:07:32 +0530 Subject: [PATCH 01/11] UN-3798 [FIX] Skip broken file-path task load for PG pluggable workers (#2197) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * UN-3798 [FIX] Skip broken file-path task load for PG pluggable workers The top-level worker.py loaded a pluggable worker's tasks.py via spec_from_file_location("tasks", ...) — a bare module name with no parent package — so the plugin's relative imports (from .clients import ...) failed with "attempted relative import with no known parent package", crash-looping every cloud PG pluggable worker (agentic_callback/UN-3754, agentic_studio/ UN-3779, bulk_download/UN-3752) on startup under `python -m pg_queue_consumer`. The tasks are already registered by that point: WorkerBuilder.build_celery_app() (called just above) verifies a pluggable type by importing pluggable_worker.{type}.worker as a proper package (_verify_pluggable_worker_exists → import_module), which runs the plugin's `from . import tasks` and registers the tasks on this same app. So the file-path load is both redundant and broken. Fix: skip the file-path load for pluggable workers; non-pluggable (top-level) workers keep it unchanged. The Celery path is unaffected — it runs via `celery -A pluggable_worker.{type}.worker` (a dotted package import) and never touches this loader; the is_pluggable() file-path branch never ran successfully. Validated on a running dev stack via `python -m pg_queue_consumer`: agentic_callback and agentic_studio now start clean ("tasks already registered … skipping file-path task load" → "ready for Celery"); general (non-pluggable) loads unchanged. Prerequisite for UN-3752 / UN-3754 / UN-3779. Co-Authored-By: Claude Opus 4.8 * UN-3798 [FIX] Address #2197 review: accurate loader comment, to_directory(), zero-task guard - Extract the task-load block into load_worker_tasks(worker_type) with guard-clause returns (depth 3 -> 1) and an accurate docstring: pluggable tasks register via WorkerBuilder's package import and Celery binds them on app finalize; the file-path load is skipped for them because it breaks any relative imports in the plugin's tasks.py. Softened the overstated "every worker crashes" and marked the cloud-plugin `from . import tasks` example as illustrative (contract, not internals). - Add WorkerType.to_directory() as the single source for the underscore->hyphen dir mapping; to_import_path() and the file-path loader both read it (no more slicing the import path). + tests. - Add a post-load zero-task check as a WARNING (not a hard raise): pluggable tasks bind on app finalize which can be after this point, so a raise would false-positive on a correctly-configured pluggable worker (the exact regression this PR fixes). - Move `import importlib.util` to the module-top imports. Deferred (per reviewer's hotfix note): the "spec_from_file_location never called for pluggable" regression test — worker.py runs infra init + builds the Celery app at import, so there is no clean seam to exercise the loader in isolation without a larger main()-guard refactor. The extraction creates that seam for a fast-follow. Co-Authored-By: Claude Opus 4.8 --------- Co-authored-by: Claude Opus 4.8 --- workers/shared/enums/worker_enums_base.py | 15 +++- workers/tests/test_worker_enums_directory.py | 30 +++++++ workers/worker.py | 84 +++++++++++++------- 3 files changed, 99 insertions(+), 30 deletions(-) create mode 100644 workers/tests/test_worker_enums_directory.py diff --git a/workers/shared/enums/worker_enums_base.py b/workers/shared/enums/worker_enums_base.py index 9c4dfbdb46..08447983af 100644 --- a/workers/shared/enums/worker_enums_base.py +++ b/workers/shared/enums/worker_enums_base.py @@ -58,13 +58,22 @@ def to_import_path(self) -> str: if self.is_pluggable(): return f"pluggable_worker.{self.value}.tasks" - # Map to actual directory structure + return f"{self.to_directory()}.tasks" + + def to_directory(self) -> str: + """Return the on-disk directory name for this (non-pluggable) worker. + + Single source of truth for the naming-convention mapping: enum values use + underscores (Python module names), but a few on-disk dirs use hyphens + (e.g. ``api-deployment``). ``to_import_path`` builds on this, and the + file-path task loader in ``worker.py`` reads it directly rather than + slicing the import path. + """ directory_mapping = { "api_deployment": "api-deployment", # All others use same name for directory and module } - directory = directory_mapping.get(self.value, self.value) - return f"{directory}.tasks" + return directory_mapping.get(self.value, self.value) def is_pluggable(self) -> bool: """Check if this worker type is a pluggable worker. diff --git a/workers/tests/test_worker_enums_directory.py b/workers/tests/test_worker_enums_directory.py new file mode 100644 index 0000000000..931fc45f1e --- /dev/null +++ b/workers/tests/test_worker_enums_directory.py @@ -0,0 +1,30 @@ +"""WorkerType.to_directory() — the single source for the on-disk dir mapping (UN-3798). + +worker.py's file-path task loader reads to_directory() directly instead of slicing +to_import_path(); these pin the hyphen/underscore mapping and that to_import_path +is built on top of it, so the two can't drift. +""" + +from celery import Celery +from shared.enums.worker_enums_base import WorkerType + +# The workers autouse conftest fixture finalizes celery's default_app around every +# test; this pure-enum test builds no Celery app of its own, so establish one. +Celery("test-worker-enums").set_default() + + +def test_to_directory_maps_underscored_value_to_hyphenated_dir(): + # The one worker whose on-disk dir differs from its enum value. + assert WorkerType.API_DEPLOYMENT.value == "api_deployment" + assert WorkerType.API_DEPLOYMENT.to_directory() == "api-deployment" + + +def test_to_directory_passthrough_when_dir_equals_value(): + assert WorkerType.GENERAL.to_directory() == "general" + + +def test_to_import_path_is_built_on_to_directory(): + # to_import_path must derive from to_directory (not a separate mapping), so the + # directory naming lives in exactly one place. + wt = WorkerType.API_DEPLOYMENT + assert wt.to_import_path() == f"{wt.to_directory()}.tasks" diff --git a/workers/worker.py b/workers/worker.py index 3822e81a16..a6cfc0d56b 100755 --- a/workers/worker.py +++ b/workers/worker.py @@ -5,6 +5,7 @@ It uses WorkerBuilder to ensure proper configuration including chord retry settings. """ +import importlib.util import logging import os import sys @@ -440,40 +441,69 @@ def on_task_postrun(sender=None, task_id=None, **kwargs): initialize_worker_infrastructure() logger.info("✅ Worker infrastructure initialized successfully") -# Import tasks from the worker-specific directory -# Determine worker path dynamically based on worker type -base_dir = os.path.dirname(os.path.abspath(__file__)) -if worker_type.is_pluggable(): - # Pluggable workers live inside workers/pluggable_worker/{worker_name} - worker_directory = os.path.join("pluggable_worker", worker_type.value) - worker_path = os.path.join(base_dir, worker_directory) -else: - # Enum values use underscores (Python module names); a few on-disk dirs - # still use hyphens (e.g. api-deployment). Derive the directory from the - # authoritative import-path map on WorkerType instead of a blind replace. - worker_directory = worker_type.to_import_path().rsplit(".", 1)[0] + +def load_worker_tasks(worker_type: WorkerType) -> None: + """Register the worker type's Celery tasks. + + Pluggable workers are already registered by this point: ``build_celery_app()`` + above verified the plugin via ``WorkerBuilder._verify_pluggable_worker_exists``, + which does ``importlib.import_module("pluggable_worker..worker")`` — a + proper PACKAGE import that runs the plugin's own task-registration code (e.g. + ``from . import tasks``, which may use relative imports such as + ``from .clients import ...``). Celery binds those tasks to the worker's app + when it finalizes, so they are registered by this point. The generic file-path + load below is therefore SKIPPED for pluggable workers — it loads ``tasks.py`` + under a bare ``"tasks"`` spec with no parent package, which breaks any relative + imports in the plugin's ``tasks.py`` ("attempted relative import with no known + parent package"). + + Non-pluggable (top-level) workers use absolute imports; their ``tasks.py`` is + loaded by file path with the worker directory on ``sys.path``. + """ + if worker_type.is_pluggable(): + logger.info( + f"✅ Pluggable worker {worker_type.value} tasks already registered via " + "WorkerBuilder (package import); skipping file-path task load" + ) + return + + base_dir = os.path.dirname(os.path.abspath(__file__)) + worker_directory = worker_type.to_directory() worker_path = os.path.join(base_dir, worker_directory) + if not os.path.exists(worker_path): + logger.error(f"❌ Worker directory not found: {worker_path}") + return -# Add worker directory to path for task imports -if os.path.exists(worker_path): sys.path.append(worker_path) logger.info(f"✅ Added {worker_directory} to Python path for task imports") - # Import tasks module to register tasks tasks_file = os.path.join(worker_path, "tasks.py") - if os.path.exists(tasks_file): - logger.info(f"📋 Loading tasks from: {tasks_file}") - # Import the tasks module to register tasks with the app - import importlib.util - - spec = importlib.util.spec_from_file_location("tasks", tasks_file) - tasks_module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(tasks_module) - logger.info(f"✅ Tasks loaded successfully from {worker_directory}") - else: + if not os.path.exists(tasks_file): logger.warning(f"⚠️ No tasks.py found at: {tasks_file}") -else: - logger.error(f"❌ Worker directory not found: {worker_path}") + return + + logger.info(f"📋 Loading tasks from: {tasks_file}") + spec = importlib.util.spec_from_file_location("tasks", tasks_file) + tasks_module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(tasks_module) + logger.info(f"✅ Tasks loaded successfully from {worker_directory}") + + +load_worker_tasks(worker_type) + +# A worker that boots with no registered tasks starts but silently processes +# nothing (Celery does not error on an empty registry). Surface that misconfig — +# as a WARNING, not a hard failure: pluggable tasks bind on app finalize +# (@shared_task / connect_on_app_finalize), which can be after this point, so a +# raise here would false-positive on a correctly-configured pluggable worker. +_registered_tasks = [name for name in app.tasks if not name.startswith("celery.")] +if not _registered_tasks: + logger.warning( + f"⚠️ No non-celery tasks registered yet for worker '{worker_type.value}' " + f"(pluggable={worker_type.is_pluggable()}). If this persists past app " + "finalize the worker will start but process nothing — check the " + "worker/plugin task registration." + ) # Log successful configuration logger.info(f"✅ Successfully loaded {worker_type} worker using WorkerBuilder") From 7148e495584424f159c1f6c1881b8a52183a3a8f Mon Sep 17 00:00:00 2001 From: ali <117142933+muhammad-ali-e@users.noreply.github.com> Date: Thu, 23 Jul 2026 20:07:36 +0530 Subject: [PATCH 02/11] UN-3753 [GATED-FEAT] Route webhook notifications through PG-queue transport (flag-gated) (#2198) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * UN-3753 [GATED-FEAT] Route webhook notifications through PG-queue transport (flag-gated) PG-queue analogue for the buffered-webhook dispatch, gated by the pg_queue_enabled Flipt flag. Flag OFF (prod default) keeps the existing Celery send_task byte-unchanged; flag ON enqueues send_webhook_notification onto the PG `notifications` queue, drained by the PG notification consumer. - Dispatch seam: notification_dispatch.py routes send_webhook_notification through resolve_transport — PG via enqueue_task when enabled for the org, else the prior celery_app.send_task (byte-identical, zero regression). _send_clubbed (the buffer-flush path) uses the seam; _org_identifier resolves the org pk -> string for the Flipt decision, keeping the pk in kwargs (the buffer/worker mark contract). + test_notification_dispatch.py (6 cases). - Sites 2 & 3 (WebhookSend / WebhookBatch internal endpoints) stay on Celery: they use countdown stagger, which the PG queue has no delayed-visibility for. Follow-up. Dev-tested on a live stack: flag-off took the Celery branch (0 PG rows); an enqueued send_webhook_notification was drained by the pg_queue_consumer (WORKER_TYPE=notification), delivered to a sink (200), and the row acked. Co-Authored-By: Claude Opus 4.8 * UN-3753 [GATED-FEAT] Address #2198 review: error-classing, org-id naming, call-site tests - Robustness (P1): _send_clubbed's broad `except` mislabeled permanent dispatch errors as broker_failure and reverted to PENDING → retry-forever + mislabeled Sentry tracebacks until the attempt cap. Branch on class: ValueError/TypeError (enqueue_task validation / payload serialization) → dead-letter now with a distinct `result=dispatch_error` metric; keep revert-to-PENDING only for genuine transport/broker exceptions. - Type design (P2): renamed the seam's routing param organization_id → org_string_id so it can't be conflated with the org pk in `kwargs["organization_id"]` (the worker buffer-mark contract) — swapping them was a silent mis-route to Celery. - Observability (P2): _org_identifier now logs a warning (with org_pk) when the lookup returns None — a dangling FK is a data anomaly (CASCADE makes it otherwise unreachable), not an expected "org deleted" path; reworded the docstring and narrowed org_pk: int. Fixed the "before the webhook HTTP call" wording. - Comment accuracy (P2): dropped the aspirational "usable by callers that surface it in an API response" from the Returns block (the sole caller discards it). - Simplify (P3): dropped the redundant `or None` on the routing arg (resolve_transport already normalizes falsy); kept the load-bearing `or ""` on enqueue_task's org_id. - Tests (P1): new test_send_clubbed.py locks the two-org-identifier contract (string routes, pk in kwargs), the transient→PENDING vs permanent→DEAD_LETTER recovery split, and _org_identifier (string id / None+warn). 11 tests pass. Co-Authored-By: Claude Opus 4.8 * UN-3753 [GATED-FEAT] Gate notification dead-letter classing to the PG path only Keep the flag-off (Celery) error flow byte-identical, per the "no change to the Celery path unless flag-gated" rule. The seam now raises a typed PermanentDispatchError ONLY on the PG branch (when enqueue_task rejects the message for a permanent reason — priority/exclusivity validation or a payload that won't JSON-serialize). _send_clubbed dead-letters on that exception; every other failure — including any Celery send_task error — falls to the transient PENDING branch exactly as before UN-3753. + seam test that a permanent enqueue ValueError/TypeError is wrapped; the call-site test now raises PermanentDispatchError (the real contract) instead of a raw ValueError. 12 tests pass. Co-Authored-By: Claude Opus 4.8 --------- Co-authored-by: Claude Opus 4.8 --- backend/notification_v2/internal_api_views.py | 73 +++++++++- .../notification_v2/notification_dispatch.py | 119 +++++++++++++++ backend/notification_v2/tests/__init__.py | 0 .../tests/test_notification_dispatch.py | 136 ++++++++++++++++++ .../tests/test_send_clubbed.py | 112 +++++++++++++++ 5 files changed, 434 insertions(+), 6 deletions(-) create mode 100644 backend/notification_v2/notification_dispatch.py create mode 100644 backend/notification_v2/tests/__init__.py create mode 100644 backend/notification_v2/tests/test_notification_dispatch.py create mode 100644 backend/notification_v2/tests/test_send_clubbed.py diff --git a/backend/notification_v2/internal_api_views.py b/backend/notification_v2/internal_api_views.py index be2bbe57e4..54b6ac3d66 100644 --- a/backend/notification_v2/internal_api_views.py +++ b/backend/notification_v2/internal_api_views.py @@ -14,6 +14,7 @@ from datetime import timedelta from typing import Any, cast +from account_v2.models import Organization from api_v2.models import APIDeployment from django.conf import settings from django.db import transaction @@ -37,6 +38,10 @@ webhook_url_hash, ) from notification_v2.models import Notification, NotificationBuffer +from notification_v2.notification_dispatch import ( + PermanentDispatchError, + dispatch_webhook_notification, +) logger = logging.getLogger(__name__) @@ -513,6 +518,34 @@ def _reclaim_stale_sending() -> int: return int(reclaimed) +def _org_identifier(org_pk: int) -> str | None: + """Resolve the string ``Organization.organization_id`` from the buffer's org pk. + + ``resolve_transport`` keys its Flipt decision on the org's string identifier, + but the buffer stores/uses the Organization pk. One indexed pk lookup per + dispatch group (post-commit) — negligible relative to the downstream webhook + dispatch. + + Data-anomaly guard: ``NotificationBuffer.organization`` is + ``on_delete=CASCADE``, so a live buffer row with a missing org is unreachable + in normal operation. A ``None`` here therefore signals a dangling FK — we log + it (the only org-traceable breadcrumb; resolve_transport's own warning is keyed + on the random dispatch uuid) and fail closed to Celery in resolve_transport. + """ + org_string_id = ( + Organization.objects.filter(pk=org_pk) + .values_list("organization_id", flat=True) + .first() + ) + if org_string_id is None: + logger.warning( + "metric=notification_org_identifier_missing_total org_pk=%s " + "(dangling FK; notification routing falls back to Celery)", + org_pk, + ) + return org_string_id + + def _send_clubbed( *, url: str, @@ -540,8 +573,12 @@ def _send_clubbed( ``buffer_row_ids`` + ``organization_id`` to the worker so it can mark them. """ try: - celery_app.send_task( - "send_webhook_notification", + # Flag-gated transport (UN-3753): PG queue when pg_queue_enabled for this + # org, else Celery (byte-identical to the prior send_task). resolve_transport + # keys on the org STRING id, but the buffer/worker contract below uses the + # org pk — hence _org_identifier(org_id) for routing, org_id in kwargs. + dispatch_webhook_notification( + celery_app=celery_app, args=[url, body, headers, settings.NOTIFICATION_TIMEOUT], kwargs={ "max_retries": max_retries, @@ -557,6 +594,7 @@ def _send_clubbed( "organization_id": org_id, }, queue="notifications", + org_string_id=_org_identifier(org_id), ) logger.info( "metric=notification_batch_dispatched_total platform=%s result=success " @@ -566,7 +604,34 @@ def _send_clubbed( webhook_url_hash(url), len(buffer_ids), ) + except PermanentDispatchError: + # PG-ONLY permanent failure — enqueue_task validation (priority range / + # reply_key+callback exclusivity) or payload serialization — fails + # identically every flush tick. Dead-letter now (distinct metric) instead + # of reverting to PENDING and re-rendering + re-dispatching + emitting a + # broker_failure traceback until the attempt cap. The Celery path never + # raises this, so the flag-off flow is unchanged: a Celery send_task failure + # is an ordinary Exception handled by the transient branch below, exactly as + # before UN-3753. Guard on SENDING so a row the worker already resolved + # isn't clobbered. + logger.exception( + "metric=notification_batch_dispatched_total platform=%s " + "result=dispatch_error org_id=%s webhook_url_hash=%s rows=%d", + platform, + org_id, + webhook_url_hash(url), + len(buffer_ids), + ) + NotificationBuffer.objects.filter( + id__in=buffer_ids, + status=BufferStatus.SENDING.value, + ).update(status=BufferStatus.DEAD_LETTER.value) except Exception: + # TRANSIENT transport/broker failure — revert to PENDING (outside the + # committed txn) so the next flush tick retries; refund the SENDING-claim + # attempt since nothing was queued or sent. Guard on SENDING so a row the + # worker already marked terminal (broker raised post-delivery) isn't + # resurrected into a duplicate. logger.exception( "metric=notification_batch_dispatched_total platform=%s " "result=broker_failure org_id=%s webhook_url_hash=%s rows=%d", @@ -575,10 +640,6 @@ def _send_clubbed( webhook_url_hash(url), len(buffer_ids), ) - # Revert to PENDING (outside the committed txn) so a transient broker - # outage retries next tick; refund the SENDING-claim attempt since nothing - # was queued or sent. Guard on SENDING so a row the worker already marked - # terminal (broker raised post-delivery) isn't resurrected into a duplicate. NotificationBuffer.objects.filter( id__in=buffer_ids, status=BufferStatus.SENDING.value, diff --git a/backend/notification_v2/notification_dispatch.py b/backend/notification_v2/notification_dispatch.py new file mode 100644 index 0000000000..aa87aa0b37 --- /dev/null +++ b/backend/notification_v2/notification_dispatch.py @@ -0,0 +1,119 @@ +"""Transport-routed dispatch for buffered webhook notifications (UN-3753). + +Routes the ``send_webhook_notification`` task through the same +:func:`resolve_transport` flag as the execution path: the PG queue when +``pg_queue_enabled`` for this org, else Celery. **Fail-closed** — with the gate +off (the production default) it resolves to Celery, behaving exactly like the +prior unconditional ``celery_app.send_task`` (zero regression). On PG the task +is drained by the PG notification consumer on the ``notifications`` queue. + +The same ``args``/``kwargs``/``queue`` are forwarded on both paths; on PG they're +JSON-normalized by ``enqueue_task`` (UUIDs/datetimes → str), so the consumer sees +the same payload the Celery worker would. +""" + +from __future__ import annotations + +import logging +import uuid +from typing import Any + +from pg_queue.producer import enqueue_task +from workflow_manager.workflow_v2.transport import resolve_transport + +from unstract.core.data_models import is_pg_transport + +logger = logging.getLogger(__name__) + +# The fired task name — mirrors the Celery task registered by the notification +# worker; kept as a local constant so the backend doesn't import the workers pkg. +WEBHOOK_NOTIFICATION_TASK = "send_webhook_notification" + + +class PermanentDispatchError(Exception): + """A dispatch failure that would fail identically on every retry. + + Raised ONLY on the PG path, when ``enqueue_task`` rejects the message for a + permanent reason (priority range / reply_key+callback exclusivity validation, + or a payload that can't be JSON-serialized). The Celery path never raises it, + so a caller can dead-letter on this exception without altering the flag-off + (Celery) error flow — a Celery ``send_task`` failure stays an ordinary + ``Exception`` the caller's transient handler owns, exactly as before. + """ + + +def dispatch_webhook_notification( + *, + celery_app: Any, + args: list[Any], + kwargs: dict[str, Any], + queue: str, + org_string_id: str | None, +) -> str: + """Dispatch ``send_webhook_notification`` on the resolved transport. + + ``args``/``kwargs``/``queue`` are forwarded unchanged on both paths, so the + flag-off (Celery) path is byte-identical to the legacy ``send_task`` call. + + Args: + celery_app: Injected Celery app (the backend's ``celery_service.app``); + passed in rather than imported so this seam stays trivially testable. + args: Positional task args, forwarded verbatim. + kwargs: Keyword task args, forwarded verbatim. For the buffered path this + carries ``organization_id`` = the buffer's org **pk** (the worker's + buffer-mark contract) — deliberately a DIFFERENT identifier from the + ``org_string_id`` param below (the two must not be conflated). + queue: Target queue name, forwarded verbatim. + org_string_id: The org's **string** identifier + (``Organization.organization_id``), used solely for the Flipt + transport decision — NOT the org pk carried in ``kwargs``. ``None`` (or + empty) fails closed to Celery. + + Returns: + A task id string — the Celery ``AsyncResult`` id on the Celery path, or + the minted PG task id on the PG path. Returned for symmetry / any future + caller; the sole current caller (fire-and-forget) discards it. + """ + # A buffered notification is a single fire-and-forget task with no natural + # sticky entity, so mint a fresh id to drive Flipt's percentage bucketing and + # to serve as the PG task id. + dispatch_id = str(uuid.uuid4()) + # resolve_transport already normalizes falsy input to Celery, so pass the id + # straight through (no `or None` needed). + transport = resolve_transport( + execution_id=dispatch_id, + organization_id=org_string_id, + ) + # Use the shared is_pg_transport() — the single source for "what counts as PG + # transport" — rather than opening a second comparison site. + if is_pg_transport(transport): + try: + enqueue_task( + task_name=WEBHOOK_NOTIFICATION_TASK, + queue=queue, + args=args, + kwargs=kwargs, + # enqueue_task's org_id is str-typed — coerce None→"" (unlike the + # routing arg above, this `or ""` is load-bearing). + org_id=org_string_id or "", + task_id=dispatch_id, + ) + except (ValueError, TypeError) as exc: + # PG-only permanent failure (enqueue_task validation / JSON encode): + # re-raise as PermanentDispatchError so the caller dead-letters it. + # A transient PG error (DB down) is NOT wrapped — it propagates as an + # ordinary Exception into the caller's retry (revert-to-PENDING) path. + raise PermanentDispatchError(str(exc)) from exc + logger.info( + "Webhook notification enqueued on PG '%s' queue (task_id=%s)", + queue, + dispatch_id, + ) + return dispatch_id + result = celery_app.send_task( + WEBHOOK_NOTIFICATION_TASK, + args=args, + kwargs=kwargs, + queue=queue, + ) + return result.id diff --git a/backend/notification_v2/tests/__init__.py b/backend/notification_v2/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/backend/notification_v2/tests/test_notification_dispatch.py b/backend/notification_v2/tests/test_notification_dispatch.py new file mode 100644 index 0000000000..e53514d5de --- /dev/null +++ b/backend/notification_v2/tests/test_notification_dispatch.py @@ -0,0 +1,136 @@ +"""Unit tests for the buffered-webhook transport routing (UN-3753). + +``resolve_transport`` + ``enqueue_task`` are patched on the module, so no Flipt / +DB is needed — these pin the routing contract: PG when the flag resolves PG, +Celery otherwise (fail-closed), with byte-identical args/kwargs/queue on both +paths. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest + +import notification_v2.notification_dispatch as nd + +_ARGS = ["https://hook.test", {"text": "hi"}, {"Content-Type": "application/json"}, 30] +_KWARGS = { + "max_retries": 3, + "retry_delay": 10, + "platform": "SLACK", + "raise_on_final_failure": True, + "buffer_row_ids": ["b1", "b2"], + "organization_id": 7, # the org pk (worker's buffer-mark contract) +} +_QUEUE = "notifications" + + +def _dispatch(celery_app, org_string_id="org_x"): + return nd.dispatch_webhook_notification( + celery_app=celery_app, + args=_ARGS, + kwargs=_KWARGS, + queue=_QUEUE, + org_string_id=org_string_id, + ) + + +class TestDispatchWebhookNotification: + def test_routes_to_pg_when_flag_resolves_pg(self): + celery = MagicMock() + with ( + patch.object(nd, "resolve_transport", return_value="pg_queue"), + patch.object(nd, "enqueue_task", return_value=42) as enqueue, + ): + task_id = _dispatch(celery) + enqueue.assert_called_once() + kwargs = enqueue.call_args.kwargs + assert kwargs["task_name"] == "send_webhook_notification" + assert kwargs["queue"] == _QUEUE + assert kwargs["args"] == _ARGS + assert kwargs["kwargs"] == _KWARGS + assert kwargs["org_id"] == "org_x" + # The minted PG task id is returned and threaded into the enqueue row. + assert kwargs["task_id"] == task_id + celery.send_task.assert_not_called() + + def test_routes_to_celery_when_flag_resolves_celery(self): + celery = MagicMock() + with ( + patch.object(nd, "resolve_transport", return_value="celery"), + patch.object(nd, "enqueue_task") as enqueue, + ): + result = _dispatch(celery) + celery.send_task.assert_called_once_with( + "send_webhook_notification", args=_ARGS, kwargs=_KWARGS, queue=_QUEUE + ) + enqueue.assert_not_called() + assert result is celery.send_task.return_value.id + + def test_pg_enqueue_failure_propagates_with_no_celery_fallback(self): + # No silent Celery fallback on a TRANSIENT PG failure — it propagates raw + # (not wrapped), so the caller (_send_clubbed) reverts rows to PENDING. + celery = MagicMock() + with ( + patch.object(nd, "resolve_transport", return_value="pg_queue"), + patch.object(nd, "enqueue_task", side_effect=RuntimeError("pg down")), + ): + with pytest.raises(RuntimeError, match="pg down"): + _dispatch(celery) + celery.send_task.assert_not_called() + + def test_pg_permanent_enqueue_error_is_wrapped(self): + # A PERMANENT enqueue error (ValueError/TypeError: validation / JSON encode) + # is re-raised as PermanentDispatchError so the caller dead-letters it — + # raised ONLY on the PG path, so the Celery flow is never affected. + celery = MagicMock() + for exc in (ValueError("priority out of range"), TypeError("not serializable")): + with ( + patch.object(nd, "resolve_transport", return_value="pg_queue"), + patch.object(nd, "enqueue_task", side_effect=exc), + ): + with pytest.raises(nd.PermanentDispatchError): + _dispatch(celery) + celery.send_task.assert_not_called() + + def test_none_org_fails_closed_to_celery(self): + # A missing org string (org deleted) must not route to PG — resolve_transport + # fails closed, and we pass organization_id=None straight through to it. + celery = MagicMock() + with ( + patch.object(nd, "resolve_transport", return_value="celery") as resolve, + patch.object(nd, "enqueue_task") as enqueue, + ): + _dispatch(celery, org_string_id=None) + assert resolve.call_args.kwargs["organization_id"] is None + celery.send_task.assert_called_once() + enqueue.assert_not_called() + + def test_resolve_transport_buckets_by_minted_dispatch_id(self): + # Fire-and-forget: entity_id is a freshly minted uuid (str), and it equals + # the PG task_id so the row and the Flipt bucket agree. + celery = MagicMock() + with ( + patch.object(nd, "resolve_transport", return_value="pg_queue") as resolve, + patch.object(nd, "enqueue_task", return_value=1) as enqueue, + ): + task_id = _dispatch(celery) + assert resolve.call_args.kwargs["execution_id"] == task_id + assert enqueue.call_args.kwargs["task_id"] == task_id + + def test_args_and_kwargs_identical_on_both_paths(self): + # The consumer must behave the same regardless of transport. + celery = MagicMock() + with patch.object(nd, "resolve_transport", return_value="celery"): + _dispatch(celery) + celery_call = celery.send_task.call_args.kwargs + celery2 = MagicMock() + with ( + patch.object(nd, "resolve_transport", return_value="pg_queue"), + patch.object(nd, "enqueue_task", return_value=1) as enqueue, + ): + _dispatch(celery2) + assert enqueue.call_args.kwargs["args"] == celery_call["args"] + assert enqueue.call_args.kwargs["kwargs"] == celery_call["kwargs"] + assert enqueue.call_args.kwargs["queue"] == celery_call["queue"] diff --git a/backend/notification_v2/tests/test_send_clubbed.py b/backend/notification_v2/tests/test_send_clubbed.py new file mode 100644 index 0000000000..4f70603d5f --- /dev/null +++ b/backend/notification_v2/tests/test_send_clubbed.py @@ -0,0 +1,112 @@ +"""Call-site tests for ``_send_clubbed`` / ``_org_identifier`` (UN-3753). + +The seam's own routing is covered by ``test_notification_dispatch``; these lock +the two highest-risk regressions at the buffer-flush call site: + +1. the two-org-identifier contract — the org **string** id routes the flag, while + the org **pk** stays in the worker kwargs (swapping them passes every seam test + yet mis-routes every org to Celery and strands the mark endpoint); +2. failure recovery — a transient/broker error refunds and reverts SENDING rows to + PENDING, while a permanent PG enqueue error (surfaced as ``PermanentDispatchError``, + raised only on the PG path) dead-letters instead of retrying forever. A Celery + send_task failure is an ordinary ``Exception`` → the PENDING path, so flag-off is + byte-identical. + +Mock-based (no broker/DB): patch the module's collaborators. +""" + +from __future__ import annotations + +from unittest.mock import patch + +from notification_v2 import internal_api_views as views +from notification_v2.enums import BufferStatus + + +def _send(**overrides): + kw = { + "url": "https://hook.test", + "body": {"text": "x"}, + "headers": {"Content-Type": "application/json"}, + "platform": "SLACK", + "max_retries": 3, + "buffer_ids": ["b1"], + "org_id": 7, + } + kw.update(overrides) + views._send_clubbed(**kw) + + +class TestSendClubbedOrgContract: + def test_routes_string_id_and_keeps_pk_in_kwargs(self): + with ( + patch.object(views, "_org_identifier", return_value="org-str-id") as ident, + patch.object(views, "dispatch_webhook_notification") as seam, + ): + _send(org_id=7) + ident.assert_called_once_with(7) + call = seam.call_args.kwargs + # Routing uses the STRING id; the worker buffer-mark contract keeps the PK. + assert call["org_string_id"] == "org-str-id" + assert call["kwargs"]["organization_id"] == 7 + assert call["queue"] == "notifications" + + +class TestSendClubbedFailureRecovery: + def test_transient_failure_reverts_sending_rows_to_pending(self): + with ( + patch.object(views, "_org_identifier", return_value="o"), + patch.object( + views, + "dispatch_webhook_notification", + side_effect=RuntimeError("broker down"), + ), + patch.object(views, "NotificationBuffer") as buf, + ): + _send(buffer_ids=["b1", "b2"]) + # Guarded on SENDING, reverted to PENDING (attempt refunded). + fkw = buf.objects.filter.call_args.kwargs + assert fkw["status"] == BufferStatus.SENDING.value + assert set(fkw["id__in"]) == {"b1", "b2"} + ukw = buf.objects.filter.return_value.update.call_args.kwargs + assert ukw["status"] == BufferStatus.PENDING.value + assert ukw["dispatched_at"] is None + + def test_permanent_pg_error_dead_letters(self): + # The PG path surfaces a permanent enqueue failure as PermanentDispatchError; + # only that dead-letters. (A raw ValueError from the Celery path can't occur, + # and would fall through to the transient PENDING branch, unchanged.) + with ( + patch.object(views, "_org_identifier", return_value="o"), + patch.object( + views, + "dispatch_webhook_notification", + side_effect=views.PermanentDispatchError("priority out of range"), + ), + patch.object(views, "NotificationBuffer") as buf, + ): + _send() + # Permanent error → terminal DEAD_LETTER, no PENDING revert / refund. + ukw = buf.objects.filter.return_value.update.call_args.kwargs + assert ukw == {"status": BufferStatus.DEAD_LETTER.value} + + +class TestOrgIdentifier: + def test_returns_string_id(self): + with patch.object(views, "Organization") as org: + chain = org.objects.filter.return_value.values_list.return_value + chain.first.return_value = "org-uuid" + assert views._org_identifier(7) == "org-uuid" + org.objects.filter.assert_called_once_with(pk=7) + + def test_missing_org_returns_none_and_warns(self): + with ( + patch.object(views, "Organization") as org, + patch.object(views.logger, "warning") as warn, + ): + chain = org.objects.filter.return_value.values_list.return_value + chain.first.return_value = None + assert views._org_identifier(7) is None + # A dangling FK is logged (org-traceable) rather than swallowed. + warn.assert_called_once() + assert "org_pk" in warn.call_args.args[0] From 02dcdcb4a19ff9c5a6f98f5859a667f74db178fd Mon Sep 17 00:00:00 2001 From: ali Date: Wed, 29 Jul 2026 13:32:42 +0530 Subject: [PATCH 03/11] UN-3445 [GATED-FEAT] Address PR #2217 review: harden loader empty-registry + notification NaN/typing/comment fixes Co-Authored-By: Claude Opus 4.8 --- backend/notification_v2/internal_api_views.py | 21 ++++-- .../notification_v2/notification_dispatch.py | 6 +- .../tests/test_send_clubbed.py | 11 +-- backend/pg_queue/producer.py | 9 ++- backend/pg_queue/tests/test_producer.py | 11 +++ workers/shared/enums/worker_enums_base.py | 4 +- workers/worker.py | 70 +++++++++++-------- 7 files changed, 87 insertions(+), 45 deletions(-) diff --git a/backend/notification_v2/internal_api_views.py b/backend/notification_v2/internal_api_views.py index 54b6ac3d66..04d77b81cb 100644 --- a/backend/notification_v2/internal_api_views.py +++ b/backend/notification_v2/internal_api_views.py @@ -538,7 +538,11 @@ def _org_identifier(org_pk: int) -> str | None: .first() ) if org_string_id is None: - logger.warning( + # Sentry-routed (logger.error): a live buffer row with no org is a data + # anomaly (dangling FK / corruption) that shouldn't happen under the + # CASCADE constraint, not routine noise. Routing still fails closed to + # Celery in resolve_transport. + logger.error( "metric=notification_org_identifier_missing_total org_pk=%s " "(dangling FK; notification routing falls back to Celery)", org_pk, @@ -554,7 +558,7 @@ def _send_clubbed( platform: str, max_retries: int, buffer_ids: list[str], - org_id: Any, + org_id: int, ) -> None: """Send the clubbed Celery task after the DB transition has committed. @@ -605,9 +609,12 @@ def _send_clubbed( len(buffer_ids), ) except PermanentDispatchError: - # PG-ONLY permanent failure — enqueue_task validation (priority range / - # reply_key+callback exclusivity) or payload serialization — fails - # identically every flush tick. Dead-letter now (distinct metric) instead + # PG-ONLY permanent failure. From this path the only reachable cause is + # payload JSON-serialization (e.g. a NaN/Infinity float that jsonb rejects + # at insert): this call passes no reply_key/callback and a default in-range + # priority, so enqueue_task's other permanent checks can't fire here (the + # full set lives on PermanentDispatchError's docstring). It fails + # identically every flush tick — dead-letter now (distinct metric) instead # of reverting to PENDING and re-rendering + re-dispatching + emitting a # broker_failure traceback until the attempt cap. The Celery path never # raises this, so the flag-off flow is unchanged: a Celery send_task failure @@ -650,7 +657,7 @@ def _send_clubbed( ) -def _penalize_render_failure(buffer_ids: list[str], org_id: Any, platform: str) -> None: +def _penalize_render_failure(buffer_ids: list[str], org_id: int, platform: str) -> None: """Charge a dispatch attempt to a group whose payloads failed to render. The SENDING-claim increment never runs on a render failure, so count it here @@ -670,7 +677,7 @@ def _penalize_render_failure(buffer_ids: list[str], org_id: Any, platform: str) def _dispatch_group( - org_id: Any, + org_id: int, webhook_url: str, auth_sig: str, platform: str, diff --git a/backend/notification_v2/notification_dispatch.py b/backend/notification_v2/notification_dispatch.py index aa87aa0b37..a7331a1618 100644 --- a/backend/notification_v2/notification_dispatch.py +++ b/backend/notification_v2/notification_dispatch.py @@ -93,8 +93,10 @@ def dispatch_webhook_notification( queue=queue, args=args, kwargs=kwargs, - # enqueue_task's org_id is str-typed — coerce None→"" (unlike the - # routing arg above, this `or ""` is load-bearing). + # enqueue_task types org_id as str; None→"" here satisfies that + # type only, not runtime — enqueue_task itself re-coerces + # ``org_id or ""`` at insert (producer.py), so this has no runtime + # effect. org_id=org_string_id or "", task_id=dispatch_id, ) diff --git a/backend/notification_v2/tests/test_send_clubbed.py b/backend/notification_v2/tests/test_send_clubbed.py index 4f70603d5f..7aeaa397ed 100644 --- a/backend/notification_v2/tests/test_send_clubbed.py +++ b/backend/notification_v2/tests/test_send_clubbed.py @@ -99,14 +99,15 @@ def test_returns_string_id(self): assert views._org_identifier(7) == "org-uuid" org.objects.filter.assert_called_once_with(pk=7) - def test_missing_org_returns_none_and_warns(self): + def test_missing_org_returns_none_and_logs_error(self): with ( patch.object(views, "Organization") as org, - patch.object(views.logger, "warning") as warn, + patch.object(views.logger, "error") as err, ): chain = org.objects.filter.return_value.values_list.return_value chain.first.return_value = None assert views._org_identifier(7) is None - # A dangling FK is logged (org-traceable) rather than swallowed. - warn.assert_called_once() - assert "org_pk" in warn.call_args.args[0] + # A dangling FK is a data anomaly: logged at error (Sentry-routed), + # org-traceable, rather than swallowed. + err.assert_called_once() + assert "org_pk" in err.call_args.args[0] diff --git a/backend/pg_queue/producer.py b/backend/pg_queue/producer.py index 3c4585fc4d..e37fdae2ce 100644 --- a/backend/pg_queue/producer.py +++ b/backend/pg_queue/producer.py @@ -50,8 +50,15 @@ def _json_safe(value: Any) -> Any: ``PgQueueMessage.message`` is a plain ``JSONField`` (no Django encoder), and the worker consumer already receives string ids on the existing PG dispatch path, so coercing here keeps both transports consistent. + + ``allow_nan=False`` rejects ``NaN``/``Infinity`` here with a ``ValueError`` + rather than letting Python's default lenient encoder emit the non-standard + ``NaN``/``Infinity`` tokens: Postgres ``jsonb`` rejects those at insert with a + ``django.db.DataError`` (a permanent failure the notification dispatcher's + ``(ValueError, TypeError)`` seam would otherwise miss, looping the row). Fail + at the intended seam instead. """ - return json.loads(json.dumps(value, default=str)) + return json.loads(json.dumps(value, default=str, allow_nan=False)) def enqueue_task( diff --git a/backend/pg_queue/tests/test_producer.py b/backend/pg_queue/tests/test_producer.py index 40df8a79fe..13a50a2e6b 100644 --- a/backend/pg_queue/tests/test_producer.py +++ b/backend/pg_queue/tests/test_producer.py @@ -106,6 +106,17 @@ def test_json_safe_coerces_datetime(self): when = model.objects.create.call_args.kwargs["message"]["kwargs"]["when"] assert isinstance(when, str) and "2026-06-18" in when + def test_json_safe_rejects_nan_with_value_error(self): + # NaN would slip past the default lenient encoder and only fail at the + # jsonb insert (DataError); allow_nan=False surfaces it as a ValueError at + # the enqueue seam so the notification dispatcher can dead-letter it. + with patch(_MODEL) as model: + model.objects.create.return_value = MagicMock(msg_id=1) + with pytest.raises(ValueError): + producer.enqueue_task( + task_name="t", queue="celery", kwargs={"score": float("nan")} + ) + def test_enqueue_failure_logs_and_propagates(self): with patch(_MODEL) as model: model.objects.create.side_effect = RuntimeError("db down") diff --git a/workers/shared/enums/worker_enums_base.py b/workers/shared/enums/worker_enums_base.py index 08447983af..97e4799768 100644 --- a/workers/shared/enums/worker_enums_base.py +++ b/workers/shared/enums/worker_enums_base.py @@ -70,10 +70,10 @@ def to_directory(self) -> str: slicing the import path. """ directory_mapping = { - "api_deployment": "api-deployment", + WorkerType.API_DEPLOYMENT: "api-deployment", # All others use same name for directory and module } - return directory_mapping.get(self.value, self.value) + return directory_mapping.get(self, self.value) def is_pluggable(self) -> bool: """Check if this worker type is a pluggable worker. diff --git a/workers/worker.py b/workers/worker.py index a6cfc0d56b..295f5c9b00 100755 --- a/workers/worker.py +++ b/workers/worker.py @@ -445,24 +445,28 @@ def on_task_postrun(sender=None, task_id=None, **kwargs): def load_worker_tasks(worker_type: WorkerType) -> None: """Register the worker type's Celery tasks. - Pluggable workers are already registered by this point: ``build_celery_app()`` - above verified the plugin via ``WorkerBuilder._verify_pluggable_worker_exists``, - which does ``importlib.import_module("pluggable_worker..worker")`` — a - proper PACKAGE import that runs the plugin's own task-registration code (e.g. - ``from . import tasks``, which may use relative imports such as - ``from .clients import ...``). Celery binds those tasks to the worker's app - when it finalizes, so they are registered by this point. The generic file-path - load below is therefore SKIPPED for pluggable workers — it loads ``tasks.py`` - under a bare ``"tasks"`` spec with no parent package, which breaks any relative - imports in the plugin's ``tasks.py`` ("attempted relative import with no known - parent package"). + For pluggable workers, ``build_celery_app()`` above already imported the + plugin package via ``WorkerBuilder._verify_pluggable_worker_exists`` + (``importlib.import_module("pluggable_worker..worker")`` — a proper + PACKAGE import that RUNS the plugin's own registration code, e.g. + ``from . import tasks`` with relative imports such as + ``from .clients import ...``). That registration's ``@shared_task`` bindings + complete when the app FINALIZES, which can be *after* this point — so the + tasks are not necessarily bound to the app yet here. The generic file-path + load below is SKIPPED for pluggable workers not because they are already + bound, but because loading ``tasks.py`` under a bare ``"tasks"`` spec (no + parent package) would break the plugin's relative imports ("attempted + relative import with no known parent package"). Non-pluggable (top-level) workers use absolute imports; their ``tasks.py`` is - loaded by file path with the worker directory on ``sys.path``. + loaded by file path with the worker directory on ``sys.path``. Their tasks + bind eagerly here — there is no later finalize step to rescue them — so a + missing directory / ``tasks.py`` is a broken deploy and hard-fails rather + than booting a no-op worker. """ if worker_type.is_pluggable(): logger.info( - f"✅ Pluggable worker {worker_type.value} tasks already registered via " + f"✅ Pluggable worker {worker_type.value} tasks registered via " "WorkerBuilder (package import); skipping file-path task load" ) return @@ -471,16 +475,16 @@ def load_worker_tasks(worker_type: WorkerType) -> None: worker_directory = worker_type.to_directory() worker_path = os.path.join(base_dir, worker_directory) if not os.path.exists(worker_path): - logger.error(f"❌ Worker directory not found: {worker_path}") - return + # Non-pluggable only (pluggable returned above): tasks bind eagerly here + # with no later finalize step, so a missing worker dir is terminal. + raise RuntimeError(f"Worker directory not found: {worker_path}") sys.path.append(worker_path) logger.info(f"✅ Added {worker_directory} to Python path for task imports") tasks_file = os.path.join(worker_path, "tasks.py") if not os.path.exists(tasks_file): - logger.warning(f"⚠️ No tasks.py found at: {tasks_file}") - return + raise RuntimeError(f"No tasks.py found for worker at: {tasks_file}") logger.info(f"📋 Loading tasks from: {tasks_file}") spec = importlib.util.spec_from_file_location("tasks", tasks_file) @@ -492,18 +496,28 @@ def load_worker_tasks(worker_type: WorkerType) -> None: load_worker_tasks(worker_type) # A worker that boots with no registered tasks starts but silently processes -# nothing (Celery does not error on an empty registry). Surface that misconfig — -# as a WARNING, not a hard failure: pluggable tasks bind on app finalize -# (@shared_task / connect_on_app_finalize), which can be after this point, so a -# raise here would false-positive on a correctly-configured pluggable worker. -_registered_tasks = [name for name in app.tasks if not name.startswith("celery.")] -if not _registered_tasks: - logger.warning( - f"⚠️ No non-celery tasks registered yet for worker '{worker_type.value}' " - f"(pluggable={worker_type.is_pluggable()}). If this persists past app " - "finalize the worker will start but process nothing — check the " - "worker/plugin task registration." +# nothing (Celery does not error on an empty registry). Surface that misconfig, +# but split by worker kind: pluggable tasks bind on app finalize (@shared_task / +# connect_on_app_finalize), which can be after this point, so a raise here would +# false-positive on a correctly-configured pluggable worker — only WARN. A +# non-pluggable worker has no later binding step (load_worker_tasks bound its +# tasks eagerly above), so an empty registry here is terminal — RAISE. +if not any(not name.startswith("celery.") for name in app.tasks): + _empty_registry_msg = ( + f"No non-celery tasks registered for worker '{worker_type.value}' " + f"(pluggable={worker_type.is_pluggable()})." ) + if worker_type.is_pluggable(): + logger.warning( + f"⚠️ {_empty_registry_msg} If this persists past app finalize the " + "worker will start but process nothing — check the plugin task " + "registration." + ) + else: + raise RuntimeError( + f"{_empty_registry_msg} The worker would start but process nothing — " + "check the worker task registration." + ) # Log successful configuration logger.info(f"✅ Successfully loaded {worker_type} worker using WorkerBuilder") From 3e7961c4ccff7ed010b0bdededcb9b5817a077a1 Mon Sep 17 00:00:00 2001 From: ali Date: Wed, 29 Jul 2026 13:57:46 +0530 Subject: [PATCH 04/11] UN-3445 [GATED-FEAT] Fix SonarCloud S5778: single throwing call in NaN test Hoist float("nan") out of the pytest.raises block so only enqueue_task is under assertion. Co-Authored-By: Claude Opus 4.8 --- backend/pg_queue/tests/test_producer.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/backend/pg_queue/tests/test_producer.py b/backend/pg_queue/tests/test_producer.py index 13a50a2e6b..e812965c63 100644 --- a/backend/pg_queue/tests/test_producer.py +++ b/backend/pg_queue/tests/test_producer.py @@ -110,12 +110,11 @@ def test_json_safe_rejects_nan_with_value_error(self): # NaN would slip past the default lenient encoder and only fail at the # jsonb insert (DataError); allow_nan=False surfaces it as a ValueError at # the enqueue seam so the notification dispatcher can dead-letter it. + nan_kwargs = {"score": float("nan")} with patch(_MODEL) as model: model.objects.create.return_value = MagicMock(msg_id=1) with pytest.raises(ValueError): - producer.enqueue_task( - task_name="t", queue="celery", kwargs={"score": float("nan")} - ) + producer.enqueue_task(task_name="t", queue="celery", kwargs=nan_kwargs) def test_enqueue_failure_logs_and_propagates(self): with patch(_MODEL) as model: From 617b2fbb0ce03df7915c3cda63290332ec989b8b Mon Sep 17 00:00:00 2001 From: ali Date: Wed, 29 Jul 2026 17:47:11 +0530 Subject: [PATCH 05/11] UN-3445 [GATED-FEAT] Pin SENDING clobber-guard in dead-letter test (CodeRabbit) Assert the dead-letter update filters on status=SENDING, mirroring the transient-revert test's guard assertion. Co-Authored-By: Claude Opus 4.8 --- backend/notification_v2/tests/test_send_clubbed.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/backend/notification_v2/tests/test_send_clubbed.py b/backend/notification_v2/tests/test_send_clubbed.py index 7aeaa397ed..4ebdc6a33c 100644 --- a/backend/notification_v2/tests/test_send_clubbed.py +++ b/backend/notification_v2/tests/test_send_clubbed.py @@ -87,6 +87,10 @@ def test_permanent_pg_error_dead_letters(self): ): _send() # Permanent error → terminal DEAD_LETTER, no PENDING revert / refund. + # The dead-letter update is guarded to only touch rows still SENDING + # (same clobber-guard the transient-revert path asserts). + fkw = buf.objects.filter.call_args.kwargs + assert fkw["status"] == BufferStatus.SENDING.value ukw = buf.objects.filter.return_value.update.call_args.kwargs assert ukw == {"status": BufferStatus.DEAD_LETTER.value} From 1eb15c6ddd0f7cb52b3f024a284b1fca6c7f521e Mon Sep 17 00:00:00 2001 From: ali Date: Wed, 29 Jul 2026 18:00:33 +0530 Subject: [PATCH 06/11] UN-3445 [GATED-FEAT] Log serialization failures with breadcrumb in enqueue_task Move the _json_safe message construction inside the try/except so a serialization ValueError (now reachable via allow_nan=False) gets the same task/queue/org breadcrumb as a DB insert failure, instead of propagating context-free on the orchestrator path. Test asserts the breadcrumb fires and the DB insert is never reached. Co-Authored-By: Claude Opus 4.8 --- backend/pg_queue/producer.py | 52 +++++++++++++------------ backend/pg_queue/tests/test_producer.py | 8 +++- 2 files changed, 34 insertions(+), 26 deletions(-) diff --git a/backend/pg_queue/producer.py b/backend/pg_queue/producer.py index e37fdae2ce..028fddfb6d 100644 --- a/backend/pg_queue/producer.py +++ b/backend/pg_queue/producer.py @@ -104,33 +104,37 @@ def enqueue_task( f"[{FAIRNESS_MIN_PRIORITY}, {FAIRNESS_MAX_PRIORITY}]: {priority!r}" ) pg_queue = queue or DEFAULT_GENERAL_QUEUE - message: TaskPayload = { - "task_name": task_name, - "args": _json_safe(list(args) if args is not None else []), - "kwargs": _json_safe(dict(kwargs) if kwargs is not None else {}), - "queue": pg_queue, - # Coerce like args/kwargs/on_success/on_error: FairnessPayload may carry a - # UUID/enum/datetime that a plain JSONField insert can't serialise, which - # would raise at enqueue and drop the task on the PG path only. - "fairness": _json_safe(fairness) if fairness is not None else fairness, - } - # Each optional key is set only when present — keeps fire-and-forget rows - # byte-identical to before these fields existed. - if reply_key is not None: - message["reply_key"] = reply_key - # Continuation specs carry a nested callback ``kwargs`` dict that may hold a - # UUID/datetime — coerce like ``args``/``kwargs`` above, else the JSONField - # insert raises at dispatch time (caller-visible). - if on_success is not None: - message["on_success"] = _json_safe(on_success) - if on_error is not None: - message["on_error"] = _json_safe(on_error) - if task_id is not None: - message["task_id"] = task_id # Mirror the worker _enqueue_pg path: log the failure with breadcrumbs before # it propagates, so a DB/constraint/serialization error isn't mislabeled by - # the caller's broad handler. + # the caller's broad handler. Message construction (the _json_safe coercions) + # lives inside the try because serialization failures surface there — e.g. a + # NaN/Infinity float raises ValueError under allow_nan=False — and on the + # orchestrator path (no silent fallback) would otherwise drop the task with no + # task/queue/org breadcrumb for on-call. try: + message: TaskPayload = { + "task_name": task_name, + "args": _json_safe(list(args) if args is not None else []), + "kwargs": _json_safe(dict(kwargs) if kwargs is not None else {}), + "queue": pg_queue, + # Coerce like args/kwargs/on_success/on_error: FairnessPayload may carry + # a UUID/enum/datetime that a plain JSONField insert can't serialise, + # which would raise at enqueue and drop the task on the PG path only. + "fairness": _json_safe(fairness) if fairness is not None else fairness, + } + # Each optional key is set only when present — keeps fire-and-forget rows + # byte-identical to before these fields existed. + if reply_key is not None: + message["reply_key"] = reply_key + # Continuation specs carry a nested callback ``kwargs`` dict that may hold a + # UUID/datetime — coerce like ``args``/``kwargs`` above, else the JSONField + # insert raises at dispatch time (caller-visible). + if on_success is not None: + message["on_success"] = _json_safe(on_success) + if on_error is not None: + message["on_error"] = _json_safe(on_error) + if task_id is not None: + message["task_id"] = task_id row = PgQueueMessage.objects.create( queue_name=pg_queue, message=message, diff --git a/backend/pg_queue/tests/test_producer.py b/backend/pg_queue/tests/test_producer.py index e812965c63..37cdb89b49 100644 --- a/backend/pg_queue/tests/test_producer.py +++ b/backend/pg_queue/tests/test_producer.py @@ -110,11 +110,15 @@ def test_json_safe_rejects_nan_with_value_error(self): # NaN would slip past the default lenient encoder and only fail at the # jsonb insert (DataError); allow_nan=False surfaces it as a ValueError at # the enqueue seam so the notification dispatcher can dead-letter it. + # The coercion runs inside the try, so the failure is logged with the + # task/queue/org breadcrumb (not dropped silently) and never reaches the + # DB insert. nan_kwargs = {"score": float("nan")} - with patch(_MODEL) as model: - model.objects.create.return_value = MagicMock(msg_id=1) + with patch(_MODEL) as model, patch.object(producer, "logger") as log: with pytest.raises(ValueError): producer.enqueue_task(task_name="t", queue="celery", kwargs=nan_kwargs) + model.objects.create.assert_not_called() + log.exception.assert_called_once() def test_enqueue_failure_logs_and_propagates(self): with patch(_MODEL) as model: From 533cf7847be437d2e3fdcb0edb502408b7c0f221 Mon Sep 17 00:00:00 2001 From: ali Date: Fri, 31 Jul 2026 09:31:47 +0530 Subject: [PATCH 07/11] UN-3445 [GATED-FEAT] Make PG consumer max_attempts env-configurable (default 5) + enrich poison-drop log Co-Authored-By: Claude Opus 4.8 --- workers/queue_backend/pg_queue/consumer.py | 10 +++-- workers/tests/test_pg_queue_consumer.py | 48 +++++++++++++++++++++- 2 files changed, 54 insertions(+), 4 deletions(-) diff --git a/workers/queue_backend/pg_queue/consumer.py b/workers/queue_backend/pg_queue/consumer.py index 3e016d2284..75a016d732 100644 --- a/workers/queue_backend/pg_queue/consumer.py +++ b/workers/queue_backend/pg_queue/consumer.py @@ -789,13 +789,17 @@ def _drop_poison_message( """ execution_id, organization_id = _pipeline_identity(payload) logger.error( - "PG-queue consumer: task %r (msg_id=%s) exceeded max_attempts=%s " - "(read_ct=%s, execution_id=%s) — poison; full payload: %r", + "PG-queue consumer: poison-dropped task %r after %s attempts " + "(msg_id=%s, read_ct=%s, queue=%s, execution_id=%s, org_id=%s) — " + "exceeded max_attempts=%s; full payload: %r", task_name, + message.read_ct, message.msg_id, - self.max_attempts, message.read_ct, + payload.get("queue"), execution_id, + organization_id or None, + self.max_attempts, payload, ) # Failure channel (request-reply / on_error) → surface there, then drop. diff --git a/workers/tests/test_pg_queue_consumer.py b/workers/tests/test_pg_queue_consumer.py index 919e94e101..636e763c93 100644 --- a/workers/tests/test_pg_queue_consumer.py +++ b/workers/tests/test_pg_queue_consumer.py @@ -272,6 +272,23 @@ def test_marks_error_then_deletes(self, monkeypatch): client.delete.assert_called_once_with(5) # then dropped client.set_vt.assert_not_called() + def test_poison_log_enriched_with_org_and_read_ct(self, monkeypatch, caplog): + # Q2 debuggability: the poison drop logs a clear "poison-dropped" message + # carrying read_ct + org_id so a dropped task is greppable/Sentry-visible. + monkeypatch.setattr( + "queue_backend.pg_queue.recovery.mark_execution_error", + lambda *a, **k: True, + ) + client = MagicMock() + client.read.return_value = [self._poison(read_ct=6)] + with caplog.at_level(logging.ERROR, logger="queue_backend.pg_queue.consumer"): + PgQueueConsumer( + ["q"], client=client, api_client=MagicMock(), max_attempts=5 + ).poll_once() + assert "poison-dropped" in caplog.text + assert "read_ct=6" in caplog.text + assert "org_id=org-1" in caplog.text # org surfaced for the dropped task + def test_positional_orchestration_poison_marks_error(self, monkeypatch): # H2 regression: a poisoned async_execute_bin carries execution_id # POSITIONALLY (args[2]) with no _barrier_context and — since its poison @@ -978,7 +995,8 @@ class TestRecordTaskStatus: pg_task_result so the REST PromptStudio.task_status poll resolves under PG. completed unless the run raised (error) or the executor reported success=False (completed rows are status-only; failed rows carry the executor error text). TTL'd - + best-effort — never wedges the ack.""" + + best-effort — never wedges the ack. + """ _RB = "queue_backend.pg_queue.consumer.PgResultBackend" _RET = 86400 @@ -1104,6 +1122,34 @@ def test_env_wires_lease_seconds(self, monkeypatch): c = mod.build_consumer_from_env() assert c.lease_seconds == 77 + def test_env_wires_max_attempts(self, monkeypatch): + # Per-worker override: an execution worker keeps the default (at-least-once), + # the interactive orchestrator sets 1 (at-most-once, no LLM re-compute). + from queue_backend.pg_queue import consumer as mod + + monkeypatch.setenv("WORKER_PG_QUEUE_CONSUMER_MAX_ATTEMPTS", "1") + with patch.object(mod, "PgQueueClient"): # no real DB connection + c = mod.build_consumer_from_env() + assert c.max_attempts == 1 + + def test_max_attempts_defaults_to_5_when_unset(self, monkeypatch): + from queue_backend.pg_queue import consumer as mod + from queue_backend.pg_queue.consumer import _DEFAULT_MAX_ATTEMPTS + + monkeypatch.delenv("WORKER_PG_QUEUE_CONSUMER_MAX_ATTEMPTS", raising=False) + with patch.object(mod, "PgQueueClient"): # no real DB connection + c = mod.build_consumer_from_env() + # The code default stays 5 so execution workers remain at-least-once. + assert c.max_attempts == _DEFAULT_MAX_ATTEMPTS == 5 + + def test_env_invalid_max_attempts_raises_named_error(self, monkeypatch): + from queue_backend.pg_queue import consumer as mod + + monkeypatch.setenv("WORKER_PG_QUEUE_CONSUMER_MAX_ATTEMPTS", "abc") + with patch.object(mod, "PgQueueClient"): + with pytest.raises(ValueError, match="WORKER_PG_QUEUE_CONSUMER_MAX_ATTEMPTS"): + mod.build_consumer_from_env() + def test_poll_claims_with_lease_not_vt(self): client = MagicMock() client.read.return_value = [] From a375e466a2aed57be88c009f1c889f5b73cd1896 Mon Sep 17 00:00:00 2001 From: ali Date: Fri, 31 Jul 2026 16:18:03 +0530 Subject: [PATCH 08/11] UN-3445 [GATED-FEAT] PG webhook: force raise_on_final_failure=False so terminal failure doesn't redeliver MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit On the PG branch of dispatch_webhook_notification, kwargs were forwarded verbatim including raise_on_final_failure=True. On PG that re-raise means the OPPOSITE of Celery: the worker already marks the buffers DEAD_LETTER, so re-raising only leaves the row for vt-expiry redelivery — re-POSTing the subscriber up to max_attempts times and tripping a false poison-drop. Override it to False on the PG branch only (the Celery branch keeps kwargs verbatim — byte-identical) so a terminal failure returns None -> the consumer acks -> single POST, matching Celery's external behaviour. Also: the poison-drop log test now asserts the source queue is surfaced (CodeRabbit). Co-Authored-By: Claude Opus 4.8 --- .../notification_v2/notification_dispatch.py | 13 +++++- .../tests/test_notification_dispatch.py | 45 ++++++++++++++++--- workers/tests/test_pg_queue_consumer.py | 7 ++- 3 files changed, 56 insertions(+), 9 deletions(-) diff --git a/backend/notification_v2/notification_dispatch.py b/backend/notification_v2/notification_dispatch.py index a7331a1618..90c0787b2d 100644 --- a/backend/notification_v2/notification_dispatch.py +++ b/backend/notification_v2/notification_dispatch.py @@ -87,12 +87,23 @@ def dispatch_webhook_notification( # Use the shared is_pg_transport() — the single source for "what counts as PG # transport" — rather than opening a second comparison site. if is_pg_transport(transport): + # ``raise_on_final_failure`` is Celery semantics: on the Celery path a + # terminal failure re-raises so the task ends FAILURE and link_error + # dead-letters the rows, with NO message-level redelivery. On the PG + # consumer the same re-raise means the OPPOSITE — the worker already marked + # the buffers DEAD_LETTER directly, so re-raising only leaves the row for + # vt-expiry redelivery (bounded by max_attempts), re-POSTing the subscriber + # on every redelivery and tripping a false poison-drop. Force it False on the + # PG branch so a terminal failure returns None → the consumer acks (deletes) + # the row → single pass, matching Celery's external behaviour. (Celery branch + # below keeps kwargs verbatim — byte-identical.) + pg_kwargs = {**kwargs, "raise_on_final_failure": False} try: enqueue_task( task_name=WEBHOOK_NOTIFICATION_TASK, queue=queue, args=args, - kwargs=kwargs, + kwargs=pg_kwargs, # enqueue_task types org_id as str; None→"" here satisfies that # type only, not runtime — enqueue_task itself re-coerces # ``org_id or ""`` at insert (producer.py), so this has no runtime diff --git a/backend/notification_v2/tests/test_notification_dispatch.py b/backend/notification_v2/tests/test_notification_dispatch.py index e53514d5de..55050779aa 100644 --- a/backend/notification_v2/tests/test_notification_dispatch.py +++ b/backend/notification_v2/tests/test_notification_dispatch.py @@ -2,8 +2,10 @@ ``resolve_transport`` + ``enqueue_task`` are patched on the module, so no Flipt / DB is needed — these pin the routing contract: PG when the flag resolves PG, -Celery otherwise (fail-closed), with byte-identical args/kwargs/queue on both -paths. +Celery otherwise (fail-closed), with identical args/queue on both paths. The +kwargs are identical EXCEPT ``raise_on_final_failure``, which is forced ``False`` +on the PG branch (the re-raise means "redeliver" on PG, not "dead-letter" — see +``test_pg_forces_raise_on_final_failure_false``). """ from __future__ import annotations @@ -49,7 +51,9 @@ def test_routes_to_pg_when_flag_resolves_pg(self): assert kwargs["task_name"] == "send_webhook_notification" assert kwargs["queue"] == _QUEUE assert kwargs["args"] == _ARGS - assert kwargs["kwargs"] == _KWARGS + # Every kwarg is forwarded verbatim EXCEPT raise_on_final_failure, which the + # PG branch forces False (see test_pg_forces_raise_on_final_failure_false). + assert kwargs["kwargs"] == {**_KWARGS, "raise_on_final_failure": False} assert kwargs["org_id"] == "org_x" # The minted PG task id is returned and threaded into the enqueue row. assert kwargs["task_id"] == task_id @@ -119,8 +123,10 @@ def test_resolve_transport_buckets_by_minted_dispatch_id(self): assert resolve.call_args.kwargs["execution_id"] == task_id assert enqueue.call_args.kwargs["task_id"] == task_id - def test_args_and_kwargs_identical_on_both_paths(self): - # The consumer must behave the same regardless of transport. + def test_args_and_queue_identical_on_both_paths(self): + # The consumer must behave the same regardless of transport — args/queue are + # byte-identical, and kwargs match apart from the transport-specific + # raise_on_final_failure override (asserted separately below). celery = MagicMock() with patch.object(nd, "resolve_transport", return_value="celery"): _dispatch(celery) @@ -132,5 +138,32 @@ def test_args_and_kwargs_identical_on_both_paths(self): ): _dispatch(celery2) assert enqueue.call_args.kwargs["args"] == celery_call["args"] - assert enqueue.call_args.kwargs["kwargs"] == celery_call["kwargs"] assert enqueue.call_args.kwargs["queue"] == celery_call["queue"] + # kwargs differ ONLY by the PG-forced raise_on_final_failure flag. + pg_kwargs = enqueue.call_args.kwargs["kwargs"] + assert pg_kwargs == {**celery_call["kwargs"], "raise_on_final_failure": False} + + def test_pg_forces_raise_on_final_failure_false(self): + # Regression (UN-3753): on the PG consumer a re-raise on retry exhaustion + # leaves the row for vt-expiry redelivery — the subscriber would be re-POSTed + # up to max_attempts times AND a false poison-drop would be logged. The + # dispatch seam must override raise_on_final_failure -> False on the PG branch + # (worker already marks buffers DEAD_LETTER directly) so the task returns None + # -> the consumer acks -> the endpoint is hit once, matching Celery's external + # behaviour. The Celery branch keeps the caller's True verbatim. + assert _KWARGS["raise_on_final_failure"] is True # guard the fixture premise + celery = MagicMock() + with ( + patch.object(nd, "resolve_transport", return_value="pg_queue"), + patch.object(nd, "enqueue_task", return_value=1) as enqueue, + ): + _dispatch(celery) + assert enqueue.call_args.kwargs["kwargs"]["raise_on_final_failure"] is False + # The caller's dict is not mutated in place (a fresh dict is enqueued). + assert _KWARGS["raise_on_final_failure"] is True + + celery2 = MagicMock() + with patch.object(nd, "resolve_transport", return_value="celery"): + _dispatch(celery2) + sent_kwargs = celery2.send_task.call_args.kwargs["kwargs"] + assert sent_kwargs["raise_on_final_failure"] is True diff --git a/workers/tests/test_pg_queue_consumer.py b/workers/tests/test_pg_queue_consumer.py index 636e763c93..cd0b641299 100644 --- a/workers/tests/test_pg_queue_consumer.py +++ b/workers/tests/test_pg_queue_consumer.py @@ -274,13 +274,15 @@ def test_marks_error_then_deletes(self, monkeypatch): def test_poison_log_enriched_with_org_and_read_ct(self, monkeypatch, caplog): # Q2 debuggability: the poison drop logs a clear "poison-dropped" message - # carrying read_ct + org_id so a dropped task is greppable/Sentry-visible. + # carrying read_ct + org_id + queue so a dropped task is greppable/ + # Sentry-visible and its source queue is identifiable at a glance. monkeypatch.setattr( "queue_backend.pg_queue.recovery.mark_execution_error", lambda *a, **k: True, ) client = MagicMock() - client.read.return_value = [self._poison(read_ct=6)] + payload = {**_callback_payload(), "queue": "agentic_callback"} + client.read.return_value = [self._poison(payload=payload, read_ct=6)] with caplog.at_level(logging.ERROR, logger="queue_backend.pg_queue.consumer"): PgQueueConsumer( ["q"], client=client, api_client=MagicMock(), max_attempts=5 @@ -288,6 +290,7 @@ def test_poison_log_enriched_with_org_and_read_ct(self, monkeypatch, caplog): assert "poison-dropped" in caplog.text assert "read_ct=6" in caplog.text assert "org_id=org-1" in caplog.text # org surfaced for the dropped task + assert "queue=agentic_callback" in caplog.text # source queue surfaced def test_positional_orchestration_poison_marks_error(self, monkeypatch): # H2 regression: a poisoned async_execute_bin carries execution_id From b46e09e03a93af69af36d6873943cea3cdbd692b Mon Sep 17 00:00:00 2001 From: ali Date: Mon, 3 Aug 2026 21:06:22 +0530 Subject: [PATCH 09/11] UN-3445 [GATED-FEAT] Address PR #2217 review: PG notification terminal-branch, task_id passthrough, credential redaction All 14 review findings. Flag-off/Celery stays byte-identical throughout. HIGH - raise_on_final_failure override was unreachable for max_retries>=1: under the consumer's eager task.apply(throw=True) request.retries is always 0, so the in-task retry guard fires on the FIRST failure and Retry propagates out of apply() - the terminal branch never ran, buffers were never dead-lettered, and the row was left for vt-expiry redelivery (re-POSTing the subscriber each time). Also force max_retries=0 on the PG branch so the terminal branch is reached. New workers-side test drives the REAL task through apply() and asserts one POST + one DEAD_LETTER mark, and pins the old broken shape as a regression guard. HIGH - docstring claimed a PG notifications consumer drains the queue; none exists in this repo. Reworded as an explicit deployment prerequisite (flag must stay off until one is deployed, else buffered webhooks are enqueued and silently lost). Idempotency - the consumer now passes the payload's stable task_id to task.apply(). Without it Celery mints a fresh uuid per delivery, so every guard keyed on request.id (fan-out claims, generation_task_id, task-complete markers) deduped nothing across a redelivery. Security - the poison drop logged the full payload, leaking subscriber Authorization / API-key headers and customer webhook bodies to stdout and Sentry. Added _redact_payload (masks auth-ish keys, summarises the body); routing metadata kept. Also: bounded the transient attempt-refund so NOTIFICATION_MAX_DISPATCH_ATTEMPTS is reachable (an unconditional refund made termination impossible); DispatchResult now carries the transport so the ramp metric can compare PG vs Celery; to_directory rejects pluggable types; corrected the empty-registry rationale (app.tasks auto-finalizes) and the stale link_error premises; fixed the conftest defect that forced a leaf test to mutate global Celery state; parametrized NaN/inf coverage over args/kwargs/fairness and the whole WorkerType->directory mapping; asserted the refund and buffer_row_ids that comments claimed but tests never checked. No regression: workers suite 208 failed / 1239 passed before, 208 failed / 1243 passed after (the 208 are pre-existing env failures - no DB, no prometheus_client). Co-Authored-By: Claude Opus 4.8 --- backend/notification_v2/internal_api_views.py | 46 ++++++-- .../notification_v2/notification_dispatch.py | 106 ++++++++++++----- .../tests/test_notification_dispatch.py | 65 +++++++---- .../tests/test_send_clubbed.py | 44 +++++-- backend/pg_queue/tests/test_producer.py | 45 +++++-- workers/notification/tasks.py | 12 +- workers/queue_backend/pg_queue/consumer.py | 80 ++++++++++++- workers/shared/enums/worker_enums_base.py | 24 +++- workers/tests/conftest.py | 10 ++ .../tests/test_notification_pg_terminal.py | 110 ++++++++++++++++++ workers/tests/test_pg_queue_consumer.py | 68 +++++++++++ workers/tests/test_worker_enums_directory.py | 40 ++++++- workers/worker.py | 46 +++++--- 13 files changed, 586 insertions(+), 110 deletions(-) create mode 100644 workers/tests/test_notification_pg_terminal.py diff --git a/backend/notification_v2/internal_api_views.py b/backend/notification_v2/internal_api_views.py index 04d77b81cb..42d0eeac97 100644 --- a/backend/notification_v2/internal_api_views.py +++ b/backend/notification_v2/internal_api_views.py @@ -48,6 +48,15 @@ # Constants for error messages INTERNAL_SERVER_ERROR_MSG = "Internal server error" +# How many early transient-failure cycles get their SENDING claim refunded before the +# claim starts sticking. Refunding forever makes net attempt progress zero, so +# ``NOTIFICATION_MAX_DISPATCH_ATTEMPTS`` becomes unreachable and a permanently +# recurring failure (not just a brief broker blip) re-dispatches every flush tick with +# no termination condition. Small by design: a genuine blip clears well inside this, +# and anything still failing after it is treated as repetitive and allowed to age +# toward the cap. +NOTIFICATION_TRANSIENT_REFUND_LIMIT = 3 + def _load_execution(execution_id: str | None) -> WorkflowExecution | None: """Best-effort lookup; returns None on missing id or unknown row.""" @@ -581,7 +590,7 @@ def _send_clubbed( # org, else Celery (byte-identical to the prior send_task). resolve_transport # keys on the org STRING id, but the buffer/worker contract below uses the # org pk — hence _org_identifier(org_id) for routing, org_id in kwargs. - dispatch_webhook_notification( + dispatched = dispatch_webhook_notification( celery_app=celery_app, args=[url, body, headers, settings.NOTIFICATION_TIMEOUT], kwargs={ @@ -600,13 +609,18 @@ def _send_clubbed( queue="notifications", org_string_id=_org_identifier(org_id), ) + # transport= makes the rollout answerable from the logs: during a percentage + # ramp the question is "are PG-routed notifications succeeding at the same + # rate as Celery-routed ones?", which result=success alone cannot answer. logger.info( "metric=notification_batch_dispatched_total platform=%s result=success " - "org_id=%s webhook_url_hash=%s rows=%d", + "transport=%s org_id=%s webhook_url_hash=%s rows=%d task_id=%s", platform, + dispatched.transport, org_id, webhook_url_hash(url), len(buffer_ids), + dispatched.task_id, ) except PermanentDispatchError: # PG-ONLY permanent failure. From this path the only reachable cause is @@ -635,10 +649,19 @@ def _send_clubbed( ).update(status=BufferStatus.DEAD_LETTER.value) except Exception: # TRANSIENT transport/broker failure — revert to PENDING (outside the - # committed txn) so the next flush tick retries; refund the SENDING-claim - # attempt since nothing was queued or sent. Guard on SENDING so a row the - # worker already marked terminal (broker raised post-delivery) isn't + # committed txn) so the next flush tick retries. Guard on SENDING so a row + # the worker already marked terminal (broker raised post-delivery) isn't # resurrected into a duplicate. + # + # The refund is BOUNDED, not unconditional: refunding the SENDING claim on + # every cycle makes net progress zero, so `NOTIFICATION_MAX_DISPATCH_ATTEMPTS` + # can never be reached and a permanently-recurring failure re-renders and + # re-dispatches forever (emitting a traceback each tick). Refunding is still + # correct for a genuinely transient blip — nothing was queued or sent — so we + # keep it while attempts are low and let the cap take over once a failure has + # proven itself repetitive. This matters more now that this branch also + # catches non-permanent PG enqueue failures (e.g. a DataError that will recur + # identically), not just "RabbitMQ is briefly down". logger.exception( "metric=notification_batch_dispatched_total platform=%s " "result=broker_failure org_id=%s webhook_url_hash=%s rows=%d", @@ -647,13 +670,20 @@ def _send_clubbed( webhook_url_hash(url), len(buffer_ids), ) - NotificationBuffer.objects.filter( + reverted = NotificationBuffer.objects.filter( id__in=buffer_ids, status=BufferStatus.SENDING.value, - ).update( + ) + # Refund only the early attempts; past the limit the claim stands so the + # counter grows and the cap can eventually dead-letter the group. + reverted.filter( + dispatch_attempts__lte=NOTIFICATION_TRANSIENT_REFUND_LIMIT + ).update(dispatch_attempts=F("dispatch_attempts") - 1) + # Every claimed row goes back to PENDING for the next flush tick, refunded + # or not. + reverted.update( status=BufferStatus.PENDING.value, dispatched_at=None, - dispatch_attempts=F("dispatch_attempts") - 1, ) diff --git a/backend/notification_v2/notification_dispatch.py b/backend/notification_v2/notification_dispatch.py index 90c0787b2d..4b027182eb 100644 --- a/backend/notification_v2/notification_dispatch.py +++ b/backend/notification_v2/notification_dispatch.py @@ -4,19 +4,29 @@ :func:`resolve_transport` flag as the execution path: the PG queue when ``pg_queue_enabled`` for this org, else Celery. **Fail-closed** — with the gate off (the production default) it resolves to Celery, behaving exactly like the -prior unconditional ``celery_app.send_task`` (zero regression). On PG the task -is drained by the PG notification consumer on the ``notifications`` queue. - -The same ``args``/``kwargs``/``queue`` are forwarded on both paths; on PG they're -JSON-normalized by ``enqueue_task`` (UUIDs/datetimes → str), so the consumer sees -the same payload the Celery worker would. +prior unconditional ``celery_app.send_task`` (zero regression). + +DEPLOYMENT PREREQUISITE (PG path): on PG the task lands on the ``notifications`` +queue and requires a pg-queue consumer configured with +``WORKER_PG_QUEUE_CONSUMER_WORKER_TYPE=notification`` / +``WORKER_PG_QUEUE_CONSUMER_QUEUE=notifications``. **No such service exists in this +repo's compose yet** — none of the ``pg-queue-consumer`` services in +``docker/docker-compose.yaml`` polls ``notifications``, and ``run-worker.sh``'s +``PG_CONSUMER_ROLES`` has no notification role. Until one is deployed the flag must +stay off for the org: enqueued rows would sit undrained (no TTL sweep covers +``pg_queue_message``) and every buffered webhook for that org would be lost. + +``args`` and ``queue`` are forwarded verbatim on both paths. ``kwargs`` differ in +exactly two keys, forced on the PG branch only (see the inline note at the branch): +``raise_on_final_failure`` → ``False`` and ``max_retries`` → ``0``. On PG the +payload is additionally JSON-normalized by ``enqueue_task`` (UUIDs/datetimes → str). """ from __future__ import annotations import logging import uuid -from typing import Any +from typing import Any, NamedTuple from pg_queue.producer import enqueue_task from workflow_manager.workflow_v2.transport import resolve_transport @@ -30,6 +40,21 @@ WEBHOOK_NOTIFICATION_TASK = "send_webhook_notification" +# Transport labels for the dispatch metric. During a percentage ramp the one thing +# you need from the logs is "are PG-routed notifications succeeding at the same rate +# as Celery-routed ones?" — a bare task id can't answer that, since a minted PG id +# and a Celery AsyncResult id are indistinguishable to the caller. +PG_TRANSPORT = "pg_queue" +CELERY_TRANSPORT = "celery" + + +class DispatchResult(NamedTuple): + """Which transport actually took the dispatch, plus the resulting task id.""" + + transport: str + task_id: str + + class PermanentDispatchError(Exception): """A dispatch failure that would fail identically on every retry. @@ -49,17 +74,23 @@ def dispatch_webhook_notification( kwargs: dict[str, Any], queue: str, org_string_id: str | None, -) -> str: +) -> DispatchResult: """Dispatch ``send_webhook_notification`` on the resolved transport. - ``args``/``kwargs``/``queue`` are forwarded unchanged on both paths, so the - flag-off (Celery) path is byte-identical to the legacy ``send_task`` call. + ``args``/``queue`` are forwarded unchanged on both paths, and on the Celery + branch ``kwargs`` too — so the flag-off path is byte-identical to the legacy + ``send_task`` call. The PG branch overrides two retry-semantics kwargs (see + below); nothing else differs. Args: celery_app: Injected Celery app (the backend's ``celery_service.app``); passed in rather than imported so this seam stays trivially testable. args: Positional task args, forwarded verbatim. - kwargs: Keyword task args, forwarded verbatim. For the buffered path this + kwargs: Keyword task args. Forwarded as-is on Celery; on PG, + ``raise_on_final_failure`` is overridden to ``False`` and + ``max_retries`` to ``0`` (the in-task retry loop is a no-op under the + consumer's eager ``apply()`` — see the inline note). For the buffered + path this carries ``organization_id`` = the buffer's org **pk** (the worker's buffer-mark contract) — deliberately a DIFFERENT identifier from the ``org_string_id`` param below (the two must not be conflated). @@ -70,9 +101,10 @@ def dispatch_webhook_notification( empty) fails closed to Celery. Returns: - A task id string — the Celery ``AsyncResult`` id on the Celery path, or - the minted PG task id on the PG path. Returned for symmetry / any future - caller; the sole current caller (fire-and-forget) discards it. + A :class:`DispatchResult` carrying which transport took the dispatch and + the resulting task id (the Celery ``AsyncResult`` id, or the minted PG task + id). The transport is what makes the caller's dispatch metric answerable + during a ramp — the two id flavours are otherwise indistinguishable. """ # A buffered notification is a single fire-and-forget task with no natural # sticky entity, so mint a fresh id to drive Flipt's percentage bucketing and @@ -87,19 +119,32 @@ def dispatch_webhook_notification( # Use the shared is_pg_transport() — the single source for "what counts as PG # transport" — rather than opening a second comparison site. if is_pg_transport(transport): - # ``raise_on_final_failure`` is Celery semantics: on the Celery path a - # terminal failure re-raises so the task ends FAILURE and link_error - # dead-letters the rows, with NO message-level redelivery. On the PG - # consumer the same re-raise means the OPPOSITE — the worker already marked - # the buffers DEAD_LETTER directly, so re-raising only leaves the row for - # vt-expiry redelivery (bounded by max_attempts), re-POSTing the subscriber - # on every redelivery and tripping a false poison-drop. Force it False on the - # PG branch so a terminal failure returns None → the consumer acks (deletes) - # the row → single pass, matching Celery's external behaviour. (Celery branch - # below keeps kwargs verbatim — byte-identical.) - pg_kwargs = {**kwargs, "raise_on_final_failure": False} + # Two PG-only kwarg overrides, both about the SAME thing: the consumer runs + # the task eagerly via ``task.apply(..., throw=True)``, where Celery's + # in-task retry loop does not work and its terminal branch is what we need. + # + # 1. ``max_retries`` → 0. Under ``apply()`` ``self.request.retries`` is ALWAYS + # 0, so the worker's ``if self.request.retries < max_retries`` guard + # (workers/notification/tasks.py) is true on the FIRST failure for any + # max_retries >= 1 and calls ``self.retry(...)``, which raises ``Retry``; + # with ``throw=True`` that propagates straight out of ``apply()`` and the + # terminal branch is never reached — so the buffers are never marked + # DEAD_LETTER and ``raise_on_final_failure`` is never even read. The + # consumer's ``except Exception`` then leaves the row for vt-expiry + # redelivery, re-POSTing the subscriber every time. Forcing 0 sends the + # task down the terminal branch on the first failure instead. + # 2. ``raise_on_final_failure`` → False. On BOTH transports the worker marks + # the buffers DEAD_LETTER over the internal API *before* it re-raises, so + # the re-raise is only a FAILURE-state signal for Celery monitoring (no + # redelivery there). On the PG consumer that same raise is treated as a + # failure and leaves the row for redelivery, so it must not raise. + # + # Together: one POST, buffers dead-lettered, task returns None → the consumer + # acks (deletes) the row. Retry spacing belongs to the PG layer, not the task. + # (The Celery branch below keeps kwargs verbatim — byte-identical.) + pg_kwargs = {**kwargs, "raise_on_final_failure": False, "max_retries": 0} try: - enqueue_task( + msg_id = enqueue_task( task_name=WEBHOOK_NOTIFICATION_TASK, queue=queue, args=args, @@ -117,16 +162,19 @@ def dispatch_webhook_notification( # A transient PG error (DB down) is NOT wrapped — it propagates as an # ordinary Exception into the caller's retry (revert-to-PENDING) path. raise PermanentDispatchError(str(exc)) from exc + # msg_id correlates this line with the producer's own ``msg_id=`` log, so a + # dropped notification can be traced across the seam/producer boundary. logger.info( - "Webhook notification enqueued on PG '%s' queue (task_id=%s)", + "Webhook notification enqueued on PG '%s' queue (task_id=%s msg_id=%s)", queue, dispatch_id, + msg_id, ) - return dispatch_id + return DispatchResult(transport=PG_TRANSPORT, task_id=dispatch_id) result = celery_app.send_task( WEBHOOK_NOTIFICATION_TASK, args=args, kwargs=kwargs, queue=queue, ) - return result.id + return DispatchResult(transport=CELERY_TRANSPORT, task_id=result.id) diff --git a/backend/notification_v2/tests/test_notification_dispatch.py b/backend/notification_v2/tests/test_notification_dispatch.py index 55050779aa..e08bf55ff1 100644 --- a/backend/notification_v2/tests/test_notification_dispatch.py +++ b/backend/notification_v2/tests/test_notification_dispatch.py @@ -3,9 +3,10 @@ ``resolve_transport`` + ``enqueue_task`` are patched on the module, so no Flipt / DB is needed — these pin the routing contract: PG when the flag resolves PG, Celery otherwise (fail-closed), with identical args/queue on both paths. The -kwargs are identical EXCEPT ``raise_on_final_failure``, which is forced ``False`` -on the PG branch (the re-raise means "redeliver" on PG, not "dead-letter" — see -``test_pg_forces_raise_on_final_failure_false``). +kwargs are identical EXCEPT two retry-semantics keys forced on the PG branch +(``raise_on_final_failure`` -> False, ``max_retries`` -> 0), because the consumer +runs tasks eagerly via ``apply()`` where the in-task retry loop cannot work — see +``test_pg_forces_terminal_branch_kwargs``. """ from __future__ import annotations @@ -45,15 +46,21 @@ def test_routes_to_pg_when_flag_resolves_pg(self): patch.object(nd, "resolve_transport", return_value="pg_queue"), patch.object(nd, "enqueue_task", return_value=42) as enqueue, ): - task_id = _dispatch(celery) + result = _dispatch(celery) + task_id = result.task_id + assert result.transport == nd.PG_TRANSPORT # rollout metric can tell them apart enqueue.assert_called_once() kwargs = enqueue.call_args.kwargs assert kwargs["task_name"] == "send_webhook_notification" assert kwargs["queue"] == _QUEUE assert kwargs["args"] == _ARGS - # Every kwarg is forwarded verbatim EXCEPT raise_on_final_failure, which the - # PG branch forces False (see test_pg_forces_raise_on_final_failure_false). - assert kwargs["kwargs"] == {**_KWARGS, "raise_on_final_failure": False} + # Every kwarg is forwarded verbatim EXCEPT the two retry-semantics keys the + # PG branch forces (see test_pg_forces_terminal_branch_kwargs). + assert kwargs["kwargs"] == { + **_KWARGS, + "raise_on_final_failure": False, + "max_retries": 0, + } assert kwargs["org_id"] == "org_x" # The minted PG task id is returned and threaded into the enqueue row. assert kwargs["task_id"] == task_id @@ -70,7 +77,8 @@ def test_routes_to_celery_when_flag_resolves_celery(self): "send_webhook_notification", args=_ARGS, kwargs=_KWARGS, queue=_QUEUE ) enqueue.assert_not_called() - assert result is celery.send_task.return_value.id + assert result.task_id is celery.send_task.return_value.id + assert result.transport == nd.CELERY_TRANSPORT def test_pg_enqueue_failure_propagates_with_no_celery_fallback(self): # No silent Celery fallback on a TRANSIENT PG failure — it propagates raw @@ -119,7 +127,7 @@ def test_resolve_transport_buckets_by_minted_dispatch_id(self): patch.object(nd, "resolve_transport", return_value="pg_queue") as resolve, patch.object(nd, "enqueue_task", return_value=1) as enqueue, ): - task_id = _dispatch(celery) + task_id = _dispatch(celery).task_id assert resolve.call_args.kwargs["execution_id"] == task_id assert enqueue.call_args.kwargs["task_id"] == task_id @@ -139,31 +147,46 @@ def test_args_and_queue_identical_on_both_paths(self): _dispatch(celery2) assert enqueue.call_args.kwargs["args"] == celery_call["args"] assert enqueue.call_args.kwargs["queue"] == celery_call["queue"] - # kwargs differ ONLY by the PG-forced raise_on_final_failure flag. + # kwargs differ ONLY by the two PG-forced retry-semantics keys. pg_kwargs = enqueue.call_args.kwargs["kwargs"] - assert pg_kwargs == {**celery_call["kwargs"], "raise_on_final_failure": False} - - def test_pg_forces_raise_on_final_failure_false(self): - # Regression (UN-3753): on the PG consumer a re-raise on retry exhaustion - # leaves the row for vt-expiry redelivery — the subscriber would be re-POSTed - # up to max_attempts times AND a false poison-drop would be logged. The - # dispatch seam must override raise_on_final_failure -> False on the PG branch - # (worker already marks buffers DEAD_LETTER directly) so the task returns None - # -> the consumer acks -> the endpoint is hit once, matching Celery's external - # behaviour. The Celery branch keeps the caller's True verbatim. + assert pg_kwargs == { + **celery_call["kwargs"], + "raise_on_final_failure": False, + "max_retries": 0, + } + + def test_pg_forces_terminal_branch_kwargs(self): + # Regression (UN-3753): the PG consumer runs the task eagerly via + # ``task.apply(..., throw=True)``, where the in-task retry loop cannot work: + # * max_retries >= 1 -> the worker's ``request.retries < max_retries`` guard + # is true on the FIRST failure (retries is always 0 under apply()), so it + # raises Retry, which propagates out of apply() -- the terminal branch + # (mark DEAD_LETTER + honour raise_on_final_failure) never runs and the row + # is left for vt-expiry redelivery, re-POSTing the subscriber each time. + # * a terminal re-raise is likewise treated as failure -> redelivery. + # So the seam must force BOTH max_retries=0 and raise_on_final_failure=False on + # the PG branch. The Celery branch keeps the caller's values verbatim. + # (The end-to-end behaviour these kwargs buy — one POST + one dead-letter mark + # through a real ``task.apply()`` — is asserted in + # workers/tests/test_notification_pg_terminal.py.) assert _KWARGS["raise_on_final_failure"] is True # guard the fixture premise + assert _KWARGS["max_retries"] == 3 # ... and that retries are configured celery = MagicMock() with ( patch.object(nd, "resolve_transport", return_value="pg_queue"), patch.object(nd, "enqueue_task", return_value=1) as enqueue, ): _dispatch(celery) - assert enqueue.call_args.kwargs["kwargs"]["raise_on_final_failure"] is False + pg_kwargs = enqueue.call_args.kwargs["kwargs"] + assert pg_kwargs["raise_on_final_failure"] is False + assert pg_kwargs["max_retries"] == 0 # The caller's dict is not mutated in place (a fresh dict is enqueued). assert _KWARGS["raise_on_final_failure"] is True + assert _KWARGS["max_retries"] == 3 celery2 = MagicMock() with patch.object(nd, "resolve_transport", return_value="celery"): _dispatch(celery2) sent_kwargs = celery2.send_task.call_args.kwargs["kwargs"] assert sent_kwargs["raise_on_final_failure"] is True + assert sent_kwargs["max_retries"] == 3 diff --git a/backend/notification_v2/tests/test_send_clubbed.py b/backend/notification_v2/tests/test_send_clubbed.py index 4ebdc6a33c..b296edfdf7 100644 --- a/backend/notification_v2/tests/test_send_clubbed.py +++ b/backend/notification_v2/tests/test_send_clubbed.py @@ -19,18 +19,24 @@ from unittest.mock import patch +from django.conf import settings from notification_v2 import internal_api_views as views from notification_v2.enums import BufferStatus +_URL = "https://hook.test" +_BODY = {"text": "x"} +_HEADERS = {"Content-Type": "application/json"} +_BUFFER_IDS = ["b1"] + def _send(**overrides): kw = { - "url": "https://hook.test", - "body": {"text": "x"}, - "headers": {"Content-Type": "application/json"}, + "url": _URL, + "body": _BODY, + "headers": _HEADERS, "platform": "SLACK", "max_retries": 3, - "buffer_ids": ["b1"], + "buffer_ids": _BUFFER_IDS, "org_id": 7, } kw.update(overrides) @@ -50,6 +56,13 @@ def test_routes_string_id_and_keeps_pk_in_kwargs(self): assert call["org_string_id"] == "org-str-id" assert call["kwargs"]["organization_id"] == 7 assert call["queue"] == "notifications" + # args carry the full webhook invocation, in order. + assert call["args"] == [_URL, _BODY, _HEADERS, settings.NOTIFICATION_TIMEOUT] + # buffer_row_ids is what REPLACED the Celery link/link_error callbacks: it is + # how the worker knows which rows to mark DISPATCHED / DEAD_LETTER. Drop it + # and every row strands in SENDING until the reclaim lease expires — silent, + # visible only as a growing SENDING backlog. + assert call["kwargs"]["buffer_row_ids"] == _BUFFER_IDS class TestSendClubbedFailureRecovery: @@ -64,13 +77,30 @@ def test_transient_failure_reverts_sending_rows_to_pending(self): patch.object(views, "NotificationBuffer") as buf, ): _send(buffer_ids=["b1", "b2"]) - # Guarded on SENDING, reverted to PENDING (attempt refunded). + # Guarded on SENDING, reverted to PENDING. fkw = buf.objects.filter.call_args.kwargs assert fkw["status"] == BufferStatus.SENDING.value assert set(fkw["id__in"]) == {"b1", "b2"} + # Every claimed row goes back to PENDING for the next tick. ukw = buf.objects.filter.return_value.update.call_args.kwargs - assert ukw["status"] == BufferStatus.PENDING.value - assert ukw["dispatched_at"] is None + assert ukw == { + "status": BufferStatus.PENDING.value, + "dispatched_at": None, + } + # The refund must be ASSERTED, not merely described in a comment: without it + # a transient outage burns no attempt (claim +1, refund -1 = net zero), so + # NOTIFICATION_MAX_DISPATCH_ATTEMPTS is unreachable and a permanently + # recurring failure re-dispatches every flush tick forever. It is also + # BOUNDED — only rows still under the limit are refunded, so a failure that + # keeps recurring eventually ages into the cap. + refund_qs = buf.objects.filter.return_value.filter + assert ( + refund_qs.call_args.kwargs["dispatch_attempts__lte"] + == views.NOTIFICATION_TRANSIENT_REFUND_LIMIT + ) + refund_kw = refund_qs.return_value.update.call_args.kwargs + assert set(refund_kw) == {"dispatch_attempts"} + assert "dispatch_attempts" in str(refund_kw["dispatch_attempts"]) # F(...) - 1 def test_permanent_pg_error_dead_letters(self): # The PG path surfaces a permanent enqueue failure as PermanentDispatchError; diff --git a/backend/pg_queue/tests/test_producer.py b/backend/pg_queue/tests/test_producer.py index 37cdb89b49..c94d219f01 100644 --- a/backend/pg_queue/tests/test_producer.py +++ b/backend/pg_queue/tests/test_producer.py @@ -5,6 +5,7 @@ """ import datetime +import logging import uuid from unittest.mock import MagicMock, patch @@ -106,19 +107,39 @@ def test_json_safe_coerces_datetime(self): when = model.objects.create.call_args.kwargs["message"]["kwargs"]["when"] assert isinstance(when, str) and "2026-06-18" in when - def test_json_safe_rejects_nan_with_value_error(self): - # NaN would slip past the default lenient encoder and only fail at the - # jsonb insert (DataError); allow_nan=False surfaces it as a ValueError at + @pytest.mark.parametrize( + "bad", [float("nan"), float("inf"), float("-inf")], ids=["nan", "inf", "-inf"] + ) + @pytest.mark.parametrize("slot", ["args", "kwargs", "fairness"]) + def test_json_safe_rejects_non_finite_floats(self, bad, slot, caplog): + # A non-finite float slips past the default lenient encoder and only fails at + # the jsonb insert (DataError); allow_nan=False surfaces it as a ValueError at # the enqueue seam so the notification dispatcher can dead-letter it. - # The coercion runs inside the try, so the failure is logged with the - # task/queue/org breadcrumb (not dropped silently) and never reaches the - # DB insert. - nan_kwargs = {"score": float("nan")} - with patch(_MODEL) as model, patch.object(producer, "logger") as log: - with pytest.raises(ValueError): - producer.enqueue_task(task_name="t", queue="celery", kwargs=nan_kwargs) - model.objects.create.assert_not_called() - log.exception.assert_called_once() + # + # Every _json_safe-coerced slot is covered, not just kwargs: the scenario this + # guard exists for is a webhook BODY carrying a non-finite float, and + # notification_dispatch puts the body in args[1]. inf/-inf are rejected by + # allow_nan=False exactly like nan. + payloads = { + "args": {"args": ["url", {"score": bad}]}, + "kwargs": {"kwargs": {"score": bad}}, + "fairness": {"fairness": {"org_id": "o", "weight": bad}}, + } + with patch(_MODEL) as model: + with caplog.at_level(logging.ERROR, logger=producer.logger.name): + with pytest.raises(ValueError): + producer.enqueue_task( + task_name="send_webhook_notification", + queue="notifications", + org_id="org-1", + **payloads[slot], + ) + model.objects.create.assert_not_called() # never reaches the DB insert + # The breadcrumb is the point of moving coercion inside the try — assert the + # rendered record actually carries it, not merely that .exception() was hit. + assert "send_webhook_notification" in caplog.text + assert "notifications" in caplog.text + assert "org-1" in caplog.text def test_enqueue_failure_logs_and_propagates(self): with patch(_MODEL) as model: diff --git a/workers/notification/tasks.py b/workers/notification/tasks.py index 41d76f5cf9..f1c780ae0c 100644 --- a/workers/notification/tasks.py +++ b/workers/notification/tasks.py @@ -252,9 +252,15 @@ def send_webhook_notification( retry_delay: The delay between retries in seconds platform: Platform type from notification config (SLACK, API, etc.) raise_on_final_failure: When True, re-raise on retry exhaustion so the - task ends in FAILURE and any Celery ``link_error`` callback runs - (used by the clubbed/buffered dispatch to dead-letter the rows). - When False (default), preserve the legacy "return None" behavior. + task ends in FAILURE (a Celery monitoring signal). NOTE: the buffer + rows are dead-lettered by ``_mark_buffer_outcome(dispatched=False)`` + over the internal API *before* this re-raise, on BOTH transports — + ``link``/``link_error`` callbacks are no longer wired (they routed to + the ``celery`` queue and were dropped as "unregistered task"), so the + raise itself dead-letters nothing. When False (default), preserve the + legacy "return None" behavior. The PG dispatch seam forces this False + (and ``max_retries`` 0) because under the consumer's eager + ``task.apply()`` a raise means redelivery, not a FAILURE state. Returns: None (matches original behavior) diff --git a/workers/queue_backend/pg_queue/consumer.py b/workers/queue_backend/pg_queue/consumer.py index 75a016d732..241a3d291e 100644 --- a/workers/queue_backend/pg_queue/consumer.py +++ b/workers/queue_backend/pg_queue/consumer.py @@ -204,6 +204,62 @@ def _pipeline_identity(payload: TaskPayload) -> tuple[str | None, str]: ) +# Header/kwarg names whose VALUES are secrets. Substring match on a lowercased key, +# so ``Authorization``/``X-Api-Key``/``auth_token``/``client_secret`` all match. +_SECRET_KEY_PARTS = ("authorization", "api-key", "api_key", "apikey", "token", "secret") + +# Payload slots that carry customer data rather than routing metadata. Logged as a +# type+size summary so a drop stays debuggable without dumping document content. +_BULKY_ARG_INDEX = 1 # args[1] is the webhook body on the notification payload + + +def _redact_secrets(value: Any) -> Any: + """Recursively mask secret-valued keys in dicts/lists; other values pass through.""" + if isinstance(value, dict): + return { + k: ( + "***REDACTED***" + if isinstance(k, str) and any(p in k.lower() for p in _SECRET_KEY_PARTS) + else _redact_secrets(v) + ) + for k, v in value.items() + } + if isinstance(value, list): + return [_redact_secrets(v) for v in value] + return value + + +def _redact_payload(payload: dict) -> dict: + """A log-safe view of a queue payload. + + Two exposures this closes, both introduced in practice once credential-bearing + notification payloads started flowing through the PG consumer: + + * **Secrets** — ``build_webhook_headers`` writes the subscriber's bearer token / + API key into the header dict carried in ``args``/``kwargs``. Masked by key name. + * **Customer data** — the clubbed webhook body is the pipeline's extracted output. + Replaced with a ```` summary rather than dumped. + + Routing metadata (task_name, queue, task_id, read counts) is preserved verbatim, + since that is what a poison drop is actually debugged from. + """ + safe = {k: v for k, v in payload.items() if k not in ("args", "kwargs")} + args = payload.get("args") + if isinstance(args, list): + redacted_args: list[Any] = [] + for i, a in enumerate(args): + if i == _BULKY_ARG_INDEX and isinstance(a, (dict, list, str)): + redacted_args.append(f"<{type(a).__name__} len={len(a)}>") + else: + redacted_args.append(_redact_secrets(a)) + safe["args"] = redacted_args + elif args is not None: + safe["args"] = _redact_secrets(args) + if payload.get("kwargs") is not None: + safe["kwargs"] = _redact_secrets(payload.get("kwargs")) + return safe + + class PgQueueConsumer: """Polls one PG queue, runs each claimed task in-process, acks on success.""" @@ -505,11 +561,22 @@ def _handle(self, message: QueueMessage) -> None: headers = {FAIRNESS_HEADER_NAME: fairness} if fairness else None # Renew the short lease while the (possibly long) task runs, so a dead # worker's claim expires fast but a live one is never redelivered. + # + # ``task_id=`` is load-bearing for idempotency: Celery's ``Task.apply`` + # does ``task_id = task_id or uuid()``, so WITHOUT this the task sees a + # FRESH ``self.request.id`` on every delivery. Any guard keyed on it + # (fan-out claim rows, ``generation_task_id`` uniqueness, task-complete + # markers) would then never collide across a redelivery and dedupe + # nothing. The producer already persisted a stable id in the payload — + # pass it through so redelivery re-runs with the SAME request id, matching + # Celery's own redelivery semantics. Falls back to Celery's uuid() when + # absent (never None, which apply() would reject). with self._lease_renewal(message.msg_id): eager = task.apply( args=payload.get("args") or [], kwargs=payload.get("kwargs") or {}, headers=headers, + task_id=payload.get("task_id") or None, throw=True, ) except Exception as exc: @@ -789,18 +856,23 @@ def _drop_poison_message( """ execution_id, organization_id = _pipeline_identity(payload) logger.error( - "PG-queue consumer: poison-dropped task %r after %s attempts " + "PG-queue consumer: poison-dropped task %r " "(msg_id=%s, read_ct=%s, queue=%s, execution_id=%s, org_id=%s) — " - "exceeded max_attempts=%s; full payload: %r", + "exceeded max_attempts=%s; payload: %r", task_name, - message.read_ct, message.msg_id, + # read_ct counts READS, not executions — a lease expiry with no run still + # increments it. Logged once, under its own name (an "after N attempts" + # phrasing duplicated this same value and misread as execution count). message.read_ct, payload.get("queue"), execution_id, + # NB: identifier flavour is payload-dependent — see _pipeline_identity. organization_id or None, self.max_attempts, - payload, + # Redacted: payloads carry subscriber credentials (webhook Authorization + # / API-key headers) and customer data. Never log them raw. + _redact_payload(payload), ) # Failure channel (request-reply / on_error) → surface there, then drop. if payload.get("reply_key") or payload.get("on_error"): diff --git a/workers/shared/enums/worker_enums_base.py b/workers/shared/enums/worker_enums_base.py index 97e4799768..65465659c0 100644 --- a/workers/shared/enums/worker_enums_base.py +++ b/workers/shared/enums/worker_enums_base.py @@ -61,14 +61,30 @@ def to_import_path(self) -> str: return f"{self.to_directory()}.tasks" def to_directory(self) -> str: - """Return the on-disk directory name for this (non-pluggable) worker. + """Return the on-disk directory name for this **non-pluggable** worker. Single source of truth for the naming-convention mapping: enum values use - underscores (Python module names), but a few on-disk dirs use hyphens - (e.g. ``api-deployment``). ``to_import_path`` builds on this, and the - file-path task loader in ``worker.py`` reads it directly rather than + underscores (Python module names), while ``API_DEPLOYMENT``'s on-disk dir + uses a hyphen (``api-deployment``). ``to_import_path`` builds on this, and + the file-path task loader in ``worker.py`` reads it directly rather than slicing the import path. + + Raises: + ValueError: if called on a pluggable worker type. Pluggable code lives + under ``pluggable_worker/``, so the mapping's passthrough + would return a bare, *plausible-looking* but wrong path. Today only + call ordering prevents that (``to_import_path`` guards with + ``is_pluggable()``, and ``worker.py`` early-returns for pluggable + workers before reaching here) — nothing enforced it. Since + ``worker.py`` now RAISES on a missing directory instead of logging, + a future caller that skips the guard would turn a silently-wrong + path into a boot crash-loop; fail loudly and locally instead. """ + if self.is_pluggable(): + raise ValueError( + f"{self!r} is a pluggable worker; its code lives under " + f"pluggable_worker/{self.value} — use to_import_path() instead." + ) directory_mapping = { WorkerType.API_DEPLOYMENT: "api-deployment", # All others use same name for directory and module diff --git a/workers/tests/conftest.py b/workers/tests/conftest.py index bd4b090b4c..faa90d7593 100644 --- a/workers/tests/conftest.py +++ b/workers/tests/conftest.py @@ -410,6 +410,16 @@ def _restore_current_celery_app(): """ from celery._state import default_app + # Guard against there being no default app at all: a module that imports no + # worker app (e.g. a pure-enum test running alone) would otherwise ERROR here on + # ``None.finalize()``. That previously forced such modules to call + # ``Celery(...).set_default()`` at import time as a workaround, which mutated + # process-global state for every later test in the session depending on + # collection order. Fix it here instead so leaf tests need no Celery app. + if default_app is None: + yield + return + default_app.finalize() default_app.set_current() try: diff --git a/workers/tests/test_notification_pg_terminal.py b/workers/tests/test_notification_pg_terminal.py new file mode 100644 index 0000000000..8c3774029d --- /dev/null +++ b/workers/tests/test_notification_pg_terminal.py @@ -0,0 +1,110 @@ +"""``send_webhook_notification`` under the PG consumer's eager ``task.apply()``. + +The PG consumer runs tasks with ``task.apply(..., throw=True)`` rather than through +a Celery worker. Under ``apply()`` Celery's in-task retry loop cannot work — +``self.request.retries`` is ALWAYS 0, so the task's +``if self.request.retries < max_retries`` guard is true on the FIRST failure for any +``max_retries >= 1`` and raises ``Retry``, which ``throw=True`` propagates straight +out of ``apply()``. The terminal branch (mark buffers DEAD_LETTER, honour +``raise_on_final_failure``) is then never reached, so the consumer sees an exception, +leaves the row for vt-expiry redelivery, and the subscriber is re-POSTed every time. + +That is why the dispatch seam (``backend/notification_v2/notification_dispatch.py``) +forces BOTH ``max_retries=0`` and ``raise_on_final_failure=False`` on the PG branch. +These tests drive the REAL task through ``apply()`` and assert the behaviour those +kwargs buy — one POST + one DEAD_LETTER mark, no raise — rather than the kwarg values +(which the seam's own unit tests already pin). +""" + +from __future__ import annotations + +from unittest.mock import patch + +import pytest +from celery.exceptions import Retry +from notification.tasks import send_webhook_notification + +_URL = "https://hook.test/x" +_BUFFER_IDS = ["b1", "b2"] +_ORG = 7 + + +class _AlwaysFailingProvider: + """Webhook provider whose send() always reports failure, counting attempts.""" + + def __init__(self) -> None: + self.posts = 0 + + def send(self, notification_data: dict) -> dict: + self.posts += 1 + return {"success": False, "message": "boom"} + + +def _run(*, max_retries: int, raise_on_final_failure: bool): + """Drive the real task through the consumer's exact call shape. + + Returns ``(provider, marks, raised)`` — POST count, the + ``_mark_buffer_outcome`` calls, and any exception propagated out of apply(). + """ + provider = _AlwaysFailingProvider() + marks: list[bool] = [] + with ( + patch( + "notification.tasks._get_webhook_provider_for_url", return_value=provider + ), + patch( + "notification.tasks._mark_buffer_outcome", + side_effect=lambda ids, org, *, dispatched: marks.append(dispatched), + ), + ): + raised: BaseException | None = None + try: + send_webhook_notification.apply( + args=[_URL, {"text": "hi"}, {"Content-Type": "application/json"}, 30], + kwargs={ + "max_retries": max_retries, + "retry_delay": 10, + "platform": None, + "raise_on_final_failure": raise_on_final_failure, + "buffer_row_ids": _BUFFER_IDS, + "organization_id": _ORG, + }, + throw=True, + ) + except BaseException as exc: # noqa: BLE001 - we assert on the type below + raised = exc + return provider, marks, raised + + +def test_pg_kwargs_give_one_post_and_one_dead_letter(): + # The seam's PG kwargs (max_retries=0, raise_on_final_failure=False): the task + # takes the terminal branch on the first failure, marks the buffers DEAD_LETTER, + # and returns None -> the consumer acks the row. Subscriber hit exactly once. + provider, marks, raised = _run(max_retries=0, raise_on_final_failure=False) + assert raised is None # nothing propagates -> consumer acks (no redelivery) + assert provider.posts == 1 # subscriber POSTed once, not once per redelivery + assert marks == [False] # exactly one DEAD_LETTER mark + + +def test_unforced_max_retries_would_escape_apply_without_dead_lettering(): + # Regression guard for the bug the override fixes: with the caller's original + # max_retries >= 1, apply() raises Retry out of the task, the terminal branch + # never runs, and NO buffer mark happens -> the consumer would redeliver and + # re-POST. This is what the seam prevents by forcing max_retries=0. + provider, marks, raised = _run(max_retries=3, raise_on_final_failure=False) + assert isinstance(raised, Retry) # escapes apply() as an exception + assert provider.posts == 1 # ... after a single POST on this delivery + assert marks == [] # buffers NEVER dead-lettered -> rows strand/redeliver + + +@pytest.mark.parametrize("raise_on_final", [True, False]) +def test_max_retries_zero_always_reaches_terminal_branch(raise_on_final): + # With max_retries=0 the terminal branch always runs, so the buffers are marked + # DEAD_LETTER regardless of raise_on_final_failure; the flag only decides whether + # the task additionally re-raises (a Celery FAILURE-state signal). + provider, marks, raised = _run( + max_retries=0, raise_on_final_failure=raise_on_final + ) + assert provider.posts == 1 + assert marks == [False] # dead-lettered on BOTH transports, before any re-raise + assert (raised is not None) is raise_on_final diff --git a/workers/tests/test_pg_queue_consumer.py b/workers/tests/test_pg_queue_consumer.py index cd0b641299..3e5d836dce 100644 --- a/workers/tests/test_pg_queue_consumer.py +++ b/workers/tests/test_pg_queue_consumer.py @@ -133,6 +133,33 @@ def test_fairness_header_rebuilt_for_run(self): assert kwargs["kwargs"] == {"k": "v"} assert kwargs["headers"] == {FAIRNESS_HEADER_NAME: fairness} + def test_payload_task_id_is_passed_to_apply(self): + # Idempotency-critical: Celery's Task.apply does ``task_id = task_id or + # uuid()``, so without an explicit task_id the task sees a FRESH request.id + # on every delivery and any guard keyed on it (fan-out claims, + # generation_task_id uniqueness, task-complete markers) dedupes nothing + # across a redelivery. The producer's stable id must reach apply(). + task = MagicMock() + app = MagicMock() + app.tasks.get.return_value = task + client = MagicMock() + client.read.return_value = [ + _msg(6, {"task_name": "t", "args": [], "kwargs": {}, "task_id": "stable-1"}) + ] + PgQueueConsumer(["q"], client=client, app=app).poll_once() + assert task.apply.call_args.kwargs["task_id"] == "stable-1" + + def test_missing_task_id_falls_back_to_celery_uuid(self): + # No task_id in the payload -> pass None so Celery mints one, rather than + # passing "" (which apply() would take literally as the request id). + task = MagicMock() + app = MagicMock() + app.tasks.get.return_value = task + client = MagicMock() + client.read.return_value = [_msg(6, {"task_name": "t", "args": [], "kwargs": {}})] + PgQueueConsumer(["q"], client=client, app=app).poll_once() + assert task.apply.call_args.kwargs["task_id"] is None + def test_ack_finding_no_row_warns(self, caplog): client = MagicMock() client.delete.return_value = False # row already gone (vt expired mid-run) @@ -291,6 +318,47 @@ def test_poison_log_enriched_with_org_and_read_ct(self, monkeypatch, caplog): assert "read_ct=6" in caplog.text assert "org_id=org-1" in caplog.text # org surfaced for the dropped task assert "queue=agentic_callback" in caplog.text # source queue surfaced + # read_ct is rendered ONCE (an earlier "after N attempts" phrasing printed + # the same value twice under a name that misread as an execution count). + assert caplog.text.count("6") >= 1 + assert "attempts (msg_id" not in caplog.text + + def test_poison_log_redacts_credentials_and_body(self, monkeypatch, caplog): + # Security: notification payloads carry the subscriber's Authorization / + # API-key headers and the clubbed webhook body (customer extracted data). + # A poison drop must not dump either to stdout/Sentry. + monkeypatch.setattr( + "queue_backend.pg_queue.recovery.mark_execution_error", + lambda *a, **k: True, + ) + payload = { + "task_name": "send_webhook_notification", + "queue": "notifications", + "args": [ + "https://hook.test", + {"secret_field": "customer extracted data"}, + {"Authorization": "Bearer super-secret", "X-Api-Key": "ak_live_123"}, + 30, + ], + "kwargs": {"organization_id": "org-1", "auth_token": "tok_abc"}, + } + client = MagicMock() + client.read.return_value = [self._poison(payload=payload, read_ct=6)] + with caplog.at_level(logging.ERROR, logger="queue_backend.pg_queue.consumer"): + PgQueueConsumer( + ["q"], client=client, api_client=MagicMock(), max_attempts=5 + ).poll_once() + # No secret value survives, by any route. + assert "super-secret" not in caplog.text + assert "ak_live_123" not in caplog.text + assert "tok_abc" not in caplog.text + # Customer body replaced by a type+size summary, not dumped. + assert "customer extracted data" not in caplog.text + assert "" in caplog.text + # Routing metadata IS preserved — that's what the drop is debugged from. + assert "send_webhook_notification" in caplog.text + assert "https://hook.test" in caplog.text + assert "REDACTED" in caplog.text def test_positional_orchestration_poison_marks_error(self, monkeypatch): # H2 regression: a poisoned async_execute_bin carries execution_id diff --git a/workers/tests/test_worker_enums_directory.py b/workers/tests/test_worker_enums_directory.py index 931fc45f1e..280a3ad59e 100644 --- a/workers/tests/test_worker_enums_directory.py +++ b/workers/tests/test_worker_enums_directory.py @@ -1,16 +1,22 @@ """WorkerType.to_directory() — the single source for the on-disk dir mapping (UN-3798). worker.py's file-path task loader reads to_directory() directly instead of slicing -to_import_path(); these pin the hyphen/underscore mapping and that to_import_path -is built on top of it, so the two can't drift. +to_import_path(); these pin the hyphen/underscore mapping and that to_import_path is +built on top of it, so the two can't drift. + +No Celery app is set up here on purpose: this is a pure-enum module, and the autouse +``_restore_current_celery_app`` fixture now tolerates a missing default app. An +earlier version called ``Celery(...).set_default()`` at import time to satisfy that +fixture, which mutated process-global Celery state for every later test in the +session depending on collection order. """ -from celery import Celery +from pathlib import Path + +import pytest from shared.enums.worker_enums_base import WorkerType -# The workers autouse conftest fixture finalizes celery's default_app around every -# test; this pure-enum test builds no Celery app of its own, so establish one. -Celery("test-worker-enums").set_default() +_WORKERS_ROOT = Path(__file__).resolve().parents[1] def test_to_directory_maps_underscored_value_to_hyphenated_dir(): @@ -28,3 +34,25 @@ def test_to_import_path_is_built_on_to_directory(): # directory naming lives in exactly one place. wt = WorkerType.API_DEPLOYMENT assert wt.to_import_path() == f"{wt.to_directory()}.tasks" + + +@pytest.mark.parametrize("wt", [w for w in WorkerType if not w.is_pluggable()]) +def test_every_worker_type_maps_to_an_existing_directory_with_tasks(wt): + # The failure mode UN-3798 exists for is enum-vs-disk DRIFT, and worker.py now + # RAISES on a missing directory (it previously logged and continued) — so drift + # is a container crash-loop across every replica, not a degraded worker. Pinning + # two members can't catch that; pin the whole mapping so adding a WorkerType + # without its directory is a red build instead of a production RuntimeError. + directory = _WORKERS_ROOT / wt.to_directory() + assert directory.is_dir(), f"{wt.value} -> {wt.to_directory()} does not exist" + assert (directory / "tasks.py").is_file(), f"{wt.to_directory()}/tasks.py missing" + + +@pytest.mark.parametrize("wt", [w for w in WorkerType if w.is_pluggable()]) +def test_to_directory_rejects_pluggable_types(wt): + # Pluggable code lives under pluggable_worker/, so the mapping's + # passthrough would return a wrong-but-plausible path. Only call ordering kept + # that unreachable; the precondition is now enforced. + with pytest.raises(ValueError, match="pluggable"): + wt.to_directory() + assert wt.to_import_path() == f"pluggable_worker.{wt.value}.tasks" diff --git a/workers/worker.py b/workers/worker.py index 295f5c9b00..077fbd9741 100755 --- a/workers/worker.py +++ b/workers/worker.py @@ -451,18 +451,18 @@ def load_worker_tasks(worker_type: WorkerType) -> None: PACKAGE import that RUNS the plugin's own registration code, e.g. ``from . import tasks`` with relative imports such as ``from .clients import ...``). That registration's ``@shared_task`` bindings - complete when the app FINALIZES, which can be *after* this point — so the - tasks are not necessarily bound to the app yet here. The generic file-path - load below is SKIPPED for pluggable workers not because they are already - bound, but because loading ``tasks.py`` under a bare ``"tasks"`` spec (no - parent package) would break the plugin's relative imports ("attempted - relative import with no known parent package"). + are applied when the app FINALIZES — which any read of ``app.tasks`` triggers + automatically (``finalize(auto=True)``), so they are in place by the time + anything inspects the registry. The generic file-path load below is SKIPPED for + pluggable workers because loading ``tasks.py`` under a bare ``"tasks"`` spec (no + parent package) would break the plugin's relative imports ("attempted relative + import with no known parent package") — not because of any binding-order + subtlety. Non-pluggable (top-level) workers use absolute imports; their ``tasks.py`` is - loaded by file path with the worker directory on ``sys.path``. Their tasks - bind eagerly here — there is no later finalize step to rescue them — so a - missing directory / ``tasks.py`` is a broken deploy and hard-fails rather - than booting a no-op worker. + loaded by file path with the worker directory on ``sys.path``. Their tasks bind + eagerly here, so a missing directory / ``tasks.py`` is a broken deploy and + hard-fails rather than booting a no-op worker. """ if worker_type.is_pluggable(): logger.info( @@ -496,12 +496,26 @@ def load_worker_tasks(worker_type: WorkerType) -> None: load_worker_tasks(worker_type) # A worker that boots with no registered tasks starts but silently processes -# nothing (Celery does not error on an empty registry). Surface that misconfig, -# but split by worker kind: pluggable tasks bind on app finalize (@shared_task / -# connect_on_app_finalize), which can be after this point, so a raise here would -# false-positive on a correctly-configured pluggable worker — only WARN. A -# non-pluggable worker has no later binding step (load_worker_tasks bound its -# tasks eagerly above), so an empty registry here is terminal — RAISE. +# nothing (Celery does not error on an empty registry). Surface that misconfig. +# +# NB: reading ``app.tasks`` below AUTO-FINALIZES the app (Celery's ``tasks`` +# property calls ``self.finalize(auto=True)``), so every pending ``@shared_task`` / +# ``connect_on_app_finalize`` registration has already been applied by the time the +# emptiness test is evaluated. The check is therefore exactly as accurate for +# pluggable workers as for non-pluggable ones — an earlier version of this comment +# claimed the opposite ("bind on finalize, which can be after this point") and used +# it to justify the split; that premise was wrong. Two real consequences: the app is +# finalized at import rather than at worker start, and an empty registry here is a +# genuine misconfiguration for BOTH kinds. +# +# The split is kept deliberately, on deployment-risk grounds rather than accuracy: +# a non-pluggable worker's tasks are bound eagerly by ``load_worker_tasks`` in this +# same module, so an empty registry is unambiguously terminal — RAISE. Pluggable +# workers register via plugin package import, whose availability is deployment +# dependent; hard-failing there would turn one absent/misconfigured plugin into a +# container crash-loop across the fleet, so it stays a loud WARN. Tightening this to +# a raise is worthwhile but needs a pass confirming every deployed pluggable worker +# registers at import time — tracked under UN-3445 hardening, not changed blind here. if not any(not name.startswith("celery.") for name in app.tasks): _empty_registry_msg = ( f"No non-celery tasks registered for worker '{worker_type.value}' " From 2a39e2694381052776c4d9a9eef3c42a361b1459 Mon Sep 17 00:00:00 2001 From: ali Date: Tue, 4 Aug 2026 11:21:47 +0530 Subject: [PATCH 10/11] UN-3893 [FIX] Make ConcurrencyMode a StrEnum so the StateStore guard survives a set env var and duplicate module copies MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit StateStore guards every set/get/clear with `cls.mode == ConcurrencyMode.THREAD` and raises RuntimeError("Unknown concurrency mode") otherwise. As a bare Enum that guard had two runtime failure modes, both invisible until production: 1. mode is read as os.environ.get("CONCURRENCY_MODE", ConcurrencyMode.THREAD), so SETTING the variable yields a plain str — and "thread" == ConcurrencyMode.THREAD is False for a bare Enum. Even the CORRECT value took the backend down. 2. The module exists in three copies (backend/utils, workers/shared/utils, workers/shared/infrastructure) and can be imported under more than one path in a merged OSS+cloud tree, producing two distinct ConcurrencyMode CLASSES. Members of different Enum classes never compare equal, so the guard raised on every call. (2) is what failed ~70 cloud integration tests across unrelated suites (test_pg_finalization_fixes, dashboard_metrics, manual_review_v2, statistics_service): the same OSS tests pass in OSS CI and cloud main is green — it only appeared once the cloud plugins were merged into the OSS tree. StrEnum members ARE strings, so both cases now compare by value and the defining class's identity stops mattering. Applied to all three copies, with a regression test that loads two copies under different names, asserts they are genuinely distinct classes, and asserts their members still compare equal — plus the env-var round trip. No behaviour change otherwise: the comparison is strictly more permissive, and the workers suite is unchanged at 208 failed / 1243 passed (pre-existing env failures: no DB, no prometheus_client). Co-Authored-By: Claude Opus 4.8 --- backend/utils/local_context.py | 25 ++++++- workers/shared/infrastructure/context.py | 25 ++++++- workers/shared/utils/local_context.py | 25 ++++++- .../tests/test_concurrency_mode_identity.py | 67 +++++++++++++++++++ 4 files changed, 136 insertions(+), 6 deletions(-) create mode 100644 workers/tests/test_concurrency_mode_identity.py diff --git a/backend/utils/local_context.py b/backend/utils/local_context.py index 0d11c0f6a1..2753fb4b4a 100644 --- a/backend/utils/local_context.py +++ b/backend/utils/local_context.py @@ -1,10 +1,31 @@ import os import threading -from enum import Enum +from enum import StrEnum from typing import Any -class ConcurrencyMode(Enum): +class ConcurrencyMode(StrEnum): + """Concurrency mode for :class:`StateStore`. + + ``StrEnum`` (not a bare ``Enum``) is load-bearing for TWO reasons, both of which + produced a hard ``RuntimeError: Unknown concurrency mode`` at runtime (UN-3893): + + 1. **A set env var used to break it.** ``mode`` is read as + ``os.environ.get("CONCURRENCY_MODE", ConcurrencyMode.THREAD)``, so setting the + variable yields a plain ``str``. Against a bare ``Enum``, ``"thread" == + ConcurrencyMode.THREAD`` is **False** — so even the CORRECT value took the + whole backend down. ``StrEnum`` members ARE strings, so this now compares equal. + 2. **Duplicate module identity.** This module exists in three places + (``backend/utils``, ``workers/shared/utils``, ``workers/shared/infrastructure``) + and can be imported under more than one path in a merged tree, producing two + distinct ``ConcurrencyMode`` CLASSES. Members of different Enum classes never + compare equal; ``StrEnum`` members compare by string value, so identity of the + defining class no longer matters. + + Keep this a ``StrEnum``. Reverting it to ``Enum`` re-opens both failures, and both + are invisible until runtime. + """ + THREAD = "thread" COROUTINE = "coroutine" diff --git a/workers/shared/infrastructure/context.py b/workers/shared/infrastructure/context.py index fea3e19739..ebf7857be3 100644 --- a/workers/shared/infrastructure/context.py +++ b/workers/shared/infrastructure/context.py @@ -6,11 +6,32 @@ import os import threading -from enum import Enum +from enum import StrEnum from typing import Any -class ConcurrencyMode(Enum): +class ConcurrencyMode(StrEnum): + """Concurrency mode for :class:`StateStore`. + + ``StrEnum`` (not a bare ``Enum``) is load-bearing for TWO reasons, both of which + produced a hard ``RuntimeError: Unknown concurrency mode`` at runtime (UN-3893): + + 1. **A set env var used to break it.** ``mode`` is read as + ``os.environ.get("CONCURRENCY_MODE", ConcurrencyMode.THREAD)``, so setting the + variable yields a plain ``str``. Against a bare ``Enum``, ``"thread" == + ConcurrencyMode.THREAD`` is **False** — so even the CORRECT value took the + whole backend down. ``StrEnum`` members ARE strings, so this now compares equal. + 2. **Duplicate module identity.** This module exists in three places + (``backend/utils``, ``workers/shared/utils``, ``workers/shared/infrastructure``) + and can be imported under more than one path in a merged tree, producing two + distinct ``ConcurrencyMode`` CLASSES. Members of different Enum classes never + compare equal; ``StrEnum`` members compare by string value, so identity of the + defining class no longer matters. + + Keep this a ``StrEnum``. Reverting it to ``Enum`` re-opens both failures, and both + are invisible until runtime. + """ + THREAD = "thread" COROUTINE = "coroutine" diff --git a/workers/shared/utils/local_context.py b/workers/shared/utils/local_context.py index 340d39b126..cca9e591b0 100644 --- a/workers/shared/utils/local_context.py +++ b/workers/shared/utils/local_context.py @@ -4,11 +4,32 @@ import os import threading -from enum import Enum +from enum import StrEnum from typing import Any -class ConcurrencyMode(Enum): +class ConcurrencyMode(StrEnum): + """Concurrency mode for :class:`StateStore`. + + ``StrEnum`` (not a bare ``Enum``) is load-bearing for TWO reasons, both of which + produced a hard ``RuntimeError: Unknown concurrency mode`` at runtime (UN-3893): + + 1. **A set env var used to break it.** ``mode`` is read as + ``os.environ.get("CONCURRENCY_MODE", ConcurrencyMode.THREAD)``, so setting the + variable yields a plain ``str``. Against a bare ``Enum``, ``"thread" == + ConcurrencyMode.THREAD`` is **False** — so even the CORRECT value took the + whole backend down. ``StrEnum`` members ARE strings, so this now compares equal. + 2. **Duplicate module identity.** This module exists in three places + (``backend/utils``, ``workers/shared/utils``, ``workers/shared/infrastructure``) + and can be imported under more than one path in a merged tree, producing two + distinct ``ConcurrencyMode`` CLASSES. Members of different Enum classes never + compare equal; ``StrEnum`` members compare by string value, so identity of the + defining class no longer matters. + + Keep this a ``StrEnum``. Reverting it to ``Enum`` re-opens both failures, and both + are invisible until runtime. + """ + THREAD = "thread" COROUTINE = "coroutine" diff --git a/workers/tests/test_concurrency_mode_identity.py b/workers/tests/test_concurrency_mode_identity.py new file mode 100644 index 0000000000..8b027dc6a2 --- /dev/null +++ b/workers/tests/test_concurrency_mode_identity.py @@ -0,0 +1,67 @@ +"""``ConcurrencyMode`` must stay a ``StrEnum`` (UN-3893). + +``StateStore`` guards every set/get/clear with ``cls.mode == ConcurrencyMode.THREAD`` +and raises ``RuntimeError("Unknown concurrency mode")`` otherwise. With a bare +``Enum`` that guard had two runtime failure modes, both invisible until production: + +1. ``mode`` is read as ``os.environ.get("CONCURRENCY_MODE", ConcurrencyMode.THREAD)``, + so SETTING the variable yields a plain ``str`` — and ``"thread" == + ConcurrencyMode.THREAD`` is False for a bare Enum. Even the CORRECT value broke it. +2. The module exists in three places and can be imported under more than one path in a + merged OSS+cloud tree, producing two distinct ``ConcurrencyMode`` CLASSES. Members + of different Enum classes never compare equal, so the guard raised for every call — + this is what took out ~70 integration tests. + +``StrEnum`` members ARE strings, so both compare by value. These tests fail if anyone +reverts it to ``Enum``. +""" + +from __future__ import annotations + +import importlib.util +import sys +from pathlib import Path + +import pytest + +_ROOT = Path(__file__).resolve().parents[2] +_COPIES = [ + _ROOT / "backend" / "utils" / "local_context.py", + _ROOT / "workers" / "shared" / "utils" / "local_context.py", + _ROOT / "workers" / "shared" / "infrastructure" / "context.py", +] + + +def _load(name: str, path: Path): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +@pytest.mark.parametrize("path", _COPIES, ids=lambda p: p.parent.name) +def test_mode_compares_equal_to_its_string_value(path): + # Failure mode 1: a SET env var yields a str; it must still satisfy the guard. + mod = _load(f"cm_str_{path.parent.name}", path) + assert mod.ConcurrencyMode.THREAD == "thread" + assert "thread" == mod.ConcurrencyMode.THREAD + + +def test_members_compare_equal_across_duplicate_module_copies(): + # Failure mode 2: the same module under two import paths => two classes. The guard + # must still hold, or every StateStore call raises in a merged tree. + a = _load("cm_dup_a", _COPIES[0]) + b = _load("cm_dup_b", _COPIES[1]) + assert a.ConcurrencyMode is not b.ConcurrencyMode # genuinely distinct classes + assert a.ConcurrencyMode.THREAD == b.ConcurrencyMode.THREAD + + +def test_state_store_round_trips_with_env_var_set(monkeypatch): + # End-to-end: the exact call that raised in CI. + monkeypatch.setenv("CONCURRENCY_MODE", "thread") + mod = _load("cm_env", _COPIES[0]) + assert mod.StateStore.mode == mod.ConcurrencyMode.THREAD + mod.StateStore.set("organization_id", "org-1") + assert mod.StateStore.get("organization_id") == "org-1" + mod.StateStore.clear("organization_id") From 425fdc1f1e9029b51b22435c43b4356b42cea3d0 Mon Sep 17 00:00:00 2001 From: ali Date: Tue, 4 Aug 2026 11:47:41 +0530 Subject: [PATCH 11/11] Revert "UN-3893 [FIX] Make ConcurrencyMode a StrEnum so the StateStore guard survives a set env var and duplicate module copies" MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This reverts 2a39e2694. Two reasons, both raised by review: OUT OF SCOPE. The bug is a pre-existing one in main with no connection to PG queue, the flag, or the aux workers. It was picked up while diagnosing why cloud CI was red and should never have been fixed inside this epic's branch: it muddies a PR whose whole promise is "PG-only, flag-gated", and it touches StateStore — multi-tenant org scoping on the request path, used by the flag-off Celery flow in staging and production. "Probably safe" is not the bar for that path. THE STATED CAUSE DOES NOT HOLD. The commit justified itself with "duplicate module identity produces two ConcurrencyMode classes". On checking: both repos import it under exactly one path (`from utils.local_context import`, 25 sites), cloud ships no overlay copy, and CONCURRENCY_MODE is set nowhere (unfiltered grep, both repos). So within a single module cls.mode and ConcurrencyMode.THREAD are the same class and should compare equal — the raise is unexplained. Shipping the change would have masked the symptom without anyone understanding the cause. Verified separately that this epic did NOT introduce the bug: our branch never touched these files, the three copies predate it (2024-03 / 2025-10), and cloud CI resolves refs/heads/main at run time so it never checks out our OSS branch at all. Tracked in UN-3893 for the owner of the merge/CI setup; cloud #1688 CI stays red on it as a blocked-by rather than something this PR introduced or can fix. Co-Authored-By: Claude Opus 4.8 --- backend/utils/local_context.py | 25 +------ workers/shared/infrastructure/context.py | 25 +------ workers/shared/utils/local_context.py | 25 +------ .../tests/test_concurrency_mode_identity.py | 67 ------------------- 4 files changed, 6 insertions(+), 136 deletions(-) delete mode 100644 workers/tests/test_concurrency_mode_identity.py diff --git a/backend/utils/local_context.py b/backend/utils/local_context.py index 2753fb4b4a..0d11c0f6a1 100644 --- a/backend/utils/local_context.py +++ b/backend/utils/local_context.py @@ -1,31 +1,10 @@ import os import threading -from enum import StrEnum +from enum import Enum from typing import Any -class ConcurrencyMode(StrEnum): - """Concurrency mode for :class:`StateStore`. - - ``StrEnum`` (not a bare ``Enum``) is load-bearing for TWO reasons, both of which - produced a hard ``RuntimeError: Unknown concurrency mode`` at runtime (UN-3893): - - 1. **A set env var used to break it.** ``mode`` is read as - ``os.environ.get("CONCURRENCY_MODE", ConcurrencyMode.THREAD)``, so setting the - variable yields a plain ``str``. Against a bare ``Enum``, ``"thread" == - ConcurrencyMode.THREAD`` is **False** — so even the CORRECT value took the - whole backend down. ``StrEnum`` members ARE strings, so this now compares equal. - 2. **Duplicate module identity.** This module exists in three places - (``backend/utils``, ``workers/shared/utils``, ``workers/shared/infrastructure``) - and can be imported under more than one path in a merged tree, producing two - distinct ``ConcurrencyMode`` CLASSES. Members of different Enum classes never - compare equal; ``StrEnum`` members compare by string value, so identity of the - defining class no longer matters. - - Keep this a ``StrEnum``. Reverting it to ``Enum`` re-opens both failures, and both - are invisible until runtime. - """ - +class ConcurrencyMode(Enum): THREAD = "thread" COROUTINE = "coroutine" diff --git a/workers/shared/infrastructure/context.py b/workers/shared/infrastructure/context.py index ebf7857be3..fea3e19739 100644 --- a/workers/shared/infrastructure/context.py +++ b/workers/shared/infrastructure/context.py @@ -6,32 +6,11 @@ import os import threading -from enum import StrEnum +from enum import Enum from typing import Any -class ConcurrencyMode(StrEnum): - """Concurrency mode for :class:`StateStore`. - - ``StrEnum`` (not a bare ``Enum``) is load-bearing for TWO reasons, both of which - produced a hard ``RuntimeError: Unknown concurrency mode`` at runtime (UN-3893): - - 1. **A set env var used to break it.** ``mode`` is read as - ``os.environ.get("CONCURRENCY_MODE", ConcurrencyMode.THREAD)``, so setting the - variable yields a plain ``str``. Against a bare ``Enum``, ``"thread" == - ConcurrencyMode.THREAD`` is **False** — so even the CORRECT value took the - whole backend down. ``StrEnum`` members ARE strings, so this now compares equal. - 2. **Duplicate module identity.** This module exists in three places - (``backend/utils``, ``workers/shared/utils``, ``workers/shared/infrastructure``) - and can be imported under more than one path in a merged tree, producing two - distinct ``ConcurrencyMode`` CLASSES. Members of different Enum classes never - compare equal; ``StrEnum`` members compare by string value, so identity of the - defining class no longer matters. - - Keep this a ``StrEnum``. Reverting it to ``Enum`` re-opens both failures, and both - are invisible until runtime. - """ - +class ConcurrencyMode(Enum): THREAD = "thread" COROUTINE = "coroutine" diff --git a/workers/shared/utils/local_context.py b/workers/shared/utils/local_context.py index cca9e591b0..340d39b126 100644 --- a/workers/shared/utils/local_context.py +++ b/workers/shared/utils/local_context.py @@ -4,32 +4,11 @@ import os import threading -from enum import StrEnum +from enum import Enum from typing import Any -class ConcurrencyMode(StrEnum): - """Concurrency mode for :class:`StateStore`. - - ``StrEnum`` (not a bare ``Enum``) is load-bearing for TWO reasons, both of which - produced a hard ``RuntimeError: Unknown concurrency mode`` at runtime (UN-3893): - - 1. **A set env var used to break it.** ``mode`` is read as - ``os.environ.get("CONCURRENCY_MODE", ConcurrencyMode.THREAD)``, so setting the - variable yields a plain ``str``. Against a bare ``Enum``, ``"thread" == - ConcurrencyMode.THREAD`` is **False** — so even the CORRECT value took the - whole backend down. ``StrEnum`` members ARE strings, so this now compares equal. - 2. **Duplicate module identity.** This module exists in three places - (``backend/utils``, ``workers/shared/utils``, ``workers/shared/infrastructure``) - and can be imported under more than one path in a merged tree, producing two - distinct ``ConcurrencyMode`` CLASSES. Members of different Enum classes never - compare equal; ``StrEnum`` members compare by string value, so identity of the - defining class no longer matters. - - Keep this a ``StrEnum``. Reverting it to ``Enum`` re-opens both failures, and both - are invisible until runtime. - """ - +class ConcurrencyMode(Enum): THREAD = "thread" COROUTINE = "coroutine" diff --git a/workers/tests/test_concurrency_mode_identity.py b/workers/tests/test_concurrency_mode_identity.py deleted file mode 100644 index 8b027dc6a2..0000000000 --- a/workers/tests/test_concurrency_mode_identity.py +++ /dev/null @@ -1,67 +0,0 @@ -"""``ConcurrencyMode`` must stay a ``StrEnum`` (UN-3893). - -``StateStore`` guards every set/get/clear with ``cls.mode == ConcurrencyMode.THREAD`` -and raises ``RuntimeError("Unknown concurrency mode")`` otherwise. With a bare -``Enum`` that guard had two runtime failure modes, both invisible until production: - -1. ``mode`` is read as ``os.environ.get("CONCURRENCY_MODE", ConcurrencyMode.THREAD)``, - so SETTING the variable yields a plain ``str`` — and ``"thread" == - ConcurrencyMode.THREAD`` is False for a bare Enum. Even the CORRECT value broke it. -2. The module exists in three places and can be imported under more than one path in a - merged OSS+cloud tree, producing two distinct ``ConcurrencyMode`` CLASSES. Members - of different Enum classes never compare equal, so the guard raised for every call — - this is what took out ~70 integration tests. - -``StrEnum`` members ARE strings, so both compare by value. These tests fail if anyone -reverts it to ``Enum``. -""" - -from __future__ import annotations - -import importlib.util -import sys -from pathlib import Path - -import pytest - -_ROOT = Path(__file__).resolve().parents[2] -_COPIES = [ - _ROOT / "backend" / "utils" / "local_context.py", - _ROOT / "workers" / "shared" / "utils" / "local_context.py", - _ROOT / "workers" / "shared" / "infrastructure" / "context.py", -] - - -def _load(name: str, path: Path): - spec = importlib.util.spec_from_file_location(name, path) - module = importlib.util.module_from_spec(spec) - sys.modules[name] = module - spec.loader.exec_module(module) - return module - - -@pytest.mark.parametrize("path", _COPIES, ids=lambda p: p.parent.name) -def test_mode_compares_equal_to_its_string_value(path): - # Failure mode 1: a SET env var yields a str; it must still satisfy the guard. - mod = _load(f"cm_str_{path.parent.name}", path) - assert mod.ConcurrencyMode.THREAD == "thread" - assert "thread" == mod.ConcurrencyMode.THREAD - - -def test_members_compare_equal_across_duplicate_module_copies(): - # Failure mode 2: the same module under two import paths => two classes. The guard - # must still hold, or every StateStore call raises in a merged tree. - a = _load("cm_dup_a", _COPIES[0]) - b = _load("cm_dup_b", _COPIES[1]) - assert a.ConcurrencyMode is not b.ConcurrencyMode # genuinely distinct classes - assert a.ConcurrencyMode.THREAD == b.ConcurrencyMode.THREAD - - -def test_state_store_round_trips_with_env_var_set(monkeypatch): - # End-to-end: the exact call that raised in CI. - monkeypatch.setenv("CONCURRENCY_MODE", "thread") - mod = _load("cm_env", _COPIES[0]) - assert mod.StateStore.mode == mod.ConcurrencyMode.THREAD - mod.StateStore.set("organization_id", "org-1") - assert mod.StateStore.get("organization_id") == "org-1" - mod.StateStore.clear("organization_id")