|
46 | 46 | task_message_update_adapter = TypeAdapter(TaskMessageUpdate) |
47 | 47 |
|
48 | 48 |
|
| 49 | +def _attach_incoming_trace_context(scope: Scope) -> Any: |
| 50 | + """Extract the W3C trace context (traceparent/tracestate) from the inbound |
| 51 | + request headers and attach it as the current OpenTelemetry context, so spans |
| 52 | + created while handling this request adopt the caller's observability trace. |
| 53 | +
|
| 54 | + Returns a detach token (pass to opentelemetry.context.detach) or None if OTel |
| 55 | + is unavailable or no context could be established. Best-effort: never raises. |
| 56 | + """ |
| 57 | + try: |
| 58 | + from opentelemetry import context as otel_context |
| 59 | + from opentelemetry.propagate import extract |
| 60 | + except ImportError: |
| 61 | + return None |
| 62 | + try: |
| 63 | + carrier = { |
| 64 | + k.decode("latin-1"): v.decode("latin-1") |
| 65 | + for k, v in scope.get("headers", []) |
| 66 | + } |
| 67 | + ctx = extract(carrier) |
| 68 | + return otel_context.attach(ctx) |
| 69 | + except Exception: # pragma: no cover - propagation must never break a request |
| 70 | + return None |
| 71 | + |
| 72 | + |
| 73 | +def _detach_trace_context(token: Any) -> None: |
| 74 | + if token is None: |
| 75 | + return |
| 76 | + try: |
| 77 | + from opentelemetry import context as otel_context |
| 78 | + |
| 79 | + otel_context.detach(token) |
| 80 | + except Exception: # pragma: no cover |
| 81 | + pass |
| 82 | + |
| 83 | + |
49 | 84 | class RequestIDMiddleware: |
50 | 85 | """Pure ASGI middleware to set request IDs without buffering streaming responses.""" |
51 | 86 |
|
52 | 87 | def __init__(self, app: ASGIApp) -> None: |
53 | 88 | self.app = app |
54 | 89 |
|
55 | 90 | async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: |
56 | | - if scope["type"] == "http": |
57 | | - headers = dict(scope.get("headers", [])) |
58 | | - raw_request_id = headers.get(b"x-request-id", b"") |
59 | | - request_id = raw_request_id.decode() if raw_request_id else uuid.uuid4().hex |
60 | | - ctx_var_request_id.set(request_id) |
61 | | - await self.app(scope, receive, send) |
| 91 | + if scope["type"] != "http": |
| 92 | + await self.app(scope, receive, send) |
| 93 | + return |
| 94 | + |
| 95 | + headers = dict(scope.get("headers", [])) |
| 96 | + raw_request_id = headers.get(b"x-request-id", b"") |
| 97 | + request_id = raw_request_id.decode() if raw_request_id else uuid.uuid4().hex |
| 98 | + ctx_var_request_id.set(request_id) |
| 99 | + |
| 100 | + # Establish the caller's W3C trace context for the duration of the |
| 101 | + # request so business spans created here (and any downstream Temporal |
| 102 | + # workflow started from this request) share the same observability trace. |
| 103 | + token = _attach_incoming_trace_context(scope) |
| 104 | + try: |
| 105 | + await self.app(scope, receive, send) |
| 106 | + finally: |
| 107 | + _detach_trace_context(token) |
62 | 108 |
|
63 | 109 |
|
64 | 110 | class BaseACPServer(FastAPI): |
|
0 commit comments