diff --git a/crates/core/src/api/llm.rs b/crates/core/src/api/llm.rs index af538a174..f0333bb91 100644 --- a/crates/core/src/api/llm.rs +++ b/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, response_codec: Option>, subscribers: Vec, + pending_publication: Option, } 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); } } diff --git a/crates/core/src/api/runtime/state.rs b/crates/core/src/api/runtime/state.rs index 0dfbcfe61..15db15dc2 100644 --- a/crates/core/src/api/runtime/state.rs +++ b/crates/core/src/api/runtime/state.rs @@ -155,6 +155,7 @@ struct GuardrailScopeCompletion<'a> { handle: Option, subscribers: &'a [EventSubscriberFn], scope_stack: ScopeStackHandle, + pending_publication: Option, } 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()); } } diff --git a/crates/core/src/api/runtime/subscriber_dispatcher.rs b/crates/core/src/api/runtime/subscriber_dispatcher.rs index 872d17aea..035831900 100644 --- a/crates/core/src/api/runtime/subscriber_dispatcher.rs +++ b/crates/core/src/api/runtime/subscriber_dispatcher.rs @@ -111,6 +111,13 @@ mod native { }, Flush { done: Sender<()>, + include_pending: bool, + }, + RegisterPending { + lineage: Arc, + }, + CompletePending { + permit: PublicationPermit, }, Barrier { publications: Receiver>, @@ -121,6 +128,7 @@ mod native { #[derive(Default)] pub(super) struct PublicationLineage { outstanding: AtomicUsize, + pending_terminal: bool, } pub(super) struct PublicationPermit(Arc); @@ -242,6 +250,22 @@ mod native { pub(super) sender: Sender>, } + pub(crate) struct PendingPublication { + sender: Sender, + permit: Option, + } + + 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 { + 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>) -> 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::>(); if lineages.is_empty() { @@ -790,8 +856,20 @@ mod native { inherited: Option<&Arc>, ) { 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::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::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() { diff --git a/crates/core/src/api/subscriber.rs b/crates/core/src/api/subscriber.rs index ba7052c8e..f09b9b535 100644 --- a/crates/core/src/api/subscriber.rs +++ b/crates/core/src/api/subscriber.rs @@ -70,8 +70,9 @@ pub fn deregister_subscriber(name: &str) -> Result { 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 diff --git a/crates/core/src/api/tool.rs b/crates/core/src/api/tool.rs index dd67a0b53..0d962549e 100644 --- a/crates/core/src/api/tool.rs +++ b/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, subscribers: Vec, scope_stack: ScopeStackHandle, + pending_publication: Option, } 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); } } diff --git a/crates/core/src/plugin.rs b/crates/core/src/plugin.rs index 8d298e396..b9c446764 100644 --- a/crates/core/src/plugin.rs +++ b/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 = { diff --git a/crates/core/tests/integration/middleware_tests.rs b/crates/core/tests/integration/middleware_tests.rs index 8fd5810ae..8bc84a9cc 100644 --- a/crates/core/tests/integration/middleware_tests.rs +++ b/crates/core/tests/integration/middleware_tests.rs @@ -82,6 +82,32 @@ use serde_json::json; // All tests share the global context, so we serialize them. static TEST_MUTEX: Mutex<()> = Mutex::new(()); +fn assert_flush_waits_for_pending_completion(complete: impl FnOnce()) { + let (flush_started_tx, flush_started_rx) = std::sync::mpsc::channel(); + let (flush_done_tx, flush_done_rx) = std::sync::mpsc::channel(); + let flush_thread = std::thread::spawn(move || { + flush_started_tx.send(()).unwrap(); + flush_done_tx.send(flush_subscribers()).unwrap(); + }); + flush_started_rx + .recv_timeout(std::time::Duration::from_secs(1)) + .expect("subscriber flush thread did not start"); + assert!( + matches!( + flush_done_rx.recv_timeout(std::time::Duration::from_millis(50)), + Err(std::sync::mpsc::RecvTimeoutError::Timeout) + ), + "flush must wait for pending managed completion" + ); + + complete(); + flush_done_rx + .recv_timeout(std::time::Duration::from_secs(1)) + .expect("subscriber flush did not complete after cancellation") + .unwrap(); + flush_thread.join().unwrap(); +} + struct CloseCallsStreamNext { next: Option, request: Option, @@ -1944,8 +1970,8 @@ async fn dropping_pending_tool_execution_closes_the_managed_lifecycle() { result = &mut execution => panic!("execution unexpectedly completed: {result:?}"), result = entered_rx => result.unwrap(), } - drop(execution); - flush_subscribers().unwrap(); + + assert_flush_waits_for_pending_completion(|| drop(execution)); let lifecycle = events .lock() @@ -2085,8 +2111,7 @@ async fn dropping_pending_conditional_closes_the_guardrail_scope() { result = &mut execution => panic!("execution unexpectedly completed: {result:?}"), result = entered_rx => result.unwrap(), } - drop(execution); - flush_subscribers().unwrap(); + assert_flush_waits_for_pending_completion(|| drop(execution)); let lifecycle = events .lock() @@ -2230,8 +2255,7 @@ async fn dropping_pending_llm_execution_closes_the_managed_lifecycle() { result = &mut execution => panic!("execution unexpectedly completed: {result:?}"), result = entered_rx => result.unwrap(), } - drop(execution); - flush_subscribers().unwrap(); + assert_flush_waits_for_pending_completion(|| drop(execution)); let lifecycle = events .lock() diff --git a/crates/core/tests/integration/pipeline_tests.rs b/crates/core/tests/integration/pipeline_tests.rs index 1ee128214..0137c1b28 100644 --- a/crates/core/tests/integration/pipeline_tests.rs +++ b/crates/core/tests/integration/pipeline_tests.rs @@ -28,6 +28,7 @@ use nemo_relay::api::registry::{ }; use nemo_relay::api::runtime::NemoRelayContextState; use nemo_relay::api::runtime::global_context; +use nemo_relay::api::runtime::subscriber_dispatcher::flush_subscribers_with_started_signal; use nemo_relay::api::runtime::{LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn}; use nemo_relay::api::runtime::{create_scope_stack, set_thread_scope_stack}; use nemo_relay::api::scope::{EmitMarkEventParams, ScopeType, event}; @@ -2087,21 +2088,28 @@ async fn test_dropped_stream_end_keeps_fifo_position_before_later_mark() { .unwrap(); let (flush_started_tx, flush_started_rx) = std::sync::mpsc::channel(); + let (flush_done_tx, flush_done_rx) = std::sync::mpsc::channel(); let flush_thread = std::thread::spawn(move || { - flush_started_tx.send(()).unwrap(); - flush_subscribers() + flush_done_tx + .send(flush_subscribers_with_started_signal(flush_started_tx)) + .unwrap(); }); flush_started_rx .recv_timeout(std::time::Duration::from_secs(2)) - .expect("subscriber flush thread did not start"); - std::thread::sleep(std::time::Duration::from_millis(50)); - let returned_before_release = flush_thread.is_finished(); - sanitizer_release.notify_one(); - flush_thread.join().unwrap().unwrap(); + .expect("subscriber flush request was not queued before waiting on the pending stream END"); assert!( - !returned_before_release, + matches!( + flush_done_rx.recv_timeout(std::time::Duration::from_millis(50)), + Err(std::sync::mpsc::RecvTimeoutError::Timeout) + ), "flush must wait for the pending stream END" ); + sanitizer_release.notify_one(); + flush_done_rx + .recv_timeout(std::time::Duration::from_secs(2)) + .expect("subscriber flush did not complete after sanitizer release") + .unwrap(); + flush_thread.join().unwrap(); let events = events.lock().unwrap(); let end_index = events diff --git a/crates/core/tests/unit/subscriber_dispatcher_tests.rs b/crates/core/tests/unit/subscriber_dispatcher_tests.rs index 38e57064e..dc242cf77 100644 --- a/crates/core/tests/unit/subscriber_dispatcher_tests.rs +++ b/crates/core/tests/unit/subscriber_dispatcher_tests.rs @@ -2,8 +2,9 @@ // SPDX-License-Identifier: Apache-2.0 use super::native::{ DispatcherLoopState, DispatcherMessage, PendingFlush, PublicationLineage, PublicationPermit, - dispatcher_sender, enqueue_dispatch_message, flush_subscribers, register_async_publication, - sanitize_event_snapshot, set_sanitizer_runtime_failure_for_test, spawn_background_publication, + dispatcher_sender, enqueue_dispatch_message, flush_queued_subscribers, flush_subscribers, + register_async_publication, register_pending_publication, sanitize_event_snapshot, + set_sanitizer_runtime_failure_for_test, spawn_background_publication, }; use super::{EventSubscriberFn, publication_context}; use crate::api::registry::RegistryRecord; @@ -50,7 +51,10 @@ fn flush_waits_for_active_but_not_later_publication_barriers() { .unwrap(); let (flush_tx, flush_rx) = mpsc::channel(); sender - .send(DispatcherMessage::Flush { done: flush_tx }) + .send(DispatcherMessage::Flush { + done: flush_tx, + include_pending: true, + }) .unwrap(); let later = register_async_publication().expect("later publication barrier"); @@ -92,6 +96,61 @@ fn flush_waits_for_active_but_not_later_publication_barriers() { flush_subscribers().unwrap(); } +#[test] +fn pending_publication_defers_flush_without_blocking_unrelated_delivery() { + let _lock = crate::shared_runtime::runtime_owner_test_mutex() + .lock() + .unwrap_or_else(|error| error.into_inner()); + flush_subscribers().unwrap(); + let pending = register_pending_publication().expect("pending publication"); + let (delivered_tx, delivered_rx) = mpsc::channel(); + let subscriber: EventSubscriberFn = Arc::new(move |_event| { + delivered_tx.send(()).unwrap(); + }); + let event = serde_json::from_value(serde_json::json!({ + "kind": "mark", + "atof_version": "0.1", + "uuid": "019c1df6-4a57-7000-8000-000000000004", + "timestamp": "2026-07-28T00:00:00Z", + "name": "unrelated-while-pending" + })) + .expect("valid event"); + enqueue_dispatch_message(DispatcherMessage::Deliver { + event: Box::new(event), + transform: None, + sanitizers: Vec::new(), + subscribers: vec![subscriber], + scope_stack: current_scope_stack(), + publication_context: None, + lineage: None, + }); + delivered_rx + .recv_timeout(std::time::Duration::from_secs(1)) + .expect("pending publication must not block unrelated delivery"); + flush_queued_subscribers().expect("queued-only flush must ignore managed work"); + + let (flush_tx, flush_rx) = mpsc::channel(); + dispatcher_sender() + .expect("dispatcher sender") + .send(DispatcherMessage::Flush { + done: flush_tx, + include_pending: true, + }) + .unwrap(); + assert!( + matches!( + flush_rx.recv_timeout(std::time::Duration::from_millis(50)), + Err(mpsc::RecvTimeoutError::Timeout) + ), + "flush must wait for work registered before it" + ); + + drop(pending); + flush_rx + .recv_timeout(std::time::Duration::from_secs(1)) + .expect("flush must complete after pending work"); +} + #[test] fn flush_does_not_wait_for_later_delivery() { let _lock = crate::shared_runtime::runtime_owner_test_mutex() @@ -102,7 +161,10 @@ fn flush_does_not_wait_for_later_delivery() { let sender = dispatcher_sender().expect("dispatcher sender"); let (flush_tx, flush_rx) = mpsc::channel(); sender - .send(DispatcherMessage::Flush { done: flush_tx }) + .send(DispatcherMessage::Flush { + done: flush_tx, + include_pending: true, + }) .unwrap(); let (release_tx, release_rx) = tokio::sync::oneshot::channel(); @@ -301,6 +363,9 @@ fn flush_waits_for_transitive_subscriber_publications_without_reordering() { flush_subscribers().unwrap(); let sender = dispatcher_sender().expect("dispatcher sender"); let delivered = Arc::new(Mutex::new(Vec::new())); + let (outer_started_tx, outer_started_rx) = mpsc::channel(); + let (release_outer_tx, release_outer_rx) = mpsc::channel(); + let release_outer_rx = Arc::new(Mutex::new(release_outer_rx)); let event = |uuid: &str, name: &str| { serde_json::from_value(serde_json::json!({ "kind": "mark", @@ -352,11 +417,20 @@ fn flush_waits_for_transitive_subscriber_publications_without_reordering() { let outer_subscriber: EventSubscriberFn = { let delivered = Arc::clone(&delivered); let child_subscriber = child_subscriber.clone(); + let release_outer_rx = Arc::clone(&release_outer_rx); Arc::new(move |event| { delivered .lock() .unwrap_or_else(|error| error.into_inner()) .push(event.name().to_string()); + outer_started_tx + .send(()) + .expect("outer subscriber start receiver was dropped"); + release_outer_rx + .lock() + .unwrap_or_else(|error| error.into_inner()) + .recv_timeout(std::time::Duration::from_secs(2)) + .expect("outer subscriber was not released within 2 seconds"); assert!(enqueue_dispatch_message(DispatcherMessage::Deliver { event: Box::new( serde_json::from_value(serde_json::json!({ @@ -398,6 +472,9 @@ fn flush_waits_for_transitive_subscriber_publications_without_reordering() { lineage: None, }) .unwrap(); + outer_started_rx + .recv_timeout(std::time::Duration::from_secs(2)) + .expect("outer subscriber did not start within 2 seconds"); sender .send(DispatcherMessage::Deliver { event: Box::new(event("019c1df6-4a57-7000-8000-000000000011", "later")), @@ -409,6 +486,9 @@ fn flush_waits_for_transitive_subscriber_publications_without_reordering() { lineage: None, }) .unwrap(); + release_outer_tx + .send(()) + .expect("outer subscriber release receiver was dropped"); flush_subscribers().unwrap(); assert_eq!( @@ -621,7 +701,10 @@ fn detached_sanitizer_tasks_cannot_inherit_a_later_publication_context() { .send(publication_context::().map(|value| value.as_str().to_string())) .unwrap(); let (done, _ignored) = mpsc::channel(); - enqueue_dispatch_message(DispatcherMessage::Flush { done }); + enqueue_dispatch_message(DispatcherMessage::Flush { + done, + include_pending: true, + }); }); Ok(fields) }) diff --git a/crates/node/README.md b/crates/node/README.md index 0de18dc68..60f742609 100644 --- a/crates/node/README.md +++ b/crates/node/README.md @@ -89,7 +89,6 @@ async function main() { }); await flushSubscribers(); - await new Promise((resolve) => setImmediate(resolve)); deregisterSubscriber("printer"); } @@ -100,11 +99,10 @@ main().catch((error) => { ``` Native subscriber delivery is asynchronous. Awaiting `flushSubscribers()` drains -the native dispatcher without blocking the Node.js event loop. JavaScript -subscribers run later through Node's callback queue, so native events they emit -are separate publications. The extra event-loop turn lets queued JavaScript -callback side effects complete before deregistration or exit; flush again if -those side effects emit native events that must also be observed. +the native dispatcher and waits for managed terminal publications registered +before the call and the JavaScript subscriber callbacks they queue, without +blocking the Node.js event loop. Native events emitted by a JavaScript subscriber +are separate publications; flush again if those events must also be observed. The main runtime API is exported from `nemo-relay-node`. Additional entry points are available at `nemo-relay-node/typed`, `nemo-relay-node/plugin`, diff --git a/crates/node/src/api/mod.rs b/crates/node/src/api/mod.rs index 12aff6e7f..ecd41f2da 100644 --- a/crates/node/src/api/mod.rs +++ b/crates/node/src/api/mod.rs @@ -485,9 +485,10 @@ pub fn push_stream_chunk(stream_id: f64, chunk: Json) -> bool { /// Signal that a stream is complete. Drops the sender so the Rust /// receiver sees the channel as closed. #[napi] -pub fn end_stream(stream_id: f64) { +pub fn end_stream(env: Env, stream_id: f64) -> napi::Result<()> { let id = stream_id as u64; finish_stream_channel(id, Ok(())); + callback_factory::expire_callback_context(&env) } /// # Safety @@ -1720,10 +1721,13 @@ pub fn create_scope_stack_from_propagation( .map_err(|error| napi::Error::from_reason(error.to_string())) } -/// Run a synchronous callback with an isolated scope stack installed. +/// Run a callback with an isolated scope stack installed. /// -/// The stack is restored before this function returns. Asynchronous callbacks -/// must not rely on this installation after their first `await`. +/// The caller's stack is restored immediately after callback invocation. When +/// the callback returns a Promise, the requested stack remains active for that +/// Promise until it settles and is then expired for inherited detached work. +/// Use this helper instead of `setThreadScopeStack` to isolate concurrent async +/// branches. #[napi] pub fn with_scope_stack( env: Env, @@ -1752,7 +1756,10 @@ pub fn current_scope_stack(env: Env) -> napi::Result { Ok(ScopeStack::from(current_scope_stack_handle())) } -/// Binds a scope stack to the current thread. +/// Binds a scope stack to the current thread or async resource. +/// +/// This mutates the current execution resource. Use `withScopeStack` when +/// concurrent asynchronous branches need isolated stack replacements. #[napi] pub fn set_thread_scope_stack(env: Env, stack: &ScopeStack) -> napi::Result<()> { if callback_factory::set_callback_scope_stack(&env, stack)? { @@ -3312,18 +3319,17 @@ pub fn deregister_subscriber(name: String) -> Result { core_subscriber_api::deregister_subscriber(&name).map_err(to_napi_err) } -/// Return a Promise that resolves when native subscriber callbacks queued -/// before this call finish. +/// Return a Promise that resolves when native and JavaScript subscriber callbacks and +/// managed terminal publications registered before this call finish. /// /// Call this function outside subscribers, event sanitizers, conditional /// guardrails, and request or execution intercepts. A queued tool or LLM /// observability sanitizer may call it, but the Promise resolves without /// waiting for its own publication. /// -/// JavaScript subscribers are queued through Node's `ThreadsafeFunction`. Awaiting this -/// Promise does not block the Node event loop while Promise-returning event sanitizers settle. -/// Native events emitted later by a JavaScript subscriber are separate publications and may -/// require another flush after the JavaScript callback runs. +/// Awaiting this Promise does not block the Node event loop while Promise-returning event +/// sanitizers settle or queued JavaScript subscriber callbacks run. Native events emitted by a +/// JavaScript subscriber are separate publications and may require another flush. /// /// The Promise rejects if the blocking task fails or the core subscriber flush returns an error. /// Callers should handle errors when awaiting it. @@ -3335,10 +3341,13 @@ pub fn flush_subscribers(env: Env) -> Result { if reentrant { return Ok(()); } - tokio::task::spawn_blocking(core_subscriber_api::flush_subscribers) - .await - .map_err(|error| to_napi_err(FlowError::Internal(error.to_string())))? - .map_err(to_napi_err) + tokio::task::spawn_blocking(|| { + core_subscriber_api::flush_subscribers()?; + callable::flush_js_subscriber_callbacks() + }) + .await + .map_err(|error| to_napi_err(FlowError::Internal(error.to_string())))? + .map_err(to_napi_err) }, |env, _| env.get_undefined(), ) diff --git a/crates/node/src/callable.rs b/crates/node/src/callable.rs index 731348696..f085fb0fa 100644 --- a/crates/node/src/callable.rs +++ b/crates/node/src/callable.rs @@ -9,9 +9,10 @@ //! handles serialization of arguments to/from JSON and manages cross-thread communication //! between the Rust async runtime and the Node.js event loop. +use std::collections::BTreeSet; use std::future::Future; use std::pin::Pin; -use std::sync::Arc; +use std::sync::{Arc, Condvar, Mutex, OnceLock}; use napi::bindgen_prelude::ToNapiValue; use napi::threadsafe_function::{ErrorStrategy, ThreadsafeFunction, ThreadsafeFunctionCallMode}; @@ -44,6 +45,47 @@ use crate::convert::{callback_json, record_callback_error, to_napi_err}; use crate::promise_call::{JsonNextFn, JsonStreamNextFn, PromiseAwareFn}; use crate::types::{EventSanitizeFields, JsEvent, event_sanitize_fields_from_json}; +#[derive(Default)] +struct JsSubscriberCallbackState { + next_id: u64, + pending: BTreeSet, +} + +fn js_subscriber_callbacks() -> &'static (Mutex, Condvar) { + static CALLBACKS: OnceLock<(Mutex, Condvar)> = OnceLock::new(); + CALLBACKS.get_or_init(Default::default) +} + +fn reserve_js_subscriber_callback() -> u64 { + let (state, _) = js_subscriber_callbacks(); + let mut state = state.lock().unwrap(); + state.next_id += 1; + let id = state.next_id; + state.pending.insert(id); + id +} + +fn complete_js_subscriber_callback(id: u64) { + let (state, completed) = js_subscriber_callbacks(); + let mut state = state.lock().unwrap(); + state.pending.remove(&id); + completed.notify_all(); +} + +pub(crate) fn flush_js_subscriber_callbacks() -> Result<()> { + let (state, completed) = js_subscriber_callbacks(); + let mut state = state + .lock() + .map_err(|error| FlowError::Internal(error.to_string()))?; + let watermark = state.next_id; + while state.pending.range(..=watermark).next().is_some() { + state = completed + .wait(state) + .map_err(|error| FlowError::Internal(error.to_string()))?; + } + Ok(()) +} + /// Structured codec identity delivered to JavaScript LLM sanitizers. #[napi(object)] #[derive(Clone)] @@ -1153,8 +1195,17 @@ pub fn wrap_js_event_subscriber( return; } }; - let status = func.call(event_json, ThreadsafeFunctionCallMode::NonBlocking); + let callback_id = reserve_js_subscriber_callback(); + let status = func.call_with_return_value( + event_json, + ThreadsafeFunctionCallMode::NonBlocking, + move |_value: JsUnknown| { + complete_js_subscriber_callback(callback_id); + Ok(()) + }, + ); if status != napi::Status::Ok { + complete_js_subscriber_callback(callback_id); record_callback_error(format!( "nemo_relay: failed to queue JS event subscriber callback: {status:?}" )); diff --git a/crates/node/src/callback_factory.rs b/crates/node/src/callback_factory.rs index 8bf1acca3..cb7988849 100644 --- a/crates/node/src/callback_factory.rs +++ b/crates/node/src/callback_factory.rs @@ -10,7 +10,7 @@ use nemo_relay::api::runtime::subscriber_dispatcher::PublicationBuffer; use crate::types::ScopeStack; -const CALLBACK_FACTORIES_PROPERTY: &str = "__nemo_relay_callback_factories_v9"; +const CALLBACK_FACTORIES_PROPERTY: &str = "__nemo_relay_callback_factories_v11"; const CALLBACK_FACTORIES_SOURCE: &str = r#"(() => { const { AsyncLocalStorage } = process.getBuiltinModule('node:async_hooks'); @@ -18,6 +18,47 @@ const CALLBACK_FACTORIES_SOURCE: &str = r#"(() => { const publicationStates = new Map(); let nextPublicationContextId = 0; + function callbackStore(publicationState, publicationContextId, scopeStack, propagationParentUuid) { + const lifecycle = { expired: false, stores: new Set() }; + const store = { + lifecycle, + publicationState, + publicationContextId, + scopeStack, + propagationParentUuid, + }; + lifecycle.stores.add(store); + return store; + } + + function replacementCallbackStore(current, scopeStack) { + const store = { + lifecycle: current.lifecycle, + publicationState: current.publicationState, + publicationContextId: current.publicationContextId, + scopeStack, + propagationParentUuid: current.propagationParentUuid, + }; + current.lifecycle.stores.add(store); + if (current.lifecycle.expired) { + store.scopeStack = null; + store.propagationParentUuid = undefined; + } + return store; + } + + function expireCallbackStore(store) { + if (store === undefined) { + return; + } + store.lifecycle.expired = true; + for (const current of store.lifecycle.stores) { + current.scopeStack = null; + current.propagationParentUuid = undefined; + } + store.lifecycle.stores.clear(); + } + function jsonValue(value, seen = new Set()) { if (value === null || typeof value === 'string' || typeof value === 'boolean') { return value; @@ -69,6 +110,7 @@ const CALLBACK_FACTORIES_SOURCE: &str = r#"(() => { publication, publicationContextId, scopeStack, + propagationParentUuid, registerAbort, ) { const controller = new AbortController(); @@ -93,17 +135,18 @@ const CALLBACK_FACTORIES_SOURCE: &str = r#"(() => { } else { publicationState = { active: false }; } - const token = { + const token = callbackStore( publicationState, publicationContextId, scopeStack, - }; + propagationParentUuid, + ); const settlePublication = () => { if (ownsPublicationState) { publicationState.active = false; publicationStates.delete(publicationContextId); } - token.scopeStack = null; + expireCallbackStore(token); }; const safeNext = next === undefined ? undefined @@ -164,6 +207,7 @@ const CALLBACK_FACTORIES_SOURCE: &str = r#"(() => { publication, publicationContextId, scopeStack, + propagationParentUuid, registerAbort, ) { if (error != null) { @@ -186,6 +230,7 @@ const CALLBACK_FACTORIES_SOURCE: &str = r#"(() => { publication, publicationContextId, scopeStack, + propagationParentUuid, registerAbort, ); }; @@ -194,12 +239,12 @@ const CALLBACK_FACTORIES_SOURCE: &str = r#"(() => { scopedStream(fn) { return function __nemo_relay_scoped_stream_wrapper(arg, scopeStack, propagationParentUuid) { const current = eventSanitizerContext.getStore(); - const token = { - publicationState: current?.publicationState ?? { active: false }, - publicationContextId: current?.publicationContextId, + const token = callbackStore( + current?.publicationState ?? { active: false }, + current?.publicationContextId, scopeStack, propagationParentUuid, - }; + ); return eventSanitizerContext.run(token, () => fn(arg)); }; }, @@ -225,20 +270,29 @@ const CALLBACK_FACTORIES_SOURCE: &str = r#"(() => { withCallbackScopeStack(scopeStack, fn) { const current = eventSanitizerContext.getStore(); - if (current === undefined) { - return { active: false }; - } - const token = { - publicationState: current.publicationState, - publicationContextId: current.publicationContextId, + const token = callbackStore( + current?.publicationState ?? { active: false }, + current?.publicationContextId, scopeStack, - propagationParentUuid: current.propagationParentUuid, - }; + current?.propagationParentUuid, + ); + const expire = () => expireCallbackStore(token); + let value; try { - return { active: true, value: eventSanitizerContext.run(token, fn) }; - } finally { - token.scopeStack = null; + value = eventSanitizerContext.run(token, fn); + } catch (error) { + expire(); + throw error; } + if ( + value !== null + && (typeof value === 'object' || typeof value === 'function') + && typeof value.then === 'function' + ) { + return { active: true, value: Promise.resolve(value).finally(expire) }; + } + expire(); + return { active: true, value }; }, setCallbackScopeStack(scopeStack) { @@ -246,14 +300,17 @@ const CALLBACK_FACTORIES_SOURCE: &str = r#"(() => { if (current === undefined) { return false; } - eventSanitizerContext.enterWith({ - publicationState: current.publicationState, - publicationContextId: current.publicationContextId, - scopeStack, - propagationParentUuid: current.propagationParentUuid, - }); + eventSanitizerContext.enterWith(replacementCallbackStore(current, scopeStack)); return true; }, + + expireCallbackContext() { + const current = eventSanitizerContext.getStore(); + if (current === undefined) { + return; + } + expireCallbackStore(current); + }, }; })()"#; @@ -361,6 +418,13 @@ pub(crate) fn callback_propagation_parent_uuid(env: &Env) -> napi::Result