scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

fix(messages): log native messages success once, after the proxy releases it (#50420)

Browse source 
* fix(messages): log native messages success once, after the proxy releases it

The proxy defers non-stream success logging past post-call guardrails, but the native host parked its log in its own slot and the release only ran for OCR, so a unary /v1/messages call served by Rust logged no success event, no SpendLogs row and no key, team, user or end-user spend. A deferred native success now parks in enqueue_deferred_logging with deferred_success_result, the same slots Python uses, so the proxy releases it like any other call and its guardrail step pruning and cost preparation see it too

Python runs sync success callbacks inside the deferred logging closure, so a post-call hook that blocks the response drops them. A parked native success now carries its sync callbacks and runs them on release

A refused first attempt and its fallback share one logger. Python logs the async success once, for the first attempt, so the native adapter now claims the async success log once across fallback attempts. Like Python, an attempt that loses the claim runs its sync callbacks only when logging is not deferred, and an internal call logs nothing

* test(messages): cover the deferred native success park and release
This commit is contained in:
Nate Armstrong 2026-11-09 16:06:53 -08:00 • committed by GitHub
parent c6794ac08b
commit 1ddc913002
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 269 additions and 102 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
litellm-rust/crates/callbacks-legacy-python/python_contract.json
View file

@ -47,9 +47,18 @@
"defers_async_logging": [
"logger"
],
"claim_async_success_log": [
"logger"
],
"mark_async_success_enqueued": [
"logger"
],
"defer_success": [
"logger",
"pending"
"pending",
"response",
"start",
"end"
],
"sync_success_for_async_call": [
"logger",

131
litellm-rust/crates/callbacks-legacy-python/src/adapter.rs
View file

@ -251,26 +251,40 @@ impl LegacyLogging {
if !self.asynchronous {
return pending().sync(py);
}
if !self.internal
&& self
if self.internal {
return Ok(());
}
let defers = logger.defers_async_logging(py);
if logger.claim_async_success_log(py)? {
let enqueues = self
.call
.kwargs()
.bind(py)
.get_item("fallbacks")?
.is_none_or(|value| value.is_none())
{
if logger.defers_async_logging(py) {
.is_none_or(|value| value.is_none());
if defers {
let pending = Py::new(
py,
PendingLogging {
pending: Some(pending()),
enqueues,
},
)?;
logger.defer_success(py, pending.bind(py).as_any())?;
} else {
return logger.defer_success(
py,
pending.bind(py).as_any(),
&self.response,
&self.start,
&self.end,
);
}
if enqueues {
pending().asynchronous(py)?;
}
}
if defers {
return Ok(());
}
logger.sync_success_for_async_call(py, &self.response, &self.start, &self.end)
}
@ -1873,10 +1887,15 @@ mod terminal_tests {
#[case::async_listened(
true,
c"",
&["async_success_handler", "enqueued", "sync_success_for_async_call"]
&["async_success_handler", "mark_async_success_enqueued", "enqueued", "sync_success_for_async_call"]
)]
#[case::async_deferred(true, c"logger._defer_async_logging = True", &["sync_success_for_async_call"])]
#[case::async_deferred(true, c"logger._defer_async_logging = True", &[])]
#[case::async_with_fallbacks(true, c"kwargs = {'fallbacks': ['other']}", &["sync_success_for_async_call"])]
#[case::async_deferred_with_fallbacks(
true,
c"logger._defer_async_logging = True\nkwargs = {'fallbacks': ['other']}",
&[]
)]
fn success_reaches_the_logging_handlers(
#[case] asynchronous: bool,
#[case] script: &CStr,
@ -1902,12 +1921,88 @@ mod terminal_tests {
&locals,
c"
assert all(value is response for name, value in logger.calls if name.endswith('_handler'))
assert hasattr(logger, '_native_pending_logging') == getattr(logger, '_defer_async_logging', False)
assert hasattr(logger, 'enqueue_deferred_logging') == getattr(logger, '_defer_async_logging', False)
assert getattr(logger, 'deferred_success_result', (response,))[0] is response
",
);
});
}
#[rstest]
#[case::accepted(c"", &["sync_success_for_async_call", "async_success_handler", "mark_async_success_enqueued", "enqueued"])]
#[case::accepted_with_fallbacks(c"kwargs = {'fallbacks': ['other']}", &["sync_success_for_async_call"])]
fn a_deferred_success_runs_sync_callbacks_only_once_released(
#[case] script: &CStr,
#[case] expected: &[&str],
) {
Python::initialize();
Python::attach(|py| {
let locals = namespace(
py,
c"response = object()\nlogger._defer_async_logging = True",
);
run(py, &locals, script);
let mut logging = logged(py, &locals, true);
succeed(py, &locals, &mut logging);
run(
py,
&locals,
c"
assert logger.names() == [], logger.calls
logger.enqueue_deferred_logging()
logger.enqueue_deferred_logging()
",
);
let names: Vec<String> = local(&locals, "logger")
.call_method0("names")
.unwrap()
.extract()
.unwrap();
assert_eq!(names, expected);
});
}
#[rstest]
#[case::deferred(
c"logger._defer_async_logging = True",
c"logger.enqueue_deferred_logging()",
&["sync_success_for_async_call", "async_success_handler", "mark_async_success_enqueued", "enqueued"]
)]
#[case::deferred_and_blocked(c"logger._defer_async_logging = True", c"", &[])]
#[case::immediate(
c"",
c"",
&["async_success_handler", "mark_async_success_enqueued", "enqueued", "sync_success_for_async_call", "sync_success_for_async_call"]
)]
fn a_fallback_attempt_on_the_same_logger_leaves_the_first_success_claimed(
#[case] script: &CStr,
#[case] release: &CStr,
#[case] expected: &[&str],
) {
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, c"refused = object()\nresponse = refused");
run(py, &locals, script);
let mut first = logged(py, &locals, true);
succeed(py, &locals, &mut first);
run(py, &locals, c"response = object()");
let mut fallback = logged(py, &locals, true);
succeed(py, &locals, &mut fallback);
run(py, &locals, release);
let names: Vec<String> = local(&locals, "logger")
.call_method0("names")
.unwrap()
.extract()
.unwrap();
assert_eq!(names, expected);
run(
py,
&locals,
c"assert all(value is refused for name, value in logger.calls if name == 'async_success_handler')",
);
});
}
#[rstest]
fn dropped_observations_preserve_deferred_success_and_response_identity() {
Python::initialize();
@ -1933,10 +2028,10 @@ assert hasattr(logger, '_native_pending_logging') == getattr(logger, '_defer_asy
));
assert_eq!(sender.dropped_events(), 1);
run(py, &locals, c"
assert logger.names() == ['sync_success_for_async_call'], logger.calls
logger._native_pending_logging.release(True)
logger._native_pending_logging.release(True)
assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'enqueued'], logger.calls
assert logger.names() == [], logger.calls
logger.enqueue_deferred_logging()
logger.enqueue_deferred_logging()
assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'mark_async_success_enqueued', 'enqueued'], logger.calls
assert logger.calls[0][1] is response
assert logger.calls[1][1] is response
");
@ -2009,7 +2104,7 @@ assert chunks[1] is last
}
#[rstest::rstest]
fn internal_async_calls_skip_the_async_success_fan_out() {
fn internal_async_calls_skip_success_logging() {
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, c"response = object()");
@ -2018,11 +2113,7 @@ assert chunks[1] is last
..logged(py, &locals, true)
};
succeed(py, &locals, &mut logging);
run(
py,
&locals,
c"assert logger.names() == ['sync_success_for_async_call'], logger.calls",
);
run(py, &locals, c"assert logger.calls == [], logger.calls");
});
}

