mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-26 16:07:53 +00:00
0f088033fe
When TraceMiddleware binds a valid inbound X-Trace-Id, worker resolution prefers that id over config.metadata.deerflow_trace_id so logs, response headers, Langfuse, and runtime context stay aligned. Also add trace-context marker reset coverage for the inbound-header flag.
120 lines
4.3 KiB
Python
120 lines
4.3 KiB
Python
"""Request trace context helpers.
|
|
|
|
The value stored here is DeerFlow's request-level correlation id. It is
|
|
separate from Langfuse's own trace id and from DeerFlow run ids.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import uuid
|
|
from collections.abc import Iterator
|
|
from contextlib import contextmanager
|
|
from contextvars import ContextVar, Token
|
|
from typing import Final
|
|
|
|
TRACE_ID_HEADER: Final[str] = "X-Trace-Id"
|
|
DEERFLOW_TRACE_METADATA_KEY: Final[str] = "deerflow_trace_id"
|
|
_MAX_TRACE_ID_LENGTH: Final[int] = 512
|
|
|
|
_current_trace_id: Final[ContextVar[str | None]] = ContextVar("deerflow_current_trace_id", default=None)
|
|
_trace_id_from_request_header: Final[ContextVar[bool]] = ContextVar(
|
|
"deerflow_trace_id_from_request_header",
|
|
default=False,
|
|
)
|
|
|
|
|
|
def generate_trace_id() -> str:
|
|
"""Return a fresh header-safe trace id."""
|
|
return uuid.uuid4().hex
|
|
|
|
|
|
def normalize_trace_id(value: object) -> str | None:
|
|
"""Return a safe trace id string, or ``None`` when *value* is unusable.
|
|
|
|
Only printable ASCII (0x20-0x7E) is accepted. Codepoints above 0x7E are
|
|
rejected because the trace id round-trips through HTTP response headers,
|
|
which Starlette encodes as latin-1: codepoints > 0xFF raise
|
|
``UnicodeEncodeError`` inside ``MutableHeaders.__setitem__`` (forcing a
|
|
500 before the response body is even dispatched), and C1 controls
|
|
(0x80-0x9F) technically encode but are stripped or rejected by hardened
|
|
intermediaries (nginx / envoy / cloudfront), silently breaking the
|
|
response. C0 controls (< 0x20) and DEL (0x7F) are rejected for the same
|
|
header-safety reason plus log-injection defense.
|
|
"""
|
|
if not isinstance(value, str):
|
|
return None
|
|
trace_id = value.strip()
|
|
if not trace_id or len(trace_id) > _MAX_TRACE_ID_LENGTH:
|
|
return None
|
|
if any(ord(ch) < 32 or ord(ch) > 126 for ch in trace_id):
|
|
return None
|
|
return trace_id
|
|
|
|
|
|
def set_current_trace_id(trace_id: str) -> Token[str | None]:
|
|
"""Bind *trace_id* to the current execution context."""
|
|
normalized = normalize_trace_id(trace_id)
|
|
if normalized is None:
|
|
normalized = generate_trace_id()
|
|
return _current_trace_id.set(normalized)
|
|
|
|
|
|
def reset_current_trace_id(token: Token[str | None]) -> None:
|
|
"""Restore the trace context captured by *token*."""
|
|
_current_trace_id.reset(token)
|
|
|
|
|
|
def get_current_trace_id() -> str | None:
|
|
"""Return the current request trace id, if one is bound."""
|
|
return _current_trace_id.get()
|
|
|
|
|
|
def mark_trace_id_from_request_header(*, from_header: bool) -> Token[bool]:
|
|
"""Record whether the current trace id came from a valid inbound header."""
|
|
return _trace_id_from_request_header.set(from_header)
|
|
|
|
|
|
def reset_trace_id_from_request_header(token: Token[bool]) -> None:
|
|
"""Restore the inbound-header flag captured by *token*."""
|
|
_trace_id_from_request_header.reset(token)
|
|
|
|
|
|
def is_trace_id_from_request_header() -> bool:
|
|
"""Return ``True`` when a valid ``X-Trace-Id`` header bound the request."""
|
|
return _trace_id_from_request_header.get()
|
|
|
|
|
|
def resolve_deerflow_trace_id(metadata_trace_id: object) -> str | None:
|
|
"""Resolve the effective ``deerflow_trace_id`` for a run.
|
|
|
|
When Gateway ``TraceMiddleware`` bound a valid inbound ``X-Trace-Id``,
|
|
that value wins over ``config.metadata.deerflow_trace_id`` so logs,
|
|
response headers, Langfuse, and runtime context stay aligned. Otherwise
|
|
caller metadata wins, then the ambient request trace context.
|
|
"""
|
|
if is_trace_id_from_request_header():
|
|
return get_current_trace_id()
|
|
return normalize_trace_id(metadata_trace_id) or get_current_trace_id()
|
|
|
|
|
|
@contextmanager
|
|
def request_trace_context(trace_id: str | None = None) -> Iterator[str]:
|
|
"""Bind a request trace id for the duration of a request or entry point."""
|
|
normalized = normalize_trace_id(trace_id) or generate_trace_id()
|
|
token = _current_trace_id.set(normalized)
|
|
try:
|
|
yield normalized
|
|
finally:
|
|
_current_trace_id.reset(token)
|
|
|
|
|
|
@contextmanager
|
|
def ensure_trace_context(trace_id: str | None = None) -> Iterator[str]:
|
|
"""Bind *trace_id*, inherit the current trace, or create a fresh one."""
|
|
normalized = normalize_trace_id(trace_id) or get_current_trace_id() or generate_trace_id()
|
|
token = _current_trace_id.set(normalized)
|
|
try:
|
|
yield normalized
|
|
finally:
|
|
_current_trace_id.reset(token)
|