From fb98e6c229c365abd5a55a2d12964205f900ee5e Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 25 Sep 2026 17:03:24 -0700 Subject: [PATCH 1/2] fix(callbacks-legacy-python): traverse and release the retained headers dict LegacyLogging keeps the headers dict it hands to pre_call and post_call, but its traverse never reported that edge to the collector and close never dropped it. A cycle a callback builds through that dict could not be collected, and a closed call kept the dict alive until the driver dropped the whole adapter. Visit and clear headers like body, with regression tests for both --- .../callbacks-legacy-python/src/adapter.rs | 144 ++++++++++++++++-- 1 file changed, 133 insertions(+), 11 deletions(-) diff --git a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index 718cc615f309..fe0f6c7dd45b 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs @@ -469,6 +469,7 @@ impl PythonLifecycle for LegacyLogging { error.write_unraisable(py, None); } self.body = None; + self.headers = None; self.context = None; self.stream = None; } @@ -486,7 +487,8 @@ impl PythonLifecycle for LegacyLogging { visit.call(&stream.chunks)?; visit.call(&stream.first_chunk)?; } - visit.call(&self.body) + visit.call(&self.body)?; + visit.call(&self.headers) } } @@ -743,6 +745,7 @@ mod payload_tests { use litellm_host::event::{MachineEvent, RawResponse, RequestContext, WireRequest}; use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle, to_py}; use proptest::prelude::*; + use pyo3::gc::{PyTraverseError, PyVisit}; use pyo3::prelude::*; use rstest::rstest; use serde_json::{Map, Value, json}; @@ -832,16 +835,7 @@ check = lambda: None headers: vec![("x-route".into(), "route".into())], body, }; - let step = logging.before_send(py, Box::new(wire), &context).unwrap(); - let raw = MachineEvent::ResponseReceived { - raw: RawResponse { - body: "raw response".into(), - }, - }; - assert!(matches!( - logging.emit(py, LifecycleEvent::Machine(&raw)).unwrap(), - LifecycleStep::Done - )); + let (_, step) = send_and_receive(py, &mut logging, wire, &context); run(py, &locals, c"check()"); let LifecycleStep::Wire(wire) = step else { panic!("before_send did not hand back the wire request"); @@ -850,6 +844,134 @@ check = lambda: None }) } + /// `before_send` over `wire`, then the provider's raw response the way the driver + /// delivers it, so `pre_call` and `post_call` have both seen the retained payload. + fn send_and_receive<'a>( + py: Python<'_>, + logging: &'a mut LegacyLogging, + wire: WireRequest, + context: &RequestContext, + ) -> (&'a mut LegacyLogging, LifecycleStep) { + let step = logging.before_send(py, Box::new(wire), context).unwrap(); + let raw = MachineEvent::ResponseReceived { + raw: RawResponse { + body: "raw response".into(), + }, + }; + assert!(matches!( + logging.emit(py, LifecycleEvent::Machine(&raw)).unwrap(), + LifecycleStep::Done + )); + (logging, step) + } + + fn route_context() -> RequestContext { + RequestContext { + model: "model".into(), + custom_llm_provider: "provider".into(), + optional_params: json!({}), + secret_fields: vec![], + api_key: Some(SecretValue::new("route-key")), + } + } + + fn route_wire() -> WireRequest { + WireRequest { + url: "https://provider.invalid/ocr".into(), + headers: vec![("x-route".into(), "route".into())], + body: json!({}), + } + } + + /// A Python object owning one `LegacyLogging`, so the interpreter's collector sees the + /// edges the adapter reports and clears them the way the driver's `Execution` does. + #[pyclass(weakref)] + struct Retained { + logging: Option, + } + + #[pymethods] + impl Retained { + fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + match &self.logging { + Some(logging) => logging.traverse(&visit), + None => Ok(()), + } + } + + fn __clear__(slf: &Bound<'_, Self>) { + drop(slf.borrow_mut().logging.take()); + } + } + + #[test] + fn a_cycle_through_the_retained_headers_is_collected() { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, PAYLOAD_LOGGER); + let mut logging = LegacyLogging { + logger: Some(PythonLogger::new(local(&locals, "logger").unbind())), + ..legacy_call(py, &locals, false) + }; + send_and_receive(py, &mut logging, route_wire(), &route_context()); + let retained = Py::new( + py, + Retained { + logging: Some(logging), + }, + ) + .unwrap(); + locals.set_item("retained", retained).unwrap(); + run( + py, + &locals, + c" +import gc +import weakref + +logger.post[2]['headers']['owner'] = retained +logger.pre = logger.post = None +reference = weakref.ref(retained) +del retained +gc.collect() +assert reference() is None +", + ); + }); + } + + #[test] + fn close_releases_the_retained_headers() { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, PAYLOAD_LOGGER); + let mut logging = LegacyLogging { + logger: Some(PythonLogger::new(local(&locals, "logger").unbind())), + ..legacy_call(py, &locals, false) + }; + send_and_receive(py, &mut logging, route_wire(), &route_context()); + run( + py, + &locals, + c" +import weakref + +class Sentinel: + pass + +sentinel = Sentinel() +logger.post[2]['headers']['sentinel'] = sentinel +logger.pre = logger.post = None +reference = weakref.ref(sentinel) +del sentinel +assert reference() is not None +", + ); + logging.close(py); + run(py, &locals, c"assert reference() is None"); + }); + } + #[rstest] #[case::caller_keyword(c" document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} From 7ef172809a1769a4d332d75e1c5f8e93f0e958fb Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 25 Sep 2026 17:09:46 -0700 Subject: [PATCH 2/2] refactor(callbacks-legacy-python): move the test support module into its own file --- .../crates/callbacks-legacy-python/src/lib.rs | 202 +----------------- .../src/test_support.rs | 199 +++++++++++++++++ 2 files changed, 200 insertions(+), 201 deletions(-) create mode 100644 litellm-rust/crates/callbacks-legacy-python/src/test_support.rs diff --git a/litellm-rust/crates/callbacks-legacy-python/src/lib.rs b/litellm-rust/crates/callbacks-legacy-python/src/lib.rs index 869c534acf4a..8fca64d1b0e6 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/lib.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/lib.rs @@ -23,204 +23,4 @@ pub(crate) use callbacks::{LegacyCallbacks, is_internal_call}; pub(crate) use logger::{DeploymentHooks, PythonLogger, finalize, setup}; #[cfg(test)] -mod test_support { - use std::ffi::CStr; - - use pyo3::prelude::*; - use pyo3::types::{PyDict, PyTuple}; - - use crate::{LegacyLogging, LegacySurface, PublicCall}; - - /// The parameters of every `callbacks_legacy_python` function, as the real module declares them. - /// `tests/unit/rust_bridge/test_callbacks_legacy_python.py` pins this file to the Python - /// signatures, and [`namespace`] binds every fake call against it. - pub(crate) const PYTHON_CONTRACT: &str = include_str!("../python_contract.json"); - - /// Stand-ins for `callbacks_legacy_python`, the only Python module the crate calls. Tests - /// share one interpreter and run concurrently, so each fake is installed idempotently and - /// forwards to the per-test `StubLogger` it is handed (directly, or as `kwargs['logger']`). - /// Every fake is bound against the contract first, so a call the real module would reject - /// fails here too. - const STUBS: &CStr = c" -import contextvars -import inspect -import json -import sys -import traceback -import types - -for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.callbacks_legacy_python'): - sys.modules.setdefault(name, types.ModuleType(name)) - -legacy = sys.modules['litellm.rust_bridge.callbacks_legacy_python'] -CONTRACT = json.loads(python_contract) - - -def contracted(name, fake): - signature = inspect.Signature( - [inspect.Parameter(parameter, inspect.Parameter.POSITIONAL_OR_KEYWORD) for parameter in CONTRACT[name]] - ) - - def checked(*args, **kwargs): - signature.bind(*args, **kwargs) - return fake(*args, **kwargs) - - return checked - - -if not hasattr(legacy, 'is_internal'): - legacy.is_internal = contextvars.ContextVar('is_internal_call', default=False) - -FAKES = { - 'setup': lambda call_type, args, kwargs, start, asynchronous: types.SimpleNamespace( - logger=kwargs['logger_factory'](kwargs) if 'logger_factory' in kwargs else kwargs['logger'], - kwargs=kwargs, - ), - 'finalize': lambda response, logger, kwargs, start, end: logger.record('finalize', response), - 'update_logging': lambda logger, kwargs, model, optional_params, litellm_params, provider: logger.update_from_kwargs( - kwargs=kwargs, - model=model, - optional_params=optional_params, - litellm_params=litellm_params, - custom_llm_provider=provider, - ), - 'pre_call': lambda logger, input, api_key, additional_args: logger.pre_call(input, api_key, additional_args), - 'post_call': lambda logger, original_response, api_key, additional_args: logger.post_call( - original_response, api_key, additional_args - ), - 'defers_async_logging': lambda logger: bool(getattr(logger, '_defer_async_logging', False)), - 'defer_success': lambda logger, pending: setattr(logger, '_native_pending_logging', pending), - 'sync_success_for_async_call': lambda logger, response, start, end: logger.handle_sync_success_callbacks_for_async_calls( - response, start, end - ), - 'failure_handler': lambda logger, error, start, end, asynchronous: ( - logger.async_failure_handler if asynchronous else logger.failure_handler - )(error, ''.join(traceback.format_exception(error)), start, end), - 'submit_success': lambda logger, response, start, end: logger.record('submit', (response, start, end)), - 'async_success_handler': lambda logger, response, start, end: logger.async_success_handler(response, start, end), - 'enqueue_logging': lambda coroutine: coroutine.enqueue(), - 'restore_context': lambda logger: logger.record('restore', None), - 'custom_pricing_fields': lambda: ('ocr_cost_per_page',), - 'is_internal_call': lambda: legacy.is_internal.get(), - 'before_deployment_call': lambda kwargs, call_type: kwargs['logger'].hook('pre', kwargs, call_type), - 'after_deployment_success': lambda kwargs, response, call_type: kwargs['logger'].hook( - 'success', response, call_type - ), - 'after_deployment_failure': lambda kwargs, error, call_type: kwargs['logger'].hook('failure', error, call_type), - 'stream_opened': lambda logger: logger.record('stream_opened', None), - 'stream_success': lambda logger, request_body, chunks, start, end, first_chunk: logger.record( - 'stream_success', list(chunks) - ), - 'stream_failure': lambda logger, request_body, chunks, error: logger.record('stream_failure', error), -} -assert FAKES.keys() == CONTRACT.keys(), sorted(FAKES.keys() ^ CONTRACT.keys()) -for name, fake in FAKES.items(): - setattr(legacy, name, contracted(name, fake)) - - -unraisable = sys.modules.setdefault( - 'litellm_test_unraisable', types.ModuleType('litellm_test_unraisable') -) -if not hasattr(unraisable, 'events'): - unraisable.events = [] - sys.unraisablehook = lambda event: unraisable.events.append((event.object, event.exc_value)) - - -def unraisable_from(owner): - return [error for source, error in unraisable.events if source is owner] - - -class StubCoroutine: - def __init__(self, logger): - self.logger = logger - - def enqueue(self): - self.logger.record('enqueued', None) - self.logger.on_enqueue(self) - - def close(self): - self.logger.record('closed', None) - - -class StubLogger: - def __init__(self): - self.calls = [] - self.hooks = {} - self.on_enqueue = lambda coroutine: None - - def record(self, name, value): - self.calls.append((name, value)) - - def names(self): - return [name for name, _ in self.calls] - - def hook(self, phase, value, call_type): - self.record(phase + '_hook', call_type) - return self.hooks.get(phase, lambda value: 'awaitable')(value) - - def failure_handler(self, error, trace, start, end): - self.record('failure_handler', error) - - def async_failure_handler(self, error, trace, start, end): - self.record('async_failure_handler', error) - return 'awaitable' - - def success_handler(self, response, start, end): - self.record('success_handler', response) - - def async_success_handler(self, response, start, end): - self.record('async_success_handler', response) - return StubCoroutine(self) - - def handle_sync_success_callbacks_for_async_calls(self, response, start, end): - self.record('sync_success_for_async_call', response) - - -logger = StubLogger() -"; - - /// A namespace with the stubs, `StubLogger` and a fresh `logger`, after `script` ran in it. - pub(crate) fn namespace<'py>(py: Python<'py>, script: &CStr) -> Bound<'py, PyDict> { - let locals = PyDict::new(py); - locals.set_item("python_contract", PYTHON_CONTRACT).unwrap(); - py.run(STUBS, Some(&locals), Some(&locals)).unwrap(); - py.run(script, Some(&locals), Some(&locals)).unwrap(); - locals - } - - pub(crate) fn run(py: Python<'_>, locals: &Bound<'_, PyDict>, code: &CStr) { - py.run(code, Some(locals), Some(locals)).unwrap(); - } - - pub(crate) fn local<'py>(locals: &Bound<'py, PyDict>, name: &str) -> Bound<'py, PyAny> { - locals.get_item(name).unwrap().unwrap() - } - - /// A legacy call over the namespace's `kwargs` (or none) and `request` (or `None`). - pub(crate) fn legacy_call( - py: Python<'_>, - locals: &Bound<'_, PyDict>, - asynchronous: bool, - ) -> LegacyLogging { - let request = locals - .get_item("request") - .unwrap() - .unwrap_or_else(|| py.None().into_bound(py)); - let kwargs = locals - .get_item("kwargs") - .unwrap() - .map(|kwargs| kwargs.cast_into::().unwrap()) - .unwrap_or_else(|| PyDict::new(py)); - let call = PublicCall::capture(&request, &PyTuple::empty(py), &kwargs).unwrap(); - LegacyLogging::new( - py, - LegacySurface { - call_type: "test", - input_description: "test input", - stream: None, - }, - call, - asynchronous, - ) - } -} +mod test_support; diff --git a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs new file mode 100644 index 000000000000..e7973a7e1a0f --- /dev/null +++ b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs @@ -0,0 +1,199 @@ +use std::ffi::CStr; + +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyTuple}; + +use crate::{LegacyLogging, LegacySurface, PublicCall}; + +/// The parameters of every `callbacks_legacy_python` function, as the real module declares them. +/// `tests/unit/rust_bridge/test_callbacks_legacy_python.py` pins this file to the Python +/// signatures, and [`namespace`] binds every fake call against it. +pub(crate) const PYTHON_CONTRACT: &str = include_str!("../python_contract.json"); + +/// Stand-ins for `callbacks_legacy_python`, the only Python module the crate calls. Tests +/// share one interpreter and run concurrently, so each fake is installed idempotently and +/// forwards to the per-test `StubLogger` it is handed (directly, or as `kwargs['logger']`). +/// Every fake is bound against the contract first, so a call the real module would reject +/// fails here too. +const STUBS: &CStr = c" +import contextvars +import inspect +import json +import sys +import traceback +import types + +for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.callbacks_legacy_python'): + sys.modules.setdefault(name, types.ModuleType(name)) + +legacy = sys.modules['litellm.rust_bridge.callbacks_legacy_python'] +CONTRACT = json.loads(python_contract) + + +def contracted(name, fake): + signature = inspect.Signature( + [inspect.Parameter(parameter, inspect.Parameter.POSITIONAL_OR_KEYWORD) for parameter in CONTRACT[name]] + ) + + def checked(*args, **kwargs): + signature.bind(*args, **kwargs) + return fake(*args, **kwargs) + + return checked + + +if not hasattr(legacy, 'is_internal'): + legacy.is_internal = contextvars.ContextVar('is_internal_call', default=False) + +FAKES = { + 'setup': lambda call_type, args, kwargs, start, asynchronous: types.SimpleNamespace( + logger=kwargs['logger_factory'](kwargs) if 'logger_factory' in kwargs else kwargs['logger'], + kwargs=kwargs, + ), + 'finalize': lambda response, logger, kwargs, start, end: logger.record('finalize', response), + 'update_logging': lambda logger, kwargs, model, optional_params, litellm_params, provider: logger.update_from_kwargs( + kwargs=kwargs, + model=model, + optional_params=optional_params, + litellm_params=litellm_params, + custom_llm_provider=provider, + ), + 'pre_call': lambda logger, input, api_key, additional_args: logger.pre_call(input, api_key, additional_args), + 'post_call': lambda logger, original_response, api_key, additional_args: logger.post_call( + original_response, api_key, additional_args + ), + 'defers_async_logging': lambda logger: bool(getattr(logger, '_defer_async_logging', False)), + 'defer_success': lambda logger, pending: setattr(logger, '_native_pending_logging', pending), + 'sync_success_for_async_call': lambda logger, response, start, end: logger.handle_sync_success_callbacks_for_async_calls( + response, start, end + ), + 'failure_handler': lambda logger, error, start, end, asynchronous: ( + logger.async_failure_handler if asynchronous else logger.failure_handler + )(error, ''.join(traceback.format_exception(error)), start, end), + 'submit_success': lambda logger, response, start, end: logger.record('submit', (response, start, end)), + 'async_success_handler': lambda logger, response, start, end: logger.async_success_handler(response, start, end), + 'enqueue_logging': lambda coroutine: coroutine.enqueue(), + 'restore_context': lambda logger: logger.record('restore', None), + 'custom_pricing_fields': lambda: ('ocr_cost_per_page',), + 'is_internal_call': lambda: legacy.is_internal.get(), + 'before_deployment_call': lambda kwargs, call_type: kwargs['logger'].hook('pre', kwargs, call_type), + 'after_deployment_success': lambda kwargs, response, call_type: kwargs['logger'].hook( + 'success', response, call_type + ), + 'after_deployment_failure': lambda kwargs, error, call_type: kwargs['logger'].hook('failure', error, call_type), + 'stream_opened': lambda logger: logger.record('stream_opened', None), + 'stream_success': lambda logger, request_body, chunks, start, end, first_chunk: logger.record( + 'stream_success', list(chunks) + ), + 'stream_failure': lambda logger, request_body, chunks, error: logger.record('stream_failure', error), +} +assert FAKES.keys() == CONTRACT.keys(), sorted(FAKES.keys() ^ CONTRACT.keys()) +for name, fake in FAKES.items(): + setattr(legacy, name, contracted(name, fake)) + + +unraisable = sys.modules.setdefault( + 'litellm_test_unraisable', types.ModuleType('litellm_test_unraisable') +) +if not hasattr(unraisable, 'events'): + unraisable.events = [] + sys.unraisablehook = lambda event: unraisable.events.append((event.object, event.exc_value)) + + +def unraisable_from(owner): + return [error for source, error in unraisable.events if source is owner] + + +class StubCoroutine: + def __init__(self, logger): + self.logger = logger + + def enqueue(self): + self.logger.record('enqueued', None) + self.logger.on_enqueue(self) + + def close(self): + self.logger.record('closed', None) + + +class StubLogger: + def __init__(self): + self.calls = [] + self.hooks = {} + self.on_enqueue = lambda coroutine: None + + def record(self, name, value): + self.calls.append((name, value)) + + def names(self): + return [name for name, _ in self.calls] + + def hook(self, phase, value, call_type): + self.record(phase + '_hook', call_type) + return self.hooks.get(phase, lambda value: 'awaitable')(value) + + def failure_handler(self, error, trace, start, end): + self.record('failure_handler', error) + + def async_failure_handler(self, error, trace, start, end): + self.record('async_failure_handler', error) + return 'awaitable' + + def success_handler(self, response, start, end): + self.record('success_handler', response) + + def async_success_handler(self, response, start, end): + self.record('async_success_handler', response) + return StubCoroutine(self) + + def handle_sync_success_callbacks_for_async_calls(self, response, start, end): + self.record('sync_success_for_async_call', response) + + +logger = StubLogger() +"; + +/// A namespace with the stubs, `StubLogger` and a fresh `logger`, after `script` ran in it. +pub(crate) fn namespace<'py>(py: Python<'py>, script: &CStr) -> Bound<'py, PyDict> { + let locals = PyDict::new(py); + locals.set_item("python_contract", PYTHON_CONTRACT).unwrap(); + py.run(STUBS, Some(&locals), Some(&locals)).unwrap(); + py.run(script, Some(&locals), Some(&locals)).unwrap(); + locals +} + +pub(crate) fn run(py: Python<'_>, locals: &Bound<'_, PyDict>, code: &CStr) { + py.run(code, Some(locals), Some(locals)).unwrap(); +} + +pub(crate) fn local<'py>(locals: &Bound<'py, PyDict>, name: &str) -> Bound<'py, PyAny> { + locals.get_item(name).unwrap().unwrap() +} + +/// A legacy call over the namespace's `kwargs` (or none) and `request` (or `None`). +pub(crate) fn legacy_call( + py: Python<'_>, + locals: &Bound<'_, PyDict>, + asynchronous: bool, +) -> LegacyLogging { + let request = locals + .get_item("request") + .unwrap() + .unwrap_or_else(|| py.None().into_bound(py)); + let kwargs = locals + .get_item("kwargs") + .unwrap() + .map(|kwargs| kwargs.cast_into::().unwrap()) + .unwrap_or_else(|| PyDict::new(py)); + let call = PublicCall::capture(&request, &PyTuple::empty(py), &kwargs).unwrap(); + LegacyLogging::new( + py, + LegacySurface { + call_type: "test", + input_description: "test input", + stream: None, + }, + call, + asynchronous, + ) +}