mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-09 04:30:20 +08:00
fix(ws): route WS planning and continuation through PII redaction
This commit is contained in:
@@ -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() {
|
||||
// 没有这个 extension,planner 会静默跳过已启用的 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`
|
||||
/// 直接用它覆盖 seed:continuation 的 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)
|
||||
|
||||
Reference in New Issue
Block a user