wip

litellm-rust/crates/python-bridge/src/routes/retained_http.rs

@@ -70,15 +70,20 @@ fn prepare(boundary: &Bound<'_, PyAny>, asynchronous: bool) -> PyResult<Py<PyAny
     )
 }
 
-fn encode(boundary: &Bound<'_, PyAny>, roots: &Bound<'_, PyAny>) -> PyResult<Request> {
-    type ByteHeaders<'py> = Vec<(Bound<'py, PyBytes>, Bound<'py, PyBytes>)>;
-    let encoded = invoke(
+#[pyfunction]
+fn encode(boundary: &Bound<'_, PyAny>, roots: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
+    invoke(
         boundary,
         BoundaryMethod::Encode.resolve(false),
         PyTuple::new(boundary.py(), [roots])?,
-    )?;
+    )
+}
+
+fn request(execution: &Bound<'_, PyAny>) -> PyResult<Request> {
+    type ByteHeaders<'py> = Vec<(Bound<'py, PyBytes>, Bound<'py, PyBytes>)>;
+    let encoded = execution.call_method0("encode")?;
     let (url, headers, body, timeout_seconds): (String, ByteHeaders<'_>, Bound<'_, PyBytes>, f64) =
-        encoded.extract(boundary.py())?;
+        encoded.extract()?;
     Ok(Request {
         url,
         headers: headers
@@ -110,12 +115,9 @@ impl<'py> IntoPyObject<'py> for Wire {
 }
 
 #[pyfunction]
-fn send<'py>(
-    boundary: &Bound<'py, PyAny>,
-    roots: &Bound<'py, PyAny>,
-) -> PyResult<Bound<'py, PyAny>> {
-    let request = encode(boundary, roots)?;
-    pyo3_async_runtimes::tokio::future_into_py(boundary.py(), async move {
+fn send<'py>(execution: &Bound<'py, PyAny>) -> PyResult<Bound<'py, PyAny>> {
+    let request = request(execution)?;
+    pyo3_async_runtimes::tokio::future_into_py(execution.py(), async move {
         let response = run_async_value(buffered_post::send(request), core_error_to_pyerr).await?;
         Ok(Wire(response))
     })
@@ -136,37 +138,65 @@ fn finish(
 
 pub(crate) fn run_sync(boundary: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
     let py = boundary.py();
-    let roots = prepare(boundary, false)?;
-    let request = encode(boundary, roots.bind(py))?;
+    let execution = execution_module(py)?
+        .getattr("RetainedExecution")?
+        .call1((boundary,))?;
+    execution.call_method0("prepare")?;
+    let request = request(&execution)?;
     let response = run_sync_value(py, buffered_post::send(request), core_error_to_pyerr)?;
     let wire = Wire(response).into_pyobject(py)?;
-    finish(boundary, wire.as_any(), false)
+    execution.call_method1("finish", (wire,)).map(Bound::unbind)
 }
 
 pub(crate) fn run_async(boundary: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
-    static DRIVER: PyOnceLock<Py<PyAny>> = PyOnceLock::new();
-    let py = boundary.py();
-    let driver = DRIVER.get_or_try_init(py, || {
-        PyModule::from_code(
-            py,
-            c"async def drive(boundary, prepare, send, finish):
-    roots = await prepare(boundary, True)
-    wire = await send(boundary, roots)
-    return await finish(boundary, wire, True)
-",
-            c"retained_http_driver.py",
-            c"_retained_http_driver",
-        )?
-        .getattr("drive")
+    let module = execution_module(boundary.py())?;
+    let execution = module.getattr("RetainedExecution")?.call1((boundary,))?;
+    module
+        .getattr("drive")?
+        .call1((execution,))
         .map(Bound::unbind)
-    })?;
-    driver.call1(
-        py,
-        (
-            boundary,
-            wrap_pyfunction!(prepare, py)?,
-            wrap_pyfunction!(send, py)?,
-            wrap_pyfunction!(finish, py)?,
-        ),
-    )
+}
+
+fn execution_module(py: Python<'_>) -> PyResult<Bound<'_, PyModule>> {
+    static MODULE: PyOnceLock<Py<PyModule>> = PyOnceLock::new();
+    let module = MODULE.get_or_try_init(py, || {
+        let module = PyModule::from_code(
+            py,
+            c"class RetainedExecution:
+    __slots__ = ('binding', 'roots')
+
+    def __init__(self, binding):
+        self.binding = binding
+        self.roots = None
+
+    def prepare(self):
+        self.roots = _prepare(self.binding, False)
+
+    async def aprepare(self):
+        self.roots = await _prepare(self.binding, True)
+
+    def encode(self):
+        return _encode(self.binding, self.roots)
+
+    def finish(self, wire):
+        return _finish(self.binding, wire, False)
+
+    async def afinish(self, wire):
+        return await _finish(self.binding, wire, True)
+
+async def drive(execution):
+    await execution.aprepare()
+    wire = await _send(execution)
+    return await execution.afinish(wire)
+",
+            c"retained_execution.py",
+            c"_retained_execution",
+        )?;
+        module.add("_prepare", wrap_pyfunction!(prepare, &module)?)?;
+        module.add("_encode", wrap_pyfunction!(encode, &module)?)?;
+        module.add("_finish", wrap_pyfunction!(finish, &module)?)?;
+        module.add("_send", wrap_pyfunction!(send, &module)?)?;
+        Ok::<_, PyErr>(module.unbind())
+    })?;
+    Ok(module.bind(py).clone())
 }

litellm-rust/crates/python-bridge/tests/fixtures/retained_http_contract.py

@@ -1,21 +1,25 @@
 import asyncio
 import contextvars
 import gc
+import http.client
 import http.server
 import inspect
 import threading
 import weakref
+from copy import deepcopy
+from typing import NamedTuple
 
 native = globals()["native"]
 started = threading.Event()
 release = threading.Event()
 server = None
+server_thread = None
 url = None
 context = contextvars.ContextVar("proof", default="unset")
 
 
 def start_server():
-    global server, url
+    global server, server_thread, url
     requests = []
 
     class Handler(http.server.BaseHTTPRequestHandler):
@@ -39,8 +43,8 @@ def start_server():
                 pass
 
     server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), Handler)
-    thread = threading.Thread(target=server.serve_forever, daemon=True)
-    thread.start()
+    server_thread = threading.Thread(target=server.serve_forever, daemon=True)
+    server_thread.start()
     url = "http://127.0.0.1:%s/ocr" % server.server_port
     return requests
 
@@ -48,6 +52,14 @@ def start_server():
 requests = start_server()
 
 
+def stop_server():
+    release.set()
+    server.shutdown()
+    server.server_close()
+    server_thread.join(5)
+    assert not server_thread.is_alive()
+
+
 class Graph(dict):
     pass
 
@@ -222,6 +234,28 @@ def run_ownership_contract():
             boundary.error.__traceback__ = None
             if phase != "prepare":
                 collected(boundary)
+        boundary = Boundary(hold=True)
+        observations = []
+
+        def observe_blocked_transport():
+            received = started.wait(2)
+            gc.collect()
+            observations.append((received, tuple(ref() is not None for ref in boundary.refs)))
+            release.set()
+
+        observer = threading.Thread(target=observe_blocked_transport)
+        observer.start()
+        try:
+            assert native.ocr_retained(boundary) is boundary.result
+        finally:
+            release.set()
+            observer.join(5)
+        assert not observer.is_alive()
+        assert observations == [(True, (True, True))]
+        assert boundary.events == ["prepare", "encode", "finish"]
+        collected(boundary)
+        started.clear()
+        release.clear()
         asyncio.run(asyncio.wait_for(exercise(), 15))
         assert requests
         assert all(
@@ -229,9 +263,7 @@ def run_ownership_contract():
             for path, headers, body in requests
         )
     finally:
-        release.set()
-        server.shutdown()
-        server.server_close()
+        stop_server()
 
 
 class TimeoutBoundary(Boundary):
@@ -248,7 +280,7 @@ class TimeoutBoundary(Boundary):
 
     def encode(self, roots):
         headers, target, body, files = roots
-        return (target, [], b"\x00", self.timeout)
+        return (target, [], b"hold", self.timeout)
 
     def finish(self, wire):
         raise AssertionError("client-side failure must not reach finish")
@@ -256,14 +288,322 @@ class TimeoutBoundary(Boundary):
 
 def run_error_contract():
     cases = [
-        ("timeout", "http://10.255.255.1:9/", 0.05),
+        ("timeout", url, 0.05),
         ("refused", "http://127.0.0.1:1/", 1.0),
     ]
-    for name, target, timeout in cases:
-        boundary = TimeoutBoundary(timeout=timeout, url=target)
-        try:
-            native.ocr_retained(boundary)
-        except RuntimeError:
-            pass
+    try:
+        for name, target, timeout in cases:
+            boundary = TimeoutBoundary(timeout=timeout, url=target)
+            try:
+                native.ocr_retained(boundary)
+            except RuntimeError:
+                pass
+            else:
+                raise AssertionError(f"{name} did not surface as RuntimeError")
+    finally:
+        stop_server()
+
+
+class FrozenObservation(NamedTuple):
+    root_tuple_identity: bool
+    roots_identity: tuple[bool, bool]
+    document_is_caller: bool
+    document_alias: bool
+    logging_envelope_identity: bool
+    logging_roots_identity: tuple[bool, bool]
+    logging_fields: tuple[str, str]
+    retained_fields: tuple[str, str]
+    caller_value: str
+    caught_error_observation: tuple[str, ...]
+    callback_count: int
+    phases: tuple[str, ...]
+    request: tuple[str, tuple[str, ...], bytes]
+    response: tuple[int, tuple[bytes, ...], bytes]
+
+
+class ContractCallback:
+    def __init__(self, scenario, caller):
+        self.scenario = scenario
+        self.caller = caller
+        self.mutate_caller = lambda: caller.__setitem__("value", "closure")
+        self.calls = 0
+        self.error = LookupError("caught callback failure")
+        self.writer = None
+        self.write = threading.Event()
+        self.written = threading.Event()
+
+    def __call__(self, view):
+        self.calls += 1
+        self.view = view
+        self.headers, self.body = view["headers"], view["body"]
+        if self.scenario == "caller_closure":
+            self.mutate_caller()
+        elif self.scenario == "field_replace":
+            document = Graph(value="replacement")
+            view["headers"] = Graph({"X-Proof": "replacement"})
+            view["body"] = Graph(document=document, alias=document)
+        elif self.scenario == "delayed_after_prepare_before_encode_read":
+
+            def writer():
+                if self.write.wait(5):
+                    self.mutate("delayed")
+                    self.written.set()
+
+            self.writer = threading.Thread(target=writer)
+            self.writer.start()
         else:
-            raise AssertionError(f"{name} did not surface as RuntimeError")
+            assert self.scenario in ("retain_mutate", "mutate_then_caught_error_observe")
+            self.mutate("mutated" if self.scenario == "retain_mutate" else "caught")
+            if self.scenario == "mutate_then_caught_error_observe":
+                raise self.error
+
+    def mutate(self, value):
+        self.headers["X-Proof"] = value
+        self.body["document"]["value"] = value
+
+    def before_encode_read(self):
+        if self.writer is not None:
+            self.write.set()
+            assert self.written.wait(2), "scheduled mutation did not complete before encoder read"
+            self.writer.join(2)
+            assert not self.writer.is_alive()
+
+    def close(self):
+        self.write.set()
+        if self.writer is not None:
+            self.writer.join(5)
+            assert not self.writer.is_alive()
+        self.error.__traceback__ = None
+
+
+class TransportContractBinding:
+    def __init__(self, caller, callback):
+        self.caller = caller
+        self.callback = callback
+        self.phases = []
+        self.caught = ()
+
+    def prepare(self):
+        self.phases.append("prepare")
+        headers = Graph({"X-Proof": "original"})
+        body = Graph(document=self.caller, alias=self.caller)
+        self.view = {"headers": headers, "body": body}
+        self.prepared_roots = (headers, url, body, None)
+        try:
+            self.callback(self.view)
+        except LookupError as error:
+            if error is not self.callback.error:
+                raise
+            self.caught = (str(error), headers["X-Proof"], body["document"]["value"])
+        return self.prepared_roots
+
+    async def aprepare(self):
+        await asyncio.sleep(0)
+        return self.prepare()
+
+    def encode(self, roots):
+        self.callback.before_encode_read()
+        self.phases.append("encode")
+        headers, target, body, files = roots
+        assert files is None
+        self.root_tuple_identity = roots is self.prepared_roots
+        self.roots_identity = (headers is self.callback.headers, body is self.callback.body)
+        self.document_is_caller = body["document"] is self.callback.caller
+        self.document_alias = body["document"] is body["alias"]
+        return (
+            target,
+            [(b"X-Proof", headers["X-Proof"].encode())],
+            b"\x00" + body["document"]["value"].encode() + b"\xff",
+            3.0,
+        )
+
+    def finish(self, wire):
+        self.phases.append("finish")
+        status, headers, content = wire
+        return (status, tuple(value for key, value in headers if key == b"x-reply"), content)
+
+    async def afinish(self, wire):
+        await asyncio.sleep(0)
+        return self.finish(wire)
+
+
+class BindingVariant:
+    def __init__(self, binding, variant):
+        self.binding = binding
+        self.variant = variant
+
+    def before_prepare(self):
+        if self.variant == "copy_caller_inputs_before_prepare":
+            self.binding.caller = deepcopy(self.binding.caller)
+
+    def at_prepare_return(self, roots):
+        if self.variant == "copy_roots_after_prepare":
+            return deepcopy(roots)
+        if self.variant == "reconstruct_root_tuple":
+            return tuple(list(roots))
+        if self.variant == "reconstruct_logging_envelope":
+            self.binding.view = dict(self.binding.view)
+        return roots
+
+    def prepare(self):
+        self.before_prepare()
+        return self.at_prepare_return(self.binding.prepare())
+
+    async def aprepare(self):
+        self.before_prepare()
+        return self.at_prepare_return(await self.binding.aprepare())
+
+    def encode(self, roots):
+        if self.variant == "encode_logging_replacements":
+            headers, target, body, files = roots
+            return self.binding.encode((self.binding.view["headers"], target, self.binding.view["body"], files))
+        return self.binding.encode(roots)
+
+    def finish(self, wire):
+        return self.binding.finish(wire)
+
+    async def afinish(self, wire):
+        return await self.binding.afinish(wire)
+
+
+def reference_post(encoded):
+    target, headers, body, timeout = encoded
+    assert target == url
+    connection = http.client.HTTPConnection("127.0.0.1", server.server_port, timeout=timeout)
+    try:
+        connection.putrequest("POST", "/ocr")
+        for key, value in headers:
+            connection.putheader(key.decode("ascii"), value.decode("ascii"))
+        connection.putheader("Content-Length", str(len(body)))
+        connection.endheaders(body)
+        response = connection.getresponse()
+        return (
+            response.status,
+            [(key.lower().encode(), value.encode()) for key, value in response.getheaders()],
+            response.read(),
+        )
+    finally:
+        connection.close()
+
+
+async def reference_async(binding):
+    roots = await binding.aprepare()
+    wire = await asyncio.to_thread(reference_post, binding.encode(roots))
+    return await binding.afinish(wire)
+
+
+def original_observation(value, *, header=None, logging=None, caught=()):
+    fields = (header or value, value)
+    return FrozenObservation(
+        True,
+        (True, True),
+        True,
+        True,
+        True,
+        (False, False) if logging else (True, True),
+        logging or fields,
+        fields,
+        value,
+        caught,
+        1,
+        ("prepare", "encode", "finish"),
+        ("/ocr", (fields[0],), b"\x00" + value.encode() + b"\xff"),
+        (429, (b"one", b"two"), b"\x00\xffR"),
+    )
+
+
+ORIGINAL_OBSERVATIONS = {
+    "retain_mutate": original_observation("mutated"),
+    "caller_closure": original_observation("closure", header="original"),
+    "delayed_after_prepare_before_encode_read": original_observation("delayed"),
+    "field_replace": original_observation("original", logging=("replacement", "replacement")),
+    "mutate_then_caught_error_observe": original_observation(
+        "caught", caught=("caught callback failure", "caught", "caught")
+    ),
+}
+
+VARIANT_OBSERVATIONS = {
+    ("retain_mutate", "copy_caller_inputs_before_prepare"): {
+        "caller_value": "original",
+        "document_is_caller": False,
+    },
+    ("caller_closure", "copy_caller_inputs_before_prepare"): {
+        "document_is_caller": False,
+        "logging_fields": ("original", "original"),
+        "retained_fields": ("original", "original"),
+        "request": ("/ocr", ("original",), b"\x00original\xff"),
+    },
+    ("retain_mutate", "copy_roots_after_prepare"): {
+        "root_tuple_identity": False,
+        "roots_identity": (False, False),
+        "document_is_caller": False,
+    },
+    ("delayed_after_prepare_before_encode_read", "copy_roots_after_prepare"): {
+        "root_tuple_identity": False,
+        "roots_identity": (False, False),
+        "document_is_caller": False,
+        "request": ("/ocr", ("original",), b"\x00original\xff"),
+    },
+    ("mutate_then_caught_error_observe", "copy_roots_after_prepare"): {
+        "root_tuple_identity": False,
+        "roots_identity": (False, False),
+        "document_is_caller": False,
+    },
+    ("field_replace", "encode_logging_replacements"): {
+        "root_tuple_identity": False,
+        "roots_identity": (False, False),
+        "document_is_caller": False,
+        "request": ("/ocr", ("replacement",), b"\x00replacement\xff"),
+    },
+    ("retain_mutate", "reconstruct_logging_envelope"): {"logging_envelope_identity": False},
+    ("field_replace", "reconstruct_logging_envelope"): {"logging_envelope_identity": False},
+    ("retain_mutate", "reconstruct_root_tuple"): {"root_tuple_identity": False},
+    ("delayed_after_prepare_before_encode_read", "reconstruct_root_tuple"): {"root_tuple_identity": False},
+    ("field_replace", "reconstruct_root_tuple"): {"root_tuple_identity": False},
+}
+
+
+def run_comparison_contract(scenario, variant):
+    try:
+        expected = ORIGINAL_OBSERVATIONS[scenario]
+        if variant != "original":
+            expected = expected._replace(**VARIANT_OBSERVATIONS[(scenario, variant)])
+        for mode in ("reference-sync", "reference-async", "native-sync", "native-async"):
+            caller = Graph(value="original")
+            callback = ContractCallback(scenario, caller)
+            binding = TransportContractBinding(caller, callback)
+            selected = binding if variant == "original" else BindingVariant(binding, variant)
+            offset = len(requests)
+            try:
+                if mode == "reference-sync":
+                    roots = selected.prepare()
+                    response = selected.finish(reference_post(selected.encode(roots)))
+                elif mode == "reference-async":
+                    response = asyncio.run(reference_async(selected))
+                elif mode == "native-sync":
+                    response = native.ocr_retained(selected)
+                else:
+                    response = asyncio.run(native.aocr_retained(selected))
+                assert len(requests) == offset + 1, (scenario, variant, mode, requests[offset:])
+                path, headers, body = requests[offset]
+                observed = FrozenObservation(
+                    binding.root_tuple_identity,
+                    binding.roots_identity,
+                    binding.document_is_caller,
+                    binding.document_alias,
+                    binding.view is callback.view,
+                    (binding.view["headers"] is callback.headers, binding.view["body"] is callback.body),
+                    (binding.view["headers"]["X-Proof"], binding.view["body"]["document"]["value"]),
+                    (callback.headers["X-Proof"], callback.body["document"]["value"]),
+                    caller["value"],
+                    binding.caught,
+                    callback.calls,
+                    tuple(binding.phases),
+                    (path, tuple(headers), body),
+                    response,
+                )
+                assert observed == expected, (scenario, variant, mode, observed, expected)
+            finally:
+                callback.close()
+    finally:
+        stop_server()

litellm-rust/crates/python-bridge/tests/retained_http.rs

@@ -1,4 +1,6 @@
 use pyo3::prelude::*;
+use rstest::rstest;
+use serial_test::serial;
 
 #[path = "support/mod.rs"]
 mod support;
@@ -6,6 +8,7 @@ mod support;
 use support::native::{native_globals, run_fixture};
 
 #[test]
+#[serial(python_interpreter)]
 fn retained_routes_preserve_callbacks_context_wire_and_ownership() -> PyResult<()> {
     Python::initialize();
     Python::attach(|py| {
@@ -28,6 +31,7 @@ fn retained_routes_preserve_callbacks_context_wire_and_ownership() -> PyResult<(
 }
 
 #[test]
+#[serial(python_interpreter)]
 fn retained_routes_surface_transport_failures_as_runtime_error() -> PyResult<()> {
     Python::initialize();
     Python::attach(|py| {
@@ -45,3 +49,57 @@ fn retained_routes_surface_transport_failures_as_runtime_error() -> PyResult<()>
         Ok(())
     })
 }
+
+#[rstest]
+#[case::original_retain_mutate("retain_mutate", "original")]
+#[case::original_caller_closure("caller_closure", "original")]
+#[case::original_delayed_after_prepare_before_encode_read(
+    "delayed_after_prepare_before_encode_read",
+    "original"
+)]
+#[case::original_field_replace("field_replace", "original")]
+#[case::original_mutate_then_caught_error_observe("mutate_then_caught_error_observe", "original")]
+#[case::copy_caller_inputs_retain_mutate("retain_mutate", "copy_caller_inputs_before_prepare")]
+#[case::copy_caller_inputs_caller_closure("caller_closure", "copy_caller_inputs_before_prepare")]
+#[case::copy_roots_retain_mutate("retain_mutate", "copy_roots_after_prepare")]
+#[case::copy_roots_delayed_after_prepare_before_encode_read(
+    "delayed_after_prepare_before_encode_read",
+    "copy_roots_after_prepare"
+)]
+#[case::copy_roots_mutate_then_caught_error_observe(
+    "mutate_then_caught_error_observe",
+    "copy_roots_after_prepare"
+)]
+#[case::encode_logging_replacements_field_replace("field_replace", "encode_logging_replacements")]
+#[case::reconstruct_logging_envelope_retain_mutate("retain_mutate", "reconstruct_logging_envelope")]
+#[case::reconstruct_logging_envelope_field_replace("field_replace", "reconstruct_logging_envelope")]
+#[case::reconstruct_root_tuple_retain_mutate("retain_mutate", "reconstruct_root_tuple")]
+#[case::reconstruct_root_tuple_delayed_after_prepare_before_encode_read(
+    "delayed_after_prepare_before_encode_read",
+    "reconstruct_root_tuple"
+)]
+#[case::reconstruct_root_tuple_field_replace("field_replace", "reconstruct_root_tuple")]
+#[serial(python_interpreter)]
+fn retained_production_transport_comparison_matrix(
+    #[case] scenario: &str,
+    #[case] variant: &str,
+) -> PyResult<()> {
+    Python::initialize();
+    Python::attach(|py| {
+        let globals = native_globals(py)?;
+        run_fixture(
+            py,
+            &globals,
+            include_str!("fixtures/retained_http_contract.py"),
+            concat!(
+                env!("CARGO_MANIFEST_DIR"),
+                "/tests/fixtures/retained_http_contract.py"
+            ),
+        )?;
+        globals
+            .get_item("run_comparison_contract")?
+            .unwrap()
+            .call1((scenario, variant))?;
+        Ok(())
+    })
+}

