From e5db478faca03bc773fde2eee678f7ca6b0d9f7e Mon Sep 17 00:00:00 2001 From: Will Killian Date: Mon, 27 Jul 2026 10:21:27 -0400 Subject: [PATCH 1/5] feat: add cross-boundary Relay propagation Signed-off-by: Will Killian --- crates/core/src/api/llm.rs | 19 ++- crates/core/src/api/runtime.rs | 10 +- crates/core/src/api/runtime/scope_stack.rs | 136 ++++++++++++++++++ crates/core/src/api/tool.rs | 4 +- .../core/src/observability/openinference.rs | 31 +++- crates/core/src/observability/otel.rs | 31 +++- .../integration/context_isolation_tests.rs | 38 ++++- crates/ffi/nemo_relay.h | 33 +++++ crates/ffi/src/api/scope_stack.rs | 123 +++++++++++++++- crates/node/src/api/mod.rs | 92 +++++++++++- crates/node/tests/context_tests.mjs | 23 +++ crates/python/src/py_api/mod.rs | 44 +++++- crates/python/src/py_types/core.rs | 47 +++++- crates/python/src/py_types/mod.rs | 1 + docs/about-nemo-relay/concepts/scopes.mdx | 18 +++ .../observability/opentelemetry.mdx | 6 + go/nemo_relay/context_test.go | 24 ++++ go/nemo_relay/nemo_relay.go | 59 ++++++++ python/nemo_relay/__init__.py | 45 ++++++ python/nemo_relay/__init__.pyi | 7 + python/nemo_relay/_native.pyi | 13 ++ python/tests/test_context_isolation.py | 15 ++ 22 files changed, 786 insertions(+), 33 deletions(-) diff --git a/crates/core/src/api/llm.rs b/crates/core/src/api/llm.rs index 29e4655a9..ea0b49af0 100644 --- a/crates/core/src/api/llm.rs +++ b/crates/core/src/api/llm.rs @@ -22,6 +22,7 @@ use crate::api::runtime::global_context; use crate::api::runtime::{ EventSubscriberFn, LlmCollectorFn, LlmExecutionNextFn, LlmFinalizerFn, LlmJsonStream, LlmSanitizeRequestContext, LlmSanitizeResponseContext, LlmStreamExecutionNextFn, + with_active_event_uuid, }; use crate::api::runtime::{ScopeStackHandle, current_scope_stack}; use crate::api::scope::event; @@ -1025,7 +1026,9 @@ pub async fn llm_call_execute(params: LlmCallExecuteParams) -> Result { emit_optimization_marks(&handle, &lifecycle_subscribers); let execution_name = name.clone(); - let execution = + let event_uuid = handle.uuid; + let execution = with_active_event_uuid( + event_uuid, scope_llm_optimization_recorder(handle.optimization_recorder.clone(), async move { let execution = { let scope_stack = current_scope_stack(); @@ -1040,8 +1043,9 @@ pub async fn llm_call_execute(params: LlmCallExecuteParams) -> Result { state.llm_build_execution_chain(&execution_name, func, &scope_locals) }; execution(intercepted_request).await - }) - .await; + }), + ) + .await; match execution { Ok(response) => { @@ -1215,7 +1219,9 @@ pub async fn llm_stream_call_execute(params: LlmStreamCallExecuteParams) -> Resu emit_optimization_marks(&handle, &lifecycle_subscribers); let execution_name = name.clone(); - let execution = + let event_uuid = handle.uuid; + let execution = with_active_event_uuid( + event_uuid, scope_llm_optimization_recorder(handle.optimization_recorder.clone(), async move { let execution = { let scope_stack = current_scope_stack(); @@ -1230,8 +1236,9 @@ pub async fn llm_stream_call_execute(params: LlmStreamCallExecuteParams) -> Resu state.llm_stream_build_execution_chain(&execution_name, func, &scope_locals) }; execution(intercepted_request).await - }) - .await; + }), + ) + .await; match execution { Ok(raw_stream) => { diff --git a/crates/core/src/api/runtime.rs b/crates/core/src/api/runtime.rs index 0261657ba..77670c804 100644 --- a/crates/core/src/api/runtime.rs +++ b/crates/core/src/api/runtime.rs @@ -19,10 +19,12 @@ pub use callbacks::{ }; pub use global::global_context; pub use scope_stack::{ - 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, with_scope_stack, + PropagationContext, ScopeStack, ScopeStackHandle, TASK_SCOPE_STACK, ThreadScopeStackBinding, + capture_propagation_context, capture_propagation_context_with_root, capture_thread_scope_stack, + create_scope_stack, create_scope_stack_from_propagation, 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, with_active_event_uuid, with_scope_stack, }; pub use state::NemoRelayContextState; pub use subscriber_dispatcher::flush_subscribers; diff --git a/crates/core/src/api/runtime/scope_stack.rs b/crates/core/src/api/runtime/scope_stack.rs index f62f7191b..ac2a70899 100644 --- a/crates/core/src/api/runtime/scope_stack.rs +++ b/crates/core/src/api/runtime/scope_stack.rs @@ -10,8 +10,10 @@ use std::cell::RefCell; use std::collections::{HashMap, HashSet}; +use std::future::Future; use std::sync::{Arc, RwLock}; +use serde::{Deserialize, Serialize}; use uuid::Uuid; use crate::api::runtime::callbacks::EventSubscriberFn; @@ -31,6 +33,51 @@ pub struct ScopeStack { stack: Vec, scope_registries: HashMap, fresh_agents: HashSet, + propagated_parent_uuid: Option, +} + +/// Versioned, transport-neutral causal context for crossing a Relay boundary. +/// +/// Applications are responsible for serializing, transporting, authenticating, +/// and trusting this value. It intentionally contains only Relay identifiers; +/// OpenTelemetry `traceparent` and `tracestate` remain transport sidecars. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct PropagationContext { + /// Wire-format version. Version 1 is the only currently supported value. + pub version: u16, + /// Stable session root when the sending application knows one. + #[serde(skip_serializing_if = "Option::is_none")] + pub root_uuid: Option, + /// Immediate Relay event or scope that caused the boundary crossing. + pub parent_uuid: Uuid, +} + +impl PropagationContext { + /// The current wire-format version. + pub const VERSION: u16 = 1; + + /// Validate a context received from an untrusted transport. + pub fn validate(&self) -> Result<()> { + if self.version != Self::VERSION { + return Err(FlowError::InvalidArgument(format!( + "unsupported propagation context version {}; expected {}", + self.version, + Self::VERSION + ))); + } + for (name, uuid) in [("parent_uuid", self.parent_uuid)] + .into_iter() + .chain(self.root_uuid.map(|uuid| ("root_uuid", uuid))) + { + let bytes = uuid.as_bytes(); + if bytes.iter().all(|byte| *byte == 0) || bytes[8..].iter().all(|byte| *byte == 0) { + return Err(FlowError::InvalidArgument(format!( + "propagation context {name} is not a usable Relay identifier" + ))); + } + } + Ok(()) + } } impl ScopeStack { @@ -49,7 +96,52 @@ impl ScopeStack { stack: vec![root], scope_registries: HashMap::new(), fresh_agents: HashSet::from([root_uuid]), + propagated_parent_uuid: None, + } + } + + fn from_propagation(context: &PropagationContext) -> Result { + context.validate()?; + let (root, parent) = match context.root_uuid { + Some(root_uuid) => { + let root = ScopeHandle::builder() + .uuid(root_uuid) + .name("propagated-root") + .scope_type(ScopeType::Agent) + .build(); + let parent = (root_uuid != context.parent_uuid).then(|| { + ScopeHandle::builder() + .uuid(context.parent_uuid) + .parent_uuid(root_uuid) + .name("propagated-parent") + .scope_type(ScopeType::Unknown) + .build() + }); + (root, parent) + } + None => ( + ScopeHandle::builder() + .uuid(context.parent_uuid) + .name("propagated-root") + .scope_type(ScopeType::Agent) + .build(), + None, + ), + }; + let root_uuid = root.uuid; + let mut stack = vec![root]; + if let Some(parent) = parent { + stack.push(parent); } + Ok(Self { + stack, + scope_registries: HashMap::new(), + fresh_agents: HashSet::from([root_uuid]), + propagated_parent_uuid: context + .root_uuid + .filter(|root_uuid| *root_uuid != context.parent_uuid) + .map(|_| context.parent_uuid), + }) } /// Push a scope handle onto the top of the stack. @@ -98,6 +190,11 @@ impl ScopeStack { .uuid } + /// Whether `uuid` is the synthetic parent imported from propagation. + pub fn is_propagated_parent(&self, uuid: Uuid) -> bool { + self.propagated_parent_uuid == Some(uuid) + } + /// Return the full ordered stack of scope handles. /// /// # Returns @@ -290,9 +387,48 @@ pub fn create_scope_stack() -> ScopeStackHandle { Arc::new(RwLock::new(ScopeStack::new())) } +/// Create an isolated scope stack rooted below a supplied propagation context. +/// +/// The imported handles are synthetic bookkeeping only; Relay never emits their +/// lifecycle events or transfers scope-local registrations across the boundary. +pub fn create_scope_stack_from_propagation( + context: &PropagationContext, +) -> Result { + Ok(Arc::new(RwLock::new(ScopeStack::from_propagation( + context, + )?))) +} + +/// Capture the current causal parent without asserting a session root. +pub fn capture_propagation_context() -> Result { + capture_propagation_context_with_root(None) +} + +/// Capture the current causal parent and an application-supplied session root. +pub fn capture_propagation_context_with_root( + root_uuid: Option, +) -> Result { + let context = PropagationContext { + version: PropagationContext::VERSION, + root_uuid, + parent_uuid: ACTIVE_EVENT_UUID + .try_with(|uuid| *uuid) + .unwrap_or_else(|_| task_scope_top().uuid), + }; + context.validate()?; + Ok(context) +} + tokio::task_local! { /// Task-local scope stack handle used by async execution contexts. pub static TASK_SCOPE_STACK: ScopeStackHandle; + /// Managed tool or LLM event currently executing in this task. + static ACTIVE_EVENT_UUID: Uuid; +} + +/// Run a future with `uuid` as the causally active managed event. +pub async fn with_active_event_uuid(uuid: Uuid, future: impl Future) -> T { + ACTIVE_EVENT_UUID.scope(uuid, future).await } thread_local! { diff --git a/crates/core/src/api/tool.rs b/crates/core/src/api/tool.rs index 2655fc857..6fbe6cd70 100644 --- a/crates/core/src/api/tool.rs +++ b/crates/core/src/api/tool.rs @@ -7,7 +7,7 @@ use crate::api::event::{BaseEvent, Event, MarkEvent, PendingMarkSpec}; use crate::api::runtime::NemoRelayContextState; use crate::api::runtime::current_scope_stack; use crate::api::runtime::global_context; -use crate::api::runtime::{EventSubscriberFn, ToolExecutionNextFn}; +use crate::api::runtime::{EventSubscriberFn, ToolExecutionNextFn, with_active_event_uuid}; use crate::api::scope::event; use crate::api::scope::{EmitMarkEventParams, ScopeHandle}; use crate::api::shared::{ @@ -551,7 +551,7 @@ pub async fn tool_call_execute(params: ToolCallExecuteParams) -> Result { state.tool_build_execution_chain(&name, func, &scope_locals) }; - match execution(intercepted_args).await { + match with_active_event_uuid(handle.uuid, execution(intercepted_args)).await { Ok(outcome) => { let ToolExecutionInterceptOutcome { result, diff --git a/crates/core/src/observability/openinference.rs b/crates/core/src/observability/openinference.rs index 362015853..05d980869 100644 --- a/crates/core/src/observability/openinference.rs +++ b/crates/core/src/observability/openinference.rs @@ -29,7 +29,7 @@ use super::{ relay_trace_id, validate_attribute_mappings, }; use crate::api::event::{Event, EventNormalizationExt, ScopeCategory}; -use crate::api::runtime::EventSubscriberFn; +use crate::api::runtime::{EventSubscriberFn, current_scope_stack}; use crate::api::scope::ScopeType; use crate::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber}; use crate::codec::request::{ @@ -42,7 +42,8 @@ use chrono::{DateTime, Utc}; use openinference_semantic_conventions::SpanKind as OpenInferenceSpanKind; use openinference_semantic_conventions::attributes as oi; use opentelemetry::trace::{ - Span as _, SpanContext, SpanKind, TraceContextExt, Tracer, TracerProvider as _, + Span as _, SpanContext, SpanKind, TraceContextExt, TraceFlags, TraceState, Tracer, + TracerProvider as _, }; use opentelemetry::{Context, KeyValue}; use opentelemetry_otlp::{Protocol, SpanExporter, WithExportConfig, WithHttpConfig}; @@ -788,11 +789,31 @@ impl OpenInferenceEventProcessor { if let Some(active_span) = self.find_parent_span(event) { return Context::new().with_remote_span_context(active_span.span_context.clone()); } - event + if let Some(span_context) = event .parent_uuid() .and_then(|uuid| self.completed_span_contexts.get(&uuid)) - .map(|span_context| Context::new().with_remote_span_context(span_context.clone())) - .unwrap_or_default() + { + return Context::new().with_remote_span_context(span_context.clone()); + } + let Some(parent_uuid) = event.parent_uuid() else { + return Context::new(); + }; + let stack = current_scope_stack(); + let stack = stack.read().expect("scope stack lock poisoned"); + if !stack.is_propagated_parent(parent_uuid) { + return Context::new(); + } + let root_uuid = stack.root_uuid(); + if parent_uuid == root_uuid { + return Context::new(); + } + Context::new().with_remote_span_context(SpanContext::new( + relay_trace_id(root_uuid), + relay_span_id(parent_uuid), + TraceFlags::SAMPLED, + true, + TraceState::default(), + )) } fn parent_span_uuid(&self, event: &Event) -> Option { diff --git a/crates/core/src/observability/otel.rs b/crates/core/src/observability/otel.rs index d9b19a1d7..841e58f03 100644 --- a/crates/core/src/observability/otel.rs +++ b/crates/core/src/observability/otel.rs @@ -29,14 +29,15 @@ use super::{ relay_trace_id, validate_attribute_mappings, }; use crate::api::event::{Event, EventNormalizationExt, ScopeCategory}; -use crate::api::runtime::EventSubscriberFn; +use crate::api::runtime::{EventSubscriberFn, current_scope_stack}; use crate::api::scope::ScopeType; use crate::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber}; use crate::codec::response::CostEstimate; use crate::error::FlowError; use chrono::{DateTime, Utc}; use opentelemetry::trace::{ - Span as _, SpanContext, SpanKind, TraceContextExt, Tracer, TracerProvider as _, + Span as _, SpanContext, SpanKind, TraceContextExt, TraceFlags, TraceState, Tracer, + TracerProvider as _, }; use opentelemetry::{Context, KeyValue}; use opentelemetry_otlp::{Protocol, SpanExporter, WithExportConfig, WithHttpConfig}; @@ -774,11 +775,31 @@ impl OtelEventProcessor { if let Some(active_span) = self.find_parent_span(event) { return Context::new().with_remote_span_context(active_span.span_context.clone()); } - event + if let Some(span_context) = event .parent_uuid() .and_then(|uuid| self.completed_span_contexts.get(&uuid)) - .map(|span_context| Context::new().with_remote_span_context(span_context.clone())) - .unwrap_or_default() + { + return Context::new().with_remote_span_context(span_context.clone()); + } + let Some(parent_uuid) = event.parent_uuid() else { + return Context::new(); + }; + let stack = current_scope_stack(); + let stack = stack.read().expect("scope stack lock poisoned"); + if !stack.is_propagated_parent(parent_uuid) { + return Context::new(); + } + let root_uuid = stack.root_uuid(); + if parent_uuid == root_uuid { + return Context::new(); + } + Context::new().with_remote_span_context(SpanContext::new( + relay_trace_id(root_uuid), + relay_span_id(parent_uuid), + TraceFlags::SAMPLED, + true, + TraceState::default(), + )) } fn parent_span_uuid(&self, event: &Event) -> Option { diff --git a/crates/core/tests/integration/context_isolation_tests.rs b/crates/core/tests/integration/context_isolation_tests.rs index 22cf33ba2..33a28e530 100644 --- a/crates/core/tests/integration/context_isolation_tests.rs +++ b/crates/core/tests/integration/context_isolation_tests.rs @@ -6,9 +6,10 @@ use std::sync::Arc; use nemo_relay::api::runtime::{ - ScopeStack, 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, + PropagationContext, ScopeStack, TASK_SCOPE_STACK, create_scope_stack, + create_scope_stack_from_propagation, 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, }; use nemo_relay::api::scope::{ PopScopeParams, PushScopeParams, ScopeHandle, ScopeType, pop_scope, push_scope, @@ -55,6 +56,37 @@ fn test_two_scope_stacks_are_independent() { assert_ne!(root_a_uuid, root_b_uuid); // scope_a != scope_b } +#[test] +fn test_propagation_context_seeds_a_synthetic_root_and_parent() { + let root_uuid = Uuid::now_v7(); + let parent_uuid = Uuid::now_v7(); + let stack = create_scope_stack_from_propagation(&PropagationContext { + version: PropagationContext::VERSION, + root_uuid: Some(root_uuid), + parent_uuid, + }) + .unwrap(); + let stack = stack.read().unwrap(); + assert_eq!(stack.root_uuid(), root_uuid); + assert_eq!(stack.top().uuid, parent_uuid); + assert_eq!(stack.scopes().len(), 2); +} + +#[test] +fn test_rootless_propagation_context_uses_the_parent_as_root() { + let parent_uuid = Uuid::now_v7(); + let stack = create_scope_stack_from_propagation(&PropagationContext { + version: PropagationContext::VERSION, + root_uuid: None, + parent_uuid, + }) + .unwrap(); + let stack = stack.read().unwrap(); + assert_eq!(stack.root_uuid(), parent_uuid); + assert_eq!(stack.top().uuid, parent_uuid); + assert_eq!(stack.scopes().len(), 1); +} + #[test] fn test_pop_scope_rejects_non_top_and_unknown_handles() { set_thread_scope_stack(create_scope_stack()); diff --git a/crates/ffi/nemo_relay.h b/crates/ffi/nemo_relay.h index 3c381015a..be8822bc4 100644 --- a/crates/ffi/nemo_relay.h +++ b/crates/ffi/nemo_relay.h @@ -2341,6 +2341,39 @@ NemoRelayStatus nemo_relay_scope_deregister_subscriber(const char *scope_uuid, c */ NemoRelayStatus nemo_relay_scope_stack_create(struct FfiScopeStack **out); +/** + * Serialize the current causal parent as a versioned propagation context. + * + * The returned JSON must be freed with `nemo_relay_string_free`. + * + * # Safety + * `out` must be a valid, writable pointer to a C-string output slot. + */ +NemoRelayStatus nemo_relay_capture_propagation_context_json(char **out); + +/** + * Serialize the current causal parent with an application-supplied root UUID. + * + * Pass null for `root_uuid` to omit the root. The returned JSON must be freed + * with `nemo_relay_string_free`. + * + * # Safety + * When non-null, `root_uuid` must point to a valid NUL-terminated C string; + * `out` must be a valid, writable pointer to a C-string output slot. + */ +NemoRelayStatus nemo_relay_capture_propagation_context_with_root_json(const char *root_uuid, + char **out); + +/** + * Create an isolated scope stack from propagation-context JSON. + * + * # Safety + * `context_json` must point to a valid NUL-terminated C string and `out` must + * be a valid, writable pointer to a scope-stack output slot. + */ +NemoRelayStatus nemo_relay_scope_stack_create_from_propagation_json(const char *context_json, + struct FfiScopeStack **out); + /** * Bind an isolated scope stack to the current OS thread. * diff --git a/crates/ffi/src/api/scope_stack.rs b/crates/ffi/src/api/scope_stack.rs index abaa761b1..3115ee74b 100644 --- a/crates/ffi/src/api/scope_stack.rs +++ b/crates/ffi/src/api/scope_stack.rs @@ -2,10 +2,15 @@ // SPDX-License-Identifier: Apache-2.0 use super::{ - FfiScopeStack, FfiThreadScopeStackBinding, NemoRelayStatus, capture_thread_scope_stack, - clear_last_error, create_scope_stack, restore_thread_scope_stack, scope_stack_active, - set_last_error, set_thread_scope_stack, + FfiScopeStack, FfiThreadScopeStackBinding, NemoRelayStatus, c_char, c_str_to_string, + capture_thread_scope_stack, clear_last_error, create_scope_stack, json_to_c_string, + restore_thread_scope_stack, scope_stack_active, set_last_error, set_thread_scope_stack, }; +use nemo_relay::api::runtime::{ + PropagationContext, capture_propagation_context, capture_propagation_context_with_root, + create_scope_stack_from_propagation, +}; +use uuid::Uuid; // --------------------------------------------------------------------------- // Scope stack isolation @@ -42,6 +47,118 @@ pub unsafe extern "C" fn nemo_relay_scope_stack_create( NemoRelayStatus::Ok } +/// Serialize the current causal parent as a versioned propagation context. +/// +/// The returned JSON must be freed with `nemo_relay_string_free`. +/// +/// # Safety +/// `out` must be a valid, writable pointer to a C-string output slot. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_capture_propagation_context_json( + out: *mut *mut c_char, +) -> NemoRelayStatus { + clear_last_error(); + if out.is_null() { + set_last_error("out pointer is null"); + return NemoRelayStatus::NullPointer; + } + match capture_propagation_context().and_then(|context| { + serde_json::to_value(context) + .map_err(|error| nemo_relay::error::FlowError::Internal(error.to_string())) + }) { + Ok(context) => { + unsafe { *out = json_to_c_string(&context) }; + NemoRelayStatus::Ok + } + Err(error) => { + set_last_error(&error.to_string()); + NemoRelayStatus::from(&error) + } + } +} + +/// Serialize the current causal parent with an application-supplied root UUID. +/// +/// Pass null for `root_uuid` to omit the root. The returned JSON must be freed +/// with `nemo_relay_string_free`. +/// +/// # Safety +/// When non-null, `root_uuid` must point to a valid NUL-terminated C string; +/// `out` must be a valid, writable pointer to a C-string output slot. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_capture_propagation_context_with_root_json( + root_uuid: *const c_char, + out: *mut *mut c_char, +) -> NemoRelayStatus { + clear_last_error(); + if out.is_null() { + set_last_error("out pointer is null"); + return NemoRelayStatus::NullPointer; + } + let root_uuid = if root_uuid.is_null() { + Ok(None) + } else { + c_str_to_string(root_uuid) + .map_err(|_| ()) + .and_then(|value| Uuid::parse_str(&value).map_err(|_| ())) + .map(Some) + }; + let Ok(root_uuid) = root_uuid else { + set_last_error("root_uuid must be a valid UUID"); + return NemoRelayStatus::InvalidArg; + }; + match capture_propagation_context_with_root(root_uuid).and_then(|context| { + serde_json::to_value(context) + .map_err(|error| nemo_relay::error::FlowError::Internal(error.to_string())) + }) { + Ok(context) => { + unsafe { *out = json_to_c_string(&context) }; + NemoRelayStatus::Ok + } + Err(error) => { + set_last_error(&error.to_string()); + NemoRelayStatus::from(&error) + } + } +} + +/// Create an isolated scope stack from propagation-context JSON. +/// +/// # Safety +/// `context_json` must point to a valid NUL-terminated C string and `out` must +/// be a valid, writable pointer to a scope-stack output slot. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_scope_stack_create_from_propagation_json( + context_json: *const c_char, + out: *mut *mut FfiScopeStack, +) -> NemoRelayStatus { + clear_last_error(); + if out.is_null() { + set_last_error("out pointer is null"); + return NemoRelayStatus::NullPointer; + } + let context = match c_str_to_string(context_json) { + Ok(value) => match serde_json::from_str::(&value) { + Ok(context) => context, + Err(error) => { + set_last_error(&format!("invalid propagation context JSON: {error}")); + return NemoRelayStatus::InvalidJson; + } + }, + Err(status) => return status, + }; + match create_scope_stack_from_propagation(&context) { + Ok(stack) => { + unsafe { *out = Box::into_raw(Box::new(FfiScopeStack(stack))) }; + NemoRelayStatus::Ok + } + Err(error) => { + set_last_error(&error.to_string()); + NemoRelayStatus::from(&error) + } + } +} + /// Bind an isolated scope stack to the current OS thread. /// /// After this call, all NeMo Relay scope operations on the current thread diff --git a/crates/node/src/api/mod.rs b/crates/node/src/api/mod.rs index 1f8384fae..58dd3f672 100644 --- a/crates/node/src/api/mod.rs +++ b/crates/node/src/api/mod.rs @@ -36,9 +36,13 @@ use nemo_relay::api::runtime::{ ToolExecutionNextFn, }; use nemo_relay::api::runtime::{ - TASK_SCOPE_STACK, create_scope_stack as create_scope_stack_handle, + TASK_SCOPE_STACK, capture_propagation_context as capture_propagation_context_handle, + capture_propagation_context_with_root as capture_propagation_context_with_root_handle, + create_scope_stack as create_scope_stack_handle, + create_scope_stack_from_propagation as create_scope_stack_from_propagation_handle, current_scope_stack as current_scope_stack_handle, scope_stack_active as scope_stack_is_active, set_thread_scope_stack as bind_thread_scope_stack, task_scope_top, + with_scope_stack as with_scope_stack_handle, }; use nemo_relay::api::scope as core_scope_api; use nemo_relay::api::scope::ScopeAttributes; @@ -1726,6 +1730,48 @@ impl Plugin for NodePlugin { // Scope stack isolation // --------------------------------------------------------------------------- +/// Transport-neutral Relay causal context for application-managed transport. +#[napi(object)] +pub struct PropagationContext { + pub version: u32, + pub root_uuid: Option, + pub parent_uuid: String, +} + +fn propagation_context_from_napi( + context: PropagationContext, +) -> napi::Result { + let root_uuid = context + .root_uuid + .as_deref() + .map(uuid::Uuid::parse_str) + .transpose() + .map_err(|error| napi::Error::from_reason(format!("invalid root UUID: {error}")))?; + let parent_uuid = uuid::Uuid::parse_str(&context.parent_uuid) + .map_err(|error| napi::Error::from_reason(format!("invalid parent UUID: {error}")))?; + let version = u16::try_from(context.version) + .map_err(|_| napi::Error::from_reason("propagation context version is out of range"))?; + let context = nemo_relay::api::runtime::PropagationContext { + version, + root_uuid, + parent_uuid, + }; + context + .validate() + .map_err(|error| napi::Error::from_reason(error.to_string()))?; + Ok(context) +} + +fn propagation_context_to_napi( + context: nemo_relay::api::runtime::PropagationContext, +) -> PropagationContext { + PropagationContext { + version: u32::from(context.version), + root_uuid: context.root_uuid.map(|uuid| uuid.to_string()), + parent_uuid: context.parent_uuid.to_string(), + } +} + /// Creates a new isolated scope stack. #[napi] pub fn create_scope_stack() -> ScopeStack { @@ -1734,6 +1780,50 @@ pub fn create_scope_stack() -> ScopeStack { } } +/// Capture the current Relay causal parent for application-managed transport. +#[napi] +pub fn capture_propagation_context() -> napi::Result { + capture_propagation_context_handle() + .map(propagation_context_to_napi) + .map_err(|error| napi::Error::from_reason(error.to_string())) +} + +/// Capture the current parent with an optional stable application session root. +#[napi] +pub fn capture_propagation_context_with_root( + root_uuid: Option, +) -> napi::Result { + let root_uuid = root_uuid + .as_deref() + .map(uuid::Uuid::parse_str) + .transpose() + .map_err(|error| napi::Error::from_reason(format!("invalid root UUID: {error}")))?; + capture_propagation_context_with_root_handle(root_uuid) + .map(propagation_context_to_napi) + .map_err(|error| napi::Error::from_reason(error.to_string())) +} + +/// Create an isolated scope stack seeded from a received propagation context. +#[napi] +pub fn create_scope_stack_from_propagation( + context: PropagationContext, +) -> napi::Result { + create_scope_stack_from_propagation_handle(&propagation_context_from_napi(context)?) + .map(ScopeStack::from) + .map_err(|error| napi::Error::from_reason(error.to_string())) +} + +/// Run a synchronous callback with an isolated scope stack installed. +/// +/// For asynchronous JavaScript request handlers, keep the stack installed with +/// `setThreadScopeStack` for the handler's lifetime instead. +#[napi] +pub fn with_scope_stack(stack: &ScopeStack, callback: JsFunction) -> napi::Result { + with_scope_stack_handle(stack.inner.clone(), || { + callback.call::(None, &[]) + }) +} + /// Returns the current execution context's scope stack handle. #[napi] pub fn current_scope_stack() -> ScopeStack { diff --git a/crates/node/tests/context_tests.mjs b/crates/node/tests/context_tests.mjs index 466145545..5b26280da 100644 --- a/crates/node/tests/context_tests.mjs +++ b/crates/node/tests/context_tests.mjs @@ -18,6 +18,8 @@ const { popScope, ScopeType, ScopeStack, + createScopeStackFromPropagation, + withScopeStack, } = lib; // =========================================================================== @@ -31,6 +33,27 @@ describe('Context isolation', () => { assert.ok(stack instanceof ScopeStack, 'Expected instance of ScopeStack'); }); + it('creates an imported stack with the propagated parent on top', () => { + const original = currentScopeStack(); + const rootUuid = '018f13f0-7c1a-7a80-8000-000000000001'; + const parentUuid = '018f13f0-7c1a-7a80-8000-000000000002'; + const stack = createScopeStackFromPropagation({ version: 1, rootUuid, parentUuid }); + setThreadScopeStack(stack); + assert.equal(getHandle().uuid, parentUuid); + setThreadScopeStack(original); + }); + + it('restores the surrounding stack after withScopeStack', () => { + const original = currentScopeStack(); + const stack = createScopeStack(); + withScopeStack(stack, () => { + pushScope('temporary-with-scope-stack', ScopeType.Agent, null, null); + assert.equal(getHandle().name, 'temporary-with-scope-stack'); + }); + assert.notEqual(getHandle().name, 'temporary-with-scope-stack'); + setThreadScopeStack(original); + }); + it('currentScopeStack returns same in same context', () => { const s1 = currentScopeStack(); const s2 = currentScopeStack(); diff --git a/crates/python/src/py_api/mod.rs b/crates/python/src/py_api/mod.rs index 64e2341e3..4983b777e 100644 --- a/crates/python/src/py_api/mod.rs +++ b/crates/python/src/py_api/mod.rs @@ -17,7 +17,10 @@ use nemo_relay::api::runtime::{ LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn, ToolExecutionNextFn, }; use nemo_relay::api::runtime::{ - TASK_SCOPE_STACK, create_scope_stack as create_scope_stack_handle, + TASK_SCOPE_STACK, capture_propagation_context as capture_propagation_context_handle, + capture_propagation_context_with_root as capture_propagation_context_with_root_handle, + create_scope_stack as create_scope_stack_handle, + create_scope_stack_from_propagation as create_scope_stack_from_propagation_handle, current_scope_stack as current_scope_stack_handle, scope_stack_active as scope_stack_is_active, set_thread_scope_stack as bind_thread_scope_stack, sync_thread_scope_stack as sync_bound_thread_scope_stack, task_scope_top, @@ -38,8 +41,8 @@ use crate::convert::{json_to_py, opt_py_to_json, opt_py_to_timestamp, py_to_json use crate::py_callable; use crate::py_types::{ PyAnnotatedLLMResponse, PyAnthropicMessagesCodec, PyLLMAttributes, PyLLMHandle, PyLLMRequest, - PyLlmStream, PyOpenAIChatCodec, PyOpenAIResponsesCodec, PyScopeAttributes, PyScopeHandle, - PyScopeStack, PyScopeType, PyToolAttributes, PyToolHandle, + PyLlmStream, PyOpenAIChatCodec, PyOpenAIResponsesCodec, PyPropagationContext, + PyScopeAttributes, PyScopeHandle, PyScopeStack, PyScopeType, PyToolAttributes, PyToolHandle, }; pub(crate) type RustJsonStream = LlmJsonStream; @@ -162,6 +165,38 @@ pub fn create_scope_stack() -> PyScopeStack { PyScopeStack(create_scope_stack_handle()) } +/// Capture a transport-neutral context from the current Relay scope stack. +#[pyfunction] +pub fn capture_propagation_context() -> PyResult { + capture_propagation_context_handle() + .map(|inner| PyPropagationContext { inner }) + .map_err(to_py_err) +} + +/// Capture a context with an application-supplied stable session root UUID. +#[pyfunction] +pub fn capture_propagation_context_with_root( + root_uuid: Option<&str>, +) -> PyResult { + let root_uuid = root_uuid + .map(Uuid::parse_str) + .transpose() + .map_err(|error| PyErr::new::(error.to_string()))?; + capture_propagation_context_with_root_handle(root_uuid) + .map(|inner| PyPropagationContext { inner }) + .map_err(to_py_err) +} + +/// Create an isolated scope stack seeded from a received propagation context. +#[pyfunction] +pub fn create_scope_stack_from_propagation( + context: &PyPropagationContext, +) -> PyResult { + create_scope_stack_from_propagation_handle(&context.inner) + .map(PyScopeStack) + .map_err(to_py_err) +} + /// Bind a ``ScopeStack`` to the current thread's thread-local storage. /// /// This ensures that subsequent NeMo Relay API calls on this thread use the given @@ -1740,6 +1775,9 @@ fn scope_deregister_subscriber(scope_uuid: &str, name: &str) -> PyResult { pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> { // Scope stack creation / binding / query m.add_function(wrap_pyfunction!(create_scope_stack, m)?)?; + m.add_function(wrap_pyfunction!(capture_propagation_context, m)?)?; + m.add_function(wrap_pyfunction!(capture_propagation_context_with_root, m)?)?; + m.add_function(wrap_pyfunction!(create_scope_stack_from_propagation, m)?)?; m.add_function(wrap_pyfunction!(set_thread_scope_stack, m)?)?; m.add_function(wrap_pyfunction!(sync_thread_scope_stack, m)?)?; m.add_function(wrap_pyfunction!(py_scope_stack_active, m)?)?; diff --git a/crates/python/src/py_types/core.rs b/crates/python/src/py_types/core.rs index d0f122311..3648e629d 100644 --- a/crates/python/src/py_types/core.rs +++ b/crates/python/src/py_types/core.rs @@ -14,7 +14,9 @@ use super::{ }; use nemo_relay::api::event::{CategoryProfile, EventCategory, PendingMarkSpec}; use nemo_relay::api::llm::LlmRequestInterceptOutcome; -use nemo_relay::api::runtime::{LlmSanitizeRequestContext, LlmSanitizeResponseContext}; +use nemo_relay::api::runtime::{ + LlmSanitizeRequestContext, LlmSanitizeResponseContext, PropagationContext, +}; use nemo_relay::api::tool::ToolExecutionInterceptOutcome; /// Structured identity of the codec active during LLM sanitization. @@ -188,6 +190,49 @@ impl PyScopeStack { } } +/// Transport-neutral causal context used to continue Relay work remotely. +#[pyclass(name = "PropagationContext", skip_from_py_object)] +#[derive(Clone)] +pub struct PyPropagationContext { + pub(crate) inner: PropagationContext, +} + +#[pymethods] +impl PyPropagationContext { + #[new] + #[pyo3(signature = (parent_uuid, root_uuid=None, version=1))] + fn new(parent_uuid: &str, root_uuid: Option<&str>, version: u16) -> PyResult { + let context = PropagationContext { + version, + root_uuid: root_uuid + .map(uuid::Uuid::parse_str) + .transpose() + .map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string()))?, + parent_uuid: uuid::Uuid::parse_str(parent_uuid) + .map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string()))?, + }; + context + .validate() + .map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string()))?; + Ok(Self { inner: context }) + } + + #[getter] + fn version(&self) -> u16 { + self.inner.version + } + + #[getter] + fn root_uuid(&self) -> Option { + self.inner.root_uuid.map(|uuid| uuid.to_string()) + } + + #[getter] + fn parent_uuid(&self) -> String { + self.inner.parent_uuid.to_string() + } +} + // --------------------------------------------------------------------------- // ScopeAttributes (bitflag wrapper) // --------------------------------------------------------------------------- diff --git a/crates/python/src/py_types/mod.rs b/crates/python/src/py_types/mod.rs index eb60374d9..9a13bef16 100644 --- a/crates/python/src/py_types/mod.rs +++ b/crates/python/src/py_types/mod.rs @@ -137,6 +137,7 @@ pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> { fn register_runtime_types(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/docs/about-nemo-relay/concepts/scopes.mdx b/docs/about-nemo-relay/concepts/scopes.mdx index 84b293f96..49ac3bd6d 100644 --- a/docs/about-nemo-relay/concepts/scopes.mdx +++ b/docs/about-nemo-relay/concepts/scopes.mdx @@ -142,6 +142,24 @@ Use this when: - The boundary cannot safely carry a native stack handle - You want a clean root scope with isolated scope-local registrations +## Cross-Process Propagation + +When work crosses a process or remote-workflow boundary, applications can carry +the versioned Relay propagation context instead of a native stack handle. The +context contains an immediate `parent_uuid` and, when the application knows a +stable session root, an optional `root_uuid`. + +The receiver creates a fresh isolated stack from that context and installs it +only for request handling. Its first local event becomes a child of +`parent_uuid`; scope-local middleware and subscribers are never transferred. +The transport is application-owned: authenticate and authorize inbound context +before importing it. Relay does not send headers, make IPC connections, or +trust remote identifiers automatically. + +Relay context is distinct from W3C propagation. An integration may carry +`traceparent` and `tracestate` alongside Relay's JSON context when it needs to +preserve OpenTelemetry sampling or vendor state. + ## Practical Guidance Use these practices when applying the concept in application or integration code. diff --git a/docs/configure-plugins/observability/opentelemetry.mdx b/docs/configure-plugins/observability/opentelemetry.mdx index 48ed44383..7cd249516 100644 --- a/docs/configure-plugins/observability/opentelemetry.mdx +++ b/docs/configure-plugins/observability/opentelemetry.mdx @@ -74,6 +74,12 @@ The collector should receive OTLP trace export requests. The tracing backend should show spans for NeMo Relay scopes, tools, LLM calls, and marks grouped by root scope. +For an imported Relay propagation context with both a root and an external +parent, Relay uses the root UUID as the OTLP trace ID and the lower 64 bits of +the parent UUID as a remote parent span ID. A rootless import begins a new +native trace. Relay does not derive W3C sampling flags or `tracestate`; carry +those separately in an integration when they are required. + The default `inherit` projection follows exporter-native handling: a mark with an active parent span is a span event, while an orphan mark is a standalone zero-duration `mark:` span. `mark_projection = "event"` explicitly selects diff --git a/go/nemo_relay/context_test.go b/go/nemo_relay/context_test.go index 81ab523a4..d89891c33 100644 --- a/go/nemo_relay/context_test.go +++ b/go/nemo_relay/context_test.go @@ -259,6 +259,30 @@ func TestCreateScopeStackCreatesFreshStack(t *testing.T) { } } +func TestNewScopeStackFromPropagationUsesParentAsCurrentHandle(t *testing.T) { + rootUUID := "018f13f0-7c1a-7a80-8000-000000000001" + parentUUID := "018f13f0-7c1a-7a80-8000-000000000002" + stack, err := NewScopeStackFromPropagation(PropagationContext{ + Version: 1, + RootUUID: &rootUUID, + ParentUUID: parentUUID, + }) + if err != nil { + t.Fatalf("NewScopeStackFromPropagation failed: %v", err) + } + defer stack.Close() + + stack.Run(func() { + handle, err := GetHandle() + if err != nil { + t.Fatalf("GetHandle failed: %v", err) + } + if handle.UUID() != parentUUID { + t.Fatalf("expected parent UUID %s, got %s", parentUUID, handle.UUID()) + } + }) +} + func TestConcurrentScopeStacksWithToolCalls(t *testing.T) { const goroutines = 5 var wg sync.WaitGroup diff --git a/go/nemo_relay/nemo_relay.go b/go/nemo_relay/nemo_relay.go index 937427e5d..60057979e 100644 --- a/go/nemo_relay/nemo_relay.go +++ b/go/nemo_relay/nemo_relay.go @@ -229,6 +229,9 @@ extern void nemo_relay_string_free(char* ptr); // Scope stack isolation extern int32_t nemo_relay_scope_stack_create(FfiScopeStack** out); +extern int32_t nemo_relay_capture_propagation_context_json(char** out); +extern int32_t nemo_relay_capture_propagation_context_with_root_json(const char* root_uuid, char** out); +extern int32_t nemo_relay_scope_stack_create_from_propagation_json(const char* context_json, FfiScopeStack** out); extern int32_t nemo_relay_scope_stack_set_thread(const FfiScopeStack* stack); extern int32_t nemo_relay_scope_stack_capture_thread(FfiThreadScopeStackBinding** out); extern int32_t nemo_relay_scope_stack_restore_thread(FfiThreadScopeStackBinding* binding); @@ -1567,12 +1570,68 @@ type ScopeStack struct { ptr *C.FfiScopeStack } +// PropagationContext is the versioned, transport-neutral causal context used +// to continue Relay work in another process. +type PropagationContext struct { + Version uint16 `json:"version"` + RootUUID *string `json:"root_uuid,omitempty"` + ParentUUID string `json:"parent_uuid"` +} + +// CapturePropagationContext captures the current Relay causal parent. +func CapturePropagationContext() (PropagationContext, error) { + var out *C.char + if err := checkStatus(C.nemo_relay_capture_propagation_context_json(&out)); err != nil { + return PropagationContext{}, err + } + defer C.nemo_relay_string_free(out) + var context PropagationContext + if err := json.Unmarshal([]byte(C.GoString(out)), &context); err != nil { + return PropagationContext{}, err + } + return context, nil +} + +// CapturePropagationContextWithRoot captures the current parent with an +// application-supplied stable session root. Pass nil when no root is known. +func CapturePropagationContextWithRoot(rootUUID *string) (PropagationContext, error) { + var cRoot *C.char + if rootUUID != nil { + cRoot = C.CString(*rootUUID) + defer C.free(unsafe.Pointer(cRoot)) + } + var out *C.char + if err := checkStatus(C.nemo_relay_capture_propagation_context_with_root_json(cRoot, &out)); err != nil { + return PropagationContext{}, err + } + defer C.nemo_relay_string_free(out) + var context PropagationContext + if err := json.Unmarshal([]byte(C.GoString(out)), &context); err != nil { + return PropagationContext{}, err + } + return context, nil +} + // NewScopeStack creates a new isolated scope stack. // The caller must call Close() when done. func NewScopeStack() (*ScopeStack, error) { return newScopeStackFunc() } +// NewScopeStackFromPropagation creates an isolated stack seeded from a +// received propagation context. The caller must call Close when done. +func NewScopeStackFromPropagation(context PropagationContext) (*ScopeStack, error) { + payload, err := json.Marshal(context) + if err != nil { + return nil, err + } + cPayload := C.CString(string(payload)) + defer C.free(unsafe.Pointer(cPayload)) + var ptr *C.FfiScopeStack + status := C.nemo_relay_scope_stack_create_from_propagation_json(cPayload, &ptr) + return checkedValue(int32(status), &ScopeStack{ptr: ptr}) +} + // Close frees the scope stack. After calling Close, the ScopeStack must not be used. func (s *ScopeStack) Close() { if s.ptr != nil { diff --git a/python/nemo_relay/__init__.py b/python/nemo_relay/__init__.py index fb8c89152..c3071e023 100644 --- a/python/nemo_relay/__init__.py +++ b/python/nemo_relay/__init__.py @@ -79,6 +79,7 @@ async def main(): import contextvars import typing from collections.abc import Callable as AbcCallable +from contextlib import contextmanager from typing import AsyncIterator, Awaitable, Callable, Literal, Optional, TypeAlias, TypedDict # Native bitflag classes exported at the top level for user code. @@ -110,6 +111,7 @@ async def main(): OpenTelemetryConfig, OpenTelemetrySubscriber, PendingMarkSpec, + PropagationContext, ScopeAttributes, ScopeEvent, ScopeHandle, @@ -119,7 +121,16 @@ async def main(): ToolExecutionInterceptOutcome, ToolHandle, ) +from nemo_relay._native import ( + capture_propagation_context as _capture_propagation_context, +) +from nemo_relay._native import ( + capture_propagation_context_with_root as _capture_propagation_context_with_root, +) from nemo_relay._native import create_scope_stack as _create_scope_stack +from nemo_relay._native import ( + create_scope_stack_from_propagation as _create_scope_stack_from_propagation, +) from nemo_relay._native import scope_stack_active as _native_scope_stack_active from nemo_relay._native import set_thread_scope_stack as _set_thread_scope_stack from nemo_relay._native import sync_thread_scope_stack as _sync_thread_scope_stack @@ -393,6 +404,35 @@ def create_scope_stack() -> ScopeStack: return _create_scope_stack() +def capture_propagation_context() -> PropagationContext: + """Capture the current Relay causal parent for application-managed transport.""" + get_scope_stack() + return _capture_propagation_context() + + +def capture_propagation_context_with_root(root_uuid: str | None) -> PropagationContext: + """Capture the current parent with an optional stable application session root.""" + get_scope_stack() + return _capture_propagation_context_with_root(root_uuid) + + +def create_scope_stack_from_propagation(context: PropagationContext) -> ScopeStack: + """Create an isolated stack seeded from a received propagation context.""" + return _create_scope_stack_from_propagation(context) + + +@contextmanager +def use_scope_stack(stack: ScopeStack): + """Temporarily install ``stack`` in the current Python context.""" + token = _scope_stack_var.set(stack) + _sync_thread_scope_stack(stack) + try: + yield stack + finally: + _scope_stack_var.reset(token) + _sync_thread_scope_stack(get_scope_stack()) + + def set_thread_scope_stack(stack: ScopeStack) -> None: """Install a scope stack into the current thread's native runtime context. @@ -461,11 +501,16 @@ def worker() -> None: "model_pricing", # Scope stack isolation "ScopeStack", + "PropagationContext", "create_scope_stack", + "capture_propagation_context", + "capture_propagation_context_with_root", + "create_scope_stack_from_propagation", "get_scope_stack", "scope_stack_active", "propagate_scope_to_thread", "set_thread_scope_stack", + "use_scope_stack", # Types "ScopeAttributes", "ToolAttributes", diff --git a/python/nemo_relay/__init__.pyi b/python/nemo_relay/__init__.pyi index 7c2372a57..951c58592 100644 --- a/python/nemo_relay/__init__.pyi +++ b/python/nemo_relay/__init__.pyi @@ -105,6 +105,9 @@ from nemo_relay._native import ( from nemo_relay._native import ( PendingMarkSpec as PendingMarkSpec, ) +from nemo_relay._native import ( + PropagationContext as PropagationContext, +) from nemo_relay._native import ( ScopeAttributes as ScopeAttributes, ) @@ -367,6 +370,10 @@ def create_scope_stack() -> ScopeStack: """ ... +def capture_propagation_context() -> PropagationContext: ... +def capture_propagation_context_with_root(root_uuid: str | None) -> PropagationContext: ... +def create_scope_stack_from_propagation(context: PropagationContext) -> ScopeStack: ... +def use_scope_stack(stack: ScopeStack): ... def set_thread_scope_stack(stack: ScopeStack) -> None: """Install a scope stack into the current thread's native runtime context. diff --git a/python/nemo_relay/_native.pyi b/python/nemo_relay/_native.pyi index a76ee5240..93f1628fd 100644 --- a/python/nemo_relay/_native.pyi +++ b/python/nemo_relay/_native.pyi @@ -1325,6 +1325,19 @@ def create_scope_stack() -> ScopeStack: """ ... +class PropagationContext: + """Transport-neutral Relay causal context.""" + def __init__(self, parent_uuid: str, root_uuid: str | None = None, version: int = 1) -> None: ... + @property + def version(self) -> int: ... + @property + def root_uuid(self) -> str | None: ... + @property + def parent_uuid(self) -> str: ... + +def capture_propagation_context() -> PropagationContext: ... +def capture_propagation_context_with_root(root_uuid: str | None) -> PropagationContext: ... +def create_scope_stack_from_propagation(context: PropagationContext) -> ScopeStack: ... def set_thread_scope_stack(stack: ScopeStack) -> None: """Install a scope stack into native thread-local storage. diff --git a/python/tests/test_context_isolation.py b/python/tests/test_context_isolation.py index 0c4146a00..252e69256 100644 --- a/python/tests/test_context_isolation.py +++ b/python/tests/test_context_isolation.py @@ -4,6 +4,7 @@ """Tests for per-request scope stack isolation via ContextVar.""" import asyncio +import uuid import nemo_relay @@ -15,6 +16,20 @@ def test_create_scope_stack_returns_scope_stack(): assert repr(stack) == "" +def test_propagation_context_installs_and_restores_a_scoped_stack(): + original = nemo_relay.get_scope_stack() + root_uuid = str(uuid.uuid4()) + parent_uuid = str(uuid.uuid4()) + context = nemo_relay.PropagationContext(parent_uuid, root_uuid) + stack = nemo_relay.create_scope_stack_from_propagation(context) + + with nemo_relay.use_scope_stack(stack): + assert nemo_relay.get_scope_stack() is stack + assert nemo_relay.scope.get_handle().uuid == parent_uuid + + assert nemo_relay.get_scope_stack() is original + + def test_get_scope_stack_returns_same_in_same_context(): """get_scope_stack returns the same instance within the same context.""" s1 = nemo_relay.get_scope_stack() From cbc211a3232329909c4e42727a02d32d1d5ac6e5 Mon Sep 17 00:00:00 2001 From: Will Killian Date: Mon, 27 Jul 2026 11:49:33 -0400 Subject: [PATCH 2/5] fix: restore propagated scope stacks safely Signed-off-by: Will Killian --- crates/core/src/api/runtime/scope_stack.rs | 7 +++ .../integration/context_isolation_tests.rs | 38 +++++++++++ crates/ffi/tests/unit/api_tests.rs | 34 ++++++++++ crates/node/src/api/mod.rs | 4 +- crates/node/tests/context_tests.mjs | 48 ++++++++------ crates/python/src/py_api/mod.rs | 8 +++ go/nemo_relay/context_test.go | 63 +++++++++++++++++++ python/nemo_relay/__init__.py | 4 +- python/nemo_relay/_native.pyi | 1 + python/tests/test_context_isolation.py | 12 ++++ 10 files changed, 196 insertions(+), 23 deletions(-) diff --git a/crates/core/src/api/runtime/scope_stack.rs b/crates/core/src/api/runtime/scope_stack.rs index ac2a70899..ff968e97d 100644 --- a/crates/core/src/api/runtime/scope_stack.rs +++ b/crates/core/src/api/runtime/scope_stack.rs @@ -373,6 +373,13 @@ pub struct ThreadScopeStackBinding { explicit: bool, } +impl ThreadScopeStackBinding { + /// Return the captured thread-local scope stack handle. + pub fn stack(&self) -> ScopeStackHandle { + self.stack.clone() + } +} + /// Create a new scope stack handle with an implicit root scope. /// /// The returned handle wraps a freshly initialized [`ScopeStack`] inside an diff --git a/crates/core/tests/integration/context_isolation_tests.rs b/crates/core/tests/integration/context_isolation_tests.rs index 33a28e530..6cfb583e4 100644 --- a/crates/core/tests/integration/context_isolation_tests.rs +++ b/crates/core/tests/integration/context_isolation_tests.rs @@ -87,6 +87,44 @@ fn test_rootless_propagation_context_uses_the_parent_as_root() { assert_eq!(stack.scopes().len(), 1); } +#[test] +fn test_propagation_context_with_root_as_parent_uses_one_synthetic_root() { + let root_uuid = Uuid::now_v7(); + let stack = create_scope_stack_from_propagation(&PropagationContext { + version: PropagationContext::VERSION, + root_uuid: Some(root_uuid), + parent_uuid: root_uuid, + }) + .unwrap(); + let stack = stack.read().unwrap(); + assert_eq!(stack.root_uuid(), root_uuid); + assert_eq!(stack.top().uuid, root_uuid); + assert_eq!(stack.scopes().len(), 1); +} + +#[test] +fn test_propagation_context_rejects_invalid_wire_values() { + for context in [ + PropagationContext { + version: PropagationContext::VERSION + 1, + root_uuid: None, + parent_uuid: Uuid::now_v7(), + }, + PropagationContext { + version: PropagationContext::VERSION, + root_uuid: None, + parent_uuid: Uuid::nil(), + }, + PropagationContext { + version: PropagationContext::VERSION, + root_uuid: Some(Uuid::from_u128(1_u128 << 64)), + parent_uuid: Uuid::now_v7(), + }, + ] { + assert!(create_scope_stack_from_propagation(&context).is_err()); + } +} + #[test] fn test_pop_scope_rejects_non_top_and_unknown_handles() { set_thread_scope_stack(create_scope_stack()); diff --git a/crates/ffi/tests/unit/api_tests.rs b/crates/ffi/tests/unit/api_tests.rs index b7661aae1..b0fede210 100644 --- a/crates/ffi/tests/unit/api_tests.rs +++ b/crates/ffi/tests/unit/api_tests.rs @@ -232,6 +232,40 @@ unsafe fn fresh_scope_stack() -> *mut FfiScopeStack { stack } +#[test] +fn propagation_context_json_round_trips_through_the_ffi() { + let _guard = lock_unpoisoned(&TEST_MUTEX); + let root_uuid = "018f13f0-7c1a-7a80-8000-000000000001"; + let root = CString::new(root_uuid).unwrap(); + let mut context_json = ptr::null_mut(); + assert_eq!( + unsafe { + crate::api::nemo_relay_capture_propagation_context_with_root_json( + root.as_ptr(), + &mut context_json, + ) + }, + NemoRelayStatus::Ok + ); + let context = unsafe { returned_json(context_json) }; + assert_eq!(context["version"], 1); + assert_eq!(context["root_uuid"], root_uuid); + + let payload = CString::new(context.to_string()).unwrap(); + let mut stack = ptr::null_mut(); + assert_eq!( + unsafe { + crate::api::nemo_relay_scope_stack_create_from_propagation_json( + payload.as_ptr(), + &mut stack, + ) + }, + NemoRelayStatus::Ok + ); + assert!(!stack.is_null()); + unsafe { nemo_relay_scope_stack_free(stack) }; +} + fn reset_globals() { lock_unpoisoned(event_log()).clear(); lock_unpoisoned(collected_chunks()).clear(); diff --git a/crates/node/src/api/mod.rs b/crates/node/src/api/mod.rs index 58dd3f672..9e5955cf3 100644 --- a/crates/node/src/api/mod.rs +++ b/crates/node/src/api/mod.rs @@ -1815,8 +1815,8 @@ pub fn create_scope_stack_from_propagation( /// Run a synchronous callback with an isolated scope stack installed. /// -/// For asynchronous JavaScript request handlers, keep the stack installed with -/// `setThreadScopeStack` for the handler's lifetime instead. +/// The stack is restored before this function returns. Asynchronous callbacks +/// must not rely on this installation after their first `await`. #[napi] pub fn with_scope_stack(stack: &ScopeStack, callback: JsFunction) -> napi::Result { with_scope_stack_handle(stack.inner.clone(), || { diff --git a/crates/node/tests/context_tests.mjs b/crates/node/tests/context_tests.mjs index 5b26280da..83d498dd2 100644 --- a/crates/node/tests/context_tests.mjs +++ b/crates/node/tests/context_tests.mjs @@ -38,20 +38,29 @@ describe('Context isolation', () => { const rootUuid = '018f13f0-7c1a-7a80-8000-000000000001'; const parentUuid = '018f13f0-7c1a-7a80-8000-000000000002'; const stack = createScopeStackFromPropagation({ version: 1, rootUuid, parentUuid }); - setThreadScopeStack(stack); - assert.equal(getHandle().uuid, parentUuid); - setThreadScopeStack(original); + try { + setThreadScopeStack(stack); + assert.equal(getHandle().uuid, parentUuid); + } finally { + setThreadScopeStack(original); + } }); it('restores the surrounding stack after withScopeStack', () => { const original = currentScopeStack(); + const originalUuid = getHandle().uuid; const stack = createScopeStack(); - withScopeStack(stack, () => { - pushScope('temporary-with-scope-stack', ScopeType.Agent, null, null); - assert.equal(getHandle().name, 'temporary-with-scope-stack'); - }); - assert.notEqual(getHandle().name, 'temporary-with-scope-stack'); - setThreadScopeStack(original); + try { + withScopeStack(stack, () => { + pushScope('temporary-with-scope-stack', ScopeType.Agent, null, null); + assert.equal(getHandle().name, 'temporary-with-scope-stack'); + }); + assert.notEqual(getHandle().name, 'temporary-with-scope-stack'); + assert.throws(() => withScopeStack(stack, () => { throw new Error('expected'); }), /expected/); + assert.equal(getHandle().uuid, originalUuid); + } finally { + setThreadScopeStack(original); + } }); it('currentScopeStack returns same in same context', () => { @@ -68,17 +77,16 @@ describe('Context isolation', () => { const original = currentScopeStack(); const newStack = createScopeStack(); - // Switch to new stack and push a scope on it - setThreadScopeStack(newStack); - const scope = pushScope('isolated_scope', ScopeType.Agent, null, null); - const handle = getHandle(); - assert.equal(handle.name, 'isolated_scope'); - popScope(scope); - - // Restore original stack — the isolated scope should not be visible - setThreadScopeStack(original); - const restored = getHandle(); - assert.notEqual(restored.name, 'isolated_scope'); + try { + setThreadScopeStack(newStack); + const scope = pushScope('isolated_scope', ScopeType.Agent, null, null); + const handle = getHandle(); + assert.equal(handle.name, 'isolated_scope'); + popScope(scope); + } finally { + setThreadScopeStack(original); + } + assert.notEqual(getHandle().name, 'isolated_scope'); }); it('scopeStackActive returns true after setThreadScopeStack', () => { diff --git a/crates/python/src/py_api/mod.rs b/crates/python/src/py_api/mod.rs index 4983b777e..6b4e471bb 100644 --- a/crates/python/src/py_api/mod.rs +++ b/crates/python/src/py_api/mod.rs @@ -19,6 +19,7 @@ use nemo_relay::api::runtime::{ use nemo_relay::api::runtime::{ TASK_SCOPE_STACK, capture_propagation_context as capture_propagation_context_handle, capture_propagation_context_with_root as capture_propagation_context_with_root_handle, + capture_thread_scope_stack as capture_thread_scope_stack_handle, create_scope_stack as create_scope_stack_handle, create_scope_stack_from_propagation as create_scope_stack_from_propagation_handle, current_scope_stack as current_scope_stack_handle, scope_stack_active as scope_stack_is_active, @@ -210,6 +211,12 @@ pub fn set_thread_scope_stack(stack: &PyScopeStack) { bind_thread_scope_stack(stack.0.clone()); } +/// Capture the scope stack currently installed in native thread-local storage. +#[pyfunction] +pub fn capture_thread_scope_stack() -> PyScopeStack { + PyScopeStack(capture_thread_scope_stack_handle().stack()) +} + /// Sync a ``ScopeStack`` to the current thread's Rust thread-local storage /// **without** marking it as explicitly set. /// @@ -1779,6 +1786,7 @@ pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(capture_propagation_context_with_root, m)?)?; m.add_function(wrap_pyfunction!(create_scope_stack_from_propagation, m)?)?; m.add_function(wrap_pyfunction!(set_thread_scope_stack, m)?)?; + m.add_function(wrap_pyfunction!(capture_thread_scope_stack, m)?)?; m.add_function(wrap_pyfunction!(sync_thread_scope_stack, m)?)?; m.add_function(wrap_pyfunction!(py_scope_stack_active, m)?)?; diff --git a/go/nemo_relay/context_test.go b/go/nemo_relay/context_test.go index d89891c33..7fc22f6e0 100644 --- a/go/nemo_relay/context_test.go +++ b/go/nemo_relay/context_test.go @@ -283,6 +283,69 @@ func TestNewScopeStackFromPropagationUsesParentAsCurrentHandle(t *testing.T) { }) } +func TestPropagationContextCaptureAndValidation(t *testing.T) { + stack, err := NewScopeStack() + if err != nil { + t.Fatalf(newScopeStackFailed, err) + } + defer stack.Close() + + stack.Run(func() { + context, err := CapturePropagationContext() + if err != nil { + t.Fatalf("CapturePropagationContext failed: %v", err) + } + if context.Version != 1 || context.RootUUID != nil || context.ParentUUID == "" { + t.Fatalf("unexpected rootless context: %+v", context) + } + + rootUUID := "018f13f0-7c1a-7a80-8000-000000000001" + withRoot, err := CapturePropagationContextWithRoot(&rootUUID) + if err != nil { + t.Fatalf("CapturePropagationContextWithRoot failed: %v", err) + } + if withRoot.RootUUID == nil || *withRoot.RootUUID != rootUUID { + t.Fatalf("expected root UUID %s, got %+v", rootUUID, withRoot.RootUUID) + } + }) + + for _, context := range []PropagationContext{ + {Version: 2, ParentUUID: "018f13f0-7c1a-7a80-8000-000000000002"}, + {Version: 1, ParentUUID: "not-a-uuid"}, + } { + if _, err := NewScopeStackFromPropagation(context); err == nil { + t.Fatalf("expected invalid context to be rejected: %+v", context) + } + } +} + +func TestPropagatedScopeStacksRemainIsolated(t *testing.T) { + rootUUID := "018f13f0-7c1a-7a80-8000-000000000001" + first, err := NewScopeStackFromPropagation(PropagationContext{Version: 1, RootUUID: &rootUUID, ParentUUID: "018f13f0-7c1a-7a80-8000-000000000002"}) + if err != nil { + t.Fatal(err) + } + defer first.Close() + second, err := NewScopeStackFromPropagation(PropagationContext{Version: 1, RootUUID: &rootUUID, ParentUUID: "018f13f0-7c1a-7a80-8000-000000000003"}) + if err != nil { + t.Fatal(err) + } + defer second.Close() + + first.Run(func() { + handle, _ := GetHandle() + if handle.UUID() != "018f13f0-7c1a-7a80-8000-000000000002" { + t.Fatalf("unexpected first propagated parent: %s", handle.UUID()) + } + }) + second.Run(func() { + handle, _ := GetHandle() + if handle.UUID() != "018f13f0-7c1a-7a80-8000-000000000003" { + t.Fatalf("unexpected second propagated parent: %s", handle.UUID()) + } + }) +} + func TestConcurrentScopeStacksWithToolCalls(t *testing.T) { const goroutines = 5 var wg sync.WaitGroup diff --git a/python/nemo_relay/__init__.py b/python/nemo_relay/__init__.py index c3071e023..75cdafa54 100644 --- a/python/nemo_relay/__init__.py +++ b/python/nemo_relay/__init__.py @@ -127,6 +127,7 @@ async def main(): from nemo_relay._native import ( capture_propagation_context_with_root as _capture_propagation_context_with_root, ) +from nemo_relay._native import capture_thread_scope_stack as _capture_thread_scope_stack from nemo_relay._native import create_scope_stack as _create_scope_stack from nemo_relay._native import ( create_scope_stack_from_propagation as _create_scope_stack_from_propagation, @@ -424,13 +425,14 @@ def create_scope_stack_from_propagation(context: PropagationContext) -> ScopeSta @contextmanager def use_scope_stack(stack: ScopeStack): """Temporarily install ``stack`` in the current Python context.""" + previous_native_stack = _capture_thread_scope_stack() token = _scope_stack_var.set(stack) _sync_thread_scope_stack(stack) try: yield stack finally: _scope_stack_var.reset(token) - _sync_thread_scope_stack(get_scope_stack()) + _sync_thread_scope_stack(previous_native_stack) def set_thread_scope_stack(stack: ScopeStack) -> None: diff --git a/python/nemo_relay/_native.pyi b/python/nemo_relay/_native.pyi index 93f1628fd..9fd3fda62 100644 --- a/python/nemo_relay/_native.pyi +++ b/python/nemo_relay/_native.pyi @@ -1338,6 +1338,7 @@ class PropagationContext: def capture_propagation_context() -> PropagationContext: ... def capture_propagation_context_with_root(root_uuid: str | None) -> PropagationContext: ... def create_scope_stack_from_propagation(context: PropagationContext) -> ScopeStack: ... +def capture_thread_scope_stack() -> ScopeStack: ... def set_thread_scope_stack(stack: ScopeStack) -> None: """Install a scope stack into native thread-local storage. diff --git a/python/tests/test_context_isolation.py b/python/tests/test_context_isolation.py index 252e69256..b933e3853 100644 --- a/python/tests/test_context_isolation.py +++ b/python/tests/test_context_isolation.py @@ -30,6 +30,18 @@ def test_propagation_context_installs_and_restores_a_scoped_stack(): assert nemo_relay.get_scope_stack() is original +def test_use_scope_stack_restores_a_previously_bound_native_stack(): + previous = nemo_relay.create_scope_stack() + replacement = nemo_relay.create_scope_stack() + nemo_relay.set_thread_scope_stack(previous) + previous_uuid = nemo_relay.scope.get_handle().uuid + + with nemo_relay.use_scope_stack(replacement): + assert nemo_relay.scope.get_handle().uuid != previous_uuid + + assert nemo_relay.scope.get_handle().uuid == previous_uuid + + def test_get_scope_stack_returns_same_in_same_context(): """get_scope_stack returns the same instance within the same context.""" s1 = nemo_relay.get_scope_stack() From 5c22b8f0df30de13f4be6a62c0b8dfcdee3e37e4 Mon Sep 17 00:00:00 2001 From: Will Killian Date: Mon, 27 Jul 2026 13:00:28 -0400 Subject: [PATCH 3/5] fix: preserve propagated root parent context Signed-off-by: Will Killian --- crates/core/src/api/runtime/scope_stack.rs | 5 +-- .../core/src/observability/openinference.rs | 3 -- crates/core/src/observability/otel.rs | 3 -- .../integration/context_isolation_tests.rs | 1 + .../unit/observability/openinference_tests.rs | 38 +++++++++++++++++-- .../tests/unit/observability/otel_tests.rs | 37 +++++++++++++++++- 6 files changed, 72 insertions(+), 15 deletions(-) diff --git a/crates/core/src/api/runtime/scope_stack.rs b/crates/core/src/api/runtime/scope_stack.rs index ff968e97d..0c56e0090 100644 --- a/crates/core/src/api/runtime/scope_stack.rs +++ b/crates/core/src/api/runtime/scope_stack.rs @@ -137,10 +137,7 @@ impl ScopeStack { stack, scope_registries: HashMap::new(), fresh_agents: HashSet::from([root_uuid]), - propagated_parent_uuid: context - .root_uuid - .filter(|root_uuid| *root_uuid != context.parent_uuid) - .map(|_| context.parent_uuid), + propagated_parent_uuid: context.root_uuid.map(|_| context.parent_uuid), }) } diff --git a/crates/core/src/observability/openinference.rs b/crates/core/src/observability/openinference.rs index 05d980869..39d3603e1 100644 --- a/crates/core/src/observability/openinference.rs +++ b/crates/core/src/observability/openinference.rs @@ -804,9 +804,6 @@ impl OpenInferenceEventProcessor { return Context::new(); } let root_uuid = stack.root_uuid(); - if parent_uuid == root_uuid { - return Context::new(); - } Context::new().with_remote_span_context(SpanContext::new( relay_trace_id(root_uuid), relay_span_id(parent_uuid), diff --git a/crates/core/src/observability/otel.rs b/crates/core/src/observability/otel.rs index 841e58f03..873383b5f 100644 --- a/crates/core/src/observability/otel.rs +++ b/crates/core/src/observability/otel.rs @@ -790,9 +790,6 @@ impl OtelEventProcessor { return Context::new(); } let root_uuid = stack.root_uuid(); - if parent_uuid == root_uuid { - return Context::new(); - } Context::new().with_remote_span_context(SpanContext::new( relay_trace_id(root_uuid), relay_span_id(parent_uuid), diff --git a/crates/core/tests/integration/context_isolation_tests.rs b/crates/core/tests/integration/context_isolation_tests.rs index 6cfb583e4..1e9121d56 100644 --- a/crates/core/tests/integration/context_isolation_tests.rs +++ b/crates/core/tests/integration/context_isolation_tests.rs @@ -100,6 +100,7 @@ fn test_propagation_context_with_root_as_parent_uses_one_synthetic_root() { assert_eq!(stack.root_uuid(), root_uuid); assert_eq!(stack.top().uuid, root_uuid); assert_eq!(stack.scopes().len(), 1); + assert!(stack.is_propagated_parent(root_uuid)); } #[test] diff --git a/crates/core/tests/unit/observability/openinference_tests.rs b/crates/core/tests/unit/observability/openinference_tests.rs index 22a3ca14d..2f0d210b9 100644 --- a/crates/core/tests/unit/observability/openinference_tests.rs +++ b/crates/core/tests/unit/observability/openinference_tests.rs @@ -8,8 +8,11 @@ use crate::api::event::{ BaseEvent, CategoryProfile, Event, EventCategory, MarkEvent, ScopeCategory, ScopeEvent, tool_attributes_to_strings, }; -use crate::api::runtime::NemoRelayContextState; -use crate::api::runtime::global_context; +use crate::api::runtime::{ + NemoRelayContextState, PropagationContext, capture_thread_scope_stack, + create_scope_stack_from_propagation, global_context, restore_thread_scope_stack, + set_thread_scope_stack, +}; use crate::api::scope::ScopeType; use crate::api::scope::{event, pop_scope, push_scope}; use crate::api::tool::ToolAttributes; @@ -24,7 +27,7 @@ use crate::codec::response::{ }; use crate::json::Json; use crate::observability::atif::{AtifAgentInfo, AtifExporter, AtifStepExtra}; -use opentelemetry::trace::Status; +use opentelemetry::trace::{Status, TraceContextExt}; use opentelemetry_sdk::trace::InMemorySpanExporterBuilder; use serde_json::json; use std::collections::HashMap; @@ -467,6 +470,35 @@ fn make_start_event( ) } +#[test] +fn propagated_root_parent_projects_as_a_remote_openinference_parent() { + let root_uuid = Uuid::now_v7(); + let previous_stack = capture_thread_scope_stack(); + let imported_stack = create_scope_stack_from_propagation(&PropagationContext { + version: PropagationContext::VERSION, + root_uuid: Some(root_uuid), + parent_uuid: root_uuid, + }) + .unwrap(); + set_thread_scope_stack(imported_stack); + + let processor = OpenInferenceEventProcessor::new(make_provider().0, "test".into()); + let parent_context = processor.parent_context(&make_start_event( + Uuid::now_v7(), + Some(root_uuid), + "receiver-tool", + ScopeType::Tool, + None, + )); + restore_thread_scope_stack(previous_stack); + + let parent_span = parent_context.span(); + let span_context = parent_span.span_context(); + assert!(span_context.is_remote()); + assert_eq!(span_context.trace_id(), relay_trace_id(root_uuid)); + assert_eq!(span_context.span_id(), relay_span_id(root_uuid)); +} + fn make_start_event_with_metadata( uuid: Uuid, parent_uuid: Option, diff --git a/crates/core/tests/unit/observability/otel_tests.rs b/crates/core/tests/unit/observability/otel_tests.rs index c162e8746..95caf01d7 100644 --- a/crates/core/tests/unit/observability/otel_tests.rs +++ b/crates/core/tests/unit/observability/otel_tests.rs @@ -8,8 +8,11 @@ use crate::api::event::{ BaseEvent, CategoryProfile, Event, EventCategory, MarkEvent, ScopeCategory, ScopeEvent, tool_attributes_to_strings, }; -use crate::api::runtime::NemoRelayContextState; -use crate::api::runtime::global_context; +use crate::api::runtime::{ + NemoRelayContextState, PropagationContext, capture_thread_scope_stack, + create_scope_stack_from_propagation, global_context, restore_thread_scope_stack, + set_thread_scope_stack, +}; use crate::api::scope::ScopeType; use crate::api::scope::{event, pop_scope, push_scope}; use crate::api::tool::ToolAttributes; @@ -20,6 +23,7 @@ use crate::codec::response::{ }; use crate::json::Json; use crate::observability::atif::{AtifAgentInfo, AtifExporter, AtifStepExtra}; +use opentelemetry::trace::TraceContextExt; use opentelemetry_sdk::trace::InMemorySpanExporterBuilder; use serde_json::json; use std::collections::HashMap; @@ -331,6 +335,35 @@ fn make_start_event( ) } +#[test] +fn propagated_root_parent_projects_as_a_remote_otel_parent() { + let root_uuid = Uuid::now_v7(); + let previous_stack = capture_thread_scope_stack(); + let imported_stack = create_scope_stack_from_propagation(&PropagationContext { + version: PropagationContext::VERSION, + root_uuid: Some(root_uuid), + parent_uuid: root_uuid, + }) + .unwrap(); + set_thread_scope_stack(imported_stack); + + let processor = OtelEventProcessor::new(make_provider().0, "test".into()); + let parent_context = processor.parent_context(&make_start_event( + Uuid::now_v7(), + Some(root_uuid), + "receiver-tool", + ScopeType::Tool, + None, + )); + restore_thread_scope_stack(previous_stack); + + let parent_span = parent_context.span(); + let span_context = parent_span.span_context(); + assert!(span_context.is_remote()); + assert_eq!(span_context.trace_id(), relay_trace_id(root_uuid)); + assert_eq!(span_context.span_id(), relay_span_id(root_uuid)); +} + fn make_start_event_with_metadata( uuid: Uuid, parent_uuid: Option, From 780df65cbfb1d8861b32bd45293750469d2bd775 Mon Sep 17 00:00:00 2001 From: Will Killian Date: Mon, 27 Jul 2026 15:44:07 -0400 Subject: [PATCH 4/5] test: strengthen propagation binding coverage Signed-off-by: Will Killian --- .../unit/observability/openinference_tests.rs | 14 +++- .../tests/unit/observability/otel_tests.rs | 14 +++- crates/ffi/tests/unit/api_tests.rs | 17 +++++ crates/python/src/py_api/mod.rs | 20 +++-- crates/python/src/py_types/core.rs | 12 +++ crates/python/src/py_types/mod.rs | 1 + .../tests/coverage/py_api_coverage_tests.rs | 19 +++++ go/nemo_relay/context_test.go | 75 +++++++++++++++++++ python/nemo_relay/__init__.py | 3 +- python/nemo_relay/_native.pyi | 6 +- python/tests/test_context_isolation.py | 60 +++++++++++++++ 11 files changed, 226 insertions(+), 15 deletions(-) diff --git a/crates/core/tests/unit/observability/openinference_tests.rs b/crates/core/tests/unit/observability/openinference_tests.rs index 2f0d210b9..a42d593e4 100644 --- a/crates/core/tests/unit/observability/openinference_tests.rs +++ b/crates/core/tests/unit/observability/openinference_tests.rs @@ -9,7 +9,7 @@ use crate::api::event::{ tool_attributes_to_strings, }; use crate::api::runtime::{ - NemoRelayContextState, PropagationContext, capture_thread_scope_stack, + NemoRelayContextState, PropagationContext, ThreadScopeStackBinding, capture_thread_scope_stack, create_scope_stack_from_propagation, global_context, restore_thread_scope_stack, set_thread_scope_stack, }; @@ -45,6 +45,14 @@ impl Drop for ResetPricingResolverGuard { } } +struct RestoreThreadScopeStackGuard(ThreadScopeStackBinding); + +impl Drop for RestoreThreadScopeStackGuard { + fn drop(&mut self) { + restore_thread_scope_stack(self.0.clone()); + } +} + fn reset_global() { let _ = spdlog::init_log_crate_proxy(); log::set_max_level(log::LevelFilter::Info); @@ -473,7 +481,7 @@ fn make_start_event( #[test] fn propagated_root_parent_projects_as_a_remote_openinference_parent() { let root_uuid = Uuid::now_v7(); - let previous_stack = capture_thread_scope_stack(); + let _restore_guard = RestoreThreadScopeStackGuard(capture_thread_scope_stack()); let imported_stack = create_scope_stack_from_propagation(&PropagationContext { version: PropagationContext::VERSION, root_uuid: Some(root_uuid), @@ -490,8 +498,6 @@ fn propagated_root_parent_projects_as_a_remote_openinference_parent() { ScopeType::Tool, None, )); - restore_thread_scope_stack(previous_stack); - let parent_span = parent_context.span(); let span_context = parent_span.span_context(); assert!(span_context.is_remote()); diff --git a/crates/core/tests/unit/observability/otel_tests.rs b/crates/core/tests/unit/observability/otel_tests.rs index 95caf01d7..e93e2357d 100644 --- a/crates/core/tests/unit/observability/otel_tests.rs +++ b/crates/core/tests/unit/observability/otel_tests.rs @@ -9,7 +9,7 @@ use crate::api::event::{ tool_attributes_to_strings, }; use crate::api::runtime::{ - NemoRelayContextState, PropagationContext, capture_thread_scope_stack, + NemoRelayContextState, PropagationContext, ThreadScopeStackBinding, capture_thread_scope_stack, create_scope_stack_from_propagation, global_context, restore_thread_scope_stack, set_thread_scope_stack, }; @@ -41,6 +41,14 @@ impl Drop for ResetPricingResolverGuard { } } +struct RestoreThreadScopeStackGuard(ThreadScopeStackBinding); + +impl Drop for RestoreThreadScopeStackGuard { + fn drop(&mut self) { + restore_thread_scope_stack(self.0.clone()); + } +} + fn empty_annotated_response() -> AnnotatedLlmResponse { AnnotatedLlmResponse { id: None, @@ -338,7 +346,7 @@ fn make_start_event( #[test] fn propagated_root_parent_projects_as_a_remote_otel_parent() { let root_uuid = Uuid::now_v7(); - let previous_stack = capture_thread_scope_stack(); + let _restore_guard = RestoreThreadScopeStackGuard(capture_thread_scope_stack()); let imported_stack = create_scope_stack_from_propagation(&PropagationContext { version: PropagationContext::VERSION, root_uuid: Some(root_uuid), @@ -355,8 +363,6 @@ fn propagated_root_parent_projects_as_a_remote_otel_parent() { ScopeType::Tool, None, )); - restore_thread_scope_stack(previous_stack); - let parent_span = parent_context.span(); let span_context = parent_span.span_context(); assert!(span_context.is_remote()); diff --git a/crates/ffi/tests/unit/api_tests.rs b/crates/ffi/tests/unit/api_tests.rs index b0fede210..2a7f41246 100644 --- a/crates/ffi/tests/unit/api_tests.rs +++ b/crates/ffi/tests/unit/api_tests.rs @@ -264,6 +264,23 @@ fn propagation_context_json_round_trips_through_the_ffi() { ); assert!(!stack.is_null()); unsafe { nemo_relay_scope_stack_free(stack) }; + + for invalid_context in [ + "not-json", + r#"{\"version\":2,\"parent_uuid\":\"018f13f0-7c1a-7a80-8000-000000000002\"}"#, + r#"{\"version\":1,\"root_uuid\":\"00000000-0000-0000-0000-000000000000\",\"parent_uuid\":\"018f13f0-7c1a-7a80-8000-000000000002\"}"#, + ] { + let payload = CString::new(invalid_context).unwrap(); + let mut rejected_stack = ptr::null_mut(); + let status = unsafe { + crate::api::nemo_relay_scope_stack_create_from_propagation_json( + payload.as_ptr(), + &mut rejected_stack, + ) + }; + assert_ne!(status, NemoRelayStatus::Ok); + assert!(rejected_stack.is_null()); + } } fn reset_globals() { diff --git a/crates/python/src/py_api/mod.rs b/crates/python/src/py_api/mod.rs index 6b4e471bb..f5c2ccd51 100644 --- a/crates/python/src/py_api/mod.rs +++ b/crates/python/src/py_api/mod.rs @@ -22,8 +22,9 @@ use nemo_relay::api::runtime::{ capture_thread_scope_stack as capture_thread_scope_stack_handle, create_scope_stack as create_scope_stack_handle, create_scope_stack_from_propagation as create_scope_stack_from_propagation_handle, - current_scope_stack as current_scope_stack_handle, scope_stack_active as scope_stack_is_active, - set_thread_scope_stack as bind_thread_scope_stack, + current_scope_stack as current_scope_stack_handle, + restore_thread_scope_stack as restore_thread_scope_stack_handle, + scope_stack_active as scope_stack_is_active, set_thread_scope_stack as bind_thread_scope_stack, sync_thread_scope_stack as sync_bound_thread_scope_stack, task_scope_top, }; use nemo_relay::api::scope as core_scope_api; @@ -43,7 +44,8 @@ use crate::py_callable; use crate::py_types::{ PyAnnotatedLLMResponse, PyAnthropicMessagesCodec, PyLLMAttributes, PyLLMHandle, PyLLMRequest, PyLlmStream, PyOpenAIChatCodec, PyOpenAIResponsesCodec, PyPropagationContext, - PyScopeAttributes, PyScopeHandle, PyScopeStack, PyScopeType, PyToolAttributes, PyToolHandle, + PyScopeAttributes, PyScopeHandle, PyScopeStack, PyScopeType, PyThreadScopeStackBinding, + PyToolAttributes, PyToolHandle, }; pub(crate) type RustJsonStream = LlmJsonStream; @@ -213,8 +215,15 @@ pub fn set_thread_scope_stack(stack: &PyScopeStack) { /// Capture the scope stack currently installed in native thread-local storage. #[pyfunction] -pub fn capture_thread_scope_stack() -> PyScopeStack { - PyScopeStack(capture_thread_scope_stack_handle().stack()) +pub fn capture_thread_scope_stack() -> PyThreadScopeStackBinding { + PyThreadScopeStackBinding(capture_thread_scope_stack_handle()) +} + +/// Restore a complete native thread binding captured by +/// [`capture_thread_scope_stack`]. +#[pyfunction] +pub fn restore_thread_scope_stack(binding: &PyThreadScopeStackBinding) { + restore_thread_scope_stack_handle(binding.0.clone()); } /// Sync a ``ScopeStack`` to the current thread's Rust thread-local storage @@ -1787,6 +1796,7 @@ pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(create_scope_stack_from_propagation, m)?)?; m.add_function(wrap_pyfunction!(set_thread_scope_stack, m)?)?; m.add_function(wrap_pyfunction!(capture_thread_scope_stack, m)?)?; + m.add_function(wrap_pyfunction!(restore_thread_scope_stack, m)?)?; m.add_function(wrap_pyfunction!(sync_thread_scope_stack, m)?)?; m.add_function(wrap_pyfunction!(py_scope_stack_active, m)?)?; diff --git a/crates/python/src/py_types/core.rs b/crates/python/src/py_types/core.rs index 3648e629d..9a9efb384 100644 --- a/crates/python/src/py_types/core.rs +++ b/crates/python/src/py_types/core.rs @@ -16,6 +16,7 @@ use nemo_relay::api::event::{CategoryProfile, EventCategory, PendingMarkSpec}; use nemo_relay::api::llm::LlmRequestInterceptOutcome; use nemo_relay::api::runtime::{ LlmSanitizeRequestContext, LlmSanitizeResponseContext, PropagationContext, + ThreadScopeStackBinding, }; use nemo_relay::api::tool::ToolExecutionInterceptOutcome; @@ -190,6 +191,17 @@ impl PyScopeStack { } } +/// Opaque captured native thread binding used to restore a Python scope context. +#[pyclass(name = "_ThreadScopeStackBinding")] +pub struct PyThreadScopeStackBinding(pub ThreadScopeStackBinding); + +#[pymethods] +impl PyThreadScopeStackBinding { + pub(crate) fn __repr__(&self) -> String { + "<_ThreadScopeStackBinding>".to_string() + } +} + /// Transport-neutral causal context used to continue Relay work remotely. #[pyclass(name = "PropagationContext", skip_from_py_object)] #[derive(Clone)] diff --git a/crates/python/src/py_types/mod.rs b/crates/python/src/py_types/mod.rs index 9a13bef16..854d4534c 100644 --- a/crates/python/src/py_types/mod.rs +++ b/crates/python/src/py_types/mod.rs @@ -137,6 +137,7 @@ pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> { fn register_runtime_types(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/crates/python/tests/coverage/py_api_coverage_tests.rs b/crates/python/tests/coverage/py_api_coverage_tests.rs index ec4c74fec..0bfe38573 100644 --- a/crates/python/tests/coverage/py_api_coverage_tests.rs +++ b/crates/python/tests/coverage/py_api_coverage_tests.rs @@ -52,6 +52,25 @@ fn py_api_helpers_and_scope_lifecycle_round_trip() { sync_thread_scope_stack(&stack); assert!(py_scope_stack_active()); + let thread_binding = capture_thread_scope_stack(); + assert_eq!(thread_binding.__repr__(), "<_ThreadScopeStackBinding>"); + let replacement = create_scope_stack(); + set_thread_scope_stack(&replacement); + restore_thread_scope_stack(&thread_binding); + assert!(py_scope_stack_active()); + + let rootless_context = capture_propagation_context().unwrap(); + assert_eq!(rootless_context.inner.version, 1); + assert!(rootless_context.inner.root_uuid.is_none()); + let propagation_root = Uuid::now_v7(); + let rooted_context = + capture_propagation_context_with_root(Some(&propagation_root.to_string())).unwrap(); + assert_eq!(rooted_context.inner.root_uuid, Some(propagation_root)); + let propagated_stack = create_scope_stack_from_propagation(&rooted_context).unwrap(); + set_thread_scope_stack(&propagated_stack); + restore_thread_scope_stack(&thread_binding); + assert!(capture_propagation_context_with_root(Some("not-a-uuid")).is_err()); + let handle = get_handle().unwrap(); assert_eq!(handle.inner.name, "root"); diff --git a/go/nemo_relay/context_test.go b/go/nemo_relay/context_test.go index 7fc22f6e0..f5f8a926f 100644 --- a/go/nemo_relay/context_test.go +++ b/go/nemo_relay/context_test.go @@ -307,8 +307,21 @@ func TestPropagationContextCaptureAndValidation(t *testing.T) { if withRoot.RootUUID == nil || *withRoot.RootUUID != rootUUID { t.Fatalf("expected root UUID %s, got %+v", rootUUID, withRoot.RootUUID) } + + withNilRoot, err := CapturePropagationContextWithRoot(nil) + if err != nil { + t.Fatalf("CapturePropagationContextWithRoot(nil) failed: %v", err) + } + if withNilRoot.RootUUID != nil || withNilRoot.ParentUUID != context.ParentUUID { + t.Fatalf("unexpected nil-root context: %+v", withNilRoot) + } }) + invalidRoot := "not-a-uuid" + if _, err := CapturePropagationContextWithRoot(&invalidRoot); err == nil { + t.Fatal("expected invalid root UUID to be rejected") + } + for _, context := range []PropagationContext{ {Version: 2, ParentUUID: "018f13f0-7c1a-7a80-8000-000000000002"}, {Version: 1, ParentUUID: "not-a-uuid"}, @@ -319,6 +332,68 @@ func TestPropagationContextCaptureAndValidation(t *testing.T) { } } +func TestNewScopeStackFromRootlessAndRootParentPropagation(t *testing.T) { + parentUUID := "018f13f0-7c1a-7a80-8000-000000000004" + for _, context := range []PropagationContext{ + {Version: 1, ParentUUID: parentUUID}, + {Version: 1, RootUUID: &parentUUID, ParentUUID: parentUUID}, + } { + stack, err := NewScopeStackFromPropagation(context) + if err != nil { + t.Fatalf("NewScopeStackFromPropagation failed: %v", err) + } + + stack.Run(func() { + handle, err := GetHandle() + if err != nil { + t.Fatalf("GetHandle failed: %v", err) + } + if handle.UUID() != parentUUID { + t.Fatalf("expected propagated parent %s, got %s", parentUUID, handle.UUID()) + } + }) + stack.Close() + } +} + +func TestPropagatedScopeStackRunRestoresOuterBinding(t *testing.T) { + outer, err := NewScopeStack() + if err != nil { + t.Fatalf(newScopeStackFailed, err) + } + defer outer.Close() + + parentUUID := "018f13f0-7c1a-7a80-8000-000000000005" + propagated, err := NewScopeStackFromPropagation(PropagationContext{Version: 1, ParentUUID: parentUUID}) + if err != nil { + t.Fatalf("NewScopeStackFromPropagation failed: %v", err) + } + defer propagated.Close() + + outer.Run(func() { + outerHandle, err := GetHandle() + if err != nil { + t.Fatalf("GetHandle failed: %v", err) + } + propagated.Run(func() { + handle, err := GetHandle() + if err != nil { + t.Fatalf("GetHandle failed: %v", err) + } + if handle.UUID() != parentUUID { + t.Fatalf("expected propagated parent %s, got %s", parentUUID, handle.UUID()) + } + }) + restored, err := GetHandle() + if err != nil { + t.Fatalf("GetHandle failed after propagated Run: %v", err) + } + if restored.UUID() != outerHandle.UUID() { + t.Fatalf("expected outer stack to be restored, got %s", restored.UUID()) + } + }) +} + func TestPropagatedScopeStacksRemainIsolated(t *testing.T) { rootUUID := "018f13f0-7c1a-7a80-8000-000000000001" first, err := NewScopeStackFromPropagation(PropagationContext{Version: 1, RootUUID: &rootUUID, ParentUUID: "018f13f0-7c1a-7a80-8000-000000000002"}) diff --git a/python/nemo_relay/__init__.py b/python/nemo_relay/__init__.py index 75cdafa54..465b3ec17 100644 --- a/python/nemo_relay/__init__.py +++ b/python/nemo_relay/__init__.py @@ -132,6 +132,7 @@ async def main(): from nemo_relay._native import ( create_scope_stack_from_propagation as _create_scope_stack_from_propagation, ) +from nemo_relay._native import restore_thread_scope_stack as _restore_thread_scope_stack from nemo_relay._native import scope_stack_active as _native_scope_stack_active from nemo_relay._native import set_thread_scope_stack as _set_thread_scope_stack from nemo_relay._native import sync_thread_scope_stack as _sync_thread_scope_stack @@ -432,7 +433,7 @@ def use_scope_stack(stack: ScopeStack): yield stack finally: _scope_stack_var.reset(token) - _sync_thread_scope_stack(previous_native_stack) + _restore_thread_scope_stack(previous_native_stack) def set_thread_scope_stack(stack: ScopeStack) -> None: diff --git a/python/nemo_relay/_native.pyi b/python/nemo_relay/_native.pyi index 9fd3fda62..b8e55e0a0 100644 --- a/python/nemo_relay/_native.pyi +++ b/python/nemo_relay/_native.pyi @@ -1338,7 +1338,11 @@ class PropagationContext: def capture_propagation_context() -> PropagationContext: ... def capture_propagation_context_with_root(root_uuid: str | None) -> PropagationContext: ... def create_scope_stack_from_propagation(context: PropagationContext) -> ScopeStack: ... -def capture_thread_scope_stack() -> ScopeStack: ... + +class _ThreadScopeStackBinding: ... + +def capture_thread_scope_stack() -> _ThreadScopeStackBinding: ... +def restore_thread_scope_stack(binding: _ThreadScopeStackBinding) -> None: ... def set_thread_scope_stack(stack: ScopeStack) -> None: """Install a scope stack into native thread-local storage. diff --git a/python/tests/test_context_isolation.py b/python/tests/test_context_isolation.py index b933e3853..546dea800 100644 --- a/python/tests/test_context_isolation.py +++ b/python/tests/test_context_isolation.py @@ -6,6 +6,8 @@ import asyncio import uuid +import pytest + import nemo_relay @@ -30,16 +32,74 @@ def test_propagation_context_installs_and_restores_a_scoped_stack(): assert nemo_relay.get_scope_stack() is original +def test_propagation_context_capture_and_constructor_validation(): + root_uuid = str(uuid.uuid4()) + + with nemo_relay.scope.scope("sender", nemo_relay.ScopeType.Agent) as sender: + rootless = nemo_relay.capture_propagation_context() + rooted = nemo_relay.capture_propagation_context_with_root(root_uuid) + + assert rootless.version == 1 + assert rootless.root_uuid is None + assert rootless.parent_uuid == sender.uuid + assert rooted.version == 1 + assert rooted.root_uuid == root_uuid + assert rooted.parent_uuid == sender.uuid + + for parent_uuid, root, version in [ + ("not-a-uuid", None, 1), + (str(uuid.uuid4()), "not-a-uuid", 1), + (str(uuid.uuid4()), None, 2), + ]: + with pytest.raises(ValueError): + nemo_relay.PropagationContext(parent_uuid, root, version) + + with pytest.raises(ValueError): + nemo_relay.capture_propagation_context_with_root("not-a-uuid") + + +def test_rootless_and_root_parent_propagation_contexts_install_current_handle(): + parent_uuid = str(uuid.uuid4()) + rootless_stack = nemo_relay.create_scope_stack_from_propagation(nemo_relay.PropagationContext(parent_uuid)) + + with nemo_relay.use_scope_stack(rootless_stack): + assert nemo_relay.scope.get_handle().uuid == parent_uuid + + root_stack = nemo_relay.create_scope_stack_from_propagation(nemo_relay.PropagationContext(parent_uuid, parent_uuid)) + with nemo_relay.use_scope_stack(root_stack): + assert nemo_relay.scope.get_handle().uuid == parent_uuid + + def test_use_scope_stack_restores_a_previously_bound_native_stack(): previous = nemo_relay.create_scope_stack() replacement = nemo_relay.create_scope_stack() nemo_relay.set_thread_scope_stack(previous) previous_uuid = nemo_relay.scope.get_handle().uuid + assert nemo_relay.scope_stack_active() with nemo_relay.use_scope_stack(replacement): assert nemo_relay.scope.get_handle().uuid != previous_uuid assert nemo_relay.scope.get_handle().uuid == previous_uuid + assert nemo_relay.scope_stack_active() + + +def test_use_scope_stack_restores_nested_and_failing_contexts(): + previous = nemo_relay.create_scope_stack() + outer = nemo_relay.create_scope_stack() + inner = nemo_relay.create_scope_stack() + nemo_relay.set_thread_scope_stack(previous) + previous_uuid = nemo_relay.scope.get_handle().uuid + + with nemo_relay.use_scope_stack(outer): + outer_uuid = nemo_relay.scope.get_handle().uuid + with pytest.raises(RuntimeError, match="expected failure"): + with nemo_relay.use_scope_stack(inner): + assert nemo_relay.scope.get_handle().uuid != outer_uuid + raise RuntimeError("expected failure") + assert nemo_relay.scope.get_handle().uuid == outer_uuid + + assert nemo_relay.scope.get_handle().uuid == previous_uuid def test_get_scope_stack_returns_same_in_same_context(): From 61af3a7c654a4d1500f6d7e581c03cbcc9aee155 Mon Sep 17 00:00:00 2001 From: Will Killian Date: Mon, 27 Jul 2026 16:30:00 -0400 Subject: [PATCH 5/5] feat: add propagation context JSON helpers Signed-off-by: Will Killian --- crates/core/src/api/runtime/scope_stack.rs | 15 ++++++ .../integration/context_isolation_tests.rs | 19 ++++++++ crates/node/src/api/mod.rs | 16 +++++++ crates/node/tests/context_tests.mjs | 30 +++++++++++- crates/python/src/py_types/core.rs | 15 ++++++ docs/about-nemo-relay/concepts/scopes.mdx | 7 +++ go/nemo_relay/context_test.go | 43 +++++++++++++++++ go/nemo_relay/nemo_relay.go | 43 +++++++++++++---- python/nemo_relay/_native.pyi | 7 +++ python/tests/test_context_isolation.py | 46 ++++++++++++++----- 10 files changed, 219 insertions(+), 22 deletions(-) diff --git a/crates/core/src/api/runtime/scope_stack.rs b/crates/core/src/api/runtime/scope_stack.rs index 0c56e0090..fe0cfa6c9 100644 --- a/crates/core/src/api/runtime/scope_stack.rs +++ b/crates/core/src/api/runtime/scope_stack.rs @@ -56,6 +56,21 @@ impl PropagationContext { /// The current wire-format version. pub const VERSION: u16 = 1; + /// Serialize this validated context for application-managed transport. + pub fn to_json(&self) -> Result { + self.validate()?; + Ok(serde_json::to_string(self).expect("PropagationContext is always JSON serializable")) + } + + /// Deserialize and validate a context received from application-managed transport. + pub fn from_json(value: &str) -> Result { + let context: Self = serde_json::from_str(value).map_err(|error| { + FlowError::InvalidArgument(format!("invalid propagation context JSON: {error}")) + })?; + context.validate()?; + Ok(context) + } + /// Validate a context received from an untrusted transport. pub fn validate(&self) -> Result<()> { if self.version != Self::VERSION { diff --git a/crates/core/tests/integration/context_isolation_tests.rs b/crates/core/tests/integration/context_isolation_tests.rs index 1e9121d56..8cdb9a1eb 100644 --- a/crates/core/tests/integration/context_isolation_tests.rs +++ b/crates/core/tests/integration/context_isolation_tests.rs @@ -126,6 +126,25 @@ fn test_propagation_context_rejects_invalid_wire_values() { } } +#[test] +fn test_propagation_context_json_round_trips_and_validates_input() { + let context = PropagationContext { + version: PropagationContext::VERSION, + root_uuid: Some(Uuid::now_v7()), + parent_uuid: Uuid::now_v7(), + }; + + let json = context.to_json().unwrap(); + assert_eq!(PropagationContext::from_json(&json).unwrap(), context); + assert!(PropagationContext::from_json("not JSON").is_err()); + assert!( + PropagationContext::from_json( + r#"{"version":2,"parent_uuid":"018f13f0-7c1a-7a80-8000-000000000002"}"#, + ) + .is_err() + ); +} + #[test] fn test_pop_scope_rejects_non_top_and_unknown_handles() { set_thread_scope_stack(create_scope_stack()); diff --git a/crates/node/src/api/mod.rs b/crates/node/src/api/mod.rs index 9e5955cf3..c45b3326d 100644 --- a/crates/node/src/api/mod.rs +++ b/crates/node/src/api/mod.rs @@ -1803,6 +1803,22 @@ pub fn capture_propagation_context_with_root( .map_err(|error| napi::Error::from_reason(error.to_string())) } +/// Serialize a Relay causal context to the JSON wire format. +#[napi] +pub fn propagation_context_to_json(context: PropagationContext) -> napi::Result { + propagation_context_from_napi(context)? + .to_json() + .map_err(|error| napi::Error::from_reason(error.to_string())) +} + +/// Deserialize and validate a Relay causal context from the JSON wire format. +#[napi] +pub fn propagation_context_from_json(value: String) -> napi::Result { + nemo_relay::api::runtime::PropagationContext::from_json(&value) + .map(propagation_context_to_napi) + .map_err(|error| napi::Error::from_reason(error.to_string())) +} + /// Create an isolated scope stack seeded from a received propagation context. #[napi] pub fn create_scope_stack_from_propagation( diff --git a/crates/node/tests/context_tests.mjs b/crates/node/tests/context_tests.mjs index 83d498dd2..41cede0e2 100644 --- a/crates/node/tests/context_tests.mjs +++ b/crates/node/tests/context_tests.mjs @@ -19,6 +19,8 @@ const { ScopeType, ScopeStack, createScopeStackFromPropagation, + propagationContextFromJson, + propagationContextToJson, withScopeStack, } = lib; @@ -46,6 +48,26 @@ describe('Context isolation', () => { } }); + it('serializes and validates propagation contexts for transport', () => { + const context = { + version: 1, + rootUuid: '018f13f0-7c1a-7a80-8000-000000000001', + parentUuid: '018f13f0-7c1a-7a80-8000-000000000002', + }; + const encoded = propagationContextToJson(context); + assert.deepEqual(JSON.parse(encoded), { + version: 1, + root_uuid: context.rootUuid, + parent_uuid: context.parentUuid, + }); + assert.deepEqual(propagationContextFromJson(encoded), context); + assert.throws(() => propagationContextFromJson('not JSON'), /invalid propagation context JSON/); + assert.throws( + () => propagationContextFromJson(`{"version":2,"parent_uuid":"${context.parentUuid}"}`), + /unsupported propagation context version 2; expected 1/, + ); + }); + it('restores the surrounding stack after withScopeStack', () => { const original = currentScopeStack(); const originalUuid = getHandle().uuid; @@ -56,7 +78,13 @@ describe('Context isolation', () => { assert.equal(getHandle().name, 'temporary-with-scope-stack'); }); assert.notEqual(getHandle().name, 'temporary-with-scope-stack'); - assert.throws(() => withScopeStack(stack, () => { throw new Error('expected'); }), /expected/); + assert.throws( + () => + withScopeStack(stack, () => { + throw new Error('expected'); + }), + /expected/, + ); assert.equal(getHandle().uuid, originalUuid); } finally { setThreadScopeStack(original); diff --git a/crates/python/src/py_types/core.rs b/crates/python/src/py_types/core.rs index 9a9efb384..3eb57bf5a 100644 --- a/crates/python/src/py_types/core.rs +++ b/crates/python/src/py_types/core.rs @@ -243,6 +243,21 @@ impl PyPropagationContext { fn parent_uuid(&self) -> String { self.inner.parent_uuid.to_string() } + + /// Serialize this context to the Relay JSON wire format. + fn to_json(&self) -> PyResult { + self.inner + .to_json() + .map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string())) + } + + /// Deserialize and validate a context from the Relay JSON wire format. + #[staticmethod] + fn from_json(value: &str) -> PyResult { + PropagationContext::from_json(value) + .map(|inner| Self { inner }) + .map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string())) + } } // --------------------------------------------------------------------------- diff --git a/docs/about-nemo-relay/concepts/scopes.mdx b/docs/about-nemo-relay/concepts/scopes.mdx index 49ac3bd6d..72d8f4ea2 100644 --- a/docs/about-nemo-relay/concepts/scopes.mdx +++ b/docs/about-nemo-relay/concepts/scopes.mdx @@ -160,6 +160,13 @@ Relay context is distinct from W3C propagation. An integration may carry `traceparent` and `tracestate` alongside Relay's JSON context when it needs to preserve OpenTelemetry sampling or vendor state. +Use the binding's JSON helpers at the transport boundary: Rust +`PropagationContext::to_json` and `PropagationContext::from_json`, Python +`context.to_json()` and `PropagationContext.from_json(...)`, Go +`context.ToJSON()` and `PropagationContextFromJSON(...)`, or Node.js +`propagationContextToJson(...)` and `propagationContextFromJson(...)`. The +helpers validate the version and UUIDs before a context is imported. + ## Practical Guidance Use these practices when applying the concept in application or integration code. diff --git a/go/nemo_relay/context_test.go b/go/nemo_relay/context_test.go index f5f8a926f..e4c39fd92 100644 --- a/go/nemo_relay/context_test.go +++ b/go/nemo_relay/context_test.go @@ -332,6 +332,49 @@ func TestPropagationContextCaptureAndValidation(t *testing.T) { } } +func TestPropagationContextJSONRoundTripAndValidation(t *testing.T) { + rootUUID := "018f13f0-7c1a-7a80-8000-000000000001" + context := PropagationContext{ + Version: 1, + RootUUID: &rootUUID, + ParentUUID: "018f13f0-7c1a-7a80-8000-000000000002", + } + + payload, err := context.ToJSON() + if err != nil { + t.Fatalf("ToJSON failed: %v", err) + } + var wire map[string]any + if err := json.Unmarshal([]byte(payload), &wire); err != nil { + t.Fatalf("serialized context was not JSON: %v", err) + } + if wire["version"] != float64(1) || wire["root_uuid"] != rootUUID || wire["parent_uuid"] != context.ParentUUID { + t.Fatalf("unexpected propagation JSON: %s", payload) + } + + decoded, err := PropagationContextFromJSON(payload) + if err != nil { + t.Fatalf("PropagationContextFromJSON failed: %v", err) + } + if decoded.Version != context.Version || decoded.RootUUID == nil || *decoded.RootUUID != rootUUID || decoded.ParentUUID != context.ParentUUID { + t.Fatalf("expected round-tripped context %+v, got %+v", context, decoded) + } + + for _, payload := range []string{ + "not JSON", + `{"version":2,"parent_uuid":"018f13f0-7c1a-7a80-8000-000000000002"}`, + `{"version":1,"parent_uuid":"not-a-uuid"}`, + } { + if _, err := PropagationContextFromJSON(payload); err == nil { + t.Fatalf("expected invalid context JSON to be rejected: %s", payload) + } + } + + if _, err := (PropagationContext{Version: 1, ParentUUID: "not-a-uuid"}).ToJSON(); err == nil { + t.Fatal("expected ToJSON to reject an invalid propagation context") + } +} + func TestNewScopeStackFromRootlessAndRootParentPropagation(t *testing.T) { parentUUID := "018f13f0-7c1a-7a80-8000-000000000004" for _, context := range []PropagationContext{ diff --git a/go/nemo_relay/nemo_relay.go b/go/nemo_relay/nemo_relay.go index 60057979e..8f4c510d9 100644 --- a/go/nemo_relay/nemo_relay.go +++ b/go/nemo_relay/nemo_relay.go @@ -1578,6 +1578,37 @@ type PropagationContext struct { ParentUUID string `json:"parent_uuid"` } +// ToJSON serializes a validated propagation context for application-managed transport. +func (context PropagationContext) ToJSON() (string, error) { + if err := validatePropagationContext(context); err != nil { + return "", err + } + // PropagationContext has only JSON-native fields, so marshaling cannot fail. + payload, _ := json.Marshal(context) + return string(payload), nil +} + +// PropagationContextFromJSON deserializes and validates a transport context. +func PropagationContextFromJSON(value string) (PropagationContext, error) { + var context PropagationContext + if err := json.Unmarshal([]byte(value), &context); err != nil { + return PropagationContext{}, err + } + if err := validatePropagationContext(context); err != nil { + return PropagationContext{}, err + } + return context, nil +} + +func validatePropagationContext(context PropagationContext) error { + stack, err := NewScopeStackFromPropagation(context) + if err != nil { + return err + } + stack.Close() + return nil +} + // CapturePropagationContext captures the current Relay causal parent. func CapturePropagationContext() (PropagationContext, error) { var out *C.char @@ -1585,11 +1616,7 @@ func CapturePropagationContext() (PropagationContext, error) { return PropagationContext{}, err } defer C.nemo_relay_string_free(out) - var context PropagationContext - if err := json.Unmarshal([]byte(C.GoString(out)), &context); err != nil { - return PropagationContext{}, err - } - return context, nil + return PropagationContextFromJSON(C.GoString(out)) } // CapturePropagationContextWithRoot captures the current parent with an @@ -1605,11 +1632,7 @@ func CapturePropagationContextWithRoot(rootUUID *string) (PropagationContext, er return PropagationContext{}, err } defer C.nemo_relay_string_free(out) - var context PropagationContext - if err := json.Unmarshal([]byte(C.GoString(out)), &context); err != nil { - return PropagationContext{}, err - } - return context, nil + return PropagationContextFromJSON(C.GoString(out)) } // NewScopeStack creates a new isolated scope stack. diff --git a/python/nemo_relay/_native.pyi b/python/nemo_relay/_native.pyi index b8e55e0a0..f91184e42 100644 --- a/python/nemo_relay/_native.pyi +++ b/python/nemo_relay/_native.pyi @@ -1334,6 +1334,13 @@ class PropagationContext: def root_uuid(self) -> str | None: ... @property def parent_uuid(self) -> str: ... + def to_json(self) -> str: + """Serialize this context to the Relay JSON wire format.""" + ... + @staticmethod + def from_json(value: str) -> PropagationContext: + """Deserialize and validate a Relay JSON wire context.""" + ... def capture_propagation_context() -> PropagationContext: ... def capture_propagation_context_with_root(root_uuid: str | None) -> PropagationContext: ... diff --git a/python/tests/test_context_isolation.py b/python/tests/test_context_isolation.py index 546dea800..49314e312 100644 --- a/python/tests/test_context_isolation.py +++ b/python/tests/test_context_isolation.py @@ -9,6 +9,17 @@ import pytest import nemo_relay +from nemo_relay._native import capture_thread_scope_stack, restore_thread_scope_stack + + +@pytest.fixture +def restore_native_scope_stack(): + """Restore the test thread's native scope binding after each test.""" + binding = capture_thread_scope_stack() + try: + yield + finally: + restore_thread_scope_stack(binding) def test_create_scope_stack_returns_scope_stack(): @@ -46,18 +57,31 @@ def test_propagation_context_capture_and_constructor_validation(): assert rooted.root_uuid == root_uuid assert rooted.parent_uuid == sender.uuid - for parent_uuid, root, version in [ - ("not-a-uuid", None, 1), - (str(uuid.uuid4()), "not-a-uuid", 1), - (str(uuid.uuid4()), None, 2), - ]: - with pytest.raises(ValueError): - nemo_relay.PropagationContext(parent_uuid, root, version) - - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="invalid character"): + nemo_relay.PropagationContext("not-a-uuid") + with pytest.raises(ValueError, match="invalid character"): + nemo_relay.PropagationContext(str(uuid.uuid4()), "not-a-uuid") + with pytest.raises(ValueError, match="unsupported propagation context version 2; expected 1"): + nemo_relay.PropagationContext(str(uuid.uuid4()), version=2) + with pytest.raises(ValueError, match="invalid character"): nemo_relay.capture_propagation_context_with_root("not-a-uuid") +def test_propagation_context_json_round_trip_and_validation(): + context = nemo_relay.PropagationContext(str(uuid.uuid4()), str(uuid.uuid4())) + + encoded = context.to_json() + decoded = nemo_relay.PropagationContext.from_json(encoded) + + assert decoded.version == context.version + assert decoded.root_uuid == context.root_uuid + assert decoded.parent_uuid == context.parent_uuid + with pytest.raises(ValueError, match="invalid propagation context JSON"): + nemo_relay.PropagationContext.from_json("not JSON") + with pytest.raises(ValueError, match="unsupported propagation context version 2; expected 1"): + nemo_relay.PropagationContext.from_json(f'{{"version":2,"parent_uuid":"{uuid.uuid4()}"}}') + + def test_rootless_and_root_parent_propagation_contexts_install_current_handle(): parent_uuid = str(uuid.uuid4()) rootless_stack = nemo_relay.create_scope_stack_from_propagation(nemo_relay.PropagationContext(parent_uuid)) @@ -70,7 +94,7 @@ def test_rootless_and_root_parent_propagation_contexts_install_current_handle(): assert nemo_relay.scope.get_handle().uuid == parent_uuid -def test_use_scope_stack_restores_a_previously_bound_native_stack(): +def test_use_scope_stack_restores_a_previously_bound_native_stack(restore_native_scope_stack): previous = nemo_relay.create_scope_stack() replacement = nemo_relay.create_scope_stack() nemo_relay.set_thread_scope_stack(previous) @@ -84,7 +108,7 @@ def test_use_scope_stack_restores_a_previously_bound_native_stack(): assert nemo_relay.scope_stack_active() -def test_use_scope_stack_restores_nested_and_failing_contexts(): +def test_use_scope_stack_restores_nested_and_failing_contexts(restore_native_scope_stack): previous = nemo_relay.create_scope_stack() outer = nemo_relay.create_scope_stack() inner = nemo_relay.create_scope_stack()