From 3654f7d034d62e1da261bc6bdce745fa7a34fb7c Mon Sep 17 00:00:00 2001 From: Brian Krabach Date: Tue, 22 Sep 2026 19:35:41 -0700 Subject: [PATCH 1/3] Fix Rust hook cleanup ordering and registration ownership --- bindings/python/src/bridges.rs | 61 ++- bindings/python/src/hooks.rs | 71 +++- bindings/python/src/session.rs | 41 +- .../python/tests/test_cleanup_ownership.py | 387 ++++++++++++++++++ python/amplifier_core/_async_compat.py | 34 ++ 5 files changed, 557 insertions(+), 37 deletions(-) create mode 100644 bindings/python/tests/test_cleanup_ownership.py diff --git a/bindings/python/src/bridges.rs b/bindings/python/src/bridges.rs index 9092c6d5..9ee9d643 100644 --- a/bindings/python/src/bridges.rs +++ b/bindings/python/src/bridges.rs @@ -44,6 +44,32 @@ pub(crate) struct PyHookHandlerBridge { unsafe impl Send for PyHookHandlerBridge {} unsafe impl Sync for PyHookHandlerBridge {} +/// into_future_with_locals does not cancel its Python task when the Rust +/// receiver is dropped. Keep exact ownership so cancellation cannot orphan a +/// hook callback or affect an unrelated task on the same event loop. +struct OwnedHookTask { + owner: Py, + event_loop: Py, + cancel_on_drop: bool, +} + +impl Drop for OwnedHookTask { + fn drop(&mut self) { + if self.cancel_on_drop { + let result = Python::try_attach(|py| -> PyResult<()> { + let cancel = self.owner.bind(py).getattr("cancel")?; + self.event_loop + .bind(py) + .call_method1("call_soon_threadsafe", (cancel,))?; + Ok(()) + }); + if let Some(Err(error)) = result { + log::warn!("Could not cancel owned Python hook task: {error}"); + } + } + } +} + impl HookHandler for PyHookHandlerBridge { fn handle( &self, @@ -100,7 +126,7 @@ impl HookHandler for PyHookHandlerBridge { // fall back to the registration loop when a native callback (such as // a blocking WASM host import) has lost Tokio task-local state. let py_result = if is_coro { - let future = Python::try_attach(|py| { + let (future, mut task) = Python::try_attach(|py| -> PyResult<_> { let locals = pyo3_async_runtimes::tokio::get_current_locals(py) .ok() .or_else(|| self.fallback_locals.clone()) @@ -109,10 +135,31 @@ impl HookHandler for PyHookHandlerBridge { "No running Python event loop available for coroutine conversion", ) })?; - pyo3_async_runtimes::into_future_with_locals( - &locals, - py_result_or_coro.into_bound(py), - ) + let owner = py + .import("amplifier_core._async_compat")? + .getattr("_OwnedHookTask")? + .call1((py_result_or_coro.bind(py),))?; + let runner = owner.call_method0("run")?; + let future = + match pyo3_async_runtimes::into_future_with_locals(&locals, runner.clone()) + { + Ok(future) => future, + Err(error) => { + // Scheduling failed (for example, the loop closed). + // Neither coroutine was adopted by a Python task. + let _ = runner.call_method0("close"); + let _ = owner.call_method0("cancel"); + return Err(error); + } + }; + Ok(( + future, + OwnedHookTask { + owner: owner.unbind(), + event_loop: locals.event_loop(py).unbind(), + cancel_on_drop: true, + }, + )) }) .ok_or_else(|| HookError::HandlerFailed { message: "Failed to attach to Python runtime for coroutine conversion" @@ -127,7 +174,9 @@ impl HookHandler for PyHookHandlerBridge { // Await OUTSIDE the GIL — drives the Python coroutine on the // current task's loop, or its registration loop after native // re-entry has lost task-local state. - future.await.map_err(|e| HookError::HandlerFailed { + let result = future.await; + task.cancel_on_drop = false; + result.map_err(|e| HookError::HandlerFailed { message: format!("Python async handler error: {e}"), handler_name: None, })? diff --git a/bindings/python/src/hooks.rs b/bindings/python/src/hooks.rs index 054c9111..dbd39794 100644 --- a/bindings/python/src/hooks.rs +++ b/bindings/python/src/hooks.rs @@ -13,6 +13,13 @@ use serde_json::Value; use crate::bridges::PyHookHandlerBridge; use crate::helpers::{json_dumps_safe, try_model_dump, wrap_future_as_coroutine}; +struct Registration { + id: uuid::Uuid, + unregister: Box, +} + +type Registrations = Arc>>>; + // --------------------------------------------------------------------------- // PyUnregisterFn — callable returned by PyHookRegistry.register() // --------------------------------------------------------------------------- @@ -24,18 +31,35 @@ use crate::helpers::{json_dumps_safe, try_model_dump, wrap_future_as_coroutine}; /// a callable that unregisters the handler when invoked. #[pyclass(name = "RustUnregisterFn")] pub(crate) struct PyUnregisterFn { - #[allow(clippy::type_complexity)] - unregister_fns: Arc>>>, + unregister_fns: Registrations, name: String, + registration_id: uuid::Uuid, } #[pymethods] impl PyUnregisterFn { fn __call__(&self) -> PyResult<()> { - if let Ok(mut fns) = self.unregister_fns.lock() { - if let Some(unreg) = fns.remove(&self.name) { - unreg(); + let registration = { + let mut fns = self + .unregister_fns + .lock() + .map_err(|e| PyErr::new::(format!("Lock poisoned: {e}")))?; + let mut owned = None; + if let Some(registrations) = fns.get_mut(&self.name) { + if let Some(index) = registrations + .iter() + .position(|r| r.id == self.registration_id) + { + owned = Some(registrations.remove(index)); + } + if registrations.is_empty() { + fns.remove(&self.name); + } } + owned + }; + if let Some(registration) = registration { + (registration.unregister)(); } Ok(()) } @@ -56,9 +80,8 @@ impl PyUnregisterFn { #[pyclass(name = "RustHookRegistry")] pub(crate) struct PyHookRegistry { pub(crate) inner: Arc, - /// Stored unregister closures keyed by handler name. - #[allow(clippy::type_complexity)] - unregister_fns: Arc>>>, + /// Names are labels, not identities: retain each registration's own closure. + unregister_fns: Registrations, } impl PyHookRegistry { @@ -118,10 +141,16 @@ impl PyHookRegistry { self.inner .register(event, bridge, priority, Some(handler_name.clone())); + let registration_id = uuid::Uuid::new_v4(); self.unregister_fns .lock() .map_err(|e| PyErr::new::(format!("Lock poisoned: {e}")))? - .insert(handler_name.clone(), unregister_fn); + .entry(handler_name.clone()) + .or_default() + .push(Registration { + id: registration_id, + unregister: unregister_fn, + }); // Return a callable that unregisters this handler when invoked. // Matches the Python HookRegistry.register() contract. @@ -130,6 +159,7 @@ impl PyHookRegistry { PyUnregisterFn { unregister_fns: self.unregister_fns.clone(), name: handler_name, + registration_id, }, )?; Ok(callable.into_any()) @@ -189,15 +219,22 @@ impl PyHookRegistry { ) } - /// Unregister a handler by name. + /// Unregister the most recent remaining registration with this name. + /// Returned callables still own their exact registration independently. fn unregister(&self, name: &str) -> PyResult<()> { - let mut fns = self - .unregister_fns - .lock() - .map_err(|e| PyErr::new::(format!("Lock poisoned: {e}")))?; - - if let Some(unreg) = fns.remove(name) { - unreg(); + let registration = { + let mut fns = self + .unregister_fns + .lock() + .map_err(|e| PyErr::new::(format!("Lock poisoned: {e}")))?; + let latest = fns.get_mut(name).and_then(Vec::pop); + if fns.get(name).is_some_and(Vec::is_empty) { + fns.remove(name); + } + latest + }; + if let Some(registration) = registration { + (registration.unregister)(); } Ok(()) } diff --git a/bindings/python/src/session.rs b/bindings/python/src/session.rs index b6233b42..1d2f71da 100644 --- a/bindings/python/src/session.rs +++ b/bindings/python/src/session.rs @@ -34,6 +34,10 @@ use crate::hooks::PyHookRegistry; pub(crate) struct PySession { /// Rust kernel session (for session_id, parent_id, initialized flag). inner: Arc>, + /// Serialize the entire end-event/drain sequence. The boolean records an + /// attempted terminal event for the current initialized lifetime, including + /// a cleanup whose waiter is cancelled during that event. + cleanup_state: Arc>, /// The PyCoordinator instance owned by this session. coordinator: Py, /// Original config dict (Python dict). @@ -183,6 +187,7 @@ impl PySession { Ok(Self { inner: Arc::new(tokio::sync::Mutex::new(session)), + cleanup_state: Arc::new(tokio::sync::Mutex::new(false)), coordinator: coord_any, config: config.clone().unbind(), is_resumed, @@ -275,7 +280,7 @@ impl PySession { } // Step 2: Extract what we need before entering the async block - let (coro_py, inner) = { + let (coro_py, inner, cleanup_state) = { let this = slf.borrow(); let helper = py.import("amplifier_core._session_init")?; let init_fn = helper.getattr("initialize_session")?; @@ -288,7 +293,7 @@ impl PySession { // Convert to an owned Py so it's 'static + Send let coro_py: Py = coro.unbind(); let inner = this.inner.clone(); - (coro_py, inner) + (coro_py, inner, this.cleanup_state.clone()) }; // Step 3: Patch the coordinator's session back-reference to point to @@ -326,8 +331,10 @@ impl PySession { // Step 5: Mark session as initialized in Rust kernel { + let mut end_emitted = cleanup_state.lock().await; let session = inner.lock().await; session.set_initialized(); + *end_emitted = false; } Ok(()) @@ -554,14 +561,17 @@ impl PySession { /// Clean up session resources. /// /// Rust controls the full cleanup lifecycle: - /// 1. Call all registered cleanup functions (reverse order, error-tolerant) - /// 2. Emit `session:end` event via hooks + /// 1. Await `session:end` once for the initialized lifetime, while hooks are live + /// 2. Call all registered cleanup functions (reverse order, error-tolerant) /// 3. Reset the initialized flag + /// Concurrent cleanup waits for this entire sequence. Uninitialized and + /// repeated cleanup still release resources but do not emit another end. /// /// Errors in cleanup functions and event emission are logged but never /// propagate — cleanup must always complete. fn cleanup<'py>(&self, py: Python<'py>) -> PyResult> { let inner = self.inner.clone(); + let cleanup_state = self.cleanup_state.clone(); // Grab references we need inside the async block let session_id = self.cached_session_id.clone(); @@ -601,8 +611,20 @@ impl PySession { wrap_future_as_coroutine( py, pyo3_async_runtimes::tokio::future_into_py(py, async move { + // Keep module cleanup behind the awaited terminal event even + // when another caller concurrently requests cleanup. + let mut end_emitted = cleanup_state.lock().await; + let initialized = inner.lock().await.is_initialized(); + if initialized && !*end_emitted { + // Claim before awaiting so a cancelled waiter cannot cause + // a subsequent cleanup to emit the terminal event twice. + *end_emitted = true; + let end_data = serde_json::json!({ "session_id": session_id }); + hooks_inner_for_end.emit("session:end", end_data).await; + } + // ---------------------------------------------------------- - // Step 1: Call all cleanup functions in reverse order + // Step 2: Call all cleanup functions in reverse order // Matches Python main's coordinator.cleanup() pattern: // if callable(fn): // if iscoroutinefunction(fn): await fn() @@ -674,15 +696,6 @@ impl PySession { } } - // ---------------------------------------------------------- - // Step 2: Emit session:end event (best-effort) - // Direct Rust emit — avoids Future/coroutine mismatch when going - // through the Python PyO3 bridge (future_into_py returns a Future, - // but into_future() expects a native coroutine). - // ---------------------------------------------------------- - let end_data = serde_json::json!({ "session_id": session_id }); - hooks_inner_for_end.emit("session:end", end_data).await; - // ---------------------------------------------------------- // Step 3: Reset the initialized flag // ---------------------------------------------------------- diff --git a/bindings/python/tests/test_cleanup_ownership.py b/bindings/python/tests/test_cleanup_ownership.py new file mode 100644 index 00000000..6936b0e5 --- /dev/null +++ b/bindings/python/tests/test_cleanup_ownership.py @@ -0,0 +1,387 @@ +"""Lifecycle ownership through the public, compiled Rust session and registry. + +The tiny filesystem modules only make initialization possible without providers. +Module loading, hook dispatch, session cleanup and unregister are all real Core. +""" + +import asyncio +import sys + +import pytest + +from amplifier_core import AmplifierSession, HookRegistry +from amplifier_core._engine import RustHookRegistry, RustSession + + +@pytest.fixture +def session_factory(tmp_path, monkeypatch): + names = [] + for module_type, mount_point in (("orchestrator", "orchestrator"), ("context", "context")): + name = f"amplifier_module_cleanup_fixture_{module_type}" + names.append(name) + package = tmp_path / name + package.mkdir() + (package / "__init__.py").write_text( + f'__amplifier_module_type__ = "{module_type}"\n' + "class FixtureModule:\n" + " async def execute(self, *args, **kwargs):\n" + " raise AssertionError('cleanup fixture must never execute a turn')\n" + "async def mount(coordinator, config=None):\n" + " if config and config.get('fail'):\n" + " raise RuntimeError('fixture mount failed')\n" + f' await coordinator.mount("{mount_point}", FixtureModule())\n' + " cleanup_log = coordinator.get_capability('fixture.cleanup')\n" + " if cleanup_log is not None:\n" + f" return lambda: cleanup_log.append('{module_type}')\n" + ) + monkeypatch.syspath_prepend(str(tmp_path)) + + def create(): + assert AmplifierSession is RustSession + return AmplifierSession( + config={ + "session": { + "orchestrator": "cleanup-fixture-orchestrator", + "context": "cleanup-fixture-context", + }, + "providers": [], + } + ) + + yield create + for name in names: + sys.modules.pop(name, None) + + +@pytest.mark.asyncio +async def test_end_is_awaited_before_module_cleanup_unregisters_it(session_factory): + session = session_factory() + await session.initialize() + calls = [] + + async def end(event, data): + await asyncio.sleep(0) + calls.append((event, data["session_id"])) + + unregister = session.coordinator.hooks.register("session:end", end, name="telemetry") + + def close(): + calls.append("close") + unregister() + + session.coordinator.register_cleanup(close) + await session.cleanup() + assert calls == [("session:end", session.session_id), "close"] + assert not session.initialized + + +@pytest.mark.asyncio +async def test_concurrent_cleanup_cannot_close_while_end_handler_is_running(session_factory): + session = session_factory() + await session.initialize() + entered, release = asyncio.Event(), asyncio.Event() + calls = [] + + async def end(event, data): + calls.append("end:entered") + entered.set() + await release.wait() + calls.append("end:finished") + + session.coordinator.hooks.register("session:end", end, name="telemetry") + session.coordinator.register_cleanup(lambda: calls.append("close")) + first = asyncio.create_task(session.cleanup()) + second = None + try: + await asyncio.wait_for(entered.wait(), timeout=2) + second = asyncio.create_task(session.cleanup()) + done, _ = await asyncio.wait({second}, timeout=0.05) + assert not done + assert calls == ["end:entered"] + finally: + release.set() + await asyncio.gather(first, *([second] if second else [])) + assert calls.count("end:entered") == 1 + assert calls.index("end:finished") < calls.index("close") + + +@pytest.mark.asyncio +async def test_cancelled_end_attempt_is_not_replayed_and_python_handler_exits(session_factory): + session = session_factory() + await session.initialize() + entered, release, exited = asyncio.Event(), asyncio.Event(), asyncio.Event() + calls = [] + + async def end(event, data): + calls.append("end:entered") + entered.set() + try: + await release.wait() + finally: + calls.append("end:exited") + exited.set() + + session.coordinator.hooks.register("session:end", end, name="telemetry") + session.coordinator.register_cleanup(lambda: calls.append("close")) + first = asyncio.create_task(session.cleanup()) + try: + await asyncio.wait_for(entered.wait(), timeout=2) + first.cancel() + with pytest.raises(asyncio.CancelledError): + await first + await asyncio.wait_for(exited.wait(), timeout=2) + await asyncio.wait_for(session.cleanup(), timeout=2) + assert calls == ["end:entered", "end:exited", "close"] + assert not session.initialized + finally: + release.set() + if not first.done(): + first.cancel() + await asyncio.gather(first, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_uninitialized_cleanup_still_releases_partial_resources(session_factory): + session = session_factory() + calls = [] + session.coordinator.hooks.register( + "session:end", lambda event, data: calls.append("end"), name="telemetry" + ) + session.coordinator.register_cleanup(lambda: calls.append("partial cleanup")) + await session.cleanup() + assert calls == ["partial cleanup"] + + +@pytest.mark.asyncio +async def test_failed_initialization_cleans_up_already_mounted_module(session_factory): + session = session_factory() + calls = [] + session.coordinator.register_capability("fixture.cleanup", calls) + session.config["session"]["context"] = { + "module": "cleanup-fixture-context", "config": {"fail": True} + } + session.coordinator.hooks.register( + "session:end", lambda event, data: calls.append("end"), name="telemetry" + ) + with pytest.raises(RuntimeError, match="fixture mount failed"): + await session.initialize() + assert not session.initialized + await session.cleanup() + assert calls == ["orchestrator"] + + +@pytest.mark.asyncio +async def test_end_once_per_initialized_lifetime(session_factory): + session = session_factory() + calls = [] + session.coordinator.hooks.register( + "session:end", lambda event, data: calls.append("end"), name="telemetry" + ) + session.coordinator.register_cleanup(lambda: calls.append("close")) + await session.initialize() + await session.cleanup() + await session.cleanup() + # Cleanup callbacks remain best-effort/idempotent resources, as before; + # a second cleanup does not fabricate another terminal lifecycle event. + assert calls == ["end", "close", "close"] + await session.initialize() + await session.cleanup() + assert calls == ["end", "close", "close", "end", "close"] + + +@pytest.mark.asyncio +async def test_end_handler_failure_does_not_skip_cleanup(session_factory): + session = session_factory() + await session.initialize() + calls = [] + + async def end(event, data): + calls.append("end") + raise RuntimeError("fixture failure") + + session.coordinator.hooks.register("session:end", end, name="telemetry") + session.coordinator.register_cleanup(lambda: calls.append("close")) + await session.cleanup() + assert calls == ["end", "close"] + assert not session.initialized + + +@pytest.mark.asyncio +@pytest.mark.parametrize("same_event", [False, True]) +async def test_unregister_owns_one_registration_even_with_repeated_names(same_event): + assert HookRegistry is RustHookRegistry + hooks = HookRegistry() + calls = [] + second_event = "first:event" if same_event else "second:event" + first = hooks.register( + "first:event", lambda event, data: calls.append("first"), name="shared" + ) + second = hooks.register( + second_event, lambda event, data: calls.append("second"), name="shared" + ) + first() + first() + for event in dict.fromkeys(("first:event", second_event)): + await hooks.emit(event, {}) + assert calls == ["second"] + second() + second() + calls.clear() + for event in dict.fromkeys(("first:event", second_event)): + await hooks.emit(event, {}) + assert calls == [] + + +@pytest.mark.asyncio +async def test_old_unregister_handle_cannot_remove_new_registration(): + hooks = HookRegistry() + calls = [] + old = hooks.register("event", lambda event, data: calls.append("old"), name="shared") + old() + new = hooks.register("event", lambda event, data: calls.append("new"), name="shared") + old() + await hooks.emit("event", {}) + assert calls == ["new"] + new() + + +@pytest.mark.asyncio +async def test_name_unregister_and_owned_handles_do_not_remove_other_registrations(): + hooks = HookRegistry() + calls = [] + old = hooks.register("event", lambda event, data: calls.append("old"), name="shared") + latest = hooks.register("event", lambda event, data: calls.append("latest"), name="shared") + # Preserve the existing name API's most-recent-registration behavior. + hooks.unregister("shared") + latest() + await hooks.emit("event", {}) + assert calls == ["old"] + old() + calls.clear() + await hooks.emit("event", {}) + assert calls == [] + + +@pytest.mark.asyncio +async def test_cancelled_hook_only_cancels_its_owned_python_task(): + hooks = HookRegistry() + entered, exited, release = asyncio.Event(), asyncio.Event(), asyncio.Event() + sibling_started = asyncio.Event() + + async def sibling(): + sibling_started.set() + await release.wait() + return "sibling completed" + + async def hook(event, data): + entered.set() + try: + await release.wait() + finally: + exited.set() + + hooks.register("event", hook, name="owned") + other = asyncio.create_task(sibling()) + pending = asyncio.create_task(hooks.emit("event", {})) + try: + await asyncio.wait_for(entered.wait(), timeout=2) + await asyncio.wait_for(sibling_started.wait(), timeout=2) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + await asyncio.wait_for(exited.wait(), timeout=2) + assert not other.done() + finally: + release.set() + await asyncio.gather(pending, return_exceptions=True) + assert await other == "sibling completed" + + +@pytest.mark.asyncio +async def test_async_hook_result_error_and_emitting_context_are_preserved(): + import contextvars + + value = contextvars.ContextVar("owned-hook-context", default="unset") + hooks = HookRegistry() + calls = [] + + async def failing(event, data): + await asyncio.sleep(0) + calls.append(value.get()) + raise RuntimeError("expected hook failure") + + async def result(event, data): + await asyncio.sleep(0) + return {"action": "deny", "reason": value.get()} + + value.set("registration") + hooks.register("event", failing, priority=0, name="failure") + hooks.register("event", result, priority=1, name="result") + value.set("emission") + answer = await hooks.emit("event", {}) + assert calls == ["emission"] + assert answer.action == "deny" + assert answer.reason == "emission" + + +@pytest.mark.asyncio +async def test_owned_hook_cancel_before_start_closes_unstarted_coroutine(): + import inspect + from amplifier_core._async_compat import _OwnedHookTask + + calls = [] + + async def hook(): + calls.append("started") + + coroutine = hook() + owner = _OwnedHookTask(coroutine) + owner.cancel() + assert inspect.getcoroutinestate(coroutine) == inspect.CORO_CLOSED + with pytest.raises(asyncio.CancelledError): + await owner.run() + assert calls == [] + + +@pytest.mark.asyncio +async def test_hook_scheduling_failure_closes_both_unawaited_coroutines(monkeypatch): + import inspect + from amplifier_core import _async_compat + + coroutines, calls, refusals = [], [], [] + + class ObservedOwner(_async_compat._OwnedHookTask): + def __init__(self, coroutine): + super().__init__(coroutine) + coroutines.append(coroutine) + + def run(self): + runner = super().run() + coroutines.append(runner) + return runner + + loop = asyncio.get_running_loop() + original = loop.call_soon_threadsafe + + def refuse_hook_schedule(callback, *args, **kwargs): + # Refuse only the actual PyO3 Python-awaitable conversion. Keep its + # result-delivery callback working so the real registry returns the + # normal fail-open HookResult for a failed hook. + if type(callback).__name__ == "PyEnsureFuture": + refusals.append(type(callback).__name__) + raise RuntimeError("fixture hook scheduling refused") + return original(callback, *args, **kwargs) + + async def hook(event, data): + calls.append("started") + + monkeypatch.setattr(_async_compat, "_OwnedHookTask", ObservedOwner) + monkeypatch.setattr(loop, "call_soon_threadsafe", refuse_hook_schedule) + hooks = HookRegistry() + hooks.register("event", hook, name="fixture") + result = await hooks.emit("event", {}) + assert result.action == "continue" + assert refusals == ["PyEnsureFuture"] + assert calls == [] + assert len(coroutines) == 2 + assert all(inspect.getcoroutinestate(c) == inspect.CORO_CLOSED for c in coroutines) diff --git a/python/amplifier_core/_async_compat.py b/python/amplifier_core/_async_compat.py index 71e7f624..c1a8d9aa 100644 --- a/python/amplifier_core/_async_compat.py +++ b/python/amplifier_core/_async_compat.py @@ -6,6 +6,40 @@ methods are drop-in compatible with the old pure-Python async def methods. """ +import asyncio + + +class _OwnedHookTask: + """Let a dropped Rust hook waiter cancel only its scheduled Python task. + + ``run`` is scheduled by into_future_with_locals, retaining the emitting + task's context. ``cancel`` must be called on that same event loop. Remember + cancellation before startup too: a dropped waiter must not start its hook + later merely because the loop has not yet run the scheduled coroutine. + """ + + def __init__(self, coroutine): + self._coroutine = coroutine + self._task = None + self._cancel_requested = False + + async def run(self): + if self._cancel_requested: + self._coroutine.close() + raise asyncio.CancelledError + self._task = asyncio.current_task() + try: + return await self._coroutine + finally: + self._task = None + + def cancel(self): + self._cancel_requested = True + if self._task is not None: + self._task.cancel() + else: + self._coroutine.close() + async def _wrap(awaitable): """Wrap a PyO3 awaitable in a proper Python coroutine.""" From 3e9cc379df41f51fabd30a36bb71df560a8ab82a Mon Sep 17 00:00:00 2001 From: Brian Krabach Date: Tue, 22 Sep 2026 20:38:05 -0700 Subject: [PATCH 2/3] Prepare lifecycle fix for versioned CLI smoke validation --- Cargo.lock | 4 +- bindings/python/Cargo.toml | 2 +- crates/amplifier-core/Cargo.toml | 2 +- pyproject.toml | 2 +- python/amplifier_core/__init__.py | 2 +- scripts/e2e-smoke-test.sh | 58 ++++++-- tests/test_e2e_smoke_script.py | 221 ++++++++++++++++++++++++++++++ 7 files changed, 273 insertions(+), 18 deletions(-) create mode 100644 tests/test_e2e_smoke_script.py diff --git a/Cargo.lock b/Cargo.lock index 505ed7e0..6dbeae54 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -40,7 +40,7 @@ checksum = "e9d4ee0d472d1cd2e28c97dfa124b3d8d992e10eb0a035f33f5d12e3a177ba3b" [[package]] name = "amplifier-core" -version = "2.0.0" +version = "2.0.1" dependencies = [ "base64", "chrono", @@ -77,7 +77,7 @@ dependencies = [ [[package]] name = "amplifier-core-py" -version = "2.0.0" +version = "2.0.1" dependencies = [ "amplifier-core", "log", diff --git a/bindings/python/Cargo.toml b/bindings/python/Cargo.toml index 9de970bd..3caf6b67 100644 --- a/bindings/python/Cargo.toml +++ b/bindings/python/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "amplifier-core-py" -version = "2.0.0" +version = "2.0.1" edition = "2021" description = "PyO3 bridge for amplifier-core Rust kernel" license = "MIT" diff --git a/crates/amplifier-core/Cargo.toml b/crates/amplifier-core/Cargo.toml index f959282e..9f0179d1 100644 --- a/crates/amplifier-core/Cargo.toml +++ b/crates/amplifier-core/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "amplifier-core" -version = "2.0.0" +version = "2.0.1" edition = "2021" description = "Pure Rust kernel for the Amplifier modular AI agent system" license = "MIT" diff --git a/pyproject.toml b/pyproject.toml index 5269da7a..32c49dbf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "amplifier-core" -version = "2.0.0" +version = "2.0.1" description = "Rust kernel with Python bindings for the Amplifier modular AI agent framework" license = "MIT" readme = "README.md" diff --git a/python/amplifier_core/__init__.py b/python/amplifier_core/__init__.py index 37279429..f2f7619e 100644 --- a/python/amplifier_core/__init__.py +++ b/python/amplifier_core/__init__.py @@ -6,7 +6,7 @@ AmplifierSession`) still give the pure-Python implementations. """ -__version__ = "2.0.0" +__version__ = "2.0.1" # --- Rust-backed primary types (THE SWITCHOVER) --- # These four were previously imported from their Python submodules. diff --git a/scripts/e2e-smoke-test.sh b/scripts/e2e-smoke-test.sh index 75f92d7e..993ce149 100755 --- a/scripts/e2e-smoke-test.sh +++ b/scripts/e2e-smoke-test.sh @@ -7,7 +7,7 @@ set -euo pipefail # # Prerequisites: # - Docker installed and running -# - ANTHROPIC_API_KEY set in environment (or in ~/.amplifier/keys.env) +# - Selected provider API key in environment (or in ~/.amplifier/keys.env) # - maturin installed (pip install maturin) # # Usage: @@ -20,6 +20,10 @@ set -euo pipefail # --local-source /path/to/amplifier-foundation # # Environment variables: +# SMOKE_PROVIDER anthropic (default) or openai +# SMOKE_MODEL Optional model override (otherwise use provider default) +# SMOKE_BUNDLE Optional bundle override (otherwise use CLI default) +# ANTHROPIC_BASE_URL / OPENAI_BASE_URL Optional selected provider endpoint # SMOKE_PROMPT Override the default test prompt # SMOKE_TIMEOUT Override the timeout in seconds (default: 360) @@ -31,6 +35,9 @@ SKIP_BUILD=false SMOKE_PROMPT="${SMOKE_PROMPT:-Ask recipe author to run one of its example recipes}" TIMEOUT_SECONDS="${SMOKE_TIMEOUT:-360}" LOCAL_SOURCES=() +SMOKE_PROVIDER="${SMOKE_PROVIDER:-anthropic}" +SMOKE_MODEL="${SMOKE_MODEL:-}" +SMOKE_BUNDLE="${SMOKE_BUNDLE:-}" # Colors (defined early so fail() works during arg parsing) RED='\033[0;31m' @@ -73,7 +80,11 @@ while [[ $# -gt 0 ]]; do echo " --local-source ../amplifier-bundle-modes/modules/hooks-mode" echo "" echo "Environment variables:" - echo " ANTHROPIC_API_KEY Required (or set in ~/.amplifier/keys.env)" + echo " SMOKE_PROVIDER anthropic (default) or openai" + echo " SMOKE_MODEL Optional model override" + echo " SMOKE_BUNDLE Optional bundle override" + echo " Selected provider API key required (environment or ~/.amplifier/keys.env)" + echo " ANTHROPIC_BASE_URL / OPENAI_BASE_URL Optional selected provider endpoint" echo " SMOKE_PROMPT Test prompt (default: 'Ask recipe author to run one of its example recipes')" echo " SMOKE_TIMEOUT Timeout in seconds (default: 360)" exit 0 @@ -92,8 +103,20 @@ trap cleanup EXIT # Step 0: Resolve API keys # --------------------------------------------------------------------------- -# If ANTHROPIC_API_KEY is not set, try to load from ~/.amplifier/keys.env -if [[ -z "${ANTHROPIC_API_KEY:-}" ]]; then +# Select the credential family before falling back to the existing keys file. +case "$SMOKE_PROVIDER" in + anthropic) + PROVIDER_KEY_ENV=ANTHROPIC_API_KEY + PROVIDER_BASE_URL_ENV=ANTHROPIC_BASE_URL + ;; + openai) + PROVIDER_KEY_ENV=OPENAI_API_KEY + PROVIDER_BASE_URL_ENV=OPENAI_BASE_URL + ;; + *) fail "SMOKE_PROVIDER must be anthropic or openai" ;; +esac + +if [[ -z "${!PROVIDER_KEY_ENV:-}" ]]; then KEYS_ENV="$HOME/.amplifier/keys.env" if [[ -f "$KEYS_ENV" ]]; then log "Loading API keys from $KEYS_ENV..." @@ -104,7 +127,13 @@ if [[ -z "${ANTHROPIC_API_KEY:-}" ]]; then fi fi -[[ -z "${ANTHROPIC_API_KEY:-}" ]] && fail "ANTHROPIC_API_KEY not set. Set it in your environment or in ~/.amplifier/keys.env" +[[ -z "${!PROVIDER_KEY_ENV:-}" ]] && fail "$PROVIDER_KEY_ENV not set. Set it in your environment or in ~/.amplifier/keys.env" +# Pass names, not values, in Docker argv. Only the selected family reaches +# noninteractive CLI auto-init, which otherwise prefers ambient Anthropic keys. +PROVIDER_ENV_ARGS=(-e "$PROVIDER_KEY_ENV") +if [[ -n "${!PROVIDER_BASE_URL_ENV:-}" ]]; then + PROVIDER_ENV_ARGS+=(-e "$PROVIDER_BASE_URL_ENV") +fi command -v docker &>/dev/null || fail "Docker not installed or not in PATH" # --------------------------------------------------------------------------- @@ -187,9 +216,7 @@ log "Pristine-import preflight passed." log "Creating isolated Docker container..." docker run -d --name "$CONTAINER_NAME" \ - -e ANTHROPIC_API_KEY="$ANTHROPIC_API_KEY" \ - -e OPENAI_API_KEY="${OPENAI_API_KEY:-}" \ - -e AZURE_OPENAI_API_KEY="${AZURE_OPENAI_API_KEY:-}" \ + "${PROVIDER_ENV_ARGS[@]}" \ python:3.12-slim \ sleep 3600 \ || fail "Container creation failed" @@ -308,10 +335,17 @@ echo "" # Run the smoke test; capture output even if timeout exits non-zero SMOKE_EXIT_CODE=0 -SMOKE_OUTPUT=$(docker exec "$CONTAINER_NAME" bash -c " - export PATH=/root/.local/bin:\$PATH - timeout $TIMEOUT_SECONDS amplifier run '$SMOKE_PROMPT' 2>&1 -" 2>&1) || SMOKE_EXIT_CODE=$? +SMOKE_OUTPUT=$(docker exec "$CONTAINER_NAME" bash -c ' + export PATH=/root/.local/bin:$PATH + smoke_args=(--provider "$2") + if [[ -n "$3" ]]; then + smoke_args+=(--model "$3") + fi + if [[ -n "$5" ]]; then + smoke_args+=(--bundle "$5") + fi + timeout "$1" amplifier run "${smoke_args[@]}" -- "$4" 2>&1 +' smoke-run "$TIMEOUT_SECONDS" "$SMOKE_PROVIDER" "$SMOKE_MODEL" "$SMOKE_PROMPT" "$SMOKE_BUNDLE" 2>&1) || SMOKE_EXIT_CODE=$? # --------------------------------------------------------------------------- # Step 8: Evaluate results diff --git a/tests/test_e2e_smoke_script.py b/tests/test_e2e_smoke_script.py new file mode 100644 index 00000000..43273152 --- /dev/null +++ b/tests/test_e2e_smoke_script.py @@ -0,0 +1,221 @@ +"""Exercise the release script with inert Docker and CLI process fixtures.""" + +from __future__ import annotations + +import json +from pathlib import Path +import shutil +import subprocess +import sys + +import pytest + + +SCRIPT = Path(__file__).resolve().parents[1] / "scripts/e2e-smoke-test.sh" +BASH = shutil.which("bash") +pytestmark = pytest.mark.skipif(BASH is None, reason="release script requires bash") + + +@pytest.fixture +def smoke(tmp_path): + repo = tmp_path / "source" + (repo / "scripts").mkdir(parents=True) + shutil.copyfile(SCRIPT, repo / "scripts/e2e-smoke-test.sh") + (repo / "dist").mkdir() + (repo / "dist/amplifier_core-2.0.0-py3-none-any.whl").write_text("inert wheel") + home = tmp_path / "home" + (home / ".amplifier").mkdir(parents=True) + commands = tmp_path / "bin" + commands.mkdir() + log = tmp_path / "calls.jsonl" + state = tmp_path / "container.json" + fixture = tmp_path / "fixture.py" + fixture.write_text( + """import json, os, pathlib, subprocess, sys +kind, *args = sys.argv[1:] +with open(os.environ['FIXTURE_LOG'], 'a') as f: + f.write(json.dumps({'kind': kind, 'args': args, + 'provider_env': {k:v for k,v in os.environ.items() + if k.startswith(('ANTHROPIC_', 'OPENAI_', 'AZURE_'))}}) + '\\n') +if kind == 'cli': + print(os.environ.get('FIXTURE_SMOKE_OUTPUT', 'inert CLI completed')) + sys.exit(int(os.environ.get('FIXTURE_SMOKE_EXIT', '0'))) +if kind == 'timeout': + sys.exit(0) +if kind != 'docker': + sys.exit(90) +if args[0] in ('rm', 'cp'): + sys.exit(0) +if args[0] == 'run': + if '-d' in args: + selected = {} + for i, arg in enumerate(args): + if arg == '-e': + key, sep, value = args[i+1].partition('=') + if sep or key in os.environ: + selected[key] = value if sep else os.environ[key] + pathlib.Path(os.environ['FIXTURE_STATE']).write_text(json.dumps(selected)) + print('inert-container') + else: + sys.exit(int(os.environ.get('FIXTURE_PREFLIGHT_EXIT', '0'))) + sys.exit(0) +if args[0] == 'exec' and args[2:4] == ['bash', '-c']: + command = args[4] + if 'amplifier run' in command: + env = {k:v for k,v in os.environ.items() + if not k.startswith(('ANTHROPIC_', 'OPENAI_', 'AZURE_'))} + env.update(json.loads(pathlib.Path(os.environ['FIXTURE_STATE']).read_text())) + result = subprocess.run([os.environ['FIXTURE_BASH'], '-c', *args[4:]], + env=env, timeout=5) + sys.exit(result.returncode) + if 'amplifier --version' in command: + print('amplifier fixture (core 2.0.0)') + elif 'uv pip install' in command: + code = int(os.environ.get('FIXTURE_INSTALL_EXIT', '0')) + print('Installed inert wheel' if code == 0 else 'install failed') + sys.exit(code) + elif 'apt-get' not in command and 'uv tool install' not in command: + sys.exit(91) + sys.exit(0) +sys.exit(92) +""" + ) + docker = commands / "docker" + docker.write_text( + f"#!{sys.executable}\n" + "import os, sys\n" + "os.execv(sys.executable, [sys.executable, os.environ['FIXTURE_PROGRAM']," + " 'docker', *sys.argv[1:]])\n" + ) + docker.chmod(0o700) + # Shell functions take precedence over PATH, including /root/.local/bin + # inserted by the actual smoke command. No installed CLI can be reached. + shell_env = tmp_path / "shell-env" + shell_env.write_text( + 'amplifier() { "$FIXTURE_PYTHON" "$FIXTURE_PROGRAM" cli "$@"; }\n' + 'timeout() { "$FIXTURE_PYTHON" "$FIXTURE_PROGRAM" timeout "$1"; ' + 'shift; "$@"; }\n' + ) + env = { + "PATH": f"{commands}:/usr/bin:/bin", + "HOME": str(home), + "BASH_ENV": str(shell_env), + "FIXTURE_PROGRAM": str(fixture), + "FIXTURE_PYTHON": sys.executable, + "FIXTURE_BASH": BASH, + "FIXTURE_LOG": str(log), + "FIXTURE_STATE": str(state), + } + + def run(extra=None, keys=None): + if keys is not None: + (home / ".amplifier/keys.env").write_text(keys) + result = subprocess.run( + [BASH, str(repo / "scripts/e2e-smoke-test.sh"), "--skip-build"], + env=env | (extra or {}), + capture_output=True, + text=True, + timeout=15, + ) + calls = [json.loads(line) for line in log.read_text().splitlines()] + return result, calls + + return run + + +def cli_call(calls): + found = [call for call in calls if call["kind"] == "cli"] + assert len(found) == 1 + return found[0] + + +def test_default_anthropic_keeps_optional_model_and_endpoint(smoke): + result, calls = smoke({"ANTHROPIC_API_KEY": "fixture-anthropic"}) + assert result.returncode == 0, result.stdout + result.stderr + call = cli_call(calls) + assert call["args"] == [ + "run", "--provider", "anthropic", "--", + "Ask recipe author to run one of its example recipes", + ] + assert call["provider_env"] == {"ANTHROPIC_API_KEY": "fixture-anthropic"} + + +@pytest.mark.parametrize("provider", ["anthropic", "openai"]) +def test_selected_family_only_with_endpoint_model_and_literal_arguments(smoke, tmp_path, provider): + prefix = provider.upper() + marker = tmp_path / "must-not-exist" + prompt = f"literal'; touch {marker}; # $(touch {marker})" + model = f"model' $(touch {marker})" + result, calls = smoke({ + "SMOKE_PROVIDER": provider, "SMOKE_MODEL": model, "SMOKE_PROMPT": prompt, + "ANTHROPIC_API_KEY": "fixture-anthropic", "OPENAI_API_KEY": "fixture-openai", + "AZURE_OPENAI_API_KEY": "fixture-azure", + "ANTHROPIC_BASE_URL": "https://anthropic.invalid/", "OPENAI_BASE_URL": "https://openai.invalid/v1", + }, keys="exit 87\n") + assert result.returncode == 0, result.stdout + result.stderr + call = cli_call(calls) + assert call["args"] == ["run", "--provider", provider, "--model", model, "--", prompt] + assert call["provider_env"] == { + f"{prefix}_API_KEY": f"fixture-{provider}", + f"{prefix}_BASE_URL": f"https://{provider}.invalid/" + ("v1" if provider == "openai" else ""), + } + assert not marker.exists() + assert "Loading API keys" not in result.stdout + + +def test_leading_dash_prompt_cannot_become_cli_help(smoke): + result, calls = smoke({"ANTHROPIC_API_KEY": "fixture-key", "SMOKE_PROMPT": "--help"}) + assert result.returncode == 0, result.stdout + result.stderr + assert cli_call(calls)["args"] == ["run", "--provider", "anthropic", "--", "--help"] + + +def test_optional_bundle_uses_supported_cli_argument(smoke): + result, calls = smoke({ + "OPENAI_API_KEY": "fixture-key", "SMOKE_PROVIDER": "openai", + "SMOKE_MODEL": "gpt-5.6-terra", "SMOKE_BUNDLE": "foundation", + "SMOKE_PROMPT": "fixture recipe", + }) + assert result.returncode == 0, result.stdout + result.stderr + assert cli_call(calls)["args"] == [ + "run", "--provider", "openai", "--model", "gpt-5.6-terra", + "--bundle", "foundation", "--", "fixture recipe", + ] + + +@pytest.mark.parametrize("provider", ["anthropic", "openai"]) +def test_selected_key_fallback_is_preserved(smoke, provider): + prefix = provider.upper() + result, calls = smoke( + {"SMOKE_PROVIDER": provider}, + keys=f'{prefix}_API_KEY="fixture-key"\n{prefix}_BASE_URL="https://fixture.invalid/"\n', + ) + assert result.returncode == 0, result.stdout + result.stderr + assert "Loading API keys" in result.stdout + assert cli_call(calls)["provider_env"] == { + f"{prefix}_API_KEY": "fixture-key", f"{prefix}_BASE_URL": "https://fixture.invalid/", + } + + +@pytest.mark.parametrize("extra", [{"SMOKE_PROVIDER": "unsupported"}, {"SMOKE_PROVIDER": "openai"}]) +def test_invalid_selection_or_missing_selected_key_stops_before_container(smoke, extra): + result, calls = smoke({"ANTHROPIC_API_KEY": "fixture-unselected"} | extra) + assert result.returncode != 0 + assert all(call["kind"] == "docker" and call["args"][0] == "rm" for call in calls) + + +@pytest.mark.parametrize( + ("extra", "message"), + [ + ({"FIXTURE_SMOKE_EXIT": "7"}, "amplifier exited with status 7"), + ({"FIXTURE_SMOKE_EXIT": "124"}, "TIMED OUT"), + ({"FIXTURE_SMOKE_OUTPUT": "Tool fixture failed: synthetic"}, "tool failure"), + ({"FIXTURE_PREFLIGHT_EXIT": "9"}, "Pristine-import preflight failed"), + ({"FIXTURE_INSTALL_EXIT": "8"}, "Wheel override failed"), + ], +) +def test_real_script_propagates_fixture_failures_and_cleans_container(smoke, extra, message): + result, calls = smoke({"ANTHROPIC_API_KEY": "fixture-key"} | extra) + assert result.returncode != 0 + assert message in result.stdout + assert "SMOKE TEST PASSED" not in result.stdout + assert calls[-1]["kind"] == "docker" and calls[-1]["args"][:2] == ["rm", "-f"] From 1e3ef7c3dba78cac74b1383115b143fef1db2265 Mon Sep 17 00:00:00 2001 From: Amplifier <240397093+microsoft-amplifier@users.noreply.github.com> Date: Wed, 23 Sep 2026 06:59:45 -0700 Subject: [PATCH 3/3] fix: clarify Python session cleanup ownership Generated with Amplifier Co-Authored-By: Amplifier <240397093+microsoft-amplifier@users.noreply.github.com> --- CONTRACTS.md | 23 ++++++++++++++++ bindings/python/src/session.rs | 26 ++++++++++++++----- .../python/tests/test_cleanup_ownership.py | 22 +++++++++++++--- .../python/tests/test_coroutine_compat.py | 8 ++++-- uv.lock | 2 +- 5 files changed, 67 insertions(+), 14 deletions(-) diff --git a/CONTRACTS.md b/CONTRACTS.md index aa65b548..77f87c83 100644 --- a/CONTRACTS.md +++ b/CONTRACTS.md @@ -235,6 +235,29 @@ but do not prevent the remaining callables from running. Source reference: `bindings/python/src/coordinator/mod.rs::cleanup()` and `bindings/python/src/coordinator/capabilities.rs::register_cleanup()`. +### Rust-backed Python Session Cleanup + +For the PyO3 `RustSession`, `cleanup()` claims and then awaits `session:end` at most +once per successfully initialized lifetime, before registered resource callbacks run. +A successful later `initialize()` resets that claim. Cancellation while a claimed +terminal dispatch is in progress may abort it: handlers not yet reached may never +receive the event, and a later cleanup does not replay it. Cancellation during +resource callbacks may interrupt remaining teardown without allowing a second +terminal attempt. The caller or host retaining and awaiting the cleanup task owns any +deadline or abandonment decision; waiter cancellation does not guarantee callbacks +drain. + +Handler cancellation is requested asynchronously. An immediate retry may run +resource callbacks before a prior terminal handler observes cancellation; the +cancellation path has no handler-drain or ordering guarantee. + +Uninitialized or partially initialized sessions still run registered resource +callbacks. Callbacks run in reverse registration order, tolerate errors, and may run +again on repeated cleanup, so they must be idempotent. Error tolerance does not make +cleanup cancellation-proof. Hooks and cleanup callbacks must not recursively await +`session.cleanup()`, because the outer cleanup is awaiting them and would form an +await-dependency cycle; this is a caller contract, not a runtime-enforced guard. + ### `on_session_ready(coordinator)` — Optional ```python diff --git a/bindings/python/src/session.rs b/bindings/python/src/session.rs index 1d2f71da..ab157459 100644 --- a/bindings/python/src/session.rs +++ b/bindings/python/src/session.rs @@ -34,9 +34,10 @@ use crate::hooks::PyHookRegistry; pub(crate) struct PySession { /// Rust kernel session (for session_id, parent_id, initialized flag). inner: Arc>, - /// Serialize the entire end-event/drain sequence. The boolean records an - /// attempted terminal event for the current initialized lifetime, including - /// a cleanup whose waiter is cancelled during that event. + /// Records whether cleanup has claimed `session:end` for the current + /// initialized lifetime. The claim happens before awaiting handlers. If + /// cancellation occurs while the claimed terminal dispatch is in progress, + /// a later cleanup does not replay it. cleanup_state: Arc>, /// The PyCoordinator instance owned by this session. coordinator: Py, @@ -561,14 +562,25 @@ impl PySession { /// Clean up session resources. /// /// Rust controls the full cleanup lifecycle: - /// 1. Await `session:end` once for the initialized lifetime, while hooks are live + /// 1. Claim then await `session:end` once for the initialized lifetime, + /// while hooks are live /// 2. Call all registered cleanup functions (reverse order, error-tolerant) /// 3. Reset the initialized flag - /// Concurrent cleanup waits for this entire sequence. Uninitialized and - /// repeated cleanup still release resources but do not emit another end. + /// Concurrent cleanup waits for this entire sequence absent cancellation. + /// Uninitialized and repeated cleanup still release resources but do not + /// emit another end. + /// Cancellation while the claimed terminal dispatch is in progress may + /// abort it; handlers not yet reached are not replayed by a later cleanup. + /// Cancellation during resource callbacks may interrupt remaining teardown + /// without allowing a second terminal attempt. The caller or host retaining + /// and awaiting the cleanup task owns any deadline or abandonment decision; + /// waiter cancellation does not guarantee callbacks drain. + /// Handler cancellation is requested asynchronously, so an immediate retry + /// may run resource callbacks before a prior terminal handler observes + /// cancellation; this path has no handler-drain or ordering guarantee. /// /// Errors in cleanup functions and event emission are logged but never - /// propagate — cleanup must always complete. + /// propagate. Error-tolerant cleanup is not cancellation-proof. fn cleanup<'py>(&self, py: Python<'py>) -> PyResult> { let inner = self.inner.clone(); let cleanup_state = self.cleanup_state.clone(); diff --git a/bindings/python/tests/test_cleanup_ownership.py b/bindings/python/tests/test_cleanup_ownership.py index 6936b0e5..0b3e495c 100644 --- a/bindings/python/tests/test_cleanup_ownership.py +++ b/bindings/python/tests/test_cleanup_ownership.py @@ -112,7 +112,7 @@ async def test_cancelled_end_attempt_is_not_replayed_and_python_handler_exits(se entered, release, exited = asyncio.Event(), asyncio.Event(), asyncio.Event() calls = [] - async def end(event, data): + async def blocking_end(event, data): calls.append("end:entered") entered.set() try: @@ -121,7 +121,15 @@ async def end(event, data): calls.append("end:exited") exited.set() - session.coordinator.hooks.register("session:end", end, name="telemetry") + async def later_end(event, data): + calls.append("end:later") + + session.coordinator.hooks.register( + "session:end", blocking_end, priority=10, name="blocking-telemetry" + ) + session.coordinator.hooks.register( + "session:end", later_end, priority=20, name="later-telemetry" + ) session.coordinator.register_cleanup(lambda: calls.append("close")) first = asyncio.create_task(session.cleanup()) try: @@ -129,15 +137,21 @@ async def end(event, data): first.cancel() with pytest.raises(asyncio.CancelledError): await first + assert "close" not in calls + await session.cleanup() await asyncio.wait_for(exited.wait(), timeout=2) - await asyncio.wait_for(session.cleanup(), timeout=2) - assert calls == ["end:entered", "end:exited", "close"] + assert calls.count("end:entered") == 1 + assert calls.count("end:exited") == 1 + assert calls.count("close") == 1 + assert "end:later" not in calls assert not session.initialized finally: release.set() if not first.done(): first.cancel() await asyncio.gather(first, return_exceptions=True) + if entered.is_set() and not exited.is_set(): + await asyncio.wait_for(exited.wait(), timeout=2) @pytest.mark.asyncio diff --git a/bindings/python/tests/test_coroutine_compat.py b/bindings/python/tests/test_coroutine_compat.py index 4ad5d887..0ff08347 100644 --- a/bindings/python/tests/test_coroutine_compat.py +++ b/bindings/python/tests/test_coroutine_compat.py @@ -27,7 +27,9 @@ async def test_hook_registry_emit_returns_coroutine(self): assert inspect.iscoroutine(result), ( f"emit() should return a coroutine, got {type(result).__name__}" ) - result.close() # cleanup + # The Rust future starts before this Python wrapper is awaited, so drain it + # before pytest tears down the event loop. + await result @pytest.mark.asyncio async def test_hook_registry_emit_and_collect_returns_coroutine(self): @@ -36,7 +38,9 @@ async def test_hook_registry_emit_and_collect_returns_coroutine(self): assert inspect.iscoroutine(result), ( f"emit_and_collect() should return a coroutine, got {type(result).__name__}" ) - result.close() + # The Rust future starts before this Python wrapper is awaited, so drain it + # before pytest tears down the event loop. + await result @pytest.mark.asyncio async def test_emit_works_with_create_task(self): diff --git a/uv.lock b/uv.lock index 1f55b0f9..fa26c1e1 100644 --- a/uv.lock +++ b/uv.lock @@ -4,7 +4,7 @@ requires-python = ">=3.11" [[package]] name = "amplifier-core" -version = "2.0.0" +version = "2.0.1" source = { editable = "." } dependencies = [ { name = "click" },