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 = "-"

unstract/core/src/unstract/core/flask/logging.py

@@ -33,11 +33,15 @@ def setup_logging(log_level: int):
             "disable_existing_loggers": False,
             "formatters": {
                 "default": {
+                    # Canonical format shared with the Django backend (``enriched``),
+                    # the workers (``WorkerLogger``) and the x2text-service so a single
+                    # gcloud query parses request_id/trace_id/span_id uniformly.
                     "format": (
                         "%(levelname)s : [%(asctime)s]"
-                        "{pid:%(process)d tid:%(thread)d request_id:%(request_id)s "
-                        + "trace_id:%(otelTraceID)s span_id:%(otelSpanID)s "
-                        + "%(name)s}:- %(message)s"
+                        "{module:%(module)s process:%(process)d thread:%(thread)d "
+                        "request_id:%(request_id)s "
+                        "trace_id:%(otelTraceID)s span_id:%(otelSpanID)s}"
+                        " :- %(message)s"
                     ),
                 },
             },

unstract/core/src/unstract/core/flask/middleware.py

@@ -13,3 +13,13 @@ def register_request_id_middleware(app: Flask):
     @app.before_request
     def assign_request_id():
         g.request_id = request.headers.get("X-Request-ID", str(uuid.uuid4()))
+
+    @app.after_request
+    def echo_request_id(response):
+        # Echo the id back so a caller that did not supply one can learn the
+        # value this service minted and correlate its own logs (mirrors the
+        # Django backend's REQUEST_ID_RESPONSE_HEADER).
+        request_id = getattr(g, "request_id", None)
+        if request_id:
+            response.headers["X-Request-ID"] = request_id
+        return response

workers/shared/clients/base_client.py

@@ -32,6 +32,20 @@
 APPLICATION_JSON = "application/json"
 
 
+def _current_request_id() -> str | None:
+    """Return the request_id bound on the current worker log context, if any.
+
+    Bound by the ``task_prerun`` handler in the logging module; used to
+    propagate ``X-Request-ID`` onto outbound calls to the backend internal API.
+    Returns ``None`` for the ``"-"`` placeholder so no empty header is sent.
+    """
+    ctx = WorkerLogger.get_context()
+    request_id = getattr(ctx, "request_id", None) if ctx else None
+    if not request_id or request_id == "-":
+        return None
+    return request_id
+
+
 # Single PG-queue rollout flag (same key as pg_queue.flags / executor_rpc).
 _PG_QUEUE_FLAG_KEY = "pg_queue_enabled"
 
