UN-2123 [FEAT] Propagate request_id across services and workers - #2229
UN-2123 [FEAT] Propagate request_id across services and workers#2229Deepak-Kesavan wants to merge 8 commits into
Conversation
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).
Summary by CodeRabbit
WalkthroughThe changes add request ID middleware and canonical logging fields, propagate request IDs through Celery task headers, bind them in workers, and forward them in outbound worker HTTP requests. ChangesRequest ID observability
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant FlaskApp
participant StateStore
participant Celery
participant WorkerLogger
participant BaseClient
FlaskApp->>FlaskApp: capture or generate request_id
FlaskApp->>StateStore: store request_id
Celery->>StateStore: read request_id before publish
StateStore-->>Celery: return request_id
Celery->>Celery: add request_id to task headers
Celery->>WorkerLogger: deliver task message
WorkerLogger->>WorkerLogger: bind request_id from headers
WorkerLogger->>BaseClient: provide current request_id
BaseClient->>BaseClient: send X-Request-ID
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@x2text-service/app/logging_util.py`:
- Around line 31-33: Update RequestIDFilter.filter to call has_request_context()
before accessing Flask’s context-local g; assign the request ID from g only when
a request context exists, otherwise retain "-".
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 6da06d9c-e764-45ab-a92a-7479f521926e
📒 Files selected for processing (7)
backend/backend/celery_service.pybackend/backend/celery_signals.pyunstract/core/src/unstract/core/flask/logging.pyworkers/shared/clients/base_client.pyworkers/shared/infrastructure/logging/logger.pyx2text-service/app/config.pyx2text-service/app/logging_util.py
|
| Filename | Overview |
|---|---|
| backend/backend/celery_signals.py | Adds defensively guarded producer and backend-worker signal handlers for binding, forwarding, and clearing request IDs. |
| backend/middleware/request_id.py | Provisions server-owned IDs for public requests while accepting canonical forwarded UUIDs on authenticated internal paths. |
| workers/shared/infrastructure/logging/logger.py | Prefers propagated task headers, tracks whether an ID may be forwarded, and resets task-scoped state after execution. |
| workers/shared/clients/base_client.py | Adds the active worker request ID to internal backend calls without altering session-level headers. |
| unstract/core/src/unstract/core/flask/middleware.py | Validates incoming request IDs, mints fallbacks, and echoes the effective ID in responses. |
| x2text-service/app/logging_util.py | Introduces self-contained request-aware logging and middleware using the shared canonical log structure. |
Sequence Diagram
sequenceDiagram
participant Client
participant Backend
participant Broker as Celery Broker
participant Worker
participant Service as Flask/x2text Service
Client->>Backend: HTTP request
Backend->>Backend: Provision request_id
Backend->>Broker: Publish task with request_id header
Broker->>Worker: Deliver task
Worker->>Worker: Bind request_id to log context
Worker->>Broker: Publish child task with request_id
Worker->>Backend: Internal callback with X-Request-ID
Worker->>Service: Service call with X-Request-ID
Service->>Service: Bind request_id to logs
Service-->>Worker: Echo X-Request-ID
Backend-->>Client: Echo X-Request-ID
Reviews (6): Last reviewed commit: "UN-2123 [FIX] Expose X-Request-ID to the..." | Re-trigger Greptile
…has_request_context()
pk-zipstack
left a comment
pk-zipstack
left a comment
There was a problem hiding this comment.
Code review — 4 findings on the request_id propagation chain. The producer/consumer plumbing itself looks correct (verified against Celery 5.5.3 that a custom key in the v2 message headers is promoted onto Task.request both as an attribute and via request.headers, and that mutating the headers dict inside before_task_publish reaches producer.publish). The issues are at the two ends of the chain and in the precedence rule.
| WorkerLogger.update_context(request_id=request_id, task_id=task_id) | ||
|
|
||
|
|
||
| def _propagate_request_id_on_publish(headers=None, **_): |
There was a problem hiding this comment.
The PG-queue transport bypasses this hop entirely, so correlation breaks there.
workers/queue_backend/dispatch.py::enqueue routes to _enqueue_pg when resolve_backend(...) is QueueBackend.PG, which serialises via to_payload() and writes straight to pg_queue_message — it never calls current_app.send_task, so before_task_publish never fires and no request_id is carried. On the consume side consumer.py:575 runs task.apply(headers={FAIRNESS_HEADER_NAME: ...}), i.e. the only headers reaching Context are the fairness ones, so _request_id_from_message returns None and the task falls back to its payload id.
Concrete scenario: an API-triggered execution whose task name is in WORKER_PG_QUEUE_ENABLED_TASKS (or that rides a backend= pipeline override) loses the HTTP request_id at the first PG hop and never regains it for the rest of the execution — exactly the chains the PG rollout is moving onto. Worth threading request_id into to_payload() / restoring it into the headers passed to task.apply(), or at minimum calling this gap out in the PR's out-of-scope list.
There was a problem hiding this comment.
Confirmed — the PG transport enqueues via PgQueueMessage.objects.create(), so before_task_publish never fires and TaskPayload carries no request_id. Keeping this out of scope for this PR (Celery is the live default; pg_queue_enabled is fail-closed). The PGMQ owner has a precise carrier spec (add request_id to TaskPayload, populate in the backend and worker PG producers, inject into the consumer apply(headers=...)). Added to the PR out-of-scope section.
chandrasekharan-zipstack
left a comment
chandrasekharan-zipstack
left a comment
There was a problem hiding this comment.
A few points on hops and guarantees that I think are still open — details inline.
| # callbacks share the originating request's request_id in logs. | ||
| request_id = _current_request_id() | ||
| if request_id: | ||
| headers["X-Request-ID"] = request_id |
There was a problem hiding this comment.
The worker → runner hop stamps a different id, and it isn't in the out-of-scope list.
This closes the worker → backend hop, but the worker → runner hop already sends an X-Request-ID and it carries the wrong value:
# unstract/tool-sandbox/src/unstract/tool_sandbox/helper.py:439 (and :592)
headers = {
"X-Request-ID": file_execution_id,
}The runner reads that header via unstract.core.flask.register_request_id_middleware and renders it as request_id:<file_execution_id> — so after this PR the runner is the one service actively logging a different id under the same field name, which is exactly what breaks the single-request_id gcloud query the PR is built for.
This is called from worker context, where WorkerLogger.get_context().request_id is already bound by _bind_task_context, so the fix is the same shape as _current_request_id() here — prefer the propagated id and keep file_execution_id as the fallback.
Either close it or add it alongside the SDK1 x2text-adapter gap in the PR's out-of-scope list; right now it reads as covered.
There was a problem hiding this comment.
Good catch. tool-sandbox -> runner sends X-Request-ID = file_execution_id, which is inconsistent with standardizing on the HTTP request_id. Reconciling it means threading the request_id through tool-sandbox, which belongs with the tool-execution / OTel follow-up (same bucket as the SDK1 x2text-adapter forwarding). Added it explicitly to the PR out-of-scope list rather than leaving it silent.
…ion, 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.
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.
chandrasekharan-zipstack
left a comment
chandrasekharan-zipstack
left a comment
There was a problem hiding this comment.
Consider making code comments and generic to avoid rot over time
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.
…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.
…ting 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.
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.
|
Unstract test resultsPer-group results
Critical paths
|
3 participants
What
Completes end-to-end
request_id(X-Request-ID) correlation so a single IDcan filter logs across the backend, Celery workers and services in gcloud.
before_task_publishsignal(
backend/backend/celery_signals.py) injects the request-scopedrequest_id(from
StateStore) into every published Celery task's message headers — noper-
send_taskchanges required.task_prerunnow prefers an explicitrequest_idfrom themessage headers over payload-derived ids (
file_execution_id, etc.), and anew
before_task_publishhandler re-propagates it onto downstreamworker→worker task chains.
X-Request-IDfrom the worker log context, so backend callbacks share theoriginating request's id.
X-Request-IDand logs it via aself-contained logging module (no new dependency).
request_id/trace_id/span_idparse identically in gcloud.Why
Debugging across services is painful when logs can't be correlated. The backend
already assigned a
request_idper request (and the frontend forwards one), butit died at the backend boundary — it was never propagated to the Celery workers
or echoed on worker→backend calls, so worker logs showed
request_id:-(or anexecution_idthat didn't match the originating request). This threads the sameid through the whole chain.
Jira: UN-2123
How
The worker-side receiving/logging machinery already existed (UN-3435) and was
explicitly designed to accept a real
request_idfrom the producer — this PRsupplies the missing producer side and closes the remaining hops:
before_task_publish(backend) readsStateStore.get(Common.REQUEST_ID)andsets
headers["request_id"]. Guarded so it can never break task publishing._bind_task_context(task_prerun) readstask.request.request_id(Celery 5.6.2 exposes custom message headers on the task
Context; a rawrequest.headersmapping is used as a version-safe fallback). Resolutionorder: message-header
request_id→ payload-derived id → Celerytask_id._propagate_request_id_on_publish(before_task_publish) forwards thebound
request_idonto tasks the worker itself publishes.base_client._make_requestaddsX-Request-IDfrom the worker log context.create_app()wiressetup_logging+register_request_id_middlewarefrom a new self-containedapp/logging_util.py.unstract/core/flask/logging.py) is aligned tothe canonical backend/worker format.
Out of scope (recommended as separate tickets):
trace_idfields in every formatter, meta-deps) is currently inert — no service runs
under
opentelemetry-instrument, exporters are hardcoded tonone, and thereis no OTLP endpoint — so
trace_id/span_idalways render as-. Activatingreal distributed tracing is a larger DevOps/observability effort. This PR uses
request_id(which the frontend/backend already surface via theX-Request-IDresponse header) as the pragmatic correlation key.execution
request_idthreaded through several tool-side layers; belongs withthe tool-execution/OTel correlation follow-up. (x2text still logs any
X-Request-IDit receives after this PR.)tool-sandbox → runnercurrently sendsX-Request-ID = file_execution_id,inconsistent with the HTTP request_id standardized here; reconciling it means
threading the request_id through tool-sandbox — same tool-execution/OTel bucket.
pg_queue_enabled) — a separate, non-Celery transport(DB-INSERT enqueue + custom poll-loop consumer) that this PR's Celery-signal
propagation does not cover. Celery is the live default; a matching carrier
(
request_idonTaskPayload, populated in the PG producers, read into theconsumer's
apply(headers=…)) is handed to the PGMQ owner so correlationstays complete as
pg_queue_enabledramps.Can this PR break any existing features. If yes, please list possible items. If no, please explain why. (PS: Admins do not merge the PR without this section filled)
Low risk — the change is additive and defensively guarded:
before_task_publishhandlers no-op safely (null-guarded,StateStoreread wrapped in try/except) so they cannot break task publishing even if
no
request_idis in scope (e.g. beat-scheduled tasks fall back to theexisting
execution_id/task_idbehaviour).task.request.request_idusesgetattr(..., None)and only overridesthe previously-working payload-derived id when a header is actually present —
so existing worker correlation is preserved when the header is absent.
the canonical backend/worker format, which means Flask-service log lines
(platform-service, runner) show
module:<source-module>instead of thelogger
name. Log content is otherwise unchanged; any dashboard keying onthe logger name string would need updating (grep-by-
request_id/trace_idisunaffected and is the point of the change).
Database Migrations
None.
Env Config
None. (x2text-service reads the existing
LOG_LEVELenv if set; defaults toINFO.)Relevant Docs
Related Issues or PRs
Dependencies Versions
None added. (x2text-service deliberately avoids taking on
unstract-core; ituses a small self-contained logging module instead.)
Notes on Testing
py_compile-clean; pre-commit (ruff, ruff-format,pycln, pyupgrade, secret-scan, test-selection) green.
exposed on the task
Contextvia both attribute and.get()access.trigger a workflow/API execution and confirm a single
X-Request-IDfrom theoriginating HTTP request appears in the worker log lines (
request_id:<id>)and on the worker→backend internal-API calls. If Celery header→Context
attribute promotion ever regressed, the
request.headersfallback covers it,but the smoke test is the real guarantee.
Screenshots
N/A (no UI surface).
Checklist
I have read and understood the Contribution Guidelines.