diff --git a/litellm-rust/crates/callbacks-legacy-python/python_contract.json b/litellm-rust/crates/callbacks-legacy-python/python_contract.json index ab887bbc014..1d4a2f7f7c2 100644 --- a/litellm-rust/crates/callbacks-legacy-python/python_contract.json +++ b/litellm-rust/crates/callbacks-legacy-python/python_contract.json @@ -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", diff --git a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index 1bb84184e4c..8d1946861aa 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs @@ -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 = 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 = 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"); }); } diff --git a/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs b/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs index 1951c9ebfb4..a428afa9ece 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs @@ -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; + + fn defer_success( + &self, + py: Python<'_>, + pending: &Bound<'_, PyAny>, + response: &Option>, + start: &Py, + end: &Option>, + ) -> 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 { + Logging::ClaimAsyncSuccess + .call(py, (self.object(py),))? + .extract() + } + + fn defer_success( + &self, + py: Python<'_>, + pending: &Bound<'_, PyAny>, + response: &Option>, + start: &Py, + end: &Option>, + ) -> 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") diff --git a/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs b/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs index 292648ba69c..ea72191ddb2 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs @@ -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, + 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::(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]) ", ); diff --git a/litellm-rust/crates/callbacks-legacy-python/src/python.rs b/litellm-rust/crates/callbacks-legacy-python/src/python.rs index 6d910d35f53..32b11988703 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")] @@ -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")] 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..9e69d0c30e2 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs @@ -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) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 8f73d47d96b..eb356491014 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -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 diff --git a/litellm/rust_bridge/host/callbacks_legacy_python.py b/litellm/rust_bridge/host/callbacks_legacy_python.py index 92590531cdc..92e5436fbb8 100644 --- a/litellm/rust_bridge/host/callbacks_legacy_python.py +++ b/litellm/rust_bridge/host/callbacks_legacy_python.py @@ -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: diff --git a/tests/test_litellm_rust/ocr/test_lifecycle.py b/tests/test_litellm_rust/ocr/test_lifecycle.py index 611316fb52c..c3846d894bb 100644 --- a/tests/test_litellm_rust/ocr/test_lifecycle.py +++ b/tests/test_litellm_rust/ocr/test_lifecycle.py @@ -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 diff --git a/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py index e9e0a585949..d2161fddeb0 100644 --- a/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py @@ -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 diff --git a/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py b/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py index 1ba0188e76d..29a6b6f56a8 100644 --- a/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py +++ b/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py @@ -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()