Skip to main content

switchyard_runner/
provider_key_redactor.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Shared output-boundary redaction for configured provider credentials.
5
6use serde_json::Value;
7
8/// Replaces configured provider keys without changing the credentials sent upstream.
9///
10/// Apply this after response translation, including to preserved provider fields.
11/// Each call is independent: keys split across stream events or separate JSON strings
12/// are not reconstructed. This is exact substring matching, not a general secret detector.
13#[derive(Default)]
14pub struct ProviderKeyRedactor {
15    text: Vec<String>,
16    json: Vec<String>,
17}
18
19impl ProviderKeyRedactor {
20    /// Registers nonempty keys, matching longer keys before their prefixes.
21    pub fn new(keys: &[String]) -> Self {
22        let mut raw: Vec<String> = keys.iter().filter(|key| !key.is_empty()).cloned().collect();
23        raw.sort_by(|a, b| b.len().cmp(&a.len()).then_with(|| a.cmp(b)));
24        raw.dedup();
25        let mut json: Vec<String> = raw
26            .iter()
27            .map(|key| {
28                let encoded = Value::String(key.clone()).to_string();
29                encoded[1..encoded.len() - 1].to_string()
30            })
31            .collect();
32        json.sort_by_key(|key| std::cmp::Reverse(key.len()));
33        // Match escaped forms before any raw-key suffix they contain. Otherwise a
34        // leading backslash can survive replacement and invalidate embedded JSON.
35        let mut text = raw;
36        text.extend(json.iter().cloned());
37        text.sort_by(|a, b| b.len().cmp(&a.len()).then_with(|| a.cmp(b)));
38        text.dedup();
39        Self { text, json }
40    }
41
42    /// Returns whether there are no configured secrets to redact.
43    pub fn is_empty(&self) -> bool {
44        self.text.is_empty()
45    }
46
47    /// Redacts JSON-escaped keys in serialized JSON, preserving the server's wire behavior.
48    pub fn json(&self, value: String) -> String {
49        replace(value, &self.json)
50    }
51
52    /// Redacts raw and JSON-escaped keys in text, including headers and error messages.
53    pub fn text(&self, value: String) -> String {
54        replace(value, &self.text)
55    }
56
57    /// Redacts raw and JSON-escaped keys in string values and member names throughout JSON.
58    pub fn value(&self, value: Value) -> Value {
59        if self.is_empty() {
60            return value;
61        }
62        match value {
63            Value::String(text) => Value::String(self.text(text)),
64            Value::Array(values) => {
65                Value::Array(values.into_iter().map(|value| self.value(value)).collect())
66            }
67            Value::Object(fields) => Value::Object(
68                fields
69                    .into_iter()
70                    .map(|(name, value)| (self.text(name), self.value(value)))
71                    .collect(),
72            ),
73            value => value,
74        }
75    }
76}
77
78fn replace(mut value: String, secrets: &[String]) -> String {
79    for secret in secrets {
80        if value.contains(secret) {
81            value = value.replace(secret, "[REDACTED]");
82        }
83    }
84    value
85}
86
87#[cfg(test)]
88mod tests {
89    use serde_json::json;
90
91    use super::*;
92
93    #[test]
94    fn escaped_credentials_are_redacted_in_names_and_nested_values() {
95        let key = "synthetic-\"provider\\key-雪";
96        let redactor = ProviderKeyRedactor::new(&[key.into()]);
97        let value = json!({
98            "error": {"debug": format!("Bearer {key}"), key: [key, "ordinary diagnostics"]},
99            "count": 3, "enabled": true, "missing": null,
100        });
101        let sanitized = redactor.value(value.clone());
102        assert_eq!(sanitized["error"]["debug"], "Bearer [REDACTED]");
103        assert_eq!(
104            sanitized["error"]["[REDACTED]"],
105            json!(["[REDACTED]", "ordinary diagnostics"])
106        );
107        assert_eq!(sanitized["count"], 3);
108        assert_eq!(sanitized["enabled"], true);
109        assert!(sanitized["missing"].is_null());
110        let wire: Value = serde_json::from_str(&redactor.json(value.to_string())).unwrap();
111        assert_eq!(wire, sanitized);
112        assert_eq!(redactor.text(format!("Bearer {key}")), "Bearer [REDACTED]");
113    }
114
115    #[test]
116    fn tool_arguments_and_preserved_extensions_are_redacted() {
117        let redactor = ProviderKeyRedactor::new(&["qa-token".into()]);
118        let value = json!({
119            "arguments": json!({"credential": "qa-token", "query": "ordinary text"}).to_string(),
120            "provider_extension": {"debug": ["Bearer qa-token"]},
121        });
122        let sanitized = redactor.value(value.clone());
123        let arguments: Value =
124            serde_json::from_str(sanitized["arguments"].as_str().unwrap()).unwrap();
125        assert_eq!(
126            arguments,
127            json!({"credential": "[REDACTED]", "query": "ordinary text"})
128        );
129        assert_eq!(
130            sanitized["provider_extension"]["debug"],
131            json!(["Bearer [REDACTED]"])
132        );
133        assert_eq!(
134            serde_json::from_str::<Value>(&redactor.json(value.to_string())).unwrap(),
135            sanitized
136        );
137    }
138
139    #[test]
140    fn json_escaped_credentials_are_redacted_in_error_text() {
141        for key in [r#"provider-"key"#, r"provider-\key", r#"\provider-"key"#] {
142            let redactor = ProviderKeyRedactor::new(&[key.into()]);
143            let detail = json!({"credential": key, "message": "ordinary diagnostics"});
144            let text = format!("raw: {key}; provider error: {detail}");
145            let expected = format!(
146                "raw: [REDACTED]; provider error: {}",
147                json!({"credential": "[REDACTED]", "message": "ordinary diagnostics"})
148            );
149            assert_eq!(redactor.text(text), expected);
150        }
151    }
152
153    #[test]
154    fn nested_json_strings_cannot_recover_escaped_credentials() {
155        for key in [r#"provider-"key"#, r"provider-\key", r#"\provider-"key"#] {
156            let redactor = ProviderKeyRedactor::new(&[key.into()]);
157            let arguments = json!({key: {"credential": key, "query": "ordinary text"}});
158            let response = json!({"result": {"arguments": arguments.to_string()}});
159            // Decode both layers, as a client consuming serialized tool arguments would.
160            let sanitized = redactor.value(response);
161            let decoded: Value =
162                serde_json::from_str(sanitized["result"]["arguments"].as_str().unwrap()).unwrap();
163            assert_eq!(
164                decoded,
165                json!({"[REDACTED]": {
166                    "credential": "[REDACTED]", "query": "ordinary text"
167                }})
168            );
169        }
170    }
171
172    #[test]
173    fn escaped_form_is_replaced_before_an_overlapping_raw_key() {
174        let key = r"\provider-key";
175        let redactor = ProviderKeyRedactor::new(&[key.into()]);
176        let text = json!({"credential": key}).to_string();
177        // Replacing the raw suffix first would leave a stray escape in the JSON string.
178        let sanitized = redactor.text(text);
179        assert_eq!(
180            serde_json::from_str::<Value>(&sanitized).unwrap(),
181            json!({"credential": "[REDACTED]"})
182        );
183    }
184
185    #[test]
186    fn empty_duplicate_and_overlapping_keys_preserve_nonsecret_content() {
187        let redactor = ProviderKeyRedactor::new(&[
188            "".into(),
189            "qa-token".into(),
190            "qa-token-extra".into(),
191            "qa-token".into(),
192        ]);
193        assert_eq!(
194            redactor.text("keep qa-token-extra / qa-token".into()),
195            "keep [REDACTED] / [REDACTED]"
196        );
197        assert!(ProviderKeyRedactor::new(&["".into()]).is_empty());
198        let value = json!({"choices": [{"delta": {"content": "ordinary text"}}]});
199        assert_eq!(redactor.value(value.clone()), value);
200        assert_eq!(ProviderKeyRedactor::default().value(value.clone()), value);
201    }
202
203    #[test]
204    fn independent_events_do_not_reassemble_split_credentials() {
205        let redactor = ProviderKeyRedactor::new(&["qa-token".into()]);
206        let events = [json!({"delta": "qa-"}), json!({"delta": "token"})];
207        for event in events {
208            assert_eq!(redactor.value(event.clone()), event);
209        }
210    }
211}