UN-2123 [FEAT] Propagate request_id across services and workers

backend/backend/celery_service.py

@@ -27,4 +27,8 @@
 app.config_from_object("backend.celery_config.CeleryConfig")
 app.autodiscover_tasks()
 
+# Register signal handlers (e.g. request_id propagation onto published tasks).
+# Importing the module connects the @before_task_publish handler.
+import backend.celery_signals  # noqa: E402, F401
+
 logger.debug(f"Celery Configuration:\n {pformat(app.conf.table(with_defaults=True))}")

backend/backend/celery_signals.py

@@ -0,0 +1,93 @@
+"""Celery signal handlers for the backend (producer side).
+
+Propagates the HTTP ``request_id`` (correlation ID assigned by
+``CustomRequestIDMiddleware``) onto every published Celery task so that worker
+logs can be correlated back to the originating request.
+
+The value is placed in the task message headers under ``request_id``. Workers
+read it from ``task.request`` in ``task_prerun`` and bind it onto their log
+context -- see ``workers/shared/infrastructure/logging/logger.py``. Using the
+``before_task_publish`` signal means this works for *every* ``send_task`` /
+``.delay`` / ``.apply_async`` call with no per-call-site changes.
+"""
+
+import logging
+
+from account_v2.constants import Common
+from celery.signals import before_task_publish, task_postrun, task_prerun
+from log_request_id import local as log_request_id_local
+from utils.local_context import StateStore
+
+logger = logging.getLogger(__name__)
+
+
+@before_task_publish.connect
+def propagate_request_id(headers=None, **kwargs):
+    """Inject the current request_id into the outgoing task's message headers.
+
+    Fires in the producer thread (the web request thread for API-triggered
+    tasks), where ``StateStore`` still holds the request_id set by
+    ``CustomRequestIDMiddleware``. No-ops when there is no request_id in scope
+    (e.g. beat-scheduled publishes), leaving the worker to fall back to its
+    own correlation id (execution_id / task_id).
+    """
+    if headers is None:
+        return
+    try:
+        request_id = StateStore.get(Common.REQUEST_ID)
+    except Exception:
+        # StateStore can raise if CONCURRENCY_MODE is misconfigured; never let
+        # correlation plumbing break task publishing.
+        logger.debug("Unable to read request_id from StateStore", exc_info=True)
+        return
+    if request_id and not headers.get(Common.REQUEST_ID):
+        headers[Common.REQUEST_ID] = request_id
+
+
+def _request_id_from_task(task) -> str | None:
+    """Read a propagated request_id off a Celery task's message context."""
+    request = getattr(task, "request", None)
+    if request is None:
+        return None
+    request_id = getattr(request, Common.REQUEST_ID, None)
+    if not request_id:
+        task_headers = getattr(request, "headers", None)
+        if isinstance(task_headers, dict):
+            request_id = task_headers.get(Common.REQUEST_ID)
+    return request_id or None
+
+
+@task_prerun.connect
+def bind_request_id(task=None, **kwargs):
+    """Bind the propagated request_id for tasks executed by the backend's OWN
+    Celery workers (beat, dashboard-metric tasks, etc.).
+
+    The separate ``workers/`` fleet has its own ``task_prerun`` reader; the
+    backend Celery app previously injected the header (``propagate_request_id``)
+    but never consumed it, so backend-executed tasks logged ``request_id:-``.
+    Binding it onto ``log_request_id``'s thread-local makes
+    ``log_request_id.filters.RequestIDFilter`` emit it, and onto ``StateStore``
+    so any task this worker itself publishes re-propagates it.
+    """
+    request_id = _request_id_from_task(task)
+    if not request_id:
+        return
+    log_request_id_local.request_id = request_id
+    try:
+        StateStore.set(Common.REQUEST_ID, request_id)
+    except Exception:
+        logger.debug("Unable to set request_id on StateStore", exc_info=True)
+
+
+@task_postrun.connect
+def clear_request_id(**kwargs):
+    """Clear the task-scoped request_id bound in ``bind_request_id``."""
+    if hasattr(log_request_id_local, "request_id"):
+        try:
+            del log_request_id_local.request_id
+        except AttributeError:
+            pass
+    try:
+        StateStore.clear(Common.REQUEST_ID)
+    except Exception:
+        logger.debug("Unable to clear request_id from StateStore", exc_info=True)

