refactor(python-bridge): make deferred logging callable

litellm-rust/crates/python-bridge/src/lifecycle/mod.rs

@@ -538,11 +538,9 @@ struct PendingLogging {
 
 #[pymethods]
 impl PendingLogging {
-    fn release(slf: &Bound<'_, Self>, py: Python<'_>, success: bool) -> PyResult<()> {
+    fn __call__(slf: &Bound<'_, Self>, py: Python<'_>) -> PyResult<()> {
         let pending = slf.borrow_mut().pending.take();
-        if let Some(pending) = pending
-            && success
-        {
+        if let Some(pending) = pending {
             match pending.asynchronous(py) {
                 Err(error) if error.is_instance_of::<PyException>(py) => {
                     error.write_unraisable(py, Some(pending.logger.object(py)));
@@ -563,10 +561,14 @@ impl PendingLogging {
         Ok(())
     }
 
-    fn __clear__(slf: &Bound<'_, Self>) {
+    fn close(slf: &Bound<'_, Self>) {
         let pending = slf.borrow_mut().pending.take();
         drop(pending);
     }
+
+    fn __clear__(slf: &Bound<'_, Self>) {
+        Self::close(slf);
+    }
 }
 
 #[cfg(test)]
@@ -989,7 +991,7 @@ sys.unraisablehook = old_hook
     }
 
     #[test]
-    fn deferred_release_uses_release_context_and_allows_reentry_once() {
+    fn deferred_logging_uses_call_context_and_allows_reentry_once() {
         let _guard = PYTHON_GLOBALS
             .lock()
             .unwrap_or_else(|error| error.into_inner());
@@ -1022,7 +1024,7 @@ class Coroutine:
 class Worker:
     def ensure_initialized_and_enqueue(self, coroutine):
         observed.append(marker.get())
-        pending.release(True)
+        pending()
         coroutine.close()
 
 class Logger:
@@ -1060,10 +1062,72 @@ logger = Logger()
             py.run(
                 pyo3::ffi::c_str!(
                     r#"
-marker.set('release')
-pending.release(True)
-pending.release(True)
-assert observed == ['created', 'release', 'closed']
+marker.set('call')
+pending()
+pending()
+assert observed == ['created', 'call', 'closed']
+"#
+                ),
+                Some(&locals),
+                Some(&locals),
+            )
+            .unwrap();
+        });
+    }
+
+    #[test]
+    fn deferred_logging_close_is_reentry_safe_and_invalidates_aliases() {
+        Python::initialize();
+        Python::attach(|py| {
+            let locals = PyDict::new(py);
+            py.run(
+                pyo3::ffi::c_str!(
+                    r#"
+observed = []
+
+class Retained:
+    def __del__(self):
+        observed.append('finalized')
+        alias()
+
+class Logger:
+    def async_success_handler(self, *args):
+        observed.append('enqueued')
+
+logger = Logger()
+retained = Retained()
+"#
+                ),
+                Some(&locals),
+                Some(&locals),
+            )
+            .unwrap();
+            let pending = Py::new(
+                py,
+                PendingLogging {
+                    pending: Some(PendingSuccess {
+                        logger: locals
+                            .get_item("logger")
+                            .unwrap()
+                            .unwrap()
+                            .extract()
+                            .unwrap(),
+                        response: Some(locals.get_item("retained").unwrap().unwrap().unbind()),
+                        start: py.None(),
+                        end: None,
+                    }),
+                },
+            )
+            .unwrap();
+            locals.set_item("pending", &pending).unwrap();
+            locals.set_item("alias", &pending).unwrap();
+            locals.del_item("retained").unwrap();
+            py.run(
+                pyo3::ffi::c_str!(
+                    r#"
+pending.close()
+alias()
+assert observed == ['finalized']
 "#
                 ),
                 Some(&locals),

litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs

@@ -1,3 +1,4 @@
+use pyo3::gc::{PyTraverseError, PyVisit};
 use pyo3::prelude::*;
 use pyo3::types::{PyDict, PyTuple};
 
@@ -46,8 +47,19 @@ struct ProjectedOcrHost {
     secret_fields: Vec<&'static str>,
     azure_ad_token_provider: Option<PythonTokenProvider>,
     pre_call: Option<callbacks::OcrLoggingFields>,
-    body: Option<Py<PyDict>>,
-    headers: Option<Py<PyDict>>,
+    payload: Option<CapturedOcrPayload>,
+}
+
+struct CapturedOcrPayload {
+    body: Py<PyDict>,
+    headers: Py<PyDict>,
+}
+
+impl CapturedOcrPayload {
+    fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        visit.call(&self.body)?;
+        visit.call(&self.headers)
+    }
 }
 
 impl PythonOcrHost {
@@ -74,8 +86,7 @@ impl PythonOcrHost {
             secret_fields: projected.secret_fields,
             azure_ad_token_provider: projected.azure_ad_token_provider,
             pre_call: None,
-            body: None,
-            headers: None,
+            payload: None,
         });
         Ok(OcrHostResult::Request(Ok((
             Box::new(projected.request),
@@ -113,15 +124,23 @@ impl PythonOcrHost {
         for (name, value) in &request.headers {
             headers.set_item(name, value)?;
         }
-        logger.pre_ocr(py, request.api_key.as_deref(), &body, &headers, &request.url)?;
+        logger.pre_ocr(
+            py,
+            request.api_key.as_deref(),
+            &body,
+            &headers,
+            &request.url,
+        )?;
         request.body = from_py(&body)?;
         request.headers = headers
             .iter()
             .map(|(name, value)| Ok((name.extract::<String>()?, value.extract::<String>()?)))
             .collect::<PyResult<Vec<_>>>()?;
         let projected = self.projected_mut()?;
-        projected.body = Some(body.unbind());
-        projected.headers = Some(headers.unbind());
+        projected.payload = Some(CapturedOcrPayload {
+            body: body.unbind(),
+            headers: headers.unbind(),
+        });
         Ok(request)
     }
 
@@ -131,11 +150,12 @@ impl PythonOcrHost {
         request: OcrPostCallRequest,
     ) -> PyResult<OcrPostCallRequest> {
         let projected = self.projected()?;
+        let payload = projected.payload.as_ref();
         self.state.logger()?.post_ocr(
             py,
             &request.original_response,
-            projected.body.as_ref(),
-            projected.headers.as_ref(),
+            payload.map(|payload| &payload.body),
+            payload.map(|payload| &payload.headers),
         )?;
         Ok(request)
     }
@@ -223,15 +243,17 @@ impl PythonRoute for PythonOcrHost {
         self.projected = None;
     }
 
-    fn traverse(&self, visit: &pyo3::gc::PyVisit<'_>) -> Result<(), pyo3::gc::PyTraverseError> {
+    fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
         let Some(projected) = &self.projected else {
             return Ok(());
         };
         if let Some(provider) = &projected.azure_ad_token_provider {
             provider.traverse(visit)?;
         }
-        visit.call(&projected.body)?;
-        visit.call(&projected.headers)
+        if let Some(payload) = &projected.payload {
+            payload.traverse(visit)?;
+        }
+        Ok(())
     }
 }
 

litellm/proxy/common_request_processing.py

@@ -3212,7 +3212,10 @@ class ProxyBaseLLMRequestProcessing:
             pending: Final = getattr(logging_obj, "_native_pending_logging", None)
             if pending is not None:
                 logging_obj._native_pending_logging = None  # rebind-ok: consume the native OCR release signal once
-                pending.release(not exception_raised)
+                if exception_raised:
+                    pending.close()
+                else:
+                    pending()
         _enqueue_fn: Final = getattr(logging_obj, "_enqueue_deferred_logging", None)
         if _enqueue_fn is None:
             return

tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py

@@ -461,10 +461,16 @@ def test_native_pending_logging_is_released_only_for_ocr(call_type: str, excepti
     )
 
     if call_type in ("ocr", "aocr"):
-        pending.release.assert_called_once_with(not exception_raised)
+        if exception_raised:
+            pending.close.assert_called_once_with()
+            pending.assert_not_called()
+        else:
+            pending.assert_called_once_with()
+            pending.close.assert_not_called()
         assert logger._native_pending_logging is None
     else:
-        pending.release.assert_not_called()
+        pending.assert_not_called()
+        pending.close.assert_not_called()
         assert logger._native_pending_logging is pending
     if exception_raised:
         enqueue.assert_not_called()