From 20f805df2f9e276fddca0998c153d78e9efaf90c Mon Sep 17 00:00:00 2001
From: Will Killian <wkillian@nvidia.com>
Date: Thu, 30 Jul 2026 18:14:57 -0400
Subject: [PATCH 1/5] fix(node): preserve async callback context

Signed-off-by: Will Killian <wkillian@nvidia.com>
---
 .../core/tests/integration/pipeline_tests.rs  | 28 ++++---
 crates/node/src/api/mod.rs                    | 17 +++--
 crates/node/src/callback_factory.rs           | 53 ++++++++++---
 crates/node/src/promise_call.rs               | 11 ++-
 crates/node/tests/context_tests.mjs           | 23 ++++++
 crates/node/tests/llm_tests.mjs               | 75 +++++++++++++++++++
 crates/node/tests/tools_tests.mjs             | 59 +++++++++++++--
 docs/about-nemo-relay/concepts/scopes.mdx     | 16 +++-
 python/tests/test_llm.py                      | 24 +++++-
 9 files changed, 267 insertions(+), 39 deletions(-)

diff --git a/crates/core/tests/integration/pipeline_tests.rs b/crates/core/tests/integration/pipeline_tests.rs
index 9cb7e67d6..f14d66c1f 100644
--- a/crates/core/tests/integration/pipeline_tests.rs
+++ b/crates/core/tests/integration/pipeline_tests.rs
@@ -2086,22 +2086,28 @@ async fn test_dropped_stream_end_keeps_fifo_position_before_later_mark() {
     )
     .unwrap();
 
-    let (release_started_tx, release_started_rx) = std::sync::mpsc::channel();
-    let release_thread = std::thread::spawn(move || {
-        release_started_tx.send(()).unwrap();
-        std::thread::sleep(std::time::Duration::from_millis(50));
-        sanitizer_release.notify_one();
+    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();
     });
-    release_started_rx
+    flush_started_rx
         .recv_timeout(std::time::Duration::from_secs(2))
-        .expect("sanitizer release thread did not start");
-    let flush_started = std::time::Instant::now();
-    flush_subscribers().unwrap();
+        .expect("subscriber flush thread did not start");
     assert!(
-        flush_started.elapsed() >= std::time::Duration::from_millis(50),
+        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"
     );
-    release_thread.join().unwrap();
+    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/node/src/api/mod.rs b/crates/node/src/api/mod.rs
index 12aff6e7f..f0721f70b 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<ScopeStack> {
     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)? {
diff --git a/crates/node/src/callback_factory.rs b/crates/node/src/callback_factory.rs
index 8bf1acca3..3c6d82d29 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_v10";
 
 const CALLBACK_FACTORIES_SOURCE: &str = r#"(() => {
   const { AsyncLocalStorage } = process.getBuiltinModule('node:async_hooks');
@@ -69,6 +69,7 @@ const CALLBACK_FACTORIES_SOURCE: &str = r#"(() => {
     publication,
     publicationContextId,
     scopeStack,
+    propagationParentUuid,
     registerAbort,
   ) {
     const controller = new AbortController();
@@ -97,6 +98,7 @@ const CALLBACK_FACTORIES_SOURCE: &str = r#"(() => {
       publicationState,
       publicationContextId,
       scopeStack,
+      propagationParentUuid,
     };
     const settlePublication = () => {
       if (ownsPublicationState) {
@@ -104,6 +106,7 @@ const CALLBACK_FACTORIES_SOURCE: &str = r#"(() => {
         publicationStates.delete(publicationContextId);
       }
       token.scopeStack = null;
+      token.propagationParentUuid = undefined;
     };
     const safeNext = next === undefined
       ? undefined
@@ -164,6 +167,7 @@ const CALLBACK_FACTORIES_SOURCE: &str = r#"(() => {
         publication,
         publicationContextId,
         scopeStack,
+        propagationParentUuid,
         registerAbort,
       ) {
         if (error != null) {
@@ -186,6 +190,7 @@ const CALLBACK_FACTORIES_SOURCE: &str = r#"(() => {
           publication,
           publicationContextId,
           scopeStack,
+          propagationParentUuid,
           registerAbort,
         );
       };
@@ -225,20 +230,32 @@ 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,
+        publicationState: current?.publicationState ?? { active: false },
+        publicationContextId: current?.publicationContextId,
         scopeStack,
-        propagationParentUuid: current.propagationParentUuid,
+        propagationParentUuid: current?.propagationParentUuid,
       };
-      try {
-        return { active: true, value: eventSanitizerContext.run(token, fn) };
-      } finally {
+      const expire = () => {
         token.scopeStack = null;
+        token.propagationParentUuid = undefined;
+      };
+      let value;
+      try {
+        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) {
@@ -254,6 +271,15 @@ const CALLBACK_FACTORIES_SOURCE: &str = r#"(() => {
       });
       return true;
     },
+
+    expireCallbackContext() {
+      const current = eventSanitizerContext.getStore();
+      if (current === undefined) {
+        return;
+      }
+      current.scopeStack = null;
+      current.propagationParentUuid = undefined;
+    },
   };
 })()"#;
 
@@ -361,6 +387,13 @@ pub(crate) fn callback_propagation_parent_uuid(env: &Env) -> napi::Result<Option
         .map(Some)
 }
 
+pub(crate) fn expire_callback_context(env: &Env) -> napi::Result<()> {
+    let factories = callback_factories(env)?;
+    let expire: JsFunction = factories.get_named_property("expireCallbackContext")?;
+    expire.call::<JsUnknown>(None, &[])?;
+    Ok(())
+}
+
 pub(crate) fn with_callback_scope_stack(
     env: &Env,
     stack: &ScopeStack,
diff --git a/crates/node/src/promise_call.rs b/crates/node/src/promise_call.rs
index ca65be2af..892bd6627 100644
--- a/crates/node/src/promise_call.rs
+++ b/crates/node/src/promise_call.rs
@@ -25,7 +25,8 @@ use nemo_relay::api::runtime::subscriber_dispatcher::{
     PublicationBuffer, capture_nested_publication_buffer, with_task_nested_publication_buffer,
 };
 use nemo_relay::api::runtime::{
-    MiddlewareContinuationContext, ScopeStackHandle, current_scope_stack,
+    MiddlewareContinuationContext, ScopeStackHandle, capture_propagation_context,
+    current_scope_stack,
 };
 use nemo_relay::error::{FlowError, Result as FlowResult};
 
@@ -89,6 +90,7 @@ struct CallArgs {
     publication_context_id: Option<String>,
     /// Scope stack captured when Relay invokes the middleware.
     scope_stack: Option<ScopeStackHandle>,
+    propagation_parent_uuid: String,
     publication_buffer: Option<PublicationBuffer>,
     continuation_context: Option<MiddlewareContinuationContext>,
     cancellation: CallCancellation,
@@ -427,6 +429,10 @@ impl PromiseAwareFn {
                         }
                         None => undefined_to_unknown(&ctx.env)?,
                     };
+                    let propagation_parent_uuid = json_to_unknown(
+                        &ctx.env,
+                        Json::String(ctx.value.propagation_parent_uuid),
+                    )?;
                     let (resolve, reject) =
                         build_completion_unknowns(&ctx.env, ctx.value.completion)?;
                     let register_abort =
@@ -440,6 +446,7 @@ impl PromiseAwareFn {
                         publication,
                         publication_context_id,
                         scope_stack,
+                        propagation_parent_uuid,
                         register_abort,
                     ])
                 })();