29
litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs
View file

@ -59,7 +59,16 @@ pub trait LegacyCallbacks {
fn defers_async_logging(&self, py: Python<'_>) -> bool;
fn defer_success(&self, py: Python<'_>, pending: &Bound<'_, PyAny>) -> PyResult<()>;
fn claim_async_success_log(&self, py: Python<'_>) -> PyResult<bool>;
fn defer_success(
&self,
py: Python<'_>,
pending: &Bound<'_, PyAny>,
response: &Option<Py<PyAny>>,
start: &Py<PyAny>,
end: &Option<Py<PyAny>>,
) -> PyResult<()>;
fn sync_success_for_async_call(
&self,
@ -203,8 +212,21 @@ impl LegacyCallbacks for PythonLogger {
.unwrap_or(false)
}
fn defer_success(&self, py: Python<'_>, pending: &Bound<'_, PyAny>) -> PyResult<()> {
Logging::DeferSuccess.call(py, (self.object(py), pending))?;
fn claim_async_success_log(&self, py: Python<'_>) -> PyResult<bool> {
Logging::ClaimAsyncSuccess
.call(py, (self.object(py),))?
.extract()
}
fn defer_success(
&self,
py: Python<'_>,
pending: &Bound<'_, PyAny>,
response: &Option<Py<PyAny>>,
start: &Py<PyAny>,
end: &Option<Py<PyAny>>,
) -> PyResult<()> {
Logging::DeferSuccess.call(py, (self.object(py), pending, response, start, end))?;
Ok(())
}
@ -252,6 +274,7 @@ impl LegacyCallbacks for PythonLogger {
) -> PyResult<()> {
let coroutine =
Logging::AsyncSuccessHandler.call(py, (self.object(py), response, start, end))?;
Logging::MarkAsyncSuccessEnqueued.call(py, (self.object(py),))?;
let enqueue = Logging::Enqueue.call(py, (&coroutine,));
if enqueue.is_err()
&& let Err(error) = coroutine.call_method0("close")

59
litellm-rust/crates/callbacks-legacy-python/src/deferred.rs
View file

@ -22,21 +22,35 @@ impl PendingSuccess {
self.logger
.enqueue_success(py, &self.response, &self.start, &self.end)
}
pub(crate) fn sync_for_async_call(&self, py: Python<'_>) -> PyResult<()> {
self.logger
.sync_success_for_async_call(py, &self.response, &self.start, &self.end)
}
}
#[pyclass]
pub(crate) struct PendingLogging {
pub(crate) pending: Option<PendingSuccess>,
pub(crate) enqueues: bool,
}
#[pymethods]
impl PendingLogging {
fn release(slf: &Bound<'_, Self>, py: Python<'_>, success: bool) -> PyResult<()> {
let pending = slf.borrow_mut().pending.take();
if let Some(pending) = pending
&& success
{
match pending.asynchronous(py) {
fn __call__(slf: &Bound<'_, Self>, py: Python<'_>) -> PyResult<()> {
let (pending, enqueues) = {
let mut slot = slf.borrow_mut();
(slot.pending.take(), slot.enqueues)
};
if let Some(pending) = pending {
let released = pending.sync_for_async_call(py).and_then(|()| {
if enqueues {
pending.asynchronous(py)
} else {
Ok(())
}
});
match released {
Err(error) if error.is_instance_of::<PyException>(py) => {
error.write_unraisable(py, Some(pending.logger.object(py)));
}
@ -87,6 +101,7 @@ mod tests {
start: py.None(),
end: Some(py.None()),
}),
enqueues: true,
},
)
.unwrap();
@ -108,7 +123,7 @@ observed = []
def on_enqueue(coroutine):
observed.append(marker.get())
pending.release(True)
pending()
logger.on_enqueue = on_enqueue
",
@ -118,28 +133,12 @@ logger.on_enqueue = on_enqueue
&locals,
c"
marker.set('release')
pending.release(True)
pending.release(True)
pending()
pending()
assert observed == ['release'], observed
assert logger.names() == ['async_success_handler', 'enqueued'], logger.calls
assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'mark_async_success_enqueued', 'enqueued'], logger.calls
assert logger.calls[0][1] is response
",
);
});
}
#[test]
fn a_blocked_release_drops_the_success_for_good() {
Python::initialize();
Python::attach(|py| {
let locals = defer(py, c"");
run(
py,
&locals,
c"
pending.release(False)
pending.release(True)
assert logger.calls == [], logger.calls
assert logger.calls[1][1] is response
",
);
});
@ -168,7 +167,7 @@ logger.on_enqueue = on_enqueue
locals
.set_item("failure", py.eval(failure, None, Some(&locals)).unwrap())
.unwrap();
let released = local(&locals, "pending").call_method1("release", (true,));
let released = local(&locals, "pending").call0();
match released {
Ok(_) => assert!(!propagates),
Err(error) => {
@ -181,8 +180,8 @@ logger.on_enqueue = on_enqueue
py,
&locals,
c"
pending.release(True)
assert logger.names() == ['async_success_handler', 'enqueued', 'closed'], logger.calls
pending()
assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'mark_async_success_enqueued', 'enqueued', 'closed'], logger.calls
assert unraisable_from(logger) == ([] if propagates else [failure])
",
);

4
litellm-rust/crates/callbacks-legacy-python/src/python.rs
View file

@ -50,6 +50,8 @@ pub(crate) enum Logging {
RecordApiWindow,
#[strum(serialize = "defers_async_logging")]
DefersAsync,
#[strum(serialize = "claim_async_success_log")]
ClaimAsyncSuccess,
#[strum(serialize = "defer_success")]
DeferSuccess,
#[strum(serialize = "sync_success_for_async_call")]
@ -58,6 +60,8 @@ pub(crate) enum Logging {
SubmitSuccess,
#[strum(serialize = "async_success_handler")]
AsyncSuccessHandler,
#[strum(serialize = "mark_async_success_enqueued")]
MarkAsyncSuccessEnqueued,
#[strum(serialize = "enqueue_logging")]
Enqueue,
#[strum(serialize = "failure_handler")]

11
litellm-rust/crates/callbacks-legacy-python/src/test_support.rs
View file

@ -72,7 +72,11 @@ FAKES = {
llm_api_duration_ms=(end - start) * 1000
),
'defers_async_logging': lambda logger: bool(getattr(logger, '_defer_async_logging', False)),
'defer_success': lambda logger, pending: setattr(logger, '_native_pending_logging', pending),
'claim_async_success_log': lambda logger: logger.claim_async_success_log(),
'defer_success': lambda logger, pending, response, start, end: (
setattr(logger, 'deferred_success_result', (response, start, end)),
setattr(logger, 'enqueue_deferred_logging', pending),
),
'sync_success_for_async_call': lambda logger, response, start, end: logger.handle_sync_success_callbacks_for_async_calls(
response, start, end
),
@ -81,6 +85,7 @@ FAKES = {
)(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),
'mark_async_success_enqueued': lambda logger: logger.record('mark_async_success_enqueued', None),
'enqueue_logging': lambda coroutine: coroutine.enqueue(),
'restore_context': lambda logger: logger.record('restore', None),
'custom_pricing_fields': lambda: ('ocr_cost_per_page',),
@ -156,6 +161,10 @@ class StubLogger:
self.record('async_success_handler', response)
return StubCoroutine(self)
def claim_async_success_log(self):
claimed, self.claimed = getattr(self, 'claimed', False), True
return not claimed
def handle_sync_success_callbacks_for_async_calls(self, response, start, end):
self.record('sync_success_for_async_call', response)

11
litellm/proxy/common_request_processing.py
View file

@ -64,7 +64,6 @@ from litellm.litellm_core_utils.bug_report import (
strip_bug_report_notice,
)
from litellm.litellm_core_utils.core_helpers import (
attribute_or_none,
finite_cost,
get_or_create_metadata_bucket,
has_attributes,
@ -72,7 +71,6 @@ from litellm.litellm_core_utils.core_helpers import (
is_expected_client_error,
proxy_request_headers,
redact_nested_match_and_regex_keys,
require_attributes,
stream_cost_injection_requested,
)
from litellm.litellm_core_utils.dd_tracing import NullTracer, tracer
@ -724,10 +722,6 @@ class _ReadsHeaderValues(Protocol):
def get(self, key: str, default: str = "") -> str: ...
class _NativeLoggingRelease(Protocol):
def release(self, success: bool, /) -> object: ...
def _as_header_reader(headers: _ReadsHeaderValues) -> _ReadsHeaderValues:
return headers
@ -4384,11 +4378,6 @@ class ProxyBaseLLMRequestProcessing:
Extracted as a static method so tests can exercise the production
gating logic directly rather than reimplementing the finally block.
"""
if getattr(logging_obj, "call_type", None) in ("ocr", "aocr"):
pending: Final = attribute_or_none(logging_obj, "_native_pending_logging")
if pending is not None:
logging_obj._native_pending_logging = None # rebind-ok: consume the native OCR release signal once
require_attributes(pending, _NativeLoggingRelease, "release").release(not exception_raised)
_enqueue_fn: Final = getattr(logging_obj, "enqueue_deferred_logging", None)
if _enqueue_fn is None:
return

28
litellm/rust_bridge/host/callbacks_legacy_python.py
View file

@ -12,7 +12,7 @@ import datetime
import traceback
import uuid
import weakref
from collections.abc import Awaitable, Coroutine, Mapping, Sequence
from collections.abc import Awaitable, Callable, Coroutine, Mapping, Sequence
from dataclasses import dataclass
from typing import (
TYPE_CHECKING,
@ -111,6 +111,9 @@ def finalize(
class LoggingSurface(Protocol):
deferred_success_result: tuple[object, datetime.datetime, datetime.datetime] | None
enqueue_deferred_logging: Callable[[], None] | None
@property
def model_call_details(self) -> Mapping[str, object]: ...
@ -139,6 +142,10 @@ class LoggingSurface(Protocol):
additional_args: dict[str, object] = ...,
) -> object: ...
def claim_async_success_log(self) -> bool: ...
def mark_async_success_log_enqueued(self) -> None: ...
def handle_sync_success_callbacks_for_async_calls(
self, result: object, start_time: datetime.datetime, end_time: datetime.datetime, cache_hit: bool | None = None
) -> None: ...
@ -269,8 +276,23 @@ def defers_async_logging(logger: LoggingSurface) -> bool:
return bool(getattr(logger, "defer_async_logging", False))
def defer_success(logger: LoggingSurface, pending: object) -> None:
setattr(logger, "_native_pending_logging", pending)
def claim_async_success_log(logger: LoggingSurface) -> bool:
return logger.claim_async_success_log()
def mark_async_success_enqueued(logger: LoggingSurface) -> None:
logger.mark_async_success_log_enqueued()
def defer_success(
logger: LoggingSurface,
pending: Callable[[], None],
response: object,
start: datetime.datetime,
end: datetime.datetime,
) -> None:
logger.deferred_success_result = (response, start, end) # rebind-ok: capture the native attempt before dispatch
logger.enqueue_deferred_logging = pending # rebind-ok: the proxy releases logging here
def _cache_hit(logger: LoggingSurface) -> Literal[True] | None:

39
tests/test_litellm_rust/ocr/test_lifecycle.py
View file

@ -211,8 +211,12 @@ async def test_deferred_logging_requires_release_and_runs_at_most_once(server: R
response: Final = await call_aocr(server, litellm_logging_obj=logger)
await drain_logging()
assert "async_log_success_event" not in recorder.names
assert logger.deferred_success_result is not None
assert logger.deferred_success_result[0] is response
assert not logger.async_success_log_enqueued
ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(logger, blocked)
ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(logger, blocked)
assert logger.async_success_log_enqueued is not blocked
await drain_logging()
events: Final = tuple(event for event in recorder.events if event.name == "async_log_success_event")
assert len(events) == int(not blocked)
@ -220,6 +224,39 @@ async def test_deferred_logging_requires_release_and_runs_at_most_once(server: R
assert events[0].response is response
@pytest.mark.asyncio
@pytest.mark.parametrize("blocked", [False, True])
async def test_deferred_sync_success_callbacks_wait_for_release(server: RecordingServer, blocked: bool) -> None:
finished: Final = threading.Event()
observations: Final = []
class Observe(CustomLogger):
def log_success_event(self, kwargs, response_obj, start_time, end_time):
observations.append(response_obj)
finished.set()
logger: Final = Logging(
model="mistral-ocr-latest",
messages=[],
stream=False,
call_type="aocr",
start_time=datetime.datetime.now(),
litellm_call_id="deferred-sync",
function_id="deferred-sync",
dynamic_success_callbacks=[Observe()],
)
logger._defer_async_logging = True
response: Final = await call_aocr(server, litellm_logging_obj=logger)
await drain_logging()
assert not await asyncio.to_thread(finished.wait, 0.2)
assert observations == []
ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(logger, blocked)
ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(logger, blocked)
await drain_logging()
assert await asyncio.to_thread(finished.wait, 0.5 if blocked else 5) is not blocked
assert observations == ([] if blocked else [response])
def test_sync_success_uses_executor_and_copied_caller_context(server: RecordingServer) -> None:
context: Final = ContextVar("sync-lifecycle", default="missing")
context.set("caller")
@ -454,7 +491,7 @@ async def test_empty_callbacks_run_deployment_hooks_and_defer_like_the_python_cl
assert logger.model_call_details["first_api_call_start_time"] <= logger.model_call_details["end_time"]
assert logger.model_call_details["response_cost"] == 0
else:
assert getattr(logger, "_native_pending_logging", None) is not None
assert logger.enqueue_deferred_logging is not None
assert "end_time" not in logger.model_call_details

32
tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py
View file

@ -416,38 +416,6 @@ async def test_deferred_anthropic_messages_bridged_to_the_responses_api_logs_the
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("call_type", ["ocr", "aocr", "completion", "acompletion", "embedding", "responses"])
@pytest.mark.parametrize("exception_raised", [False, True])
def test_native_pending_logging_is_released_only_for_ocr(call_type: str, exception_raised: bool) -> None:
pending: Final = MagicMock()
enqueue: Final = MagicMock()
logger: Final = MagicMock(
call_type=call_type,
_native_pending_logging=pending,
enqueue_deferred_logging=enqueue,
)
ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(
logging_obj=logger,
exception_raised=exception_raised,
)
ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(
logging_obj=logger,
exception_raised=exception_raised,
)
if call_type in ("ocr", "aocr"):
pending.release.assert_called_once_with(not exception_raised)
assert logger._native_pending_logging is None
else:
pending.release.assert_not_called()
assert logger._native_pending_logging is pending
if exception_raised:
enqueue.assert_not_called()
else:
enqueue.assert_called_once_with()
def test_flush_deferred_async_logging_fires_on_success():
"""
Happy path: with no exception, the production flush helper invokes the

16
tests/unit/rust_bridge/host/test_callbacks_legacy_python.py
View file

@ -101,6 +101,22 @@ def _kwargs_with_a_budget_reservation(reservation: dict) -> dict[str, object]:
return {**_OCR_KWARGS, "metadata": {"user_api_key_budget_reservation": reservation}}
def test_a_deferred_success_waits_for_the_proxy_release() -> None:
logger: Final = _supplied_logger()
released: Final[list[str]] = []
start: Final = datetime.datetime.now()
end: Final = start + datetime.timedelta(seconds=1)
legacy.defer_success(logger, lambda: released.append("released"), "response", start, end)
assert logger.deferred_success_result == ("response", start, end)
assert released == []
release: Final = logger.enqueue_deferred_logging
assert release is not None
release()
assert released == ["released"]
def test_setup_claims_the_budget_reservation_for_an_async_call() -> None:
reservation: Final = _budget_reservation()