Skip to main content

starweaver_model/transport/
audit.rs

1use std::{
2    collections::{BTreeMap, BTreeSet},
3    sync::{Arc, Mutex},
4};
5
6use serde::{Deserialize, Serialize};
7use serde_json::{Map, Value};
8
9use super::{HttpMethod, HttpRequest};
10
11/// Shared provider request audit recorder.
12pub type DynProviderRequestAuditRecorder = Arc<dyn ProviderRequestAuditRecorder>;
13
14/// Provider request payload capture level.
15#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
16#[serde(rename_all = "snake_case")]
17pub enum ProviderRequestAuditPayloadPolicy {
18    /// Do not capture this payload section.
19    #[default]
20    Omit,
21    /// Capture this payload section after redacting configured sensitive keys.
22    Redacted,
23    /// Capture this payload section without audit-level redaction.
24    Full,
25}
26
27/// Explicit policy for provider request audit snapshots.
28#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
29pub struct ProviderRequestAuditPolicy {
30    /// Header capture policy.
31    #[serde(default)]
32    pub headers: ProviderRequestAuditPayloadPolicy,
33    /// JSON body capture policy.
34    #[serde(default)]
35    pub body: ProviderRequestAuditPayloadPolicy,
36    /// Case-insensitive header names redacted by [`ProviderRequestAuditPayloadPolicy::Redacted`].
37    #[serde(default = "default_sensitive_header_keys")]
38    pub sensitive_header_keys: BTreeSet<String>,
39    /// Case-insensitive JSON object keys redacted by [`ProviderRequestAuditPayloadPolicy::Redacted`].
40    #[serde(default = "default_sensitive_body_keys")]
41    pub sensitive_body_keys: BTreeSet<String>,
42    /// Replacement value used for redacted fields.
43    #[serde(default = "default_redaction_value")]
44    pub redaction_value: Value,
45}
46
47impl Default for ProviderRequestAuditPolicy {
48    fn default() -> Self {
49        Self {
50            headers: ProviderRequestAuditPayloadPolicy::Omit,
51            body: ProviderRequestAuditPayloadPolicy::Omit,
52            sensitive_header_keys: default_sensitive_header_keys(),
53            sensitive_body_keys: default_sensitive_body_keys(),
54            redaction_value: default_redaction_value(),
55        }
56    }
57}
58
59impl ProviderRequestAuditPolicy {
60    /// Capture method, URL, timeout, and request metadata only.
61    #[must_use]
62    pub fn metadata_only() -> Self {
63        Self::default()
64    }
65
66    /// Capture headers and JSON body after audit-level sensitive-field redaction.
67    #[must_use]
68    pub fn redacted_payloads() -> Self {
69        Self {
70            headers: ProviderRequestAuditPayloadPolicy::Redacted,
71            body: ProviderRequestAuditPayloadPolicy::Redacted,
72            ..Self::default()
73        }
74    }
75
76    /// Capture full headers and JSON body. Use only for explicit local debugging or fixture work.
77    #[must_use]
78    pub fn full_payloads() -> Self {
79        Self {
80            headers: ProviderRequestAuditPayloadPolicy::Full,
81            body: ProviderRequestAuditPayloadPolicy::Full,
82            ..Self::default()
83        }
84    }
85
86    /// Add a case-insensitive sensitive header name.
87    #[must_use]
88    pub fn with_sensitive_header_key(mut self, key: impl AsRef<str>) -> Self {
89        self.sensitive_header_keys
90            .insert(normalize_key(key.as_ref()));
91        self
92    }
93
94    /// Add a case-insensitive sensitive JSON object key.
95    #[must_use]
96    pub fn with_sensitive_body_key(mut self, key: impl AsRef<str>) -> Self {
97        self.sensitive_body_keys.insert(normalize_key(key.as_ref()));
98        self
99    }
100}
101
102/// Captured provider request snapshot stored outside redacted observability spans.
103#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
104pub struct ProviderRequestAuditSnapshot {
105    /// Provider adapter name.
106    pub provider_name: String,
107    /// Model name used by the adapter.
108    pub model_name: String,
109    /// Whether the request was sent through the streaming path.
110    pub stream: bool,
111    /// HTTP method.
112    pub method: HttpMethod,
113    /// Absolute endpoint URL.
114    pub url: String,
115    /// Captured headers, when enabled by policy.
116    #[serde(default, skip_serializing_if = "Option::is_none")]
117    pub headers: Option<BTreeMap<String, String>>,
118    /// Captured JSON body, when enabled by policy.
119    #[serde(default, skip_serializing_if = "Option::is_none")]
120    pub body: Option<Value>,
121    /// Request timeout in milliseconds.
122    #[serde(default, skip_serializing_if = "Option::is_none")]
123    pub timeout_ms: Option<u64>,
124    /// Request metadata for audit lookup and correlation.
125    #[serde(default, skip_serializing_if = "Map::is_empty")]
126    pub metadata: Map<String, Value>,
127}
128
129impl ProviderRequestAuditSnapshot {
130    /// Build a policy-filtered snapshot from an HTTP request.
131    #[must_use]
132    pub fn from_request(
133        provider_name: impl Into<String>,
134        model_name: impl Into<String>,
135        stream: bool,
136        request: &HttpRequest,
137        policy: &ProviderRequestAuditPolicy,
138    ) -> Self {
139        Self {
140            provider_name: provider_name.into(),
141            model_name: model_name.into(),
142            stream,
143            method: request.method,
144            url: request.url.clone(),
145            headers: capture_headers(&request.headers, policy),
146            body: capture_body(&request.body, policy),
147            timeout_ms: request
148                .timeout
149                .and_then(|duration| u64::try_from(duration.as_millis()).ok()),
150            metadata: request.metadata.clone(),
151        }
152    }
153}
154
155/// Recorder for provider request audit snapshots.
156pub trait ProviderRequestAuditRecorder: Send + Sync {
157    /// Record a provider request snapshot.
158    fn record_provider_request(&self, snapshot: ProviderRequestAuditSnapshot);
159}
160
161/// Deterministic in-memory provider request audit recorder.
162#[derive(Default)]
163pub struct InMemoryProviderRequestAuditRecorder {
164    snapshots: Mutex<Vec<ProviderRequestAuditSnapshot>>,
165}
166
167impl InMemoryProviderRequestAuditRecorder {
168    /// Create an empty recorder.
169    #[must_use]
170    pub fn new() -> Self {
171        Self::default()
172    }
173
174    /// Return captured snapshots.
175    #[must_use]
176    pub fn snapshots(&self) -> Vec<ProviderRequestAuditSnapshot> {
177        self.snapshots
178            .lock()
179            .map_or_else(|_| Vec::new(), |v| v.clone())
180    }
181
182    /// Remove all captured snapshots and return them.
183    #[must_use]
184    pub fn take_snapshots(&self) -> Vec<ProviderRequestAuditSnapshot> {
185        self.snapshots
186            .lock()
187            .map_or_else(|_| Vec::new(), |mut v| std::mem::take(&mut *v))
188    }
189}
190
191impl ProviderRequestAuditRecorder for InMemoryProviderRequestAuditRecorder {
192    fn record_provider_request(&self, snapshot: ProviderRequestAuditSnapshot) {
193        if let Ok(mut snapshots) = self.snapshots.lock() {
194            snapshots.push(snapshot);
195        }
196    }
197}
198
199#[derive(Clone)]
200pub struct ProviderRequestAuditCapture {
201    recorder: DynProviderRequestAuditRecorder,
202    policy: ProviderRequestAuditPolicy,
203}
204
205impl ProviderRequestAuditCapture {
206    pub(crate) fn new(
207        recorder: DynProviderRequestAuditRecorder,
208        policy: ProviderRequestAuditPolicy,
209    ) -> Self {
210        Self { recorder, policy }
211    }
212
213    pub(crate) fn record(
214        &self,
215        provider_name: &str,
216        model_name: &str,
217        stream: bool,
218        request: &HttpRequest,
219    ) {
220        self.recorder
221            .record_provider_request(ProviderRequestAuditSnapshot::from_request(
222                provider_name,
223                model_name,
224                stream,
225                request,
226                &self.policy,
227            ));
228    }
229}
230
231fn capture_headers(
232    headers: &BTreeMap<String, String>,
233    policy: &ProviderRequestAuditPolicy,
234) -> Option<BTreeMap<String, String>> {
235    match policy.headers {
236        ProviderRequestAuditPayloadPolicy::Omit => None,
237        ProviderRequestAuditPayloadPolicy::Full => Some(headers.clone()),
238        ProviderRequestAuditPayloadPolicy::Redacted => Some(
239            headers
240                .iter()
241                .map(|(key, value)| {
242                    if policy.sensitive_header_keys.contains(&normalize_key(key)) {
243                        (key.clone(), policy.redaction_value.to_string())
244                    } else {
245                        (key.clone(), value.clone())
246                    }
247                })
248                .collect(),
249        ),
250    }
251}
252
253fn capture_body(body: &Value, policy: &ProviderRequestAuditPolicy) -> Option<Value> {
254    match policy.body {
255        ProviderRequestAuditPayloadPolicy::Omit => None,
256        ProviderRequestAuditPayloadPolicy::Full => Some(body.clone()),
257        ProviderRequestAuditPayloadPolicy::Redacted => {
258            let mut body = body.clone();
259            redact_json_value(&mut body, policy);
260            Some(body)
261        }
262    }
263}
264
265fn redact_json_value(value: &mut Value, policy: &ProviderRequestAuditPolicy) {
266    match value {
267        Value::Object(object) => {
268            for (key, value) in object {
269                if policy.sensitive_body_keys.contains(&normalize_key(key)) {
270                    *value = policy.redaction_value.clone();
271                } else {
272                    redact_json_value(value, policy);
273                }
274            }
275        }
276        Value::Array(values) => {
277            for value in values {
278                redact_json_value(value, policy);
279            }
280        }
281        Value::Null | Value::Bool(_) | Value::Number(_) | Value::String(_) => {}
282    }
283}
284
285fn normalize_key(key: &str) -> String {
286    key.to_ascii_lowercase()
287}
288
289fn default_sensitive_header_keys() -> BTreeSet<String> {
290    [
291        "authorization",
292        "proxy-authorization",
293        "x-api-key",
294        "api-key",
295        "cookie",
296        "set-cookie",
297    ]
298    .into_iter()
299    .map(str::to_string)
300    .collect()
301}
302
303fn default_sensitive_body_keys() -> BTreeSet<String> {
304    [
305        "api_key",
306        "apikey",
307        "authorization",
308        "access_token",
309        "refresh_token",
310        "id_token",
311        "token",
312        "password",
313        "secret",
314    ]
315    .into_iter()
316    .map(str::to_string)
317    .collect()
318}
319
320fn default_redaction_value() -> Value {
321    serde_json::json!({ "redacted": true })
322}