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..fe0cfa6c9 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,66 @@ 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; + + /// 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 { + 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 +111,49 @@ 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.map(|_| context.parent_uuid), + }) } /// Push a scope handle onto the top of the stack. @@ -98,6 +202,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 @@ -276,6 +385,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 @@ -290,9 +406,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..39d3603e1 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,28 @@ 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(); + 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..873383b5f 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,28 @@ 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(); + 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..8cdb9a1eb 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,95 @@ 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_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); + assert!(stack.is_propagated_parent(root_uuid)); +} + +#[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_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/core/tests/unit/observability/openinference_tests.rs b/crates/core/tests/unit/observability/openinference_tests.rs index 22a3ca14d..a42d593e4 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, ThreadScopeStackBinding, 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; @@ -42,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); @@ -467,6 +478,33 @@ fn make_start_event( ) } +#[test] +fn propagated_root_parent_projects_as_a_remote_openinference_parent() { + let root_uuid = Uuid::now_v7(); + 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), + 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, + )); + 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..e93e2357d 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, ThreadScopeStackBinding, 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; @@ -37,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, @@ -331,6 +343,33 @@ fn make_start_event( ) } +#[test] +fn propagated_root_parent_projects_as_a_remote_otel_parent() { + let root_uuid = Uuid::now_v7(); + 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), + 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, + )); + 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/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/ffi/tests/unit/api_tests.rs b/crates/ffi/tests/unit/api_tests.rs index b7661aae1..2a7f41246 100644 --- a/crates/ffi/tests/unit/api_tests.rs +++ b/crates/ffi/tests/unit/api_tests.rs @@ -232,6 +232,57 @@ 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) }; + + 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() { 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 1f8384fae..c45b3326d 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,66 @@ 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())) +} + +/// 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( + 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. +/// +/// 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(), || { + 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..41cede0e2 100644 --- a/crates/node/tests/context_tests.mjs +++ b/crates/node/tests/context_tests.mjs @@ -18,6 +18,10 @@ const { popScope, ScopeType, ScopeStack, + createScopeStackFromPropagation, + propagationContextFromJson, + propagationContextToJson, + withScopeStack, } = lib; // =========================================================================== @@ -31,6 +35,62 @@ 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 }); + try { + setThreadScopeStack(stack); + assert.equal(getHandle().uuid, parentUuid); + } finally { + setThreadScopeStack(original); + } + }); + + 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; + const stack = createScopeStack(); + 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', () => { const s1 = currentScopeStack(); const s2 = currentScopeStack(); @@ -45,17 +105,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 64e2341e3..f5c2ccd51 100644 --- a/crates/python/src/py_api/mod.rs +++ b/crates/python/src/py_api/mod.rs @@ -17,9 +17,14 @@ 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, - 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_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, + 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; @@ -38,8 +43,9 @@ 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, PyThreadScopeStackBinding, + PyToolAttributes, PyToolHandle, }; pub(crate) type RustJsonStream = LlmJsonStream; @@ -162,6 +168,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 @@ -175,6 +213,19 @@ 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() -> 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 /// **without** marking it as explicitly set. /// @@ -1740,7 +1791,12 @@ 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!(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 d0f122311..3eb57bf5a 100644 --- a/crates/python/src/py_types/core.rs +++ b/crates/python/src/py_types/core.rs @@ -14,7 +14,10 @@ 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, + ThreadScopeStackBinding, +}; use nemo_relay::api::tool::ToolExecutionInterceptOutcome; /// Structured identity of the codec active during LLM sanitization. @@ -188,6 +191,75 @@ 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)] +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() + } + + /// 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())) + } +} + // --------------------------------------------------------------------------- // ScopeAttributes (bitflag wrapper) // --------------------------------------------------------------------------- diff --git a/crates/python/src/py_types/mod.rs b/crates/python/src/py_types/mod.rs index eb60374d9..854d4534c 100644 --- a/crates/python/src/py_types/mod.rs +++ b/crates/python/src/py_types/mod.rs @@ -137,6 +137,8 @@ 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::()?; 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/docs/about-nemo-relay/concepts/scopes.mdx b/docs/about-nemo-relay/concepts/scopes.mdx index 84b293f96..72d8f4ea2 100644 --- a/docs/about-nemo-relay/concepts/scopes.mdx +++ b/docs/about-nemo-relay/concepts/scopes.mdx @@ -142,6 +142,31 @@ 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. + +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/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..e4c39fd92 100644 --- a/go/nemo_relay/context_test.go +++ b/go/nemo_relay/context_test.go @@ -259,6 +259,211 @@ 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 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) + } + + 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"}, + } { + if _, err := NewScopeStackFromPropagation(context); err == nil { + t.Fatalf("expected invalid context to be rejected: %+v", context) + } + } +} + +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{ + {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"}) + 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/go/nemo_relay/nemo_relay.go b/go/nemo_relay/nemo_relay.go index 937427e5d..8f4c510d9 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,91 @@ 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"` +} + +// 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 + if err := checkStatus(C.nemo_relay_capture_propagation_context_json(&out)); err != nil { + return PropagationContext{}, err + } + defer C.nemo_relay_string_free(out) + return PropagationContextFromJSON(C.GoString(out)) +} + +// 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) + return PropagationContextFromJSON(C.GoString(out)) +} + // 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..465b3ec17 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,18 @@ 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 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, +) +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 @@ -393,6 +406,36 @@ 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.""" + 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) + _restore_thread_scope_stack(previous_native_stack) + + def set_thread_scope_stack(stack: ScopeStack) -> None: """Install a scope stack into the current thread's native runtime context. @@ -461,11 +504,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..f91184e42 100644 --- a/python/nemo_relay/_native.pyi +++ b/python/nemo_relay/_native.pyi @@ -1325,6 +1325,31 @@ 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 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: ... +def create_scope_stack_from_propagation(context: PropagationContext) -> 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 0c4146a00..49314e312 100644 --- a/python/tests/test_context_isolation.py +++ b/python/tests/test_context_isolation.py @@ -4,8 +4,22 @@ """Tests for per-request scope stack isolation via ContextVar.""" import asyncio +import uuid + +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(): @@ -15,6 +29,103 @@ 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_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 + + 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)) + + 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(restore_native_scope_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(restore_native_scope_stack): + 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(): """get_scope_stack returns the same instance within the same context.""" s1 = nemo_relay.get_scope_stack()