fix: make Python plugin teardown asyncio-safe

crates/python/src/py_plugin.rs

@@ -3,11 +3,11 @@
 
 //! Python-facing generic plugin configuration and registration helpers.
 
+#[cfg(test)]
+use std::collections::HashSet;
 use std::future::Future;
 use std::pin::Pin;
-use std::sync::{Arc, Mutex};
-#[cfg(test)]
-use std::{collections::HashSet, sync::LazyLock};
+use std::sync::{Arc, LazyLock, Mutex};
 
 use pyo3::prelude::*;
 use serde_json::{Map, Value as Json};
@@ -685,6 +685,7 @@ fn initialize_plugins_py<'py>(
         .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
     pyo3_async_runtimes::tokio::future_into_py(py, async move {
         let report = initialize_plugins(config).await.map_err(to_py_err)?;
+        reset_plugin_configuration_clear_state();
         Python::attach(|py| {
             let report = serde_json::to_value(&report)
                 .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
@@ -756,6 +757,35 @@ impl PluginTeardownError {
 
 type PluginTeardownResult = std::result::Result<(), PluginTeardownError>;
 
+struct PluginTeardownCompletion {
+    result: tokio::sync::watch::Sender<Option<PluginTeardownResult>>,
+}
+
+impl PluginTeardownCompletion {
+    fn new() -> Self {
+        let (result, _) = tokio::sync::watch::channel(None);
+        Self { result }
+    }
+
+    fn finish(&self, result: PluginTeardownResult) {
+        self.result.send_replace(Some(result));
+    }
+
+    async fn wait(&self, operation: &'static str) -> PluginTeardownResult {
+        let mut result = self.result.subscribe();
+        loop {
+            if let Some(result) = result.borrow().clone() {
+                return result;
+            }
+            if result.changed().await.is_err() {
+                return Err(PluginTeardownError::runtime(format!(
+                    "{operation} result channel closed unexpectedly"
+                )));
+            }
+        }
+    }
+}
+
 enum PluginHostCloseStatus {
     Active(Option<PluginHostActivation>),
     Closing,
@@ -764,15 +794,14 @@ enum PluginHostCloseStatus {
 
 struct PluginHostCloseState {
     status: Mutex<PluginHostCloseStatus>,
-    completion: tokio::sync::watch::Sender<Option<PluginTeardownResult>>,
+    completion: PluginTeardownCompletion,
 }
 
 impl PluginHostCloseState {
     fn new(activation: PluginHostActivation) -> Self {
-        let (completion, _) = tokio::sync::watch::channel(None);
         Self {
             status: Mutex::new(PluginHostCloseStatus::Active(Some(activation))),
-            completion,
+            completion: PluginTeardownCompletion::new(),
         }
     }
 
@@ -858,24 +887,84 @@ impl PluginHostCloseState {
             .status
             .lock()
             .unwrap_or_else(|poisoned| poisoned.into_inner()) = PluginHostCloseStatus::Closed;
-        self.completion.send_replace(Some(result));
+        reset_plugin_configuration_clear_state();
+        self.completion.finish(result);
     }
 
     async fn wait_for_close(&self) -> PluginTeardownResult {
-        let mut completion = self.completion.subscribe();
-        loop {
-            if let Some(result) = completion.borrow().clone() {
-                return result;
-            }
-            if completion.changed().await.is_err() {
-                return Err(PluginTeardownError::runtime(
-                    "dynamic plugin teardown result channel closed unexpectedly",
-                ));
+        self.completion.wait("dynamic plugin teardown").await
+    }
+}
+
+struct PluginConfigurationClearState {
+    started: Mutex<bool>,
+    completion: PluginTeardownCompletion,
+}
+
+impl PluginConfigurationClearState {
+    fn new() -> Self {
+        Self {
+            started: Mutex::new(false),
+            completion: PluginTeardownCompletion::new(),
+        }
+    }
+
+    fn begin_clear(self: &Arc<Self>) {
+        let should_start = {
+            let mut started = self
+                .started
+                .lock()
+                .unwrap_or_else(|poisoned| poisoned.into_inner());
+            if *started {
+                false
+            } else {
+                *started = true;
+                true
             }
+        };
+        if !should_start {
+            return;
         }
+
+        let clear_state = Arc::clone(self);
+        let spawn = std::thread::Builder::new()
+            .name("nemo-relay-python-plugin-clear".into())
+            .spawn(move || {
+                let result = std::panic::catch_unwind(clear_plugin_configuration)
+                    .map_err(|_| PluginTeardownError::runtime("plugin teardown task panicked"))
+                    .and_then(|result| result.map_err(PluginTeardownError::from_plugin_error));
+                clear_state.completion.finish(result);
+            });
+        if let Err(error) = spawn {
+            self.completion
+                .finish(Err(PluginTeardownError::runtime(format!(
+                    "failed to start plugin teardown task: {error}"
+                ))));
+        }
+    }
+
+    async fn wait_for_clear(&self) -> PluginTeardownResult {
+        self.completion.wait("plugin teardown").await
     }
 }
 
+static PLUGIN_CONFIGURATION_CLEAR_STATE: LazyLock<Mutex<Arc<PluginConfigurationClearState>>> =
+    LazyLock::new(|| Mutex::new(Arc::new(PluginConfigurationClearState::new())));
+
+fn plugin_configuration_clear_state() -> Arc<PluginConfigurationClearState> {
+    PLUGIN_CONFIGURATION_CLEAR_STATE
+        .lock()
+        .unwrap_or_else(|poisoned| poisoned.into_inner())
+        .clone()
+}
+
+fn reset_plugin_configuration_clear_state() {
+    *PLUGIN_CONFIGURATION_CLEAR_STATE
+        .lock()
+        .unwrap_or_else(|poisoned| poisoned.into_inner()) =
+        Arc::new(PluginConfigurationClearState::new());
+}
+
 #[pymethods]
 impl PyPluginHostActivation {
     /// Return the activation report captured during initialization.
@@ -935,6 +1024,7 @@ fn initialize_with_dynamic_plugins_py<'py>(
             PluginHostActivation::activate_with_discovered_config(config, dynamic_plugins)
                 .await
                 .map_err(plugin_error_to_py_err)?;
+        reset_plugin_configuration_clear_state();
         Python::attach(|py| {
             Py::new(
                 py,
@@ -953,6 +1043,19 @@ fn clear_plugin_configuration_py(py: Python<'_>) -> PyResult<()> {
     py.detach(clear_plugin_configuration).map_err(to_py_err)
 }
 
+#[pyfunction(name = "clear_plugin_configuration_async")]
+#[pyo3(signature = () -> "object", text_signature = "() -> object")]
+fn clear_plugin_configuration_async_py<'py>(py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+    let clear_state = plugin_configuration_clear_state();
+    clear_state.begin_clear();
+    pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        clear_state
+            .wait_for_clear()
+            .await
+            .map_err(|error| error.to_py_err())
+    })
+}
+
 #[pyfunction(name = "active_plugin_report")]
 #[pyo3(signature = () -> "object", text_signature = "() -> object")]
 fn active_plugin_report_py(py: Python<'_>) -> PyResult<Py<PyAny>> {
@@ -997,6 +1100,7 @@ pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(initialize_plugins_py, m)?)?;
     m.add_function(wrap_pyfunction!(initialize_with_dynamic_plugins_py, m)?)?;
     m.add_function(wrap_pyfunction!(clear_plugin_configuration_py, m)?)?;
+    m.add_function(wrap_pyfunction!(clear_plugin_configuration_async_py, m)?)?;
     m.add_function(wrap_pyfunction!(active_plugin_report_py, m)?)?;
     m.add_function(wrap_pyfunction!(list_plugin_kinds_py, m)?)?;
     m.add_function(wrap_pyfunction!(register_plugin_py, m)?)?;

crates/python/tests/coverage/nemo_guardrails_coverage_tests.rs

@@ -709,7 +709,7 @@ async def run_case():
     else:
         raise AssertionError("expected streamed output block")
 
-    nemo_relay.plugin.clear()
+    await nemo_relay.plugin.clear_async()
     await nemo_relay.plugin.initialize(plugin_config("stream_first_false"))
     try:
         await run_stream(request)

crates/python/tests/coverage/py_plugin_coverage_tests.rs

@@ -112,6 +112,7 @@ fn register_adds_plugin_management_bindings() {
             "initialize_plugins",
             "initialize_with_dynamic_plugins",
             "clear_plugin_configuration",
+            "clear_plugin_configuration_async",
             "active_plugin_report",
             "list_plugin_kinds",
             "register_plugin",
@@ -159,6 +160,30 @@ fn register_adds_plugin_management_bindings() {
     });
 }
 
+#[test]
+fn async_clear_binding_completes_on_python_event_loop() {
+    let _python = crate::test_support::init_python_test();
+    let _plugin_test_state = lock_plugin_test_state_for_tests();
+    Python::attach(|py| {
+        let module = PyModule::new(py, "_plugin_async_clear").unwrap();
+        register(&module).unwrap();
+        let helpers = load_module(
+            py,
+            r#"
+async def clear(module):
+    await module.clear_plugin_configuration_async()
+"#,
+        );
+        with_event_loop(py, |event_loop| {
+            let clear = helpers.getattr("clear").unwrap().call1((module,)).unwrap();
+            event_loop
+                .call_method1("run_until_complete", (clear,))
+                .unwrap();
+        });
+        assert!(active_plugin_report_py(py).unwrap().bind(py).is_none());
+    });
+}
+
 #[test]
 fn plugin_context_registers_all_runtime_hooks_and_drains_registrations() {
     let _python = crate::test_support::init_python_test();

docs/build-plugins/language-binding/code-examples.mdx

@@ -85,7 +85,7 @@ async def main() -> None:
         print("Activation report:", active_report)
         # Run instrumented application work here.
     finally:
-        nemo_relay.plugin.clear()
+        await nemo_relay.plugin.clear_async()
 
 
 if __name__ == "__main__":
@@ -216,7 +216,7 @@ async def main() -> None:
         print("Activation report:", active_report)
         # Run managed tool or LLM work here to log lifecycle events.
     finally:
-        nemo_relay.plugin.clear()
+        await nemo_relay.plugin.clear_async()
 
 
 if __name__ == "__main__":
@@ -335,7 +335,7 @@ async def main() -> None:
         await nemo_relay.plugin.initialize(config)
         # Run managed tool or LLM work here.
     finally:
-        nemo_relay.plugin.clear()
+        await nemo_relay.plugin.clear_async()
 
 
 if __name__ == "__main__":

docs/build-plugins/language-binding/register-behavior.mdx

@@ -278,7 +278,7 @@ async def main() -> None:
         print("Activation report:", active_report)
         print("Available kinds:", nemo_relay.plugin.list_kinds())
     finally:
-        nemo_relay.plugin.clear()
+        await nemo_relay.plugin.clear_async()
 
 
 if __name__ == "__main__":

docs/configure-plugins/adaptive/acg.mdx

@@ -92,7 +92,7 @@ async def main():
         # Run instrumented application work here.
         pass
     finally:
-        nemo_relay.plugin.clear()
+        await nemo_relay.plugin.clear_async()
 
 asyncio.run(main())
 ```

docs/configure-plugins/adaptive/adaptive-hints.mdx

@@ -89,7 +89,7 @@ async def main():
         # Run instrumented application work here.
         pass
     finally:
-        nemo_relay.plugin.clear()
+        await nemo_relay.plugin.clear_async()
 
 asyncio.run(main())
 ```

docs/configure-plugins/adaptive/response-cache.mdx

@@ -127,7 +127,7 @@ async def main():
         await nemo_relay.llm.execute("openai", request, call_model)
         await nemo_relay.llm.execute("openai", request, call_model)
     finally:
-        nemo_relay.plugin.clear()
+        await nemo_relay.plugin.clear_async()
 
 asyncio.run(main())
 ```

docs/configure-plugins/observability/about.mdx

@@ -85,7 +85,8 @@ Use the barrier that matches how you manage the exporter lifecycle:
 - For a manual exporter, call its documented `force_flush()`, `export()`, or
   `shutdown()` operation.
 - For plugin-managed export, call `plugin.clear()` or
-  `clear_plugin_configuration()` during graceful shutdown.
+  `clear_plugin_configuration()` during graceful shutdown. From a Python
+  `asyncio` task, await `plugin.clear_async()` instead.
 
 If the process terminates before those operations complete, the output can omit
 queued telemetry even when the instrumented application work completed

docs/configure-plugins/observability/atif.mdx

@@ -323,7 +323,7 @@ async def main():
         # Run instrumented application work here.
         pass
     finally:
-        plugin.clear()
+        await plugin.clear_async()
 
 asyncio.run(main())
 ```

docs/configure-plugins/observability/atof.mdx

@@ -149,8 +149,10 @@ events delivered after the exporter closes. A timeout is logged and does not by
 itself cause `shutdown()` to return an error.
 
 During graceful shutdown, `plugin.clear()` or `clear_plugin_configuration()`
-uses that same terminal stream-sink operation. It does not guarantee that a
-worker which timed out drained or closed.
+uses that same terminal stream-sink operation. From a Python `asyncio` task,
+await `plugin.clear_async()` so queued sanitizers can finish on the running
+event loop. Teardown does not guarantee that a worker which timed out drained
+or closed.
 
 ## Migrate From Version 1
 
@@ -260,7 +262,7 @@ async def main():
         # Run instrumented application work here.
         pass
     finally:
-        plugin.clear()
+        await plugin.clear_async()
 
 asyncio.run(main())
 ```

python/nemo_relay/_native.pyi

@@ -2385,6 +2385,17 @@ def clear_plugin_configuration() -> None:
     """
     ...
 
+def clear_plugin_configuration_async() -> Awaitable[None]:
+    """Clear active plugin configuration without blocking the Python event loop.
+
+    Returns:
+        Awaitable resolving when native teardown completes.
+
+    Exceptional flow:
+        Native cleanup and teardown worker errors propagate through the awaitable.
+    """
+    ...
+
 def active_plugin_report() -> Optional[_JsonObject]:
     """Return the active plugin report.