fix(ws): route WS planning and continuation through PII redaction

This commit is contained in:
AAEE86
2026-08-17 14:52:49 +08:00
committed by ZheFox
parent 3b036299d4
commit 9d80281b53
8 changed files with 674 additions and 12 deletions
+3 -2
View File
@@ -56,8 +56,9 @@ pub(crate) use self::planner::{
maybe_build_stream_plan_payload, maybe_build_sync_decision_payload,
maybe_build_sync_plan_payload, planner_is_matching_stream_request, provider_key_pool_score_id,
provider_key_pool_score_scope, read_candidate_transport_snapshot,
record_local_runtime_candidate_skip_reason, resolve_tunnel_scheduler_affinity_context,
resolve_upstream_is_stream_for_provider, set_local_openai_chat_execution_exhausted_diagnostic,
record_local_runtime_candidate_skip_reason, resolve_provider_chat_pii_redaction,
resolve_tunnel_scheduler_affinity_context, resolve_upstream_is_stream_for_provider,
set_local_openai_chat_execution_exhausted_diagnostic,
set_local_openai_image_execution_exhausted_diagnostic, validate_final_openai_provider_request,
CandidateFailureDiagnostic, CandidateFailureDiagnosticKind, EligibleLocalExecutionCandidate,
GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot, LocalExecutionAttemptSource,
@@ -49,6 +49,7 @@ pub(crate) use self::plan_builders::{
pub(crate) use self::pool_scores::{
build_provider_key_pool_score_upsert, provider_key_pool_score_id, provider_key_pool_score_scope,
};
pub(crate) use self::redaction::resolve_provider_chat_pii_redaction;
pub(crate) use self::request_gzip::resolve_transport_request_encoding_policy;
pub(crate) use self::route::is_matching_stream_request as planner_is_matching_stream_request;
pub(crate) use self::runtime_miss::{
@@ -13,6 +13,7 @@ use super::lifecycle::{
send_responses_websocket_turn_start_error, ActiveResponsesWebSocketTurn,
};
use super::quota::{mark_active_response_retry_unsafe, send_previous_response_not_found};
use super::redaction::redact_responses_websocket_client_event;
use super::request::{
build_planning_parts, changed_followup_response_create_model,
continuation_requires_same_upstream, normalize_followup_response_create,
@@ -173,6 +174,48 @@ pub(super) async fn forward_client_message(
.await;
return RelayDisposition::Continue;
};
// 这一轮的 planning Parts 只构造一次(它携带 per-turn 的
// RedactionSessionSlot),并且客户端事件也只在这里脱敏一次:
// 复用已绑定 upstream 的 continuation 根本不进 planner,只靠 planner
// 内部脱敏拦不住它。之后 re-plan / continuation / 配额重试都只看脱敏
// 后的事件,上游请求体与审计 original_request_body 因此一致。
let planning_parts = build_planning_parts(context);
let redacted_client_event = redact_responses_websocket_client_event(
state,
&planning_parts,
&context.decision,
&client_event,
)
.await;
let client_event = match redacted_client_event {
Ok(Some(redacted)) => redacted,
Ok(None) => client_event,
Err(error) => {
warn!(
event_name = "responses_websocket_followup_redaction_failed",
log_type = "ops",
transport = WEBSOCKET_LOG_TRANSPORT,
websocket = true,
trace_id = %context.trace_id,
error = ?error,
"gateway could not apply chat PII redaction to a Responses WebSocket turn"
);
send_gateway_error_with_status(
client_socket,
500,
"responses_websocket_redaction_unavailable",
"Gateway could not apply the configured PII redaction",
)
.await;
close_client_socket(
client_socket,
CLOSE_INTERNAL_ERROR,
"responses_websocket_redaction_unavailable",
)
.await;
return RelayDisposition::Close;
}
};
if bound.upstream.is_none() {
if response_create_has_previous_response_id(&client_event) {
send_previous_response_not_found(client_socket).await;
@@ -199,6 +242,7 @@ pub(super) async fn forward_client_message(
client_socket,
state,
context,
&planning_parts,
client_event,
requested_model,
)
@@ -223,6 +267,7 @@ pub(super) async fn forward_client_message(
client_socket,
state,
context,
&planning_parts,
client_event,
requested_model,
)
@@ -234,6 +279,7 @@ pub(super) async fn forward_client_message(
client_socket,
state,
context,
&planning_parts,
client_event,
bound.client_model.clone(),
)
@@ -295,7 +341,6 @@ pub(super) async fn forward_client_message(
&logical_turn_id,
1,
);
let planning_parts = build_planning_parts(context);
let mut turn = match begin_responses_websocket_turn(
state,
&planning_parts,
@@ -398,15 +443,20 @@ pub(super) async fn forward_client_message(
}
}
/// 重新规划一轮 `response.create`(换模型或独立轮)。
///
/// `planning_parts` 与 `client_event` 都由调用方准备:事件已经过请求侧脱敏,
/// Parts 携带这一轮的 `RedactionSessionSlot`,所以 planner 里的候选级脱敏对
/// 已脱敏内容是幂等的 no-op,上游请求体与审计 body 都保持脱敏态。
async fn forward_replanned_response_create(
bound: &mut BoundResponsesConnection,
client_socket: &mut WebSocket,
state: &AppState,
context: &WebSocketRequestContext,
planning_parts: &http::request::Parts,
client_event: Value,
requested_model: String,
) -> RelayDisposition {
let planning_parts = build_planning_parts(context);
let client_event_text = match serde_json::to_vec(&client_event) {
Ok(value) => Bytes::from(value),
Err(_) => {
@@ -475,7 +525,7 @@ async fn forward_replanned_response_create(
(!excluded_codex_account_ids.is_empty()).then_some(&excluded_codex_account_ids);
let planned = match maybe_build_responses_websocket_decision(
state,
&planning_parts,
planning_parts,
&turn_request_id,
&context.decision,
&client_event,
@@ -569,7 +619,7 @@ async fn forward_replanned_response_create(
);
let mut turn = match begin_responses_websocket_turn(
state,
&planning_parts,
planning_parts,
&context.decision,
turn_decision,
&client_event,
@@ -15,6 +15,7 @@ mod connection;
mod frame;
mod lifecycle;
mod quota;
mod redaction;
mod relay_policy;
mod request;
mod session;
@@ -0,0 +1,480 @@
//! Responses WebSocket 请求侧 PII 脱敏。
//!
//! HTTP 路径在前门建 `RedactionSessionSlot` 并塞进 `parts.extensions`planner
//! 只有拿到这个 slot 才会脱敏。WS 的 planning Parts 是合成的:四个规划入口
//! (首轮、换模型 re-plan、独立轮、配额透明重试)靠 `build_planning_parts` 注入
//! slot 就能复用 planner 的脱敏;但复用已绑定 upstream 的 continuation 根本不进
//! planner,必须在这里先把客户端事件脱敏,再交给协议归一化、上游发送和审计。
//!
//! 因此约定:**进入任何下游用途之前,客户端 `response.create` 只在这里脱敏一次**
//! 之后所有路径都只看脱敏后的事件。
use serde_json::Value;
use crate::ai_serving::{
resolve_local_decision_execution_runtime_auth_context, resolve_provider_chat_pii_redaction,
};
use crate::control::GatewayControlDecision;
use crate::{AppState, GatewayError};
/// Responses WebSocket 只承载 `openai:responses`,脱敏规则按这个客户端格式选取。
const RESPONSES_WEBSOCKET_CLIENT_API_FORMAT: &str = "openai:responses";
/// WS 在选出候选之前就要脱敏,所以脱敏 session 先记在这个固定 key 下。
///
/// slot 是 per-turn 的(见 `build_planning_parts`),这一轮之后即随 slot 一起丢弃;
/// planner 后续用真实 candidate_id 再取一次配置时,body 已是脱敏态、不会重复写入。
const WEBSOCKET_TURN_REDACTION_CANDIDATE_ID: &str = "responses_websocket_turn";
/// 对一条客户端 `response.create` 做请求侧脱敏。
///
/// 返回 `Some(脱敏后的事件)` 仅当脱敏真正命中;`None` 表示未启用或没有命中,调用方
/// 继续用原事件即可(避免未开启脱敏时多一次整包 clone)。
///
/// 脱敏只改写 `instructions` / `input`(见 `privacy::mask_openai_responses_request_value`),
/// `type` / `model` / `previous_response_id` / `generate` 等协议字段原样保留,所以脱敏后的
/// 事件仍可直接用于协议归一化和上游发送。
///
/// 出错必须让这一轮失败:脱敏已启用却读不到配置或加密密钥时,把原文发上游就是
/// 静默旁路,正是本次要修的问题。
pub(super) async fn redact_responses_websocket_client_event(
state: &AppState,
parts: &http::request::Parts,
control_decision: &GatewayControlDecision,
client_event: &Value,
) -> Result<Option<Value>, GatewayError> {
let Some(auth_context) =
resolve_local_decision_execution_runtime_auth_context(control_decision)
else {
return Ok(None);
};
let redaction = resolve_provider_chat_pii_redaction(
state,
parts,
client_event,
&auth_context,
RESPONSES_WEBSOCKET_CLIENT_API_FORMAT,
WEBSOCKET_TURN_REDACTION_CANDIDATE_ID,
)
.await?;
if !redaction.redacted {
return Ok(None);
}
Ok(Some(redaction.body_json.into_owned()))
}
#[cfg(test)]
mod tests {
use std::net::SocketAddr;
use std::sync::Arc;
use aether_crypto::DEVELOPMENT_ENCRYPTION_KEY;
use aether_data::repository::auth::{
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeyExportRecord,
};
use axum::http::{HeaderMap, Uri};
use serde_json::{json, Value};
use super::super::request::{
build_planning_parts, normalize_followup_response_create, planned_response_create_event,
};
use super::super::state::ActiveResponsesWebSocketRequest;
use super::super::turn::prepare_responses_websocket_turn_decision;
use super::redact_responses_websocket_client_event;
use crate::ai_serving::{AiExecutionDecision, ResponsesWebSocketBodyNormalization};
use crate::control::{GatewayControlAuthContext, GatewayControlDecision};
use crate::handlers::proxy::websocket::ingress::WebSocketRequestContext;
use crate::AppState;
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";
fn auth_export_record() -> StoredAuthApiKeyExportRecord {
StoredAuthApiKeyExportRecord::new(
TEST_USER_ID.to_string(),
TEST_API_KEY_ID.to_string(),
"hash-responses-ws-redaction".to_string(),
None,
Some("ws".to_string()),
None,
None,
None,
None,
None,
None,
true,
None,
false,
0,
0,
0.0,
false,
)
.expect("auth api key export record should build")
.with_feature_settings(Some(json!({
"chat_pii_redaction": {"enabled": true}
})))
}
/// 只装脱敏真正需要的东西:系统配置开关 + 规则、加密密钥、带 feature settings
/// 的 API Key 导出记录。候选/上游都不需要,这条链路在 planner 之前。
fn redaction_enabled_state() -> AppState {
let auth_repository = Arc::new(
InMemoryAuthApiKeySnapshotRepository::seed(vec![])
.with_export_records(vec![auth_export_record()]),
);
let data_state =
crate::data::GatewayDataState::with_auth_api_key_reader_for_tests(auth_repository)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY)
.with_system_config_values_for_tests(vec![
("module.chat_pii_redaction.enabled".to_string(), json!(true)),
(
"module.chat_pii_redaction.rules".to_string(),
json!([{
"id": "email",
"name": "邮箱",
"pattern": r"(?i)[A-Z0-9._%+-]{1,64}@[A-Z0-9.-]{1,253}\.[A-Z]{2,63}",
"enabled": true,
"features": {"validator": "email"},
"system": true
}]),
),
(
"module.chat_pii_redaction.cache_ttl_seconds".to_string(),
json!(300),
),
]);
AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(data_state)
}
fn control_decision() -> GatewayControlDecision {
let mut decision = GatewayControlDecision::synthetic(
"/v1/responses".to_string(),
Some("ai_public".to_string()),
Some("openai".to_string()),
Some("responses_websocket".to_string()),
Some("openai:responses".to_string()),
);
decision.auth_context = Some(GatewayControlAuthContext {
user_id: TEST_USER_ID.to_string(),
api_key_id: TEST_API_KEY_ID.to_string(),
username: Some("ws".to_string()),
api_key_name: Some("ws".to_string()),
balance_remaining: None,
access_allowed: true,
user_rate_limit: None,
api_key_rate_limit: None,
api_key_is_standalone: false,
admin_bypass_limits: false,
local_rejection: None,
allowed_models: None,
ip_rules: None,
});
decision
}
fn websocket_context(decision: GatewayControlDecision) -> WebSocketRequestContext {
WebSocketRequestContext {
trace_id: "trace-responses-ws-redaction".to_string(),
headers: HeaderMap::new(),
uri: Uri::from_static("/v1/responses"),
remote_addr: "127.0.0.1:65000"
.parse::<SocketAddr>()
.expect("remote address should parse"),
decision,
rpm_bypassed: false,
websocket_connection_permit: None,
}
}
fn client_event() -> Value {
json!({
"type": "response.create",
"model": "public-model",
"previous_response_id": "resp-previous",
"generate": false,
"input": [{
"role": "user",
"content": [{"type": "input_text", "text": format!("mail {TEST_EMAIL}")}]
}]
})
}
#[tokio::test]
async fn websocket_client_event_is_redacted_without_losing_protocol_fields() {
let state = redaction_enabled_state();
let context = websocket_context(control_decision());
let parts = build_planning_parts(&context);
let event = client_event();
let redacted =
redact_responses_websocket_client_event(&state, &parts, &context.decision, &event)
.await
.expect("redaction should resolve")
.expect("an email in the request should be redacted");
let serialized = serde_json::to_string(&redacted).expect("event should serialize");
assert!(!serialized.contains(TEST_EMAIL), "{serialized}");
assert!(serialized.contains("<AETHER:EMAIL:"), "{serialized}");
// 协议字段必须原样保留,否则 continuation 链路会断。
assert_eq!(redacted["type"], "response.create");
assert_eq!(redacted["model"], "public-model");
assert_eq!(redacted["previous_response_id"], "resp-previous");
assert_eq!(redacted["generate"], false);
}
#[tokio::test]
async fn redacting_an_already_redacted_event_is_a_no_op() {
// re-plan 与配额重试路径会把已脱敏的事件再交给 planner,planner 内部会对
// 同一个 body 再跑一遍 mask。占位符本身不该被任何规则命中,否则会被二次
// 替换、破坏与上游已有 previous_response_id 链的一致性。
let state = redaction_enabled_state();
let context = websocket_context(control_decision());
let parts = build_planning_parts(&context);
let event = client_event();
let redacted =
redact_responses_websocket_client_event(&state, &parts, &context.decision, &event)
.await
.expect("redaction should resolve")
.expect("an email in the request should be redacted");
// 复用同一个 parts/slot,和 re-plan 在同一 turn 内二次脱敏的情形一致。
let second_pass =
redact_responses_websocket_client_event(&state, &parts, &context.decision, &redacted)
.await
.expect("second redaction pass should resolve");
assert!(
second_pass.is_none(),
"already redacted event should stay byte-identical: {second_pass:?}"
);
}
#[tokio::test]
async fn redaction_is_skipped_without_a_local_auth_context() {
let state = redaction_enabled_state();
let mut decision = control_decision();
decision.auth_context = None;
let context = websocket_context(decision);
let parts = build_planning_parts(&context);
let event = client_event();
let redacted =
redact_responses_websocket_client_event(&state, &parts, &context.decision, &event)
.await
.expect("redaction should resolve");
assert!(redacted.is_none());
}
/// 真跑一遍脱敏,拿到这一轮的「生效事件」。
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)
.await
.expect("redaction should resolve")
.expect("an email in the request should be redacted")
}
/// 只有 `action` 没有 serde 默认值,其余字段都能省略。
fn decision_template(
provider_request_body: Value,
report_context: Value,
) -> AiExecutionDecision {
serde_json::from_value(json!({
"action": "local",
"candidate_id": "candidate-responses-ws",
"provider_request_body": provider_request_body,
"report_context": report_context,
}))
.expect("decision template should deserialize")
}
/// planner 在脱敏 body 上做模型映射后的 provider body。
fn provider_body_from(effective_event: &Value) -> Value {
let mut provider_body = effective_event.clone();
provider_body["model"] = json!("provider-model");
provider_body
}
/// 绑定那一轮留下的 report_context seed:故意带上原始 PII,用来证明这一轮
/// 会用脱敏后的 body 覆盖它,而不是把原文带进审计。
fn seed_report_context_with_raw_pii() -> Value {
json!({
"request_id": "connection",
"candidate_id": "candidate-responses-ws",
"original_request_body": {
"type": "response.create",
"model": "public-model",
"input": format!("mail {TEST_EMAIL}")
}
})
}
fn assert_redacted_json(value: &Value, label: &str) {
let serialized = serde_json::to_string(value).expect("value should serialize");
assert!(
!serialized.contains(TEST_EMAIL),
"{label} must not carry raw PII: {serialized}"
);
assert!(
serialized.contains("<AETHER:EMAIL:"),
"{label} must carry the redaction sentinel: {serialized}"
);
}
#[tokio::test]
async fn first_turn_upstream_and_audit_bodies_are_redacted() {
let state = redaction_enabled_state();
let decision = control_decision();
let effective_event = redacted_client_event(&state, &decision).await;
let template = decision_template(
provider_body_from(&effective_event),
seed_report_context_with_raw_pii(),
);
// 首轮实际发上游的事件由 decision.provider_request_body 派生。
let provider_event: Value = serde_json::from_str(
&planned_response_create_event(&template, &effective_event)
.expect("first provider event should serialize"),
)
.expect("first provider event should parse");
let turn_decision = prepare_responses_websocket_turn_decision(
&template,
"turn-1".to_string(),
true,
&effective_event,
&provider_event,
"connection",
1,
"logical-turn-1",
1,
);
assert_redacted_json(&provider_event, "first turn upstream event");
assert_redacted_json(
turn_decision
.provider_request_body
.as_ref()
.expect("turn decision should carry a provider body"),
"first turn provider request body",
);
let report_context = turn_decision
.report_context
.as_ref()
.expect("turn decision should carry a report context");
assert_redacted_json(
&report_context["original_request_body"],
"first turn audit body",
);
// 整个 report_context 都不该残留原文(seed 里的原始 body 必须被覆盖)。
assert_redacted_json(report_context, "first turn report context");
assert_eq!(provider_event["type"], "response.create");
assert_eq!(provider_event["model"], "provider-model");
}
#[tokio::test]
async fn continuation_upstream_and_audit_bodies_are_redacted() {
let state = redaction_enabled_state();
let decision = control_decision();
let effective_event = redacted_client_event(&state, &decision).await;
// continuation 复用已绑定的 upstream:不再规划,直接重放归一化器。
let outbound = normalize_followup_response_create(
&effective_event,
"provider-model",
&ResponsesWebSocketBodyNormalization::for_tests("provider-model"),
)
.expect("continuation should normalize");
let provider_event: Value =
serde_json::from_str(&outbound).expect("continuation event should parse");
let template = decision_template(
provider_body_from(&effective_event),
seed_report_context_with_raw_pii(),
);
let turn_decision = prepare_responses_websocket_turn_decision(
&template,
"turn-2".to_string(),
false,
&effective_event,
&provider_event,
"connection",
2,
"logical-turn-2",
1,
);
assert!(
!outbound.contains(TEST_EMAIL),
"continuation upstream frame must not carry raw PII: {outbound}"
);
assert!(
outbound.contains("<AETHER:EMAIL:"),
"continuation upstream frame must carry the sentinel: {outbound}"
);
assert_eq!(provider_event["previous_response_id"], "resp-previous");
let report_context = turn_decision
.report_context
.as_ref()
.expect("turn decision should carry a report context");
assert_redacted_json(
&report_context["original_request_body"],
"continuation audit body",
);
assert_redacted_json(report_context, "continuation report context");
}
#[tokio::test]
async fn quota_retry_replays_the_redacted_event() {
let state = redaction_enabled_state();
let decision = control_decision();
let effective_event = redacted_client_event(&state, &decision).await;
// 配额透明重试重放 active_response_create 里保存的事件,所以保存的必须
// 已经是脱敏版,否则重试会把原文发给新的上游账号。
let active = ActiveResponsesWebSocketRequest::new(
effective_event.clone(),
2,
"logical-turn-2".to_string(),
);
assert_redacted_json(&active.client_event, "quota retry replay event");
let template = decision_template(
provider_body_from(&active.client_event),
seed_report_context_with_raw_pii(),
);
let provider_event: Value = serde_json::from_str(
&planned_response_create_event(&template, &active.client_event)
.expect("retry provider event should serialize"),
)
.expect("retry provider event should parse");
let turn_decision = prepare_responses_websocket_turn_decision(
&template,
"turn-2-retry".to_string(),
true,
&active.client_event,
&provider_event,
"connection",
active.turn_index,
"logical-turn-2",
2,
);
assert_redacted_json(&provider_event, "quota retry upstream event");
let report_context = turn_decision
.report_context
.as_ref()
.expect("turn decision should carry a report context");
assert_eq!(report_context["websocket_turn_attempt"], 2);
assert_redacted_json(
&report_context["original_request_body"],
"quota retry audit body",
);
assert_redacted_json(report_context, "quota retry report context");
}
}
@@ -11,7 +11,14 @@ use serde_json::Value;
use crate::ai_serving::{AiExecutionDecision, ResponsesWebSocketBodyNormalization};
use crate::handlers::proxy::websocket::ingress::WebSocketRequestContext;
use crate::headers::request_origin_from_headers_and_remote_addr;
use crate::privacy::RedactionSessionSlot;
/// 把一条 WebSocket turn 还原成 planner 需要的 HTTP 形状请求头部。
///
/// 这里必须和 HTTP 前门(`handlers/proxy/mod.rs`)保持同一份 extension 契约:
/// planner 只在 `parts.extensions` 里拿到 `RedactionSessionSlot` 时才做请求脱敏
/// `ai_serving/planner/redaction.rs`),少插这一项等于整条 WS 链路静默绕过
/// 已启用的 PII 脱敏。
pub(super) fn build_planning_parts(context: &WebSocketRequestContext) -> http::request::Parts {
let mut request = http::Request::builder()
.method(Method::POST)
@@ -40,6 +47,15 @@ pub(super) fn build_planning_parts(context: &WebSocketRequestContext) -> http::r
&context.headers,
&context.remote_addr,
));
// slot 必须每个 turn 新建,不能按连接复用:planner 侧的请求脱敏缓存键是
// `{format:?}:{body_json 指针地址}``ai_serving/planner/redaction.rs:169`),
// 连接级复用同一个 slot 时,上一轮 client_event 释放后这一轮的 `Value` 很可能
// 落在同一地址,会命中上一轮缓存,把上一轮的脱敏 body 当成这一轮的发出去。
// 每个 `response.create` 本身就是独立计费/审计请求,per-turn 也正好对应
// HTTP 前门「一个请求一个 slot」的语义。
request
.extensions_mut()
.insert(RedactionSessionSlot::default());
request.into_parts().0
}
@@ -198,10 +214,66 @@ pub(super) fn normalize_followup_response_create(
#[cfg(test)]
mod tests {
use std::net::SocketAddr;
use axum::http::{HeaderMap, Uri};
use serde_json::json;
use super::{normalize_followup_response_create, response_create_has_previous_response_id};
use super::{
build_planning_parts, normalize_followup_response_create,
response_create_has_previous_response_id,
};
use crate::ai_serving::ResponsesWebSocketBodyNormalization;
use crate::control::GatewayControlDecision;
use crate::handlers::proxy::websocket::ingress::WebSocketRequestContext;
use crate::privacy::RedactionSessionSlot;
fn websocket_context() -> WebSocketRequestContext {
WebSocketRequestContext {
trace_id: "trace-planning-parts".to_string(),
headers: HeaderMap::new(),
uri: Uri::from_static("/v1/responses"),
remote_addr: "127.0.0.1:65001"
.parse::<SocketAddr>()
.expect("remote address should parse"),
decision: GatewayControlDecision::synthetic(
"/v1/responses".to_string(),
Some("ai_public".to_string()),
Some("openai".to_string()),
Some("responses_websocket".to_string()),
Some("openai:responses".to_string()),
),
rpm_bypassed: false,
websocket_connection_permit: None,
}
}
#[test]
fn planning_parts_carry_a_fresh_redaction_session_slot_per_turn() {
// 没有这个 extensionplanner 会静默跳过已启用的 PII 脱敏
// ai_serving/planner/redaction.rs),整条 WS 链路都按原文发上游。
let context = websocket_context();
let first = build_planning_parts(&context);
let second = build_planning_parts(&context);
let first_slot = first
.extensions
.get::<RedactionSessionSlot>()
.expect("planning parts must carry a redaction session slot");
let second_slot = second
.extensions
.get::<RedactionSessionSlot>()
.expect("planning parts must carry a redaction session slot");
// 每轮必须是独立 slot:slot 内的请求缓存以 body 指针地址为键,跨轮共享会
// 命中上一轮缓存。用缓存条目相互不可见来证明两者不是同一个 slot。
first_slot.put_cached_request_redaction(
"turn-1",
crate::privacy::CachedRequestRedaction::unredacted(),
);
assert!(first_slot.cached_request_redaction("turn-1").is_some());
assert!(second_slot.cached_request_redaction("turn-1").is_none());
}
fn normalized_continuation(
event: &serde_json::Value,
@@ -22,6 +22,7 @@ use super::lifecycle::{
await_turn_finalization_handle, finalize_unbound_turn, responses_websocket_turn_start_close,
send_responses_websocket_turn_start_error, ActiveResponsesWebSocketTurn,
};
use super::redaction::redact_responses_websocket_client_event;
use super::request::{build_planning_parts, planned_response_create_event};
use super::state::ActiveResponsesWebSocketRequest;
use super::turn::{
@@ -216,6 +217,46 @@ pub(super) async fn run_responses_websocket(
}
}
// 请求侧脱敏必须在规划之前完成,而且这一轮只在这里做一次:planner 会把这份
// body 写进 upstream 请求体和审计 original_request_body,绑定上游的首条
// response.create 也从它派生。脱敏失败时直接断开,绝不退回原文发上游。
let redacted_first_event = redact_responses_websocket_client_event(
&state,
&planning_parts,
&context.decision,
&first_event,
)
.await;
let first_event = match redacted_first_event {
Ok(Some(redacted)) => redacted,
Ok(None) => first_event,
Err(error) => {
warn!(
event_name = "responses_websocket_redaction_failed",
log_type = "ops",
transport = WEBSOCKET_LOG_TRANSPORT,
websocket = true,
trace_id = %context.trace_id,
error = ?error,
"gateway could not apply chat PII redaction to the initial Responses WebSocket event"
);
send_gateway_error_with_status(
&mut client_socket,
500,
"responses_websocket_redaction_unavailable",
"Gateway could not apply the configured PII redaction",
)
.await;
close_client_socket(
&mut client_socket,
CLOSE_INTERNAL_ERROR,
"responses_websocket_redaction_unavailable",
)
.await;
return;
}
};
let planned = match maybe_build_responses_websocket_decision(
&state,
&planning_parts,
@@ -300,11 +300,18 @@ pub(super) struct ResponsesWebSocketTurn {
terminal_error_body: Option<String>,
}
/// 组装一轮 turn 的 decision。
///
/// `effective_client_event` 必须是**已经过请求侧脱敏**的客户端事件(见
/// `super::redaction`),`provider_event` 由它派生。审计里的 `original_request_body`
/// 直接用它覆盖 seedcontinuation 的 seed 来自绑定那一轮的 report_context,不覆盖
/// 会记成上一轮的 body;但覆盖成 raw 事件又会把已脱敏的审计内容换回原文,等于
/// 脱敏在审计侧失效。
pub(super) fn prepare_responses_websocket_turn_decision(
template: &AiExecutionDecision,
request_id: String,
reuse_selected_candidate: bool,
client_event: &Value,
effective_client_event: &Value,
provider_event: &Value,
connection_trace_id: &str,
turn_index: u64,
@@ -322,7 +329,7 @@ pub(super) fn prepare_responses_websocket_turn_decision(
decision.report_context.take(),
request_id.as_str(),
reuse_selected_candidate,
client_event,
effective_client_event,
provider_event,
connection_trace_id,
turn_index,
@@ -1003,11 +1010,17 @@ pub(super) async fn spawn_responses_websocket_turn_finalization(
})
}
/// 把一轮 turn 的事实写进审计/用量 report context。
///
/// `effective_client_event` 是脱敏后的客户端事件(未启用脱敏时就是原事件)。
/// HTTP 路径的约定是「脱敏生效时审计记录脱敏后的 body」
/// `ai_serving/planner/standard/openai/responses/decision/payload.rs`),
/// WS 这里必须保持一致,否则上游收到的是脱敏内容、审计里却留着原始 PII。
fn prepare_websocket_report_context(
report_context: Option<Value>,
request_id: &str,
reuse_selected_candidate: bool,
client_event: &Value,
effective_client_event: &Value,
provider_event: &Value,
connection_trace_id: &str,
turn_index: u64,
@@ -1040,8 +1053,11 @@ fn prepare_websocket_report_context(
object.remove(field);
}
}
object.insert("original_request_body".to_string(), client_event.clone());
if let Some(model) = client_event
object.insert(
"original_request_body".to_string(),
effective_client_event.clone(),
);
if let Some(model) = effective_client_event
.get("model")
.and_then(Value::as_str)
.map(str::trim)