From e217f4dd8c7612f6235e9b767f95796db4cf27c3 Mon Sep 17 00:00:00 2001 From: Ghost Jake <89829542+Deepak-Kesavan@users.noreply.github.com> Date: Fri, 7 Aug 2026 10:33:09 +0530 Subject: [PATCH 1/8] UN-2123 [FEAT] Propagate request_id across services and workers Complete end-to-end request_id (X-Request-ID) correlation so a single ID can be used to filter logs across the backend, Celery workers and services in gcloud. - backend: before_task_publish signal injects the request-scoped request_id (from StateStore) into every published Celery task's message headers, with no per-call-site changes. - workers: task_prerun now prefers an explicit request_id from the message headers over payload-derived ids; a before_task_publish handler re-propagates it onto downstream worker->worker task chains. - workers: internal-API HTTP client now forwards X-Request-ID from the worker log context so backend callbacks share the originating request's id. - x2text-service: reads/mints X-Request-ID and logs it (self-contained, no new dependency). - logging: unified all services onto one canonical log format so request_id/ trace_id/span_id parse identically in gcloud. LLMW adoption and full OpenTelemetry trace activation are intentionally out of scope (separate tickets). --- backend/backend/celery_service.py | 4 + backend/backend/celery_signals.py | 43 +++++++++ .../core/src/unstract/core/flask/logging.py | 10 +- workers/shared/clients/base_client.py | 20 ++++ .../shared/infrastructure/logging/logger.py | 66 +++++++++++-- x2text-service/app/config.py | 8 ++ x2text-service/app/logging_util.py | 92 +++++++++++++++++++ 7 files changed, 230 insertions(+), 13 deletions(-) create mode 100644 backend/backend/celery_signals.py create mode 100644 x2text-service/app/logging_util.py diff --git a/backend/backend/celery_service.py b/backend/backend/celery_service.py index 9e4f697170..4b0484bb66 100644 --- a/backend/backend/celery_service.py +++ b/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))}") diff --git a/backend/backend/celery_signals.py b/backend/backend/celery_signals.py new file mode 100644 index 0000000000..cf9a27b801 --- /dev/null +++ b/backend/backend/celery_signals.py @@ -0,0 +1,43 @@ +"""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 +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 diff --git a/unstract/core/src/unstract/core/flask/logging.py b/unstract/core/src/unstract/core/flask/logging.py index d131cb92fe..d6848a4a7e 100644 --- a/unstract/core/src/unstract/core/flask/logging.py +++ b/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" ), }, }, diff --git a/workers/shared/clients/base_client.py b/workers/shared/clients/base_client.py index 017a1d62ec..33beb79337 100644 --- a/workers/shared/clients/base_client.py +++ b/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 diff --git a/workers/shared/infrastructure/logging/logger.py b/workers/shared/infrastructure/logging/logger.py index c82619acf9..74acaea349 100644 --- a/workers/shared/infrastructure/logging/logger.py +++ b/workers/shared/infrastructure/logging/logger.py @@ -703,24 +703,69 @@ 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. + Resolution order: an explicit request_id propagated on the message headers + (cross-service correlation), then a payload-derived id + (``_extract_request_id``), then the Celery ``task_id``. + Catches any extraction failure so a malformed payload can never leave the previous task's id bound on the thread. """ - try: - request_id = _extract_request_id(args or (), kwargs or {}, task) or task_id - 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 + request_id = _request_id_from_message(task) + if not request_id: + try: + 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 = None + request_id = request_id or task_id WorkerLogger.update_context(request_id=request_id, task_id=task_id) +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 the correlation id flowing across worker->worker task chains (e.g. a + file-processing task enqueuing a callback). Reads the request_id bound onto + the thread-local log context by ``_bind_task_context``; 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() + request_id = _coerce_id(getattr(ctx, "request_id", None)) if ctx else None + if request_id: + headers["request_id"] = request_id + + def _clear_task_context(**_): """Celery ``task_postrun`` handler: reset task-scoped fields only. @@ -739,13 +784,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) diff --git a/x2text-service/app/config.py b/x2text-service/app/config.py index 2fca6a6c4d..5823a27c71 100644 --- a/x2text-service/app/config.py +++ b/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) diff --git a/x2text-service/app/logging_util.py b/x2text-service/app/logging_util.py new file mode 100644 index 0000000000..de9d0a0c66 --- /dev/null +++ b/x2text-service/app/logging_util.py @@ -0,0 +1,92 @@ +"""Request-id-aware logging for the x2text-service. + +Self-contained mirror of the shared ``unstract.core.flask`` logging pattern so +the service participates in cross-service correlation (a single ``request_id`` +in every log line) without taking on the ``unstract-core`` dependency. + +The format string is kept identical to the Django backend and the workers so a +single gcloud query parses ``request_id`` / ``trace_id`` / ``span_id`` uniformly +across every service. +""" + +import logging +import uuid +from logging.config import dictConfig + +from flask import Flask, g, request + +# Canonical log format shared with the Django backend (``enriched``) and the +# workers (``WorkerLogger``). Keep these in sync. +LOG_FORMAT = ( + "%(levelname)s : [%(asctime)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" +) + + +class RequestIDFilter(logging.Filter): + """Inject the current request's ``request_id`` into log records.""" + + def filter(self, record: logging.LogRecord) -> bool: + record.request_id = getattr(g, "request_id", "-") if g else "-" + return True + + +class OTelFieldFilter(logging.Filter): + """Default OpenTelemetry id fields to ``"-"`` when not populated.""" + + def filter(self, record: logging.LogRecord) -> bool: + for attr in ("otelTraceID", "otelSpanID"): + if not getattr(record, attr, None): + setattr(record, attr, "-") + return True + + +def setup_logging(log_level: int = logging.INFO) -> None: + """Configure root/werkzeug/gunicorn loggers with the standardized format.""" + dictConfig( + { + "version": 1, + "disable_existing_loggers": False, + "formatters": {"default": {"format": LOG_FORMAT}}, + "filters": { + "request_id": {"()": RequestIDFilter}, + "otel_ids": {"()": OTelFieldFilter}, + }, + "handlers": { + "wsgi": { + "class": "logging.StreamHandler", + "stream": "ext://flask.logging.wsgi_errors_stream", + "formatter": "default", + "filters": ["request_id", "otel_ids"], + }, + }, + "loggers": { + "werkzeug": { + "level": log_level, + "handlers": ["wsgi"], + "propagate": False, + }, + "gunicorn.access": { + "level": log_level, + "handlers": ["wsgi"], + "propagate": False, + }, + "gunicorn.error": { + "level": log_level, + "handlers": ["wsgi"], + "propagate": False, + }, + }, + "root": {"level": log_level, "handlers": ["wsgi"]}, + } + ) + + +def register_request_id_middleware(app: Flask) -> None: + """Read ``X-Request-ID`` from each request (or mint one) onto Flask ``g``.""" + + @app.before_request + def _assign_request_id() -> None: + g.request_id = request.headers.get("X-Request-ID", str(uuid.uuid4())) From 2fd885a3c782f8de77b228e627bdcf66179c8037 Mon Sep 17 00:00:00 2001 From: Ghost Jake <89829542+Deepak-Kesavan@users.noreply.github.com> Date: Tue, 11 Aug 2026 13:17:27 +0530 Subject: [PATCH 2/8] UN-2123 [FIX] Address PR review: guard x2text request_id filter with has_request_context() --- x2text-service/app/logging_util.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/x2text-service/app/logging_util.py b/x2text-service/app/logging_util.py index de9d0a0c66..c88cf1f774 100644 --- a/x2text-service/app/logging_util.py +++ b/x2text-service/app/logging_util.py @@ -13,7 +13,7 @@ import uuid from logging.config import dictConfig -from flask import Flask, g, request +from flask import Flask, g, has_request_context, request # Canonical log format shared with the Django backend (``enriched``) and the # workers (``WorkerLogger``). Keep these in sync. @@ -29,7 +29,11 @@ class RequestIDFilter(logging.Filter): """Inject the current request's ``request_id`` into log records.""" def filter(self, record: logging.LogRecord) -> bool: - record.request_id = getattr(g, "request_id", "-") if g else "-" + # Only touch the request-scoped ``g`` inside an active request context; + # outside one (e.g. gunicorn startup logs) fall back to the placeholder. + record.request_id = ( + getattr(g, "request_id", "-") if has_request_context() else "-" + ) return True From 713f9a474ed173b3678106725f00dab3412a0df8 Mon Sep 17 00:00:00 2001 From: Ghost Jake <89829542+Deepak-Kesavan@users.noreply.github.com> Date: Wed, 12 Aug 2026 13:27:00 +0530 Subject: [PATCH 3/8] UN-2123 [FIX] Address review: honor incoming header, gate re-propagation, backend reader - settings: LOG_REQUEST_ID_HEADER must be the WSGI META key HTTP_X_REQUEST_ID, else django-log-request-id never reads the incoming X-Request-ID and mints a fresh id every request (broke frontend + worker->backend correlation). - workers: run all of _bind_task_context's resolution inside the try (restores the 'never leave prior task's id bound' guarantee); mark only header-origin ids propagatable so _propagate_request_id_on_publish never stamps a local fallback (task_id/execution_id) onto child tasks and clobbers their own file_execution_id correlation (beat/scheduler pipelines). - backend celery: add task_prerun/postrun reader so the backend's OWN Celery workers (beat, dashboard tasks) consume the injected header instead of logging request_id:- (the injection had no reader on that side). - flask (core + x2text): echo X-Request-ID on responses so callers can learn a minted id (mirrors backend REQUEST_ID_RESPONSE_HEADER). - x2text: drop the stale module-import logging.basicConfig superseded by setup_logging's dictConfig. --- backend/backend/celery_signals.py | 52 +++++++++++++- backend/backend/settings/base.py | 6 +- .../src/unstract/core/flask/middleware.py | 10 +++ .../shared/infrastructure/logging/logger.py | 67 +++++++++++++------ x2text-service/app/controllers/controller.py | 6 +- x2text-service/app/logging_util.py | 9 +++ 6 files changed, 122 insertions(+), 28 deletions(-) diff --git a/backend/backend/celery_signals.py b/backend/backend/celery_signals.py index cf9a27b801..d14f6618ab 100644 --- a/backend/backend/celery_signals.py +++ b/backend/backend/celery_signals.py @@ -14,7 +14,8 @@ import logging from account_v2.constants import Common -from celery.signals import before_task_publish +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__) @@ -41,3 +42,52 @@ def propagate_request_id(headers=None, **kwargs): 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) diff --git a/backend/backend/settings/base.py b/backend/backend/settings/base.py index 2f28090240..75e2526461 100644 --- a/backend/backend/settings/base.py +++ b/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 = "-" diff --git a/unstract/core/src/unstract/core/flask/middleware.py b/unstract/core/src/unstract/core/flask/middleware.py index 2f31fc0b36..8463426589 100644 --- a/unstract/core/src/unstract/core/flask/middleware.py +++ b/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 diff --git a/workers/shared/infrastructure/logging/logger.py b/workers/shared/infrastructure/logging/logger.py index 74acaea349..18b745cc16 100644 --- a/workers/shared/infrastructure/logging/logger.py +++ b/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): @@ -728,40 +734,55 @@ def _bind_task_context(task_id, task, args, kwargs, **_): """Celery ``task_prerun`` handler: bind request_id onto the log context. Resolution order: an explicit request_id propagated on the message headers - (cross-service correlation), then a payload-derived id - (``_extract_request_id``), then the Celery ``task_id``. - - Catches any extraction failure so a malformed payload can never leave - the previous task's id bound on the thread. + (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. """ - request_id = _request_id_from_message(task) - if not request_id: - try: + propagatable = False + try: + 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 = None + 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 = None request_id = request_id or task_id - WorkerLogger.update_context(request_id=request_id, task_id=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 the correlation id flowing across worker->worker task chains (e.g. a - file-processing task enqueuing a callback). Reads the request_id bound onto - the thread-local log context by ``_bind_task_context``; no-ops when absent - or when the caller already set the header. + 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() - request_id = _coerce_id(getattr(ctx, "request_id", None)) if ctx else None + 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 @@ -773,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) diff --git a/x2text-service/app/controllers/controller.py b/x2text-service/app/controllers/controller.py index 195cef9682..e1fcd57f6f 100644 --- a/x2text-service/app/controllers/controller.py +++ b/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" diff --git a/x2text-service/app/logging_util.py b/x2text-service/app/logging_util.py index c88cf1f774..fd83897bcd 100644 --- a/x2text-service/app/logging_util.py +++ b/x2text-service/app/logging_util.py @@ -94,3 +94,12 @@ def register_request_id_middleware(app: Flask) -> None: @app.before_request def _assign_request_id() -> None: 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 (mirrors the backend's response header). + request_id = getattr(g, "request_id", None) + if request_id: + response.headers["X-Request-ID"] = request_id + return response From 9bd918a3db30d79d98c67a6d106174104088cd5f Mon Sep 17 00:00:00 2001 From: Ghost Jake <89829542+Deepak-Kesavan@users.noreply.github.com> Date: Mon, 17 Aug 2026 08:32:19 +0530 Subject: [PATCH 4/8] UN-2123 [FIX] Address PR review: add request_id correlation tests pk-zipstack asked for a test asserting the backend honors an incoming X-Request-ID and echoes it back -- the LOG_REQUEST_ID_HEADER fix had no regression guard, and the bug it fixed was silent (a fresh uuid4 minted per request, so every service logged a different id and nothing errored). - backend/middleware/test_request_id.py: drives a real middleware chain and asserts the sent id is adopted as request.id and echoed on the response. Settings are deliberately not overridden, so reverting LOG_REQUEST_ID_HEADER to the raw header name fails these (verified: 3 failures reproducing the original bug). - backend/backend/test_celery_signals.py: covers the publish/prerun/ postrun handlers -- header injection from StateStore, the no-clobber and no-op paths, binding from both the Context attribute and the raw headers mapping, re-propagation from a backend worker, and postrun clearing so a pooled thread can't leak the previous task's id. Both are DB-free so they land in the fast unit tier. --- backend/backend/test_celery_signals.py | 161 +++++++++++++++++++++++++ backend/middleware/test_request_id.py | 121 +++++++++++++++++++ 2 files changed, 282 insertions(+) create mode 100644 backend/backend/test_celery_signals.py create mode 100644 backend/middleware/test_request_id.py diff --git a/backend/backend/test_celery_signals.py b/backend/backend/test_celery_signals.py new file mode 100644 index 0000000000..430e8c7e93 --- /dev/null +++ b/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 diff --git a/backend/middleware/test_request_id.py b/backend/middleware/test_request_id.py new file mode 100644 index 0000000000..a94261cb2c --- /dev/null +++ b/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", "") + 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"]) From 858fac16b9b70415a99682d4adea267618384d7e Mon Sep 17 00:00:00 2001 From: Ghost Jake <89829542+Deepak-Kesavan@users.noreply.github.com> Date: Tue, 18 Aug 2026 17:07:11 +0530 Subject: [PATCH 5/8] UN-2123 [FIX] Address PR review: keep code comments rot-resistant Comments now carry only constraints that are not derivable from the code. Dropped: narration of the change's own history ("previously injected but never consumed"), Celery version/protocol pins that go stale on the next bump while the fallback keeps working, a hardcoded cross-file path, and two "keep these in sync" pleas that nothing enforces. Kept and reworded rather than removed: the HTTP_X_REQUEST_ID footgun note in settings, and the request_id_propagatable rationale -- both explain why a surprising thing is deliberate, which is exactly what does not rot. --- backend/backend/celery_signals.py | 39 +++++++----------- backend/backend/test_celery_signals.py | 4 +- .../core/src/unstract/core/flask/logging.py | 8 ++-- .../shared/infrastructure/logging/logger.py | 41 +++++++------------ 4 files changed, 36 insertions(+), 56 deletions(-) diff --git a/backend/backend/celery_signals.py b/backend/backend/celery_signals.py index d14f6618ab..ba2a76aa6a 100644 --- a/backend/backend/celery_signals.py +++ b/backend/backend/celery_signals.py @@ -1,14 +1,8 @@ -"""Celery signal handlers for the backend (producer side). +"""Celery signal handlers carrying the HTTP ``request_id`` onto published tasks. -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. +The id travels in the task message headers under ``Common.REQUEST_ID``. Hooking +``before_task_publish`` rather than each producer is deliberate: it covers every +``send_task`` / ``.delay`` / ``.apply_async`` with no per-call-site change. """ import logging @@ -25,11 +19,9 @@ 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). + Relies on firing in the *producer* thread, where ``StateStore`` still holds + the id set by ``CustomRequestIDMiddleware``. No-ops without one (e.g. beat + publishes), leaving the worker to derive its own correlation id. """ if headers is None: return @@ -45,7 +37,6 @@ def propagate_request_id(headers=None, **kwargs): 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 @@ -59,15 +50,13 @@ def _request_id_from_task(task) -> str | 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.). + """Bind the propagated request_id for tasks run by the backend's own workers. - 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. + Two writes, both required: ``log_request_id``'s thread-local is the only + thing ``log_request_id.filters.RequestIDFilter`` reads (the HTTP middleware + is otherwise its sole writer, so a task would log ``request_id:-``), and + ``StateStore`` is what ``propagate_request_id`` reads, so tasks this worker + itself publishes carry the id onward. """ request_id = _request_id_from_task(task) if not request_id: @@ -81,7 +70,7 @@ def bind_request_id(task=None, **kwargs): @task_postrun.connect def clear_request_id(**kwargs): - """Clear the task-scoped request_id bound in ``bind_request_id``.""" + """Clear the task-scoped request_id -- worker threads are pooled and reused.""" if hasattr(log_request_id_local, "request_id"): try: del log_request_id_local.request_id diff --git a/backend/backend/test_celery_signals.py b/backend/backend/test_celery_signals.py index 430e8c7e93..9d426ee549 100644 --- a/backend/backend/test_celery_signals.py +++ b/backend/backend/test_celery_signals.py @@ -102,7 +102,7 @@ def test_publish_tolerates_missing_headers(): def test_prerun_binds_id_from_task_attribute(): - """Protocol v2 promotes custom headers to attributes on the Context.""" + """Celery surfaces the custom header as an attribute on the Context.""" bind_request_id(task=_Task(_Request(request_id=REQUEST_ID))) assert log_request_id_local.request_id == REQUEST_ID @@ -110,7 +110,7 @@ def test_prerun_binds_id_from_task_attribute(): def test_prerun_falls_back_to_raw_headers_mapping(): - """Version-safe path for when the header is not promoted to an attribute.""" + """Fallback for when Celery leaves the header only in the raw mapping.""" bind_request_id(task=_Task(_Request(headers={Common.REQUEST_ID: REQUEST_ID}))) assert log_request_id_local.request_id == REQUEST_ID diff --git a/unstract/core/src/unstract/core/flask/logging.py b/unstract/core/src/unstract/core/flask/logging.py index d6848a4a7e..7b588aa223 100644 --- a/unstract/core/src/unstract/core/flask/logging.py +++ b/unstract/core/src/unstract/core/flask/logging.py @@ -33,9 +33,11 @@ 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. + # Canonical cross-service log format -- this module owns it. + # A single gcloud query parses request_id/trace_id/span_id + # across every service only while the copies agree; the known + # copies are the Django backend ``enriched`` formatter, the + # workers' ``WorkerLogger``, and ``x2text-service``. "format": ( "%(levelname)s : [%(asctime)s]" "{module:%(module)s process:%(process)d thread:%(thread)d " diff --git a/workers/shared/infrastructure/logging/logger.py b/workers/shared/infrastructure/logging/logger.py index 18b745cc16..0b857063e6 100644 --- a/workers/shared/infrastructure/logging/logger.py +++ b/workers/shared/infrastructure/logging/logger.py @@ -33,11 +33,10 @@ 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). + # Gates worker->worker re-propagation: only an id received from upstream may + # be stamped onto child tasks. A locally-derived fallback (task_id, or a + # payload id) must not be, or it would override the child's own + # file_execution_id correlation -- e.g. on beat-scheduled pipelines. request_id_propagatable: bool = False @@ -712,12 +711,8 @@ def _extract_request_id( 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.). + Celery may surface a custom header either as a direct attribute on + ``task.request`` or only in its raw ``headers`` mapping, so both are tried. """ request = getattr(task, "request", None) if request is None: @@ -733,15 +728,12 @@ def _request_id_from_message(task: Any) -> str | None: def _bind_task_context(task_id, task, args, kwargs, **_): """Celery ``task_prerun`` handler: bind request_id onto the log context. - 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. + Resolution order: message header, then a payload-derived id, then the Celery + ``task_id``. Only the header source is marked propagatable. - 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. + The whole resolution stays inside the ``try``: an exception escaping here + would skip ``update_context`` and leave the *previous* task's id bound on a + pooled thread. """ propagatable = False try: @@ -769,13 +761,10 @@ 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. + Keeps an upstream correlation id flowing across worker->worker chains (e.g. + a file-processing task enqueuing a callback). Gated on + ``request_id_propagatable`` -- see ``LogContext`` for why a fallback id must + not fan out. No-ops when absent or when the caller already set the header. """ if headers is None or headers.get("request_id"): return From 272d66d99c87c56a75f06bf7689748d3b4f0bd4f Mon Sep 17 00:00:00 2001 From: Ghost Jake <89829542+Deepak-Kesavan@users.noreply.github.com> Date: Tue, 18 Aug 2026 17:07:36 +0530 Subject: [PATCH 6/8] UN-2123 [FIX] Self-review: validate incoming X-Request-ID, cover the worker hop Two findings from a self-review pass, neither raised on the PR. 1. Making LOG_REQUEST_ID_HEADER resolve is what turned the incoming header into an attacker-controlled value for the first time -- before the fix it never matched, so request.id was always server-minted. It was then adopted verbatim and fanned out to every log line, Celery message header, child task and outbound internal-API call, with no validation anywhere. gunicorn rejects only NUL/CR/LF and caps a field at ~8KB, so an ANSI escape erases the real prefix of a log record and lets the caller forge the rest (CWE-117), and an 8000-char id is re-stamped onto every child task of an execution. All three ingress points now accept an id only in a shape that cannot forge or bloat a record, and mint a fresh one otherwise. 2. The worker half of the chain had no test at all: the propagatable gate could be deleted, precedence flipped back to payload-first, or the teardown reset dropped, and the whole suite stayed green. Adds 17 tests over the real handlers. Verified both regressions now fail. Backend: 9 tests (2 fail without the sanitizer). Workers: 17 tests. --- backend/middleware/request_id.py | 14 + backend/middleware/test_request_id.py | 41 ++- .../src/unstract/core/flask/middleware.py | 15 +- workers/tests/test_request_id_signals.py | 240 ++++++++++++++++++ x2text-service/app/logging_util.py | 24 +- 5 files changed, 320 insertions(+), 14 deletions(-) create mode 100644 workers/tests/test_request_id_signals.py diff --git a/backend/middleware/request_id.py b/backend/middleware/request_id.py index c7d0bdb82a..b8df8053e0 100644 --- a/backend/middleware/request_id.py +++ b/backend/middleware/request_id.py @@ -1,8 +1,22 @@ +import re import uuid from log_request_id.middleware import RequestIDMiddleware +# An incoming X-Request-ID reaches every log line, every published Celery message +# header and every outbound internal-API call, so it is only ever accepted in a +# shape that cannot forge a log record or bloat a message. gunicorn rejects only +# \0\r\n and caps a header field at ~8KB, which leaves ANSI escapes and 8000-char +# values arriving intact; anything not matching here is discarded for a fresh id. +SAFE_REQUEST_ID = re.compile(r"\A[A-Za-z0-9._:-]{1,128}\Z") + class CustomRequestIDMiddleware(RequestIDMiddleware): + def _get_request_id(self, request): + request_id = super()._get_request_id(request) + if request_id and SAFE_REQUEST_ID.match(request_id): + return request_id + return self._generate_id() + def _generate_id(self): return str(uuid.uuid4()) diff --git a/backend/middleware/test_request_id.py b/backend/middleware/test_request_id.py index a94261cb2c..60520d8994 100644 --- a/backend/middleware/test_request_id.py +++ b/backend/middleware/test_request_id.py @@ -1,8 +1,8 @@ """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 +The correlation chain (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 @@ -119,3 +119,38 @@ def test_uuid_module_still_backs_the_generator(self): # 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"]) + + def test_ansi_escapes_in_incoming_id_are_rejected(self): + r"""A caller-supplied id lands unescaped in every log line, so one + carrying terminal control codes could erase the real prefix of a record + and forge the rest. gunicorn only rejects ``\\0\\r\\n``, so the escape + arrives intact and the middleware is the place that has to drop it. + """ + sent = "\x1b[2K\x1b[1000Ddeadbeef} :- SPOOFED: admin deleted org 42" + + response = self.client.get(ECHO_URL, headers={"x-request-id": sent}) + + bound = response[REQUEST_ID_ECHO_HEADER] + self.assertNotEqual(bound, sent) + self.assertRegex(bound, _UUID4_RE) + + def test_overlong_incoming_id_is_rejected(self): + """The id is re-stamped onto every published Celery message and every + log line of an execution, so an unbounded one amplifies: gunicorn + accepts ~8KB, which one API call can fan out across N file tasks. + """ + sent = "A" * 8000 + + response = self.client.get(ECHO_URL, headers={"x-request-id": sent}) + + self.assertRegex(response[REQUEST_ID_ECHO_HEADER], _UUID4_RE) + + def test_incoming_id_at_the_length_limit_is_kept(self): + """The bound is on hostile length, not on legitimate upstream formats -- + an id up to 128 chars is still adopted verbatim. + """ + sent = "b" * 128 + + response = self.client.get(ECHO_URL, headers={"x-request-id": sent}) + + self.assertEqual(response[REQUEST_ID_ECHO_HEADER], sent) diff --git a/unstract/core/src/unstract/core/flask/middleware.py b/unstract/core/src/unstract/core/flask/middleware.py index 8463426589..64655a78c8 100644 --- a/unstract/core/src/unstract/core/flask/middleware.py +++ b/unstract/core/src/unstract/core/flask/middleware.py @@ -1,7 +1,20 @@ +import re import uuid from flask import Flask, g, request +# An incoming X-Request-ID is caller-supplied and lands in every log line and in +# the echoed response header, so it is only accepted in a shape that cannot forge +# a log record (ANSI/control characters) or bloat one (unbounded length). +SAFE_REQUEST_ID = re.compile(r"\A[A-Za-z0-9._:-]{1,128}\Z") + + +def _incoming_request_id() -> str: + request_id = request.headers.get("X-Request-ID") + if request_id and SAFE_REQUEST_ID.match(request_id): + return request_id + return str(uuid.uuid4()) + def register_request_id_middleware(app: Flask): """Adds request ID to each request @@ -12,7 +25,7 @@ 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())) + g.request_id = _incoming_request_id() @app.after_request def echo_request_id(response): diff --git a/workers/tests/test_request_id_signals.py b/workers/tests/test_request_id_signals.py new file mode 100644 index 0000000000..1bc797a62c --- /dev/null +++ b/workers/tests/test_request_id_signals.py @@ -0,0 +1,240 @@ +"""Worker half of the ``X-Request-ID`` correlation chain. + +The backend's tests cover the producer side (injecting the id into a published +task's message headers). These cover what the worker does with it: preferring an +upstream header over a payload-derived id, refusing to fan out an id it derived +locally, clearing per-task state on a pooled thread, and re-emitting the id on +outbound internal-API calls. + +Two of these pin regressions that were caught by review rather than by a test -- +the propagatable gate, and the reset of that flag on task teardown. +""" + +import pytest + +from shared.clients.base_client import _current_request_id +from shared.infrastructure.logging.logger import ( + LogContext, + WorkerLogger, + _bind_task_context, + _clear_task_context, + _propagate_request_id_on_publish, + _request_id_from_message, +) + +HEADER_ID = "11111111-2222-4333-8444-555555555555" +PAYLOAD_ID = "99999999-8888-4777-8666-555555555555" +TASK_ID = "celery-task-id-0001" + + +class _Request: + """Stands in for ``celery.app.task.Context``. + + Celery may expose a custom message header as an attribute or leave it only + in the raw ``headers`` mapping, so both shapes are constructible here. + """ + + def __init__(self, request_id=None, headers=None, **payload): + if request_id is not None: + self.request_id = request_id + self.headers = headers + self.__dict__.update(payload) + + +class _Task: + def __init__(self, request): + self.request = request + self.name = "workers.test.task" + + +@pytest.fixture(autouse=True) +def _isolate_context(): + """Each test starts and ends on a clean thread-local context.""" + WorkerLogger.clear_context() + yield + WorkerLogger.clear_context() + + +# --------------------------------------------------------------------------- +# Reading the id off the message +# --------------------------------------------------------------------------- + + +def test_header_attribute_is_read(): + assert _request_id_from_message(_Task(_Request(request_id=HEADER_ID))) == HEADER_ID + + +def test_raw_headers_mapping_is_the_fallback(): + task = _Task(_Request(headers={"request_id": HEADER_ID})) + + assert _request_id_from_message(task) == HEADER_ID + + +def test_absent_header_reads_as_none(): + assert _request_id_from_message(_Task(_Request())) is None + + +def test_missing_request_object_does_not_raise(): + class _Bare: + request = None + + assert _request_id_from_message(_Bare()) is None + + +# --------------------------------------------------------------------------- +# Binding: precedence and the propagatable gate +# --------------------------------------------------------------------------- + + +def test_header_id_wins_over_payload_and_is_propagatable(): + """An upstream id is the authoritative correlation key, so it takes + precedence over anything derivable from the payload. + """ + task = _Task(_Request(request_id=HEADER_ID)) + + _bind_task_context(TASK_ID, task, (), {"file_execution_id": PAYLOAD_ID}) + + ctx = WorkerLogger.get_context() + assert ctx.request_id == HEADER_ID + assert ctx.request_id_propagatable is True + + +def test_payload_id_is_used_but_is_not_propagatable(): + """A payload-derived id correlates this task only. Marking it propagatable + would stamp it onto child tasks and override *their* own file_execution_id. + """ + task = _Task(_Request()) + + _bind_task_context(TASK_ID, task, (), {"file_execution_id": PAYLOAD_ID}) + + ctx = WorkerLogger.get_context() + assert ctx.request_id == PAYLOAD_ID + assert ctx.request_id_propagatable is False + + +def test_task_id_is_the_last_resort_and_is_not_propagatable(): + _bind_task_context(TASK_ID, _Task(_Request()), (), {}) + + ctx = WorkerLogger.get_context() + assert ctx.request_id == TASK_ID + assert ctx.request_id_propagatable is False + + +def test_a_raising_request_object_still_binds_the_task_id(): + """The whole resolution sits inside a try precisely so a malformed message + cannot leave the *previous* task's id bound on a reused thread. + """ + + class _Exploding: + @property + def request(self): + raise RuntimeError("malformed message") + + WorkerLogger.set_context(LogContext(request_id="stale-previous-task")) + + _bind_task_context(TASK_ID, _Exploding(), (), {}) + + ctx = WorkerLogger.get_context() + assert ctx.request_id == TASK_ID + assert ctx.request_id_propagatable is False + + +# --------------------------------------------------------------------------- +# Re-publishing to child tasks +# --------------------------------------------------------------------------- + + +def test_propagatable_id_is_stamped_onto_a_child_task(): + WorkerLogger.set_context( + LogContext(request_id=HEADER_ID, request_id_propagatable=True) + ) + headers = {} + + _propagate_request_id_on_publish(headers=headers) + + assert headers["request_id"] == HEADER_ID + + +def test_locally_derived_id_is_not_stamped_onto_a_child_task(): + """The regression the gate exists for: without it a scheduler-originated + task fans its own task_id out over every child's own correlation id. + """ + WorkerLogger.set_context( + LogContext(request_id=TASK_ID, request_id_propagatable=False) + ) + headers = {} + + _propagate_request_id_on_publish(headers=headers) + + assert headers == {} + + +def test_an_id_the_caller_already_set_is_left_alone(): + WorkerLogger.set_context( + LogContext(request_id=HEADER_ID, request_id_propagatable=True) + ) + headers = {"request_id": "explicitly-set-by-caller"} + + _propagate_request_id_on_publish(headers=headers) + + assert headers["request_id"] == "explicitly-set-by-caller" + + +def test_no_headers_mapping_is_a_no_op(): + WorkerLogger.set_context( + LogContext(request_id=HEADER_ID, request_id_propagatable=True) + ) + + _propagate_request_id_on_publish(headers=None) # does not raise + + +# --------------------------------------------------------------------------- +# Teardown +# --------------------------------------------------------------------------- + + +def test_teardown_resets_the_propagatable_flag(): + """Prefork/thread workers reuse the thread. A flag left True would let the + next task's locally-derived id fan out as though it came from upstream. + """ + _bind_task_context(TASK_ID, _Task(_Request(request_id=HEADER_ID)), (), {}) + + _clear_task_context() + + ctx = WorkerLogger.get_context() + assert ctx.request_id is None + assert ctx.request_id_propagatable is False + + +def test_a_child_published_after_teardown_inherits_nothing(): + _bind_task_context(TASK_ID, _Task(_Request(request_id=HEADER_ID)), (), {}) + _clear_task_context() + headers = {} + + _propagate_request_id_on_publish(headers=headers) + + assert headers == {} + + +# --------------------------------------------------------------------------- +# Outbound internal-API calls +# --------------------------------------------------------------------------- + + +def test_bound_id_is_offered_to_outbound_calls(): + _bind_task_context(TASK_ID, _Task(_Request(request_id=HEADER_ID)), (), {}) + + assert _current_request_id() == HEADER_ID + + +def test_placeholder_id_sends_no_header(): + """``"-"`` is the formatter's empty rendering, not an id -- sending it would + put a meaningless X-Request-ID on the wire. + """ + WorkerLogger.set_context(LogContext(request_id="-")) + + assert _current_request_id() is None + + +def test_no_context_sends_no_header(): + assert _current_request_id() is None diff --git a/x2text-service/app/logging_util.py b/x2text-service/app/logging_util.py index fd83897bcd..97aca91706 100644 --- a/x2text-service/app/logging_util.py +++ b/x2text-service/app/logging_util.py @@ -1,22 +1,23 @@ """Request-id-aware logging for the x2text-service. -Self-contained mirror of the shared ``unstract.core.flask`` logging pattern so -the service participates in cross-service correlation (a single ``request_id`` -in every log line) without taking on the ``unstract-core`` dependency. - -The format string is kept identical to the Django backend and the workers so a -single gcloud query parses ``request_id`` / ``trace_id`` / ``span_id`` uniformly -across every service. +Deliberately a self-contained copy of ``unstract.core.flask``'s logging rather +than an import: this service does not take the ``unstract-core`` dependency. """ import logging +import re import uuid from logging.config import dictConfig from flask import Flask, g, has_request_context, request -# Canonical log format shared with the Django backend (``enriched``) and the -# workers (``WorkerLogger``). Keep these in sync. +# See ``unstract.core.flask.middleware``: a caller-supplied id reaches every log +# line and the echoed response header, so only a shape that cannot forge or bloat +# a record is accepted. +SAFE_REQUEST_ID = re.compile(r"\A[A-Za-z0-9._:-]{1,128}\Z") + +# Copy of the canonical format owned by ``unstract.core.flask.logging``; a +# divergence silently splits this service out of the cross-service log query. LOG_FORMAT = ( "%(levelname)s : [%(asctime)s]" "{module:%(module)s process:%(process)d thread:%(thread)d " @@ -93,7 +94,10 @@ def register_request_id_middleware(app: Flask) -> None: @app.before_request def _assign_request_id() -> None: - g.request_id = request.headers.get("X-Request-ID", str(uuid.uuid4())) + request_id = request.headers.get("X-Request-ID") + if not (request_id and SAFE_REQUEST_ID.match(request_id)): + request_id = str(uuid.uuid4()) + g.request_id = request_id @app.after_request def _echo_request_id(response): From 52ed22a5c82b6f5d396f4b8e6d569fe83e8594aa Mon Sep 17 00:00:00 2001 From: Ghost Jake <89829542+Deepak-Kesavan@users.noreply.github.com> Date: Tue, 18 Aug 2026 17:21:47 +0530 Subject: [PATCH 7/8] UN-2123 [FIX] Provision the request_id in the backend instead of trusting callers The previous commit validated the inbound X-Request-ID's shape. That stops log forgery and the 8KB amplification, but not the actual problem: a caller can send one valid id on every request, and a log query for it then returns unrelated requests across tenants -- which reads as though correlation worked, and is worse than returning nothing. So public routes no longer adopt a caller's id at all; the backend provisions a uuid4 and echoes it. The frontend's interceptor still sends one, but nothing is lost -- getRequestIdFromError reads the response header first, so the id it surfaces on an error is the one the backend actually logged. /internal/ is the exception: there the caller is our own worker forwarding the id of the request that started the execution, which is the hop that makes worker logs correlate back to the originating call. Only a canonical UUID is accepted, since that path is authenticated downstream rather than in the middleware. The Flask services keep adopting a forwarded id -- they are not internet-facing and every caller is another Unstract service -- but now only in the canonical UUID shape, which bounds length and rules out log-forging control characters. Backend: 21 tests. Workers: 17. Reverting the internal-only gate fails 2. --- backend/middleware/request_id.py | 39 +++- backend/middleware/test_request_id.py | 179 ++++++++++-------- .../src/unstract/core/flask/middleware.py | 22 ++- x2text-service/app/logging_util.py | 13 +- 4 files changed, 147 insertions(+), 106 deletions(-) diff --git a/backend/middleware/request_id.py b/backend/middleware/request_id.py index b8df8053e0..197fc3f19c 100644 --- a/backend/middleware/request_id.py +++ b/backend/middleware/request_id.py @@ -1,21 +1,40 @@ -import re import uuid from log_request_id.middleware import RequestIDMiddleware -# An incoming X-Request-ID reaches every log line, every published Celery message -# header and every outbound internal-API call, so it is only ever accepted in a -# shape that cannot forge a log record or bloat a message. gunicorn rejects only -# \0\r\n and caps a header field at ~8KB, which leaves ANSI escapes and 8000-char -# values arriving intact; anything not matching here is discarded for a fresh id. -SAFE_REQUEST_ID = re.compile(r"\A[A-Za-z0-9._:-]{1,128}\Z") +# The internal service boundary: our own workers call back here while executing +# a workflow, forwarding the request_id of the HTTP call that started it. +INTERNAL_PATH_PREFIX = "/internal/" + + +def _canonical_uuid(value: str | None) -> str | None: + """Return ``value`` only if it is a canonical hyphenated UUID.""" + try: + return value if str(uuid.UUID(value)) == value else None + except (AttributeError, TypeError, ValueError): + return None class CustomRequestIDMiddleware(RequestIDMiddleware): + """Provisions the request id here rather than trusting the caller's. + + Adopting a caller-supplied id lets any client repeat one value across + unrelated requests. That does not merely lose correlation -- a log query for + that id returns other tenants' requests and reads as though it worked, which + is worse than returning nothing. The id also reaches every log line, Celery + message header and outbound internal-API call, so an unvalidated one can + forge a log record outright. + + The internal boundary is the exception: there the caller is our own worker + forwarding the id of the request that started the execution, which is the + hop that makes worker logs correlate back to the originating call. + """ + def _get_request_id(self, request): - request_id = super()._get_request_id(request) - if request_id and SAFE_REQUEST_ID.match(request_id): - return request_id + if request.path.startswith(INTERNAL_PATH_PREFIX): + forwarded = _canonical_uuid(request.META.get(self.request_id_header)) + if forwarded: + return forwarded return self._generate_id() def _generate_id(self): diff --git a/backend/middleware/test_request_id.py b/backend/middleware/test_request_id.py index 60520d8994..bfee1f7a30 100644 --- a/backend/middleware/test_request_id.py +++ b/backend/middleware/test_request_id.py @@ -1,18 +1,21 @@ -"""Request-level tests for ``X-Request-ID`` correlation at the HTTP boundary. - -The correlation chain (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. +"""Request-level tests for ``X-Request-ID`` provisioning at the HTTP boundary. + +Two properties are pinned here, and they pull in opposite directions. + +*The id is server-provisioned.* A caller-supplied id is ignored on public +routes. Honouring one would let a client send the same value on every request, +and a log query for it would then return unrelated requests across tenants -- +which reads as though correlation worked. The id also reaches every log line and +every published Celery message, so an unvalidated one can forge a log record. + +*Except across the internal boundary,* where the caller is our own worker +forwarding the id of the request that started the execution. That hop is what +makes worker logs correlate back to the originating call, and it hinges on a +single setting with a silent failure mode: 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 and mints a fresh id instead. Nothing errors -- +every service just logs a different id. Settings are deliberately *not* +overridden here, so reverting the production value fails these tests. No DB is touched (``SimpleTestCase`` + a bare middleware list), keeping this in the fast unit tier. @@ -38,8 +41,13 @@ def _echo_view(request): return response -urlpatterns = [path("echo/", _echo_view)] -ECHO_URL = "/echo/" +PUBLIC_URL = "/echo/" +INTERNAL_URL = "/internal/echo/" + +urlpatterns = [ + path("echo/", _echo_view), + path("internal/echo/", _echo_view), +] # Just the middleware under test: no auth, tenancy or session, so nothing here # reaches the database. @@ -51,106 +59,115 @@ def _echo_view(request): @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. +class PublicRequestIDTest(SimpleTestCase): + """Public routes: the id is ours, whatever the caller sent.""" - 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. + def test_caller_supplied_id_is_ignored(self): + """The frontend attaches a uuid4 of its own to every call. It is still + the backend's id that is authoritative -- the frontend reads the value + back off the response header, so nothing is lost by ignoring it. """ sent = "11111111-2222-4333-8444-555555555555" - response = self.client.get(ECHO_URL, headers={"x-request-id": sent}) + response = self.client.get(PUBLIC_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", - ) + bound = response[REQUEST_ID_ECHO_HEADER] + self.assertNotEqual(bound, sent) + self.assertRegex(bound, _UUID4_RE) - 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. + def test_repeated_caller_id_does_not_collapse_requests(self): + """The reason a caller's id is not adopted: one repeated across requests + would make a single log query return unrelated requests, and read as + though it had worked. """ sent = "aaaaaaaa-bbbb-4ccc-8ddd-eeeeeeeeeeee" + headers = {"x-request-id": sent} - response = self.client.get(ECHO_URL, headers={"x-request-id": sent}) + first = self.client.get(PUBLIC_URL, headers=headers) + second = self.client.get(PUBLIC_URL, headers=headers) - self.assertEqual(response["X-Request-ID"], sent) + self.assertNotEqual(first[REQUEST_ID_ECHO_HEADER], second[REQUEST_ID_ECHO_HEADER]) - 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. + def test_hostile_id_never_reaches_a_log_line(self): + """An id lands unescaped in every log line, so one carrying terminal + control codes could erase the real prefix of a record and forge the + rest. gunicorn rejects only NUL/CR/LF, so the escape arrives intact. """ - response = self.client.get(ECHO_URL) + sent = "\x1b[2K\x1b[1000Ddeadbeef} :- SPOOFED: admin deleted org 42" - minted = response["X-Request-ID"] - self.assertRegex(minted, _UUID4_RE) - self.assertEqual(response[REQUEST_ID_ECHO_HEADER], minted) + response = self.client.get(PUBLIC_URL, headers={"x-request-id": sent}) - def test_distinct_callers_get_distinct_ids(self): - """Two header-less requests must not share an id, or correlation - collapses instead of merely missing. + self.assertRegex(response[REQUEST_ID_ECHO_HEADER], _UUID4_RE) + + def test_response_echoes_the_provisioned_id(self): + """The response carries the id back, which is how the frontend surfaces + it on an error and how anyone debugging a live call retrieves it. """ - first = self.client.get(ECHO_URL)["X-Request-ID"] - second = self.client.get(ECHO_URL)["X-Request-ID"] + response = self.client.get(PUBLIC_URL) - self.assertNotEqual(first, second) + echoed = response["X-Request-ID"] + self.assertRegex(echoed, _UUID4_RE) + self.assertEqual(response[REQUEST_ID_ECHO_HEADER], echoed) - 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. + def test_distinct_callers_get_distinct_ids(self): + """Two requests must not share an id, or correlation collapses instead + of merely missing. """ - sent = "edge-lb-7f3a91" - - response = self.client.get(ECHO_URL, headers={"x-request-id": sent}) + first = self.client.get(PUBLIC_URL)["X-Request-ID"] + second = self.client.get(PUBLIC_URL)["X-Request-ID"] - self.assertEqual(response[REQUEST_ID_ECHO_HEADER], sent) - self.assertEqual(response["X-Request-ID"], sent) + self.assertNotEqual(first, second) 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) + response = self.client.get(PUBLIC_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"]) - def test_ansi_escapes_in_incoming_id_are_rejected(self): - r"""A caller-supplied id lands unescaped in every log line, so one - carrying terminal control codes could erase the real prefix of a record - and forge the rest. gunicorn only rejects ``\\0\\r\\n``, so the escape - arrives intact and the middleware is the place that has to drop it. + +@override_settings(ROOT_URLCONF=__name__, MIDDLEWARE=_MIDDLEWARE) +class InternalRequestIDTest(SimpleTestCase): + """The internal boundary: a worker's forwarded id is honoured.""" + + def test_forwarded_id_is_adopted(self): + """The hop that silently broke: workers send the originating request's + id on their callbacks, and before the fix the backend discarded it. """ - sent = "\x1b[2K\x1b[1000Ddeadbeef} :- SPOOFED: admin deleted org 42" + sent = "11111111-2222-4333-8444-555555555555" - response = self.client.get(ECHO_URL, headers={"x-request-id": sent}) + response = self.client.get(INTERNAL_URL, headers={"x-request-id": sent}) - bound = response[REQUEST_ID_ECHO_HEADER] - self.assertNotEqual(bound, sent) - self.assertRegex(bound, _UUID4_RE) + self.assertEqual( + response[REQUEST_ID_ECHO_HEADER], + sent, + "backend minted a new id instead of adopting the worker's -- " + "LOG_REQUEST_ID_HEADER must be the WSGI META key HTTP_X_REQUEST_ID, " + "not the raw header name", + ) + self.assertEqual(response["X-Request-ID"], sent) - def test_overlong_incoming_id_is_rejected(self): - """The id is re-stamped onto every published Celery message and every - log line of an execution, so an unbounded one amplifies: gunicorn - accepts ~8KB, which one API call can fan out across N file tasks. + def test_forwarded_id_must_be_a_uuid(self): + """The boundary is authenticated downstream, not here, so an unauthorised + caller can still reach this path. Only the shape our own services emit is + accepted, which bounds both length and character set. """ - sent = "A" * 8000 - - response = self.client.get(ECHO_URL, headers={"x-request-id": sent}) + response = self.client.get(INTERNAL_URL, headers={"x-request-id": "edge-lb-7f"}) self.assertRegex(response[REQUEST_ID_ECHO_HEADER], _UUID4_RE) - def test_incoming_id_at_the_length_limit_is_kept(self): - """The bound is on hostile length, not on legitimate upstream formats -- - an id up to 128 chars is still adopted verbatim. + def test_overlong_forwarded_id_is_rejected(self): + """An id is re-stamped onto every published Celery message and every log + line of an execution, so an unbounded one amplifies: gunicorn accepts + ~8KB, which one call can fan out across N file tasks. """ - sent = "b" * 128 + response = self.client.get(INTERNAL_URL, headers={"x-request-id": "A" * 8000}) - response = self.client.get(ECHO_URL, headers={"x-request-id": sent}) + self.assertRegex(response[REQUEST_ID_ECHO_HEADER], _UUID4_RE) - self.assertEqual(response[REQUEST_ID_ECHO_HEADER], sent) + def test_id_is_provisioned_when_nothing_is_forwarded(self): + response = self.client.get(INTERNAL_URL) + + self.assertRegex(response[REQUEST_ID_ECHO_HEADER], _UUID4_RE) diff --git a/unstract/core/src/unstract/core/flask/middleware.py b/unstract/core/src/unstract/core/flask/middleware.py index 64655a78c8..dbde819c67 100644 --- a/unstract/core/src/unstract/core/flask/middleware.py +++ b/unstract/core/src/unstract/core/flask/middleware.py @@ -1,18 +1,24 @@ -import re import uuid from flask import Flask, g, request -# An incoming X-Request-ID is caller-supplied and lands in every log line and in -# the echoed response header, so it is only accepted in a shape that cannot forge -# a log record (ANSI/control characters) or bloat one (unbounded length). -SAFE_REQUEST_ID = re.compile(r"\A[A-Za-z0-9._:-]{1,128}\Z") - def _incoming_request_id() -> str: + """Adopt the caller's id, or mint one. + + Unlike the Django backend these services are not internet-facing: every + caller is another Unstract service forwarding an id, so honouring it is the + whole point. The id still lands unescaped in every log line and in the echoed + response header, so only the canonical UUID our services emit is accepted -- + which bounds length and rules out the control characters that would let a + caller forge a log record. + """ request_id = request.headers.get("X-Request-ID") - if request_id and SAFE_REQUEST_ID.match(request_id): - return request_id + try: + if request_id and str(uuid.UUID(request_id)) == request_id: + return request_id + except (AttributeError, TypeError, ValueError): + pass return str(uuid.uuid4()) diff --git a/x2text-service/app/logging_util.py b/x2text-service/app/logging_util.py index 97aca91706..d2b51066f7 100644 --- a/x2text-service/app/logging_util.py +++ b/x2text-service/app/logging_util.py @@ -5,17 +5,11 @@ """ import logging -import re import uuid from logging.config import dictConfig from flask import Flask, g, has_request_context, request -# See ``unstract.core.flask.middleware``: a caller-supplied id reaches every log -# line and the echoed response header, so only a shape that cannot forge or bloat -# a record is accepted. -SAFE_REQUEST_ID = re.compile(r"\A[A-Za-z0-9._:-]{1,128}\Z") - # Copy of the canonical format owned by ``unstract.core.flask.logging``; a # divergence silently splits this service out of the cross-service log query. LOG_FORMAT = ( @@ -94,8 +88,13 @@ def register_request_id_middleware(app: Flask) -> None: @app.before_request def _assign_request_id() -> None: + # Only the canonical UUID our services emit is adopted; see + # ``unstract.core.flask.middleware`` for why the shape is checked at all. request_id = request.headers.get("X-Request-ID") - if not (request_id and SAFE_REQUEST_ID.match(request_id)): + try: + if not (request_id and str(uuid.UUID(request_id)) == request_id): + request_id = str(uuid.uuid4()) + except (AttributeError, TypeError, ValueError): request_id = str(uuid.uuid4()) g.request_id = request_id From fc9ba72a6bd4cf4aece58823cc58fd81de903f6f Mon Sep 17 00:00:00 2001 From: Ghost Jake <89829542+Deepak-Kesavan@users.noreply.github.com> Date: Tue, 18 Aug 2026 17:23:15 +0530 Subject: [PATCH 8/8] UN-2123 [FIX] Expose X-Request-ID to the browser via CORS The frontend is a separate origin, so the browser hides every response header not named in CORS_EXPOSE_HEADERS -- which was never set, so it defaulted to empty. getRequestIdFromError therefore always fell through to the id the axios interceptor *sent*, and that happened to be right only because the backend used to adopt it. Now that the backend provisions its own id, that fallback would surface an id appearing in no log line: the error toast would show something a support engineer could search for and never find, with nothing erroring to signal it. Naming the header here makes the echo readable in JS, so the toast shows the id the backend actually logged. Pinned by a test through the real CORS middleware, which fails if the setting is dropped. --- backend/backend/settings/base.py | 4 ++++ backend/middleware/test_request_id.py | 30 +++++++++++++++++++++++++++ 2 files changed, 34 insertions(+) diff --git a/backend/backend/settings/base.py b/backend/backend/settings/base.py index 75e2526461..439a8e9c25 100644 --- a/backend/backend/settings/base.py +++ b/backend/backend/settings/base.py @@ -284,6 +284,10 @@ def get_required_setting(setting_key: str, default: str | None = None) -> str | REQUEST_ID_RESPONSE_HEADER = "X-Request-ID" GENERATE_REQUEST_ID_IF_NOT_IN_HEADER = True NO_REQUEST_ID = "-" +# The frontend is a separate origin, so the browser hides every response header +# not named here -- without this the id the backend logged is unreadable in JS +# and the error toast falls back to showing an id that appears in no log. +CORS_EXPOSE_HEADERS = [REQUEST_ID_RESPONSE_HEADER] class OTelFieldFilter(logging.Filter): diff --git a/backend/middleware/test_request_id.py b/backend/middleware/test_request_id.py index bfee1f7a30..596f3a75a2 100644 --- a/backend/middleware/test_request_id.py +++ b/backend/middleware/test_request_id.py @@ -171,3 +171,33 @@ def test_id_is_provisioned_when_nothing_is_forwarded(self): response = self.client.get(INTERNAL_URL) self.assertRegex(response[REQUEST_ID_ECHO_HEADER], _UUID4_RE) + + +@override_settings( + ROOT_URLCONF=__name__, + MIDDLEWARE=["corsheaders.middleware.CorsMiddleware"] + _MIDDLEWARE, +) +class RequestIDIsReadableByTheBrowserTest(SimpleTestCase): + """The frontend is a separate origin, so the echo only reaches JS if the + header is named in ``CORS_EXPOSE_HEADERS``. + + Without it the browser hides the header, ``getRequestIdFromError`` falls + through to the id the interceptor *sent*, and -- now that the backend + provisions its own -- the error toast shows an id that appears in no log. + Nothing errors; the id is simply wrong, which is why this is pinned. + """ + + def test_request_id_is_exposed_to_a_cross_origin_caller(self): + origin = "http://localhost:3000" + + response = self.client.get( + PUBLIC_URL, headers={"origin": origin, "x-request-id": "ignored"} + ) + + exposed = response.get("Access-Control-Expose-Headers", "") + self.assertIn( + "X-Request-ID", + [h.strip() for h in exposed.split(",")], + "the browser cannot read the id the backend logged", + ) + self.assertRegex(response["X-Request-ID"], _UUID4_RE)