Commit graph

10 commits

Author SHA1 Message Date
Deepanshu Lulla
9ce96c2d34
feat(logging): add opt-in session_id and trace_id correlation to JSON log records via contextvars (#34418)
* feat(logging): add opt-in session_id/trace_id correlation to JSON log records via contextvars

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 <noreply@anthropic.com>

* fix(logging): guard session_id/trace_id injection against overwriting caller-supplied extra fields

* fix(logging): always reset session_id_var to empty string when no session_id provided

* feat: gate request correlation IDs in logs behind request_correlation_in_logs flag

* refactor: move correlation ID injection into CorrelationContextFilter

* feat(logging): extend request_correlation_in_logs to plaintext logs and StandardLoggingPayload

Plaintext log lines (json_logs off) now get the same trace_id/session_id
suffix as JSON logs via a new CorrelationPlainFormatter, so the flag has a
visible effect regardless of log format.

StandardLoggingPayload gets a new independent session_id field, populated
from litellm_session_id. trace_id's existing session_id-first fallback is
preserved when request_correlation_in_logs is off; with the flag on, an
explicit litellm_trace_id now takes priority over litellm_session_id so
the two fields carry genuinely independent values.

* fix(logging): restore correlation context after nested calls; sanitize correlation ids

Addresses two review findings on this PR.

CorrelationContextFilter's trace_id/session_id contextvars were set on every
Logging.__init__ but never reset, so a nested LiteLLM call sharing the same
asyncio Task as an outer request (e.g. a guardrail's own LLM-as-judge call,
an MCP sampling call) would leave the outer request's subsequent log lines
stamped with the nested call's ids instead of its own. set_trace_id/
set_session_id now return their contextvars.Token, and Logging stores them
and resets both once its own success/failure handler actually completes,
via a new idempotent _restore_correlation_context() called from all four
terminal handlers.

set_trace_id/set_session_id also now strip control characters and bound
length before storing a caller-controlled trace_id/session_id, since these
values can originate from request input (litellm_session_id, x-litellm-
trace-id) and get interpolated into plain-text log lines - without this, a
caller could embed \r/\n or escape sequences to forge fake log entries.

* fix(logging): restore correlation context after nested calls, not before

The previous commit called _restore_correlation_context() as the first
line of each terminal handler, before that handler's own callback
dispatch loop runs. That's backwards: a nested LiteLLM call triggered
from within a callback (e.g. a guardrail's own LLM-as-judge call) would
then capture the *already-reset* value as its own pre-call baseline,
and its own reset would restore to that instead of the true outer
value - verified live to still leak.

success_handler/async_success_handler/failure_handler/async_failure_handler
are now thin wrappers: the original bodies move to
_success_handler_body/etc, called inside a try/finally that restores
context only once the full body - including any nested calls its own
callback dispatch triggers - has actually finished, mirroring proper
stack-scoped nesting semantics.

* test(logging): cover async_failure_handler's correlation-context restore

Codecov flagged the new async_failure_handler wrapper (try/finally around
_async_failure_handler_body) as uncovered - the method had no direct test
at all before this PR's refactor split it into a wrapper. Adds a test that
awaits it directly and asserts both that async_log_failure_event still
fires and that _restore_correlation_context() puts the pre-call
trace_id/session_id back.

* fix(logging): restore correlation context by value, not by contextvars.Token

veria-ai correctly flagged that contextvars.Token.reset() only works in the
exact Context it was created in, and litellm's async success path (and
streaming failure path) dispatch async_success_handler/async_failure_handler
via asyncio.create_task and the global logging worker - a different Context
than Logging.__init__ ran in. reset_trace_id/reset_session_id silently
swallowed the resulting ValueError, so the restore was a no-op for exactly
those paths. Verified independently: reproduced the raw contextvars
behavior, then confirmed litellm's async success dispatch really does go
through asyncio.create_task + GLOBAL_LOGGING_WORKER (litellm/utils.py).

Logging now captures the pre-call *value* (not a Token) and restores via a
plain set_trace_id()/set_session_id() call, which works regardless of which
Task/Context calls it. reset_trace_id/reset_session_id are removed as
dead/unreliable code. Added a regression test that spawns __init__ and the
restore in different asyncio Tasks - confirmed it fails against the prior
Token-based commit and passes here.

* fix(logging): restore correlation context in the originating task too

Greptile's re-review correctly identified a remaining gap: for a
successful acompletion(), async_success_handler is dispatched via
asyncio.create_task + the global logging worker into a *different* Task
than the one wrapper_async/Logging.__init__ ran in. The prior fix (43c164a)
only restored the handler's own (detached, throwaway) Task - it never
touched the originating request Task, which keeps this call's trace_id/
session_id set for the rest of its own execution (e.g. nested calls made
via the same Task).

wrapper()/wrapper_async() in litellm/utils.py now restore the originating
Task's correlation context in a finally block once the whole call is done,
regardless of what detached logging tasks it spawned along the way. Since
the wrapped body rebinds its own `kwargs` local via function_setup(),
sharing the dict object doesn't work here; a small mutable holder carries
the constructed Logging instance back out to the outer wrapper instead.

_restore_correlation_context() is no longer guarded against repeat calls:
with value-based (not Token-based) restoration, each distinct Task that
calls it needs its own restore to take effect in that Task's own view of
the contextvars, so multiple calls (once per Task involved in an attempt)
are required, not just tolerated.

Added a regression test using mock_response to exercise the real success
dispatch path (asyncio.create_task + GLOBAL_LOGGING_WORKER) without a live
provider call, asserting the *test's own* (originating) task context is
restored after the call - this is exactly the case Greptile flagged and
the prior commit didn't cover.

* fix(logging): restore correlation context when function_setup itself fails

Greptile's 4th finding: if function_setup() constructs Logging() (whose
__init__ already mutates trace_id_var/session_id_var) and then raises
before returning - e.g. update_environment_variables() throws - the
caller's wrapper()/wrapper_async() never receives a logging_obj reference,
so its own restore-on-finally never fires. The correlation ids leak into
every subsequent log line on that thread/task with no way to clear them.

function_setup()'s own except block now restores the context itself in
that case, using whatever logging_obj it managed to construct before
failing (locals().get(), safe against the earlier failure modes where
logging_obj was never assigned at all).