@@ -560,6 +567,7 @@ impl PromiseAwareFn {
         let continuation_context = next
             .as_ref()
             .map(|_| MiddlewareContinuationContext::capture());
+        let propagation_parent_uuid = capture_propagation_context()?.parent_uuid.to_string();
         let tsfn = self
             .tsfn
             .lock()
@@ -578,6 +586,7 @@ impl PromiseAwareFn {
                 // Publication context also lets queued tool/LLM observability
                 // sanitizers avoid waiting on their own publication.
                 scope_stack: Some(current_scope_stack()),
+                propagation_parent_uuid,
                 publication_buffer: capture_nested_publication_buffer(),
                 continuation_context,
                 cancellation,
diff --git a/crates/node/tests/context_tests.mjs b/crates/node/tests/context_tests.mjs
index 41cede0e2..3180bf244 100644
--- a/crates/node/tests/context_tests.mjs
+++ b/crates/node/tests/context_tests.mjs
@@ -91,6 +91,29 @@ describe('Context isolation', () => {
     }
   });
 
+  it('keeps withScopeStack active until an async callback settles', async () => {
+    const originalUuid = getHandle().uuid;
+    const stack = createScopeStack();
+    const stackUuid = withScopeStack(stack, () => getHandle().uuid);
+    let resolveDetached;
+    const detached = new Promise((resolve) => {
+      resolveDetached = resolve;
+    });
+
+    const execution = withScopeStack(stack, async () => {
+      assert.equal(getHandle().uuid, stackUuid);
+      await new Promise((resolve) => setImmediate(resolve));
+      assert.equal(getHandle().uuid, stackUuid);
+      setImmediate(() => resolveDetached(getHandle().uuid));
+      return 'done';
+    });
+
+    assert.equal(getHandle().uuid, originalUuid);
+    assert.equal(await execution, 'done');
+    assert.equal(getHandle().uuid, originalUuid);
+    assert.equal(await detached, originalUuid);
+  });
+
   it('currentScopeStack returns same in same context', () => {
     const s1 = currentScopeStack();
     const s2 = currentScopeStack();
diff --git a/crates/node/tests/llm_tests.mjs b/crates/node/tests/llm_tests.mjs
index cb9c0db4a..4abdd98de 100644
--- a/crates/node/tests/llm_tests.mjs
+++ b/crates/node/tests/llm_tests.mjs
@@ -1125,6 +1125,56 @@ describe('LLM guardrails', () => {
 // ===========================================================================
 
 describe('LLM intercepts', () => {
+  it('execution callbacks preserve the managed propagation parent across await', async () => {
+    const events = [];
+    const observed = [];
+    registerSubscriber('node_llm_exec_propagation_parent', (event) => events.push(event));
+    registerLlmExecutionIntercept('node_llm_exec_propagation_parent', 10, async (request, next) => {
+      observed.push(['intercept-before', lib.capturePropagationContext().parentUuid]);
+      await new Promise((resolve) => setImmediate(resolve));
+      observed.push(['intercept-after', lib.capturePropagationContext().parentUuid]);
+      return next(request);
+    });
+    try {
+      const result = await llmCallExecuteAsync(
+        'propagation_parent_llm',
+        makeNative(),
+        async () => {
+          observed.push(['provider-before', lib.capturePropagationContext().parentUuid]);
+          await new Promise((resolve) => setImmediate(resolve));
+          observed.push(['provider-after', lib.capturePropagationContext().parentUuid]);
+          return { ok: true };
+        },
+        null,
+        null,
+        null,
+        null,
+        null,
+      );
+      assert.deepEqual(result, { ok: true });
+      await waitForSubscriberCallbacks(() =>
+        events.some(
+          (event) =>
+            event.name === 'propagation_parent_llm' && event.kind === 'scope' && event.scope_category === 'end',
+        ),
+      );
+      const start = events.find(
+        (event) =>
+          event.name === 'propagation_parent_llm' && event.kind === 'scope' && event.scope_category === 'start',
+      );
+      assert.ok(start, 'expected managed LLM start event');
+      assert.deepEqual(observed, [
+        ['intercept-before', start.uuid],
+        ['intercept-after', start.uuid],
+        ['provider-before', start.uuid],
+        ['provider-after', start.uuid],
+      ]);
+    } finally {
+      deregisterLlmExecutionIntercept('node_llm_exec_propagation_parent');
+      deregisterSubscriber('node_llm_exec_propagation_parent');
+    }
+  });
+
   it('request intercept', () => {
     registerLlmRequestIntercept('node_llm_req_int', 10, false, ({ name, request, annotated }) => {
       request.intercepted = true;
@@ -1599,6 +1649,31 @@ describe('LLM intercepts', () => {
     }
   });
 
+  it('default lazy stream expires its callback context when the provider ends', async () => {
+    const baseline = {
+      active: lib.scopeStackActive(),
+      parentUuid: lib.capturePropagationContext().parentUuid,
+    };
+    let resolveLateContext;
+    const lateContext = new Promise((resolve) => {
+      resolveLateContext = resolve;
+    });
+
+    const stream = await llmStreamCallExecute('node_default_lazy_stream_expiry', makeNative(), (wrapper) => {
+      lib.pushStreamChunk(wrapper.__nemo_relay_stream_id, { token: 'done' });
+      setImmediate(() => {
+        resolveLateContext({
+          active: lib.scopeStackActive(),
+          parentUuid: lib.capturePropagationContext().parentUuid,
+        });
+      });
+      lib.endStream(wrapper.__nemo_relay_stream_id);
+    });
+    assert.deepEqual(await stream.next(), { token: 'done' });
+    assert.equal(await stream.next(), null);
+    assert.deepEqual(await lateContext, baseline);
+  });
+
   it('stream execution next honors concurrent per-call scope-stack replacements', async () => {
     const firstStack = lib.createScopeStack();
     const secondStack = lib.createScopeStack();
diff --git a/crates/node/tests/tools_tests.mjs b/crates/node/tests/tools_tests.mjs
index 196dad76f..c4158cc62 100644
--- a/crates/node/tests/tools_tests.mjs
+++ b/crates/node/tests/tools_tests.mjs
@@ -1007,6 +1007,47 @@ describe('Tool intercepts', () => {
     }
   });
 
+  it('execution callbacks preserve the managed propagation parent across await', async () => {
+    const events = [];
+    const observed = [];
+    registerSubscriber('node_tool_exec_propagation_parent', (event) => events.push(event));
+    registerToolExecutionIntercept('node_tool_exec_propagation_parent', 10, async (args, next) => {
+      observed.push(['intercept-before', lib.capturePropagationContext().parentUuid]);
+      await new Promise((resolve) => setImmediate(resolve));
+      observed.push(['intercept-after', lib.capturePropagationContext().parentUuid]);
+      return { result: await next(args) };
+    });
+    try {
+      const result = await toolCallExecuteAsync('propagation_parent_tool', {}, async () => {
+        observed.push(['provider-before', lib.capturePropagationContext().parentUuid]);
+        await new Promise((resolve) => setImmediate(resolve));
+        observed.push(['provider-after', lib.capturePropagationContext().parentUuid]);
+        return { ok: true };
+      });
+      assert.deepEqual(result, { ok: true });
+      await waitForSubscriberCallbacks(() =>
+        events.some(
+          (event) =>
+            event.name === 'propagation_parent_tool' && event.kind === 'scope' && event.scope_category === 'end',
+        ),
+      );
+      const start = events.find(
+        (event) =>
+          event.name === 'propagation_parent_tool' && event.kind === 'scope' && event.scope_category === 'start',
+      );
+      assert.ok(start, 'expected managed tool start event');
+      assert.deepEqual(observed, [
+        ['intercept-before', start.uuid],
+        ['intercept-after', start.uuid],
+        ['provider-before', start.uuid],
+        ['provider-after', start.uuid],
+      ]);
+    } finally {
+      deregisterToolExecutionIntercept('node_tool_exec_propagation_parent');
+      deregisterSubscriber('node_tool_exec_propagation_parent');
+    }
+  });
+
   it('execution intercept rejects a detached next call after settlement', async () => {
     let releaseLateNext;
     const lateGate = new Promise((resolve) => {
@@ -1134,21 +1175,24 @@ describe('Tool intercepts', () => {
     const secondInstalled = new Promise((resolve) => {
       secondStackInstalled = resolve;
     });
+    let parentScope;
 
     registerToolExecutionIntercept('node_tool_exec_concurrent_scope_replacements', 10, async (args, next) => {
-      const first = (async () => {
-        lib.setThreadScopeStack(firstStack);
+      parentScope = lib.getHandle().uuid;
+      const first = lib.withScopeStack(firstStack, async () => {
         firstStackInstalled();
         await secondInstalled;
         return next({ ...args, branch: 'first' });
-      })();
-      const second = (async () => {
+      });
+      const second = lib.withScopeStack(secondStack, async () => {
         await firstInstalled;
-        lib.setThreadScopeStack(secondStack);
         secondStackInstalled();
         return next({ ...args, branch: 'second' });
-      })();
-      return { result: await Promise.all([first, second]) };
+      });
+      const branches = await Promise.all([first, second]);
+      assert.equal(lib.getHandle().uuid, parentScope);
+      const parent = await next({ ...args, branch: 'parent' });
+      return { result: [...branches, parent] };
     });
     try {
       const result = await toolCallExecuteAsync('concurrent_scope_replacements_tool', {}, async (args) => ({
@@ -1158,6 +1202,7 @@ describe('Tool intercepts', () => {
       assert.deepEqual(result, [
         { branch: 'first', scope: firstScope },
         { branch: 'second', scope: secondScope },
+        { branch: 'parent', scope: parentScope },
       ]);
     } finally {
       deregisterToolExecutionIntercept('node_tool_exec_concurrent_scope_replacements');
diff --git a/docs/about-nemo-relay/concepts/scopes.mdx b/docs/about-nemo-relay/concepts/scopes.mdx
index 10098f978..4aa9fe948 100644
--- a/docs/about-nemo-relay/concepts/scopes.mdx
+++ b/docs/about-nemo-relay/concepts/scopes.mdx
@@ -184,9 +184,19 @@ tokio::spawn(TASK_SCOPE_STACK.scope(stack, async {
 </Tab>
 <Tab title="Node.js" language="javascript">
 
-Node.js does not yet provide an asynchronous scope-forking helper. Its
-`withScopeStack()` API is synchronous and must not be used to bind work across
-an `await` boundary.
+Use `withScopeStack()` to bind each asynchronous branch to its isolated stack.
+The stack remains active until the branch's returned Promise settles, while the
+calling context is restored immediately.
+
+```javascript
+const [first, second] = await Promise.all([
+  withScopeStack(firstStack, async () => runWorker("first")),
+  withScopeStack(secondStack, async () => runWorker("second")),
+]);
+```
+
+`setThreadScopeStack()` mutates the current thread or async resource. Do not use
+it to isolate concurrent branches that begin in the same synchronous callback.
 
 </Tab>
 </Tabs>
diff --git a/python/tests/test_llm.py b/python/tests/test_llm.py
index d45dc77fd..7d5929c66 100644
--- a/python/tests/test_llm.py
+++ b/python/tests/test_llm.py
@@ -635,7 +635,16 @@ def provider(request):
             with pytest.raises(asyncio.CancelledError):
                 await execution
             await asyncio.wait_for(cancelled.wait(), timeout=1)
-            await subscribers.flush_async()
+
+            async def wait_for_end_publication():
+                while not any(
+                    isinstance(event, ScopeEvent) and event.name == "cancel_llm" and event.scope_category == "end"
+                    for event in events
+                ):
+                    await subscribers.flush_async()
+                    await asyncio.sleep(0)
+
+            await asyncio.wait_for(wait_for_end_publication(), timeout=1)
         finally:
             release.set()
             intercepts.deregister_llm_execution("py_llm_cancel_intercept")
@@ -688,7 +697,18 @@ async def generate():
             with pytest.raises(asyncio.CancelledError):
                 await execution
             await asyncio.wait_for(cancelled.wait(), timeout=1)
-            await subscribers.flush_async()
+
+            async def wait_for_end_publication():
+                while not any(
+                    isinstance(event, ScopeEvent)
+                    and event.name == "cancel_stream_llm"
+                    and event.scope_category == "end"
+                    for event in events
+                ):
+                    await subscribers.flush_async()
+                    await asyncio.sleep(0)
+
+            await asyncio.wait_for(wait_for_end_publication(), timeout=1)
         finally:
             release.set()
             intercepts.deregister_llm_stream_execution("py_llm_stream_cancel_intercept")

From 65c3f19b9f487a7d92fd960be8afefaa67c912b8 Mon Sep 17 00:00:00 2001
From: Will Killian <wkillian@nvidia.com>
Date: Thu, 30 Jul 2026 19:03:26 -0400
Subject: [PATCH 2/5] fix(runtime): include pending completions in subscriber
 flush

Signed-off-by: Will Killian <wkillian@nvidia.com>
---
 crates/core/src/api/llm.rs                    |  10 +-
 crates/core/src/api/runtime/state.rs          |   6 +
 .../src/api/runtime/subscriber_dispatcher.rs  | 112 ++++++++++++++++--
 crates/core/src/api/subscriber.rs             |   5 +-
 crates/core/src/api/tool.rs                   |  10 +-
 crates/core/src/plugin.rs                     |   2 +-
 .../tests/integration/middleware_tests.rs     |  36 +++++-
 .../tests/unit/subscriber_dispatcher_tests.rs |  75 +++++++++++-
 crates/node/README.md                         |  11 +-
 crates/node/src/api/mod.rs                    |   4 +-
 python/nemo_relay/_native.pyi                 |   2 +-
 python/nemo_relay/subscribers.py              |  11 +-
 python/tests/test_llm.py                      |  24 +---
 13 files changed, 249 insertions(+), 59 deletions(-)

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<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);
     }
 }
 
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<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());
     }
 }
 
diff --git a/crates/core/src/api/runtime/subscriber_dispatcher.rs b/crates/core/src/api/runtime/subscriber_dispatcher.rs
index 872d17aea..5c8707c91 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<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) -> Result<()> {
         if in_dispatcher_callback() {
             return Ok(());
         }
@@ -567,7 +614,10 @@ 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}"))
             })?;
@@ -577,6 +627,14 @@ mod native {
         Ok(())
     }
 
+    pub(super) fn flush_subscribers() -> Result<()> {
+        flush_subscribers_inner(true)
+    }
+
+    pub(super) fn flush_queued_subscribers() -> Result<()> {
+        flush_subscribers_inner(false)
+    }
+
     pub(super) fn in_dispatcher_callback() -> bool {
         IN_DISPATCHER.with(Cell::get)
             || ASYNC_PUBLICATION_BUFFER.try_with(|_| ()).is_ok()
@@ -740,11 +798,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 +849,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 +904,9 @@ mod native {
                 }
                 drop(permit);
             }
-            DispatcherMessage::Flush { .. } => unreachable!(),
+            DispatcherMessage::Flush { .. }
+            | DispatcherMessage::RegisterPending { .. }
+            | DispatcherMessage::CompletePending { .. } => unreachable!(),
         }
         state.complete_ready_flushes();
     }
@@ -1128,6 +1201,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 +1235,21 @@ 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 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<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
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<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);
     }
 }
 
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<LlmStreamExecutionNextFn>,
     request: Option<LlmRequest>,
@@ -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/unit/subscriber_dispatcher_tests.rs b/crates/core/tests/unit/subscriber_dispatcher_tests.rs
index 38e57064e..3339d8b84 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();
@@ -621,7 +683,10 @@ fn detached_sanitizer_tasks_cannot_inherit_a_later_publication_context() {
                     .send(publication_context::<String>().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..e51418719 100644
--- a/crates/node/README.md
+++ b/crates/node/README.md
@@ -100,11 +100,12 @@ 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 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 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 f0721f70b..d569e5514 100644
--- a/crates/node/src/api/mod.rs
+++ b/crates/node/src/api/mod.rs
@@ -3319,8 +3319,8 @@ pub fn deregister_subscriber(name: String) -> Result<bool> {
     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 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
diff --git a/python/nemo_relay/_native.pyi b/python/nemo_relay/_native.pyi
index 512d311c8..790fa9619 100644
--- a/python/nemo_relay/_native.pyi
+++ b/python/nemo_relay/_native.pyi
@@ -1976,7 +1976,7 @@ def deregister_subscriber(name: str) -> bool:
     ...
 
 def flush_subscribers() -> None:
-    """Wait for queued subscriber callbacks and their transitive publications.
+    """Wait for queued callbacks and registered managed terminal publications.
 
     Call this function outside subscribers, event sanitizers, conditional
     guardrails, and request or execution intercepts. The public Python wrapper
diff --git a/python/nemo_relay/subscribers.py b/python/nemo_relay/subscribers.py
index 92c9baabe..715ce440c 100644
--- a/python/nemo_relay/subscribers.py
+++ b/python/nemo_relay/subscribers.py
@@ -187,11 +187,13 @@ def deregister(name: str) -> bool:
 
 
 def flush() -> None:
-    """Wait for queued subscriber callbacks and their transitive publications.
+    """Wait for queued callbacks and registered managed terminal publications.
 
     Native NeMo Relay event APIs enqueue subscriber callbacks and return without
     waiting for observer work. Use this barrier in tests and shutdown paths when
-    captured subscriber output must be complete before continuing.
+    captured subscriber output must be complete before continuing. The barrier
+    also waits for tool, LLM, and guardrail-scope terminal events from managed
+    work that started before the flush.
 
     Call this function outside subscribers, event sanitizers, conditional
     guardrails, and request or execution intercepts. A queued tool or LLM
@@ -217,11 +219,12 @@ def flush() -> None:
 
 
 async def flush_async() -> None:
-    """Wait asynchronously for queued callbacks and transitive publications.
+    """Wait asynchronously for queued callbacks and managed terminal events.
 
     Use this barrier from an ``asyncio`` task. A process-local daemon bridge
     thread coalesces concurrent barriers and waits for the native dispatcher
-    without blocking the Python event loop.
+    without blocking the Python event loop. Managed tool, LLM, and guardrail
+    work registered before the flush is included through its terminal event.
     """
     if _publication_callback_active():
         return None
diff --git a/python/tests/test_llm.py b/python/tests/test_llm.py
index 7d5929c66..d45dc77fd 100644
--- a/python/tests/test_llm.py
+++ b/python/tests/test_llm.py
@@ -635,16 +635,7 @@ def provider(request):
             with pytest.raises(asyncio.CancelledError):
                 await execution
             await asyncio.wait_for(cancelled.wait(), timeout=1)
-
-            async def wait_for_end_publication():
-                while not any(
-                    isinstance(event, ScopeEvent) and event.name == "cancel_llm" and event.scope_category == "end"
-                    for event in events
-                ):
-                    await subscribers.flush_async()
-                    await asyncio.sleep(0)
-
-            await asyncio.wait_for(wait_for_end_publication(), timeout=1)
+            await subscribers.flush_async()
         finally:
             release.set()
             intercepts.deregister_llm_execution("py_llm_cancel_intercept")
@@ -697,18 +688,7 @@ async def generate():
             with pytest.raises(asyncio.CancelledError):
                 await execution
             await asyncio.wait_for(cancelled.wait(), timeout=1)
-
-            async def wait_for_end_publication():
-                while not any(
-                    isinstance(event, ScopeEvent)
-                    and event.name == "cancel_stream_llm"
-                    and event.scope_category == "end"
-                    for event in events
-                ):
-                    await subscribers.flush_async()
-                    await asyncio.sleep(0)
-
-            await asyncio.wait_for(wait_for_end_publication(), timeout=1)
+            await subscribers.flush_async()
         finally:
             release.set()
             intercepts.deregister_llm_stream_execution("py_llm_stream_cancel_intercept")

From 504d5e5d676e59a786b289e26cbe47bfe197be46 Mon Sep 17 00:00:00 2001
From: Will Killian <wkillian@nvidia.com>
Date: Thu, 30 Jul 2026 19:25:19 -0400
Subject: [PATCH 3/5] fix(node): await JavaScript subscribers when flushing

Signed-off-by: Will Killian <wkillian@nvidia.com>
---
 crates/node/README.md                 |  9 ++--
 crates/node/src/api/mod.rs            | 22 ++++----
 crates/node/src/callable.rs           | 55 ++++++++++++++++++-
 crates/node/tests/llm_tests.mjs       | 72 ++++---------------------
 crates/node/tests/scope_tests.mjs     | 30 ++++++++---
 crates/node/tests/test_support.mjs    | 24 ---------
 crates/node/tests/timestamp_tests.mjs | 10 +---
 crates/node/tests/tools_tests.mjs     | 76 ++++-----------------------
 crates/node/tests/typed_tests.mjs     | 57 ++------------------
 9 files changed, 115 insertions(+), 240 deletions(-)
 delete mode 100644 crates/node/tests/test_support.mjs

diff --git a/crates/node/README.md b/crates/node/README.md
index e51418719..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");
 }
 
@@ -101,11 +100,9 @@ main().catch((error) => {
 
 Native subscriber delivery is asynchronous. Awaiting `flushSubscribers()` drains
 the native dispatcher and waits for managed terminal publications registered
-before the call 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.
+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 d569e5514..ecd41f2da 100644
--- a/crates/node/src/api/mod.rs
+++ b/crates/node/src/api/mod.rs
@@ -3319,18 +3319,17 @@ pub fn deregister_subscriber(name: String) -> Result<bool> {
     core_subscriber_api::deregister_subscriber(&name).map_err(to_napi_err)
 }
 
-/// Return a Promise that resolves when native subscriber callbacks and managed
-/// terminal publications registered 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.
@@ -3342,10 +3341,13 @@ pub fn flush_subscribers(env: Env) -> Result<JsObject> {
             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<u64>,
+}
+
+fn js_subscriber_callbacks() -> &'static (Mutex<JsSubscriberCallbackState>, Condvar) {
+    static CALLBACKS: OnceLock<(Mutex<JsSubscriberCallbackState>, 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/tests/llm_tests.mjs b/crates/node/tests/llm_tests.mjs
index 4abdd98de..250aa0c7d 100644
--- a/crates/node/tests/llm_tests.mjs
+++ b/crates/node/tests/llm_tests.mjs
@@ -8,7 +8,6 @@ import { createRequire } from 'node:module';
 import { readFileSync } from 'node:fs';
 import path from 'node:path';
 import { fileURLToPath } from 'node:url';
-import { waitForSubscriberCallbacks } from './test_support.mjs';
 
 const require = createRequire(import.meta.url);
 const lib = require('../index.js');
@@ -177,10 +176,7 @@ describe('LLM lifecycle', () => {
       const native = makeNative();
       const handle = llmCall('evt_llm', native, null, null, null, null, null);
       llmCallEnd(handle, {}, null, null);
-      const deadline = Date.now() + 2000;
-      while (events.length < 2 && Date.now() < deadline) {
-        await new Promise((r) => setTimeout(r, 10));
-      }
+      await flushSubscribers();
       assert.ok(events.length >= 2, 'Expected at least 2 events');
     } finally {
       deregisterSubscriber('node_llm_evt_sub');
@@ -306,11 +302,7 @@ describe('LLM execute', () => {
         /llm status failure/,
       );
 
-      await waitForSubscriberCallbacks(
-        () =>
-          events.some((e) => e.name === 'exec_status_ok_llm' && e.scope_category === 'end') &&
-          events.some((e) => e.name === 'exec_status_error_llm' && e.scope_category === 'end'),
-      );
+      await flushSubscribers();
       const okEnd = events.find(
         (e) =>
           e.name === 'exec_status_ok_llm' && e.kind === 'scope' && e.category === 'llm' && e.scope_category === 'end',
@@ -405,11 +397,7 @@ describe('LLM guardrails', () => {
       assert.deepEqual(result, { ok: true });
       assert.equal(requestContextChecked, true);
       assert.equal(responseContextChecked, true);
-      await waitForSubscriberCallbacks(
-        () =>
-          events.some((event) => event.name === 'contextual_sanitize_llm' && event.scope_category === 'start') &&
-          events.some((event) => event.name === 'contextual_sanitize_llm' && event.scope_category === 'end'),
-      );
+      await flushSubscribers();
       const start = events.find(
         (event) => event.name === 'contextual_sanitize_llm' && event.scope_category === 'start',
       );
@@ -798,16 +786,7 @@ describe('LLM guardrails', () => {
       assert.deepEqual(result, {
         model: 'test-model',
       });
-      const deadline = Date.now() + 2000;
-      while (
-        !events.some(
-          (e) =>
-            e.name === 'san_req_evt_llm' && e.kind === 'scope' && e.category === 'llm' && e.scope_category === 'start',
-        ) &&
-        Date.now() < deadline
-      ) {
-        await new Promise((r) => setTimeout(r, 10));
-      }
+      await flushSubscribers();
       const start = events.find(
         (e) =>
           e.name === 'san_req_evt_llm' && e.kind === 'scope' && e.category === 'llm' && e.scope_category === 'start',
@@ -862,15 +841,7 @@ describe('LLM guardrails', () => {
     try {
       const request = makeNative();
       await llmCallExecute('llm_san_req_throw', request, () => ({ ok: true }), null, null, null, null, null);
-      await waitForSubscriberCallbacks(() =>
-        events.some(
-          (event) =>
-            event.name === 'llm_san_req_throw' &&
-            event.kind === 'scope' &&
-            event.category === 'llm' &&
-            event.scope_category === 'start',
-        ),
-      );
+      await flushSubscribers();
       const start = events.find(
         (event) =>
           event.name === 'llm_san_req_throw' &&
@@ -958,16 +929,7 @@ describe('LLM guardrails', () => {
       assert.deepEqual(result, {
         ok: true,
       });
-      const deadline = Date.now() + 2000;
-      while (
-        !events.some(
-          (e) =>
-            e.name === 'san_resp_evt_llm' && e.kind === 'scope' && e.category === 'llm' && e.scope_category === 'end',
-        ) &&
-        Date.now() < deadline
-      ) {
-        await new Promise((r) => setTimeout(r, 10));
-      }
+      await flushSubscribers();
       const end = events.find(
         (e) =>
           e.name === 'san_resp_evt_llm' && e.kind === 'scope' && e.category === 'llm' && e.scope_category === 'end',
@@ -992,15 +954,7 @@ describe('LLM guardrails', () => {
     try {
       const response = { ok: true };
       await llmCallExecute('llm_san_resp_throw', makeNative(), () => response, null, null, null, null, null);
-      await waitForSubscriberCallbacks(() =>
-        events.some(
-          (event) =>
-            event.name === 'llm_san_resp_throw' &&
-            event.kind === 'scope' &&
-            event.category === 'llm' &&
-            event.scope_category === 'end',
-        ),
-      );
+      await flushSubscribers();
       const end = events.find(
         (event) =>
           event.name === 'llm_san_resp_throw' &&
@@ -1152,12 +1106,7 @@ describe('LLM intercepts', () => {
         null,
       );
       assert.deepEqual(result, { ok: true });
-      await waitForSubscriberCallbacks(() =>
-        events.some(
-          (event) =>
-            event.name === 'propagation_parent_llm' && event.kind === 'scope' && event.scope_category === 'end',
-        ),
-      );
+      await flushSubscribers();
       const start = events.find(
         (event) =>
           event.name === 'propagation_parent_llm' && event.kind === 'scope' && event.scope_category === 'start',
@@ -2084,10 +2033,7 @@ describe('LLM event fields', () => {
         },
       );
 
-      const deadline = Date.now() + 2000;
-      while (events.filter((e) => e.name === 'field_llm').length < 2 && Date.now() < deadline) {
-        await new Promise((r) => setTimeout(r, 10));
-      }
+      await flushSubscribers();
 
       const start = events.find(
         (e) => e.name === 'field_llm' && e.kind === 'scope' && e.category === 'llm' && e.scope_category === 'start',
diff --git a/crates/node/tests/scope_tests.mjs b/crates/node/tests/scope_tests.mjs
index 3a184bfa7..c0a721893 100644
--- a/crates/node/tests/scope_tests.mjs
+++ b/crates/node/tests/scope_tests.mjs
@@ -4,7 +4,6 @@
 import { describe, it } from 'node:test';
 import assert from 'node:assert/strict';
 import { createRequire } from 'node:module';
-import { waitForSubscriberCallbacks } from './test_support.mjs';
 
 const require = createRequire(import.meta.url);
 const lib = require('../index.js');
@@ -98,7 +97,7 @@ describe('Scope operations', () => {
     try {
       const scope = pushScope('pop_metadata_scope', ScopeType.Agent, null, null, null, { a: 1, b: 2, c: 3 });
       popScope(scope, null, null, { c: 3.5, d: 4 });
-      await waitForSubscriberCallbacks(() => events.some((e) => e.name === 'pop_metadata_scope' && e.scope_category === 'end'));
+      await flushSubscribers();
 
       const end = events.find(
         (e) => e.name === 'pop_metadata_scope' && e.kind === 'scope' && e.scope_category === 'end',
@@ -202,7 +201,7 @@ describe('withScope', () => {
       await withScope('with_scope_ok_status', ScopeType.Function, () => ({ ok: true }), null, null, null, {
         caller: 'node',
       });
-      await waitForSubscriberCallbacks(() => events.some((e) => e.name === 'with_scope_ok_status' && e.scope_category === 'end'));
+      await flushSubscribers();
 
       const end = events.find(
         (e) => e.name === 'with_scope_ok_status' && e.kind === 'scope' && e.scope_category === 'end',
@@ -254,7 +253,7 @@ describe('withScope', () => {
           }),
         /node status failure/,
       );
-      await waitForSubscriberCallbacks(() => events.some((e) => e.name === 'with_scope_error_status' && e.scope_category === 'end'));
+      await flushSubscribers();
 
       const end = events.find(
         (e) => e.name === 'with_scope_error_status' && e.kind === 'scope' && e.scope_category === 'end',
@@ -349,7 +348,7 @@ describe('Subscribers', () => {
     try {
       const scope = pushScope('sub_test', ScopeType.Agent, null, null);
       popScope(scope);
-      await waitForSubscriberCallbacks(() => events.length > 0);
+      await flushSubscribers();
     } finally {
       deregisterSubscriber('node_event_collector');
     }
@@ -360,13 +359,28 @@ describe('Subscribers', () => {
     registerSubscriber('node_flush_collector', (e) => events.push(e));
     try {
       event('node_flush_mark', null, null, null);
-      await waitForSubscriberCallbacks(() => events.some((e) => e.kind === 'mark' && e.name === 'node_flush_mark'));
+      await flushSubscribers();
       assert.ok(events.some((e) => e.kind === 'mark' && e.name === 'node_flush_mark'));
     } finally {
       deregisterSubscriber('node_flush_collector');
     }
   });
 
+  it('flushSubscribers waits for JavaScript callbacks without inspecting their return values', async () => {
+    let called = false;
+    registerSubscriber('node_flush_js_callback', () => {
+      called = true;
+      return 1n;
+    });
+    try {
+      event('node_flush_js_callback_mark', null, null, null);
+      await flushSubscribers();
+      assert.equal(called, true);
+    } finally {
+      deregisterSubscriber('node_flush_js_callback');
+    }
+  });
+
   it('subscriber event properties', async () => {
     let captured = null;
     registerSubscriber('node_prop_collector', (e) => {
@@ -375,7 +389,7 @@ describe('Subscribers', () => {
     try {
       const scope = pushScope('prop_test', ScopeType.Function, null, null);
       popScope(scope);
-      await waitForSubscriberCallbacks(() => captured !== null);
+      await flushSubscribers();
       assert.ok(captured, 'Expected an event');
       assert.ok(typeof captured.uuid === 'string');
       assert.ok(typeof captured.timestamp === 'string');
@@ -398,7 +412,7 @@ describe('Subscribers', () => {
         },
         null,
       );
-      await waitForSubscriberCallbacks(() => events.some((e) => e.kind === 'mark'));
+      await flushSubscribers();
       const found = events.some((e) => e.kind === 'mark');
       assert.ok(found, 'Expected a Mark event');
     } finally {
diff --git a/crates/node/tests/test_support.mjs b/crates/node/tests/test_support.mjs
deleted file mode 100644
index 17274dcf0..000000000
--- a/crates/node/tests/test_support.mjs
+++ /dev/null
@@ -1,24 +0,0 @@
-// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-// SPDX-License-Identifier: Apache-2.0
-
-import { createRequire } from 'node:module';
-
-const require = createRequire(import.meta.url);
-const { flushSubscribers } = require('../index.js');
-
-export async function waitForSubscriberCallbacks(predicate, timeoutMs = 15000) {
-  await flushSubscribers();
-  // flushSubscribers() waits for Relay's Rust subscriber dispatcher, but JS
-  // subscriber callbacks are queued onto Node's event loop through N-API
-  // ThreadsafeFunction. Yield event-loop turns until the observed JS-side
-  // callback state is ready, with a timeout to avoid hanging the test forever.
-  const deadline = Date.now() + timeoutMs;
-  while (!predicate()) {
-    await flushSubscribers();
-    if (Date.now() >= deadline) {
-      throw new Error('timed out waiting for subscriber callbacks');
-    }
-    await new Promise((resolve) => setImmediate(resolve));
-  }
-  await flushSubscribers();
-}
diff --git a/crates/node/tests/timestamp_tests.mjs b/crates/node/tests/timestamp_tests.mjs
index 55243e6b4..086355f1f 100644
--- a/crates/node/tests/timestamp_tests.mjs
+++ b/crates/node/tests/timestamp_tests.mjs
@@ -10,6 +10,7 @@ const {
   ScopeType,
   deregisterSubscriber,
   event,
+  flushSubscribers,
   llmCall,
   llmCallEnd,
   popScope,
@@ -27,13 +28,6 @@ const parseTimestampMicros = (value) => {
   return Date.parse(`${base}${zone}`) * 1000 + Number(fraction.padEnd(6, '0').slice(0, 6));
 };
 
-async function waitForTimestampEvents(events) {
-  const deadline = Date.now() + 2000;
-  while (events.filter((event) => event.name.startsWith('node_ts_')).length < 7 && Date.now() < deadline) {
-    await new Promise((resolve) => setTimeout(resolve, 10));
-  }
-}
-
 test('manual lifecycle APIs accept optional timestamp arguments', async () => {
   const events = [];
   const timestamps = [
@@ -63,7 +57,7 @@ test('manual lifecycle APIs accept optional timestamp arguments', async () => {
   );
   llmCallEnd(llm, { ok: true }, null, null, timestamps[5]);
   popScope(scope, null, timestamps[6]);
-  await waitForTimestampEvents(events);
+  await flushSubscribers();
   deregisterSubscriber(subscriberName);
 
   assert.deepEqual(
diff --git a/crates/node/tests/tools_tests.mjs b/crates/node/tests/tools_tests.mjs
index c4158cc62..5bcb86358 100644
--- a/crates/node/tests/tools_tests.mjs
+++ b/crates/node/tests/tools_tests.mjs
@@ -4,7 +4,6 @@
 import { describe, it } from 'node:test';
 import assert from 'node:assert/strict';
 import { createRequire } from 'node:module';
-import { waitForSubscriberCallbacks } from './test_support.mjs';
 
 const require = createRequire(import.meta.url);
 const lib = require('../index.js');
@@ -135,10 +134,7 @@ describe('Tool lifecycle', () => {
     try {
       const handle = toolCall('evt_tool', {}, null, null, null, null);
       toolCallEnd(handle, {}, null, null);
-      const deadline = Date.now() + 2000;
-      while (events.length < 2 && Date.now() < deadline) {
-        await new Promise((r) => setTimeout(r, 10));
-      }
+      await flushSubscribers();
       assert.ok(events.length >= 2, 'Expected at least 2 events');
     } finally {
       deregisterSubscriber('node_tool_evt_sub');
@@ -180,10 +176,7 @@ describe('Tool lifecycle', () => {
         },
       );
 
-      const deadline = Date.now() + 2000;
-      while (events.filter((e) => e.name === 'field_tool').length < 2 && Date.now() < deadline) {
-        await new Promise((r) => setTimeout(r, 10));
-      }
+      await flushSubscribers();
 
       const start = events.find(
         (e) => e.name === 'field_tool' && e.kind === 'scope' && e.category === 'tool' && e.scope_category === 'start',
@@ -403,23 +396,7 @@ describe('Tool execute', () => {
         /tool status failure/,
       );
 
-      await waitForSubscriberCallbacks(
-        () =>
-          events.some(
-            (e) =>
-              e.name === 'exec_status_ok_tool' &&
-              e.kind === 'scope' &&
-              e.category === 'tool' &&
-              e.scope_category === 'end',
-          ) &&
-          events.some(
-            (e) =>
-              e.name === 'exec_status_error_tool' &&
-              e.kind === 'scope' &&
-              e.category === 'tool' &&
-              e.scope_category === 'end',
-          ),
-      );
+      await flushSubscribers();
       const okEnd = events.find(
         (e) =>
           e.name === 'exec_status_ok_tool' && e.kind === 'scope' && e.category === 'tool' && e.scope_category === 'end',
@@ -520,19 +497,7 @@ describe('Tool guardrails', () => {
       assert.deepEqual(result, {
         x: 1,
       });
-      const deadline = Date.now() + 2000;
-      while (
-        !events.some(
-          (e) =>
-            e.name === 'san_req_evt_tool' &&
-            e.kind === 'scope' &&
-            e.category === 'tool' &&
-            e.scope_category === 'start',
-        ) &&
-        Date.now() < deadline
-      ) {
-        await new Promise((r) => setTimeout(r, 10));
-      }
+      await flushSubscribers();
       const start = events.find(
         (e) =>
           e.name === 'san_req_evt_tool' && e.kind === 'scope' && e.category === 'tool' && e.scope_category === 'start',
@@ -579,16 +544,7 @@ describe('Tool guardrails', () => {
       assert.deepEqual(result, {
         ok: true,
       });
-      const deadline = Date.now() + 2000;
-      while (
-        !events.some(
-          (e) =>
-            e.name === 'san_resp_evt_tool' && e.kind === 'scope' && e.category === 'tool' && e.scope_category === 'end',
-        ) &&
-        Date.now() < deadline
-      ) {
-        await new Promise((r) => setTimeout(r, 10));
-      }
+      await flushSubscribers();
       const end = events.find(
         (e) =>
           e.name === 'san_resp_evt_tool' && e.kind === 'scope' && e.category === 'tool' && e.scope_category === 'end',
@@ -733,7 +689,7 @@ describe('Tool guardrails', () => {
     });
     try {
       await toolCallExecute('tool_sanitize_throw', { original: true }, (args) => args, null, null, null, null);
-      await waitForSubscriberCallbacks(() => events.some((event) => event.name === 'tool_sanitize_throw'));
+      await flushSubscribers();
       const start = events.find(
         (event) =>
           event.name === 'tool_sanitize_throw' &&
@@ -768,11 +724,7 @@ describe('Tool guardrails', () => {
     try {
       const handle = toolCall('node_manual_tool_flush', { original: true });
       toolCallEnd(handle, { ok: true });
-      await waitForSubscriberCallbacks(
-        () =>
-          events.some((event) => event.name === 'node_manual_tool_flush' && event.scope_category === 'start') &&
-          events.some((event) => event.name === 'node_manual_tool_flush' && event.scope_category === 'end'),
-      );
+      await flushSubscribers();
     } finally {
       deregisterToolSanitizeRequestGuardrail('node_manual_tool_flush_request');
       deregisterToolSanitizeResponseGuardrail('node_manual_tool_flush_response');
@@ -982,12 +934,7 @@ describe('Tool intercepts', () => {
       );
       assert.equal(result.original, false);
       assert.equal(result.wrapped, true);
-      await waitForSubscriberCallbacks(
-        () =>
-          events.some(
-            (event) => event.name === 'replaced_tool' && event.kind === 'scope' && event.scope_category === 'end',
-          ) && events.some((event) => event.name === 'node.tool.execution'),
-      );
+      await flushSubscribers();
       const start = events.find(
         (event) => event.name === 'replaced_tool' && event.kind === 'scope' && event.scope_category === 'start',
       );
@@ -1025,12 +972,7 @@ describe('Tool intercepts', () => {
         return { ok: true };
       });
       assert.deepEqual(result, { ok: true });
-      await waitForSubscriberCallbacks(() =>
-        events.some(
-          (event) =>
-            event.name === 'propagation_parent_tool' && event.kind === 'scope' && event.scope_category === 'end',
-        ),
-      );
+      await flushSubscribers();
       const start = events.find(
         (event) =>
           event.name === 'propagation_parent_tool' && event.kind === 'scope' && event.scope_category === 'start',
diff --git a/crates/node/tests/typed_tests.mjs b/crates/node/tests/typed_tests.mjs
index c15156a08..668e738e1 100644
--- a/crates/node/tests/typed_tests.mjs
+++ b/crates/node/tests/typed_tests.mjs
@@ -51,6 +51,7 @@ const {
   deregisterLlmExecutionIntercept,
   registerSubscriber,
   deregisterSubscriber,
+  flushSubscribers,
 } = lib;
 
 // ===========================================================================
@@ -372,19 +373,7 @@ describe('typedToolExecute', () => {
         },
       );
 
-      const deadline = Date.now() + 2000;
-      while (
-        !events.some(
-          (event) =>
-            event.kind === 'scope' &&
-            event.category === 'llm' &&
-            event.scope_category === 'start' &&
-            event.name === 'typed_node_falsy_opts_llm',
-        ) &&
-        Date.now() < deadline
-      ) {
-        await new Promise((resolve) => setTimeout(resolve, 10));
-      }
+      await flushSubscribers();
 
       const startEvent = events.find(
         (event) =>
@@ -672,19 +661,7 @@ describe('typedLlmExecute', () => {
 
       assert.equal(result.content[0].text, 'Anthropic hello');
 
-      const deadline = Date.now() + 2000;
-      while (
-        !events.some(
-          (event) =>
-            event.kind === 'scope' &&
-            event.category === 'llm' &&
-            event.scope_category === 'end' &&
-            event.name === 'typed_anthropic_codec_llm',
-        ) &&
-        Date.now() < deadline
-      ) {
-        await new Promise((resolve) => setTimeout(resolve, 10));
-      }
+      await flushSubscribers();
 
       const endEvent = events.find(
         (event) =>
@@ -952,19 +929,7 @@ describe('typedLlmStreamExecute', () => {
       assert.equal(interceptedAnnotated.instructions, 'Be terse.');
       assert.equal(interceptedAnnotated.messages[0].role, 'user');
 
-      const deadline = Date.now() + 2000;
-      while (
-        !events.some(
-          (event) =>
-            event.kind === 'scope' &&
-            event.category === 'llm' &&
-            event.scope_category === 'end' &&
-            event.name === 'typed_responses_stream_llm',
-        ) &&
-        Date.now() < deadline
-      ) {
-        await new Promise((resolve) => setTimeout(resolve, 10));
-      }
+      await flushSubscribers();
 
       const endEvent = events.find(
         (event) =>
@@ -1016,19 +981,7 @@ describe('typedLlmStreamExecute', () => {
         // Drain stream
       }
 
-      const deadline = Date.now() + 2000;
-      while (
-        !events.some(
-          (event) =>
-            event.kind === 'scope' &&
-            event.category === 'llm' &&
-            event.scope_category === 'end' &&
-            event.name === 'typed_stream_bad_response_codec_llm',
-        ) &&
-        Date.now() < deadline
-      ) {
-        await new Promise((resolve) => setTimeout(resolve, 10));
-      }
+      await flushSubscribers();
 
       const endEvent = events.find(
         (event) =>

From 040c3be44e3c2ce8ebf8f073b5287c25a984c2db Mon Sep 17 00:00:00 2001
From: Will Killian <wkillian@nvidia.com>
Date: Thu, 30 Jul 2026 20:09:16 -0400
Subject: [PATCH 4/5] test(rust): synchronize transitive FIFO regression

Signed-off-by: Will Killian <wkillian@nvidia.com>
---
 crates/core/tests/unit/subscriber_dispatcher_tests.rs | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/crates/core/tests/unit/subscriber_dispatcher_tests.rs b/crates/core/tests/unit/subscriber_dispatcher_tests.rs
index 3339d8b84..97035ed68 100644
--- a/crates/core/tests/unit/subscriber_dispatcher_tests.rs
+++ b/crates/core/tests/unit/subscriber_dispatcher_tests.rs
@@ -10,7 +10,7 @@ use super::{EventSubscriberFn, publication_context};
 use crate::api::registry::RegistryRecord;
 use crate::api::runtime::EventSanitizeFn;
 use crate::api::runtime::scope_stack::current_scope_stack;
-use std::sync::{Arc, Mutex, mpsc};
+use std::sync::{Arc, Barrier, Mutex, mpsc};
 
 #[test]
 fn flush_waits_for_active_but_not_later_publication_barriers() {
@@ -363,6 +363,8 @@ 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 = Arc::new(Barrier::new(2));
+    let release_outer = Arc::new(Barrier::new(2));
     let event = |uuid: &str, name: &str| {
         serde_json::from_value(serde_json::json!({
             "kind": "mark",
@@ -414,11 +416,15 @@ 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 outer_started = Arc::clone(&outer_started);
+        let release_outer = Arc::clone(&release_outer);
         Arc::new(move |event| {
             delivered
                 .lock()
                 .unwrap_or_else(|error| error.into_inner())
                 .push(event.name().to_string());
+            outer_started.wait();
+            release_outer.wait();
             assert!(enqueue_dispatch_message(DispatcherMessage::Deliver {
                 event: Box::new(
                     serde_json::from_value(serde_json::json!({
@@ -460,6 +466,7 @@ fn flush_waits_for_transitive_subscriber_publications_without_reordering() {
             lineage: None,
         })
         .unwrap();
+    outer_started.wait();
     sender
         .send(DispatcherMessage::Deliver {
             event: Box::new(event("019c1df6-4a57-7000-8000-000000000011", "later")),
@@ -471,6 +478,7 @@ fn flush_waits_for_transitive_subscriber_publications_without_reordering() {
             lineage: None,
         })
         .unwrap();
+    release_outer.wait();
 
     flush_subscribers().unwrap();
     assert_eq!(

From 22763d05a259b445c2c3aed7e7f9b60e7525f316 Mon Sep 17 00:00:00 2001
From: Will Killian <wkillian@nvidia.com>
Date: Thu, 30 Jul 2026 20:50:35 -0400
Subject: [PATCH 5/5] fix(node): expire replacement callback contexts

Signed-off-by: Will Killian <wkillian@nvidia.com>
---
 .../src/api/runtime/subscriber_dispatcher.rs  | 20 ++++-
 .../core/tests/integration/pipeline_tests.rs  |  8 +-
 crates/node/src/callback_factory.rs           | 83 +++++++++++++------
 crates/node/tests/tools_tests.mjs             | 26 ++++++
 4 files changed, 105 insertions(+), 32 deletions(-)

diff --git a/crates/core/src/api/runtime/subscriber_dispatcher.rs b/crates/core/src/api/runtime/subscriber_dispatcher.rs
index 5c8707c91..035831900 100644
--- a/crates/core/src/api/runtime/subscriber_dispatcher.rs
+++ b/crates/core/src/api/runtime/subscriber_dispatcher.rs
@@ -595,7 +595,7 @@ mod native {
         })
     }
 
-    fn flush_subscribers_inner(include_pending: bool) -> Result<()> {
+    fn flush_subscribers_inner(include_pending: bool, started: Option<Sender<()>>) -> Result<()> {
         if in_dispatcher_callback() {
             return Ok(());
         }
@@ -621,6 +621,9 @@ mod native {
             .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}")))?;
@@ -628,11 +631,15 @@ mod native {
     }
 
     pub(super) fn flush_subscribers() -> Result<()> {
-        flush_subscribers_inner(true)
+        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)
+        flush_subscribers_inner(false, None)
     }
 
     pub(super) fn in_dispatcher_callback() -> bool {
@@ -1242,6 +1249,13 @@ 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
diff --git a/crates/core/tests/integration/pipeline_tests.rs b/crates/core/tests/integration/pipeline_tests.rs
index f14d66c1f..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};
@@ -2089,12 +2090,13 @@ async fn test_dropped_stream_end_keeps_fifo_position_before_later_mark() {
     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_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");
+        .expect("subscriber flush request was not queued before waiting on the pending stream END");
     assert!(
         matches!(
             flush_done_rx.recv_timeout(std::time::Duration::from_millis(50)),
diff --git a/crates/node/src/callback_factory.rs b/crates/node/src/callback_factory.rs
index 3c6d82d29..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_v10";
+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;
@@ -94,19 +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;
-      token.propagationParentUuid = undefined;
+      expireCallbackStore(token);
     };
     const safeNext = next === undefined
       ? undefined
@@ -199,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));
       };
     },
@@ -230,16 +270,13 @@ const CALLBACK_FACTORIES_SOURCE: &str = r#"(() => {
 
     withCallbackScopeStack(scopeStack, fn) {
       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: current?.propagationParentUuid,
-      };
-      const expire = () => {
-        token.scopeStack = null;
-        token.propagationParentUuid = undefined;
-      };
+        current?.propagationParentUuid,
+      );
+      const expire = () => expireCallbackStore(token);
       let value;
       try {
         value = eventSanitizerContext.run(token, fn);
@@ -263,12 +300,7 @@ 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;
     },
 
@@ -277,8 +309,7 @@ const CALLBACK_FACTORIES_SOURCE: &str = r#"(() => {
       if (current === undefined) {
         return;
       }
-      current.scopeStack = null;
-      current.propagationParentUuid = undefined;
+      expireCallbackStore(current);
     },
   };
 })()"#;
diff --git a/crates/node/tests/tools_tests.mjs b/crates/node/tests/tools_tests.mjs
index 5bcb86358..d40a7f6e2 100644
--- a/crates/node/tests/tools_tests.mjs
+++ b/crates/node/tests/tools_tests.mjs
@@ -990,6 +990,32 @@ describe('Tool intercepts', () => {
     }
   });
 
+  it('execution settlement expires scope replacements inherited by detached work', async () => {
+    const baseline = {
+      active: lib.scopeStackActive(),
+      parentUuid: lib.capturePropagationContext().parentUuid,
+    };
+    const replacementStack = lib.createScopeStack();
+    let releaseLateContext;
+    const lateGate = new Promise((resolve) => {
+      releaseLateContext = resolve;
+    });
+    let lateContext;
+
+    await toolCallExecuteAsync('scope_replacement_expiry_tool', {}, async () => {
+      lib.setThreadScopeStack(replacementStack);
+      lateContext = lateGate.then(() => ({
+        active: lib.scopeStackActive(),
+        parentUuid: lib.capturePropagationContext().parentUuid,
+      }));
+      await new Promise((resolve) => setImmediate(resolve));
+      return { ok: true };
+    });
+
+    releaseLateContext();
+    assert.deepEqual(await lateContext, baseline);
+  });
+
   it('execution intercept rejects a detached next call after settlement', async () => {
     let releaseLateNext;
     const lateGate = new Promise((resolve) => {