Fix Rust hook cleanup ordering and registration ownership

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

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"

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<PyAny>,
+    event_loop: Py<PyAny>,
+    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,
                 })?

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<dyn Fn() + Send + Sync>,
+}
+
+type Registrations = Arc<std::sync::Mutex<HashMap<String, Vec<Registration>>>>;
+
 // ---------------------------------------------------------------------------
 // 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<std::sync::Mutex<HashMap<String, Box<dyn Fn() + Send + Sync>>>>,
+    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::<PyRuntimeError, _>(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<amplifier_core::HookRegistry>,
-    /// Stored unregister closures keyed by handler name.
-    #[allow(clippy::type_complexity)]
-    unregister_fns: Arc<std::sync::Mutex<HashMap<String, Box<dyn Fn() + Send + Sync>>>>,
+    /// 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::<PyRuntimeError, _>(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::<PyRuntimeError, _>(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::<PyRuntimeError, _>(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(())
     }

bindings/python/src/session.rs

@@ -34,6 +34,11 @@ use crate::hooks::PyHookRegistry;
 pub(crate) struct PySession {
     /// Rust kernel session (for session_id, parent_id, initialized flag).
     inner: Arc<tokio::sync::Mutex<amplifier_core::Session>>,
+    /// 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<tokio::sync::Mutex<bool>>,
     /// The PyCoordinator instance owned by this session.
     coordinator: Py<PyAny>,
     /// Original config dict (Python dict).
@@ -183,6 +188,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 +281,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 +294,7 @@ impl PySession {
             // Convert to an owned Py<PyAny> so it's 'static + Send
             let coro_py: Py<PyAny> = 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 +332,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 +562,28 @@ 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. 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 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<Bound<'py, PyAny>> {
         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 +623,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 +708,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
                 // ----------------------------------------------------------