switchyard_llm_client/
payload_capture.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Explicit opt-in for HTTP body capture by a host's tracing subscriber.
5
6use std::fmt;
7use std::future::Future;
8
9tokio::task_local! {
10    static CAPTURE_ENABLED: ();
11}
12
13/// Enables raw HTTP body events while polling `future` on the current task.
14///
15/// The host must install a subscriber for `switchyard.payload` and apply its
16/// own filtering, size limits, and redaction. Spawned tasks do not inherit capture.
17pub async fn with_payload_capture<T>(future: impl Future<Output = T>) -> T {
18    CAPTURE_ENABLED.scope((), future).await
19}
20
21pub(crate) fn record<D: fmt::Display>(phase: &'static str, body: impl FnOnce() -> D) {
22    if CAPTURE_ENABLED.try_with(|_| ()).is_ok() {
23        tracing::trace!(target: "switchyard.payload", phase, body = %body());
24    }
25}
26
27#[cfg(test)]
28mod tests {
29    use std::sync::Arc;
30
31    use parking_lot::Mutex;
32    use tracing::field::Field;
33    use tracing::instrument::WithSubscriber;
34    use tracing_subscriber::Layer;
35    use tracing_subscriber::layer::{Context, SubscriberExt};
36
37    use super::*;
38
39    struct Recorder(Arc<Mutex<Vec<String>>>);
40
41    impl<S: tracing::Subscriber> Layer<S> for Recorder {
42        fn on_event(&self, event: &tracing::Event<'_>, _: Context<'_, S>) {
43            event.record(&mut |field: &Field, value: &dyn fmt::Debug| {
44                if field.name() == "body" {
45                    self.0.lock().push(format!("{value:?}"));
46                }
47            });
48        }
49    }
50
51    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
52    async fn capture_requires_opt_in_and_does_not_leak_between_tasks_or_after_cancellation() {
53        let events = Arc::new(Mutex::new(Vec::new()));
54        let subscriber = tracing_subscriber::registry().with(Recorder(Arc::clone(&events)));
55        async {
56            record("request", || -> &str {
57                panic!("disabled capture must not evaluate the body")
58            });
59            tokio::join!(
60                with_payload_capture(async {
61                    tokio::spawn(
62                        async {
63                            record("request", || "spawned");
64                        }
65                        .with_current_subscriber(),
66                    )
67                    .await
68                    .unwrap();
69                    tokio::task::yield_now().await;
70                    record("request", || "enabled");
71                }),
72                async {
73                    tokio::task::yield_now().await;
74                    record("request", || "other-task");
75                },
76            );
77            let mut pending = Box::pin(with_payload_capture(std::future::pending::<()>()));
78            assert!(futures::poll!(&mut pending).is_pending());
79            drop(pending);
80            record("response", || "after-cancellation");
81        }
82        .with_subscriber(subscriber)
83        .await;
84        assert_eq!(*events.lock(), ["enabled"]);
85    }
86}
switchyard_llm_client/payload_capture.rs

switchyard_llm_client/
payload_capture.rs
