diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/client.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/client.rs index 3165965d8..6abfb51d6 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/client.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/client.rs @@ -189,7 +189,11 @@ pub(super) async fn forward_client_message( ) .await; let client_event = match redacted_client_event { - Ok(Some(redacted)) => redacted, + Ok(Some(redaction)) => { + // 这一轮的映射登记到连接上，响应帧才能在最后一跳还原回真实值。 + bound.redaction_restorer.register(redaction.session); + redaction.client_event + } Ok(None) => client_event, Err(error) => { warn!( diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/connection.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/connection.rs index ecc6937b5..40902a5c5 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/connection.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/connection.rs @@ -2,7 +2,7 @@ use std::time::Duration; -use axum::extract::ws::WebSocket; +use axum::extract::ws::{Message as AxumWsMessage, WebSocket}; use futures_util::{SinkExt, StreamExt}; use serde_json::Value; use tokio::time::sleep; @@ -486,10 +486,22 @@ pub(super) async fn relay_bound_connection( detach_exhausted_upstream(bound, directive, &context.trace_id).await; continue; } - if let Err(error) = send_client_message( - client_socket, - upstream_message_to_client(upstream_message.clone()), - ).await { + // 响应侧还原：HTTP 在把响应体交给客户端之前会把占位符换回真实值 + // （`privacy::restore_sync_response_body` / + // `privacy::StreamingResponseRestorer`），这里是 WS 的同一个位置 + // ——最后一跳之前，并且在 `capture_client_frame` 之前，所以审计与 + // 终态观测继续消费脱敏态的事件。没有命中还原时保持上游原字节。 + let restored_client_frame = parsed_upstream_frame + .as_ref() + .map(ParsedResponsesWebSocketFrame::event) + .and_then(|event| { + bound.redaction_restorer.restore_provider_frame_text(event) + }); + let client_frame = match restored_client_frame { + Some(restored) => AxumWsMessage::Text(restored.into()), + None => upstream_message_to_client(upstream_message.clone()), + }; + if let Err(error) = send_client_message(client_socket, client_frame).await { warn!( event_name = "responses_websocket_client_send_failed", log_type = "ops", diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/redaction.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/redaction.rs index d3ae81540..575dd4eaa 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/redaction.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/redaction.rs @@ -1,4 +1,4 @@ -//! Responses WebSocket 请求侧 PII 脱敏。 +//! Responses WebSocket 两侧的 PII 脱敏：请求侧 mask + 响应侧 restore。 //! //! HTTP 路径在前门建 `RedactionSessionSlot` 并塞进 `parts.extensions`，planner //! 只有拿到这个 slot 才会脱敏。WS 的 planning Parts 是合成的：四个规划入口 @@ -8,6 +8,37 @@ //! //! 因此约定：**进入任何下游用途之前，客户端 `response.create` 只在这里脱敏一次**， //! 之后所有路径都只看脱敏后的事件。 +//! +//! # 响应侧 +//! +//! 只 mask 不 restore 是半个实现：HTTP 在把响应交给客户端之前会把占位符换回真实值 +//! （`privacy::restore_sync_response_body` / `privacy::StreamingResponseRestorer`）， +//! WS 少了这一步，客户端就会直接看到 ``。 +//! [`ResponsesWebSocketRedactionRestorer`] 补上这一跳，语义与 HTTP 完全一致： +//! 复用 `privacy::restore_json_strings`，只还原本连接自己 mask 出来的映射， +//! 未映射的占位符原样透传。 +//! +//! ## session 为什么活在连接上而不是活在这一轮里 +//! +//! mask session 由 planner 写进 per-turn 的 slot，而 slot 随 planning Parts 在 +//! 规划结束时就被丢弃，响应帧到达时已经无处可取。可选的存活范围有两个： +//! +//! * 挂在 `LogicalTurn` 上：这一轮结束即释放，是 HTTP「一个请求一个 session」的 +//! 直译。但 WS 的会话历史留在上游：continuation 只发增量输入，第 1 轮的 +//! `input` 不会在第 3 轮重发。于是第 3 轮的响应里若回显了第 1 轮的占位符 +//! （"你刚才给我的邮箱是……"），本轮 session 里没有这条映射，占位符就漏给客户端。 +//! HTTP 不会漏，是因为它每次都重发整段历史，重新 mask 同一个值会派生出同一个 +//! sentinel（HMAC over 规则 + bucket + 值），所以映射天然齐备。 +//! * 挂在连接上（当前实现）：每轮仍然各自 mask、各自持有独立 session +//! （per-turn 语义不变），连接只是把最近若干轮的 session 留下来一起参与还原， +//! 凑出的映射集合正好等于「等价 HTTP 请求会拥有的那一份」。 +//! +//! 选后者。代价是每帧最多对 [`MAX_RETAINED_TURN_REDACTION_SESSIONS`] 个 session +//! 各扫一遍，以及这些 session 的映射会驻留到连接结束；用有界 FIFO 兜住上限。 +//! 窗口不够用或每帧成本变高时，正确的下一步是在 `privacy` 侧提供跨 session 的 +//! 合并匹配器，而不是把这个窗口调大。 + +use std::collections::VecDeque; use serde_json::Value; @@ -15,6 +46,7 @@ use crate::ai_serving::{ resolve_local_decision_execution_runtime_auth_context, resolve_provider_chat_pii_redaction, }; use crate::control::GatewayControlDecision; +use crate::privacy::{restore_json_strings, RedactionSession, RedactionSessionSlot}; use crate::{AppState, GatewayError}; /// Responses WebSocket 只承载 `openai:responses`，脱敏规则按这个客户端格式选取。 @@ -26,9 +58,24 @@ const RESPONSES_WEBSOCKET_CLIENT_API_FORMAT: &str = "openai:responses"; /// planner 后续用真实 candidate_id 再取一次配置时，body 已是脱敏态、不会重复写入。 const WEBSOCKET_TURN_REDACTION_CANDIDATE_ID: &str = "responses_websocket_turn"; +/// 一条连接最多留几轮的 mask session 用于响应侧还原。 +/// +/// 取值权衡见模块文档：调大会线性增加每帧还原成本和常驻映射量，调小则更容易漏还原 +/// 上游历史里更早那几轮的占位符。8 覆盖的是「上游最可能回显的最近窗口」。 +const MAX_RETAINED_TURN_REDACTION_SESSIONS: usize = 8; + +/// 一轮客户端 `response.create` 的请求侧脱敏结果。 +#[derive(Debug)] +pub(super) struct ResponsesWebSocketTurnRedaction { + /// 脱敏后的客户端事件；这一轮之后所有下游路径都只看它。 + pub(super) client_event: Value, + /// 这一轮 mask 出来的映射表，响应侧还原只能靠它。 + pub(super) session: RedactionSession, +} + /// 对一条客户端 `response.create` 做请求侧脱敏。 /// -/// 返回 `Some(脱敏后的事件)` 仅当脱敏真正命中；`None` 表示未启用或没有命中，调用方 +/// 返回 `Some(..)` 仅当脱敏真正命中；`None` 表示未启用或没有命中，调用方 /// 继续用原事件即可（避免未开启脱敏时多一次整包 clone）。 /// /// 脱敏只改写 `instructions` / `input`（见 `privacy::mask_openai_responses_request_value`）， @@ -42,7 +89,7 @@ pub(super) async fn redact_responses_websocket_client_event( parts: &http::request::Parts, control_decision: &GatewayControlDecision, client_event: &Value, -) -> Result, GatewayError> { +) -> Result, GatewayError> { let Some(auth_context) = resolve_local_decision_execution_runtime_auth_context(control_decision) else { @@ -60,7 +107,73 @@ pub(super) async fn redact_responses_websocket_client_event( if !redaction.redacted { return Ok(None); } - Ok(Some(redaction.body_json.into_owned())) + // mask 命中时 `resolve_provider_chat_pii_redaction` 必定把 session 写进 slot。 + // 取不到就是内部契约被破坏了，此时继续下发意味着这一轮的响应无法还原、占位符 + // 会漏给客户端；按本模块既有的「脱敏链路出错就让这一轮失败」处理，不做降级。 + let Some(session) = parts + .extensions + .get::() + .and_then(|slot| slot.take_for_candidate(Some(WEBSOCKET_TURN_REDACTION_CANDIDATE_ID))) + else { + return Err(GatewayError::Internal( + "chat pii redaction masked a Responses WebSocket turn without retaining its session" + .to_string(), + )); + }; + Ok(Some(ResponsesWebSocketTurnRedaction { + client_event: redaction.body_json.into_owned(), + session, + })) +} + +/// 一条连接上「我们 mask 过哪些映射」的留存集合，供响应侧还原使用。 +/// +/// 每轮一个独立 session（per-turn mask 语义不变），连接按 FIFO 留最近 +/// [`MAX_RETAINED_TURN_REDACTION_SESSIONS`] 轮。上游重绑不清空：客户端仍在同一段 +/// 对话里，旧占位符可能随重发的输入再次出现。 +#[derive(Default)] +pub(super) struct ResponsesWebSocketRedactionRestorer { + sessions: VecDeque, +} + +impl ResponsesWebSocketRedactionRestorer { + /// 登记这一轮的 mask session。 + pub(super) fn register(&mut self, session: RedactionSession) { + if session.mapping_count() == 0 { + return; + } + self.sessions.push_back(session); + while self.sessions.len() > MAX_RETAINED_TURN_REDACTION_SESSIONS { + self.sessions.pop_front(); + } + } + + /// 把一帧 provider 事件里的占位符换回真实值，返回要发给客户端的帧文本。 + /// + /// `None` 表示这一帧没有任何东西要还原，调用方必须原样转发上游字节：未启用 + /// 脱敏（没有任何 session）时连 clone 都不做。 + /// + /// 入参只读：审计与终态观测继续消费脱敏态的事件，还原只作用于发往客户端的 + /// 那一份拷贝，和 HTTP 侧「审计存脱敏体、线上还原」保持一致。 + pub(super) fn restore_provider_frame_text(&self, event: &Value) -> Option { + if self.sessions.is_empty() { + return None; + } + let mut restored_event = event.clone(); + let mut restored = false; + for session in &self.sessions { + // 逐 session 还原而不是合并映射：每个 session 只认自己 mask 过的 + // sentinel（`RedactionSession::restore_text`），跨 session 合并会绕开 + // 这条边界。同一个值在不同轮派生出的 sentinel 相同，所以顺序无关。 + restored |= restore_json_strings(&mut restored_event, session); + } + if !restored { + return None; + } + // 刚从 JSON 解析出来的 Value 再序列化不会失败；真失败时宁可让客户端看到 + // 占位符，也不能丢掉这一帧——丢帧会让客户端的协议状态机卡死。 + serde_json::to_string(&restored_event).ok() + } } #[cfg(test)] @@ -80,7 +193,10 @@ mod tests { }; use super::super::turn_state::LogicalTurn; use super::super::turn::prepare_responses_websocket_turn_decision; - use super::redact_responses_websocket_client_event; + use super::{ + redact_responses_websocket_client_event, ResponsesWebSocketRedactionRestorer, + ResponsesWebSocketTurnRedaction, MAX_RETAINED_TURN_REDACTION_SESSIONS, + }; use crate::ai_serving::{AiExecutionDecision, ResponsesWebSocketBodyNormalization}; use crate::control::{GatewayControlAuthContext, GatewayControlDecision}; use crate::handlers::proxy::websocket::ingress::WebSocketRequestContext; @@ -89,6 +205,11 @@ mod tests { const TEST_USER_ID: &str = "user-responses-ws-redaction"; const TEST_API_KEY_ID: &str = "api-key-responses-ws-redaction"; const TEST_EMAIL: &str = "ws.user@example.com"; + /// 另一轮用的 PII，用来证明连接级还原覆盖到更早的轮次。 + const OTHER_TEST_EMAIL: &str = "ws.other@example.com"; + /// 不是本连接 mask 出来的占位符：格式合法（符合 sentinel 正则），但没有任何 + /// session 记过它，必须原样透传。 + const FOREIGN_SENTINEL: &str = ""; fn auth_export_record() -> StoredAuthApiKeyExportRecord { StoredAuthApiKeyExportRecord::new( @@ -191,6 +312,10 @@ mod tests { } fn client_event() -> Value { + client_event_with_email(TEST_EMAIL) + } + + fn client_event_with_email(email: &str) -> Value { json!({ "type": "response.create", "model": "public-model", @@ -198,11 +323,46 @@ mod tests { "generate": false, "input": [{ "role": "user", - "content": [{"type": "input_text", "text": format!("mail {TEST_EMAIL}")}] + "content": [{"type": "input_text", "text": format!("mail {email}")}] }] }) } + /// 真跑一遍请求侧脱敏，拿到这一轮的生效事件和 mask session。 + async fn turn_redaction( + state: &AppState, + decision: &GatewayControlDecision, + email: &str, + ) -> ResponsesWebSocketTurnRedaction { + let context = websocket_context(decision.clone()); + let parts = build_planning_parts(&context); + let event = client_event_with_email(email); + redact_responses_websocket_client_event(state, &parts, &context.decision, &event) + .await + .expect("redaction should resolve") + .expect("an email in the request should be redacted") + } + + /// 这一轮为 `email` 派生出的占位符。 + fn sentinel_for(redaction: &ResponsesWebSocketTurnRedaction, email: &str) -> String { + redaction + .session + .sentinel_for_original(email) + .expect("a masked email must have a sentinel") + .to_string() + } + + /// 上游回显占位符的一帧 provider 事件。 + fn provider_delta_frame(text: &str) -> Value { + json!({ + "type": "response.output_text.delta", + "item_id": "msg_ws", + "output_index": 0, + "content_index": 0, + "delta": text, + }) + } + #[tokio::test] async fn websocket_client_event_is_redacted_without_losing_protocol_fields() { let state = redaction_enabled_state(); @@ -214,7 +374,8 @@ mod tests { redact_responses_websocket_client_event(&state, &parts, &context.decision, &event) .await .expect("redaction should resolve") - .expect("an email in the request should be redacted"); + .expect("an email in the request should be redacted") + .client_event; let serialized = serde_json::to_string(&redacted).expect("event should serialize"); assert!(!serialized.contains(TEST_EMAIL), "{serialized}"); @@ -240,7 +401,8 @@ mod tests { redact_responses_websocket_client_event(&state, &parts, &context.decision, &event) .await .expect("redaction should resolve") - .expect("an email in the request should be redacted"); + .expect("an email in the request should be redacted") + .client_event; // 复用同一个 parts/slot，和 re-plan 在同一 turn 内二次脱敏的情形一致。 let second_pass = @@ -273,13 +435,9 @@ mod tests { /// 真跑一遍脱敏，拿到这一轮的「生效事件」。 async fn redacted_client_event(state: &AppState, decision: &GatewayControlDecision) -> Value { - let context = websocket_context(decision.clone()); - let parts = build_planning_parts(&context); - let event = client_event(); - redact_responses_websocket_client_event(state, &parts, &context.decision, &event) + turn_redaction(state, decision, TEST_EMAIL) .await - .expect("redaction should resolve") - .expect("an email in the request should be redacted") + .client_event } /// 只有 `action` 没有 serde 默认值，其余字段都能省略。 @@ -477,4 +635,220 @@ mod tests { ); assert_redacted_json(report_context, "quota retry report context"); } + + // ----------------------------------------------------------------------- + // 响应侧还原 + // ----------------------------------------------------------------------- + + /// 本次修复的核心：上游把占位符回显在事件里，客户端必须拿到真实值。 + #[tokio::test] + async fn provider_frame_placeholders_are_restored_before_client_delivery() { + let state = redaction_enabled_state(); + let redaction = turn_redaction(&state, &control_decision(), TEST_EMAIL).await; + let sentinel = sentinel_for(&redaction, TEST_EMAIL); + let mut restorer = ResponsesWebSocketRedactionRestorer::default(); + restorer.register(redaction.session); + + let frame = provider_delta_frame(&format!("your mail is {sentinel}")); + let restored = restorer + .restore_provider_frame_text(&frame) + .expect("a frame echoing this turn's sentinel must be restored"); + + assert!( + restored.contains(TEST_EMAIL), + "the client must receive the real value: {restored}" + ); + assert!( + !restored.contains(&sentinel), + "no sentinel may survive to the client: {restored}" + ); + // 协议字段不受影响，客户端的状态机照旧。 + let restored: Value = serde_json::from_str(&restored).expect("restored frame is JSON"); + assert_eq!(restored["type"], "response.output_text.delta"); + assert_eq!(restored["item_id"], "msg_ws"); + assert_eq!(restored["output_index"], 0); + } + + /// Codex 把多个事件批量塞进 `{"chunks":[...]}`，还原必须走进批量里。 + #[tokio::test] + async fn placeholders_batched_inside_a_chunks_envelope_are_restored() { + let state = redaction_enabled_state(); + let redaction = turn_redaction(&state, &control_decision(), TEST_EMAIL).await; + let sentinel = sentinel_for(&redaction, TEST_EMAIL); + let mut restorer = ResponsesWebSocketRedactionRestorer::default(); + restorer.register(redaction.session); + + let frame = json!({ + "chunks": [ + provider_delta_frame("plain delta"), + provider_delta_frame(&format!("mail {sentinel}")), + ] + }); + let restored = restorer + .restore_provider_frame_text(&frame) + .expect("a batched sentinel must be restored"); + + assert!(restored.contains(TEST_EMAIL), "{restored}"); + assert!(!restored.contains(&sentinel), "{restored}"); + } + + /// 只还原本连接 mask 过的映射，和 `RedactionSession::restore_text` 一致： + /// 别处来的占位符（比如客户端自己发的、或上一条连接的）保持原样。 + #[tokio::test] + async fn an_unmapped_placeholder_is_left_untouched() { + let state = redaction_enabled_state(); + let redaction = turn_redaction(&state, &control_decision(), TEST_EMAIL).await; + let sentinel = sentinel_for(&redaction, TEST_EMAIL); + let mut restorer = ResponsesWebSocketRedactionRestorer::default(); + restorer.register(redaction.session); + + let frame = provider_delta_frame(&format!("{FOREIGN_SENTINEL} and {sentinel}")); + let restored = restorer + .restore_provider_frame_text(&frame) + .expect("the mapped sentinel is still restored"); + + assert!(restored.contains(TEST_EMAIL), "{restored}"); + assert!( + restored.contains(FOREIGN_SENTINEL), + "an unmapped placeholder must survive verbatim: {restored}" + ); + } + + /// 没有命中还原时必须让调用方原样转发上游字节。 + #[tokio::test] + async fn a_frame_without_known_placeholders_is_not_rewritten() { + let state = redaction_enabled_state(); + let redaction = turn_redaction(&state, &control_decision(), TEST_EMAIL).await; + let mut restorer = ResponsesWebSocketRedactionRestorer::default(); + restorer.register(redaction.session); + + assert!( + restorer + .restore_provider_frame_text(&provider_delta_frame("nothing to restore")) + .is_none(), + "a frame with no mapped sentinel must be relayed byte-for-byte" + ); + assert!( + restorer + .restore_provider_frame_text(&provider_delta_frame(FOREIGN_SENTINEL)) + .is_none(), + "a frame that only carries unmapped placeholders must not be rewritten" + ); + } + + /// 未启用脱敏（或这条连接从没 mask 到东西）时，还原器必须完全不介入： + /// 连 clone 都不做，输出就是上游原字节。 + #[tokio::test] + async fn a_restorer_without_sessions_never_rewrites_a_frame() { + let restorer = ResponsesWebSocketRedactionRestorer::default(); + + assert!(restorer + .restore_provider_frame_text(&provider_delta_frame(FOREIGN_SENTINEL)) + .is_none()); + assert!(restorer + .restore_provider_frame_text(&provider_delta_frame(TEST_EMAIL)) + .is_none()); + } + + /// 空 session（启用了脱敏但这一轮没命中任何规则）不该被留下来白扫每一帧。 + #[tokio::test] + async fn a_session_without_mappings_is_not_retained() { + let state = redaction_enabled_state(); + let hmac_key = state + .encryption_key() + .expect("the test state carries an encryption key") + .as_bytes() + .to_vec(); + let empty_session = crate::privacy::RedactionSession::new( + crate::privacy::RedactionSessionConfig::default_ttl(hmac_key, 0), + ); + let mut restorer = ResponsesWebSocketRedactionRestorer::default(); + restorer.register(empty_session); + + assert!(restorer + .restore_provider_frame_text(&provider_delta_frame(FOREIGN_SENTINEL)) + .is_none()); + } + + /// 还原只作用于发给客户端的那一份拷贝：审计和终态观测消费的事件必须保持脱敏态。 + #[tokio::test] + async fn restoring_does_not_mutate_the_event_the_audit_path_keeps() { + let state = redaction_enabled_state(); + let redaction = turn_redaction(&state, &control_decision(), TEST_EMAIL).await; + let sentinel = sentinel_for(&redaction, TEST_EMAIL); + let mut restorer = ResponsesWebSocketRedactionRestorer::default(); + restorer.register(redaction.session); + + let frame = provider_delta_frame(&format!("mail {sentinel}")); + let before = frame.clone(); + let _ = restorer + .restore_provider_frame_text(&frame) + .expect("the frame is restored for the client"); + + assert_eq!( + frame, before, + "capture_client_frame / 终态观测拿到的事件必须仍是脱敏态" + ); + } + + /// 连接级持有的意义：WS 的会话历史留在上游，continuation 只发增量输入， + /// 所以第 2 轮的响应可能回显第 1 轮的占位符。per-turn 持有会漏掉这一条。 + #[tokio::test] + async fn a_later_turn_restores_a_placeholder_first_masked_by_an_earlier_turn() { + let state = redaction_enabled_state(); + let decision = control_decision(); + let first = turn_redaction(&state, &decision, TEST_EMAIL).await; + let second = turn_redaction(&state, &decision, OTHER_TEST_EMAIL).await; + let first_sentinel = sentinel_for(&first, TEST_EMAIL); + let second_sentinel = sentinel_for(&second, OTHER_TEST_EMAIL); + assert_ne!(first_sentinel, second_sentinel); + + let mut restorer = ResponsesWebSocketRedactionRestorer::default(); + restorer.register(first.session); + restorer.register(second.session); + + let frame = provider_delta_frame(&format!("{first_sentinel} then {second_sentinel}")); + let restored = restorer + .restore_provider_frame_text(&frame) + .expect("both turns' sentinels are restorable on this connection"); + + assert!(restored.contains(TEST_EMAIL), "{restored}"); + assert!(restored.contains(OTHER_TEST_EMAIL), "{restored}"); + assert!(!restored.contains(&first_sentinel), "{restored}"); + assert!(!restored.contains(&second_sentinel), "{restored}"); + } + + /// 留存窗口是有界的：长连接不能无限累积映射，代价是更早的轮次会退回 + /// 「占位符原样透传」而不是被错误还原成别的值。 + #[tokio::test] + async fn the_retained_session_window_is_bounded() { + let state = redaction_enabled_state(); + let decision = control_decision(); + let oldest = turn_redaction(&state, &decision, TEST_EMAIL).await; + let oldest_sentinel = sentinel_for(&oldest, TEST_EMAIL); + let mut restorer = ResponsesWebSocketRedactionRestorer::default(); + restorer.register(oldest.session); + + // 再灌满整个窗口，最老的那一轮必须被挤出去。 + let mut newest_sentinel = String::new(); + for index in 0..MAX_RETAINED_TURN_REDACTION_SESSIONS { + let email = format!("ws.turn{index}@example.com"); + let redaction = turn_redaction(&state, &decision, &email).await; + newest_sentinel = sentinel_for(&redaction, &email); + restorer.register(redaction.session); + } + + assert!( + restorer + .restore_provider_frame_text(&provider_delta_frame(&oldest_sentinel)) + .is_none(), + "the evicted turn's sentinel is relayed verbatim, never mis-restored" + ); + assert!( + restorer + .restore_provider_frame_text(&provider_delta_frame(&newest_sentinel)) + .is_some(), + "the most recent turns stay restorable" + ); + } } diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/session.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/session.rs index 64f9cb201..8a8d0dbf1 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/session.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/session.rs @@ -227,9 +227,11 @@ pub(super) async fn run_responses_websocket( &first_event, ) .await; - let first_event = match redacted_first_event { - Ok(Some(redacted)) => redacted, - Ok(None) => first_event, + // 首轮的 mask session 要活到响应帧还原，但连接此刻还没绑定，只能先接住， + // 等 `bind_responses_upstream` 之后登记到连接上。 + let (first_event, first_turn_redaction_session) = match redacted_first_event { + Ok(Some(redaction)) => (redaction.client_event, Some(redaction.session)), + Ok(None) => (first_event, None), Err(error) => { warn!( event_name = "responses_websocket_redaction_failed", @@ -413,6 +415,9 @@ pub(super) async fn run_responses_websocket( }; first_turn.mark_upstream_request_sent(); first_turn.set_provider_response_headers(bound.upstream_response_headers.clone()); + if let Some(session) = first_turn_redaction_session { + bound.redaction_restorer.register(session); + } bound.turn_state.begin( LogicalTurn::new(first_event, 1, first_logical_turn_id), ActiveProviderAttempt::new(&state, first_turn), @@ -524,6 +529,7 @@ mod tests { observe_active_response_rebind_safety, record_exhausted_bound_key, should_request_full_continuation_retry, }; + use super::super::redaction::ResponsesWebSocketRedactionRestorer; use super::super::request::{ changed_followup_response_create_model, continuation_requires_same_upstream, normalize_followup_response_create, planned_response_create_event, @@ -1212,6 +1218,7 @@ mod tests { pending_adapter_observation: None, exhausted_exclusions: ExhaustedResponsesWebSocketExclusions::default(), pending_turn_finalization: None, + redaction_restorer: ResponsesWebSocketRedactionRestorer::default(), } } diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/state.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/state.rs index 35851e981..df123c8e1 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/state.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/state.rs @@ -9,6 +9,7 @@ use tokio::task::JoinHandle; use super::adapter::{ResponsesWebSocketDrainDirective, ResponsesWebSocketProtocolAdapter}; use super::binding::UpstreamBindingIdentity; +use super::redaction::ResponsesWebSocketRedactionRestorer; use super::turn_state::ResponsesTurnState; use crate::ai_serving::{AiExecutionDecision, ResponsesWebSocketBodyNormalization}; @@ -28,6 +29,12 @@ pub(super) struct BoundResponsesConnection { pub(super) binding_identity: UpstreamBindingIdentity, /// 这条连接上「有没有正在进行的 logical turn」的唯一事实来源。 pub(super) turn_state: ResponsesTurnState, + /// 这条连接迄今 mask 出来的映射，用于把 provider 事件里的占位符换回真实值。 + /// + /// 刻意按连接持有而不是按 turn 持有：WS 的会话历史留在上游，continuation 只发 + /// 增量输入，所以后面几轮的响应可能回显更早那几轮的占位符（理由详见 + /// [`super::redaction`]）。上游重绑时不重置。 + pub(super) redaction_restorer: ResponsesWebSocketRedactionRestorer, pub(super) next_turn_index: u64, pub(super) upstream_response_headers: BTreeMap, pub(super) pending_adapter_drain: Option, diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/upstream.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/upstream.rs index 97b018ba5..e0dc1fc52 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/upstream.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/upstream.rs @@ -7,6 +7,7 @@ use wreq::ws::message::Message as WreqWsMessage; use super::adapter::ResponsesWebSocketProtocolAdapter; use super::binding::{UpstreamBindingIdentity, UpstreamBindingIdentityError}; +use super::redaction::ResponsesWebSocketRedactionRestorer; use super::request::planned_response_create_event; use super::state::{BoundResponsesConnection, ExhaustedResponsesWebSocketExclusions}; use super::turn_state::ResponsesTurnState; @@ -122,6 +123,8 @@ async fn bind_responses_upstream_inner( // 首条 response.create 已经发出，但这一轮的 logical turn 和 attempt 由调用方 // 通过 `ResponsesTurnState::begin` 装上：绑定本身不持有记账状态。 turn_state: ResponsesTurnState::Idle, + // 同理，这一轮的 mask session 也由调用方登记：绑定看不到脱敏链路。 + redaction_restorer: ResponsesWebSocketRedactionRestorer::default(), next_turn_index: 2, upstream_response_headers: upstream.response_headers, pending_adapter_drain: None, diff --git a/apps/aether-gateway/src/privacy/mod.rs b/apps/aether-gateway/src/privacy/mod.rs index 3e5ddc858..b00c8de01 100644 --- a/apps/aether-gateway/src/privacy/mod.rs +++ b/apps/aether-gateway/src/privacy/mod.rs @@ -2523,7 +2523,14 @@ fn restore_json_response_body( }) } -fn restore_json_strings(value: &mut Value, session: &RedactionSession) -> bool { +/// 递归把 JSON 里的占位符换回真实值，只认本 `session` 记录过的映射。 +/// +/// 同步响应体（[`restore_sync_response_body`]）和 Responses WebSocket 的 +/// provider 事件帧（`handlers::proxy::websocket::responses::redaction`）共用它， +/// 两边因此保持同一套还原语义：未映射的占位符原样保留，`type` / `model` / `id` +/// 这类协议字段虽然也被遍历，但它们不可能包含本 session 派生出的 sentinel， +/// 所以不会被改写。 +pub(crate) fn restore_json_strings(value: &mut Value, session: &RedactionSession) -> bool { match value { Value::String(text) => { let restored = session.restore_text(text); diff --git a/crates/aether-testing/integration/tests/responses_websocket_e2e.rs b/crates/aether-testing/integration/tests/responses_websocket_e2e.rs index cf73f917d..a2c0b1223 100644 --- a/crates/aether-testing/integration/tests/responses_websocket_e2e.rs +++ b/crates/aether-testing/integration/tests/responses_websocket_e2e.rs @@ -336,6 +336,73 @@ async fn provider_quota_exhaustion_transparently_retries_onto_another_key() -> R Ok(()) } +/// 脱敏的另一半：请求侧把真实 PII 换成占位符发给上游，响应侧必须在推给客户端之前 +/// 换回真实值。 +/// +/// 上游把收到的 `input` 原样回显，所以它回来的就是占位符——这一条同时钉住了两个 +/// 方向：上游不能看到原文，客户端不能看到占位符。 +#[tokio::test] +async fn redacted_pii_is_restored_before_the_client_sees_a_provider_frame() -> Result<(), BoxError> +{ + const CLIENT_EMAIL: &str = "responses.ws.pii@example.com"; + + let harness = Harness::start_with_pii_redaction(UpstreamBehavior::EchoInputBack).await?; + let mut client = harness.connect().await?; + + client + .send(response_create( + json!({"input": format!("my mail is {CLIENT_EMAIL}")}), + )) + .await?; + + let delta = receive_event(&mut client, "response.output_text.delta").await?; + let delta_text = delta + .get("delta") + .and_then(Value::as_str) + .ok_or("the provider delta must carry text")?; + assert!( + delta_text.contains(CLIENT_EMAIL), + "the client must receive the restored value: {delta_text}" + ); + assert!( + !delta_text.contains(" bool { audit.status.eq_ignore_ascii_case("pending") } @@ -384,21 +451,65 @@ impl ProviderFixture { } } +/// 这条用例要不要打开 chat PII 脱敏模块。 +/// +/// 默认关闭：其余用例都靠原文 body 断言上游看到了什么，打开脱敏会把断言目标换成 +/// 占位符。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum PiiRedaction { + Disabled, + Enabled, +} + +impl PiiRedaction { + const fn is_enabled(self) -> bool { + matches!(self, Self::Enabled) + } +} + impl Harness { async fn start(behavior: UpstreamBehavior) -> Result { - Self::start_with_fixture(behavior, ProviderFixture::SingleOpenAiKey).await + Self::start_with( + behavior, + ProviderFixture::SingleOpenAiKey, + PiiRedaction::Disabled, + ) + .await } async fn start_with_fixture( behavior: UpstreamBehavior, fixture: ProviderFixture, + ) -> Result { + Self::start_with(behavior, fixture, PiiRedaction::Disabled).await + } + + async fn start_with_pii_redaction(behavior: UpstreamBehavior) -> Result { + Self::start_with( + behavior, + ProviderFixture::SingleOpenAiKey, + PiiRedaction::Enabled, + ) + .await + } + + async fn start_with( + behavior: UpstreamBehavior, + fixture: ProviderFixture, + redaction: PiiRedaction, ) -> Result { let upstream = Arc::new(MockUpstreamState::new(behavior)); let upstream_server = SpawnedServer::start(mock_upstream_router(Arc::clone(&upstream))).await?; let database = TemporarySqlite::new(); - prepare_and_seed_database(&database.config, upstream_server.base_url(), fixture).await?; + prepare_and_seed_database( + &database.config, + upstream_server.base_url(), + fixture, + redaction, + ) + .await?; let data_config = GatewayDataConfig::from_database_config(database.config.clone()) .with_encryption_key(DEVELOPMENT_ENCRYPTION_KEY); @@ -618,6 +729,11 @@ enum UpstreamBehavior { /// 第一轮刻意不发 `response.created`：任何标准 `response.*` 事件都会让 /// codex adapter 把这一轮判成 replay-unsafe，透明重试就不会发生。 QuotaExhaustedThenComplete, + /// 把收到的 `input` 原样回显成一个 delta，再正常完成。 + /// + /// 上游看到的是脱敏后的 body，所以回显出来的就是占位符——正是响应侧还原要处理 + /// 的形状。 + EchoInputBack, } #[derive(Debug)] @@ -688,6 +804,11 @@ async fn run_mock_upstream( if event.get("type").and_then(Value::as_str) != Some("response.create") { continue; } + let echoed_input = event + .get("input") + .and_then(Value::as_str) + .unwrap_or_default() + .to_string(); let turn = { let mut events = state.events.lock().await; events.push(event); @@ -719,6 +840,14 @@ async fn run_mock_upstream( break; } } + UpstreamBehavior::EchoInputBack => { + if send_mock_turn_with_delta(&mut socket, &response_id, &echoed_input) + .await + .is_err() + { + break; + } + } } } AxumWsMessage::Ping(payload) => { @@ -766,13 +895,21 @@ async fn send_mock_created(socket: &mut WebSocket, response_id: &str) -> Result< } async fn send_mock_turn(socket: &mut WebSocket, response_id: &str) -> Result<(), axum::Error> { + send_mock_turn_with_delta(socket, response_id, "hello").await +} + +async fn send_mock_turn_with_delta( + socket: &mut WebSocket, + response_id: &str, + delta: &str, +) -> Result<(), axum::Error> { send_mock_created(socket, response_id).await?; send_mock_event( socket, json!({ "type": "response.output_text.delta", "response_id": response_id, - "delta": "hello" + "delta": delta }), ) .await?; @@ -848,6 +985,7 @@ async fn prepare_and_seed_database( database: &SqlDatabaseConfig, upstream_base_url: &str, fixture: ProviderFixture, + redaction: PiiRedaction, ) -> Result<(), BoxError> { let backends = DataBackends::from_config(DataLayerConfig::from_database(database.clone()))?; let pending = backends @@ -862,6 +1000,9 @@ async fn prepare_and_seed_database( seed_models(&backends).await?; let user_id = seed_user(&backends).await?; seed_client_api_key(&backends, &user_id).await?; + if redaction.is_enabled() { + seed_chat_pii_redaction(&backends).await?; + } let candidates = backends .read() @@ -1114,6 +1255,27 @@ async fn seed_client_api_key(backends: &DataBackends, user_id: &str) -> Result<( Ok(()) } +/// 打开 chat PII 脱敏：系统模块开关 + 这把 client key 的 feature 开关。 +/// +/// 规则集刻意不写：缺省即内置规则（含 email 规则），和生产上「只打开开关」的最小 +/// 配置一致。 +async fn seed_chat_pii_redaction(backends: &DataBackends) -> Result<(), BoxError> { + backends + .upsert_system_config_entry("module.chat_pii_redaction.enabled", &json!(true), None) + .await?; + backends + .write() + .auth_api_keys() + .ok_or("auth API key writer unavailable")? + .set_standalone_api_key_feature_settings( + API_KEY_ID, + Some(json!({"chat_pii_redaction": {"enabled": true}})), + ) + .await? + .ok_or("failed to enable chat PII redaction on the E2E API key")?; + Ok(()) +} + fn sha256_hex(value: &str) -> String { let mut hasher = sha2::Sha256::new(); hasher.update(value.as_bytes());