From 391bb4cc74ccde1a3bd0424c2ef28a169bd5d943 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 6 Sep 2026 19:11:28 -0700 Subject: [PATCH] wip --- .../python-bridge/src/routes/retained_http.rs | 106 ++-- .../tests/fixtures/retained_http_contract.py | 370 +++++++++++++- .../python-bridge/tests/retained_http.rs | 58 +++ .../tests/callback_lifecycle.rs | 134 +++-- .../tests/fixtures/callback_controls.py | 477 ++++++++++++++++++ .../tests/fixtures/callback_integrations.py | 104 +++- .../tests/fixtures/callback_lifecycle.py | 14 +- 7 files changed, 1164 insertions(+), 99 deletions(-) create mode 100644 litellm-rust/crates/python-interop/tests/fixtures/callback_controls.py diff --git a/litellm-rust/crates/python-bridge/src/routes/retained_http.rs b/litellm-rust/crates/python-bridge/src/routes/retained_http.rs index 0ac67a50bb6..350f1545d6f 100644 --- a/litellm-rust/crates/python-bridge/src/routes/retained_http.rs +++ b/litellm-rust/crates/python-bridge/src/routes/retained_http.rs @@ -70,15 +70,20 @@ fn prepare(boundary: &Bound<'_, PyAny>, asynchronous: bool) -> PyResult, roots: &Bound<'_, PyAny>) -> PyResult { - type ByteHeaders<'py> = Vec<(Bound<'py, PyBytes>, Bound<'py, PyBytes>)>; - let encoded = invoke( +#[pyfunction] +fn encode(boundary: &Bound<'_, PyAny>, roots: &Bound<'_, PyAny>) -> PyResult> { + invoke( boundary, BoundaryMethod::Encode.resolve(false), PyTuple::new(boundary.py(), [roots])?, - )?; + ) +} + +fn request(execution: &Bound<'_, PyAny>) -> PyResult { + 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> { - let request = encode(boundary, roots)?; - pyo3_async_runtimes::tokio::future_into_py(boundary.py(), async move { +fn send<'py>(execution: &Bound<'py, PyAny>) -> PyResult> { + 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> { 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> { - static DRIVER: PyOnceLock> = 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> { + static MODULE: PyOnceLock> = 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()) } diff --git a/litellm-rust/crates/python-bridge/tests/fixtures/retained_http_contract.py b/litellm-rust/crates/python-bridge/tests/fixtures/retained_http_contract.py index a3399ca5d55..1b3c14995c0 100644 --- a/litellm-rust/crates/python-bridge/tests/fixtures/retained_http_contract.py +++ b/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() diff --git a/litellm-rust/crates/python-bridge/tests/retained_http.rs b/litellm-rust/crates/python-bridge/tests/retained_http.rs index e7046250ed1..1ff3514debe 100644 --- a/litellm-rust/crates/python-bridge/tests/retained_http.rs +++ b/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(()) + }) +} diff --git a/litellm-rust/crates/python-interop/tests/callback_lifecycle.rs b/litellm-rust/crates/python-interop/tests/callback_lifecycle.rs index 37fc2fad15f..e4628fd784e 100644 --- a/litellm-rust/crates/python-interop/tests/callback_lifecycle.rs +++ b/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, + 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, + #[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, + #[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, + 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<()> { diff --git a/litellm-rust/crates/python-interop/tests/fixtures/callback_controls.py b/litellm-rust/crates/python-interop/tests/fixtures/callback_controls.py new file mode 100644 index 00000000000..9463e34d0f1 --- /dev/null +++ b/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()) diff --git a/litellm-rust/crates/python-interop/tests/fixtures/callback_integrations.py b/litellm-rust/crates/python-interop/tests/fixtures/callback_integrations.py index 27aac2fbc59..210d9f848c6 100644 --- a/litellm-rust/crates/python-interop/tests/fixtures/callback_integrations.py +++ b/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 diff --git a/litellm-rust/crates/python-interop/tests/fixtures/callback_lifecycle.py b/litellm-rust/crates/python-interop/tests/fixtures/callback_lifecycle.py index d9601f429a2..ed5201d37d7 100644 --- a/litellm-rust/crates/python-interop/tests/fixtures/callback_lifecycle.py +++ b/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))