backend/backend/settings/base.py

@@ -276,7 +276,11 @@ def get_required_setting(setting_key: str, default: str | None = None) -> str |
 CORS_ALLOW_ALL_ORIGINS = False
 
 # Request ID middleware settings
-LOG_REQUEST_ID_HEADER = "X-Request-ID"
+# django-log-request-id resolves this via request.META.get(...), where WSGI exposes
+# the incoming "X-Request-ID" header as the HTTP_-prefixed key HTTP_X_REQUEST_ID.
+# It MUST be the META key, not the raw header name, or the incoming id is never read
+# and a fresh one is minted on every request (breaking client/worker correlation).
+LOG_REQUEST_ID_HEADER = "HTTP_X_REQUEST_ID"
 REQUEST_ID_RESPONSE_HEADER = "X-Request-ID"
 GENERATE_REQUEST_ID_IF_NOT_IN_HEADER = True
 NO_REQUEST_ID = "-"

backend/backend/test_celery_signals.py

@@ -0,0 +1,161 @@
+"""Unit checks for the backend's Celery request_id signal handlers.
+
+These are the producer half of the correlation chain: the HTTP request binds an
+id (see ``middleware/test_request_id.py``), and these handlers carry it onto
+every task the backend publishes and back off the message for tasks the
+backend's *own* Celery workers run.
+
+The handlers are written to fail open -- correlation plumbing must never break
+task publishing -- which means a regression here is silent by construction: the
+header simply stops being set and every worker log line reverts to
+``request_id:-``. Nothing raises, so only assertions catch it.
+
+Pure-logic: ``StateStore`` is a thread-local and the signals are called
+directly, so no broker, worker or DB is involved.
+"""
+
+import pytest
+from account_v2.constants import Common
+from log_request_id import local as log_request_id_local
+from utils.local_context import StateStore
+
+from backend.celery_signals import (
+    bind_request_id,
+    clear_request_id,
+    propagate_request_id,
+)
+
+REQUEST_ID = "6f1c2d3e-4a5b-4c6d-8e9f-0a1b2c3d4e5f"
+
+
+class _Request:
+    """Stands in for a Celery task ``Context``."""
+
+    def __init__(self, request_id=None, headers=None):
+        if request_id is not None:
+            self.request_id = request_id
+        self.headers = headers
+
+
+class _Task:
+    def __init__(self, request=None):
+        self.request = request
+
+
+@pytest.fixture(autouse=True)
+def _clean_context():
+    """Both stores are thread-local and shared across tests in a worker, so a
+    leaked id would make a later test pass for the wrong reason.
+    """
+    yield
+    for store_key in (Common.REQUEST_ID,):
+        try:
+            StateStore.clear(store_key)
+        except AttributeError:
+            pass
+    if hasattr(log_request_id_local, "request_id"):
+        del log_request_id_local.request_id
+
+
+# Producer side: before_task_publish
+
+
+def test_publish_injects_request_id_from_state_store():
+    """The id bound by the HTTP middleware rides out on the task message."""
+    StateStore.set(Common.REQUEST_ID, REQUEST_ID)
+    headers = {}
+
+    propagate_request_id(headers=headers)
+
+    assert headers[Common.REQUEST_ID] == REQUEST_ID
+
+
+def test_publish_is_a_noop_without_a_request_id():
+    """Beat-scheduled publishes have no request in scope; the worker falls back
+    to its own execution_id/task_id rather than receiving an empty header.
+    """
+    headers = {}
+
+    propagate_request_id(headers=headers)
+
+    assert headers == {}
+
+
+def test_publish_does_not_clobber_an_explicit_header():
+    """A caller that set the header deliberately outranks the ambient id."""
+    StateStore.set(Common.REQUEST_ID, REQUEST_ID)
+    headers = {Common.REQUEST_ID: "caller-supplied"}
+
+    propagate_request_id(headers=headers)
+
+    assert headers[Common.REQUEST_ID] == "caller-supplied"
+
+
+def test_publish_tolerates_missing_headers():
+    """``before_task_publish`` must never raise -- it would break the publish
+    itself, taking the actual task down with the correlation plumbing.
+    """
+    propagate_request_id(headers=None)  # does not raise
+
+
+# Consumer side: task_prerun / task_postrun on the backend's own workers
+
+
+def test_prerun_binds_id_from_task_attribute():
+    """Protocol v2 promotes custom headers to attributes on the Context."""
+    bind_request_id(task=_Task(_Request(request_id=REQUEST_ID)))
+
+    assert log_request_id_local.request_id == REQUEST_ID
+    assert StateStore.get(Common.REQUEST_ID) == REQUEST_ID
+
+
+def test_prerun_falls_back_to_raw_headers_mapping():
+    """Version-safe path for when the header is not promoted to an attribute."""
+    bind_request_id(task=_Task(_Request(headers={Common.REQUEST_ID: REQUEST_ID})))
+
+    assert log_request_id_local.request_id == REQUEST_ID
+
+
+def test_prerun_binding_makes_the_id_re_propagate():
+    """Second-order effect that motivated the handler: a task the backend
+    worker itself publishes must carry the inherited id onward, or the chain
+    dies at the first backend-worker hop.
+    """
+    bind_request_id(task=_Task(_Request(request_id=REQUEST_ID)))
+    headers = {}
+
+    propagate_request_id(headers=headers)
+
+    assert headers[Common.REQUEST_ID] == REQUEST_ID
+
+
+def test_prerun_without_an_id_leaves_stores_untouched():
+    """No header means no correlation to inherit -- and, importantly, no empty
+    string bound over whatever the logger would otherwise show.
+    """
+    bind_request_id(task=_Task(_Request()))
+
+    assert not hasattr(log_request_id_local, "request_id")
+    assert StateStore.get(Common.REQUEST_ID) is None
+
+
+def test_prerun_tolerates_a_task_without_a_request():
+    bind_request_id(task=_Task(request=None))  # does not raise
+
+    assert not hasattr(log_request_id_local, "request_id")
+
+
+def test_postrun_clears_the_bound_id():
+    """Celery reuses worker threads; a surviving id would mislabel the *next*
+    task's logs with the previous task's correlation id.
+    """
+    bind_request_id(task=_Task(_Request(request_id=REQUEST_ID)))
+
+    clear_request_id()
+
+    assert not hasattr(log_request_id_local, "request_id")
+    assert StateStore.get(Common.REQUEST_ID) is None
+
+
+def test_postrun_is_safe_when_nothing_was_bound():
+    clear_request_id()  # does not raise

