fix(node): preserve async context and flush callbacks

crates/core/src/api/llm.rs

@@ -22,7 +22,8 @@ use crate::api::runtime::NemoRelayContextState;
 use crate::api::runtime::global_context;
 use crate::api::runtime::state::contextualize_stream;
 use crate::api::runtime::subscriber_dispatcher::{
-    dispatch_reserved_sanitized_event, dispatch_sanitized_event, dispatch_transformed_event,
+    PendingPublication, dispatch_reserved_sanitized_event, dispatch_sanitized_event,
+    dispatch_transformed_event, register_pending_publication,
 };
 use crate::api::runtime::{
     EventSubscriberFn, LlmCollectorFn, LlmExecutionNextFn, LlmFinalizerFn, LlmJsonStream,
@@ -1209,6 +1210,7 @@ struct ManagedLlmCompletion {
     metadata: Option<Json>,
     response_codec: Option<Arc<dyn LlmResponseCodec>>,
     subscribers: Vec<EventSubscriberFn>,
+    pending_publication: Option<PendingPublication>,
 }
 
 impl ManagedLlmCompletion {
@@ -1223,16 +1225,21 @@ impl ManagedLlmCompletion {
             metadata,
             response_codec,
             subscribers: subscribers.to_vec(),
+            pending_publication: (!subscribers.is_empty())
+                .then(register_pending_publication)
+                .flatten(),
         }
     }
 
     fn disarm(&mut self) {
         self.handle = None;
+        drop(self.pending_publication.take());
     }
 }
 
 impl Drop for ManagedLlmCompletion {
     fn drop(&mut self) {
+        let pending_publication = self.pending_publication.take();
         let Some(handle) = self.handle.take() else {
             return;
         };
@@ -1312,6 +1319,7 @@ impl Drop for ManagedLlmCompletion {
             &subscribers,
             scope_stack,
         );
+        drop(pending_publication);
     }
 }
 

crates/core/src/api/runtime/state.rs

@@ -155,6 +155,7 @@ struct GuardrailScopeCompletion<'a> {
     handle: Option<ScopeHandle>,
     subscribers: &'a [EventSubscriberFn],
     scope_stack: ScopeStackHandle,
+    pending_publication: Option<subscriber_dispatcher::PendingPublication>,
 }
 
 impl GuardrailScopeCompletion<'_> {
@@ -167,6 +168,9 @@ impl GuardrailScopeCompletion<'_> {
             handle: Some(handle),
             subscribers,
             scope_stack,
+            pending_publication: (!subscribers.is_empty())
+                .then(subscriber_dispatcher::register_pending_publication)
+                .flatten(),
         }
     }
 
@@ -178,6 +182,7 @@ impl GuardrailScopeCompletion<'_> {
             self.subscribers,
             self.scope_stack.clone(),
         );
+        drop(self.pending_publication.take());
     }
 }
 
@@ -196,6 +201,7 @@ impl Drop for GuardrailScopeCompletion<'_> {
             self.subscribers,
             self.scope_stack.clone(),
         );
+        drop(self.pending_publication.take());
     }
 }
 

crates/core/src/api/runtime/subscriber_dispatcher.rs

@@ -111,6 +111,13 @@ mod native {
         },
         Flush {
             done: Sender<()>,
+            include_pending: bool,
+        },
+        RegisterPending {
+            lineage: Arc<PublicationLineage>,
+        },
+        CompletePending {
+            permit: PublicationPermit,
         },
         Barrier {
             publications: Receiver<Vec<DispatcherMessage>>,
@@ -121,6 +128,7 @@ mod native {
     #[derive(Default)]
     pub(super) struct PublicationLineage {
         outstanding: AtomicUsize,
+        pending_terminal: bool,
     }
 
     pub(super) struct PublicationPermit(Arc<PublicationLineage>);
@@ -242,6 +250,22 @@ mod native {
         pub(super) sender: Sender<Vec<DispatcherMessage>>,
     }
 
+    pub(crate) struct PendingPublication {
+        sender: Sender<DispatcherMessage>,
+        permit: Option<PublicationPermit>,
+    }
+
+    impl Drop for PendingPublication {
+        fn drop(&mut self) {
+            let Some(permit) = self.permit.take() else {
+                return;
+            };
+            let _ = self
+                .sender
+                .send(DispatcherMessage::CompletePending { permit });
+        }
+    }
+
     fn process_state() -> &'static ProcessState {
         let mut state = PROCESS_STATE.load(Ordering::Acquire);
         if state.is_null() {
@@ -331,7 +355,9 @@ mod native {
                     current
                 }
             }
-            DispatcherMessage::Flush { .. } => current,
+            DispatcherMessage::Flush { .. }
+            | DispatcherMessage::RegisterPending { .. }
+            | DispatcherMessage::CompletePending { .. } => current,
         }
     }
 
