fix: preserve Go scope stacks across OS threads (NVIDIA#100)

#### Overview

Fix intermittent Windows Go test failures caused by Go goroutines moving between OS threads while NeMo Flow scope state is stored in Rust thread-local fallback storage.

- [x] I confirm this contribution is my own work, or I have the right to submit it under this project's license.
- [x] I searched existing issues and open pull requests, and this does not duplicate existing work.

#### Details

- Add Rust runtime helpers to capture and restore the current thread-local scope stack binding.
- Expose the capture/restore helpers through the FFI API and generated header.
- Update Go `ScopeStack.Run` to restore the previous Rust thread-local binding before releasing the locked OS thread.
- Run the flaky ATOF Go lifecycle test inside an explicit `ScopeStack.Run` so `PushScope`, `EmitEvent`, and `PopScope` stay on one OS thread.
- Extend Go coverage to assert `ScopeStackActive()` is restored after `Run`.

#### Where should the reviewer start?

Start with `go/nemo_flow/nemo_flow.go` for the Go runtime behavior, then review `crates/core/src/api/runtime/scope_stack.rs` and `crates/ffi/src/api/scope_stack.rs` for the capture/restore bridge. The regression coverage is in `go/nemo_flow/atof_test.go` and `go/nemo_flow/context_test.go`.

#### Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)

- Relates to: none




## Summary by CodeRabbit

* **New Features**
  * Added scope-stack binding capture and restoration functionality for managing thread-local state preservation.
  * Extended FFI support for thread scope stack operations.

* **Tests**
  * Updated scope-stack lifecycle tests to use the new binding management approach.



[![Review Change Stack](https://storage.googleapis.com/coderabbit_public_assets/review-stack-in-coderabbit-ui.svg)](https://app.coderabbit.ai/change-stack/NVIDIA/NeMo-Flow/pull/100)

Authors:
  - Will Killian (https://github.com/willkill07)

Approvers:
  - David Gardner (https://github.com/dagardner-nv)

URL: NVIDIA#100

Author: willkill07

crates/core/src/api/runtime.rs

@@ -17,8 +17,9 @@ pub use callbacks::{
 };
 pub use global::global_context;
 pub use scope_stack::{
-    ScopeStack, ScopeStackHandle, TASK_SCOPE_STACK, create_scope_stack, current_scope_stack,
-    propagate_scope_to_thread, scope_stack_active, set_thread_scope_stack, sync_thread_scope_stack,
-    task_scope_push, task_scope_remove, task_scope_top,
+    ScopeStack, ScopeStackHandle, TASK_SCOPE_STACK, ThreadScopeStackBinding,
+    capture_thread_scope_stack, create_scope_stack, current_scope_stack, propagate_scope_to_thread,
+    restore_thread_scope_stack, scope_stack_active, set_thread_scope_stack,
+    sync_thread_scope_stack, task_scope_push, task_scope_remove, task_scope_top,
 };
 pub use state::NemoFlowContextState;

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

@@ -236,6 +236,16 @@ impl Default for ScopeStack {
 /// concurrent readers.
 pub type ScopeStackHandle = Arc<RwLock<ScopeStack>>;
 
+/// Captured thread-local scope stack binding.
+///
+/// This preserves both the visible scope stack handle and whether it was
+/// explicitly installed on the current thread.
+#[derive(Clone)]
+pub struct ThreadScopeStackBinding {
+    stack: ScopeStackHandle,
+    explicit: bool,
+}
+
 /// Create a new scope stack handle with an implicit root scope.
 ///
 /// The returned handle wraps a freshly initialized [`ScopeStack`] inside an
@@ -297,6 +307,23 @@ pub fn set_thread_scope_stack(handle: ScopeStackHandle) {
     THREAD_SCOPE_STACK_EXPLICIT.with(|flag| flag.set(true));
 }
 
+/// Capture the current thread-local scope stack binding.
+///
+/// This is intended for foreign runtimes that temporarily bind a scope stack to
+/// an OS thread and need to restore the exact previous state before releasing
+/// that thread back to their scheduler.
+pub fn capture_thread_scope_stack() -> ThreadScopeStackBinding {
+    let stack = THREAD_SCOPE_STACK.with(|stack| stack.borrow().clone());
+    let explicit = THREAD_SCOPE_STACK_EXPLICIT.with(|flag| flag.get());
+    ThreadScopeStackBinding { stack, explicit }
+}
+
+/// Restore a previously captured thread-local scope stack binding.
+pub fn restore_thread_scope_stack(binding: ThreadScopeStackBinding) {
+    THREAD_SCOPE_STACK.with(|stack| *stack.borrow_mut() = binding.stack);
+    THREAD_SCOPE_STACK_EXPLICIT.with(|flag| flag.set(binding.explicit));
+}
+
 /// Synchronize the thread-local scope stack without marking it explicit.
 ///
 /// This updates the thread-local slot used by native runtime code while

crates/ffi/nemo_flow.h

@@ -184,6 +184,11 @@ typedef struct FfiScopeStack FfiScopeStack;
  */
 typedef struct FfiStream FfiStream;
 
+/**
+ * Opaque handle to a captured thread-local scope stack binding.
+ */
+typedef struct FfiThreadScopeStackBinding FfiThreadScopeStackBinding;
+
 /**
  * Opaque handle representing an active tool call.
  */
@@ -1770,6 +1775,29 @@ NemoFlowStatus nemo_flow_scope_stack_create(struct FfiScopeStack **out);
  */
 NemoFlowStatus nemo_flow_scope_stack_set_thread(const struct FfiScopeStack *stack);
 
+/**
+ * Capture the current thread-local scope stack binding.
+ *
+ * The returned binding must be restored with
+ * `nemo_flow_scope_stack_restore_thread`.
+ *
+ * # Parameters
+ * - `out`: On success, receives a heap-allocated binding handle.
+ *
+ * # Safety
+ * `out` must be a valid, non-null pointer.
+ */
+NemoFlowStatus nemo_flow_scope_stack_capture_thread(struct FfiThreadScopeStackBinding **out);
+
+/**
+ * Restore and free a captured thread-local scope stack binding.
+ *
+ * # Safety
+ * `binding` must be a valid pointer returned by
+ * `nemo_flow_scope_stack_capture_thread`.
+ */
+NemoFlowStatus nemo_flow_scope_stack_restore_thread(struct FfiThreadScopeStackBinding *binding);
+
 /**
  * Returns whether the current execution context has an explicitly-initialized
  * scope stack.

crates/ffi/src/api/mod.rs

@@ -36,8 +36,8 @@ use crate::error::{
 };
 use crate::types::{
     FfiAtifExporter, FfiAtofExporter, FfiCodecHandle, FfiLLMHandle, FfiOpenInferenceSubscriber,
-    FfiOpenTelemetrySubscriber, FfiPluginContext, FfiScopeHandle, FfiScopeStack, FfiToolHandle,
-    NemoFlowScopeType,
+    FfiOpenTelemetrySubscriber, FfiPluginContext, FfiScopeHandle, FfiScopeStack,
+    FfiThreadScopeStackBinding, FfiToolHandle, NemoFlowScopeType,
 };
 pub use crate::types::{nemo_flow_openinference_subscriber_free, nemo_flow_otel_subscriber_free};
 use libc::c_char;
@@ -46,8 +46,8 @@ use nemo_flow::api::llm::{LlmAttributes, LlmRequest};
 use nemo_flow::api::registry as core_registry_api;
 use nemo_flow::api::runtime::{LlmExecutionNextFn, LlmStreamExecutionNextFn, ToolExecutionNextFn};
 use nemo_flow::api::runtime::{
-    TASK_SCOPE_STACK, create_scope_stack, current_scope_stack, scope_stack_active,
-    set_thread_scope_stack,
+    TASK_SCOPE_STACK, capture_thread_scope_stack, create_scope_stack, current_scope_stack,
+    restore_thread_scope_stack, scope_stack_active, set_thread_scope_stack,
 };
 use nemo_flow::api::scope as core_scope_api;
 use nemo_flow::api::scope::ScopeAttributes;

crates/ffi/src/api/scope_stack.rs

@@ -2,7 +2,8 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use super::{
-    FfiScopeStack, NemoFlowStatus, clear_last_error, create_scope_stack, scope_stack_active,
+    FfiScopeStack, FfiThreadScopeStackBinding, NemoFlowStatus, capture_thread_scope_stack,
+    clear_last_error, create_scope_stack, restore_thread_scope_stack, scope_stack_active,
     set_last_error, set_thread_scope_stack,
 };
 
@@ -75,6 +76,49 @@ pub unsafe extern "C" fn nemo_flow_scope_stack_set_thread(
     NemoFlowStatus::Ok
 }
 
+/// Capture the current thread-local scope stack binding.
+///
+/// The returned binding must be restored with
+/// `nemo_flow_scope_stack_restore_thread`.
+///
+/// # Parameters
+/// - `out`: On success, receives a heap-allocated binding handle.
+///
+/// # Safety
+/// `out` must be a valid, non-null pointer.
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn nemo_flow_scope_stack_capture_thread(
+    out: *mut *mut FfiThreadScopeStackBinding,
+) -> NemoFlowStatus {
+    clear_last_error();
+    if out.is_null() {
+        set_last_error("out pointer is null");
+        return NemoFlowStatus::NullPointer;
+    }
+    let binding = capture_thread_scope_stack();
+    unsafe { *out = Box::into_raw(Box::new(FfiThreadScopeStackBinding(binding))) };
+    NemoFlowStatus::Ok
+}
+
+/// Restore and free a captured thread-local scope stack binding.
+///
+/// # Safety
+/// `binding` must be a valid pointer returned by
+/// `nemo_flow_scope_stack_capture_thread`.
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn nemo_flow_scope_stack_restore_thread(
+    binding: *mut FfiThreadScopeStackBinding,
+) -> NemoFlowStatus {
+    clear_last_error();
+    if binding.is_null() {
+        set_last_error("binding pointer is null");
+        return NemoFlowStatus::NullPointer;
+    }
+    let binding = unsafe { Box::from_raw(binding) };
+    restore_thread_scope_stack(binding.0);
+    NemoFlowStatus::Ok
+}
+
 /// Returns whether the current execution context has an explicitly-initialized
 /// scope stack.
 ///

crates/ffi/src/types/mod.rs

@@ -11,7 +11,7 @@
 //! with their corresponding `nemo_flow_*_free` function.
 
 use libc::c_char;
-use nemo_flow::api::runtime::ScopeStackHandle;
+use nemo_flow::api::runtime::{ScopeStackHandle, ThreadScopeStackBinding};
 use nemo_flow::plugin::PluginRegistrationContext;
 use serde_json::Value as Json;
 
@@ -48,6 +48,8 @@ pub struct FfiLLMRequest(pub LlmRequest);
 pub struct FfiEvent(pub Event);
 /// Opaque handle to an isolated scope stack for per-request/per-task isolation.
 pub struct FfiScopeStack(pub ScopeStackHandle);
+/// Opaque handle to a captured thread-local scope stack binding.
+pub struct FfiThreadScopeStackBinding(pub ThreadScopeStackBinding);
 /// Opaque ATIF exporter handle.
 pub struct FfiAtifExporter(pub nemo_flow::observability::atif::AtifExporter);
 /// Opaque ATOF JSONL exporter handle.

go/nemo_flow/atof_test.go

@@ -44,10 +44,23 @@ func TestAtofExporterLifecycleWritesRawJSONL(t *testing.T) {
 
 	name := "go_atof_" + time.Now().Format("150405.000000")
 	requireNoError(t, exporter.Register(name), "Register failed")
-	handle, err := PushScope("atof_scope", ScopeTypeAgent, WithInput(json.RawMessage(`{"scope":true}`)))
-	requireNoError(t, err, "PushScope failed")
-	requireNoError(t, EmitEvent("atof_mark", WithEventParent(handle), WithEventData(json.RawMessage(`{"step":1}`))), "EmitEvent failed")
-	requireNoError(t, PopScope(handle, WithOutput(json.RawMessage(`{"done":true}`))), "PopScope failed")
+	stack, err := NewScopeStack()
+	requireNoError(t, err, "NewScopeStack failed")
+	defer stack.Close()
+	var runErr error
+	stack.Run(func() {
+		handle, err := PushScope("atof_scope", ScopeTypeAgent, WithInput(json.RawMessage(`{"scope":true}`)))
+		if err != nil {
+			runErr = err
+			return
+		}
+		if err := EmitEvent("atof_mark", WithEventParent(handle), WithEventData(json.RawMessage(`{"step":1}`))); err != nil {
+			runErr = err
+			return
+		}
+		runErr = PopScope(handle, WithOutput(json.RawMessage(`{"done":true}`)))
+	})
+	requireNoError(t, runErr, "scope lifecycle failed")
 	requireNoError(t, exporter.Deregister(name), "Deregister failed")
 	requireNoError(t, exporter.Deregister(name), "repeated Deregister should be safe")
 	requireNoError(t, exporter.ForceFlush(), "ForceFlush failed")

go/nemo_flow/context_test.go

@@ -6,6 +6,7 @@ package nemo_flow
 import (
 	"encoding/json"
 	"fmt"
+	"runtime"
 	"sync"
 	"testing"
 )
@@ -123,13 +124,21 @@ func TestScopeStackActiveInsideRun(t *testing.T) {
 	defer stack.Close()
 
 	var active bool
+	runtime.LockOSThread()
+	defer runtime.UnlockOSThread()
+	if ScopeStackActive() {
+		t.Fatal("expected ScopeStackActive() to be false before Run")
+	}
 	stack.Run(func() {
 		active = ScopeStackActive()
 	})
 
 	if !active {
 		t.Error("expected ScopeStackActive() to be true inside Run")
 	}
+	if ScopeStackActive() {
+		t.Error("expected ScopeStackActive() to be restored after Run")
+	}
 }
 
 func TestScopeStackRunIsolation(t *testing.T) {

go/nemo_flow/nemo_flow.go

@@ -31,6 +31,7 @@ package nemo_flow
 
 typedef struct FfiScopeHandle FfiScopeHandle;
 typedef struct FfiScopeStack FfiScopeStack;
+typedef struct FfiThreadScopeStackBinding FfiThreadScopeStackBinding;
 typedef struct FfiToolHandle FfiToolHandle;
 typedef struct FfiLLMHandle FfiLLMHandle;
 typedef struct FfiLLMRequest FfiLLMRequest;
@@ -211,6 +212,8 @@ extern void nemo_flow_string_free(char* ptr);
 // Scope stack isolation
 extern int32_t nemo_flow_scope_stack_create(FfiScopeStack** out);
 extern int32_t nemo_flow_scope_stack_set_thread(const FfiScopeStack* stack);
+extern int32_t nemo_flow_scope_stack_capture_thread(FfiThreadScopeStackBinding** out);
+extern int32_t nemo_flow_scope_stack_restore_thread(FfiThreadScopeStackBinding* binding);
 extern _Bool nemo_flow_scope_stack_active(void);
 extern void nemo_flow_scope_stack_free(FfiScopeStack* ptr);
 
@@ -1500,8 +1503,22 @@ func (s *ScopeStack) Close() {
 // This is the canonical way to propagate a scope stack to a worker goroutine.
 func (s *ScopeStack) Run(fn func()) {
 	runtime.LockOSThread()
-	defer runtime.UnlockOSThread()
-	C.nemo_flow_scope_stack_set_thread(s.ptr)
+	var binding *C.FfiThreadScopeStackBinding
+	if err := checkStatus(C.nemo_flow_scope_stack_capture_thread(&binding)); err != nil {
+		runtime.UnlockOSThread()
+		panic(err)
+	}
+	defer func() {
+		status := C.nemo_flow_scope_stack_restore_thread(binding)
+		if err := checkStatus(status); err != nil {
+			runtime.UnlockOSThread()
+			panic(err)
+		}
+		runtime.UnlockOSThread()
+	}()
+	if err := checkStatus(C.nemo_flow_scope_stack_set_thread(s.ptr)); err != nil {
+		panic(err)
+	}
 	fn()
 }