Added a regression test that monkeypatches Logging.update_environment_variables
to raise after construction, confirmed it fails without this fix (the
leaked ids show up directly in the raised exception's own log line) and
passes with it. Broader sweep (test_utils.py, test_router.py,
test_main_module_header.py, streaming handler tests, plus all
logging-specific tests): 722 passed.

* fix(logging): don't assume every litellm_logging_obj is a real Logging instance

CI caught a real regression from the last commit: tests/test_litellm/llms/xai/test_xai_key_fallback.py
injects a minimal FakeLogging stand-in (only implementing
update_from_kwargs) as litellm_logging_obj for a narrow realtime-config
unit test, bypassing the real Logging class entirely. wrapper()/
wrapper_async()'s finally block and function_setup()'s except block both
unconditionally called _restore_correlation_context() on whatever ended up
in the holder, which doesn't exist on that stand-in.

_restore_correlation_context is new plumbing specific to this PR's
feature, not part of any pre-existing stand-in's expected interface, so
callers of it can't assume every object playing the litellm_logging_obj
role implements it. Added _restore_correlation_context_if_supported(),
a small getattr-guarded helper, and used it at all three call sites.

* fix(logging): don't restore context too early on setup failure or streaming

Two more findings from Greptile's 5th review round.

1. function_setup()'s except block restored correlation context *after*
   logging the "Error in function_setup" exception, so that diagnostic log
   line itself was stamped with the doomed call's ids instead of the outer
   ids - misleading, since the failed call never produces anything else to
   attribute those ids to. Restore now happens before the log call.

2. wrapper()/wrapper_async() restored the originating task's context as
   soon as a streaming call returned, before the caller ever starts
   iterating the CustomStreamWrapper it just got back. Any log lines
   emitted while iterating (in the same thread/task) incorrectly showed
   the pre-call ids instead of this call's own ones. The wrapper finally
   block now skips the restore when the return value is a stream wrapper,
   deferring to the terminal handler that already fires once the stream is
   actually assembled/exhausted.

Both verified with tests that fail against the prior commit and pass
against this one. Broader sweep unchanged at 829 passing.

* fix(logging): best-effort correlation cleanup on abandoned streams

Greptile's 7th finding: if a caller returns a streaming response and never
fully consumes it - stops iterating early, drops the reference, cancels
it - the terminal handler that normally restores the originating task's
trace_id/session_id never fires, since it only runs once the stream is
actually assembled/exhausted. The ids leak into every subsequent log line
in that thread/task with no bound.

There's no reliable Python hook for "this was abandoned without being
closed" - CustomStreamWrapper has no close()/__aexit__/context-manager
convention today, and the only automatic option is __del__, whose timing
is inherently unpredictable (delayed by cyclic GC, not guaranteed at
interpreter shutdown, can run on a different thread). This is a best-effort
safety net, not a guarantee, and is documented as such in the docstring.

Testing this via real garbage collection proved unreliable in practice:
per-chunk logging submits work to a thread pool executor whose worker
thread transiently holds its own bound-method reference to the wrapper
until that task completes, so refcount doesn't hit zero on a
deterministic schedule even with polling. Tests call __del__ directly
instead - a plain method, safe to invoke early - which exercises exactly
the restore logic real garbage collection would eventually trigger,
plus a case confirming a broken logging_obj can never make __del__ raise.

* fix(logging): restore consumer's context at every real stream exit point

Two more findings from this round.

Veria AI: even a *fully consumed* stream never restored the actual
consuming thread/task's correlation context. The terminal success dispatch
(dispatch_success_handlers via asyncio.create_task for async, or
success_handler via the shared executor for sync) only restores whatever
detached context it runs in - never the caller's own thread/task that's
running the for/async for loop. Same root cause as the wrapper-level fix
two rounds ago, just missed for the streaming-completion path.

Greptile: explicit aclose() (client disconnect, router fallback aborting
a partial stream) closed the underlying stream without restoring
correlation context either, since request wrappers intentionally skip
restoration for returned streams and no terminal handler runs on this
path.

Added CustomStreamWrapper._restore_consumer_correlation_context(), called
from every point control genuinely returns to the consumer: the final
raise StopIteration/StopAsyncIteration on natural exhaustion (both sync
branches, both async branches), _handle_stream_fallback_error (the shared
choke point for all three failure-raising call sites), and aclose(). __del__
now delegates to the same helper instead of duplicating it.

Verified with tests extending the existing streaming-exhaustion cases to
assert the consuming context is restored after the loop completes (fails
against the prior commit, passes now), plus a dedicated aclose() test.
Broader sweep: 832 passing.

* fix(logging): don't let a delayed __del__ finalizer clobber a newer active call

If an abandoned stream's __del__ fires late (after cyclic GC delay), a
different call may have already taken over the correlation contextvars in
the same Task/thread. Restoring unconditionally would stomp that active
call's trace_id/session_id with the abandoned stream's stale pre-call
snapshot. __del__ now only restores when the contextvars still hold the
ids this call itself set.

* fix(logging): compare sanitized ids in the __del__ ownership guard

set_trace_id()/set_session_id() sanitize (strip control chars, bound length)
before storing, so the contextvar's value can differ from the raw
litellm_trace_id/litellm_session_id. The __del__ ownership guard was
comparing against the raw values, so a caller-supplied id containing control
characters or exceeding 256 chars would never match, permanently skipping
cleanup. Capture what set_trace_id()/set_session_id() actually stored and
compare against that instead.

* fix(logging): restore consumer context on the synthesized finish_reason chunk

Both __next__ and _finalize_completed_stream() have a branch that fires when
the underlying stream ends without ever emitting an explicit finish_reason
chunk: they synthesize one via finish_reason_handler() and return it. A
consumer that stops as soon as it sees finish_reason - a common pattern -
never calls __next__()/__anext__() again, so the existing restore in the
sent_last_chunk-is-True StopIteration branch never runs for them. The
underlying stream is already exhausted at this point regardless of whether
the caller keeps iterating, so restoring here is safe.

* fix(logging): don't restore correlation context before the caller receives the final chunk

The previous fix (5147c69186) restored context immediately before returning
the synthesized finish_reason chunk from __next__/_finalize_completed_stream,
reasoning that completion_stream was already exhausted. But that chunk is
still this call's own data, and the caller's own application-level log
statements processing it run in the same synchronous frame right after the
return - restoring first made those lines carry the wrong (outer) ids,
exactly what wrapper()/wrapper_async() deliberately avoid by not restoring
while a stream is being iterated.

Revert to not restoring there. A caller that keeps iterating still gets a
correct, deterministic restore on its very next __next__()/__anext__() call
(completion_stream is exhausted, so that immediately re-raises
StopIteration/StopAsyncIteration through the already-restoring branch). A
caller that stops right after finish_reason relies on aclose() or the
best-effort __del__ guard, same as any other stream the caller doesn't fully
exhaust.

* refactor(logging): hoist a safely-hoistable function-body import to module top

CorrelationContextFilter.filter()'s `import litellm` was a function-body
import; verified it can move to module top without a circular-import failure
(litellm/__init__.py already imports from litellm._logging before setting
request_correlation_in_logs, but a bare `import litellm` only binds the
already-in-sys.modules module object - the attribute itself isn't read until
filter() actually runs, by which point litellm is fully initialized).

* test(logging): move correlation tests into their conventionally-mapped files

tests/test_litellm/ mirrors litellm/ in a parallel path. Correlation tests
for the Logging class (litellm_logging.py), function_setup/wrapper_async
(utils.py), and CustomStreamWrapper (streaming_handler.py) had all landed in
test_logging.py, which only maps to litellm/_logging.py itself. Moving each
group to its correctly-mapped file: test_litellm_logging.py (Logging class
init/restore), test_utils.py (function_setup, wrapper_async), and
test_streaming_handler.py (CustomStreamWrapper) in the next commit.
test_logging.py keeps only what actually exercises _logging.py's own
contextvars/filters/formatters/sanitization. No behavior change - same
assertions, same coverage, just relocated.

* fix(logging): restore correlation context unconditionally in wrapper()'s sync path

Blocking finding from review: a caller-visible correlation feature was
silently misattributing one request's logs to a different, unrelated one on
the sync/threaded path. wrapper()/wrapper_async() both left trace_id/session_id
"open" across a stream's entire iteration so the caller's own log lines while
consuming it would carry the right ids. That's safe for wrapper_async(): each
async call gets its own asyncio Task with its own copy of the contextvars,
and Tasks are never recycled across requests, so a leftover value can only
ever affect that one already-abandoned Task.

It is not safe for wrapper() (sync): a plain OS thread has no such per-call
isolation, and a thread pool's worker threads *are* recycled across unrelated
requests. If a sync stream was abandoned (client disconnect, early break, an
uncaught exception) without ever being exhausted or closed, nothing restored
its contextvars, and a pool could later hand that same thread to a completely
different call, which would inherit the abandoned request's ids as its own
"pre-call" baseline and then restore back to that poison when it finished -
permanently misattributing every subsequent log line on that thread,
including its own, to the abandoned request. Strengthening the __del__
finalizer already added for this can't fix it: finalizer timing is exactly
what a permanently-reused thread can't rely on.

wrapper() now restores unconditionally in its own finally, before a sync
stream is ever handed back to the caller. The trade-off: a sync stream
consumer's own application-level log statements while iterating no longer
automatically carry this call's ids (litellm's own internal per-chunk
logging is unaffected, since it's dispatched separately). That's an
acceptable cost for eliminating a silent cross-request misattribution bug.
wrapper_async() keeps the existing conditional (skip-if-streaming) behavior,
justified by the Task-isolation argument above; CustomStreamWrapper's
__del__/aclose()/next-iteration restore machinery remains meaningful and
necessary there.

This also simplifies wrapper()/wrapper_async() back toward their original
shape: both previously used a mutable-dict-holder split into a separate
_body function to smuggle logging_obj/result out to an outer finally,
working around function_setup() rebinding its own local `kwargs`. That
restructuring is no longer needed - `logging_obj` (and, for wrapper_async(),
`result`) were already function-level locals in scope for a plain
try/finally; three of wrapper_async()'s retry-return statements now assign
through `result` first so it accurately reflects what's actually returned
even on a retry path.

Regression test: test_abandoned_sync_stream_does_not_contaminate_a_later_call_on_the_same_thread
in test_streaming_handler.py reproduces the exact reported scenario with a
real single-worker ThreadPoolExecutor - confirmed it fails with the prior
(skip-restore-on-stream) wrapper() and passes with this fix.

* refactor(logging): use Mapping instead of bare dict for read-only params

_get_standard_logging_payload_trace_id/_session_id only read litellm_params
(.get() calls, no mutation) - annotate it as Mapping[str, Any] rather than a
bare mutable dict, per the repo's no-mutable-collection-in-annotation rule.

* fix(logging): scope request_correlation_in_logs to the async/proxy path only

Blocking review finding: wrapper() (the sync entry point) used the same
skip-restore-on-stream design as wrapper_async(), but a plain OS thread has
no per-call context isolation the way an asyncio Task does, and a thread
pool's worker threads are recycled across unrelated requests - an abandoned
sync stream could leave its ids stuck on a thread a pool later hands to a
completely different request, misattributing that request's logs. A fix
existed and was tested (restore unconditionally in wrapper()'s own finally),
but it doesn't benefit this feature's primary consumer - the proxy only ever
calls the async entry point - and carries sync-specific complexity this PR
doesn't need.

Scope the feature to async only instead: Logging.__init__() takes a new
supports_correlation_logging parameter (default True), threaded down from a
new function_setup(..., is_async_call: bool = True) parameter. wrapper() is
the one caller that passes is_async_call=False; every other function_setup()
call site (wrapper_async(), the router, and proxy/MCP-internal call sites)
is already async and keeps the default. With
supports_correlation_logging=False, Logging.__init__() never calls
set_trace_id()/set_session_id() at all, so a sync call has nothing to leak
in the first place. wrapper() reverts to its pre-review shape with no
correlation-specific code at all.

StandardLoggingPayload's own trace_id/session_id fields are unaffected
either way - they're a deterministic per-call read of
self.litellm_trace_id/self.litellm_session_id, not ambient contextvar state,
so they were never exposed to the cross-request bug.

Full sync/direct-SDK support (stamping + its own safe-restore mechanism) is
deferred to a follow-up PR; the fix and its regression test already exist in
this branch's history at commit 9f3a20f4b2 and can be resurrected there.

Tests: replaced the two wrapper()-level tests with ones proving the new
invariant (sync calls, streaming and non-streaming, never touch
trace_id_var/session_id_var even when the caller explicitly passes
litellm_trace_id/litellm_session_id), and added a direct unit test for the
supports_correlation_logging=False gate on Logging.__init__ itself. Verified
live: a real proxy (Postgres-backed, real OpenAI calls) shows clean
trace_id/session_id isolation across two concurrent sessions with no
cross-contamination; a standalone script confirms real sync SDK calls
against a real model never touch the correlation contextvars.

* feat(logging): fall back to W3C traceparent/baggage for trace_id/session_id

request_correlation_in_logs previously only resolved trace_id/session_id from
litellm-specific sources: x-litellm-trace-id/x-litellm-session-id headers, a
generic x-<vendor>-session-id header, or Anthropic-style metadata.user_id. If
none were present, trace_id fell back to an auto-generated UUID unrelated to
anything else, and session_id stayed empty - even when the caller already had
real distributed-tracing instrumentation sending the actual industry-standard
headers for this.

Add a fallback to the W3C Trace Context traceparent header (trace-id
component) and W3C Baggage header (session.id entry), so a request already
carrying real OpenTelemetry trace context correlates litellm's own logs with
the same trace in the caller's observability backend (Datadog, Honeycomb,
Tempo, etc.) instead of getting an unrelated generated id. Precedence is
unchanged for existing sources: explicit litellm headers and the Anthropic
metadata path both still win over this new fallback, which only fires when
neither found anything. trace_id and session_id are resolved independently
here (unlike the existing chain_id mechanism, which uses one shared value for
both), since traceparent and baggage are semantically distinct W3C concepts.

New helpers _trace_id_from_traceparent/_session_id_from_baggage in
litellm_pre_call_utils.py parse the header formats directly (no new
dependency - both are simple fixed-width/delimited strings), wired into
LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers() only
when the corresponding litellm_trace_id/litellm_session_id key isn't already
set by the existing paths.

Verified live against a real proxy: a bare traceparent header produces a log
trace_id exactly matching its trace-id component; a traceparent alongside an
explicit x-litellm-trace-id header (different value) produces a log showing
the explicit header's value, proving precedence.

* fix(logging): reserve trace_id/session_id in JsonFormatter against message-content spoofing

JsonFormatter merges keys parsed from the message body before applying extra
record attributes, and the extra-attributes loop skips a key that's already
present. A caller-controlled log message that happens to parse as JSON/dict
with a "trace_id"/"session_id" key (e.g. the proxy logging a raw request-header
dict) could therefore make the JSON record carry the attacker-supplied value
instead of the real correlation context set via CorrelationContextFilter.

trace_id/session_id are now applied from the LogRecord's own attributes after
message-content parsing, unconditionally overwriting anything the message body
claimed for those two keys.

* style(logging): fix import order (ruff I001) in _logging.py and litellm_logging.py

- _logging.py: import litellm belongs after the stdlib from-imports, grouped
  with the other litellm.* imports, not before them.
- litellm_logging.py: the refactor to Mapping introduced a second, separate
  `from collections.abc import Mapping` instead of merging it into the
  existing `from collections.abc import Callable` import.

Caught by the strict-rule budget gate (ruff-strict-budget.json caps I001 at
0 new violations); both auto-fixed with `ruff check --fix --select I001`.

* style(logging): freeze mutable-collection constructions flagged by LIT002

Five sites in this PR's diff built a mutable list/dict literal instead of a
frozen value: a plain list of optional strings in CorrelationPlainFormatter,
a `kwargs or {}` fallback, a `metadata or {}` fallback, two `[...]` candidate
orderings, and a `dict(headers)` copy feeding a dict comprehension. Each is
build-once/read-only, so this rewrites them as tuples, MappingProxyType, or a
plain conditional `.get()` instead of seeding then reading a fresh mutable
collection - no behavior change, confirmed by the existing test suite.

Caught by the type-discipline budget gate (LIT002 capped at 0 new
violations).

* fix(logging): reserve trace_id/session_id even when no correlation context is active

Live-proxy verification surfaced a gap in the earlier message-content-spoofing
fix (7f390a57fc): that fix only overwrites trace_id/session_id from the
LogRecord's own attribute, so it does nothing for a log line emitted before
CorrelationContextFilter has stamped anything on this record (e.g. the
"Request Headers" debug line, which fires before Logging.__init__() runs for
the request). On such a record, a caller-supplied header literally named
trace_id/session_id still got promoted into the JSON output via the embedded
JSON/dict-repr parser, since there was no genuine value to protect.

Fixed at the source: trace_id/session_id are now excluded unconditionally from
the message-content-parsing promotion step, not just superseded afterward.
Verified live against a real proxy - the exact adversarial request (headers
literally named trace_id/session_id) no longer leaks into any JSON log record.
Added a regression test for this no-active-context variant specifically,
confirmed it fails against the prior commit and passes now.

Also fixes an unrelated basedpyright regression from an earlier rebase's
conflict resolution: litellm/utils.py's `logging_obj` was incorrectly
re-annotated `Final` at its second assignment in function_setup() (it's first
declared `None` a few lines earlier), which basedpyright correctly rejects.

* fix(proxy): stop logging the raw W3C baggage session_id value

_session_id_from_baggage() extracts the caller-controlled session.id entry
verbatim - it isn't sanitized until set_session_id() runs later in
Logging.__init__(). The debug log line for this extraction interpolated the
raw value directly, so a caller could embed terminal control characters or
ANSI escape sequences that forge/alter plaintext log output for anyone
tailing the proxy's logs.

Verified live: a baggage header with an embedded ANSI escape reached the
terminal as a real, unescaped control sequence before this fix. Drops the
value from the log line entirely (the extraction succeeding is enough signal
on its own) rather than sanitizing-then-logging, matching veria-ai's
suggestion. Added a regression test using caplog that fails against the prior
commit and passes now.

* fix(logging): restore consumer context only after stream-failure exception mapping

_map_anthropic_exception/_map_aleph_alpha_exception synchronously log a debug
diagnostic (the raw status code) as part of exception_type()'s mapping.
_handle_stream_fallback_error restored the consumer's outer correlation
context before calling exception_type(), so that diagnostic log line carried
the outer (or empty) trace_id/session_id instead of the failing stream's own -
flagged by Greptile.

Moved the restore to run after mapping completes, matching the same
restore-after-not-before pattern already applied elsewhere in this file for
success/finish_reason handling. Added a regression test that captures the
correlation context live during a mocked exception_type() call; fails against
the prior commit, passes now.

* fix(logging): restore consumer context only after aclose()'s stream close completes

aclose() restored the consumer's outer correlation context as its first
statement, before awaiting the underlying provider stream's own aclose()/
close(). If that close attempt raises, the except branch's debug diagnostic
ran under the already-restored outer context instead of the closing stream's
own trace_id/session_id - flagged by Greptile, same restore-too-early pattern
as the stream-failure fix in f1cf9589d6.

Moved the restore to the end of aclose(), after the close attempt (and its
diagnostic logging) completes. Added a regression test with a fake stream
whose aclose() raises, capturing the correlation context live during the
diagnostic log call; fails against the prior commit, passes now.

* style(logging): satisfy new strict-lint budgets introduced upstream (Final, ANN401, S110, TRY300, kwargs typing)

Rebasing onto litellm_internal_staging pulled in 116 upstream commits that
introduced/tightened several lint gates this PR's own code now trips:

- LIT010 (every local/module-level variable must be Final): added Final
  annotations across _logging.py, litellm_logging.py, streaming_handler.py,
  litellm_pre_call_utils.py, and utils.py. Where a name is genuinely
  reassigned (logging_obj: starts None, later set to the real object) or
  branch-assigned, either restructured into a single ternary expression
  (ordered_candidates) or suppressed with `# rebind-ok: <reason>` matching
  this repo's documented escape hatch.
- LIT011 (parameter mutation): suppressed the two new `data[key] = value`
  writes in litellm_pre_call_utils.py with `# rebind-ok`, matching the
  unsuppressed precedent already used for every other `data[...]` write in
  the same function - `data` is an intentional out-param there.
- ANN001/ANN003/ANN202 (missing parameter/return type annotations): fully
  typed success_handler/_success_handler_body, their async twins, and
  failure_handler/_failure_handler_body/async variants in litellm_logging.py,
  plus function_setup in utils.py (added Rules to its existing TYPE_CHECKING
  block for the rules_obj: Rules annotation).
- ANN401 (explicit Any disallowed): suppressed with `# noqa: ANN401` on the
  handful of genuinely-heterogeneous result/*args/**kwargs parameters, since
  ordinary suppression is this repo's documented path.
- S110 (try/except/pass): added to the existing BLE001 noqa on the one
  best-effort correlation-cleanup try/except this PR added.
- TRY300 (return inside try): moved two `return result` statements into
  `else:` blocks in the retry-fallback paths this PR's own diff touched.
- reportPrivateUsage (basedpyright): renamed the two new
  StandardLoggingPayloadSetup static methods (get_standard_logging_payload_
  trace_id/session_id) to drop their leading underscore, since they're
  genuinely called from a sibling module-level function in the same file.

No behavior change - confirmed by the full existing test suite (819 passed)
plus all four lint gates (ruff format, ruff-strict, type-discipline,
basedpyright) passing clean.

* fix(lint): stop RUF100 flagging noqa suppressions the strict gate needs

CI's plain "ruff check" job uses the default ruff.toml, a narrower config
than ruff-strict.toml (used only by the strict-rule budget gate). ANN401 and
S110 aren't enabled in the default config, so RUF100 (unused-noqa) flagged
the `# noqa: ANN401`/`# noqa: ...,S110` suppressions this PR added as pointless
under that config, even though they're genuinely needed under ruff-strict.toml.

- ANN401: added to ruff.toml's existing `lint.external` list (same mechanism
  already used for C901/TID251, enforced by the strict gate but not by this
  config) - these Any usages are genuinely dynamic/forwarded, so the
  suppression itself is correct and just needed registering.
- S110: fixed the underlying code instead of registering another external
  code - the try/except/pass in
  CustomStreamWrapper._restore_consumer_correlation_context now logs at
  debug level on failure (matching the existing best-effort-cleanup pattern
  in _record_partial_usage_for_failure elsewhere in this file), which
  satisfies S110's own suggestion directly and needs no suppression at all.

Verified against both ruff.toml and ruff-strict.toml directly, plus all
three other gates (ruff format, type-discipline, basedpyright) and the full
test suite (821 passed).

* fix(lint): scope the ANN401 exemption to file level instead of a repo-wide noqa

Ruff has no per-line-scoped way to register a noqa code across configs (that
requires the default ruff.toml's lint.external list, which is repo-wide in
scope even though the noqa itself is per-line). Since ruff does support
file-level exemptions via per-file-ignores, and ANN401 only needed exempting
in exactly two files, moved the exemption there instead:

- ruff-strict.toml: added [lint.per-file-ignores] disabling ANN401 for
  litellm_logging.py and utils.py specifically, with a comment explaining
  why (heterogeneous response/forwarded-args parameters with no fitting
  concrete type - already verified by trying CostResponseTypes and hitting
  a real basedpyright mismatch).
- ruff.toml: reverted the ANN401 entry from lint.external - no longer
  needed, since there's no `# noqa: ANN401` left anywhere for RUF100 to
  second-guess.
- Removed the now-redundant `# noqa: ANN401` from the 10 affected
  parameters in both files, keeping the existing kwargs-ok reasons and
  adding a short inline comment on the `result`/`*args` lines pointing at
  the ruff-strict.toml exemption for context.

Verified against both configs directly (ANN401 clean under ruff-strict.toml
for these files, RUF100 clean under the default config), all four gates
(ruff format, ruff-strict, type-discipline, basedpyright), and the full
test suite (821 passed).

* fix(logging): redact credential-shaped trace_id/session_id before stamping log records

CorrelationContextFilter stamps trace_id/session_id onto a LogRecord after
SecretRedactionFilter has already run, so a caller-controlled value (e.g. via
x-litellm-trace-id or a W3C baggage header) that happens to look like a real
credential reached JSON and plaintext logs unredacted. Apply the same
credential redaction already used elsewhere in this module at
_sanitize_correlation_id(), the single choke point both set_trace_id() and
set_session_id() route through, so every caller-facing entry point is covered
without depending on filter ordering.

* fix(logging): restore correlation context when a stream's max-duration timeout fires

CustomStreamWrapper.__anext__() called _check_max_streaming_duration() before
entering its try block, so the litellm.Timeout it raises bypassed the except
Exception -> _handle_stream_fallback_error path entirely, leaking the timed-out
stream's own trace_id/session_id into whatever the consumer's task logs next.
Move the check inside the try so it flows through the same restoration path
every other stream failure already uses.

* test(streaming): make dispatch_failure_handlers mock awaitable for the async max-duration test

Moving _check_max_streaming_duration() inside __anext__()'s try block (prior
commit) means a max-duration Timeout now dispatches failure handlers through
the same path every other stream failure already uses, instead of bypassing
it entirely. dispatch_failure_handlers is async on the real Logging class;
the test's plain MagicMock logging_obj made asyncio.create_task() choke on a
non-coroutine return value once that path actually got exercised.

---------

Co-authored-by: Deepanshu <deepanshu.lulla@alpha-sense.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-08-10 10:40:13 -07:00
mateo-berri
c3536c29a0 fix(lint): cover every base-owned ruff rule in the strict gate's external list 2026-08-07 21:38:53 -07:00
mateo-berri
557d14cc71 fix(lint): make strict-gate noqas survive base ruff and flag stale ones 2026-08-07 21:27:00 -07:00
Mateo Wang
818e319b53
chore: make it more concise 2026-08-06 03:38:28 -07:00
mateo-berri
d7ca4dc77f chore(lint): say annotate Mapping[...] instead of Mapping alias in typing.Dict ban 2026-08-06 03:34:34 -07:00
mateo-berri
0c1a4b127d chore(lint): name MappingProxyType in the mutable-collection fix messages
LIT001/LIT002 and the typing.Dict ban all steered dict-shaped values to
frozen dataclasses or suppression even though the checker already accepts
MappingProxyType as a freezing wrapper; the messages now name it so the
dict-shaped freeze path is actually discoverable at fix time.
2026-08-06 03:31:40 -07:00
Mateo Wang
92d0788da2
chore(lint): widen ANN slack to 10% of baseline and drop PLR0913 from the strict gate (#31335)
* chore(lint): widen ruff budget slack to 10% of baseline for high-volume ANN rules and PLR0913

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>

* chore(lint): drop PLR0913 from strict gate to roll out rules gradually

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>

* fix(lint): ratchet-guard rising baselines even when slack is cut to mask them

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>

---------

Co-authored-by: Cursor Agent <cursoragent@cursor.com>
Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>
2026-06-25 14:43:45 -07:00
Mateo Wang
be4fa702e7
ci(lint): ratcheted type-discipline gate (mutable collections, casts, guards, kwargs, suppressions) (#30500)
* ci(lint): enforce type-discipline budget for casts and type guards

Add a ratcheted gate that blocks net-new typing.cast() usage and bans
TypeGuard/TypeIs outright, layered on the existing ruff-strict budget setup.

- ruff-strict.toml: ban cast/TypeGuard/TypeIs (typing + typing_extensions)
  via flake8-tidy-imports banned-api (TID251) for a coarse import-level freeze.
- ruff-strict-budget.json: bump TID251 baseline 2404 -> 2662 to absorb the
  ~258 pre-existing usages now matched by the new banned-api entries.
- scripts/check_type_discipline.py: AST checker adding LIT006 (cast call sites,
  suppress with `# cast-ok: <reason>`) and LIT007 (TypeGuard/TypeIs annotations,
  suppress with `# guard-ok: <reason>`) for per-call-site granularity.
- scripts/type_discipline_gate.py: baseline+slack gate with delta-vs-base,
  mirroring ruff_strict_gate.py.
- type-discipline-budget.json: LIT006 baseline 1013 (slack 10), LIT007 0/0.
- test-linting.yml: run the gate in CI against the PR base SHA.

* ci(lint): enforce suppression-reason budgets and guard budgets against loosening

- wire the **kwargs ban (LIT008) into the vendored type-discipline checker so it
  matches the budget that already referenced it
- freeze LIT003/LIT004 (noqa / type-ignore without codes or reason) and LIT005
  (*-ok suppression without a reason) at slack 0 so any net-new unexplained
  suppression trips the type-discipline gate
- add scripts/budget_ratchet_check.py and a separate, non-gating budget-ratchet CI
  job that turns red when any *-budget.json ceiling is raised, a rule is dropped,
  or a budget file is deleted

* ci(lint): ban mutable collections in annotations and all mutable construction

Expand LIT001 from coarse builtins at interfaces to any mutable collection
in any annotation (builtins, typing aliases, collections concretes, mutable
ABCs) across signatures, class attributes, locals, and globals. Add LIT009 to
flag mutable-collection construction (literals, comprehensions, constructors)
so the unannotated seed-then-mutate pattern is caught too. Enumerate any-ok in
LIT005 so its reason requirement holds even when only the stdlib checker runs.
Budget LIT001 (21452) and LIT009 (25222) with slack 10 to ratchet down.

* ci(lint): recommend pydantic at boundaries and add functional-refactor guidance

Drop the msgspec mention from the cast banned-api messages so the recommended
validation path matches the codebase's primary pattern (pydantic). Add a note to
CLAUDE.md that lint / type-discipline failures should be resolved by refactoring
to functional, immutable patterns rather than reaching for mutable structures or
`# mutable-ok`.

* style: make CLAUDE.md more concise

* chore: update CLAUDE.md guidelines

* ci(lint): renumber mutable construction LIT009 -> LIT002 next to LIT001

Group the mutable-collection family together: LIT001 (mutable collection in any
annotation) and the construction rule now sit adjacent at LIT001/LIT002. The
freed LIT009 slot is taken by the sibling Any gate (check_any_discipline.py,
#30379), which moves its Any-typed-value rule LIT002 -> LIT009 in lockstep so
the shared LIT namespace stays contiguous with no holes. Budget, gate docstring,
and the checker's own docstring/messages are updated to match.

* fix: numbering in CLAUDE.md

* test(lint): test type-discipline checker, scope LIT007 to return types

Add regression tests for check_type_discipline.py (every LIT rule, its
suppression, and the comment scanner) and for budget_ratchet_check.py.

Confine LIT007 to function return annotations, the only place TypeGuard/TypeIs
are valid, so a runtime name that merely reads those identifiers is no longer
flagged. Switch scan_comments to io.StringIO(source).readline, the standard
readline that returns '' at EOF, dropping the iter(...).__next__ idiom.

* fix(lint): best-effort worktree teardown so cleanup can't mask the real error

base_counts ran `git worktree remove` through the raising `_run` in its finally,
so a failed `git worktree add` (or a failure in the body) was masked by a second
SystemExit from the cleanup. Tear the worktree down best-effort, like the sibling
rmtree, so the original error propagates.

* fix(lint): ratchet fails loudly on an unresolvable base; drop dead checker state

Verify the merge-base ref resolves to a commit before trusting a missing-file
result from git show, so an invalid or empty BASE_SHA now turns the budget-ratchet
guard red instead of skipping every budget and passing vacuously

Also drop the unused Comments.by_line field and the phantom --changed-only usage
line from check_type_discipline's docstring, and cover the ref handling with tests

* fix(lint): degrade malformed source to LIT000 instead of crashing the checker

tokenize.generate_tokens raises IndentationError (a SyntaxError subclass) on a dedent
mismatch, which escaped scan_comments' tokenize.TokenError handler and crashed the whole
checker run, zeroing the gate for that invocation. Catch SyntaxError too so the file
falls through to ast.parse and is reported as LIT000, matching the checker's
graceful-degradation contract. Also add the trailing newline ruff-strict.toml lacked

* perf(lint): skip the base worktree scan when no rule is over its ceiling

cmd_check created a git worktree and re-scanned the base tree on every run, but a
rule can only breach when its head count is already over baseline + slack; when none
are, the base comparison cannot change the verdict. Short-circuit to OK in that case,
which is every green PR, roughly halving the gate's work. Extract over_ceiling and
cover it (and evaluate's drift-safety) with tests

* fix(lint): exempt .dict()/.list()/.set() method calls from LIT002

_construction_kind matched dict/list/set as constructors via func.attr too, flagging
common method calls like pydantic's model.dict() as mutable construction; 200 such
false positives existed in litellm. Recognize dict/list/set construction only when
unqualified while keeping the collections concretes (deque/defaultdict/...) matchable
as attributes, since those are rarely method names. Ratchet the LIT002 baseline down
25222 -> 25022 to reflect the removed false positives

* chore(lint): bump basedpyright ceilings to absorb staging base drift

The basedpyright gate added in #30379 is a total-count check against
basedpyright-code-budget.json and the linting workflow runs only on
pull_request, so pushes to litellm_internal_staging never re-baseline it.
Merging staging into this branch surfaced that drift: seven
reportAny/reportUnknown* rules sit 10-149 errors above their committed ceiling
even though this PR changes no files under litellm/, the only path basedpyright
scans (pyrightconfig include is litellm). The new baselines match the counts CI
measured on the merge commit, with the existing per-rule slack preserved

* fix(lint): ratchet guard watches every budget file, not just two

DEFAULT_BUDGETS only listed ruff-strict-budget.json and
type-discipline-budget.json, so mypy-code-budget.json and
basedpyright-code-budget.json were unguarded and their ceilings could rise with
no signal, which is exactly the failure mode this guard exists to prevent. The
gap became concrete when this PR bumped basedpyright-code-budget.json to absorb
staging drift. All four budgets are now watched, so the budget-ratchet job
surfaces that basedpyright bump for human review the same way it surfaces the
TID251 raise. A regression test pins that every *-budget.json on disk is in
DEFAULT_BUDGETS, failing loudly if a future budget escapes the ratchet

* fix: add a lot more slack

* fix(lint): restore LIT003 frozen slack to 0

The blanket slack bump set LIT003 (bare # noqa without codes or a reason) to a
slack of 50, which contradicts the documented zero-tolerance invariant: the gate
docstring and the PR description table both freeze LIT003/LIT004/LIT005 at slack
0 so any net-new unexplained suppression trips the gate. Slack 50 would let 50
new bare noqas through silently. The actual LIT003 count is 397, well under the
516 baseline, so restoring slack to 0 keeps the gate green while putting the
freeze back. LIT004/LIT005/LIT007 were already correct at 0

* fix(lint): restore documented slack 10 for the buffered LIT rules

The slack bump left LIT001/LIT002/LIT006/LIT008 at 2000/2500/100/100, 10-250x the
"/ 10" the PR description table and the gate docstring document. That buffer was
never needed: the gate already blames a rule only when its count exceeds the
ceiling and grew vs the merge-base, so the violations the staging merge added in
litellm/ sit in both head and base and are never charged to this PR. With slack
back at the documented 10 the gate stays green, and the ceiling is tight again
(LIT006 no longer waves through 99 net-new cast() calls). Baselines are
unchanged; only the slack returns to its documented value

* fix(lint): ratchet LIT003 baseline down to its actual count

The LIT003 baseline was 516 while the current bare-noqa count is 397, leaving
~119 units of headroom that undercut the documented zero-tolerance freeze: the
gate docstring claims any net-new bare noqa trips the gate, but with cap 516 a PR
could add over a hundred first. Drop the baseline to the measured 397 so the
freeze is exact (cap = 397 + slack 0), the same hard-zero-at-the-boundary shape
LIT005 and LIT007 already use and pass in CI. PR table row updated to 397 / 0

* fix: increase slack

* fix: increase slack

* docs(lint): align gate docstring with buffered LIT003/LIT004 slack

The budget now gives LIT003/LIT004 nonzero slack, so the gate's prose no
longer claims they are frozen at slack 0; LIT005 remains the reasonless-
suppression freeze and LIT007 the hard zero.
2026-06-16 16:59:21 -07:00
Mateo Wang
d0c2e87810
ci: ratchet lint and type-check gates (ruff preview, ANN, mypy, basedpyright) (#30379)
* ci: enable ruff preview rules under the budgeted strict gate

Turn on ruff preview in the strict-budget lane (ruff-strict.toml) only,
leaving the clean gate (ruff.toml) untouched so make lint-ruff stays at
zero. Enumerate the 118 firing codes explicitly with
explicit-preview-rules so the gate is deterministic and stable across
ruff upgrades rather than depending on preview auto-selecting the broad
catalog.

Grandfather the existing 58438 violations into ruff-strict-budget.json
as per-rule baselines with headroom, so only net-new violations fail CI.
The existing ten rules keep their hand-tuned slack; the new rules get
slack 10 when the baseline is 50 or more and 3 otherwise.

* ci: add ANN return-type rules to the budgeted strict gate

Add ANN201/202/204/205/206 (missing return annotations) to the strict
lane and grandfather the existing counts into ruff-strict-budget.json so
the codebase ratchets toward explicit return types without breaking CI.

* ci: add mypy (disallow_untyped_defs) and basedpyright strict gates with baselines

Add two type-check gates, each grandfathering the current tree so only
net-new violations fail CI, matching the ruff strict-budget ratchet.

mypy gains disallow_untyped_defs in litellm/mypy.ini (the config the CI
invocation actually reads; the root [tool.mypy] is not picked up from the
litellm/ working dir). The 4885 existing missing-annotation errors are
captured in litellm/.mypy-baseline.txt and the run is piped through
mypy-baseline filter so new untyped defs are rejected.

basedpyright runs in strict mode over litellm/, with
enableTypeIgnoreComments disabled so it only honors '# pyright: ignore'
and never polices mypy's '# type: ignore'. The existing strict diagnostics
are grandfathered into .basedpyright/baseline.json.

Both tools are pinned in the dev group and uv.lock; the lint workflow and
Makefile run them filtered through their baselines, with
lint-mypy-baseline-update and lint-basedpyright-baseline-update to ratchet.

* ci: raise lint job timeout to 15m for the basedpyright strict pass

* ci: pin pythonVersion 3.12 and regenerate baselines against merged base

Merge litellm_internal_staging so the baselines cover code the CI merge
includes (e.g. the cisco_ai_defense guardrail), which otherwise tripped
the mypy gate with 3 ungrandfathered no-untyped-def errors. Pin
pythonVersion 3.12 in pyrightconfig so basedpyright's strict analysis is
reproducible across interpreter versions (CI runs 3.12).

* ci: regenerate basedpyright baseline against the frozen lint env

The previous baseline was generated with optional provider deps (azure,
google, anthropic, mcp, numpydoc, google-genai) installed locally, so CI's
dev-only env surfaced ~3500 reportUnknown*/reportMissingTypeStubs errors
not in the baseline. Regenerate after uv sync --frozen so the baseline
reflects the same dependency set the lint job sees.

* ci: regenerate basedpyright baseline on python 3.12 frozen env

The prior baseline still carried proxy-dev packages (e.g. prisma) that the
lint job's dev-only, python 3.12 env lacks, leaving 2 unresolved-import
errors ungrandfathered. Regenerate in a python 3.12 venv synced to the
frozen lock with default groups only, so the baseline matches exactly what
CI sees.

* ci: replace type-check baselines with per-file count budgets

The mypy and basedpyright baselines were position-sensitive (and the
basedpyright one was a 27MB file), so ordinary line shifts churned them.
Replace both with a per-file count gate: scripts/type_check_gate.py reduces
each tool's output to errors-per-file and checks it against a committed
{file: max} budget, ignoring line and column numbers. A file fails only
when it gains more errors than its ceiling; debt can't be shuffled between
files because each file has its own cap and new files default to zero.

Budgets (mypy-file-budget.json 48K, basedpyright-file-budget.json 96K) are
generated in the python 3.12 frozen lint env so they match CI. Drops the
mypy-baseline dependency; basedpyright runs without its native baseline.
ratchet via make lint-mypy-budget-update / lint-basedpyright-budget-update.

* ci: add a small per-file slack to the type-check gate

Allow each file to drift PER_FILE_SLACK (5) errors past its recorded count
before failing, so a basedpyright inference ripple in an unrelated file
doesn't break the build over a couple of errors. Budgets still record exact
counts; the tolerance is applied at check time.

* ci: move type-check slack into the budget json and trim lint timeout

Make slack declarative: the budget is now {"slack": N, "files": {path: count}}
so the tolerance is tuned in JSON without editing the script, mirroring how
ruff-strict-budget.json carries its slack. --update preserves the existing
slack. Also drop the lint job timeout from 15m to 10m; the mypy and
basedpyright passes add ~2m, leaving the job around 4-5m, so 10m is a
comfortable margin.

* ci: collapse fully-adopted ruff categories and drop inert preview flag

ANN (all nine non-removed rules) and BLE (its only rule) were spelled out
code-by-code; replace each with its category selector, which is exactly
equivalent in 0.15.3 (the removed ANN101/ANN102 are skipped by a category
selector and error when named explicitly). explicit-preview-rules was inert:
every selected rule is stable and nothing is selected by category, so the flag
had nothing to gate. Verified the strict-rule counts are identical before and
after (62379 each, zero per-rule drift), so no budget change.

* ci: drop redundant pyright dev dependency

Nothing invokes bare pyright in the Makefile, the linting workflow, or
scripts; the basedpyright gate added on this branch is the only type
checker that runs. basedpyright is a superset fork that reads the same
pyrightconfig.json and honors the same "# pyright: ignore" comments, so
pyright==1.1.408 in the ci group was dead weight. Regenerated uv.lock
under the same exclude-newer cutoff so the only change is removing
pyright and its package stanza

* ci: un-weaken mypy and error on Any in basedpyright

mypy: enable warn_return_any, drop the valid-type silencer, and stop globally ignoring missing first-party imports via [mypy-litellm.*] ignore_missing_imports = False, which surfaced eight real broken litellm.* imports the blanket ignore was hiding; third-party imports stay ignored. The per-file budget moves 4888 -> 5799 (902 no-any-return, 1 valid-type, 8 import-not-found), all grandfathered so only net-new errors fail and the ceilings ratchet down

basedpyright: error on reportExplicitAny and reportAny. The per-file budget moves 117033 -> 148946 (6931 explicit-Any, 24954 Any-typed expressions), grandfathered the same way

* ci: add Any-discipline gate on changed lines under litellm/

Add scripts/check_any_discipline.py, a type-aware gate that fails when a
changed line holds a value typed Any -- including the X | Any unions that
mypy --strict / basedpyright accept (e.g. re.Match.group() -> str | Any,
json.loads() -> Any, bare dict -> dict[Any, Any]).

It reuses the repo's mypyc-compiled mypy 1.19 via a custom generic AST
walker (mypyc precludes subclassing TraverserVisitor), loads litellm/mypy.ini
for parity with lint-mypy, and uses a dedicated incremental cache
(.mypy_cache_any) with mtime+hash invalidation to force re-checks. Scope is
changed-lines-only so editing a legacy file never forces cleaning its
existing Any debt; suppress a genuine typed/untyped boundary with
# any-ok: <reason> (ANY002 requires the reason).

Wire it into the Makefile (lint-any, lint, lint-dev), a parallel
any-discipline CI job with its own actions/cache, .gitignore, and the
CLAUDE.md / CONTRIBUTING.md docs.

* ci: move Any-gate codes into the shared LIT namespace

Renumber the Any-discipline checker into the LIT*** scheme owned by
scripts/check_type_discipline.py (PR #30500) so the two checkers share one
rule namespace and suppression convention:

  ANY001 -> LIT002  (Any-typed value; LIT002 was the retired/free slot)
  ANY002 -> LIT005  (any-ok without a reason; the shared suppression-reason code)
  ANY000 -> LIT000  (setup/build/read error; the shared error code)

Messages and behavior are unchanged; LIT005's text already matches the
"<token> requires a reason" shape used for cast-ok/guard-ok.

* ci: gate mypy and basedpyright per error rule, not per file

Switch the mypy/basedpyright budget gate from per-file error counts to
per-rule-code totals, mirroring the {rule: {baseline, slack}} shape of
ruff-strict-budget.json. A rule fails when its codebase-wide error count
exceeds baseline + slack, so violations are tracked by category rather
than by file location.

scripts/type_check_gate.py now parses mypy from its text output (trailing
[code]) and basedpyright from --outputjson (the JSON `rule` field), since
basedpyright's wrapped text diagnostics mis-attribute the rule on
continuation lines. Replace the *-file-budget.json files with freshly
captured *-code-budget.json baselines and update the Makefile, CI, and
CLAUDE.md accordingly.

* docs: prefer Pydantic validation over any-ok suppression

Point the Any-discipline guidance at validating Any with Pydantic (a model
or TypeAdapter that returns a typed value or raises) and frame
# any-ok as a last resort that should ideally never be used.

* chore: remove extraneous comment

* chore: make the CLAUDE.md more concise

* chore: clean up bloated CONTRIBUTING.md additions

* chore: make Makefile more concise

* ci: add the lint-budget-update target CLAUDE.md references

CLAUDE.md tells contributors to run make lint-budget-update, but the
target was never defined. Add it as an aggregate that re-captures the
ruff, mypy, and basedpyright budgets in one shot.

* ci: recapture mypy and basedpyright budgets in the lint env

The per-rule baselines were captured in a richer dependency env than the
CI lint job's uv sync --frozen, so CI resolved fewer types and reported
more errors than the budgets allowed (no-any-return 902 over cap 900, plus
several basedpyright reportUnknown* rules). Regenerate both in the frozen
env so they grandfather the true CI debt: mypy 5786 -> 5799 (no-any-return
890 -> 902, valid-type 1 restored), basedpyright 146213 -> 148942.

* ci: check out PR head sha in lint and any-discipline jobs

The default pull_request checkout uses refs/pull/N/merge, which folds the
latest base commits into HEAD. The diff-based gates (ruff delta, Any
discipline) then diff against the event's older base.sha and blame base's
own new commits on this branch; staging's otel-v2 and streaming changes
(#30326, #30485) tripped the Any gate on files this branch never touched.
Checking out the PR head sha makes the gates diff the real branch tip
against base, and pins the tree the mypy/basedpyright budgets were captured
against so their counts stay deterministic as the base advances.

* ci(lint): renumber Any-typed-value rule LIT002 -> LIT009

Free up LIT002 for the sibling type-discipline gate (check_type_discipline.py,
#30500), which groups its mutable-collection family at LIT001 (annotation) and
LIT002 (construction). This gate's Any-typed-value rule moves to LIT009 so the
shared LIT namespace stays contiguous with no holes; LIT000 and LIT005 are
unchanged.

* style: rename lint-strict-budget -> lint-ruff-budget

* ci: harden type-check gates against silent passes (greptile review)

type_check_gate.py: refuse to certify a vacuous run. The CI pipe swallows
the tool's exit code ('tool || true'), so a crashed mypy/basedpyright that
emits nothing would parse to zero errors, breach no ceiling, and pass.
is_vacuous_run() now fails when nothing was parsed but the budget expects
errors. Also wrap basedpyright's json.loads in a JSONDecodeError handler
that prints the offending output instead of dumping a raw traceback.

check_any_discipline.py: ALL_LINES was None, which dict.get() also returns
for a path absent from the line map, so a path-normalisation mismatch could
let a violation on an unchanged file pass the scope filter. Make ALL_LINES a
distinct sentinel object so 'whole file' and 'path missing' are unambiguous.

Adds tests for all three.
2026-06-16 12:07:46 -07:00
ryan-crabbe-berri
c90eb7e96f
feat: ruff strict-rule suppressions baseline gate (#30303)
* feat: add ruff strict-rule suppressions baseline gate

Introduce a stricter ruff rule set (typed params, no Any, complexity and
arg-count caps, mutable-default and global-rebinding checks) grandfathered
against the current tree and enforced as a budget rather than zero-tolerance

ruff-strict.toml defines the 9 rules separately from ruff.toml so the existing
ruff check stays green. scripts/ruff_suppressions.py builds the per-file,
per-rule baseline in ruff-suppressions.json and gates CI by failing when the
total grows past the baseline plus a 0.5% slack margin. The baseline ratchets
down via `make lint-suppressions-update` after fixes

* fix: surface per-file drift as a warning on a passing suppressions check

Greptile flagged that cmd_check computed per-file regressions but only printed
them on failure, so violations shifted between files (or a brand-new file under
the slack) passed with a silent OK. Print them as a non-fatal warning on the
pass path too; pass/fail behavior is unchanged

* refactor: gate strict ruff rules on the delta vs base, not a frozen baseline

The committed total-count baseline went stale against a moving base. CI lints the
PR merged with the current staging tip, so violations merged by other PRs counted
against this PR and tripped the budget even though nothing here touched them

Replace it with a drift-proof gate. scripts/ruff_strict_gate.py runs ruff on the
head, keeps only violations on lines this change adds relative to the merge-base,
and fails when a rule exceeds its per-rule allowance in ruff-strict-budget.json
(all 0 today). Because the base is measured live, base drift cancels out and only
what the change introduces is gated. Drops ruff-suppressions.json and the old
suppressions script

* chore: allow 5 new ANN001/ANN003/ANN401 per change

Give the three annotation-completeness rules a small per-change allowance so a
large new module is not blocked over a few untyped params or kwargs, while the
correctness and structural rules (B006, C901, PLR0913, PLW0603, RUF012, ANN002)
stay at 0

* feat: add TID251 typing.Any/Dict import ban and widen annotation budgets

Add TID251 (flake8-tidy-imports banned-api) to ruff-strict.toml, banning new
imports of typing.Any and typing.Dict and steering new code toward structured
types. It counts the import site, about one per file, so it is set non-blocking
at 50 as a forward-looking signal

Widen the annotation-completeness budgets so they nudge rather than block:
ANN001 50, ANN401 50, ANN003 25. Correctness and structural rules stay at 0

* refactor: make the strict gate a drift-safe per-rule total ceiling

Switch the gate from a per-change allowance to a hard ceiling on each rule's
total count across the codebase. The ceiling is baseline + slack in
ruff-strict-budget.json, with baseline captured from today's tree

To stay drift-safe, the gate counts each rule on the head and on the merge-base
(via a throwaway git worktree) and fails a rule only when its head total is over
the ceiling and higher than the base, so base drift never blames a change that
did not add to that rule. Annotation rules keep generous slack (ANN001 and
ANN401 50, ANN003 25, TID251 50); structural and correctness rules are frozen at
today's count. Add make lint-strict-budget-update to re-capture baselines

* chore: give the structural strict rules a cushion of 3

To be liberal to start, B006, C901, PLR0913, PLW0603, RUF012, and ANN002 each get
a slack of 3 instead of 0, so an occasional legitimate case is not hard-blocked.
The annotation budgets are unchanged, and these ratchet down later

* feat: ban more typing collection aliases and tighten annotation slack to 10

Add typing.List, typing.Set, typing.MutableSequence, and typing.MutableMapping to
the TID251 banned-api list, steering new code toward tuple, Sequence, Mapping,
frozenset, and frozen dataclasses. This raises TID251's baseline to 2404

Bring the three rules that were at slack 50 (ANN001, ANN401, TID251) down to 10

* docs: document the strict-gate ratchet and Any-avoidance in CLAUDE.md

Add a line on running make lint-strict-budget-update to knock baselines down
after fixes, and a line on validating untyped inputs in the caller rather than
spending the Any budget

* feat: make it a bit more strict

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
2026-06-12 20:14:45 -07:00