"""Unified context management for request correlation and distributed tracing.
This module provides a single source of truth for request context propagation
across the foundry-mcp codebase, including:
- Correlation ID generation and propagation
- W3C Trace Context (traceparent/tracestate) support
- Thread-safe context variables via contextvars
- Both async and sync context managers
Usage:
from foundry_mcp.core.context import (
request_context,
sync_request_context,
get_correlation_id,
get_current_context,
generate_correlation_id,
)
# Async usage
async with request_context(client_id="user123") as ctx:
print(ctx.correlation_id) # e.g., "req_a1b2c3d4e5f6"
# Sync usage
with sync_request_context() as ctx:
print(ctx.correlation_id)
# Manual ID generation
corr_id = generate_correlation_id(prefix="task") # "task_a1b2c3d4e5f6"
"""
from __future__ import annotations
import re
import secrets
import time
from contextlib import contextmanager
from contextvars import ContextVar
from dataclasses import dataclass, field
from typing import Any, Dict, Generator, Optional
__all__ = [
# Context variables (for advanced use)
"correlation_id_var",
"client_id_var",
"start_time_var",
"trace_context_var",
# Dataclasses
"W3CTraceContext",
"RequestContext",
# ID generation
"generate_correlation_id",
# Context managers
"request_context",
"sync_request_context",
# Accessors
"get_correlation_id",
"get_client_id",
"get_start_time",
"get_trace_context",
"get_current_context",
# W3C helpers
"parse_traceparent",
"format_traceparent",
]
# -----------------------------------------------------------------------------
# Context Variables
# -----------------------------------------------------------------------------
correlation_id_var: ContextVar[str] = ContextVar("correlation_id", default="")
"""Request correlation ID for tracing requests across components."""
client_id_var: ContextVar[str] = ContextVar("client_id", default="anonymous")
"""Identifier for the client making the request."""
start_time_var: ContextVar[float] = ContextVar("start_time", default=0.0)
"""Request start time as Unix timestamp."""
trace_context_var: ContextVar[Optional["W3CTraceContext"]] = ContextVar(
"trace_context", default=None
)
"""W3C Trace Context for distributed tracing integration."""
# -----------------------------------------------------------------------------
# W3C Trace Context Support
# -----------------------------------------------------------------------------
# W3C traceparent format: version-trace_id-parent_id-flags
# Example: 00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01
_TRACEPARENT_REGEX = re.compile(
r"^(?P<version>[0-9a-f]{2})-"
r"(?P<trace_id>[0-9a-f]{32})-"
r"(?P<parent_id>[0-9a-f]{16})-"
r"(?P<flags>[0-9a-f]{2})$"
)
@dataclass(frozen=True)
class W3CTraceContext:
"""W3C Trace Context representation for distributed tracing.
Implements the W3C Trace Context specification:
https://www.w3.org/TR/trace-context/
Attributes:
version: Trace context version (currently "00")
trace_id: 32-character hex trace identifier
parent_id: 16-character hex span/parent identifier
flags: 2-character hex flags (01 = sampled)
tracestate: Optional vendor-specific trace state
"""
version: str = "00"
trace_id: str = ""
parent_id: str = ""
flags: str = "00"
tracestate: Optional[str] = None
@classmethod
def parse(
cls,
traceparent: Optional[str] = None,
tracestate: Optional[str] = None,
) -> Optional["W3CTraceContext"]:
"""Parse W3C traceparent and tracestate headers.
Args:
traceparent: The traceparent header value
tracestate: Optional tracestate header value
Returns:
Parsed W3CTraceContext or None if parsing fails
"""
if not traceparent:
return None
match = _TRACEPARENT_REGEX.match(traceparent.lower().strip())
if not match:
return None
return cls(
version=match.group("version"),
trace_id=match.group("trace_id"),
parent_id=match.group("parent_id"),
flags=match.group("flags"),
tracestate=tracestate,
)
@classmethod
def generate(cls, sampled: bool = True) -> "W3CTraceContext":
"""Generate a new W3C Trace Context.
Args:
sampled: Whether this trace should be sampled
Returns:
New W3CTraceContext with generated IDs
"""
return cls(
version="00",
trace_id=secrets.token_hex(16), # 32 hex chars
parent_id=secrets.token_hex(8), # 16 hex chars
flags="01" if sampled else "00",
)
@property
def is_sampled(self) -> bool:
"""Check if trace is sampled based on flags."""
try:
return (int(self.flags, 16) & 0x01) == 0x01
except ValueError:
return False
@property
def traceparent(self) -> str:
"""Format as W3C traceparent header value."""
return f"{self.version}-{self.trace_id}-{self.parent_id}-{self.flags}"
def with_new_parent(self, sampled: Optional[bool] = None) -> "W3CTraceContext":
"""Create child context with new parent_id, preserving trace_id.
Args:
sampled: Override sampling decision (None = inherit)
Returns:
New W3CTraceContext with same trace_id but new parent_id
"""
flags = self.flags
if sampled is not None:
flags = "01" if sampled else "00"
return W3CTraceContext(
version=self.version,
trace_id=self.trace_id,
parent_id=secrets.token_hex(8),
flags=flags,
tracestate=self.tracestate,
)
def to_dict(self) -> Dict[str, Any]:
"""Convert to dictionary for serialization."""
result = {
"version": self.version,
"trace_id": self.trace_id,
"parent_id": self.parent_id,
"flags": self.flags,
"sampled": self.is_sampled,
}
if self.tracestate:
result["tracestate"] = self.tracestate
return result
def parse_traceparent(header: Optional[str]) -> Optional[W3CTraceContext]:
"""Parse a traceparent header into W3CTraceContext.
Args:
header: The traceparent header value
Returns:
Parsed context or None
"""
return W3CTraceContext.parse(header)
def format_traceparent(ctx: W3CTraceContext) -> str:
"""Format W3CTraceContext as traceparent header.
Args:
ctx: The trace context
Returns:
Formatted traceparent header value
"""
return ctx.traceparent
# -----------------------------------------------------------------------------
# Correlation ID Generation
# -----------------------------------------------------------------------------
def generate_correlation_id(prefix: str = "req") -> str:
"""Generate a unique correlation ID with optional prefix.
Format: {prefix}_{12_hex_chars}
Example: "req_a1b2c3d4e5f6"
Args:
prefix: ID prefix (default: "req")
Returns:
Unique correlation ID string
"""
return f"{prefix}_{secrets.token_hex(6)}"
# -----------------------------------------------------------------------------
# Request Context
# -----------------------------------------------------------------------------
@dataclass
class RequestContext:
"""Immutable snapshot of the current request context.
This dataclass captures all context variables at a point in time,
providing convenient access to timing information and serialization.
Attributes:
correlation_id: Unique request identifier
client_id: Client/user identifier
start_time: Request start timestamp
trace_context: Optional W3C distributed tracing context
"""
correlation_id: str = ""
client_id: str = "anonymous"
start_time: float = field(default_factory=time.time)
trace_context: Optional[W3CTraceContext] = None
@property
def elapsed_seconds(self) -> float:
"""Calculate elapsed time since request start in seconds."""
if self.start_time <= 0:
return 0.0
return time.time() - self.start_time
@property
def elapsed_ms(self) -> float:
"""Calculate elapsed time since request start in milliseconds."""
return self.elapsed_seconds * 1000
def to_dict(self) -> Dict[str, Any]:
"""Convert context to dictionary for logging/serialization.
Returns:
Dictionary with all context fields
"""
result: Dict[str, Any] = {
"correlation_id": self.correlation_id,
"client_id": self.client_id,
"start_time": self.start_time,
"elapsed_ms": round(self.elapsed_ms, 2),
}
if self.trace_context:
result["trace"] = self.trace_context.to_dict()
return result
# -----------------------------------------------------------------------------
# Context Managers
# -----------------------------------------------------------------------------
@contextmanager
def sync_request_context(
*,
correlation_id: Optional[str] = None,
client_id: Optional[str] = None,
traceparent: Optional[str] = None,
tracestate: Optional[str] = None,
) -> Generator[RequestContext, None, None]:
"""Synchronous context manager for request context.
Sets up context variables for the duration of the with block,
automatically cleaning up on exit.
Args:
correlation_id: Request ID (auto-generated if None)
client_id: Client identifier (default: "anonymous")
traceparent: W3C traceparent header for distributed tracing
tracestate: W3C tracestate header
Yields:
RequestContext snapshot
Example:
with sync_request_context(client_id="user123") as ctx:
logger.info(f"Processing request {ctx.correlation_id}")
"""
# Generate or use provided correlation ID
corr_id = correlation_id or generate_correlation_id()
client = client_id or "anonymous"
start = time.time()
# Parse trace context if provided
trace_ctx = W3CTraceContext.parse(traceparent, tracestate)
# Set context variables
token_corr = correlation_id_var.set(corr_id)
token_client = client_id_var.set(client)
token_start = start_time_var.set(start)
token_trace = trace_context_var.set(trace_ctx)
try:
yield RequestContext(
correlation_id=corr_id,
client_id=client,
start_time=start,
trace_context=trace_ctx,
)
finally:
# Reset context variables
correlation_id_var.reset(token_corr)
client_id_var.reset(token_client)
start_time_var.reset(token_start)
trace_context_var.reset(token_trace)
async def request_context(
*,
correlation_id: Optional[str] = None,
client_id: Optional[str] = None,
traceparent: Optional[str] = None,
tracestate: Optional[str] = None,
) -> RequestContext:
"""Async context manager for request context.
This is an async generator that can be used with `async with`.
Sets up context variables for the duration of the async with block.
Args:
correlation_id: Request ID (auto-generated if None)
client_id: Client identifier (default: "anonymous")
traceparent: W3C traceparent header for distributed tracing
tracestate: W3C tracestate header
Returns:
RequestContext snapshot (use with async context manager)
Example:
async with request_context(client_id="user123") as ctx:
logger.info(f"Processing request {ctx.correlation_id}")
"""
# This uses the sync implementation since contextvars work across
# sync/async boundaries in Python 3.7+
with sync_request_context(
correlation_id=correlation_id,
client_id=client_id,
traceparent=traceparent,
tracestate=tracestate,
) as ctx:
return ctx
# Make request_context work as both async and sync context manager
class _AsyncContextManager:
"""Wrapper to make request_context work as async context manager."""
def __init__(
self,
correlation_id: Optional[str] = None,
client_id: Optional[str] = None,
traceparent: Optional[str] = None,
tracestate: Optional[str] = None,
):
self.correlation_id = correlation_id
self.client_id = client_id
self.traceparent = traceparent
self.tracestate = tracestate
self._sync_cm: Optional[Generator[RequestContext, None, None]] = None
self._ctx: Optional[RequestContext] = None
async def __aenter__(self) -> RequestContext:
self._sync_cm = sync_request_context(
correlation_id=self.correlation_id,
client_id=self.client_id,
traceparent=self.traceparent,
tracestate=self.tracestate,
)
self._ctx = self._sync_cm.__enter__()
return self._ctx
async def __aexit__(self, exc_type, exc_val, exc_tb):
if self._sync_cm:
self._sync_cm.__exit__(exc_type, exc_val, exc_tb)
return False
def async_request_context(
*,
correlation_id: Optional[str] = None,
client_id: Optional[str] = None,
traceparent: Optional[str] = None,
tracestate: Optional[str] = None,
) -> _AsyncContextManager:
"""Create an async context manager for request context.
Args:
correlation_id: Request ID (auto-generated if None)
client_id: Client identifier (default: "anonymous")
traceparent: W3C traceparent header for distributed tracing
tracestate: W3C tracestate header
Returns:
Async context manager yielding RequestContext
Example:
async with async_request_context(client_id="user123") as ctx:
await some_async_operation()
logger.info(f"Completed request {ctx.correlation_id}")
"""
return _AsyncContextManager(
correlation_id=correlation_id,
client_id=client_id,
traceparent=traceparent,
tracestate=tracestate,
)
# -----------------------------------------------------------------------------
# Context Accessors
# -----------------------------------------------------------------------------
def get_correlation_id() -> str:
"""Get the current correlation ID from context.
Returns:
Current correlation ID or empty string if not set
"""
return correlation_id_var.get()
def get_client_id() -> str:
"""Get the current client ID from context.
Returns:
Current client ID or "anonymous" if not set
"""
return client_id_var.get()
def get_start_time() -> float:
"""Get the request start time from context.
Returns:
Start time as Unix timestamp or 0.0 if not set
"""
return start_time_var.get()
def get_trace_context() -> Optional[W3CTraceContext]:
"""Get the W3C trace context if set.
Returns:
Current W3CTraceContext or None
"""
return trace_context_var.get()
def get_current_context() -> RequestContext:
"""Get a snapshot of all current context values.
Returns:
RequestContext with current values from context variables
"""
return RequestContext(
correlation_id=get_correlation_id(),
client_id=get_client_id(),
start_time=get_start_time(),
trace_context=get_trace_context(),
)
# -----------------------------------------------------------------------------
# Backward Compatibility Aliases
# -----------------------------------------------------------------------------
# For backward compatibility with concurrency.py
request_id = correlation_id_var
"""Alias for correlation_id_var (backward compatibility)."""
def get_request_id() -> str:
"""Alias for get_correlation_id() (backward compatibility)."""
return get_correlation_id()