litellm-rust/crates/python-interop/tests/callback_lifecycle.rs

@@ -59,18 +59,7 @@ fn lifecycle_contract(
     #[case] scenario: &str,
     #[values(false, true)] retained: bool,
 ) -> PyResult<()> {
-    Python::attach(|py| {
-        scenario_scope
-            .bind(py)
-            .get_item("run_scenario")?
-            .unwrap()
-            .call1((
-                scenario,
-                retained,
-                scenario_scope.bind(py).get_item("factory")?.unwrap(),
-            ))?;
-        Ok(())
-    })
+    run_scenario_fixture(scenario_scope, scenario, retained, None)
 }
 
 #[rstest]
@@ -89,28 +78,23 @@ fn component_contract(
     #[case] scenario: &str,
     #[values(false, true)] retained: bool,
 ) -> PyResult<()> {
-    Python::attach(|py| {
-        let globals = scenario_scope.bind(py);
-        run_fixture(
-            py,
-            globals,
+    run_scenario_fixture(
+        scenario_scope,
+        scenario,
+        retained,
+        Some((
             include_str!("fixtures/callback_components.py"),
             concat!(
                 env!("CARGO_MANIFEST_DIR"),
                 "/tests/fixtures/callback_components.py"
             ),
-        )?;
-        globals.get_item("run_scenario")?.unwrap().call1((
-            scenario,
-            retained,
-            scenario_scope.bind(py).get_item("factory")?.unwrap(),
-        ))?;
-        Ok(())
-    })
+        )),
+    )
 }
 
 #[rstest]
 #[case::real_logging_queue_chain("real_logging_queue_chain")]
