"""Context tracking for workflow execution (correlation ID and session ID)
Provides context variable support for automatic ID propagation across async
execution paths without requiring function signature changes.
Correlation ID: Tracks a single workflow execution for logging/debugging
Session ID: Groups related operations for asset/job isolation
Example:
>>> from src.utils.context import set_correlation_id, set_session_id
>>> from src.utils.context import get_correlation_id, get_session_id
>>>
>>> # At tool entry point
>>> correlation_id = generate_correlation_id()
>>> set_correlation_id(correlation_id)
>>>
>>> session_id = generate_session_id()
>>> set_session_id(session_id)
>>>
>>> # Anywhere in the execution chain
>>> current_corr_id = get_correlation_id()
>>> current_sess_id = get_session_id()
>>> logger.info(f"[{current_corr_id}] Session {current_sess_id}: Processing...")
"""
import uuid
from contextvars import ContextVar
# Context variable for correlation ID (thread-safe, async-aware)
_correlation_id_var: ContextVar[str | None] = ContextVar("correlation_id", default=None)
# Context variable for session ID (thread-safe, async-aware)
_session_id_var: ContextVar[str | None] = ContextVar("session_id", default=None)
def get_correlation_id() -> str | None:
"""Get the current correlation ID from context.
Returns:
Current correlation ID or None if not set
"""
return _correlation_id_var.get()
def set_correlation_id(correlation_id: str) -> None:
"""Set the correlation ID in the current context.
Args:
correlation_id: Correlation ID to set (typically a UUID string)
"""
_correlation_id_var.set(correlation_id)
def generate_correlation_id() -> str:
"""Generate a new correlation ID (UUID4).
Returns:
New correlation ID as a string
"""
return str(uuid.uuid4())
def clear_correlation_id() -> None:
"""Clear the correlation ID from the current context.
Useful for cleanup or testing.
"""
_correlation_id_var.set(None)
def get_session_id() -> str | None:
"""Get the current session ID from context.
Returns:
Current session ID or None if not set
"""
return _session_id_var.get()
def set_session_id(session_id: str) -> None:
"""Set the session ID in the current context.
Args:
session_id: Session ID to set (typically a UUID string)
"""
_session_id_var.set(session_id)
def generate_session_id() -> str:
"""Generate a new session ID (UUID4).
Returns:
New session ID as a string
"""
return str(uuid.uuid4())
def clear_session_id() -> None:
"""Clear the session ID from the current context.
Useful for cleanup or testing.
"""
_session_id_var.set(None)
