From 3558431a186b356ee1e2ac60163e7700bcaf5999 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Thu, 5 Nov 2026 20:52:57 -0800 Subject: [PATCH 1/4] fix(observability): observe routing streams before buffering Signed-off-by: Bryan Bednarski --- crates/libsy-llm-client/src/run.rs | 18 +- .../libsy-llm-client/tests/observability.rs | 168 +++++++++++++++ .../scripts/test_buffered_routing.py | 192 ++++++++++++++++++ 3 files changed, 373 insertions(+), 5 deletions(-) create mode 100644 crates/switchyard-nemo-relay-plugin/scripts/test_buffered_routing.py diff --git a/crates/libsy-llm-client/src/run.rs b/crates/libsy-llm-client/src/run.rs index 5646344aa1..b0609227cf 100644 --- a/crates/libsy-llm-client/src/run.rs +++ b/crates/libsy-llm-client/src/run.rs @@ -456,23 +456,31 @@ async fn call_one( let started = Instant::now(); let mut routing_usage = None; let result = async { - let mut response = client?.call(request).await?; + let response = async { client?.call(request).await } + .await + .map_err(|source| LibsyError::client_call(model_id.clone(), source))?; + // Observe the upstream stream; replay must not retain or re-observe the call. + let mut response = observability::observe_client_call(Ok(response))?; if is_routing && let LlmResponse::Stream(chunks) = response.llm_response { - let (chunks, usage) = buffer_routing_stream(chunks).await?; + let (chunks, usage) = buffer_routing_stream(chunks) + .await + .map_err(|source| LibsyError::client_call(model_id.clone(), source))?; routing_usage = usage; response.llm_response = LlmResponse::Stream(chunks); } Ok(response) } - .await - .map_err(|source| LibsyError::client_call(model_id.clone(), source)); + .await; let duration = started.elapsed(); let result = result.map(|mut response| { response.set_served_model(model_id); response }); - let result = observability::observe_client_call(result); + let result = match result { + Ok(response) => Ok(response), + Err(error) => observability::observe_client_call(Err(error)), + }; let result = if is_routing { result } else { diff --git a/crates/libsy-llm-client/tests/observability.rs b/crates/libsy-llm-client/tests/observability.rs index 1177a25325..359c69876f 100644 --- a/crates/libsy-llm-client/tests/observability.rs +++ b/crates/libsy-llm-client/tests/observability.rs @@ -1457,6 +1457,174 @@ impl RoutedLlmClient for StreamingUsageClient { } } +#[tokio::test] +async fn buffered_routing_closes_observation_before_replay() -> switchyard_libsy::Result<()> { + use tracing::Instrument; + + let _guard = serialize_test().lock().await; + let (_, _, _, exporter, _) = telemetry(); + for drain in [false, true] { + let model = if drain { + "routing-replay-drained" + } else { + "routing-replay-dropped" + }; + let workflow = format!("workflow-{model}"); + let mut request = request_with_metadata(model, model); + request.llm_request.stream = true; + let mut upstream_request = request.clone(); + upstream_request.llm_request.model = Some(model.into()); + let original = StreamingUsageClient + .call(upstream_request) + .await + .expect("valid upstream fixture"); + let LlmResponse::Stream(original) = original.llm_response else { + return Err(test_error("expected upstream stream")); + }; + let expected = original + .map(|event| event.expect("valid fixture")) + .collect::>() + .await; + let (_, response) = run( + Arc::new(RoutingCallAlgo { + name: workflow.clone(), + target: model.into(), + }), + Arc::new(StreamingUsageClient), + request, + ) + .instrument(tracing::info_span!("buffered-workflow", workflow = %workflow)) + .await?; + + // The original I/O has finished even though nobody has polled the replay. + let before = exporter.get_finished_spans().expect("exported spans"); + let call = before + .iter() + .find(|span| span.name == format!("chat {model}")) + .expect("client span must close before replay starts"); + assert_eq!( + otel_attribute(call, "gen_ai.usage.input_tokens"), + Some(&OtelValue::I64(21)) + ); + assert_eq!( + otel_attribute(call, "gen_ai.usage.output_tokens"), + Some(&OtelValue::I64(5)) + ); + let mut ancestor = call; + while ancestor.name != "buffered-workflow" { + ancestor = before + .iter() + .find(|span| span.span_context.span_id() == ancestor.parent_span_id) + .expect("every parent up to the workflow must already be exported"); + } + assert_eq!( + otel_attribute(ancestor, "workflow"), + Some(&OtelValue::String(workflow.into())) + ); + + let LlmResponse::Stream(mut replay) = response.llm_response else { + return Err(test_error("expected buffered stream replay")); + }; + if drain { + // Provider events and their trailing usage survive buffering unchanged. + let mut events = Vec::new(); + while let Some(event) = replay.next().await { + events.push(event.expect("valid replay")); + } + assert_eq!( + events, expected, + "normalized chunks and preserved provider data must survive replay" + ); + } + drop(replay); + assert_eq!( + exporter.get_finished_spans().expect("exported spans").len(), + before.len(), + "replaying or dropping buffered data must not export another observation" + ); + } + Ok(()) +} + +struct InterruptedRoutingStream { + stall: bool, +} + +#[async_trait] +impl RoutedLlmClient for InterruptedRoutingStream { + async fn call(&self, _request: Request) -> Result { + let usage = LlmResponseChunk::Usage(Usage { + input_tokens: Some(11), + output_tokens: Some(3), + ..Usage::default() + }); + let prefix = futures::stream::iter([Ok(usage.into())]); + let tail: switchyard_protocol::LlmResponseStream = if self.stall { + futures::stream::pending().boxed() + } else { + futures::stream::iter([Err(LlmClientError::ResponseTranslation( + "broken stream".into(), + ))]) + .boxed() + }; + Ok(Response { + llm_response: LlmResponse::Stream(prefix.chain(tail).boxed()), + metadata: None, + upstream_headers: http::HeaderMap::new(), + }) + } +} + +#[tokio::test] +async fn interrupted_routing_buffer_retains_observed_usage() -> switchyard_libsy::Result<()> { + let _guard = serialize_test().lock().await; + let (store, _, _, exporter, _) = telemetry(); + for stall in [false, true] { + let model = if stall { + "routing-buffer-cancelled" + } else { + "routing-buffer-failed" + }; + let mut request = request_with_metadata(model, model); + request.llm_request.stream = true; + let mut execution = Box::pin(run( + Arc::new(RoutingCallAlgo { + name: model.into(), + target: model.into(), + }), + Arc::new(InterruptedRoutingStream { stall }), + request, + )); + if stall { + assert!( + tokio::time::timeout(Duration::from_millis(50), execution.as_mut()) + .await + .is_err() + ); + drop(execution); + } else { + assert!(execution.await.is_err()); + } + let span = find_span(&store.spans(), "libsy.client_call", "selected_model", model); + for (field, value) in [ + ("gen_ai.usage.input_tokens", "11"), + ("gen_ai.usage.output_tokens", "3"), + ("outcome", if stall { "cancelled" } else { "error" }), + ] { + assert_eq!(span.fields.get(field).map(String::as_str), Some(value)); + } + let exported = exporter.get_finished_spans().expect("exported spans"); + assert_eq!( + exported + .iter() + .filter(|span| span.name == format!("chat {model}")) + .count(), + 1 + ); + } + Ok(()) +} + struct TimeoutClient; #[async_trait] diff --git a/crates/switchyard-nemo-relay-plugin/scripts/test_buffered_routing.py b/crates/switchyard-nemo-relay-plugin/scripts/test_buffered_routing.py new file mode 100644 index 0000000000..6b1f74586e --- /dev/null +++ b/crates/switchyard-nemo-relay-plugin/scripts/test_buffered_routing.py @@ -0,0 +1,192 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Check buffered routing attribution with a packaged plugin and Relay 0.10. + +Run with an isolated Relay configuration and a loopback-only network policy. +The synthetic judge streams even though the classifier requests a verdict. +""" + +import argparse +import asyncio +import json +import math +import threading +from collections import Counter +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path + +import nemo_relay as relay +from nemo_relay import codecs, model_pricing, plugin, subscribers + + +class Provider(BaseHTTPRequestHandler): + calls = Counter() + lock = threading.Lock() + + def do_POST(self): + request = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) + model = request["model"] + with self.lock: + self.calls[model] += 1 + judge = model == "test/judge" + content = '{"target":"efficient"}' if judge else "answer" + usage = {"prompt_tokens": 11 if judge else 7, "completion_tokens": 3 if judge else 2} + if judge: + assert request["stream"] is True + streaming = request.get("stream", False) + if streaming: + frames = [ + {"id": "judge", "model": model, "choices": [ + {"index": 0, "delta": {"role": "assistant", "content": content}}, + ]}, + {"id": "judge", "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, + {"id": "judge", "choices": [], "usage": usage}, + ] + body = ("".join("data: " + json.dumps(frame) + "\n\n" for frame in frames) + + "data: [DONE]\n\n").encode() + else: + body = json.dumps({"id": "answer", "model": model, "usage": usage, "choices": [ + {"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop"}, + ]}).encode() + self.send_response(200) + self.send_header("Content-Type", "text/event-stream" if streaming else "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, *_args): + pass + + +def check_workflow(events, label): + ends = {e["uuid"]: e for e in events if e.get("kind") == "scope" and e.get("scope_category") == "end"} + workflow = next(e for e in ends.values() if e.get("name") == label) + outer, = [e for e in ends.values() if e.get("name") == "openai.chat_completions" + and e.get("parent_uuid") == workflow["uuid"]] + calls = [e for e in ends.values() if e.get("name") == "libsy.client_call" + and e.get("parent_uuid") == outer["uuid"]] + assert len(calls) == 2, (label, calls) + assert {e["metadata"]["selected_model"] for e in calls} == {"test/judge", "test/answer"} + total = 0.0 + for call in calls: + # Exactly one start/end pair: replay must not create another observation. + pair = [e for e in events if e.get("uuid") == call["uuid"]] + assert Counter(e["scope_category"] for e in pair) == {"start": 1, "end": 1} + assert all(e["parent_uuid"] == outer["uuid"] for e in pair) + assert call["metadata"]["outcome"] == "ok" + judge = call["metadata"]["selected_model"] == "test/judge" + usage = codecs.AnthropicMessagesCodec().decode_response(call["data"]).usage + assert usage["prompt_tokens"] == (11 if judge else 7), usage + assert usage["completion_tokens"] == (3 if judge else 2), usage + assert usage["cost"]["currency"] == "USD", usage + total += usage["cost"]["total"] + assert math.isclose(total, 0.028), total + return {"workflow_uuid": workflow["uuid"], "calls": len(calls), "logical_call_cost_usd": total} + + +async def run(manifest, output): + server = ThreadingHTTPServer(("127.0.0.1", 0), Provider) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + events = [] + subscribers.register("buffered-routing-test", lambda event: events.append(event.to_dict())) + activation = None + try: + deployment = output / "switchyard.toml" + deployment.write_text(f'''schema_version = 1 +[llm_clients.local] +format = "openai_chat" +base_url = "http://127.0.0.1:{server.server_port}/v1" +max_retries = 0 +[targets.judge] +id = "test/judge" +llm_client = "local" +extra_body = {{stream = true}} +[targets.answer] +id = "test/answer" +llm_client = "local" +[routes.test] +id = "switchyard/test" +type = "llm_classifier" +mode = "custom" +models = {{judge = ["judge"], efficient = ["answer"], any = ["answer"]}} +default_target = "efficient" +prompt = "Select efficient." +response_schema = '{{"type":"object","properties":{{"target":{{"type":"string","enum":["efficient"]}}}},"required":["target"],"additionalProperties":false}}' +policy = {{type = "target_selector", selector = "/target"}} +''') + config = output / "plugins.toml" + config.write_text(f'''[[plugins.dynamic]] +manifest = {json.dumps(str(manifest))} +[plugins.dynamic.config] +switchyard_config_path = {json.dumps(str(deployment))} +[plugins.policy.overrides."nvidia.switchyard"] +attestation = "integrity_only" +''') + pricing = model_pricing.ComponentSpec(model_pricing.PricingConfig(sources=[ + model_pricing.InlineSource(model_pricing.PricingCatalog(entries=[ + model_pricing.ModelPricing( + provider="anthropic", model_id=model, pricing_as_of="2026-11-05", + pricing_source="synthetic", rates=model_pricing.TokenPricingRates( + input_per_million=1000, output_per_million=2000, + ), + ) for model in ["test/judge", "test/answer"] + ])), + ])) + activation = await plugin.initialize(plugin.PluginConfig(components=[pricing]), config) + + async def forbidden(_request): + raise AssertionError("plugin did not intercept the request") + + async def forbidden_stream(_request): + raise AssertionError("plugin did not intercept the stream") + yield # Make this an async generator, as required by Relay. + + async def execute(label, streaming): + chunks = [] + request = relay.LLMRequest({}, {"model": "switchyard/test", "stream": streaming, + "messages": [{"role": "user", "content": label}]}) + with relay.use_scope_stack(relay.create_scope_stack()): + with relay.scope.scope(label, relay.ScopeType.Agent): + if streaming: + stream = await relay.llm.stream_execute( + "openai.chat_completions", request, forbidden_stream, chunks.append, + lambda: {"model": "test/answer", "choices": []}, + codec=codecs.OpenAIChatCodec(), response_codec=codecs.OpenAIChatCodec(), + ) + async for _ in stream: + pass + assert chunks + else: + await relay.llm.execute("openai.chat_completions", request, forbidden, + codec=codecs.OpenAIChatCodec(), response_codec=codecs.OpenAIChatCodec()) + + await execute("buffered", False) + await execute("streamed", True) + await asyncio.gather(execute("concurrent-a", True), execute("concurrent-b", True)) + await subscribers.flush_async() + results = {label: check_workflow(events, label) + for label in ["buffered", "streamed", "concurrent-a", "concurrent-b"]} + assert Provider.calls == {"test/judge": 4, "test/answer": 4}, Provider.calls + all_calls = [e for e in events if e.get("name") == "libsy.client_call" and e.get("scope_category") == "end"] + assert len(all_calls) == 8, "unexpected detached or duplicated model call" + (output / "events.jsonl").write_text("".join(json.dumps(event) + "\n" for event in events)) + (output / "results.json").write_text(json.dumps(results, indent=2)) + print(json.dumps(results, indent=2)) + finally: + if activation is not None: + await activation.close() + subscribers.deregister("buffered-routing-test") + server.shutdown() + server.server_close() + thread.join() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--manifest", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + args.output.mkdir(parents=True, exist_ok=True) + asyncio.run(run(args.manifest.resolve(), args.output.resolve())) From 85eab26a91ff4f6dbb5b61679c201f08ed1ccbce Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Fri, 6 Nov 2026 12:48:17 -0800 Subject: [PATCH 2/4] test(relay): cover final-answer replay and incremental delivery Signed-off-by: Bryan Bednarski --- .../scripts/test_buffered_routing.py | 117 ++++++++++++++---- 1 file changed, 93 insertions(+), 24 deletions(-) diff --git a/crates/switchyard-nemo-relay-plugin/scripts/test_buffered_routing.py b/crates/switchyard-nemo-relay-plugin/scripts/test_buffered_routing.py index 6b1f74586e..d786ed0fe0 100644 --- a/crates/switchyard-nemo-relay-plugin/scripts/test_buffered_routing.py +++ b/crates/switchyard-nemo-relay-plugin/scripts/test_buffered_routing.py @@ -4,7 +4,7 @@ """Check buffered routing attribution with a packaged plugin and Relay 0.10. Run with an isolated Relay configuration and a loopback-only network policy. -The synthetic judge streams even though the classifier requests a verdict. +The fixtures exercise a streamed classifier, VGR answer replay, and gated delivery. """ import argparse @@ -23,6 +23,17 @@ class Provider(BaseHTTPRequestHandler): calls = Counter() lock = threading.Lock() + consumer_ready = {} + + @staticmethod + def frames(model, content, usage): + return [ + {"id": model, "model": model, "choices": [ + {"index": 0, "delta": {"role": "assistant", "content": content}}, + ]}, + {"id": model, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, + {"id": model, "choices": [], "usage": usage}, + ] def do_POST(self): request = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) @@ -30,44 +41,58 @@ def do_POST(self): with self.lock: self.calls[model] += 1 judge = model == "test/judge" - content = '{"target":"efficient"}' if judge else "answer" - usage = {"prompt_tokens": 11 if judge else 7, "completion_tokens": 3 if judge else 2} + verifier = model == "test/verifier" + if judge: + content = '{"target":"efficient"}' + elif verifier: + content = "yes" if request.get("logprobs") else "chat" + else: + content = "answer" + usage = {"prompt_tokens": 11 if judge or verifier else 7, + "completion_tokens": 3 if judge or verifier else 2} if judge: assert request["stream"] is True streaming = request.get("stream", False) if streaming: - frames = [ - {"id": "judge", "model": model, "choices": [ - {"index": 0, "delta": {"role": "assistant", "content": content}}, - ]}, - {"id": "judge", "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, - {"id": "judge", "choices": [], "usage": usage}, - ] + frames = self.frames(model, content, usage) body = ("".join("data: " + json.dumps(frame) + "\n\n" for frame in frames) + "data: [DONE]\n\n").encode() else: body = json.dumps({"id": "answer", "model": model, "usage": usage, "choices": [ - {"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop"}, + {"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop", + **({"logprobs": {"content": [{"top_logprobs": [ + {"token": "yes", "logprob": math.log(0.99)}, + {"token": "no", "logprob": math.log(0.01)}, + ]}]}} if request.get("logprobs") else {})}, ]}).encode() self.send_response(200) self.send_header("Content-Type", "text/event-stream" if streaming else "application/json") self.send_header("Content-Length", str(len(body))) self.end_headers() - self.wfile.write(body) + if streaming and model == "test/answer": + # Withhold EOF until the consumer sees content: eager buffering deadlocks. + prefix = ("data: " + json.dumps(frames[0]) + "\n\n").encode() + self.wfile.write(prefix) + self.wfile.flush() + label = request["messages"][-1]["content"] + assert self.consumer_ready[label].wait(3), "first chunk did not reach consumer" + self.wfile.write(body[len(prefix):]) + else: + self.wfile.write(body) def log_message(self, *_args): pass -def check_workflow(events, label): +def check_workflow(events, label, replay=False): ends = {e["uuid"]: e for e in events if e.get("kind") == "scope" and e.get("scope_category") == "end"} workflow = next(e for e in ends.values() if e.get("name") == label) outer, = [e for e in ends.values() if e.get("name") == "openai.chat_completions" and e.get("parent_uuid") == workflow["uuid"]] calls = [e for e in ends.values() if e.get("name") == "libsy.client_call" and e.get("parent_uuid") == outer["uuid"]] - assert len(calls) == 2, (label, calls) - assert {e["metadata"]["selected_model"] for e in calls} == {"test/judge", "test/answer"} + expected = {"test/replay": 1, "test/verifier": 2} if replay else {"test/judge": 1, "test/answer": 1} + assert Counter(e["metadata"]["selected_model"] for e in calls) == expected, (label, calls) total = 0.0 for call in calls: # Exactly one start/end pair: replay must not create another observation. @@ -75,13 +100,15 @@ def check_workflow(events, label): assert Counter(e["scope_category"] for e in pair) == {"start": 1, "end": 1} assert all(e["parent_uuid"] == outer["uuid"] for e in pair) assert call["metadata"]["outcome"] == "ok" - judge = call["metadata"]["selected_model"] == "test/judge" + judge = call["metadata"]["selected_model"] in {"test/judge", "test/verifier"} + if replay: + assert call["metadata"]["switchyard.call_phase"] == "routing" usage = codecs.AnthropicMessagesCodec().decode_response(call["data"]).usage assert usage["prompt_tokens"] == (11 if judge else 7), usage assert usage["completion_tokens"] == (3 if judge else 2), usage assert usage["cost"]["currency"] == "USD", usage total += usage["cost"]["total"] - assert math.isclose(total, 0.028), total + assert math.isclose(total, 0.045 if replay else 0.028), total return {"workflow_uuid": workflow["uuid"], "calls": len(calls), "logical_call_cost_usd": total} @@ -106,6 +133,19 @@ async def run(manifest, output): [targets.answer] id = "test/answer" llm_client = "local" +[targets.replay] +id = "test/replay" +llm_client = "local" +[targets.verifier] +id = "test/verifier" +llm_client = "local" +[routes.replay] +id = "switchyard/replay" +type = "vgr" +efficient_target = "replay" +capable_target = "answer" +judge_target = "verifier" +mode = "evaluate" [routes.test] id = "switchyard/test" type = "llm_classifier" @@ -131,7 +171,7 @@ async def run(manifest, output): pricing_source="synthetic", rates=model_pricing.TokenPricingRates( input_per_million=1000, output_per_million=2000, ), - ) for model in ["test/judge", "test/answer"] + ) for model in ["test/judge", "test/answer", "test/replay", "test/verifier"] ])), ])) activation = await plugin.initialize(plugin.PluginConfig(components=[pricing]), config) @@ -143,9 +183,11 @@ async def forbidden_stream(_request): raise AssertionError("plugin did not intercept the stream") yield # Make this an async generator, as required by Relay. - async def execute(label, streaming): + async def execute(label, streaming, replay=False, abandon=False): chunks = [] - request = relay.LLMRequest({}, {"model": "switchyard/test", "stream": streaming, + delivered = [] + Provider.consumer_ready[label] = threading.Event() + request = relay.LLMRequest({}, {"model": "switchyard/replay" if replay else "switchyard/test", "stream": streaming, "messages": [{"role": "user", "content": label}]}) with relay.use_scope_stack(relay.create_scope_stack()): with relay.scope.scope(label, relay.ScopeType.Agent): @@ -155,9 +197,32 @@ async def execute(label, streaming): lambda: {"model": "test/answer", "choices": []}, codec=codecs.OpenAIChatCodec(), response_codec=codecs.OpenAIChatCodec(), ) - async for _ in stream: - pass + try: + async for chunk in stream: + delivered.append(chunk) + if chunk.get("choices") and chunk["choices"][0].get("delta", {}).get("content"): + Provider.consumer_ready[label].set() + if replay: + await subscribers.flush_async() + # The live calls are settled before the first replayed content. + workflow = next(e for e in events if e.get("name") == label) + outer = next(e for e in events if e.get("name") == "openai.chat_completions" + and e.get("parent_uuid") == workflow["uuid"]) + settled = [e for e in events if e.get("name") == "libsy.client_call" + and e.get("scope_category") == "end" + and e.get("parent_uuid") == outer["uuid"]] + assert len(settled) == 3, settled + if abandon: + break + finally: + await stream.aclose() assert chunks + if replay: + expected = Provider.frames("test/replay", "answer", {"prompt_tokens": 7, "completion_tokens": 2}) + assert delivered == (expected[:1] if abandon else expected), delivered + # Relay may prefetch callbacks even after the consumer stops. + assert 1 <= len(chunks) <= len(expected) + assert chunks == expected[:len(chunks)], chunks else: await relay.llm.execute("openai.chat_completions", request, forbidden, codec=codecs.OpenAIChatCodec(), response_codec=codecs.OpenAIChatCodec()) @@ -165,12 +230,16 @@ async def execute(label, streaming): await execute("buffered", False) await execute("streamed", True) await asyncio.gather(execute("concurrent-a", True), execute("concurrent-b", True)) + await execute("replay-drained", True, replay=True) + await execute("replay-abandoned", True, replay=True, abandon=True) await subscribers.flush_async() results = {label: check_workflow(events, label) for label in ["buffered", "streamed", "concurrent-a", "concurrent-b"]} - assert Provider.calls == {"test/judge": 4, "test/answer": 4}, Provider.calls + results.update({label: check_workflow(events, label, replay=True) + for label in ["replay-drained", "replay-abandoned"]}) + assert Provider.calls == {"test/judge": 4, "test/answer": 4, "test/replay": 2, "test/verifier": 4}, Provider.calls all_calls = [e for e in events if e.get("name") == "libsy.client_call" and e.get("scope_category") == "end"] - assert len(all_calls) == 8, "unexpected detached or duplicated model call" + assert len(all_calls) == 14, "unexpected detached or duplicated model call" (output / "events.jsonl").write_text("".join(json.dumps(event) + "\n" for event in events)) (output / "results.json").write_text(json.dumps(results, indent=2)) print(json.dumps(results, indent=2)) From c012eee09fc2c45f6ec79bba8712e4125d6cc20e Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Fri, 6 Nov 2026 12:54:51 -0800 Subject: [PATCH 3/4] test(observability): retain only buffered replay regression Signed-off-by: Bryan Bednarski --- .../libsy-llm-client/tests/observability.rs | 161 ++--------- .../scripts/test_buffered_routing.py | 261 ------------------ 2 files changed, 28 insertions(+), 394 deletions(-) delete mode 100644 crates/switchyard-nemo-relay-plugin/scripts/test_buffered_routing.py diff --git a/crates/libsy-llm-client/tests/observability.rs b/crates/libsy-llm-client/tests/observability.rs index 359c69876f..025a001432 100644 --- a/crates/libsy-llm-client/tests/observability.rs +++ b/crates/libsy-llm-client/tests/observability.rs @@ -1469,157 +1469,52 @@ async fn buffered_routing_closes_observation_before_replay() -> switchyard_libsy } else { "routing-replay-dropped" }; - let workflow = format!("workflow-{model}"); let mut request = request_with_metadata(model, model); request.llm_request.stream = true; - let mut upstream_request = request.clone(); - upstream_request.llm_request.model = Some(model.into()); - let original = StreamingUsageClient - .call(upstream_request) - .await - .expect("valid upstream fixture"); - let LlmResponse::Stream(original) = original.llm_response else { - return Err(test_error("expected upstream stream")); - }; - let expected = original - .map(|event| event.expect("valid fixture")) - .collect::>() - .await; let (_, response) = run( Arc::new(RoutingCallAlgo { - name: workflow.clone(), + name: model.into(), target: model.into(), }), Arc::new(StreamingUsageClient), request, ) - .instrument(tracing::info_span!("buffered-workflow", workflow = %workflow)) + .instrument(tracing::info_span!( + "buffered-workflow", + gen_ai.request.model = model + )) .await?; - // The original I/O has finished even though nobody has polled the replay. - let before = exporter.get_finished_spans().expect("exported spans"); - let call = before - .iter() - .find(|span| span.name == format!("chat {model}")) - .expect("client span must close before replay starts"); - assert_eq!( - otel_attribute(call, "gen_ai.usage.input_tokens"), - Some(&OtelValue::I64(21)) - ); - assert_eq!( - otel_attribute(call, "gen_ai.usage.output_tokens"), - Some(&OtelValue::I64(5)) - ); - let mut ancestor = call; - while ancestor.name != "buffered-workflow" { - ancestor = before - .iter() - .find(|span| span.span_context.span_id() == ancestor.parent_span_id) - .expect("every parent up to the workflow must already be exported"); + // Upstream I/O is complete, so usage and parentage must export before replay. + let call = find_otel_span(exporter, &format!("chat {model}"), model); + let workflow = find_otel_span(exporter, "buffered-workflow", model); + assert_eq!(call.parent_span_id, workflow.span_context.span_id()); + for (field, tokens) in [ + ("gen_ai.usage.input_tokens", 21), + ("gen_ai.usage.output_tokens", 5), + ] { + assert_eq!(otel_attribute(&call, field), Some(&OtelValue::I64(tokens))); } - assert_eq!( - otel_attribute(ancestor, "workflow"), - Some(&OtelValue::String(workflow.into())) - ); - - let LlmResponse::Stream(mut replay) = response.llm_response else { - return Err(test_error("expected buffered stream replay")); - }; + let count = exporter.get_finished_spans().expect("exported spans").len(); if drain { - // Provider events and their trailing usage survive buffering unchanged. - let mut events = Vec::new(); - while let Some(event) = replay.next().await { - events.push(event.expect("valid replay")); - } + let replay = response + .llm_response + .into_agg() + .await + .expect("valid replay"); assert_eq!( - events, expected, - "normalized chunks and preserved provider data must survive replay" - ); - } - drop(replay); - assert_eq!( - exporter.get_finished_spans().expect("exported spans").len(), - before.len(), - "replaying or dropping buffered data must not export another observation" - ); - } - Ok(()) -} - -struct InterruptedRoutingStream { - stall: bool, -} - -#[async_trait] -impl RoutedLlmClient for InterruptedRoutingStream { - async fn call(&self, _request: Request) -> Result { - let usage = LlmResponseChunk::Usage(Usage { - input_tokens: Some(11), - output_tokens: Some(3), - ..Usage::default() - }); - let prefix = futures::stream::iter([Ok(usage.into())]); - let tail: switchyard_protocol::LlmResponseStream = if self.stall { - futures::stream::pending().boxed() - } else { - futures::stream::iter([Err(LlmClientError::ResponseTranslation( - "broken stream".into(), - ))]) - .boxed() - }; - Ok(Response { - llm_response: LlmResponse::Stream(prefix.chain(tail).boxed()), - metadata: None, - upstream_headers: http::HeaderMap::new(), - }) - } -} - -#[tokio::test] -async fn interrupted_routing_buffer_retains_observed_usage() -> switchyard_libsy::Result<()> { - let _guard = serialize_test().lock().await; - let (store, _, _, exporter, _) = telemetry(); - for stall in [false, true] { - let model = if stall { - "routing-buffer-cancelled" - } else { - "routing-buffer-failed" - }; - let mut request = request_with_metadata(model, model); - request.llm_request.stream = true; - let mut execution = Box::pin(run( - Arc::new(RoutingCallAlgo { - name: model.into(), - target: model.into(), - }), - Arc::new(InterruptedRoutingStream { stall }), - request, - )); - if stall { - assert!( - tokio::time::timeout(Duration::from_millis(50), execution.as_mut()) - .await - .is_err() + replay.outputs[0].content, + vec![ContentBlock::Text { + text: "hello".into() + }] ); - drop(execution); } else { - assert!(execution.await.is_err()); + drop(response); } - let span = find_span(&store.spans(), "libsy.client_call", "selected_model", model); - for (field, value) in [ - ("gen_ai.usage.input_tokens", "11"), - ("gen_ai.usage.output_tokens", "3"), - ("outcome", if stall { "cancelled" } else { "error" }), - ] { - assert_eq!(span.fields.get(field).map(String::as_str), Some(value)); - } - let exported = exporter.get_finished_spans().expect("exported spans"); assert_eq!( - exported - .iter() - .filter(|span| span.name == format!("chat {model}")) - .count(), - 1 + exporter.get_finished_spans().expect("exported spans").len(), + count, + "consuming or dropping replay must not export another observation" ); } Ok(()) diff --git a/crates/switchyard-nemo-relay-plugin/scripts/test_buffered_routing.py b/crates/switchyard-nemo-relay-plugin/scripts/test_buffered_routing.py deleted file mode 100644 index d786ed0fe0..0000000000 --- a/crates/switchyard-nemo-relay-plugin/scripts/test_buffered_routing.py +++ /dev/null @@ -1,261 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Check buffered routing attribution with a packaged plugin and Relay 0.10. - -Run with an isolated Relay configuration and a loopback-only network policy. -The fixtures exercise a streamed classifier, VGR answer replay, and gated delivery. -""" - -import argparse -import asyncio -import json -import math -import threading -from collections import Counter -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from pathlib import Path - -import nemo_relay as relay -from nemo_relay import codecs, model_pricing, plugin, subscribers - - -class Provider(BaseHTTPRequestHandler): - calls = Counter() - lock = threading.Lock() - consumer_ready = {} - - @staticmethod - def frames(model, content, usage): - return [ - {"id": model, "model": model, "choices": [ - {"index": 0, "delta": {"role": "assistant", "content": content}}, - ]}, - {"id": model, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, - {"id": model, "choices": [], "usage": usage}, - ] - - def do_POST(self): - request = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) - model = request["model"] - with self.lock: - self.calls[model] += 1 - judge = model == "test/judge" - verifier = model == "test/verifier" - if judge: - content = '{"target":"efficient"}' - elif verifier: - content = "yes" if request.get("logprobs") else "chat" - else: - content = "answer" - usage = {"prompt_tokens": 11 if judge or verifier else 7, - "completion_tokens": 3 if judge or verifier else 2} - if judge: - assert request["stream"] is True - streaming = request.get("stream", False) - if streaming: - frames = self.frames(model, content, usage) - body = ("".join("data: " + json.dumps(frame) + "\n\n" for frame in frames) - + "data: [DONE]\n\n").encode() - else: - body = json.dumps({"id": "answer", "model": model, "usage": usage, "choices": [ - {"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop", - **({"logprobs": {"content": [{"top_logprobs": [ - {"token": "yes", "logprob": math.log(0.99)}, - {"token": "no", "logprob": math.log(0.01)}, - ]}]}} if request.get("logprobs") else {})}, - ]}).encode() - self.send_response(200) - self.send_header("Content-Type", "text/event-stream" if streaming else "application/json") - self.send_header("Content-Length", str(len(body))) - self.end_headers() - if streaming and model == "test/answer": - # Withhold EOF until the consumer sees content: eager buffering deadlocks. - prefix = ("data: " + json.dumps(frames[0]) + "\n\n").encode() - self.wfile.write(prefix) - self.wfile.flush() - label = request["messages"][-1]["content"] - assert self.consumer_ready[label].wait(3), "first chunk did not reach consumer" - self.wfile.write(body[len(prefix):]) - else: - self.wfile.write(body) - - def log_message(self, *_args): - pass - - -def check_workflow(events, label, replay=False): - ends = {e["uuid"]: e for e in events if e.get("kind") == "scope" and e.get("scope_category") == "end"} - workflow = next(e for e in ends.values() if e.get("name") == label) - outer, = [e for e in ends.values() if e.get("name") == "openai.chat_completions" - and e.get("parent_uuid") == workflow["uuid"]] - calls = [e for e in ends.values() if e.get("name") == "libsy.client_call" - and e.get("parent_uuid") == outer["uuid"]] - expected = {"test/replay": 1, "test/verifier": 2} if replay else {"test/judge": 1, "test/answer": 1} - assert Counter(e["metadata"]["selected_model"] for e in calls) == expected, (label, calls) - total = 0.0 - for call in calls: - # Exactly one start/end pair: replay must not create another observation. - pair = [e for e in events if e.get("uuid") == call["uuid"]] - assert Counter(e["scope_category"] for e in pair) == {"start": 1, "end": 1} - assert all(e["parent_uuid"] == outer["uuid"] for e in pair) - assert call["metadata"]["outcome"] == "ok" - judge = call["metadata"]["selected_model"] in {"test/judge", "test/verifier"} - if replay: - assert call["metadata"]["switchyard.call_phase"] == "routing" - usage = codecs.AnthropicMessagesCodec().decode_response(call["data"]).usage - assert usage["prompt_tokens"] == (11 if judge else 7), usage - assert usage["completion_tokens"] == (3 if judge else 2), usage - assert usage["cost"]["currency"] == "USD", usage - total += usage["cost"]["total"] - assert math.isclose(total, 0.045 if replay else 0.028), total - return {"workflow_uuid": workflow["uuid"], "calls": len(calls), "logical_call_cost_usd": total} - - -async def run(manifest, output): - server = ThreadingHTTPServer(("127.0.0.1", 0), Provider) - thread = threading.Thread(target=server.serve_forever, daemon=True) - thread.start() - events = [] - subscribers.register("buffered-routing-test", lambda event: events.append(event.to_dict())) - activation = None - try: - deployment = output / "switchyard.toml" - deployment.write_text(f'''schema_version = 1 -[llm_clients.local] -format = "openai_chat" -base_url = "http://127.0.0.1:{server.server_port}/v1" -max_retries = 0 -[targets.judge] -id = "test/judge" -llm_client = "local" -extra_body = {{stream = true}} -[targets.answer] -id = "test/answer" -llm_client = "local" -[targets.replay] -id = "test/replay" -llm_client = "local" -[targets.verifier] -id = "test/verifier" -llm_client = "local" -[routes.replay] -id = "switchyard/replay" -type = "vgr" -efficient_target = "replay" -capable_target = "answer" -judge_target = "verifier" -mode = "evaluate" -[routes.test] -id = "switchyard/test" -type = "llm_classifier" -mode = "custom" -models = {{judge = ["judge"], efficient = ["answer"], any = ["answer"]}} -default_target = "efficient" -prompt = "Select efficient." -response_schema = '{{"type":"object","properties":{{"target":{{"type":"string","enum":["efficient"]}}}},"required":["target"],"additionalProperties":false}}' -policy = {{type = "target_selector", selector = "/target"}} -''') - config = output / "plugins.toml" - config.write_text(f'''[[plugins.dynamic]] -manifest = {json.dumps(str(manifest))} -[plugins.dynamic.config] -switchyard_config_path = {json.dumps(str(deployment))} -[plugins.policy.overrides."nvidia.switchyard"] -attestation = "integrity_only" -''') - pricing = model_pricing.ComponentSpec(model_pricing.PricingConfig(sources=[ - model_pricing.InlineSource(model_pricing.PricingCatalog(entries=[ - model_pricing.ModelPricing( - provider="anthropic", model_id=model, pricing_as_of="2026-11-05", - pricing_source="synthetic", rates=model_pricing.TokenPricingRates( - input_per_million=1000, output_per_million=2000, - ), - ) for model in ["test/judge", "test/answer", "test/replay", "test/verifier"] - ])), - ])) - activation = await plugin.initialize(plugin.PluginConfig(components=[pricing]), config) - - async def forbidden(_request): - raise AssertionError("plugin did not intercept the request") - - async def forbidden_stream(_request): - raise AssertionError("plugin did not intercept the stream") - yield # Make this an async generator, as required by Relay. - - async def execute(label, streaming, replay=False, abandon=False): - chunks = [] - delivered = [] - Provider.consumer_ready[label] = threading.Event() - request = relay.LLMRequest({}, {"model": "switchyard/replay" if replay else "switchyard/test", "stream": streaming, - "messages": [{"role": "user", "content": label}]}) - with relay.use_scope_stack(relay.create_scope_stack()): - with relay.scope.scope(label, relay.ScopeType.Agent): - if streaming: - stream = await relay.llm.stream_execute( - "openai.chat_completions", request, forbidden_stream, chunks.append, - lambda: {"model": "test/answer", "choices": []}, - codec=codecs.OpenAIChatCodec(), response_codec=codecs.OpenAIChatCodec(), - ) - try: - async for chunk in stream: - delivered.append(chunk) - if chunk.get("choices") and chunk["choices"][0].get("delta", {}).get("content"): - Provider.consumer_ready[label].set() - if replay: - await subscribers.flush_async() - # The live calls are settled before the first replayed content. - workflow = next(e for e in events if e.get("name") == label) - outer = next(e for e in events if e.get("name") == "openai.chat_completions" - and e.get("parent_uuid") == workflow["uuid"]) - settled = [e for e in events if e.get("name") == "libsy.client_call" - and e.get("scope_category") == "end" - and e.get("parent_uuid") == outer["uuid"]] - assert len(settled) == 3, settled - if abandon: - break - finally: - await stream.aclose() - assert chunks - if replay: - expected = Provider.frames("test/replay", "answer", {"prompt_tokens": 7, "completion_tokens": 2}) - assert delivered == (expected[:1] if abandon else expected), delivered - # Relay may prefetch callbacks even after the consumer stops. - assert 1 <= len(chunks) <= len(expected) - assert chunks == expected[:len(chunks)], chunks - else: - await relay.llm.execute("openai.chat_completions", request, forbidden, - codec=codecs.OpenAIChatCodec(), response_codec=codecs.OpenAIChatCodec()) - - await execute("buffered", False) - await execute("streamed", True) - await asyncio.gather(execute("concurrent-a", True), execute("concurrent-b", True)) - await execute("replay-drained", True, replay=True) - await execute("replay-abandoned", True, replay=True, abandon=True) - await subscribers.flush_async() - results = {label: check_workflow(events, label) - for label in ["buffered", "streamed", "concurrent-a", "concurrent-b"]} - results.update({label: check_workflow(events, label, replay=True) - for label in ["replay-drained", "replay-abandoned"]}) - assert Provider.calls == {"test/judge": 4, "test/answer": 4, "test/replay": 2, "test/verifier": 4}, Provider.calls - all_calls = [e for e in events if e.get("name") == "libsy.client_call" and e.get("scope_category") == "end"] - assert len(all_calls) == 14, "unexpected detached or duplicated model call" - (output / "events.jsonl").write_text("".join(json.dumps(event) + "\n" for event in events)) - (output / "results.json").write_text(json.dumps(results, indent=2)) - print(json.dumps(results, indent=2)) - finally: - if activation is not None: - await activation.close() - subscribers.deregister("buffered-routing-test") - server.shutdown() - server.server_close() - thread.join() - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--manifest", type=Path, required=True) - parser.add_argument("--output", type=Path, required=True) - args = parser.parse_args() - args.output.mkdir(parents=True, exist_ok=True) - asyncio.run(run(args.manifest.resolve(), args.output.resolve())) From 8e5a74d2a2f2ad129133ddec0ef6926100bfd87b Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Mon, 9 Nov 2026 08:08:26 -0800 Subject: [PATCH 4/4] fix(observability): record sanitized stream failures once Signed-off-by: Bryan Bednarski --- crates/libsy-llm-client/src/observability.rs | 9 +- crates/libsy-llm-client/src/run.rs | 10 +-- .../libsy-llm-client/tests/observability.rs | 83 +++++++++++++++++++ 3 files changed, 93 insertions(+), 9 deletions(-) diff --git a/crates/libsy-llm-client/src/observability.rs b/crates/libsy-llm-client/src/observability.rs index 26a916343c..e1b92bc13d 100644 --- a/crates/libsy-llm-client/src/observability.rs +++ b/crates/libsy-llm-client/src/observability.rs @@ -296,6 +296,11 @@ impl ClientStreamObserver { let error_type = llm_client_error_type(error); match error { LlmClientError::Host { .. } => record_host_error(&self.span), + LlmClientError::UpstreamHttp { status, .. } => record_client_error( + &self.span, + &error_type, + &format_args!("upstream HTTP {}", status.as_u16()), + ), _ => record_client_error(&self.span, &error_type, error), } self.outcome = Outcome::Failed; @@ -319,8 +324,8 @@ impl ClientStreamObserver { record_client_error(&self.span, "response_translation", message); self.outcome = Outcome::Failed; } - LlmResponseChunk::StreamError { message } => { - record_client_error(&self.span, "502", message); + LlmResponseChunk::StreamError { .. } => { + record_client_error(&self.span, "502", &"upstream HTTP 502"); self.outcome = Outcome::Failed; } _ => {} diff --git a/crates/libsy-llm-client/src/run.rs b/crates/libsy-llm-client/src/run.rs index b0609227cf..132819641f 100644 --- a/crates/libsy-llm-client/src/run.rs +++ b/crates/libsy-llm-client/src/run.rs @@ -458,9 +458,9 @@ async fn call_one( let result = async { let response = async { client?.call(request).await } .await - .map_err(|source| LibsyError::client_call(model_id.clone(), source))?; - // Observe the upstream stream; replay must not retain or re-observe the call. - let mut response = observability::observe_client_call(Ok(response))?; + .map_err(|source| LibsyError::client_call(model_id.clone(), source)); + // Observe dispatch and live I/O once; buffering must not re-observe replay. + let mut response = observability::observe_client_call(response)?; if is_routing && let LlmResponse::Stream(chunks) = response.llm_response { let (chunks, usage) = buffer_routing_stream(chunks) .await @@ -477,10 +477,6 @@ async fn call_one( response.set_served_model(model_id); response }); - let result = match result { - Ok(response) => Ok(response), - Err(error) => observability::observe_client_call(Err(error)), - }; let result = if is_routing { result } else { diff --git a/crates/libsy-llm-client/tests/observability.rs b/crates/libsy-llm-client/tests/observability.rs index 025a001432..fdf191ea68 100644 --- a/crates/libsy-llm-client/tests/observability.rs +++ b/crates/libsy-llm-client/tests/observability.rs @@ -1520,6 +1520,89 @@ async fn buffered_routing_closes_observation_before_replay() -> switchyard_libsy Ok(()) } +struct FailingStreamClient(bool); + +#[async_trait] +impl RoutedLlmClient for FailingStreamClient { + async fn call(&self, _request: Request) -> Result { + let failure = if self.0 { + Err(LlmClientError::UpstreamHttp { + status: http::StatusCode::BAD_GATEWAY, + body: LEAKED_CONTENT.into(), + }) + } else { + Ok(LlmResponseChunk::StreamError { + message: LEAKED_CONTENT.into(), + } + .into()) + }; + Ok(Response { + llm_response: LlmResponse::Stream( + futures::stream::iter([ + Ok(LlmResponseChunk::Usage(Usage { + input_tokens: Some(11), + output_tokens: Some(3), + ..Usage::default() + }) + .into()), + failure, + ]) + .boxed(), + ), + metadata: None, + upstream_headers: http::HeaderMap::new(), + }) + } +} + +#[tokio::test] +async fn routing_stream_error_is_recorded_once_without_upstream_body() { + let _guard = serialize_test().lock().await; + let (_, _, _, exporter, _) = telemetry(); + for http_error in [false, true] { + let model = if http_error { + "routing-http-stream-error" + } else { + "routing-in-band-error" + }; + let mut request = request_with_metadata(model, model); + request.llm_request.stream = true; + assert!( + run( + Arc::new(RoutingCallAlgo { + name: model.into(), + target: model.into() + }), + Arc::new(FailingStreamClient(http_error)), + request, + ) + .await + .is_err() + ); + let call = find_otel_span(exporter, &format!("chat {model}"), model); + for (key, expected) in [ + ("outcome", OtelValue::String("error".into())), + ("error.type", OtelValue::String("502".into())), + ("error", OtelValue::String("upstream HTTP 502".into())), + ("gen_ai.usage.input_tokens", OtelValue::I64(11)), + ("gen_ai.usage.output_tokens", OtelValue::I64(3)), + ] { + let values: Vec<_> = call + .attributes + .iter() + .filter(|attr| attr.key.as_str() == key) + .map(|attr| attr.value.clone()) + .collect(); + assert_eq!( + values, + vec![expected], + "{key} must be exported exactly once" + ); + } + assert!(!format!("{:?}", call.attributes).contains(LEAKED_CONTENT)); + } +} + struct TimeoutClient; #[async_trait]