From ade2e2afd9ebcbb750d26760b997953e19ea2250 Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Sun, 11 Oct 2026 14:14:01 -0700 Subject: [PATCH 1/3] feat(relay): export nested routing spans Signed-off-by: Alex Fournier --- Cargo.lock | 2 + README.md | 2 +- crates/libsy-llm-client/src/run.rs | 47 +- .../libsy-llm-client/tests/observability.rs | 50 +- .../switchyard-nemo-relay-plugin/Cargo.toml | 2 + crates/switchyard-nemo-relay-plugin/README.md | 6 + .../switchyard-nemo-relay-plugin/src/lib.rs | 13 +- .../src/redaction_tests.rs | 18 +- .../src/runtime.rs | 50 +- .../src/span_capture.rs | 543 ++++++++++++++++++ docs/integrations/nemo_relay.md | 32 +- 11 files changed, 732 insertions(+), 33 deletions(-) create mode 100644 crates/switchyard-nemo-relay-plugin/src/span_capture.rs diff --git a/Cargo.lock b/Cargo.lock index c1664fc838..ec5215206c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3362,6 +3362,8 @@ dependencies = [ "tempfile", "tokio", "toml 1.1.3+spec-1.1.0", + "tracing", + "tracing-subscriber", "wiremock", ] diff --git a/README.md b/README.md index ab16fee2cd..9bad346903 100644 --- a/README.md +++ b/README.md @@ -24,7 +24,7 @@ Evaluate the complete agent, model pool, and routing configuration against your | --- | --- | --- | | **OpenRouter** | [Use Switchyard through OpenRouter](https://openrouter.ai/docs/guides/routing/routers/switchyard-router) | Hosted by OpenRouter. Set `model` to `nvidia/switchyard` in API requests. | | **LiteLLM** | [Run the Switchyard routing-plugin example](examples/litellm/README.md#quick-start-with-the-local-proxy) | Experimental and checkout-only. The example pins LiteLLM 1.102.0 and supports Stage routing based on request history, plus Random routing. It cannot service intermediate model calls required by classifier or escalation algorithms. | -| **NeMo Relay** | [Build and configure the native plugin](crates/switchyard-nemo-relay-plugin/README.md#build-from-source) | Requires Relay `>=0.8.0, <1.0.0`. The source-build path requires a Rust toolchain and Python 3. | +| **NeMo Relay** | [Build and configure the native plugin](crates/switchyard-nemo-relay-plugin/README.md#build-from-source) | Requires Relay `>=0.10.0, <1.0.0`. The source-build path requires a Rust toolchain and Python 3. | Relay 0.8.x and 0.9.0 can lose upstream error status and details when the plugin is enabled, even for models outside its routes. Review the diff --git a/crates/libsy-llm-client/src/run.rs b/crates/libsy-llm-client/src/run.rs index ab8a3e9cb2..aaca962e63 100644 --- a/crates/libsy-llm-client/src/run.rs +++ b/crates/libsy-llm-client/src/run.rs @@ -232,7 +232,14 @@ async fn serve( let algorithm = call.algorithm.clone(); let candidates = call.models.len(); call.respond(call_one( - &clients, &target, request, &algorithm, &observe, 0, candidates, true, + &clients, + &target, + request, + &algorithm, + &observe, + 0, + candidates, + CallPhase::Routing, )) .await } @@ -255,7 +262,7 @@ async fn call_first_available( observe, index, models.len(), - false, + CallPhase::Completion, ) .await { @@ -275,6 +282,33 @@ async fn call_first_available( Err(LibsyError::NoTargets) } +/// The stage of the Switchyard run in which a model call was issued. +/// +/// A routing-phase call may still supply the final answer. This is deliberately +/// distinct from the `routing`/`answer` role assigned after the outcome is known. +#[derive(Clone, Copy)] +enum CallPhase { + Routing, + Completion, +} + +impl CallPhase { + const fn as_str(self) -> &'static str { + match self { + Self::Routing => "routing", + Self::Completion => "completion", + } + } + + const fn buffers_stream(self) -> bool { + matches!(self, Self::Routing) + } + + const fn records_routed_request(self) -> bool { + matches!(self, Self::Completion) + } +} + /// Call one candidate model and record its observation and span. #[allow(clippy::too_many_arguments)] #[tracing::instrument( @@ -286,6 +320,7 @@ async fn call_first_available( switchyard.algorithm = algorithm, switchyard.candidate = index + 1, switchyard.candidate_count = count, + switchyard.call_phase = phase.as_str(), selected_model = %model_id, otel.kind = "client", otel.name = %format_args!("chat {model_id}"), @@ -325,7 +360,7 @@ async fn call_one( index: usize, // count is for span log count: usize, - buffer: bool, + phase: CallPhase, ) -> Result { let span = tracing::Span::current(); observability::record_gen_ai_request(&span, &request.llm_request); @@ -342,7 +377,9 @@ async fn call_one( let started = Instant::now(); let result = async { let mut response = client?.call(request).await?; - if buffer && let LlmResponse::Stream(chunks) = response.llm_response { + if phase.buffers_stream() + && let LlmResponse::Stream(chunks) = response.llm_response + { response.llm_response = LlmResponse::Stream(buffer_routing_stream(chunks).await?); } Ok(response) @@ -356,7 +393,7 @@ async fn call_one( response }); let result = observability::observe_client_call(result); - let result = if !buffer { + let result = if phase.records_routed_request() { metrics::observe_routed_request(algorithm, model_id, Some(duration), result) } else { result diff --git a/crates/libsy-llm-client/tests/observability.rs b/crates/libsy-llm-client/tests/observability.rs index 37a71f3dbe..8666264b7a 100644 --- a/crates/libsy-llm-client/tests/observability.rs +++ b/crates/libsy-llm-client/tests/observability.rs @@ -1387,6 +1387,45 @@ async fn observed_run_reports_one_successful_routed_call() -> switchyard_libsy:: Ok(()) } +#[tokio::test] +async fn answered_routing_call_keeps_its_execution_phase() -> switchyard_libsy::Result<()> { + let _guard = serialize_test().lock().await; + let (store, _, _, _, _) = telemetry(); + let observations = Arc::new(Mutex::new(Vec::new())); + let observed = Arc::clone(&observations); + let observer: RunObserver = Arc::new(move |observation| observed.lock().push(observation)); + const ALGO: &str = "answered-routing-call"; + const MODEL: &str = "answer-from-routing"; + + switchyard_llm_client::run( + Arc::new(RoutingCallAlgo { + name: ALGO.to_string(), + target: MODEL.into(), + }), + ClientRouter::single(Arc::new(UsageClient { + usage: Usage::default(), + })), + request_with_metadata("answered-routing-session", "answered-routing-correlation"), + Arc::new(RuntimeModels::default()), + Some(observer), + ) + .await?; + + let span = find_span(&store.spans(), "libsy.client_call", "selected_model", MODEL); + assert_eq!( + span.fields.get("switchyard.call_phase").map(String::as_str), + Some("routing") + ); + assert!( + observations + .lock() + .iter() + .any(|observation| matches!(observation, RunObservation::AnswerCall(_))), + "the post-outcome observation should still classify the call as the answer" + ); + Ok(()) +} + /// A streamed response keeps the client span available until terminal usage arrives. struct StreamingUsageClient; @@ -1839,7 +1878,7 @@ async fn failed_call_records_metrics_without_error_details() -> switchyard_libsy #[tokio::test] async fn classifier_metrics_count_routing_and_answer_calls_once() -> switchyard_libsy::Result<()> { let _guard = serialize_test().lock().await; - let (_store, exporter, provider, _, _) = telemetry(); + let (store, exporter, provider, _, _) = telemetry(); let before = flushed_metrics(exporter, provider); let total_requests_before = u64_gauge_value(&before, "switchyard.total_requests").unwrap_or_default(); @@ -1854,6 +1893,15 @@ async fn classifier_metrics_count_routing_and_answer_calls_once() -> switchyard_ assert_eq!(selected_model, "weak"); + let spans = store.spans(); + for (model, phase) in [("classifier", "routing"), ("weak", "completion")] { + let span = find_span(&spans, "libsy.client_call", "selected_model", model); + assert_eq!( + span.fields.get("switchyard.call_phase").map(String::as_str), + Some(phase) + ); + } + let snapshots = flushed_metrics(exporter, provider); assert_eq!( u64_counter_value( diff --git a/crates/switchyard-nemo-relay-plugin/Cargo.toml b/crates/switchyard-nemo-relay-plugin/Cargo.toml index ac9099f7f7..1e893593df 100644 --- a/crates/switchyard-nemo-relay-plugin/Cargo.toml +++ b/crates/switchyard-nemo-relay-plugin/Cargo.toml @@ -27,6 +27,8 @@ switchyard-runner.workspace = true switchyard-translation.workspace = true tokio.workspace = true toml = "1.1" +tracing.workspace = true +tracing-subscriber.workspace = true [dev-dependencies] tempfile = "3" diff --git a/crates/switchyard-nemo-relay-plugin/README.md b/crates/switchyard-nemo-relay-plugin/README.md index 208de4386a..34fa3dbab2 100644 --- a/crates/switchyard-nemo-relay-plugin/README.md +++ b/crates/switchyard-nemo-relay-plugin/README.md @@ -231,6 +231,12 @@ marks. Call marks distinguish routing from answer calls; decisions distinguish selected from served models. Token metrics cover both call roles, while Relay retains ownership of the outer LLM lifecycle. +The plugin also captures the existing `libsy.client_call` and +`libsy.upstream_attempt` spans for each routed request. It replays completed +spans beneath Relay's outer LLM scope with their original timestamps. The +projection includes allowlisted, payload-free routing metadata only; it excludes +prompts, responses, provider error bodies, and credentials. + ## Observability When Relay is configured with OTLP logs and metrics exporters, the plugin emits diff --git a/crates/switchyard-nemo-relay-plugin/src/lib.rs b/crates/switchyard-nemo-relay-plugin/src/lib.rs index ba8db5de67..15ffc430b7 100644 --- a/crates/switchyard-nemo-relay-plugin/src/lib.rs +++ b/crates/switchyard-nemo-relay-plugin/src/lib.rs @@ -3,6 +3,7 @@ mod config; mod runtime; +mod span_capture; mod translation; use std::sync::Arc; @@ -76,7 +77,17 @@ fn register_buffered( return next.call(request).await; } let decoded = runtime.decode_request(inbound, request, false)?; - let execution = runtime.execute_buffered(inbound, decoded).await; + let span_plugin_runtime = plugin_runtime.clone(); + let span_runtime = Arc::clone(&runtime); + let execution = runtime + .execute_buffered( + inbound, + decoded, + Arc::new(move |event| { + emit_event(&span_plugin_runtime, span_runtime.redactor(), event) + }), + ) + .await; emit_events(&plugin_runtime, runtime.redactor(), execution.events); execution.result } diff --git a/crates/switchyard-nemo-relay-plugin/src/redaction_tests.rs b/crates/switchyard-nemo-relay-plugin/src/redaction_tests.rs index e25e0ad7de..e30b4ad39a 100644 --- a/crates/switchyard-nemo-relay-plugin/src/redaction_tests.rs +++ b/crates/switchyard-nemo-relay-plugin/src/redaction_tests.rs @@ -142,7 +142,14 @@ async fn reflected_response(format: WireFormat, streaming: bool, fail: bool) { assert_eq!(events.iter().any(Result::is_err), fail); format!("{events:?}") } else { - let execution = runtime.execute_buffered(format, request).await; + let observed = Arc::clone(&captured); + let execution = runtime + .execute_buffered( + format, + request, + Arc::new(move |event| observed.lock().unwrap().push(event)), + ) + .await; captured.lock().unwrap().extend(execution.events); assert_eq!(execution.result.is_err(), fail); format!("{:?}", execution.result) @@ -155,7 +162,14 @@ async fn reflected_response(format: WireFormat, streaming: bool, fail: bool) { assert!(output.contains("ordinary diagnostics"), "{output}"); assert!(output.contains("[REDACTED]"), "{output}"); } - let telemetry = format!("{:?}", captured.lock().unwrap()); + let captured = captured.lock().unwrap(); + assert!( + captured + .iter() + .any(|event| matches!(event, RoutingEvent::Span(_))), + "client span was not emitted" + ); + let telemetry = format!("{captured:?}"); assert!(!telemetry.contains(KEY), "{telemetry}"); assert!(telemetry.contains("[REDACTED]"), "{telemetry}"); } diff --git a/crates/switchyard-nemo-relay-plugin/src/runtime.rs b/crates/switchyard-nemo-relay-plugin/src/runtime.rs index 4c1b7ab057..74cd0a606f 100644 --- a/crates/switchyard-nemo-relay-plugin/src/runtime.rs +++ b/crates/switchyard-nemo-relay-plugin/src/runtime.rs @@ -22,6 +22,7 @@ use switchyard_runner::{ use switchyard_translation::{TranslationEngine, encode_stream_with_extensions}; use crate::config::SwitchyardConfig; +use crate::span_capture::{CapturedSpan, capture_client_spans}; use crate::translation; #[cfg(test)] @@ -59,6 +60,7 @@ pub(crate) struct RoutingMetric { pub(crate) enum RoutingEvent { Mark(RoutingMark), Metric(RoutingMetric), + Span(CapturedSpan), } struct MetricDescriptor<'a> { @@ -139,9 +141,11 @@ impl SwitchyardRuntime { &self, inbound: WireFormat, request: Request, + emit_event: RoutingEventEmitter, ) -> Execution { let request_extensions = request.llm_request.extensions.clone(); - let Execution { result, mut events } = self.execute(inbound, request).await; + let emit_event = self.sanitized_emitter(emit_event); + let Execution { result, mut events } = self.execute(inbound, request, emit_event).await; let (result, finalization_failed) = match result { Ok(routed) => { let result = finalize_buffered_response( @@ -170,19 +174,17 @@ impl SwitchyardRuntime { request: Request, emit_event: RoutingEventEmitter, ) -> Execution { - let redactor = Arc::clone(&self.redactor); - let emit_event: RoutingEventEmitter = Arc::new(move |event| { - emit_event(sanitize_event(&redactor, event)); - }); + let emit_event = self.sanitized_emitter(emit_event); let request_extensions = request.llm_request.extensions.clone(); - let Execution { result, mut events } = self.execute(inbound, request).await; + let Execution { result, mut events } = + self.execute(inbound, request, emit_event.clone()).await; let (result, finalization_failed) = match result { Ok(routed) => { let metadata = events .iter() .find_map(|event| match event { RoutingEvent::Mark(mark) => Some(mark.metadata.clone()), - RoutingEvent::Metric(_) => None, + RoutingEvent::Metric(_) | RoutingEvent::Span(_) => None, }) .unwrap_or_else(|| Json::Object(Map::new())); let result = returned_events( @@ -224,7 +226,17 @@ impl SwitchyardRuntime { } } - async fn execute(&self, inbound: WireFormat, request: Request) -> Execution { + fn sanitized_emitter(&self, emit_event: RoutingEventEmitter) -> RoutingEventEmitter { + let redactor = Arc::clone(&self.redactor); + Arc::new(move |event| emit_event(sanitize_event(&redactor, event))) + } + + async fn execute( + &self, + inbound: WireFormat, + request: Request, + emit_event: RoutingEventEmitter, + ) -> Execution { let Some(route) = self.route(&request) else { return Execution { result: Err("Switchyard has no route for this request model".into()), @@ -254,7 +266,11 @@ impl SwitchyardRuntime { .unwrap_or_else(|poisoned| poisoned.into_inner()) .push(observation); }); - match route.execute(request, Some(observer)).await { + match capture_client_spans(route.execute(request, Some(observer)), move |span| { + emit_event(RoutingEvent::Span(span)); + }) + .await + { Ok(output) => { let outcome_fields = self.emit_observations( &mut events, @@ -496,6 +512,9 @@ pub(crate) fn emit_event( RoutingEvent::Metric(metric) => runtime .emit_metric(&metric.name, metric.measurements, Some(&metric.metadata)) .map_err(|error| ("routing metric", metric.name, error)), + RoutingEvent::Span(span) => span + .emit(runtime) + .map_err(|error| ("routing span", "libsy".into(), error)), }; if let Err((kind, name, error)) = result { let message = redactor.text(format!( @@ -521,6 +540,7 @@ fn sanitize_event(redactor: &ProviderKeyRedactor, mut event: RoutingEvent) -> Ro .map(|value| redactor.value(value)); } } + RoutingEvent::Span(span) => span.sanitize(redactor), } event } @@ -690,7 +710,7 @@ fn route_execution_error_events( fn event_metadata(events: &[RoutingEvent]) -> Option { events.iter().find_map(|event| match event { RoutingEvent::Mark(mark) => Some(mark.metadata.clone()), - RoutingEvent::Metric(_) => None, + RoutingEvent::Metric(_) | RoutingEvent::Span(_) => None, }) } @@ -1048,7 +1068,7 @@ mod tests { ); let execution = runtime - .execute_buffered(WireFormat::OpenAiResponses, decoded) + .execute_buffered(WireFormat::OpenAiResponses, decoded, Arc::new(|_| {})) .await; let decision = execution .events @@ -1057,7 +1077,7 @@ mod tests { RoutingEvent::Mark(mark) if mark.name == "switchyard.routing.decision" => { Some(&mark.data) } - RoutingEvent::Mark(_) | RoutingEvent::Metric(_) => None, + RoutingEvent::Mark(_) | RoutingEvent::Metric(_) | RoutingEvent::Span(_) => None, }) .expect("decision mark should be emitted"); assert_eq!(decision["selected_model"], "target/model"); @@ -1130,7 +1150,7 @@ mod tests { .expect("request should decode"); let execution = runtime - .execute_buffered(WireFormat::OpenAiChat, request) + .execute_buffered(WireFormat::OpenAiChat, request, Arc::new(|_| {})) .await; assert!(execution.result.is_err()); @@ -1141,7 +1161,7 @@ mod tests { RoutingEvent::Mark(mark) if mark.name == "switchyard.routing.error" => { Some(&mark.data) } - RoutingEvent::Mark(_) | RoutingEvent::Metric(_) => None, + RoutingEvent::Mark(_) | RoutingEvent::Metric(_) | RoutingEvent::Span(_) => None, }) .expect("error mark should be emitted"); assert!(error["outcome_id"].is_string()); @@ -1189,7 +1209,7 @@ mod tests { .expect("namespaced Responses request should decode"); let response = runtime - .execute_buffered(WireFormat::OpenAiResponses, request) + .execute_buffered(WireFormat::OpenAiResponses, request, Arc::new(|_| {})) .await .result .expect("buffered target call should succeed"); diff --git a/crates/switchyard-nemo-relay-plugin/src/span_capture.rs b/crates/switchyard-nemo-relay-plugin/src/span_capture.rs new file mode 100644 index 0000000000..3559f63fa0 --- /dev/null +++ b/crates/switchyard-nemo-relay-plugin/src/span_capture.rs @@ -0,0 +1,543 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::fmt; +use std::future::Future; +use std::pin::Pin; +use std::sync::{Arc, OnceLock}; +use std::task::{Context as TaskContext, Poll}; +use std::time::{Instant, SystemTime}; + +use nemo_relay_plugin::{Json, PluginRuntime, ScopeType}; +use serde_json::{Map, Value}; +use switchyard_runner::ProviderKeyRedactor; +use tracing::field::{Field, Visit}; +use tracing::metadata::LevelFilter; +use tracing::span::{Attributes, Id, Record}; +use tracing::subscriber::Interest; +use tracing::{Metadata, Subscriber}; +use tracing_subscriber::Layer; +use tracing_subscriber::layer::Context; +use tracing_subscriber::prelude::*; +use tracing_subscriber::registry::LookupSpan; + +const CLIENT_CALL: &str = "libsy.client_call"; +const UPSTREAM_ATTEMPT: &str = "libsy.upstream_attempt"; + +type SpanEmitter = Arc; + +// One shared dispatch avoids global callsite registration on every request. Each +// open span keeps its request emitter so a stream can close after this scope ends. +tokio::task_local! { + static SPAN_EMITTER: SpanEmitter; +} + +#[derive(Debug)] +pub(crate) struct CapturedSpan { + name: &'static str, + fields: Map, + started_at: SystemTime, + ended_at: SystemTime, + children: Vec, +} + +struct OpenSpan { + name: &'static str, + fields: Map, + started_at: SystemTime, + started: Instant, + children: Vec, + emit: SpanEmitter, +} + +struct CaptureLayer; + +impl Layer for CaptureLayer +where + S: Subscriber + for<'lookup> LookupSpan<'lookup>, +{ + fn register_callsite(&self, metadata: &'static Metadata<'static>) -> Interest { + if captures(metadata) { + Interest::always() + } else { + Interest::never() + } + } + + fn enabled(&self, metadata: &Metadata<'_>, _ctx: Context<'_, S>) -> bool { + captures(metadata) + } + + fn max_level_hint(&self) -> Option { + Some(LevelFilter::DEBUG) + } + + fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, ctx: Context<'_, S>) { + let Ok(emit) = SPAN_EMITTER.try_with(Arc::clone) else { + return; + }; + let name = attrs.metadata().name(); + let mut fields = Map::new(); + attrs.record(&mut FieldVisitor(&mut fields)); + if let Some(span) = ctx.span(id) { + span.extensions_mut().insert(OpenSpan { + name, + fields, + started_at: SystemTime::now(), + started: Instant::now(), + children: Vec::new(), + emit, + }); + } + } + + fn on_record(&self, id: &Id, values: &Record<'_>, ctx: Context<'_, S>) { + let Some(span) = ctx.span(id) else { + return; + }; + let mut extensions = span.extensions_mut(); + if let Some(open) = extensions.get_mut::() { + values.record(&mut FieldVisitor(&mut open.fields)); + } + } + + fn on_close(&self, id: Id, ctx: Context<'_, S>) { + let Some(span) = ctx.span(&id) else { + return; + }; + let parent_id = span.parent().map(|parent| parent.id().clone()); + let Some(open) = span.extensions_mut().remove::() else { + return; + }; + let OpenSpan { + name, + mut fields, + started_at, + started, + children, + emit, + } = open; + // A span that closes before its terminal bookkeeping ran was cancelled. + fields + .entry("outcome") + .or_insert_with(|| Value::String("cancelled".into())); + let captured = CapturedSpan { + name, + fields, + started_at, + ended_at: started_at + .checked_add(started.elapsed()) + .unwrap_or_else(SystemTime::now), + children, + }; + + if captured.name == UPSTREAM_ATTEMPT + && let Some(parent) = parent_id.and_then(|parent| ctx.span(&parent)) + && let Some(open) = parent.extensions_mut().get_mut::() + { + open.children.push(captured); + return; + } + emit(captured); + } +} + +impl CapturedSpan { + pub(crate) fn sanitize(&mut self, redactor: &ProviderKeyRedactor) { + if let Json::Object(fields) = redactor.value(Json::Object(std::mem::take(&mut self.fields))) + { + self.fields = fields; + } + for child in &mut self.children { + child.sanitize(redactor); + } + } + + pub(crate) fn emit(self, runtime: &PluginRuntime) -> Result<(), String> { + let metadata = Json::Object(self.fields); + let status = span_status(&metadata); + let mut scope = runtime.scope_at( + self.name, + relay_scope_type(self.name), + None, + Some(&metadata), + None, + self.started_at, + )?; + let mut child_result = Ok(()); + for child in self.children { + child_result = child_result.and(child.emit(runtime)); + } + let close_result = scope.close_at(None, status.as_ref(), self.ended_at); + child_result.and(close_result) + } +} + +/// Captures Switchyard client spans while polling one route execution. +/// +/// This capture dispatch temporarily replaces any subscriber configured inside the plugin cdylib. +pub(crate) async fn capture_client_spans( + future: impl Future, + emit: impl Fn(CapturedSpan) + Send + Sync + 'static, +) -> T { + SPAN_EMITTER + .scope(Arc::new(emit), CaptureFuture::new(future)) + .await +} + +struct CaptureFuture { + inner: Option>>, + dispatch: &'static tracing::Dispatch, +} + +impl CaptureFuture { + fn new(inner: F) -> Self { + Self { + inner: Some(Box::pin(inner)), + dispatch: capture_dispatch(), + } + } +} + +impl Future for CaptureFuture { + type Output = F::Output; + + fn poll(self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll { + let this = self.get_mut(); + let Some(inner) = this.inner.as_mut() else { + return Poll::Pending; + }; + tracing::dispatcher::with_default(this.dispatch, || inner.as_mut().poll(cx)) + } +} + +impl Drop for CaptureFuture { + fn drop(&mut self) { + let Some(inner) = self.inner.take() else { + return; + }; + // Releasing a child-held parent reference uses the current dispatch. + tracing::dispatcher::with_default(self.dispatch, || drop(inner)); + } +} + +fn capture_dispatch() -> &'static tracing::Dispatch { + static DISPATCH: OnceLock = OnceLock::new(); + DISPATCH + .get_or_init(|| tracing::Dispatch::new(tracing_subscriber::registry().with(CaptureLayer))) +} + +struct FieldVisitor<'a>(&'a mut Map); + +impl FieldVisitor<'_> { + fn insert(&mut self, field: &Field, value: impl FnOnce() -> Value) { + if is_safe_field(field.name()) { + self.0.insert(field.name().into(), value()); + } + } +} + +impl Visit for FieldVisitor<'_> { + fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) { + self.insert(field, || Value::String(format!("{value:?}"))); + } + + fn record_str(&mut self, field: &Field, value: &str) { + self.insert(field, || Value::String(value.into())); + } + + fn record_bool(&mut self, field: &Field, value: bool) { + self.insert(field, || Value::Bool(value)); + } + + fn record_i64(&mut self, field: &Field, value: i64) { + self.insert(field, || Value::from(value)); + } + + fn record_u64(&mut self, field: &Field, value: u64) { + self.insert(field, || Value::from(value)); + } + + fn record_f64(&mut self, field: &Field, value: f64) { + self.insert(field, || Value::from(value)); + } +} + +fn captures(metadata: &Metadata<'_>) -> bool { + metadata.target() == "libsy" && matches!(metadata.name(), CLIENT_CALL | UPSTREAM_ATTEMPT) +} + +fn relay_scope_type(name: &str) -> ScopeType { + if name == CLIENT_CALL { + ScopeType::Llm + } else { + ScopeType::Custom + } +} + +fn is_safe_field(name: &str) -> bool { + matches!( + name, + "algorithm" + | "switchyard.candidate" + | "switchyard.candidate_count" + | "switchyard.call_phase" + | "selected_model" + | "model" + | "wire_format" + | "attempt" + | "max_attempts" + | "retry" + | "outcome" + | "status_code" + | "will_retry" + | "retry_delay_ms" + | "otel.status_code" + | "error.type" + ) +} + +fn span_status(metadata: &Json) -> Option { + let outcome = metadata.get("outcome").and_then(Json::as_str); + let status = metadata.get("otel.status_code").and_then(Json::as_str); + match (status, outcome) { + (Some("ERROR"), _) | (_, Some("error" | "cancelled")) => { + Some(serde_json::json!({"otel.status_code": "ERROR"})) + } + (Some("OK"), _) | (_, Some("ok")) => Some(serde_json::json!({"otel.status_code": "OK"})), + _ => None, + } +} + +#[cfg(test)] +mod tests { + use std::fmt; + use std::sync::Mutex; + + use tracing::Instrument; + + use super::*; + + struct PanicOnDebug; + + impl fmt::Debug for PanicOnDebug { + fn fmt(&self, _formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + panic!("disallowed fields must not be formatted"); + } + } + + #[tokio::test] + async fn captures_only_allowlisted_client_and_attempt_fields() { + let captured = Arc::new(Mutex::new(Vec::new())); + let observed = Arc::clone(&captured); + let result = capture_client_spans( + async { + async {} + .instrument(tracing::info_span!( + target: "other", + CLIENT_CALL, + outcome = "ignored", + )) + .await; + async { + async {} + .instrument(tracing::debug_span!( + target: "libsy", + UPSTREAM_ATTEMPT, + model = "target/model", + attempt = 1_u64, + outcome = "error", + error = "provider body", + )) + .await; + tracing::Span::current().record("outcome", "ok"); + } + .instrument(tracing::info_span!( + target: "libsy", + CLIENT_CALL, + selected_model = "target/model", + switchyard.call_phase = "routing", + outcome = tracing::field::Empty, + error = ?PanicOnDebug, + )) + .await; + 7 + }, + move |span| { + observed + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .push(span); + }, + ) + .await; + + assert_eq!(result, 7); + let mut spans = captured + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let [call] = spans.as_mut_slice() else { + panic!("expected one client call: {spans:?}"); + }; + assert_eq!(call.name, CLIENT_CALL); + assert!(matches!(relay_scope_type(call.name), ScopeType::Llm)); + assert_eq!(call.fields["selected_model"], "target/model"); + assert_eq!(call.fields["switchyard.call_phase"], "routing"); + assert_eq!(call.fields["outcome"], "ok"); + assert!(!call.fields.contains_key("error")); + let [attempt] = call.children.as_slice() else { + panic!("expected one upstream attempt: {call:?}"); + }; + assert_eq!(attempt.name, UPSTREAM_ATTEMPT); + assert!(matches!(relay_scope_type(attempt.name), ScopeType::Custom)); + assert_eq!(attempt.fields["model"], "target/model"); + assert_eq!(attempt.fields["attempt"], 1); + assert_eq!(attempt.fields["outcome"], "error"); + assert!(!attempt.fields.contains_key("error")); + assert!(call.started_at <= attempt.started_at); + assert!(attempt.ended_at <= call.ended_at); + + call.sanitize(&ProviderKeyRedactor::new(&["target/model".into()])); + assert_eq!(call.fields["selected_model"], "[REDACTED]"); + assert_eq!(call.children[0].fields["model"], "[REDACTED]"); + } + + #[tokio::test] + async fn captures_spans_that_close_after_the_scoped_future_returns() { + let captured = Arc::new(Mutex::new(Vec::new())); + let observed = Arc::clone(&captured); + let call = capture_client_spans( + async { + let call = tracing::info_span!( + target: "libsy", + CLIENT_CALL, + selected_model = "target/model", + outcome = "cancelled", + ); + let attempt = tracing::debug_span!( + target: "libsy", + parent: &call, + UPSTREAM_ATTEMPT, + attempt = 1_u64, + outcome = "cancelled", + ); + drop(attempt); + call + }, + move |span| { + observed + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .push(span); + }, + ) + .await; + + assert!(captured.lock().unwrap().is_empty()); + drop(call); + + let spans = captured.lock().unwrap(); + let [call] = spans.as_slice() else { + panic!("expected one deferred client call: {spans:?}"); + }; + assert_eq!(call.fields["outcome"], "cancelled"); + assert_eq!(call.children.len(), 1); + assert_eq!(call.children[0].fields["outcome"], "cancelled"); + } + + #[tokio::test] + async fn concurrent_captures_keep_their_emitters_isolated() { + let first = Arc::new(Mutex::new(Vec::new())); + let first_observed = Arc::clone(&first); + let second = Arc::new(Mutex::new(Vec::new())); + let second_observed = Arc::clone(&second); + + tokio::join!( + capture_client_spans( + async { + tokio::task::yield_now().await; + async {} + .instrument(tracing::info_span!( + target: "libsy", + CLIENT_CALL, + selected_model = "first", + )) + .await; + }, + move |span| first_observed.lock().unwrap().push(span), + ), + capture_client_spans( + async { + tokio::task::yield_now().await; + async {} + .instrument(tracing::info_span!( + target: "libsy", + CLIENT_CALL, + selected_model = "second", + )) + .await; + }, + move |span| second_observed.lock().unwrap().push(span), + ), + ); + + let first = first.lock().unwrap(); + let second = second.lock().unwrap(); + let [first] = first.as_slice() else { + panic!("expected one first client call: {first:?}"); + }; + let [second] = second.as_slice() else { + panic!("expected one second client call: {second:?}"); + }; + assert_eq!(first.fields["selected_model"], "first"); + assert_eq!(second.fields["selected_model"], "second"); + } + + #[tokio::test] + async fn dropped_in_flight_spans_are_cancelled() { + let captured = Arc::new(Mutex::new(Vec::new())); + let observed = Arc::clone(&captured); + + let mut future = Box::pin(capture_client_spans( + async move { + let call = tracing::info_span!( + target: "libsy", + CLIENT_CALL, + outcome = tracing::field::Empty, + ); + async move { + let attempt = tracing::debug_span!( + target: "libsy", + UPSTREAM_ATTEMPT, + outcome = tracing::field::Empty, + ); + async move { + std::future::pending::<()>().await; + } + .instrument(attempt) + .await; + } + .instrument(call) + .await; + }, + move |span| observed.lock().unwrap().push(span), + )); + + assert!(futures_util::poll!(&mut future).is_pending()); + drop(future); + + let spans = captured.lock().unwrap(); + let [call] = spans.as_slice() else { + panic!("expected one cancelled client call: {spans:?}"); + }; + assert_eq!(call.fields["outcome"], "cancelled"); + assert_eq!( + span_status(&Json::Object(call.fields.clone())), + Some(serde_json::json!({"otel.status_code": "ERROR"})) + ); + let [attempt] = call.children.as_slice() else { + panic!("expected one cancelled upstream attempt: {call:?}"); + }; + assert_eq!(attempt.fields["outcome"], "cancelled"); + } +} diff --git a/docs/integrations/nemo_relay.md b/docs/integrations/nemo_relay.md index baa6bbec18..83bba4022a 100644 --- a/docs/integrations/nemo_relay.md +++ b/docs/integrations/nemo_relay.md @@ -305,17 +305,29 @@ can receive these records, but each output format presents them differently. | Output | What the current integration shows | | --- | --- | | ATIF | The request and response seen by the agent. Internal routing is not added as separate steps or included in `final_metrics`. Relay-managed ATIF files may retain the raw records under `extra.observed_events`. | -| OpenTelemetry traces, including OpenInference | Relay's LLM span plus eligible routing marks, depending on the projection. Marks are point-in-time records, not duration spans. | +| OpenTelemetry traces, including OpenInference | Relay's outer LLM span, nested `libsy.client_call` and `libsy.upstream_attempt` duration scopes, plus eligible routing marks. | | OpenTelemetry metrics | Switchyard request, routing-call, latency, token, and failure measurements. Combine the reported routing usage with Relay's answer cost to estimate the observed routed cost. | | OTLP logs | Non-metric Switchyard marks that meet the configured severity threshold. | -Switchyard creates `libsy.run`, `libsy.llm_call`, `libsy.client_call`, and -`libsy.upstream_attempt` spans internally. Together they cover the algorithm -run, model calls requested by the algorithm, candidate attempts including -fallbacks, and individual HTTP attempts. They are not one connected hierarchy -today, and the plugin does not send them to Relay. As a result, Relay traces do -not show those internal operations as nested duration spans or give an internal -stream its own cancellation lifecycle. +The plugin captures two existing Switchyard spans. Each routing or answer LLM +call becomes a `libsy.client_call` scope. Physical HTTP attempts and retries +become `libsy.upstream_attempt` child scopes. Buffered calls close with the response; +streaming calls close when the stream finishes, fails, or is dropped. The +projection contains allowlisted, payload-free routing metadata only and does not +copy prompts, responses, provider error bodies, or credentials into Relay. + +`switchyard.call_phase` records when Switchyard issued a client call: +`routing` while the algorithm was still deciding, or `completion` after a +target was selected. A routing-phase response can later become the final +answer, so this is intentionally different from the `call_role` on +`switchyard.routing.llm_call` marks, which records the post-outcome +classification as `routing` or `answer`. + +| Switchyard signal | Relay record | Parent | Lifetime | +| --- | --- | --- | --- | +| Routing or answer model call | `libsy.client_call` scope | Relay's managed request scope | Starts with the Switchyard call and closes with the buffered response or terminal stream event. | +| Physical provider request or retry | `libsy.upstream_attempt` scope | Its `libsy.client_call` scope | Starts before the HTTP attempt and closes on response, error, or cancellation. | +| Routing request, decision, overhead, or failure | Existing `switchyard.routing.*` mark | Active Relay scope at emission | Point-in-time mark using the existing outcome fields. | ### How Routing Appears in Traces @@ -334,12 +346,16 @@ omits marks. With `mark_projection = "tool"`, the trace has this shape: flowchart LR agent["Relay agent scope"] llm["LLM call"] + call["libsy.client_call
routing or completion"] + attempt["libsy.upstream_attempt
one per HTTP attempt"] requested["mark:switchyard.routing.requested
zero duration"] calls["mark:switchyard.routing.llm_call
zero duration"] overhead["mark:switchyard.routing.overhead
zero duration"] decision["mark:switchyard.routing.decision
zero duration"] error["mark:switchyard.routing.error
zero duration"] agent --> llm + llm --> call + call --> attempt agent --> requested agent --> calls agent --> overhead From eea4bc8563c86922b8d79fd9a3077daf49fc74e3 Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Mon, 12 Oct 2026 13:21:55 -0700 Subject: [PATCH 2/3] refactor(relay): simplify routing span capture Signed-off-by: Alex Fournier --- crates/libsy-llm-client/src/run.rs | 19 +- .../libsy-llm-client/tests/observability.rs | 24 +- .../switchyard-nemo-relay-plugin/src/lib.rs | 70 +++-- .../src/redaction_tests.rs | 23 +- .../src/runtime.rs | 50 +-- .../src/span_capture.rs | 297 ++++++------------ 6 files changed, 179 insertions(+), 304 deletions(-) diff --git a/crates/libsy-llm-client/src/run.rs b/crates/libsy-llm-client/src/run.rs index aaca962e63..8ce6e8ceca 100644 --- a/crates/libsy-llm-client/src/run.rs +++ b/crates/libsy-llm-client/src/run.rs @@ -299,14 +299,6 @@ impl CallPhase { Self::Completion => "completion", } } - - const fn buffers_stream(self) -> bool { - matches!(self, Self::Routing) - } - - const fn records_routed_request(self) -> bool { - matches!(self, Self::Completion) - } } /// Call one candidate model and record its observation and span. @@ -362,6 +354,7 @@ async fn call_one( count: usize, phase: CallPhase, ) -> Result { + let is_routing = matches!(phase, CallPhase::Routing); let span = tracing::Span::current(); observability::record_gen_ai_request(&span, &request.llm_request); if let Some(session_id) = request @@ -377,9 +370,7 @@ async fn call_one( let started = Instant::now(); let result = async { let mut response = client?.call(request).await?; - if phase.buffers_stream() - && let LlmResponse::Stream(chunks) = response.llm_response - { + if is_routing && let LlmResponse::Stream(chunks) = response.llm_response { response.llm_response = LlmResponse::Stream(buffer_routing_stream(chunks).await?); } Ok(response) @@ -393,10 +384,10 @@ async fn call_one( response }); let result = observability::observe_client_call(result); - let result = if phase.records_routed_request() { - metrics::observe_routed_request(algorithm, model_id, Some(duration), result) - } else { + let result = if is_routing { result + } else { + metrics::observe_routed_request(algorithm, model_id, Some(duration), result) }; observe(ModelCallObservation { selected_model: model_id.clone(), diff --git a/crates/libsy-llm-client/tests/observability.rs b/crates/libsy-llm-client/tests/observability.rs index 8666264b7a..46ff1b13df 100644 --- a/crates/libsy-llm-client/tests/observability.rs +++ b/crates/libsy-llm-client/tests/observability.rs @@ -1391,9 +1391,6 @@ async fn observed_run_reports_one_successful_routed_call() -> switchyard_libsy:: async fn answered_routing_call_keeps_its_execution_phase() -> switchyard_libsy::Result<()> { let _guard = serialize_test().lock().await; let (store, _, _, _, _) = telemetry(); - let observations = Arc::new(Mutex::new(Vec::new())); - let observed = Arc::clone(&observations); - let observer: RunObserver = Arc::new(move |observation| observed.lock().push(observation)); const ALGO: &str = "answered-routing-call"; const MODEL: &str = "answer-from-routing"; @@ -1407,7 +1404,7 @@ async fn answered_routing_call_keeps_its_execution_phase() -> switchyard_libsy:: })), request_with_metadata("answered-routing-session", "answered-routing-correlation"), Arc::new(RuntimeModels::default()), - Some(observer), + None, ) .await?; @@ -1416,13 +1413,6 @@ async fn answered_routing_call_keeps_its_execution_phase() -> switchyard_libsy:: span.fields.get("switchyard.call_phase").map(String::as_str), Some("routing") ); - assert!( - observations - .lock() - .iter() - .any(|observation| matches!(observation, RunObservation::AnswerCall(_))), - "the post-outcome observation should still classify the call as the answer" - ); Ok(()) } @@ -1894,13 +1884,11 @@ async fn classifier_metrics_count_routing_and_answer_calls_once() -> switchyard_ assert_eq!(selected_model, "weak"); let spans = store.spans(); - for (model, phase) in [("classifier", "routing"), ("weak", "completion")] { - let span = find_span(&spans, "libsy.client_call", "selected_model", model); - assert_eq!( - span.fields.get("switchyard.call_phase").map(String::as_str), - Some(phase) - ); - } + let span = find_span(&spans, "libsy.client_call", "selected_model", "weak"); + assert_eq!( + span.fields.get("switchyard.call_phase").map(String::as_str), + Some("completion") + ); let snapshots = flushed_metrics(exporter, provider); assert_eq!( diff --git a/crates/switchyard-nemo-relay-plugin/src/lib.rs b/crates/switchyard-nemo-relay-plugin/src/lib.rs index 15ffc430b7..6f0c543119 100644 --- a/crates/switchyard-nemo-relay-plugin/src/lib.rs +++ b/crates/switchyard-nemo-relay-plugin/src/lib.rs @@ -8,13 +8,15 @@ mod translation; use std::sync::Arc; -use nemo_relay_plugin::{ - ConfigDiagnostic, DiagnosticLevel, Json, NativePlugin, PluginContext, PluginRuntime, -}; +use nemo_relay_plugin::{ConfigDiagnostic, DiagnosticLevel, Json, NativePlugin, PluginContext}; use serde_json::Map; use crate::config::{SwitchyardConfig, protocol_from_call}; -use crate::runtime::{SwitchyardRuntime, emit_event, emit_events}; +use crate::runtime::{ + RoutingEvent, RoutingEventEmitter, SwitchyardRuntime, emit_event as emit_routing_event, + emit_events, +}; +use crate::span_capture::SpanCapture; #[derive(Default)] struct SwitchyardPlugin; @@ -51,8 +53,24 @@ impl NativePlugin for SwitchyardPlugin { let priority = config.priority; let runtime = Arc::new(SwitchyardRuntime::new(config)?); let plugin_runtime = ctx.runtime(); - register_buffered(ctx, priority, Arc::clone(&runtime), plugin_runtime.clone())?; - register_stream(ctx, priority, runtime, plugin_runtime)?; + let event_runtime = Arc::clone(&runtime); + let emit_event: RoutingEventEmitter = Arc::new(move |event| { + emit_routing_event(&plugin_runtime, event_runtime.redactor(), event) + }); + let span_runtime = Arc::clone(&runtime); + let span_emit_event = Arc::clone(&emit_event); + let span_capture = SpanCapture::new(move |mut span| { + span.sanitize(span_runtime.redactor()); + span_emit_event(RoutingEvent::Span(span)); + }); + register_buffered( + ctx, + priority, + Arc::clone(&runtime), + Arc::clone(&emit_event), + span_capture.clone(), + )?; + register_stream(ctx, priority, runtime, emit_event, span_capture)?; Ok(()) } } @@ -61,14 +79,16 @@ fn register_buffered( ctx: &mut PluginContext<'_>, priority: i32, runtime: Arc, - plugin_runtime: PluginRuntime, + emit_event: RoutingEventEmitter, + span_capture: SpanCapture, ) -> Result<(), String> { ctx.register_llm_execution_intercept( "switchyard.runner.buffered", priority, move |name, request, _context, next| { let runtime = Arc::clone(&runtime); - let plugin_runtime = plugin_runtime.clone(); + let emit_event = Arc::clone(&emit_event); + let span_capture = span_capture.clone(); async move { let Some(inbound) = protocol_from_call(&name) else { return next.call(request).await; @@ -77,18 +97,10 @@ fn register_buffered( return next.call(request).await; } let decoded = runtime.decode_request(inbound, request, false)?; - let span_plugin_runtime = plugin_runtime.clone(); - let span_runtime = Arc::clone(&runtime); - let execution = runtime - .execute_buffered( - inbound, - decoded, - Arc::new(move |event| { - emit_event(&span_plugin_runtime, span_runtime.redactor(), event) - }), - ) + let execution = span_capture + .capture(runtime.execute_buffered(inbound, decoded)) .await; - emit_events(&plugin_runtime, runtime.redactor(), execution.events); + emit_events(&emit_event, execution.events); execution.result } }, @@ -99,14 +111,16 @@ fn register_stream( ctx: &mut PluginContext<'_>, priority: i32, runtime: Arc, - plugin_runtime: PluginRuntime, + emit_event: RoutingEventEmitter, + span_capture: SpanCapture, ) -> Result<(), String> { ctx.register_llm_stream_execution_intercept( "switchyard.runner.streaming", priority, move |name, request, _context, next| { let runtime = Arc::clone(&runtime); - let plugin_runtime = plugin_runtime.clone(); + let emit_event = Arc::clone(&emit_event); + let span_capture = span_capture.clone(); async move { let Some(inbound) = protocol_from_call(&name) else { return next.call(request).await; @@ -115,18 +129,10 @@ fn register_stream( return next.call(request).await; } let decoded = runtime.decode_request(inbound, request, true)?; - let stream_plugin_runtime = plugin_runtime.clone(); - let stream_runtime = Arc::clone(&runtime); - let execution = runtime - .execute_stream( - inbound, - decoded, - Arc::new(move |event| { - emit_event(&stream_plugin_runtime, stream_runtime.redactor(), event) - }), - ) + let execution = span_capture + .capture(runtime.execute_stream(inbound, decoded, Arc::clone(&emit_event))) .await; - emit_events(&plugin_runtime, runtime.redactor(), execution.events); + emit_events(&emit_event, execution.events); execution.result } }, diff --git a/crates/switchyard-nemo-relay-plugin/src/redaction_tests.rs b/crates/switchyard-nemo-relay-plugin/src/redaction_tests.rs index e30b4ad39a..d3a42c72e0 100644 --- a/crates/switchyard-nemo-relay-plugin/src/redaction_tests.rs +++ b/crates/switchyard-nemo-relay-plugin/src/redaction_tests.rs @@ -8,6 +8,8 @@ use std::process::Command; use wiremock::matchers::{header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; +use crate::span_capture::SpanCapture; + use super::*; const KEY_ENV: &str = "SWITCHYARD_RELAY_REDACTION_TEST_KEY"; @@ -122,16 +124,22 @@ async fn reflected_response(format: WireFormat, streaming: bool, fail: bool) { }), }; let captured = Arc::new(Mutex::new(Vec::new())); + let span_observed = Arc::clone(&captured); + let redactor = Arc::clone(&runtime.redactor); + let span_capture = SpanCapture::new(move |mut span| { + span.sanitize(&redactor); + span_observed.lock().unwrap().push(RoutingEvent::Span(span)); + }); let output = if streaming { let observed = Arc::clone(&captured); - let execution = runtime - .execute_stream( + let execution = span_capture + .capture(runtime.execute_stream( format, request, Arc::new(move |event| { observed.lock().unwrap().push(event); }), - ) + )) .await; assert!( captured.lock().unwrap().is_empty(), @@ -142,13 +150,8 @@ async fn reflected_response(format: WireFormat, streaming: bool, fail: bool) { assert_eq!(events.iter().any(Result::is_err), fail); format!("{events:?}") } else { - let observed = Arc::clone(&captured); - let execution = runtime - .execute_buffered( - format, - request, - Arc::new(move |event| observed.lock().unwrap().push(event)), - ) + let execution = span_capture + .capture(runtime.execute_buffered(format, request)) .await; captured.lock().unwrap().extend(execution.events); assert_eq!(execution.result.is_err(), fail); diff --git a/crates/switchyard-nemo-relay-plugin/src/runtime.rs b/crates/switchyard-nemo-relay-plugin/src/runtime.rs index 74cd0a606f..0870348323 100644 --- a/crates/switchyard-nemo-relay-plugin/src/runtime.rs +++ b/crates/switchyard-nemo-relay-plugin/src/runtime.rs @@ -22,7 +22,7 @@ use switchyard_runner::{ use switchyard_translation::{TranslationEngine, encode_stream_with_extensions}; use crate::config::SwitchyardConfig; -use crate::span_capture::{CapturedSpan, capture_client_spans}; +use crate::span_capture::CapturedSpan; use crate::translation; #[cfg(test)] @@ -141,11 +141,9 @@ impl SwitchyardRuntime { &self, inbound: WireFormat, request: Request, - emit_event: RoutingEventEmitter, ) -> Execution { let request_extensions = request.llm_request.extensions.clone(); - let emit_event = self.sanitized_emitter(emit_event); - let Execution { result, mut events } = self.execute(inbound, request, emit_event).await; + let Execution { result, mut events } = self.execute(inbound, request).await; let (result, finalization_failed) = match result { Ok(routed) => { let result = finalize_buffered_response( @@ -176,17 +174,10 @@ impl SwitchyardRuntime { ) -> Execution { let emit_event = self.sanitized_emitter(emit_event); let request_extensions = request.llm_request.extensions.clone(); - let Execution { result, mut events } = - self.execute(inbound, request, emit_event.clone()).await; + let Execution { result, mut events } = self.execute(inbound, request).await; let (result, finalization_failed) = match result { Ok(routed) => { - let metadata = events - .iter() - .find_map(|event| match event { - RoutingEvent::Mark(mark) => Some(mark.metadata.clone()), - RoutingEvent::Metric(_) | RoutingEvent::Span(_) => None, - }) - .unwrap_or_else(|| Json::Object(Map::new())); + let metadata = event_metadata(&events).unwrap_or_else(|| Json::Object(Map::new())); let result = returned_events( routed.response, inbound, @@ -231,12 +222,7 @@ impl SwitchyardRuntime { Arc::new(move |event| emit_event(sanitize_event(&redactor, event))) } - async fn execute( - &self, - inbound: WireFormat, - request: Request, - emit_event: RoutingEventEmitter, - ) -> Execution { + async fn execute(&self, inbound: WireFormat, request: Request) -> Execution { let Some(route) = self.route(&request) else { return Execution { result: Err("Switchyard has no route for this request model".into()), @@ -266,11 +252,7 @@ impl SwitchyardRuntime { .unwrap_or_else(|poisoned| poisoned.into_inner()) .push(observation); }); - match capture_client_spans(route.execute(request, Some(observer)), move |span| { - emit_event(RoutingEvent::Span(span)); - }) - .await - { + match route.execute(request, Some(observer)).await { Ok(output) => { let outcome_fields = self.emit_observations( &mut events, @@ -481,13 +463,9 @@ fn evidence_for_mark(evidence: Option) -> Option { (!evidence.is_empty()).then_some(Json::Object(evidence)) } -pub(crate) fn emit_events( - runtime: &PluginRuntime, - redactor: &ProviderKeyRedactor, - events: Vec, -) { +pub(crate) fn emit_events(emit_event: &RoutingEventEmitter, events: Vec) { for event in events { - emit_event(runtime, redactor, event); + emit_event(event); } } @@ -710,7 +688,7 @@ fn route_execution_error_events( fn event_metadata(events: &[RoutingEvent]) -> Option { events.iter().find_map(|event| match event { RoutingEvent::Mark(mark) => Some(mark.metadata.clone()), - RoutingEvent::Metric(_) | RoutingEvent::Span(_) => None, + _ => None, }) } @@ -1068,7 +1046,7 @@ mod tests { ); let execution = runtime - .execute_buffered(WireFormat::OpenAiResponses, decoded, Arc::new(|_| {})) + .execute_buffered(WireFormat::OpenAiResponses, decoded) .await; let decision = execution .events @@ -1077,7 +1055,7 @@ mod tests { RoutingEvent::Mark(mark) if mark.name == "switchyard.routing.decision" => { Some(&mark.data) } - RoutingEvent::Mark(_) | RoutingEvent::Metric(_) | RoutingEvent::Span(_) => None, + _ => None, }) .expect("decision mark should be emitted"); assert_eq!(decision["selected_model"], "target/model"); @@ -1150,7 +1128,7 @@ mod tests { .expect("request should decode"); let execution = runtime - .execute_buffered(WireFormat::OpenAiChat, request, Arc::new(|_| {})) + .execute_buffered(WireFormat::OpenAiChat, request) .await; assert!(execution.result.is_err()); @@ -1161,7 +1139,7 @@ mod tests { RoutingEvent::Mark(mark) if mark.name == "switchyard.routing.error" => { Some(&mark.data) } - RoutingEvent::Mark(_) | RoutingEvent::Metric(_) | RoutingEvent::Span(_) => None, + _ => None, }) .expect("error mark should be emitted"); assert!(error["outcome_id"].is_string()); @@ -1209,7 +1187,7 @@ mod tests { .expect("namespaced Responses request should decode"); let response = runtime - .execute_buffered(WireFormat::OpenAiResponses, request, Arc::new(|_| {})) + .execute_buffered(WireFormat::OpenAiResponses, request) .await .result .expect("buffered target call should succeed"); diff --git a/crates/switchyard-nemo-relay-plugin/src/span_capture.rs b/crates/switchyard-nemo-relay-plugin/src/span_capture.rs index 3559f63fa0..ec4fb9900d 100644 --- a/crates/switchyard-nemo-relay-plugin/src/span_capture.rs +++ b/crates/switchyard-nemo-relay-plugin/src/span_capture.rs @@ -4,7 +4,6 @@ use std::fmt; use std::future::Future; use std::pin::Pin; -use std::sync::{Arc, OnceLock}; use std::task::{Context as TaskContext, Poll}; use std::time::{Instant, SystemTime}; @@ -24,13 +23,7 @@ use tracing_subscriber::registry::LookupSpan; const CLIENT_CALL: &str = "libsy.client_call"; const UPSTREAM_ATTEMPT: &str = "libsy.upstream_attempt"; -type SpanEmitter = Arc; - -// One shared dispatch avoids global callsite registration on every request. Each -// open span keeps its request emitter so a stream can close after this scope ends. -tokio::task_local! { - static SPAN_EMITTER: SpanEmitter; -} +type SpanEmitter = Box; #[derive(Debug)] pub(crate) struct CapturedSpan { @@ -47,10 +40,11 @@ struct OpenSpan { started_at: SystemTime, started: Instant, children: Vec, - emit: SpanEmitter, } -struct CaptureLayer; +struct CaptureLayer { + emit: SpanEmitter, +} impl Layer for CaptureLayer where @@ -73,9 +67,6 @@ where } fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, ctx: Context<'_, S>) { - let Ok(emit) = SPAN_EMITTER.try_with(Arc::clone) else { - return; - }; let name = attrs.metadata().name(); let mut fields = Map::new(); attrs.record(&mut FieldVisitor(&mut fields)); @@ -86,7 +77,6 @@ where started_at: SystemTime::now(), started: Instant::now(), children: Vec::new(), - emit, }); } } @@ -115,7 +105,6 @@ where started_at, started, children, - emit, } = open; // A span that closes before its terminal bookkeeping ran was cancelled. fields @@ -138,7 +127,7 @@ where open.children.push(captured); return; } - emit(captured); + (self.emit)(captured); } } @@ -173,32 +162,39 @@ impl CapturedSpan { } } -/// Captures Switchyard client spans while polling one route execution. -/// -/// This capture dispatch temporarily replaces any subscriber configured inside the plugin cdylib. -pub(crate) async fn capture_client_spans( - future: impl Future, - emit: impl Fn(CapturedSpan) + Send + Sync + 'static, -) -> T { - SPAN_EMITTER - .scope(Arc::new(emit), CaptureFuture::new(future)) - .await +/// Captures Switchyard client spans for one plugin registration. +#[derive(Clone)] +pub(crate) struct SpanCapture { + dispatch: tracing::Dispatch, } -struct CaptureFuture { - inner: Option>>, - dispatch: &'static tracing::Dispatch, -} - -impl CaptureFuture { - fn new(inner: F) -> Self { +impl SpanCapture { + pub(crate) fn new(emit: impl Fn(CapturedSpan) + Send + Sync + 'static) -> Self { + let layer = CaptureLayer { + emit: Box::new(emit), + }; Self { - inner: Some(Box::pin(inner)), - dispatch: capture_dispatch(), + dispatch: tracing::Dispatch::new(tracing_subscriber::registry().with(layer)), + } + } + + /// Polls a future with the capture subscriber as its default. + pub(crate) async fn capture(&self, future: impl Future) -> T { + CaptureFuture { + inner: Some(Box::pin(future)), + dispatch: self.dispatch.clone(), } + .await } } +// Standard subscriber wrappers restore the ambient dispatch before drop. This +// wrapper also drops the inner future under capture so cancelled spans close here. +struct CaptureFuture { + inner: Option>>, + dispatch: tracing::Dispatch, +} + impl Future for CaptureFuture { type Output = F::Output; @@ -207,7 +203,7 @@ impl Future for CaptureFuture { let Some(inner) = this.inner.as_mut() else { return Poll::Pending; }; - tracing::dispatcher::with_default(this.dispatch, || inner.as_mut().poll(cx)) + tracing::dispatcher::with_default(&this.dispatch, || inner.as_mut().poll(cx)) } } @@ -216,17 +212,10 @@ impl Drop for CaptureFuture { let Some(inner) = self.inner.take() else { return; }; - // Releasing a child-held parent reference uses the current dispatch. - tracing::dispatcher::with_default(self.dispatch, || drop(inner)); + tracing::dispatcher::with_default(&self.dispatch, || drop(inner)); } } -fn capture_dispatch() -> &'static tracing::Dispatch { - static DISPATCH: OnceLock = OnceLock::new(); - DISPATCH - .get_or_init(|| tracing::Dispatch::new(tracing_subscriber::registry().with(CaptureLayer))) -} - struct FieldVisitor<'a>(&'a mut Map); impl FieldVisitor<'_> { @@ -312,7 +301,7 @@ fn span_status(metadata: &Json) -> Option { #[cfg(test)] mod tests { use std::fmt; - use std::sync::Mutex; + use std::sync::mpsc; use tracing::Instrument; @@ -326,59 +315,44 @@ mod tests { } } + fn recorder() -> (SpanCapture, mpsc::Receiver) { + let (sender, receiver) = mpsc::channel(); + let capture = SpanCapture::new(move |span| sender.send(span).unwrap()); + (capture, receiver) + } + #[tokio::test] async fn captures_only_allowlisted_client_and_attempt_fields() { - let captured = Arc::new(Mutex::new(Vec::new())); - let observed = Arc::clone(&captured); - let result = capture_client_spans( - async { - async {} - .instrument(tracing::info_span!( - target: "other", - CLIENT_CALL, - outcome = "ignored", - )) - .await; - async { - async {} - .instrument(tracing::debug_span!( - target: "libsy", - UPSTREAM_ATTEMPT, - model = "target/model", - attempt = 1_u64, - outcome = "error", - error = "provider body", - )) - .await; - tracing::Span::current().record("outcome", "ok"); - } - .instrument(tracing::info_span!( + let (capture, captured) = recorder(); + capture + .capture(async { + drop(tracing::info_span!(target: "other", CLIENT_CALL)); + let call = tracing::info_span!( target: "libsy", CLIENT_CALL, selected_model = "target/model", switchyard.call_phase = "routing", outcome = tracing::field::Empty, error = ?PanicOnDebug, - )) - .await; - 7 - }, - move |span| { - observed - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .push(span); - }, - ) - .await; - - assert_eq!(result, 7); - let mut spans = captured - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let [call] = spans.as_mut_slice() else { - panic!("expected one client call: {spans:?}"); - }; + ); + { + let _entered = call.enter(); + drop(tracing::debug_span!( + target: "libsy", + UPSTREAM_ATTEMPT, + model = "target/model", + attempt = 1_u64, + outcome = "error", + error = "provider body", + )); + call.record("outcome", "ok"); + } + drop(call); + }) + .await; + + let mut call = captured.try_recv().expect("expected one client call"); + assert!(captured.try_recv().is_err()); assert_eq!(call.name, CLIENT_CALL); assert!(matches!(relay_scope_type(call.name), ScopeType::Llm)); assert_eq!(call.fields["selected_model"], "target/model"); @@ -404,132 +378,67 @@ mod tests { #[tokio::test] async fn captures_spans_that_close_after_the_scoped_future_returns() { - let captured = Arc::new(Mutex::new(Vec::new())); - let observed = Arc::clone(&captured); - let call = capture_client_spans( - async { - let call = tracing::info_span!( - target: "libsy", - CLIENT_CALL, - selected_model = "target/model", - outcome = "cancelled", - ); - let attempt = tracing::debug_span!( - target: "libsy", - parent: &call, - UPSTREAM_ATTEMPT, - attempt = 1_u64, - outcome = "cancelled", - ); - drop(attempt); + let (capture, captured) = recorder(); + let call = capture + .capture(async { + let call = tracing::info_span!(target: "libsy", CLIENT_CALL); + drop(tracing::debug_span!(target: "libsy", parent: &call, UPSTREAM_ATTEMPT)); call - }, - move |span| { - observed - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .push(span); - }, - ) - .await; - - assert!(captured.lock().unwrap().is_empty()); + }) + .await; + + assert!(captured.try_recv().is_err()); + drop(capture); drop(call); - let spans = captured.lock().unwrap(); - let [call] = spans.as_slice() else { - panic!("expected one deferred client call: {spans:?}"); - }; - assert_eq!(call.fields["outcome"], "cancelled"); + let call = captured + .try_recv() + .expect("expected one deferred client call"); assert_eq!(call.children.len(), 1); - assert_eq!(call.children[0].fields["outcome"], "cancelled"); } #[tokio::test] - async fn concurrent_captures_keep_their_emitters_isolated() { - let first = Arc::new(Mutex::new(Vec::new())); - let first_observed = Arc::clone(&first); - let second = Arc::new(Mutex::new(Vec::new())); - let second_observed = Arc::clone(&second); + async fn one_capture_handles_concurrent_futures() { + let (capture, captured) = recorder(); tokio::join!( - capture_client_spans( - async { - tokio::task::yield_now().await; - async {} - .instrument(tracing::info_span!( - target: "libsy", - CLIENT_CALL, - selected_model = "first", - )) - .await; - }, - move |span| first_observed.lock().unwrap().push(span), - ), - capture_client_spans( - async { - tokio::task::yield_now().await; - async {} - .instrument(tracing::info_span!( - target: "libsy", - CLIENT_CALL, - selected_model = "second", - )) - .await; - }, - move |span| second_observed.lock().unwrap().push(span), - ), + capture.capture(async { + tokio::task::yield_now().await; + drop(tracing::info_span!(target: "libsy", CLIENT_CALL, selected_model = "first")); + }), + capture.capture(async { + tokio::task::yield_now().await; + drop(tracing::info_span!(target: "libsy", CLIENT_CALL, selected_model = "second")); + }), ); - let first = first.lock().unwrap(); - let second = second.lock().unwrap(); - let [first] = first.as_slice() else { - panic!("expected one first client call: {first:?}"); - }; - let [second] = second.as_slice() else { - panic!("expected one second client call: {second:?}"); - }; - assert_eq!(first.fields["selected_model"], "first"); - assert_eq!(second.fields["selected_model"], "second"); + let mut models = captured + .try_iter() + .map(|span| span.fields["selected_model"].as_str().unwrap().to_owned()) + .collect::>(); + models.sort_unstable(); + assert_eq!(models, ["first", "second"]); } #[tokio::test] async fn dropped_in_flight_spans_are_cancelled() { - let captured = Arc::new(Mutex::new(Vec::new())); - let observed = Arc::clone(&captured); + let (capture, captured) = recorder(); - let mut future = Box::pin(capture_client_spans( - async move { - let call = tracing::info_span!( - target: "libsy", - CLIENT_CALL, - outcome = tracing::field::Empty, - ); - async move { - let attempt = tracing::debug_span!( - target: "libsy", - UPSTREAM_ATTEMPT, - outcome = tracing::field::Empty, - ); - async move { - std::future::pending::<()>().await; - } - .instrument(attempt) - .await; - } + let mut future = Box::pin(capture.capture(async move { + let call = tracing::info_span!(target: "libsy", CLIENT_CALL); + let attempt = tracing::debug_span!(target: "libsy", parent: &call, UPSTREAM_ATTEMPT); + std::future::pending::<()>() + .instrument(attempt) .instrument(call) .await; - }, - move |span| observed.lock().unwrap().push(span), - )); + })); assert!(futures_util::poll!(&mut future).is_pending()); drop(future); - let spans = captured.lock().unwrap(); - let [call] = spans.as_slice() else { - panic!("expected one cancelled client call: {spans:?}"); - }; + let call = captured + .try_recv() + .expect("expected one cancelled client call"); assert_eq!(call.fields["outcome"], "cancelled"); assert_eq!( span_status(&Json::Object(call.fields.clone())), From 9ae208bf3b09b73d96b4268ef8ebf35dea2e60a7 Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Mon, 12 Oct 2026 14:14:53 -0700 Subject: [PATCH 3/3] refactor(relay): use built-in span filter Signed-off-by: Alex Fournier --- .../src/span_capture.rs | 22 +++---------------- 1 file changed, 3 insertions(+), 19 deletions(-) diff --git a/crates/switchyard-nemo-relay-plugin/src/span_capture.rs b/crates/switchyard-nemo-relay-plugin/src/span_capture.rs index ec4fb9900d..10fc1c4734 100644 --- a/crates/switchyard-nemo-relay-plugin/src/span_capture.rs +++ b/crates/switchyard-nemo-relay-plugin/src/span_capture.rs @@ -11,11 +11,10 @@ use nemo_relay_plugin::{Json, PluginRuntime, ScopeType}; use serde_json::{Map, Value}; use switchyard_runner::ProviderKeyRedactor; use tracing::field::{Field, Visit}; -use tracing::metadata::LevelFilter; use tracing::span::{Attributes, Id, Record}; -use tracing::subscriber::Interest; use tracing::{Metadata, Subscriber}; use tracing_subscriber::Layer; +use tracing_subscriber::filter::{LevelFilter, filter_fn}; use tracing_subscriber::layer::Context; use tracing_subscriber::prelude::*; use tracing_subscriber::registry::LookupSpan; @@ -50,22 +49,6 @@ impl Layer for CaptureLayer where S: Subscriber + for<'lookup> LookupSpan<'lookup>, { - fn register_callsite(&self, metadata: &'static Metadata<'static>) -> Interest { - if captures(metadata) { - Interest::always() - } else { - Interest::never() - } - } - - fn enabled(&self, metadata: &Metadata<'_>, _ctx: Context<'_, S>) -> bool { - captures(metadata) - } - - fn max_level_hint(&self) -> Option { - Some(LevelFilter::DEBUG) - } - fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, ctx: Context<'_, S>) { let name = attrs.metadata().name(); let mut fields = Map::new(); @@ -172,7 +155,8 @@ impl SpanCapture { pub(crate) fn new(emit: impl Fn(CapturedSpan) + Send + Sync + 'static) -> Self { let layer = CaptureLayer { emit: Box::new(emit), - }; + } + .with_filter(filter_fn(captures).with_max_level_hint(LevelFilter::DEBUG)); Self { dispatch: tracing::Dispatch::new(tracing_subscriber::registry().with(layer)), }