From 8ab758623da793b7293f80b2556bfc289ce62f24 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Wed, 1 Jul 2026 08:48:01 -0400 Subject: [PATCH 1/5] feat(logging): add opt-in session_id/trace_id correlation to JSON log records via contextvars MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds two ContextVar instances (session_id_var, trace_id_var) to litellm/_logging.py and two setter functions (set_session_id, set_trace_id). Logging.__init__() now calls both setters after assigning litellm_trace_id so every JSON log record emitted within the async request context carries trace_id and, when provided, session_id — enabling log correlation in Loki, CloudWatch Logs Insights, and other structured-log sinks without any changes to individual log call sites. Co-Authored-By: Claude Sonnet 4.6 --- litellm/_logging.py | 20 +++ litellm/litellm_core_utils/litellm_logging.py | 14 +- tests/test_litellm/test_logging.py | 138 ++++++++++++++++++ 3 files changed, 171 insertions(+), 1 deletion(-) diff --git a/litellm/_logging.py b/litellm/_logging.py index 5f3c483869d9..3954b9d7917e 100644 --- a/litellm/_logging.py +++ b/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." @@ -184,6 +197,13 @@ def format(self, record): if "logger" not in json_record: json_record["logger"] = f"{record.filename}:{record.lineno}" + session_id = session_id_var.get() + if session_id: + json_record["session_id"] = session_id + trace_id = trace_id_var.get() + if trace_id: + json_record["trace_id"] = trace_id + if record.exc_info: json_record["stacktrace"] = record.exc_text or self.formatException(record.exc_info) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index db7c1c7dfb4b..aa2d58a5de6f 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/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,12 @@ 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") + if _sid: + set_session_id(str(_sid)) + 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 diff --git a/tests/test_litellm/test_logging.py b/tests/test_litellm/test_logging.py index fbed044445b4..1908d650a962 100644 --- a/tests/test_litellm/test_logging.py +++ b/tests/test_litellm/test_logging.py @@ -18,6 +18,10 @@ 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 +332,137 @@ 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] = [] + + 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(): + """trace_id set via set_trace_id() appears in every JSON record in that context.""" + 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(): + """session_id set via set_session_id() appears in JSON record.""" + 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_does_not_set_session_id_when_absent(): + """Logging.__init__() must NOT mutate session_id_var when no session_id is provided.""" + 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() == "preexisting-sid" + session_id_var.set("") From b42250afa70c161023ff3e92838c5830c8e6a47d Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Tue, 14 Jul 2026 18:25:36 -0400 Subject: [PATCH 2/5] fix(logging): guard session_id/trace_id injection against overwriting caller-supplied extra fields --- litellm/_logging.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/_logging.py b/litellm/_logging.py index 3954b9d7917e..b08dc92bfcc1 100644 --- a/litellm/_logging.py +++ b/litellm/_logging.py @@ -198,10 +198,10 @@ def format(self, record): json_record["logger"] = f"{record.filename}:{record.lineno}" session_id = session_id_var.get() - if session_id: + if session_id and "session_id" not in json_record: json_record["session_id"] = session_id trace_id = trace_id_var.get() - if trace_id: + if trace_id and "trace_id" not in json_record: json_record["trace_id"] = trace_id if record.exc_info: From b0b13c9f15ca2f10a97f272ca9beecb33c84306d Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Wed, 15 Jul 2026 10:54:06 -0400 Subject: [PATCH 3/5] fix(logging): always reset session_id_var to empty string when no session_id provided --- litellm/litellm_core_utils/litellm_logging.py | 3 +-- tests/test_litellm/test_logging.py | 8 ++++---- 2 files changed, 5 insertions(+), 6 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index aa2d58a5de6f..9bf15c943b11 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -354,8 +354,7 @@ def __init__( set_trace_id(self.litellm_trace_id) _sid = (kwargs or {}).get("litellm_session_id") - if _sid: - set_session_id(str(_sid)) + set_session_id(str(_sid) if _sid else "") self.function_id = function_id self.streaming_chunks: List[Any] = [] # for generating complete stream response diff --git a/tests/test_litellm/test_logging.py b/tests/test_litellm/test_logging.py index 1908d650a962..2e61be4602c9 100644 --- a/tests/test_litellm/test_logging.py +++ b/tests/test_litellm/test_logging.py @@ -448,8 +448,9 @@ def test_logging_init_sets_session_id_when_provided(): assert session_id_var.get() == "my-session-99" -def test_logging_init_does_not_set_session_id_when_absent(): - """Logging.__init__() must NOT mutate session_id_var when no session_id is provided.""" +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") @@ -464,5 +465,4 @@ def test_logging_init_does_not_set_session_id_when_absent(): function_id="fn-003", kwargs={}, ) - assert session_id_var.get() == "preexisting-sid" - session_id_var.set("") + assert session_id_var.get() == "" From 18670ec4d8426a40e901d209c2d3dbaa39540591 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Thu, 23 Jul 2026 10:47:05 -0400 Subject: [PATCH 4/5] feat: gate request correlation IDs in logs behind request_correlation_in_logs flag --- litellm/__init__.py | 1 + litellm/_logging.py | 15 +++++++++------ tests/test_litellm/test_logging.py | 30 ++++++++++++++++++++++++++++-- 3 files changed, 38 insertions(+), 8 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index 59bc8c37fd2f..16e9f01d22b3 100644 --- a/litellm/__init__.py +++ b/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 diff --git a/litellm/_logging.py b/litellm/_logging.py index b08dc92bfcc1..80dec96c07ac 100644 --- a/litellm/_logging.py +++ b/litellm/_logging.py @@ -197,12 +197,15 @@ def format(self, record): if "logger" not in json_record: json_record["logger"] = f"{record.filename}:{record.lineno}" - session_id = session_id_var.get() - if session_id and "session_id" not in json_record: - json_record["session_id"] = session_id - trace_id = trace_id_var.get() - if trace_id and "trace_id" not in json_record: - json_record["trace_id"] = trace_id + import litellm + + if litellm.request_correlation_in_logs: + session_id = session_id_var.get() + if session_id and "session_id" not in json_record: + json_record["session_id"] = session_id + trace_id = trace_id_var.get() + if trace_id and "trace_id" not in json_record: + json_record["trace_id"] = trace_id if record.exc_info: json_record["stacktrace"] = record.exc_text or self.formatException(record.exc_info) diff --git a/tests/test_litellm/test_logging.py b/tests/test_litellm/test_logging.py index 2e61be4602c9..7839fee6a35f 100644 --- a/tests/test_litellm/test_logging.py +++ b/tests/test_litellm/test_logging.py @@ -352,8 +352,9 @@ def _make_capture_logger(name: str) -> tuple[logging.Logger, _JsonCapture]: return lg, cap -def test_trace_id_injected_into_json_record(): +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: @@ -364,8 +365,9 @@ def test_trace_id_injected_into_json_record(): trace_id_var.set("") -def test_session_id_injected_when_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: @@ -466,3 +468,27 @@ def test_logging_init_resets_session_id_to_empty_when_absent(): 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("") From 07b24cb44eabfd17a28376ff69121b876f2ee055 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Thu, 23 Jul 2026 11:12:32 -0400 Subject: [PATCH 5/5] refactor: move correlation ID injection into CorrelationContextFilter --- litellm/_logging.py | 37 ++++++++++++++++++++++-------- tests/test_litellm/test_logging.py | 2 ++ 2 files changed, 29 insertions(+), 10 deletions(-) diff --git a/litellm/_logging.py b/litellm/_logging.py index 80dec96c07ac..938c74ec344a 100644 --- a/litellm/_logging.py +++ b/litellm/_logging.py @@ -90,6 +90,30 @@ 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") @@ -97,6 +121,7 @@ def filter(self, record: logging.LogRecord) -> bool: 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]]: @@ -197,16 +222,6 @@ def format(self, record): if "logger" not in json_record: json_record["logger"] = f"{record.filename}:{record.lineno}" - import litellm - - if litellm.request_correlation_in_logs: - session_id = session_id_var.get() - if session_id and "session_id" not in json_record: - json_record["session_id"] = session_id - trace_id = trace_id_var.get() - if trace_id and "trace_id" not in json_record: - json_record["trace_id"] = trace_id - if record.exc_info: json_record["stacktrace"] = record.exc_text or self.formatException(record.exc_info) @@ -219,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): @@ -335,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 diff --git a/tests/test_litellm/test_logging.py b/tests/test_litellm/test_logging.py index 7839fee6a35f..c003f7b298df 100644 --- a/tests/test_litellm/test_logging.py +++ b/tests/test_litellm/test_logging.py @@ -15,6 +15,7 @@ import litellm from litellm._logging import ( ALL_LOGGERS, + CorrelationContextFilter, JsonFormatter, _initialize_loggers_with_handler, _turn_on_json, @@ -339,6 +340,7 @@ 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)))