feat(logging): add opt-in session_id and trace_id correlation to JSON log records via contextvars

litellm/__init__.py

@@ -196,6 +196,7 @@ def _dev_env_hot_reload_enabled() -> bool:
     None  # Fields to exclude from StandardLoggingPayload before callbacks receive it
 )
 log_raw_request_response: bool = False
+request_correlation_in_logs: bool = False
 redact_messages_in_exceptions: Optional[bool] = False
 redact_user_api_key_info: Optional[bool] = False
 # When True (default — preserves historical behavior), the Router appends

litellm/_logging.py

@@ -1,4 +1,5 @@
 import ast
+import contextvars
 import logging
 import os
 import sys
@@ -12,6 +13,18 @@
 
 set_verbose = False
 
+session_id_var: contextvars.ContextVar[str] = contextvars.ContextVar("session_id", default="")
+trace_id_var: contextvars.ContextVar[str] = contextvars.ContextVar("trace_id", default="")
+
+
+def set_session_id(session_id: str) -> None:
+    session_id_var.set(session_id)
+
+
+def set_trace_id(trace_id: str) -> None:
+    trace_id_var.set(trace_id)
+
+
 if set_verbose is True:
     logging.warning(
         "`litellm.set_verbose` is deprecated. Please set `os.environ['LITELLM_LOG'] = 'DEBUG'` for debug logs."
@@ -77,13 +90,38 @@ def filter(self, record: logging.LogRecord) -> bool:
 _secret_filter = SecretRedactionFilter()
 
 
+class CorrelationContextFilter(logging.Filter):
+    """Stamps each log record with the current request's trace_id and session_id from contextvars.
+
+    Works in tandem with JsonFormatter: the formatter's record.__dict__ loop picks up these
+    attributes as first-class JSON fields without any formatter-level code.
+    """
+
+    def filter(self, record: logging.LogRecord) -> bool:
+        import litellm
+
+        if not litellm.request_correlation_in_logs:
+            return True
+        trace_id = trace_id_var.get()
+        if trace_id:
+            record.trace_id = trace_id
+        session_id = session_id_var.get()
+        if session_id:
+            record.session_id = session_id
+        return True
+
+
+_correlation_filter = CorrelationContextFilter()
+
+
 json_logs = bool(os.getenv("JSON_LOGS", False))
 # Create a handler for the logger (you may need to adapt this based on your needs)
 log_level = os.getenv("LITELLM_LOG", "DEBUG")
 numeric_level: str = getattr(logging, log_level.upper())
 handler = logging.StreamHandler()
 handler.setLevel(numeric_level)
 handler.addFilter(_secret_filter)
+handler.addFilter(_correlation_filter)
 
 
 def _try_parse_json_message(message: str) -> Optional[Dict[str, Any]]:
@@ -196,6 +234,7 @@ def _setup_json_exception_handlers(formatter):
     error_handler = logging.StreamHandler()
     error_handler.setFormatter(formatter)
     error_handler.addFilter(_secret_filter)
+    error_handler.addFilter(_correlation_filter)
 
     # Setup excepthook for uncaught exceptions
     def json_excepthook(exc_type, exc_value, exc_traceback):
@@ -312,6 +351,7 @@ def _initialize_loggers_with_handler(handler: logging.Handler):
     - Prevents bubbling to parent/root (critical to prevent duplicate JSON logs)
     """
     handler.addFilter(_secret_filter)
+    handler.addFilter(_correlation_filter)
     for lg in _get_loggers_to_initialize():
         lg.handlers.clear()  # remove any existing handlers
         lg.addHandler(handler)  # add JSON formatter handler

litellm/litellm_core_utils/litellm_logging.py

@@ -36,7 +36,13 @@
     log_raw_request_response,
     turn_off_message_logging,
 )
-from litellm._logging import _is_debugging_on, _redact_string, verbose_logger
+from litellm._logging import (
+    _is_debugging_on,
+    _redact_string,
+    set_session_id,
+    set_trace_id,
+    verbose_logger,
+)
 from litellm.exceptions import (
     validate_rate_limit_category,
     validate_rate_limit_type,
@@ -345,6 +351,11 @@ def __init__(
         self.call_type = call_type
         self.litellm_call_id = litellm_call_id
         self.litellm_trace_id: str = litellm_trace_id if litellm_trace_id else str(uuid.uuid4())
+
+        set_trace_id(self.litellm_trace_id)
+        _sid = (kwargs or {}).get("litellm_session_id")
+        set_session_id(str(_sid) if _sid else "")
+
         self.function_id = function_id
         self.streaming_chunks: List[Any] = []  # for generating complete stream response
         self.sync_streaming_chunks: List[Any] = []  # for generating complete stream response

tests/test_litellm/test_logging.py

@@ -15,9 +15,14 @@
 import litellm
 from litellm._logging import (
     ALL_LOGGERS,
+    CorrelationContextFilter,
     JsonFormatter,
     _initialize_loggers_with_handler,
     _turn_on_json,
+    session_id_var,
+    set_session_id,
+    set_trace_id,
+    trace_id_var,
     verbose_logger,
     verbose_proxy_logger,
     verbose_router_logger,
@@ -328,3 +333,164 @@ async def test_cache_hit_includes_custom_llm_provider():
         # Clean up
         litellm.callbacks = original_callbacks
         litellm.cache = None
+
+
+class _JsonCapture(logging.Handler):
+    def __init__(self):
+        super().__init__()
+        self.formatter = JsonFormatter()
+        self.records: list[dict] = []
+        self.addFilter(CorrelationContextFilter())
+
+    def emit(self, record):
+        self.records.append(json.loads(self.formatter.format(record)))
+
+
+def _make_capture_logger(name: str) -> tuple[logging.Logger, _JsonCapture]:
+    lg = logging.getLogger(name)
+    cap = _JsonCapture()
+    lg.addHandler(cap)
+    lg.setLevel(logging.DEBUG)
+    return lg, cap
+
+
+def test_trace_id_injected_into_json_record(monkeypatch):
+    """trace_id set via set_trace_id() appears in every JSON record in that context."""
+    monkeypatch.setattr(litellm, "request_correlation_in_logs", True)
+    lg, cap = _make_capture_logger("test.trace_inject")
+    set_trace_id("trace-abc-123")
+    try:
+        lg.info("test message")
+        assert len(cap.records) == 1
+        assert cap.records[0]["trace_id"] == "trace-abc-123"
+    finally:
+        trace_id_var.set("")
+
+
+def test_session_id_injected_when_set(monkeypatch):
+    """session_id set via set_session_id() appears in JSON record."""
+    monkeypatch.setattr(litellm, "request_correlation_in_logs", True)
+    lg, cap = _make_capture_logger("test.session_inject")
+    set_session_id("sess-xyz-456")
+    try:
+        lg.info("another message")
+        assert cap.records[0]["session_id"] == "sess-xyz-456"
+    finally:
+        session_id_var.set("")
+
+
+def test_session_id_absent_when_not_set():
+    """session_id must NOT appear in JSON record when not set for this context."""
+    lg, cap = _make_capture_logger("test.no_session")
+    session_id_var.set("")
+    lg.info("no session message")
+    assert "session_id" not in cap.records[0]
+
+
+def test_trace_id_absent_when_not_set():
+    """trace_id must NOT appear when not set."""
+    lg, cap = _make_capture_logger("test.no_trace")
+    trace_id_var.set("")
+    lg.info("no trace message")
+    assert "trace_id" not in cap.records[0]
+
+
+@pytest.mark.asyncio
+async def test_contextvar_isolation_between_tasks():
+    """Two concurrent async tasks each see only their own trace_id."""
+    results: dict[str, str] = {}
+
+    async def task(task_id: str, trace_id: str) -> None:
+        set_trace_id(trace_id)
+        await asyncio.sleep(0)
+        results[task_id] = trace_id_var.get()
+
+    await asyncio.gather(
+        task("A", "trace-for-A"),
+        task("B", "trace-for-B"),
+    )
+
+    assert results["A"] == "trace-for-A"
+    assert results["B"] == "trace-for-B"
+
+
+def test_logging_init_sets_trace_id():
+    """Logging.__init__() must call set_trace_id with self.litellm_trace_id."""
+    from litellm.litellm_core_utils.litellm_logging import Logging
+
+    trace_id_var.set("")
+
+    log_obj = Logging(
+        model="gpt-3.5-turbo",
+        messages=[{"role": "user", "content": "hi"}],
+        stream=False,
+        call_type="completion",
+        start_time=None,
+        litellm_call_id="call-001",
+        function_id="fn-001",
+        kwargs={},
+    )
+    assert trace_id_var.get() == log_obj.litellm_trace_id
+
+
+def test_logging_init_sets_session_id_when_provided():
+    """Logging.__init__() must call set_session_id when litellm_session_id is in kwargs."""
+    from litellm.litellm_core_utils.litellm_logging import Logging
+
+    session_id_var.set("")
+
+    Logging(
+        model="gpt-3.5-turbo",
+        messages=[{"role": "user", "content": "hi"}],
+        stream=False,
+        call_type="completion",
+        start_time=None,
+        litellm_call_id="call-002",
+        function_id="fn-002",
+        kwargs={"litellm_session_id": "my-session-99"},
+    )
+    assert session_id_var.get() == "my-session-99"
+
+
+def test_logging_init_resets_session_id_to_empty_when_absent():
+    """When no session_id is in kwargs, Logging.__init__() must reset session_id_var to ""
+    so a prior request's session_id does not leak into subsequent log records."""
+    from litellm.litellm_core_utils.litellm_logging import Logging
+
+    session_id_var.set("preexisting-sid")
+
+    Logging(
+        model="gpt-3.5-turbo",
+        messages=[{"role": "user", "content": "hi"}],
+        stream=False,
+        call_type="completion",
+        start_time=None,
+        litellm_call_id="call-003",
+        function_id="fn-003",
+        kwargs={},
+    )
+    assert session_id_var.get() == ""
+
+
+def test_trace_id_not_in_log_when_flag_disabled(monkeypatch):
+    """When request_correlation_in_logs is False (default), trace_id must not appear in JSON records even when set."""
+    monkeypatch.setattr(litellm, "request_correlation_in_logs", False)
+    lg, cap = _make_capture_logger("test.no_trace_gated")
+    set_trace_id("trace-should-not-appear")
+    try:
+        lg.info("message")
+        assert "trace_id" not in cap.records[0]
+    finally:
+        trace_id_var.set("")
+
+
+def test_session_id_not_in_log_when_flag_disabled(monkeypatch):
+    """When request_correlation_in_logs is False (default), session_id must not appear in JSON records even when set."""
+    monkeypatch.setattr(litellm, "request_correlation_in_logs", False)
+    lg, cap = _make_capture_logger("test.no_session_gated")
+    set_session_id("sess-should-not-appear")
+    try:
+        lg.info("message")
+        assert "session_id" not in cap.records[0]
+    finally:
+        session_id_var.set("")