+#[case::real_logging_queue_copy_control("real_logging_queue_copy_control")]
 #[case::real_crowdstrike_translator_identity("real_crowdstrike_translator_identity")]
 #[case::real_rubrik_block_lifecycle("real_rubrik_block_lifecycle")]
 #[case::real_parallel_guardrail_snapshots("real_parallel_guardrail_snapshots")]
@@ -122,17 +106,31 @@ fn integration_contract(
     #[case] scenario: &str,
     #[values(false, true)] retained: bool,
 ) -> PyResult<()> {
-    Python::attach(|py| {
-        let globals = scenario_scope.bind(py);
-        run_fixture(
-            py,
-            globals,
+    run_scenario_fixture(
+        scenario_scope,
+        scenario,
+        retained,
+        Some((
             include_str!("fixtures/callback_integrations.py"),
             concat!(
                 env!("CARGO_MANIFEST_DIR"),
                 "/tests/fixtures/callback_integrations.py"
             ),
-        )?;
+        )),
+    )
+}
+
+fn run_scenario_fixture(
+    scenario_scope: Py<PyDict>,
+    scenario: &str,
+    retained: bool,
+    fixture: Option<(&str, &str)>,
+) -> PyResult<()> {
+    Python::attach(|py| {
+        let globals = scenario_scope.bind(py);
+        if let Some((source, filename)) = fixture {
+            run_fixture(py, globals, source, filename)?;
+        }
         globals.get_item("run_scenario")?.unwrap().call1((
             scenario,
             retained,
@@ -142,6 +140,80 @@ fn integration_contract(
     })
 }
 
+#[rstest]
+#[case::original_arguments("argument_identity", "identity")]
+#[case::envelope_arguments("argument_identity", "envelope")]
+#[case::shallow_arguments("argument_identity", "shallow_payload")]
+#[case::copied_graph("argument_identity", "deep_graph")]
+#[case::independent_copies("argument_identity", "deep_separate")]
+#[case::original_read_timing("mutation_timing", "identity")]
+#[case::envelope_read_timing("mutation_timing", "envelope")]
+#[case::shallow_read_timing("mutation_timing", "shallow_payload")]
+#[case::deep_read_timing("mutation_timing", "deep_graph")]
+#[case::independent_read_timing("mutation_timing", "deep_separate")]
+#[case::original_result("result_identity", "identity")]
+#[case::passthrough_result("result_identity", "result_passthrough")]
+#[case::shallow_result("result_identity", "result_shallow")]
+#[case::deep_result("result_identity", "result_deep")]
+#[case::retained_lifetime("deferred_lifetime", "identity")]
+#[case::expired_borrow("deferred_lifetime", "weak")]
+#[case::prepared_ownership("deferred_lifetime", "missing_handoff")]
+#[case::externally_owned_retained("borrowed_lifetime", "identity")]
+#[case::externally_owned_borrow("borrowed_lifetime", "weak")]
+#[case::original_coroutine("direct_coroutine", "identity")]
+#[case::passthrough_coroutine("direct_coroutine", "result_passthrough")]
+#[serial(python_interpreter)]
+fn control_contract(
+    scenario_scope: Py<PyDict>,
+    #[case] witness: &str,
+    #[case] control: &str,
+    #[values(false, true)] retained: bool,
+    #[values(false, true)] awaited: bool,
+) -> PyResult<()> {
+    run_control_fixture(scenario_scope, witness, control, retained, awaited)
+}
+
+#[rstest]
+#[case::retained("identity")]
+#[case::missing_handoff("missing_handoff")]
+#[serial(python_interpreter)]
+fn pending_handoff_control(
+    scenario_scope: Py<PyDict>,
+    #[case] control: &str,
+    #[values(false, true)] retained: bool,
+) -> PyResult<()> {
+    run_control_fixture(scenario_scope, "pending_handoff", control, retained, true)
+}
+
+fn run_control_fixture(
+    scenario_scope: Py<PyDict>,
+    witness: &str,
+    control: &str,
+    retained: bool,
+    awaited: bool,
+) -> PyResult<()> {
+    Python::attach(|py| {
+        let globals = scenario_scope.bind(py);
+        run_fixture(
+            py,
+            globals,
+            include_str!("fixtures/callback_controls.py"),
+            concat!(
+                env!("CARGO_MANIFEST_DIR"),
+                "/tests/fixtures/callback_controls.py"
+            ),
+        )?;
+        globals.get_item("run_control")?.unwrap().call1((
+            witness,
+            control,
+            retained,
+            awaited,
+            globals.get_item("factory")?.unwrap(),
+        ))?;
+        Ok(())
+    })
+}
+
 #[rstest]
 #[parallel(python_interpreter)]
 fn detached_release(initialized_python: &InitializedPython) -> PyResult<()> {

litellm-rust/crates/python-interop/tests/fixtures/callback_controls.py

@@ -0,0 +1,477 @@
+import asyncio
+import copy
+import gc
+import inspect
+import weakref
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass
+from typing import Protocol, cast
+
+
+class PreparedInvocation(Protocol):
+    def invoke(self) -> object: ...
+
+    def close(self) -> None: ...
+
+
+class CallFactory(Protocol):
+    def prepare(
+        self,
+        callable: Callable[..., object],
+        positional: tuple[object, ...],
+        keywords: dict[str, object] | None = None,
+        awaited: bool = False,
+    ) -> PreparedInvocation: ...
+
+
+class LiveCallFactory(CallFactory, Protocol):
+    @property
+    def live(self) -> int: ...
+
+
+Arguments = tuple[tuple[object, ...], dict[str, object] | None]
+ArgumentTransform = Callable[[tuple[object, ...], dict[str, object] | None], Arguments]
+
+
+def reconstruct_envelope(positional: tuple[object, ...], keywords: dict[str, object] | None) -> Arguments:
+    return tuple(value for value in positional), None if keywords is None else dict(keywords)
+
+
+def shallow_selected_payload(positional: tuple[object, ...], keywords: dict[str, object] | None) -> Arguments:
+    return (copy.copy(positional[0]), *positional[1:]), keywords
+
+
+def deepcopy_graph(positional: tuple[object, ...], keywords: dict[str, object] | None) -> Arguments:
+    return copy.deepcopy((positional, keywords))
+
+
+def deepcopy_separate(positional: tuple[object, ...], keywords: dict[str, object] | None) -> Arguments:
+    return copy.deepcopy(positional), copy.deepcopy(keywords)
+
+
+def unchanged_result(value: object) -> object:
+    return value
+
+
+@dataclass(frozen=True, slots=True)
+class ArgumentTransformFactory:
+    inner: CallFactory
+    transform: ArgumentTransform
+
+    def prepare(
+        self,
+        callable: Callable[..., object],
+        positional: tuple[object, ...],
+        keywords: dict[str, object] | None = None,
+        awaited: bool = False,
+    ) -> PreparedInvocation:
+        args, kwargs = self.transform(positional, keywords)
+        return self.inner.prepare(callable, args, kwargs, awaited=awaited)
+
+
+@dataclass(frozen=True, slots=True)
+class ResultTransformInvocation:
+    inner: PreparedInvocation
+    transform: Callable[[object], object]
+    awaited: bool
+
+    def invoke(self) -> object:
+        result = self.inner.invoke()
+        transform = self.transform
+        if not self.awaited:
+            return transform(result)
+
+        async def run() -> object:
+            return transform(await cast(Awaitable[object], result))
+
+        return run()
+
+    def close(self) -> None:
+        self.inner.close()
+
+
+@dataclass(frozen=True, slots=True)
+class ResultTransformFactory:
+    inner: CallFactory
+    transform: Callable[[object], object]
+
+    def prepare(
+        self,
+        callable: Callable[..., object],
+        positional: tuple[object, ...],
+        keywords: dict[str, object] | None = None,
+        awaited: bool = False,
+    ) -> PreparedInvocation:
+        return ResultTransformInvocation(
+            self.inner.prepare(callable, positional, keywords, awaited=awaited), self.transform, awaited
+        )
+
+
+@dataclass(frozen=True, slots=True)
+class ExpiredBorrow:
+    edges: tuple[str, ...]
+
+
+@dataclass(frozen=True, slots=True)
+class CheckedWeakInvocation:
+    callback: weakref.ReferenceType[Callable[..., object]]
+    positional: tuple[weakref.ReferenceType[object], ...]
+    keywords: tuple[tuple[str, weakref.ReferenceType[object]], ...]
+    awaited: bool
+
+    def resolve(self) -> tuple[Callable[..., object], tuple[object, ...], dict[str, object]] | ExpiredBorrow:
+        callback = self.callback()
+        positional = tuple(reference() for reference in self.positional)
+        keywords = {name: reference() for name, reference in self.keywords}
+        expired = (
+            *(("callable",) if callback is None else ()),
+            *(f"positional:{index}" for index, value in enumerate(positional) if value is None),
+            *(f"keyword:{name}" for name, value in keywords.items() if value is None),
+        )
+        if expired:
+            return ExpiredBorrow(expired)
+        assert callback is not None
+        return callback, positional, keywords
+
+    def invoke(self) -> object:
+        if not self.awaited:
+            resolved = self.resolve()
+            if isinstance(resolved, ExpiredBorrow):
+                return resolved
+            callback, positional, keywords = resolved
+            return callback(*positional, **keywords)
+
+        async def run() -> object:
+            resolved = self.resolve()
+            if isinstance(resolved, ExpiredBorrow):
+                return resolved
+            callback, positional, keywords = resolved
+            return await cast(Awaitable[object], callback(*positional, **keywords))
+
+        return run()
+
+    def close(self) -> None:
+        pass
+
+
+class CheckedWeakFactory:
+    def prepare(
+        self,
+        callable: Callable[..., object],
+        positional: tuple[object, ...],
+        keywords: dict[str, object] | None = None,
+        awaited: bool = False,
+    ) -> PreparedInvocation:
+        return CheckedWeakInvocation(
+            weakref.ref(callable),
+            tuple(weakref.ref(value) for value in positional),
+            tuple((name, weakref.ref(value)) for name, value in (keywords or {}).items()),
+            awaited,
+        )
+
+
+@dataclass(frozen=True, slots=True)
+class MissingHandoffInvocation:
+    inner: PreparedInvocation
+    borrowed: PreparedInvocation
+
+    def invoke(self) -> object:
+        return self.borrowed.invoke()
+
+    def close(self) -> None:
+        try:
+            self.inner.close()
+        finally:
+            self.borrowed.close()
+
+
+@dataclass(frozen=True, slots=True)
+class MissingHandoffFactory:
+    inner: CallFactory
+
+    def prepare(
+        self,
+        callable: Callable[..., object],
+        positional: tuple[object, ...],
+        keywords: dict[str, object] | None = None,
+        awaited: bool = False,
+    ) -> PreparedInvocation:
+        borrowed = CheckedWeakFactory().prepare(callable, positional, keywords, awaited=awaited)
+        return MissingHandoffInvocation(self.inner.prepare(callable, positional, keywords, awaited=awaited), borrowed)
+
+
+def control_factory(control: str, inner: CallFactory) -> CallFactory:
+    if control == "identity":
+        return inner
+    if control == "weak":
+        return CheckedWeakFactory()
+    if control == "missing_handoff":
+        return MissingHandoffFactory(inner)
+    if control in ("result_passthrough", "result_shallow", "result_deep"):
+        return ResultTransformFactory(
+            inner,
+            {"result_passthrough": unchanged_result, "result_shallow": copy.copy, "result_deep": copy.deepcopy}[
+                control
+            ],
+        )
+    return ArgumentTransformFactory(
+        inner,
+        {
+            "envelope": reconstruct_envelope,
+            "shallow_payload": shallow_selected_payload,
+            "deep_graph": deepcopy_graph,
+            "deep_separate": deepcopy_separate,
+        }[control],
+    )
+
+
+@dataclass
+class ControlNode:
+    stage: int = 0
+
+
+@dataclass
+class ControlPayload:
+    nested: ControlNode
+    stage: int = 0
+
+
+@dataclass(frozen=True, slots=True)
+class IdentityObservation:
+    root: bool
+    nested: bool
+    cross_argument: bool
+
+
+async def argument_identity(owners: CallFactory, awaited: bool) -> IdentityObservation:
+    nested = ControlNode()
+    original = ControlPayload(nested)
+
+    def observe(value: ControlPayload, *, alias: ControlNode) -> IdentityObservation:
+        return IdentityObservation(value is original, value.nested is nested, value.nested is alias)
+
+    async def observe_async(value: ControlPayload, *, alias: ControlNode) -> IdentityObservation:
+        return observe(value, alias=alias)
+
+    owner = owners.prepare(observe_async if awaited else observe, (original,), {"alias": nested}, awaited=awaited)
+    try:
+        pending = owner.invoke()
+        return await pending if awaited else pending
+    finally:
+        owner.close()
+
+
+@dataclass(frozen=True, slots=True)
+class TimingObservation:
+    root: int
+    nested: int
+    alias: int
+
+
+async def mutation_timing(owners: CallFactory, awaited: bool) -> TimingObservation:
+    nested = ControlNode()
+    original = ControlPayload(nested)
+
+    def observe(value: ControlPayload, *, alias: ControlNode) -> TimingObservation:
+        return TimingObservation(value.stage, value.nested.stage, alias.stage)
+
+    async def observe_async(value: ControlPayload, *, alias: ControlNode) -> TimingObservation:
+        return observe(value, alias=alias)
+
+    owner = owners.prepare(observe_async if awaited else observe, (original,), {"alias": nested}, awaited=awaited)
+    try:
+        original.stage = nested.stage = 1
+        pending = owner.invoke()
+        original.stage = nested.stage = 2
+        return await pending if awaited else pending
+    finally:
+        owner.close()
+
+
+async def result_identity(owners: CallFactory, awaited: bool) -> IdentityObservation:
+    original = ControlPayload(ControlNode())
+
+    def callback() -> ControlPayload:
+        return original
+
+    async def callback_async() -> ControlPayload:
+        return original
+
+    owner = owners.prepare(callback_async if awaited else callback, (), awaited=awaited)
+    try:
+        pending = owner.invoke()
+        result = await pending if awaited else pending
+        return IdentityObservation(result is original, result.nested is original.nested, True)
+    finally:
+        owner.close()
+
+
+@dataclass
+class LifetimeCallback:
+    awaited: bool
+
+    def __call__(self, value: ControlNode, *, alias: ControlNode) -> object:
+        if self.awaited:
+            return self.run(value, alias=alias)
+        return value.stage + alias.stage
+
+    async def run(self, value: ControlNode, *, alias: ControlNode) -> int:
+        await asyncio.sleep(0)
+        return value.stage + alias.stage
+
+
+def prepare_released(
+    owners: CallFactory, awaited: bool
+) -> tuple[
+    PreparedInvocation,
+    tuple[
+        weakref.ReferenceType[LifetimeCallback], weakref.ReferenceType[ControlNode], weakref.ReferenceType[ControlNode]
+    ],
+]:
+    callback = LifetimeCallback(awaited)
+    value = ControlNode(13)
+    alias = ControlNode(29)
+    return (
+        owners.prepare(callback, (value,), {"alias": alias}, awaited=awaited),
+        (weakref.ref(callback), weakref.ref(value), weakref.ref(alias)),
+    )
+
+
+@dataclass(frozen=True, slots=True)
+class LifetimeObservation:
+    alive: tuple[bool, bool, bool]
+    result: int | ExpiredBorrow
+
+
+async def deferred_lifetime(owners: CallFactory, awaited: bool) -> LifetimeObservation:
+    owner, references = prepare_released(owners, awaited)
+    try:
+        gc.collect()
+        alive = tuple(reference() is not None for reference in references)
+        pending = owner.invoke()
+        result = await pending if awaited else pending
+        observation = LifetimeObservation(alive, result)
+    finally:
+        owner.close()
+    gc.collect()
+    assert all(reference() is None for reference in references)
+    return observation
+
+
+@dataclass(frozen=True, slots=True)
+class BorrowedObservation:
+    positional_identity: bool
+    keyword_identity: bool
+    positional_stage: int
+    keyword_stage: int
+
+
+async def borrowed_lifetime(owners: CallFactory, awaited: bool) -> BorrowedObservation:
+    value, alias = ControlNode(), ControlNode()
+    value_ref, alias_ref = weakref.ref(value), weakref.ref(alias)
+
+    def observe(value: ControlNode, *, alias: ControlNode) -> BorrowedObservation:
+        return BorrowedObservation(value is value_ref(), alias is alias_ref(), value.stage, alias.stage)
+
+    async def observe_async(value: ControlNode, *, alias: ControlNode) -> BorrowedObservation:
+        return observe(value, alias=alias)
+
+    owner = owners.prepare(observe_async if awaited else observe, (value,), {"alias": alias}, awaited=awaited)
+    try:
+        value.stage, alias.stage = 13, 29
+        pending = owner.invoke()
+        value.stage, alias.stage = 17, 31
+        return await pending if awaited else pending
+    finally:
+        owner.close()
+
+
+async def pending_handoff(owners: CallFactory, awaited: bool) -> LifetimeObservation:
+    assert awaited
+    owner, references = prepare_released(owners, awaited)
+    try:
+        pending = owner.invoke()
+        try:
+            owner.close()
+            gc.collect()
+            alive = tuple(reference() is not None for reference in references)
+            observation = LifetimeObservation(alive, await pending)
+        finally:
+            pending.close()
+    finally:
+        owner.close()
+    gc.collect()
+    assert all(reference() is None for reference in references)
+    return observation
+
+
+async def direct_coroutine(owners: CallFactory, awaited: bool) -> bool:
+    async def body() -> int:
+        return 73
+
+    original = body()
+    try:
+        owner = owners.prepare(lambda: original, (), awaited=awaited)
+        try:
+            pending = owner.invoke()
+            if awaited:
+                assert await pending == 73
+                nested = body()
+                try:
+
+                    async def returns_coroutine() -> Awaitable[int]:
+                        return nested
+
+                    nested_owner = owners.prepare(returns_coroutine, (), awaited=True)
+                    try:
+                        result = await nested_owner.invoke()
+                        return result is nested and inspect.getcoroutinestate(nested) == inspect.CORO_CREATED
+                    finally:
+                        nested_owner.close()
+                finally:
+                    nested.close()
+            try:
+                return pending is original and inspect.getcoroutinestate(original) == inspect.CORO_CREATED
+            finally:
+                if inspect.iscoroutine(pending):
+                    pending.close()
+        finally:
+            owner.close()
+    finally:
+        original.close()
+
+
+def expected_control(witness: str, control: str, awaited: bool) -> object:
+    if witness == "argument_identity":
+        return IdentityObservation(
+            control in ("identity", "envelope"),
+            control in ("identity", "envelope", "shallow_payload"),
+            control != "deep_separate",
+        )
+    if witness == "mutation_timing":
+        stage = 2 if awaited else 1
+        if control in ("deep_graph", "deep_separate"):
+            return TimingObservation(0, 0, 0)
+        return TimingObservation(0 if control == "shallow_payload" else stage, stage, stage)
+    if witness == "result_identity":
+        return IdentityObservation(control in ("identity", "result_passthrough"), control != "result_deep", True)
+    if witness in ("deferred_lifetime", "pending_handoff"):
+        if control == "weak" or (witness == "pending_handoff" and control == "missing_handoff"):
+            return LifetimeObservation(
+                (False, False, False), ExpiredBorrow(("callable", "positional:0", "keyword:alias"))
+            )
+        return LifetimeObservation((True, True, True), 42)
+    if witness == "borrowed_lifetime":
+        return BorrowedObservation(True, True, 17 if awaited else 13, 31 if awaited else 29)
+    return {"direct_coroutine": True}[witness]
+
+
+def run_control(witness: str, control: str, retained: bool, awaited: bool, factory: LiveCallFactory) -> None:
+    inner = factory if retained else cast(Callable[[], LiveCallFactory], globals()["ReferenceFactory"])()
+    owners = control_factory(control, inner)
+
+    async def run() -> None:
+        observed = await globals()[witness](owners, awaited)
+        assert observed == expected_control(witness, control, awaited), (witness, control, awaited, observed)
+
+    globals()["run_checked"](inner, run())

litellm-rust/crates/python-interop/tests/fixtures/callback_integrations.py

@@ -4,8 +4,10 @@ import gzip
 import json
 import threading
 from collections import OrderedDict
+from dataclasses import dataclass
 from datetime import datetime
 from functools import wraps
+from typing import Literal
 from unittest import TestCase
 
 import httpx
@@ -13,6 +15,7 @@ import httpx
 import litellm
 from litellm.integrations.custom_batch_logger import CustomBatchLogger
 from litellm.integrations.custom_guardrail import CustomGuardrail, ModifyResponseException
+from litellm.integrations.custom_logger import CustomLogger
 from litellm.integrations.datadog.datadog import DataDogLogger
 from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger
 from litellm.integrations.literal_ai import LiteralAILogger
@@ -51,9 +54,34 @@ def integration_callback_scope(scenario):
 
 @integration_callback_scope
 async def real_logging_queue_chain(owners):
-    entered, release = asyncio.Event(), asyncio.Event()
-    uploads = []
+    return await integration_logging_queue_case(owners)
 
+
+@dataclass(frozen=True, slots=True)
+class QueueObservation:
+    gcs_model_parameters: str
+    datadog_snapshot: str
+    literal_prepared_settings: str
+
+
+@integration_callback_scope
+async def real_logging_queue_copy_control(owners):
+    baseline = await integration_logging_queue_case(owners)
+    copied = await integration_logging_queue_case(owners, literal_copy="payload")
+    envelope = await integration_logging_queue_case(owners, literal_copy="envelope")
+    assert envelope == baseline
+    assert json.loads(baseline.gcs_model_parameters) == {"stream": True, "temperature": 0.25}
+    assert copied.datadog_snapshot == baseline.datadog_snapshot
+    assert copied.literal_prepared_settings == baseline.literal_prepared_settings
+    assert json.loads(copied.literal_prepared_settings) == {"stream": True}
+    assert json.loads(copied.gcs_model_parameters) == {
+        **json.loads(baseline.gcs_model_parameters),
+        "tools": [{"type": "function", "function": {"name": "lookup"}}],
+    }
+    return copied
+
+
+def queue_loggers(entered, release, uploads):
     class VertexTransport:
         async def _ensure_access_token_async(self, **kwargs):
             entered.set()
@@ -84,6 +112,36 @@ async def real_logging_queue_chain(owners):
     CustomBatchLogger.__init__(literal, batch_size=100, flush_lock=asyncio.Lock())
     literal.literalai_api_url, literal.headers = "https://literal.invalid", {}
     literal.async_httpx_client = Transport()
+    return datadog, gcs, literal
+
+
+async def integration_logging_queue_case(owners, *, literal_copy: Literal["direct", "envelope", "payload"] = "direct"):
+    entered, release = asyncio.Event(), asyncio.Event()
+    uploads = []
+    datadog, gcs, literal = queue_loggers(entered, release, uploads)
+    copy_payload = literal_copy == "payload"
+
+    class LiteralCallback(CustomLogger):
+        def __init__(self, delegate):
+            super().__init__()
+            self.delegate = delegate
+            self.calls = self.completed = 0
+
+        async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
+            self.calls += 1
+            self.received = kwargs
+            self.forwarded = {
+                **kwargs,
+                "standard_logging_object": (
+                    copy.deepcopy(kwargs["standard_logging_object"])
+                    if copy_payload
+                    else kwargs["standard_logging_object"]
+                ),
+            }
+            await self.delegate.async_log_failure_event(self.forwarded, response_obj, start_time, end_time)
+            self.completed += 1
+
+    literal_callback = literal if literal_copy == "direct" else LiteralCallback(literal)
 
     payload = create_dummy_standard_logging_payload()
     payload.update(status="failure", error_str="x" * 10001)
@@ -100,25 +158,37 @@ async def real_logging_queue_chain(owners):
         start_time=now,
         litellm_call_id="fixture-queue",
         function_id="fixture",
-        dynamic_async_failure_callbacks=[datadog, gcs, literal],
+        dynamic_async_failure_callbacks=[datadog, gcs, literal_callback],
     )
     error = RuntimeError("fixture failure")
     kwargs = logging.model_call_details
     kwargs.update(standard_logging_object=payload, model="fixture-model", exception=error, end_time=now)
     await integration_invoke(owners, logging.async_failure_handler, error, "fixture traceback", now, now)
+    if literal_copy != "direct":
+        assert literal_callback.calls == literal_callback.completed == 1
+        assert literal_callback.received is kwargs and literal_callback.forwarded is not kwargs
+        assert (literal_callback.forwarded["standard_logging_object"] is payload) is (not copy_payload)
+    assert len(datadog.log_queue) == gcs.log_queue.qsize() == len(literal.log_queue) == 1
     assert kwargs["standard_logging_object"] is payload
     assert payload["messages"] is messages and payload["model_parameters"] is settings
     assert payload["error_str"].endswith("truncated by litellm, this logger does not support large content")
-    assert "tools" not in settings
+    assert ("tools" in settings) is copy_payload
     dd_snapshot = json.loads(datadog.log_queue[0]["message"])
     assert dd_snapshot["model_parameters"]["tools"] == tools
     queued = gcs.log_queue.get_nowait()
     assert queued["payload"] is payload and queued["kwargs"] is kwargs and queued["response_obj"] is None
     gcs.log_queue.put_nowait(queued)
     generation = literal.log_queue[0]["generation"]
-    assert generation["settings"] is settings and generation["tools"] is tools
-    assert generation["messages"] is messages and generation["messageCompletion"] is completion
-    assert literal.log_queue[0]["metadata"] is metadata
+    prepared_settings = json.dumps(generation["settings"], sort_keys=True)
+    assert "tools" not in generation["settings"] and generation["tools"] == tools
+    if copy_payload:
+        assert generation["settings"] is not settings and generation["tools"] is not tools
+        assert generation["messages"] is not messages and generation["messageCompletion"] is not completion
+        assert literal.log_queue[0]["metadata"] is not metadata
+    else:
+        assert generation["settings"] is settings and generation["tools"] is tools
+        assert generation["messages"] is messages and generation["messageCompletion"] is completion
+        assert literal.log_queue[0]["metadata"] is metadata
 
     flush = asyncio.create_task(integration_invoke(owners, gcs.flush_queue))
     try:
@@ -148,13 +218,25 @@ async def real_logging_queue_chain(owners):
     assert json.loads(sent_dd[0]["message"]) == dd_snapshot
     literal_wire = uploads[2][1]["json"]
     sent_generation = literal_wire["variables"]["generation_0"]
-    assert sent_generation["messages"] == messages
     assert sent_generation["messages"] != gcs_snapshot["messages"]
-    assert sent_generation["settings"]["temperature"] == 0.25
-    assert sent_generation["messageCompletion"]["content"] == "late completion"
+    if copy_payload:
+        assert sent_generation["messages"] == dd_snapshot["messages"]
+        assert json.dumps(sent_generation["settings"], sort_keys=True) == prepared_settings
+        assert sent_generation["messageCompletion"] == dd_snapshot["response"]["choices"][0]["message"]
+    else:
+        assert sent_generation["messages"] == messages
+        assert sent_generation["settings"]["temperature"] == 0.25
+        assert sent_generation["messageCompletion"]["content"] == "late completion"
     messages[0]["content"] = "after serialization"
-    assert sent_generation["messages"][0]["content"] == "mutated before serialization"
+    assert sent_generation["messages"][0]["content"] == (
+        "Hello, world!" if copy_payload else "mutated before serialization"
+    )
     assert dd_snapshot["messages"][0]["content"] == "Hello, world!"
+    return QueueObservation(
+        gcs_model_parameters=json.dumps(gcs_snapshot["model_parameters"], sort_keys=True),
+        datadog_snapshot=json.dumps(dd_snapshot, sort_keys=True),
+        literal_prepared_settings=prepared_settings,
+    )
 
 
 @integration_callback_scope

litellm-rust/crates/python-interop/tests/fixtures/callback_lifecycle.py

@@ -594,15 +594,21 @@ async def detached_work_after_error(owners):
         assert reference() is None
 
 
-def run_scenario(name, retained, factory):
-    owners = factory if retained else ReferenceFactory()
+def run_checked(owners, scenario):
+    baseline = owners.live
 
     async def run():
-        await asyncio.wait_for(globals()[name](owners), timeout=15)
-        assert owners.live == 0
+        await asyncio.wait_for(scenario, timeout=15)
+        assert owners.live == baseline
         pending = asyncio.all_tasks() - {asyncio.current_task()}
         assert not pending, f"undrained tasks: {pending}"
 
     asyncio.run(run())
     gc.collect()
+    assert owners.live == baseline
+
+
+def run_scenario(name, retained, factory):
+    owners = factory if retained else ReferenceFactory()
     assert owners.live == 0
+    run_checked(owners, globals()[name](owners))