@@ -548,7 +574,28 @@ mod native {
         })
     }
 
-    pub(super) fn flush_subscribers() -> Result<()> {
+    /// Track asynchronous work without blocking unrelated dispatcher delivery.
+    ///
+    /// A flush observed after this registration waits until the returned handle
+    /// is dropped. Dropping the handle queues completion behind any terminal
+    /// publications sent by that work.
+    pub(super) fn register_pending_publication() -> Option<PendingPublication> {
+        let sender = dispatcher_sender().ok()?;
+        let lineage = Arc::new(PublicationLineage {
+            outstanding: AtomicUsize::new(0),
+            pending_terminal: true,
+        });
+        let permit = PublicationPermit::new(Arc::clone(&lineage));
+        sender
+            .send(DispatcherMessage::RegisterPending { lineage })
+            .ok()?;
+        Some(PendingPublication {
+            sender,
+            permit: Some(permit),
+        })
+    }
+
+    fn flush_subscribers_inner(include_pending: bool, started: Option<Sender<()>>) -> Result<()> {
         if in_dispatcher_callback() {
             return Ok(());
         }
@@ -567,16 +614,34 @@ mod native {
         };
         let (done_tx, done_rx) = mpsc::channel();
         sender
-            .send(DispatcherMessage::Flush { done: done_tx })
+            .send(DispatcherMessage::Flush {
+                done: done_tx,
+                include_pending,
+            })
             .map_err(|error| {
                 FlowError::Internal(format!("failed to queue subscriber flush: {error}"))
             })?;
+        if let Some(started) = started {
+            let _ = started.send(());
+        }
         done_rx
             .recv()
             .map_err(|error| FlowError::Internal(format!("subscriber flush failed: {error}")))?;
         Ok(())
     }
 
+    pub(super) fn flush_subscribers() -> Result<()> {
+        flush_subscribers_inner(true, None)
+    }
+
+    pub(super) fn flush_subscribers_with_started_signal(started: Sender<()>) -> Result<()> {
+        flush_subscribers_inner(true, Some(started))
+    }
+
+    pub(super) fn flush_queued_subscribers() -> Result<()> {
+        flush_subscribers_inner(false, None)
+    }
+
     pub(super) fn in_dispatcher_callback() -> bool {
         IN_DISPATCHER.with(Cell::get)
             || ASYNC_PUBLICATION_BUFFER.try_with(|_| ()).is_ok()
@@ -740,11 +805,12 @@ mod native {
             }
         }
 
-        fn defer_or_complete_flush(&mut self, done: Sender<()>) {
+        fn defer_or_complete_flush(&mut self, done: Sender<()>, include_pending: bool) {
             let lineages = self
                 .active_lineages
                 .iter()
                 .filter_map(Weak::upgrade)
+                .filter(|lineage| include_pending || !lineage.pending_terminal)
                 .filter(|lineage| lineage.outstanding.load(Ordering::Acquire) > 0)
                 .collect::<Vec<_>>();
             if lineages.is_empty() {
@@ -790,8 +856,20 @@ mod native {
         inherited: Option<&Arc<PublicationLineage>>,
     ) {
         let lineage = match message {
-            DispatcherMessage::Flush { done } => {
-                state.defer_or_complete_flush(done);
+            DispatcherMessage::Flush {
+                done,
+                include_pending,
+            } => {
+                state.defer_or_complete_flush(done, include_pending);
+                return;
+            }
+            DispatcherMessage::RegisterPending { lineage } => {
+                state.register(&lineage);
+                return;
+            }
+            DispatcherMessage::CompletePending { permit } => {
+                state.register(&permit.lineage());
+                drop(permit);
                 return;
             }
             _ => attach_publication_lineage(&mut message, inherited),
@@ -833,7 +911,9 @@ mod native {
                 }
                 drop(permit);
             }
-            DispatcherMessage::Flush { .. } => unreachable!(),
+            DispatcherMessage::Flush { .. }
+            | DispatcherMessage::RegisterPending { .. }
+            | DispatcherMessage::CompletePending { .. } => unreachable!(),
         }
         state.complete_ready_flushes();
     }
@@ -1128,6 +1208,18 @@ pub(crate) fn register_async_publication() -> Option<native::AsyncPublication> {
     native::register_async_publication()
 }
 
+pub(crate) use native::PendingPublication;
+
+/// Register pending asynchronous work that must finish before a later
+/// subscriber flush can complete.
+///
+/// Unlike an async publication barrier, this does not block unrelated
+/// dispatcher delivery. Dropping the returned handle releases pending flushes
+/// after previously queued terminal publications have been delivered.
+pub(crate) fn register_pending_publication() -> Option<native::PendingPublication> {
+    native::register_pending_publication()
+}
+
 /// Run asynchronous middleware as part of an already-registered publication,
 /// buffering the finalization publications explicitly assigned to its reserved
 /// FIFO position.
@@ -1150,12 +1242,28 @@ where
     native::spawn_background_publication(future)
 }
 
-/// Wait for all queued subscriber callbacks submitted before this call,
-/// including publications emitted transitively by those callbacks.
+/// Wait for all queued subscriber callbacks and managed terminal publications
+/// registered before this call, including publications emitted transitively by
+/// those callbacks.
 pub fn flush_subscribers() -> Result<()> {
     native::flush_subscribers()
 }
 
+/// Wait for subscriber completion and signal once the flush request has been
+/// queued, immediately before waiting for the dispatcher to acknowledge it.
+#[doc(hidden)]
+pub fn flush_subscribers_with_started_signal(started: std::sync::mpsc::Sender<()>) -> Result<()> {
+    native::flush_subscribers_with_started_signal(started)
+}
+
+/// Wait only for subscriber publications already queued on the dispatcher.
+///
+/// Plugin teardown uses this legacy barrier so it can detach registries while
+/// managed callbacks from an earlier snapshot remain in flight.
+pub(crate) fn flush_queued_subscribers() -> Result<()> {
+    native::flush_queued_subscribers()
+}
+
 /// Acquire process-local dispatcher resources before a Unix `fork`.
 #[doc(hidden)]
 pub fn prepare_for_fork() {

crates/core/src/api/subscriber.rs

@@ -70,8 +70,9 @@ pub fn deregister_subscriber(name: &str) -> Result<bool> {
     Ok(state.event_subscribers.remove(name).is_some())
 }
 
-/// Wait for all subscriber callbacks queued before this call to finish,
-/// including publications emitted transitively by those callbacks.
+/// Wait for all subscriber callbacks and managed terminal publications
+/// registered before this call to finish, including publications emitted
+/// transitively by those callbacks.
 ///
 /// A direct re-entrant call from queued publication middleware returns without
 /// waiting. Publication middleware must not move such a flush into

crates/core/src/api/tool.rs

@@ -8,7 +8,8 @@ use crate::api::runtime::NemoRelayContextState;
 use crate::api::runtime::current_scope_stack;
 use crate::api::runtime::global_context;
 use crate::api::runtime::subscriber_dispatcher::{
-    dispatch_sanitized_event, dispatch_transformed_event,
+    PendingPublication, dispatch_sanitized_event, dispatch_transformed_event,
+    register_pending_publication,
 };
 use crate::api::runtime::{
     EventSubscriberFn, ScopeStackHandle, ToolExecutionNextFn, with_active_event_uuid,
@@ -595,6 +596,7 @@ struct ManagedToolCompletion {
     metadata: Option<Json>,
     subscribers: Vec<EventSubscriberFn>,
     scope_stack: ScopeStackHandle,
+    pending_publication: Option<PendingPublication>,
 }
 
 impl ManagedToolCompletion {
@@ -609,16 +611,21 @@ impl ManagedToolCompletion {
             metadata,
             subscribers: subscribers.to_vec(),
             scope_stack,
+            pending_publication: (!subscribers.is_empty())
+                .then(register_pending_publication)
+                .flatten(),
         }
     }
 
     fn disarm(&mut self) {
         self.handle = None;
+        drop(self.pending_publication.take());
     }
 }
 
 impl Drop for ManagedToolCompletion {
     fn drop(&mut self) {
+        let pending_publication = self.pending_publication.take();
         let Some(handle) = self.handle.take() else {
             return;
         };
@@ -633,6 +640,7 @@ impl Drop for ManagedToolCompletion {
             &self.subscribers,
             self.scope_stack.clone(),
         );
+        drop(pending_publication);
     }
 }
 

crates/core/src/plugin.rs

@@ -1821,7 +1821,7 @@ pub(crate) struct PluginHostClearOutcome {
 }
 
 fn clear_plugin_configuration_inner() -> PluginHostClearOutcome {
-    let flush_error = crate::api::runtime::flush_subscribers()
+    let flush_error = crate::api::runtime::subscriber_dispatcher::flush_queued_subscribers()
         .err()
         .map(|error| error.to_string());
     let previous = {