backend/middleware/test_request_id.py

@@ -0,0 +1,121 @@
+"""Request-level tests for ``X-Request-ID`` correlation at the HTTP boundary.
+
+The correlation chain this PR builds (backend -> Celery workers -> internal API
+callbacks) is only anchored if the backend actually *adopts* the id its caller
+sent. That hinges on a single setting, and the failure mode is silent: when
+``LOG_REQUEST_ID_HEADER`` is the raw header name rather than the WSGI ``META``
+key, ``django-log-request-id`` looks up a key that never exists, falls through
+to ``GENERATE_REQUEST_ID_IF_NOT_IN_HEADER`` and mints a fresh uuid4 per
+request. Nothing errors -- every service just logs a different id, and the
+single-``request_id`` log query the feature exists for returns one hop.
+
+So these assert the contract end-to-end through a real middleware chain rather
+than asserting on the setting's value: send a header, get the same id back.
+Settings are deliberately *not* overridden here -- ``backend.settings.test``
+re-exports ``base``, so reverting the production value fails these tests.
+
+No DB is touched (``SimpleTestCase`` + a bare middleware list), keeping this in
+the fast unit tier.
+"""
+
+import re
+import uuid
+
+from django.http import HttpResponse
+from django.test import SimpleTestCase, override_settings
+from django.urls import path
+
+# Set by CustomRequestIDMiddleware; echoed back so the view can assert on it.
+REQUEST_ID_ECHO_HEADER = "X-Seen-Request-Id"
+
+
+def _echo_view(request):
+    """Reports the request id the middleware bound, so the test can compare it
+    against both the id sent in and the id echoed on the response.
+    """
+    response = HttpResponse("ok")
+    response[REQUEST_ID_ECHO_HEADER] = getattr(request, "id", "<unset>")
+    return response
+
+
+urlpatterns = [path("echo/", _echo_view)]
+ECHO_URL = "/echo/"
+
+# Just the middleware under test: no auth, tenancy or session, so nothing here
+# reaches the database.
+_MIDDLEWARE = ["middleware.request_id.CustomRequestIDMiddleware"]
+
+_UUID4_RE = re.compile(
+    r"\A[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}\Z"
+)
+
+
+@override_settings(ROOT_URLCONF=__name__, MIDDLEWARE=_MIDDLEWARE)
+class IncomingRequestIDTest(SimpleTestCase):
+    def test_incoming_header_is_adopted_as_request_id(self):
+        """The id a caller sends becomes ``request.id`` verbatim.
+
+        This is the hop that silently broke: the frontend and the workers both
+        send ``X-Request-ID``, and before the fix the backend discarded it.
+        """
+        sent = "11111111-2222-4333-8444-555555555555"
+
+        response = self.client.get(ECHO_URL, headers={"x-request-id": sent})
+
+        self.assertEqual(
+            response[REQUEST_ID_ECHO_HEADER],
+            sent,
+            "backend minted a new id instead of adopting the caller's -- "
+            "LOG_REQUEST_ID_HEADER must be the WSGI META key HTTP_X_REQUEST_ID, "
+            "not the raw header name",
+        )
+
+    def test_response_echoes_the_incoming_id(self):
+        """The response carries the same id back, which is how the frontend and
+        anyone debugging a live call retrieve it.
+        """
+        sent = "aaaaaaaa-bbbb-4ccc-8ddd-eeeeeeeeeeee"
+
+        response = self.client.get(ECHO_URL, headers={"x-request-id": sent})
+
+        self.assertEqual(response["X-Request-ID"], sent)
+
+    def test_id_is_minted_and_echoed_when_header_absent(self):
+        """With no incoming header the backend still assigns an id and returns
+        it, so a caller that sends nothing can still correlate afterwards.
+        """
+        response = self.client.get(ECHO_URL)
+
+        minted = response["X-Request-ID"]
+        self.assertRegex(minted, _UUID4_RE)
+        self.assertEqual(response[REQUEST_ID_ECHO_HEADER], minted)
+
+    def test_distinct_callers_get_distinct_ids(self):
+        """Two header-less requests must not share an id, or correlation
+        collapses instead of merely missing.
+        """
+        first = self.client.get(ECHO_URL)["X-Request-ID"]
+        second = self.client.get(ECHO_URL)["X-Request-ID"]
+
+        self.assertNotEqual(first, second)
+
+    def test_non_uuid_incoming_id_is_preserved(self):
+        """Ids are opaque: an upstream proxy's own format is adopted as-is
+        rather than being normalised or replaced.
+        """
+        sent = "edge-lb-7f3a91"
+
+        response = self.client.get(ECHO_URL, headers={"x-request-id": sent})
+
+        self.assertEqual(response[REQUEST_ID_ECHO_HEADER], sent)
+        self.assertEqual(response["X-Request-ID"], sent)
+
+    def test_uuid_module_still_backs_the_generator(self):
+        """Guards the custom ``_generate_id`` override, which exists so ids are
+        plain uuid4 strings rather than the library's default hex form.
+        """
+        response = self.client.get(ECHO_URL)
+
+        # Parses without raising, and round-trips to the same string.
+        parsed = uuid.UUID(response["X-Request-ID"])
+        self.assertEqual(str(parsed), response["X-Request-ID"])