From bc8937593ec7651f0b970648dccc69c523947742 Mon Sep 17 00:00:00 2001 From: narmstrong-berri Date: Fri, 6 Nov 2026 08:59:03 +0000 Subject: [PATCH] fix(messages): claim the async success log once across native fallback attempts Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_013FDwtkaNTifWAP4W7JZavE --- .../python_contract.json | 3 ++ .../callbacks-legacy-python/src/adapter.rs | 41 +++++++++++++++++++ .../callbacks-legacy-python/src/callbacks.rs | 8 ++++ .../callbacks-legacy-python/src/python.rs | 2 + .../src/test_support.rs | 5 +++ .../host/callbacks_legacy_python.py | 6 +++ 6 files changed, 65 insertions(+) diff --git a/litellm-rust/crates/callbacks-legacy-python/python_contract.json b/litellm-rust/crates/callbacks-legacy-python/python_contract.json index ab887bbc014..22210f14023 100644 --- a/litellm-rust/crates/callbacks-legacy-python/python_contract.json +++ b/litellm-rust/crates/callbacks-legacy-python/python_contract.json @@ -47,6 +47,9 @@ "defers_async_logging": [ "logger" ], + "claim_async_success_log": [ + "logger" + ], "defer_success": [ "logger", "pending" diff --git a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index 1bb84184e4c..414f1994a49 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs @@ -252,6 +252,7 @@ impl LegacyLogging { return pending().sync(py); } if !self.internal + && logger.claim_async_success_log(py)? && self .call .kwargs() @@ -1908,6 +1909,46 @@ assert hasattr(logger, '_native_pending_logging') == getattr(logger, '_defer_asy }); } + #[rstest] + #[case::deferred( + c"logger._defer_async_logging = True", + c"logger._native_pending_logging.release(True)", + &["sync_success_for_async_call", "sync_success_for_async_call", "async_success_handler", "enqueued"] + )] + #[case::immediate( + c"", + c"", + &["async_success_handler", "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 = local(&locals, "logger") + .call_method0("names") + .unwrap() + .extract() + .unwrap(); + assert_eq!(names, expected); + run( + py, + &locals, + c"assert [value for name, value in logger.calls if name == 'async_success_handler'] == [refused]", + ); + }); + } + #[rstest] fn dropped_observations_preserve_deferred_success_and_response_identity() { Python::initialize(); diff --git a/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs b/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs index 1951c9ebfb4..b97d33deb5a 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs @@ -59,6 +59,8 @@ pub trait LegacyCallbacks { fn defers_async_logging(&self, py: Python<'_>) -> bool; + fn claim_async_success_log(&self, py: Python<'_>) -> PyResult; + fn defer_success(&self, py: Python<'_>, pending: &Bound<'_, PyAny>) -> PyResult<()>; fn sync_success_for_async_call( @@ -203,6 +205,12 @@ impl LegacyCallbacks for PythonLogger { .unwrap_or(false) } + fn claim_async_success_log(&self, py: Python<'_>) -> PyResult { + Logging::ClaimAsyncSuccess + .call(py, (self.object(py),))? + .extract() + } + fn defer_success(&self, py: Python<'_>, pending: &Bound<'_, PyAny>) -> PyResult<()> { Logging::DeferSuccess.call(py, (self.object(py), pending))?; Ok(()) diff --git a/litellm-rust/crates/callbacks-legacy-python/src/python.rs b/litellm-rust/crates/callbacks-legacy-python/src/python.rs index 6d910d35f53..aa6cb80a283 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/python.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/python.rs @@ -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")] diff --git a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs index 8acfa7a898b..26fcd2898a9 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs @@ -72,6 +72,7 @@ FAKES = { llm_api_duration_ms=(end - start) * 1000 ), 'defers_async_logging': lambda logger: bool(getattr(logger, '_defer_async_logging', False)), + 'claim_async_success_log': lambda logger: logger.claim_async_success_log(), '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 @@ -156,6 +157,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) diff --git a/litellm/rust_bridge/host/callbacks_legacy_python.py b/litellm/rust_bridge/host/callbacks_legacy_python.py index 92590531cdc..c3f04a255b6 100644 --- a/litellm/rust_bridge/host/callbacks_legacy_python.py +++ b/litellm/rust_bridge/host/callbacks_legacy_python.py @@ -139,6 +139,8 @@ class LoggingSurface(Protocol): additional_args: dict[str, object] = ..., ) -> object: ... + def claim_async_success_log(self) -> bool: ... + 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,6 +271,10 @@ def defers_async_logging(logger: LoggingSurface) -> bool: return bool(getattr(logger, "defer_async_logging", False)) +def claim_async_success_log(logger: LoggingSurface) -> bool: + return logger.claim_async_success_log() + + def defer_success(logger: LoggingSurface, pending: object) -> None: setattr(logger, "_native_pending_logging", pending)