From 092f29aa6763fbd60cb1aeea56df2b46e39e0a55 Mon Sep 17 00:00:00 2001 From: nachiketb Date: Mon, 5 Oct 2026 17:21:07 -0700 Subject: [PATCH] fix(libsy): retain model call spans through completion Signed-off-by: nachiketb --- .../libsy-llm-client/tests/observability.rs | 7 +++- crates/libsy/src/core/algorithm.rs | 33 ++++++++++--------- crates/libsy/src/observability.rs | 13 +++----- 3 files changed, 28 insertions(+), 25 deletions(-) diff --git a/crates/libsy-llm-client/tests/observability.rs b/crates/libsy-llm-client/tests/observability.rs index 5c3207a49..30366dbb2 100644 --- a/crates/libsy-llm-client/tests/observability.rs +++ b/crates/libsy-llm-client/tests/observability.rs @@ -2150,7 +2150,7 @@ async fn in_flight_gauge_reads_a_run_parked_on_an_unanswered_routing_call() #[tokio::test] async fn in_flight_gauge_clears_when_a_run_is_abandoned() -> switchyard_libsy::Result<()> { let _guard = serialize_test().lock().await; - let (_, exporter, provider, _, _) = telemetry(); + let (store, exporter, provider, _, _) = telemetry(); const ALGO: &str = "obs-abandoned-algo"; const MODEL: &str = "obs-abandoned-model"; let algorithm = Arc::new(RoutingCallAlgo { @@ -2195,5 +2195,10 @@ async fn in_flight_gauge_clears_when_a_run_is_abandoned() -> switchyard_libsy::R Some(0), "an abandoned run must not strand the gauge above zero" ); + let call_span = find_span(&store.spans(), "libsy.llm_call", "selected_model", MODEL); + assert_eq!( + call_span.fields.get("outcome").map(String::as_str), + Some("error") + ); Ok(()) } diff --git a/crates/libsy/src/core/algorithm.rs b/crates/libsy/src/core/algorithm.rs index 249805a01..133f8fbe4 100644 --- a/crates/libsy/src/core/algorithm.rs +++ b/crates/libsy/src/core/algorithm.rs @@ -116,6 +116,8 @@ pub struct CallModel { /// How to send the response back to the algorithm. `None` once the call is recorded. reply: Option>>, started: Instant, + // Retain the originating span even if the waiting algorithm is cancelled first. + span: tracing::Span, } impl CallModel { @@ -124,6 +126,9 @@ impl CallModel { /// can only be fulfilled once. pub fn respond(mut self, result: Result) -> Result<()> { self.record(result.is_ok()); + if let Ok(response) = &result { + observability::record_llm_response(response, &self.span); + } self.reply .take() .ok_or(DriverError::ResponseDropped)? @@ -149,6 +154,8 @@ impl CallModel { self.started.elapsed(), is_ok, ); + self.span + .record("outcome", if is_ok { "ok" } else { "error" }); } } @@ -331,9 +338,9 @@ impl Driver { /// Errors if the stream is closed or the call failed. /// The await is wrapped in a `libsy.llm_call` span measuring *fulfillment* as /// the algorithm observes it (host queueing/serving included; a streamed - /// response resolves when its stream handle arrives). The host records call metrics - /// through [`CallModel::respond`] or [`CallModel::fail`]; outcome and token usage - /// are recorded on the span when the promise resolves. The provider call itself is the + /// response resolves when its stream handle arrives). The host records call metrics, + /// outcome, and buffered token usage when it completes or drops the call, even if the + /// waiting algorithm has been cancelled. The provider call itself is the /// host's, and is instrumented by whoever makes it. pub async fn call_model(&self, request: Request, models: Vec) -> Result { self.call_model_with_error_recovery(request, models, false) @@ -375,19 +382,15 @@ impl Driver { recover_errors, reply: Some(reply), started, + span: tracing::Span::current(), }; - let result = async { - self.step_tx - .send(Ok(Step::CallModel(Box::new(call)))) - .await - .map_err(|_| DriverError::StreamClosed)?; - response - .await - .map_err(|_| LibsyError::from(DriverError::ResponseDropped))? - } - .await; - observability::record_llm_call_span(&result, &tracing::Span::current()); - result + self.step_tx + .send(Ok(Step::CallModel(Box::new(call)))) + .await + .map_err(|_| DriverError::StreamClosed)?; + response + .await + .map_err(|_| LibsyError::from(DriverError::ResponseDropped))? } /// Override the request's model with `model` and wait for the host to reply. diff --git a/crates/libsy/src/observability.rs b/crates/libsy/src/observability.rs index 4775256c3..6cacfd4ca 100644 --- a/crates/libsy/src/observability.rs +++ b/crates/libsy/src/observability.rs @@ -256,15 +256,10 @@ pub(crate) fn record_llm_call( .record(duration.as_secs_f64() * 1000.0, &call_attributes); } -/// Records the outcome and token fields on the algorithm's call span. -pub(crate) fn record_llm_call_span(result: &Result, span: &Span) { - span.record("outcome", outcome_value(result)); - if let Ok(response) = result { - // Token usage exists only once a response is buffered; a streamed - // response resolves before its usage is known, so none is recorded. - let Some(usage) = response.llm_response.as_agg().map(|agg| &agg.usage) else { - return; - }; +/// Records buffered token usage on the algorithm's call span. +pub(crate) fn record_llm_response(response: &Response, span: &Span) { + // A streamed response resolves before its usage is known. + if let Some(usage) = response.llm_response.as_agg().map(|agg| &agg.usage) { record_call_usage(usage, span); } }