@@ -316,6 +330,12 @@ def _make_request(
                 if current_org_id:
                     headers["X-Organization-ID"] = current_org_id
 
+                # Propagate the correlation id back to the backend so worker
+                # callbacks share the originating request's request_id in logs.
+                request_id = _current_request_id()
+                if request_id:
+                    headers["X-Request-ID"] = request_id
+
                 if headers:
                     kwargs["headers"] = headers
 

workers/shared/infrastructure/logging/logger.py

@@ -33,6 +33,12 @@ class LogContext:
     organization_id: str | None = None
     correlation_id: str | None = None
     request_id: str | None = None
+    # True only when request_id came from an upstream message header (a genuine
+    # cross-service correlation id), not a locally-derived payload id or the
+    # task_id fallback. Gates worker->worker re-propagation so a fallback id is
+    # never stamped onto child tasks (which would override their own
+    # file_execution_id correlation).
+    request_id_propagatable: bool = False
 
 
 class RequestIDFilter(logging.Filter):
@@ -703,22 +709,82 @@ def _extract_request_id(
     return None
 
 
+def _request_id_from_message(task: Any) -> str | None:
+    """Read an explicit request_id propagated via Celery message headers.
+
+    The task producer (backend ``before_task_publish`` handler, or a worker
+    re-publishing a downstream task) injects ``request_id`` into the message
+    headers. Celery exposes custom headers on ``task.request`` -- as a direct
+    attribute under protocol v2, and via the raw ``headers`` mapping as a
+    version-safe fallback. This is the authoritative cross-service correlation
+    id and takes precedence over payload-derived ids (file_execution_id, etc.).
+    """
+    request = getattr(task, "request", None)
+    if request is None:
+        return None
+    value = getattr(request, "request_id", None)
+    if not value:
+        headers = getattr(request, "headers", None)
+        if isinstance(headers, Mapping):
+            value = headers.get("request_id")
+    return _coerce_id(value)
+
+
 def _bind_task_context(task_id, task, args, kwargs, **_):
     """Celery ``task_prerun`` handler: bind request_id onto the log context.
 
-    Catches any extraction failure so a malformed payload can never leave
-    the previous task's id bound on the thread.
+    Resolution order: an explicit request_id propagated on the message headers
+    (genuine cross-service correlation), then a payload-derived id
+    (``_extract_request_id``), then the Celery ``task_id``. Only the first
+    (header) source is marked propagatable, so a locally-derived fallback is
+    never re-stamped onto child tasks.
+
+    The whole resolution runs inside the ``try`` so a malformed payload -- or a
+    surprising ``task.request`` -- can never raise and leave the previous task's
+    id bound on the thread.
     """
+    propagatable = False
     try:
-        request_id = _extract_request_id(args or (), kwargs or {}, task) or task_id
+        request_id = _request_id_from_message(task)
+        if request_id:
+            propagatable = True
+        else:
+            request_id = _extract_request_id(args or (), kwargs or {}, task)
     except Exception:
         logging.getLogger(__name__).debug(
             "request_id extraction failed for task %s; falling back to task_id",
             task_id,
             exc_info=True,
         )
-        request_id = task_id
-    WorkerLogger.update_context(request_id=request_id, task_id=task_id)
+        request_id = None
+    request_id = request_id or task_id
+    WorkerLogger.update_context(
+        request_id=request_id,
+        task_id=task_id,
+        request_id_propagatable=propagatable,
+    )
+
+
+def _propagate_request_id_on_publish(headers=None, **_):
+    """Celery ``before_task_publish`` handler (worker side): forward the current
+    request_id onto tasks this worker publishes.
+
+    Keeps a genuine cross-service correlation id flowing across worker->worker
+    task chains (e.g. a file-processing task enqueuing a callback). Only
+    propagates when the current id came from an upstream header
+    (``request_id_propagatable``) -- never a locally-derived payload id or the
+    ``task_id`` fallback, which would otherwise override the child task's own
+    ``file_execution_id`` correlation (e.g. for beat/scheduler-originated
+    pipelines). No-ops when absent or when the caller already set the header.
+    """
+    if headers is None or headers.get("request_id"):
+        return
+    ctx = WorkerLogger.get_context()
+    if not ctx or not getattr(ctx, "request_id_propagatable", False):
+        return
+    request_id = _coerce_id(getattr(ctx, "request_id", None))
+    if request_id:
+        headers["request_id"] = request_id
 
 
 def _clear_task_context(**_):
@@ -728,7 +794,9 @@ def _clear_task_context(**_):
     ``WorkerLogger.configure()``; only nulls out the per-task fields bound
     in ``_bind_task_context``.
     """
-    WorkerLogger.update_context(request_id=None, task_id=None)
+    WorkerLogger.update_context(
+        request_id=None, task_id=None, request_id_propagatable=False
+    )
 
 
 @functools.lru_cache(maxsize=1)
@@ -739,13 +807,14 @@ def _install_celery_request_id_signals() -> None:
     debug log if Celery is not importable (e.g. unit tests).
     """
     try:
-        from celery.signals import task_postrun, task_prerun
+        from celery.signals import before_task_publish, task_postrun, task_prerun
     except ImportError as exc:
         logging.getLogger(__name__).debug(
             "celery.signals not importable; request_id signal install skipped: %s",
             exc,
         )
         return
+    before_task_publish.connect(_propagate_request_id_on_publish, weak=False)
     task_prerun.connect(_bind_task_context, weak=False)
     task_postrun.connect(_clear_task_context, weak=False)
 

x2text-service/app/config.py

@@ -1,17 +1,25 @@
+import logging
 from os import environ as env
 
 from dotenv import load_dotenv
 from flask import Flask
 
 from app.controllers import api
+from app.logging_util import register_request_id_middleware, setup_logging
 from app.models import X2TextAudit, be_db
 
 load_dotenv()
 
 
 def create_app() -> Flask:
+    log_level = getattr(logging, env.get("LOG_LEVEL", "INFO").upper(), logging.INFO)
+    setup_logging(log_level)
+
     app = Flask(__name__)
 
+    # Assign/propagate a request_id (X-Request-ID) for cross-service log correlation.
+    register_request_id_middleware(app)
+
     api_url_prefix = env.get("API_URL_PREFIX", "/api/v1")
     app.register_blueprint(api, url_prefix=api_url_prefix)
 

x2text-service/app/controllers/controller.py

@@ -15,10 +15,8 @@
 from app.util import X2TextUtil
 
 basic = Blueprint("basic", __name__)
-# Configure the logging format and level
-logging.basicConfig(
-    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
-)
+# Logging is configured centrally in app.logging_util.setup_logging() (called from
+# create_app) with the request_id-aware canonical format shared across services.
 
 UNSTRUCTURED_URL = "unstructured-url"
 UNSTRUCTURED_API_KEY = "unstructured-api-key"