i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-12 04:09:48 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(security): harden gateway boundaries and usage policies

Browse Source 
Consolidate subscription usage policy enforcement, privacy-safe persistence, and gateway security hardening into one reviewable change.

Includes bounded HTTP and execution envelopes, header and protocol guards, DNS and relay validation, authentication and secret projection hardening, secure backup/install paths, and regression coverage.
This commit is contained in:
elky
2026-09-04 03:45:52 +08:00
parent ddcbeb3ae9
commit 579f2c7cc1
1019 changed files with 190437 additions and 26080 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/handlers/proxy/websocket/responses/client.rs
+178 -60
View File
@@ -17,6 +17,9 @@ use super::ownership::{
await_owned_responses_websocket_plan, begin_responses_websocket_turn_with_planned_lease,
spawn_owned_responses_websocket_plan, OwnedResponsesWebSocketDecision,
};
use super::plan_admission::{
acquire_responses_websocket_plan_admission, send_responses_websocket_plan_admission_error,
};
use super::quota::mark_active_response_retry_unsafe;
use super::redaction::redact_responses_websocket_client_event_with_reasoning_replay_policy;
use super::request::{
@@ -35,6 +38,8 @@ use super::turn::{
use super::turn_state::LogicalTurn;
use super::upstream::{
bind_responses_upstream, decision_bound_upstream_change_fields, decision_reuses_bound_upstream,
send_responses_websocket_upstream_message, ResponsesWebSocketUpstreamBindError,
ResponsesWebSocketUpstreamSendError,
};
use crate::ai_serving::ResponsesWebSocketPinnedCandidate;
use crate::clock::current_unix_secs;
@@ -44,8 +49,9 @@ use crate::handlers::proxy::websocket::session::{CLOSE_INTERNAL_ERROR, WEBSOCKET
use crate::handlers::proxy::websocket::transport::{
close_client_socket, close_upstream_socket, send_client_message, send_gateway_error,
send_gateway_error_with_status, send_gateway_error_with_stream_id,
send_responses_websocket_error_with_param, send_upstream_message,
send_responses_websocket_error_with_param,
};
use crate::plan_usage_policy::PlanUsagePolicySnapshot;
use crate::privacy::RedactionSession;
use crate::rate_limit::FrontdoorUserRpmOutcome;
use crate::AppState;
@@ -69,6 +75,7 @@ pub(super) enum RelayDisposition {
Continue,
Close,
UpstreamError(&'static str),
PlanUsagePermitLost,
}
pub(super) fn adapter_drain_ready(
@@ -317,6 +324,37 @@ pub(super) async fn forward_client_message(
}
}
// Subscription-plan concurrency is admitted per logical
// response.create. Keep the permit and policy snapshot attached
// to this turn through any transparent provider rebind.
let plan_usage_admission = match acquire_responses_websocket_plan_admission(
state,
&turn_control.decision,
&logical_turn_id,
)
.await
{
Ok(admission) => admission,
Err(error) => {
warn!(
event_name = "responses_websocket_followup_plan_usage_rejected",
log_type = "event",
transport = WEBSOCKET_LOG_TRANSPORT,
websocket = true,
trace_id = %context.trace_id,
logical_turn_id = %logical_turn_id,
error = ?error,
"gateway rejected a Responses WebSocket follow-up at its subscription plan limit"
);
send_responses_websocket_plan_admission_error(client_socket, &error).await;
return RelayDisposition::Continue;
}
};
let crate::plan_usage_policy::PlanUsageAdmission {
permit: plan_usage_permit,
policy_snapshot: plan_usage_policy_snapshot,
} = plan_usage_admission;
let response_id_owned_by_connection = client_event
.get("previous_response_id")
.and_then(Value::as_str)
@@ -465,6 +503,8 @@ pub(super) async fn forward_client_message(
codex_fingerprint_context,
turn_control,
turn_redaction_session,
plan_usage_permit,
plan_usage_policy_snapshot,
)
.await;
}
@@ -482,6 +522,8 @@ pub(super) async fn forward_client_message(
raw_responses_lite_static_config
.expect("independent turns always retain their raw static config"),
turn_redaction_session,
plan_usage_permit,
plan_usage_policy_snapshot,
)
.await
}
@@ -529,6 +571,8 @@ async fn forward_pinned_continuation(
codex_fingerprint_context: CodexFingerprintConvergenceContext,
turn_control: ResponsesWebSocketTurnControl,
turn_redaction_session: Option<RedactionSession>,
plan_usage_permit: Option<aether_runtime::AdmissionPermit>,
plan_usage_policy_snapshot: Option<PlanUsagePolicySnapshot>,
) -> RelayDisposition {
let Some(pinned_candidate) =
ResponsesWebSocketPinnedCandidate::from_decision(&bound.decision_template)
@@ -725,6 +769,7 @@ async fn forward_pinned_continuation(
&turn_control.decision,
turn_decision,
&client_event,
plan_usage_policy_snapshot.clone(),
planned_lease,
)
.await
@@ -755,21 +800,42 @@ async fn forward_pinned_continuation(
.await;
return RelayDisposition::UpstreamError("responses_websocket_send_failed");
};
if send_upstream_message(upstream, WreqWsMessage::text(outbound))
.await
.is_err()
match send_responses_websocket_upstream_message(
upstream,
WreqWsMessage::text(outbound),
plan_usage_permit.as_ref(),
|request_state| turn.record_upstream_request_state(request_state),
)
.await
{
queue_turn_finalization(
bound,
state,
turn,
ResponsesWebSocketTurnOutcome::upstream_send_failed(),
)
.await;
return RelayDisposition::UpstreamError("responses_websocket_send_failed");
Ok(()) => {}
Err(ResponsesWebSocketUpstreamSendError::PlanUsagePermitLost) => {
turn.release_plan_usage_cost_before_upstream_send(
state,
"continuation_plan_usage_permit_lost",
)
.await;
queue_turn_finalization(
bound,
state,
turn,
ResponsesWebSocketTurnOutcome::connection_admission_lost(),
)
.await;
return RelayDisposition::PlanUsagePermitLost;
}
Err(ResponsesWebSocketUpstreamSendError::Transport(_)) => {
queue_turn_finalization(
bound,
state,
turn,
ResponsesWebSocketTurnOutcome::upstream_send_failed(),
)
.await;
return RelayDisposition::UpstreamError("responses_websocket_send_failed");
}
}
turn.mark_upstream_request_sent();
turn.set_provider_response_headers(bound.upstream_response_headers.clone());
if let Some(session) = turn_redaction_session {
bound.redaction_restorer.register(session);
@@ -781,7 +847,9 @@ async fn forward_pinned_continuation(
LogicalTurn::new(client_event, turn_index, logical_turn_id)
.with_codex_fingerprint_context(codex_fingerprint_context)
.with_provider_store(provider_event.get("store") == Some(&Value::Bool(true)))
.with_turn_control(turn_control),
.with_turn_control(turn_control)
.with_plan_usage_permit(plan_usage_permit)
.with_plan_usage_policy_snapshot(plan_usage_policy_snapshot),
turn,
);
bound.next_turn_index = bound.next_turn_index.saturating_add(1);
@@ -819,6 +887,8 @@ async fn forward_replanned_response_create(
turn_control: ResponsesWebSocketTurnControl,
raw_responses_lite_static_config: ResponsesLiteStaticConfig,
turn_redaction_session: Option<RedactionSession>,
plan_usage_permit: Option<aether_runtime::AdmissionPermit>,
plan_usage_policy_snapshot: Option<PlanUsagePolicySnapshot>,
) -> RelayDisposition {
let turn_request_id = Uuid::new_v4().to_string();
let now_unix_secs = current_unix_secs();
@@ -919,6 +989,7 @@ async fn forward_replanned_response_create(
&turn_control.decision,
turn_decision,
&client_event,
plan_usage_policy_snapshot.clone(),
planned_lease,
)
.await
@@ -970,18 +1041,40 @@ async fn forward_replanned_response_create(
.await;
return RelayDisposition::UpstreamError("responses_websocket_send_failed");
};
if send_upstream_message(upstream, WreqWsMessage::text(outbound))
.await
.is_err()
match send_responses_websocket_upstream_message(
upstream,
WreqWsMessage::text(outbound),
plan_usage_permit.as_ref(),
|request_state| turn.record_upstream_request_state(request_state),
)
.await
{
queue_turn_finalization(
bound,
state,
turn,
ResponsesWebSocketTurnOutcome::upstream_send_failed(),
)
.await;
return RelayDisposition::UpstreamError("responses_websocket_send_failed");
Ok(()) => {}
Err(ResponsesWebSocketUpstreamSendError::PlanUsagePermitLost) => {
turn.release_plan_usage_cost_before_upstream_send(
state,
"replanned_reuse_plan_usage_permit_lost",
)
.await;
queue_turn_finalization(
bound,
state,
turn,
ResponsesWebSocketTurnOutcome::connection_admission_lost(),
)
.await;
return RelayDisposition::PlanUsagePermitLost;
}
Err(ResponsesWebSocketUpstreamSendError::Transport(_)) => {
queue_turn_finalization(
bound,
state,
turn,
ResponsesWebSocketTurnOutcome::upstream_send_failed(),
)
.await;
return RelayDisposition::UpstreamError("responses_websocket_send_failed");
}
}
// A response.create without previous_response_id starts a new chain.
@@ -991,7 +1084,6 @@ async fn forward_replanned_response_create(
bound
.redaction_restorer
.start_new_chain(turn_redaction_session);
turn.mark_upstream_request_sent();
turn.set_provider_response_headers(bound.upstream_response_headers.clone());
let provider_model =
provider_model_from_decision(&decision).unwrap_or_else(|| bound.provider_model.clone());
@@ -1009,7 +1101,9 @@ async fn forward_replanned_response_create(
LogicalTurn::new(client_event.clone(), turn_index, logical_turn_id.clone())
.with_codex_fingerprint_context(codex_fingerprint_context.clone())
.with_provider_store(provider_event.get("store") == Some(&Value::Bool(true)))
.with_turn_control(turn_control),
.with_turn_control(turn_control)
.with_plan_usage_permit(plan_usage_permit)
.with_plan_usage_policy_snapshot(plan_usage_policy_snapshot),
turn,
);
bound.next_turn_index = bound.next_turn_index.saturating_add(1);
@@ -1031,42 +1125,64 @@ async fn forward_replanned_response_create(
return RelayDisposition::Continue;
}
let mut replacement =
match bind_responses_upstream(&decision, normalization, &client_event, adapter).await {
Ok(connection) => connection,
Err(code) => {
queue_turn_finalization(
bound,
state,
turn,
ResponsesWebSocketTurnOutcome::upstream_connect_failed(code),
)
.await;
warn!(
event_name = "responses_websocket_followup_model_rebind_failed",
log_type = "ops",
transport = WEBSOCKET_LOG_TRANSPORT,
websocket = true,
trace_id = %context.trace_id,
requested_model = %requested_model,
error_code = code,
"gateway failed to rebind Responses WebSocket follow-up model"
);
send_gateway_error_with_status(
client_socket,
502,
code,
"Gateway could not establish the requested model",
)
.await;
return RelayDisposition::Continue;
}
};
let mut replacement = match bind_responses_upstream(
&decision,
normalization,
&client_event,
adapter,
plan_usage_permit.as_ref(),
|request_state| turn.record_upstream_request_state(request_state),
)
.await
{
Ok(connection) => connection,
Err(ResponsesWebSocketUpstreamBindError::PlanUsagePermitLost) => {
turn.release_plan_usage_cost_before_upstream_send(
state,
"replanned_rebind_plan_usage_permit_lost",
)
.await;
queue_turn_finalization(
bound,
state,
turn,
ResponsesWebSocketTurnOutcome::connection_admission_lost(),
)
.await;
return RelayDisposition::PlanUsagePermitLost;
}
Err(ResponsesWebSocketUpstreamBindError::Transport(code)) => {
queue_turn_finalization(
bound,
state,
turn,
ResponsesWebSocketTurnOutcome::upstream_connect_failed(code),
)
.await;
warn!(
event_name = "responses_websocket_followup_model_rebind_failed",
log_type = "ops",
transport = WEBSOCKET_LOG_TRANSPORT,
websocket = true,
trace_id = %context.trace_id,
requested_model = %requested_model,
error_code = code,
"gateway failed to rebind Responses WebSocket follow-up model"
);
send_gateway_error_with_status(
client_socket,
502,
code,
"Gateway could not establish the requested model",
)
.await;
return RelayDisposition::Continue;
}
};
if replacement.responses_lite_static_config.is_some() {
replacement.responses_lite_static_config = Some(raw_responses_lite_static_config);
}
turn.mark_upstream_request_sent();
turn.set_provider_response_headers(replacement.upstream_response_headers.clone());
let previous_client_model = bound.client_model.clone();
let previous_provider_model = bound.provider_model.clone();
@@ -1094,7 +1210,9 @@ async fn forward_replanned_response_create(
LogicalTurn::new(client_event, turn_index, logical_turn_id)
.with_codex_fingerprint_context(codex_fingerprint_context)
.with_provider_store(provider_event.get("store") == Some(&Value::Bool(true)))
.with_turn_control(turn_control),
.with_turn_control(turn_control)
.with_plan_usage_permit(plan_usage_permit)
.with_plan_usage_policy_snapshot(plan_usage_policy_snapshot),
turn,
);
bound.next_turn_index = bound.next_turn_index.saturating_add(1);
apps/aether-gateway/src/handlers/proxy/websocket/responses/connection.rs
+65 -27
View File
@@ -17,6 +17,7 @@ use super::lifecycle::{
await_pending_adapter_observation, finalize_active_turn, queue_turn_finalization,
settle_turn_finalization, spawn_bounded_adapter_observation, PreviousAttemptSettled,
};
use super::plan_admission::terminate_responses_websocket_for_plan_permit_loss;
use super::quota::{
detach_exhausted_upstream, is_usage_limit_error_event, mark_active_response_retry_unsafe,
observe_active_response_rebind_safety, retry_active_turn_after_quota_exhaustion,
@@ -144,6 +145,19 @@ pub(super) async fn relay_bound_connection(
).await;
break;
}
RelayDisposition::PlanUsagePermitLost => {
warn!(
event_name = "responses_websocket_plan_usage_concurrency_lost_before_send",
log_type = "ops",
transport = WEBSOCKET_LOG_TRANSPORT,
websocket = true,
trace_id = %context.trace_id,
"gateway stopped a Responses WebSocket turn before its upstream send after the subscription plan concurrency lease became unhealthy"
);
close_bound_upstream(bound).await;
terminate_responses_websocket_for_plan_permit_loss(client_socket).await;
break;
}
RelayDisposition::UpstreamError(code) => {
warn!(
event_name = "responses_websocket_upstream_send_failed",
@@ -469,18 +483,33 @@ pub(super) async fn relay_bound_connection(
)
.await
}
None => PreviousAttemptSettled::nothing_to_settle(),
None => Some(PreviousAttemptSettled::nothing_to_settle()),
};
// Planning and binding a replacement carries the complete
// scheduler/provider state machine. Keep that large future
// off the relay task's stack; the default Tokio/test worker
// stack is otherwise easy to exhaust on this rare branch.
if Box::pin(retry_active_turn_after_quota_exhaustion(
bound, state, context, settled,
))
.await
{
continue;
if let Some(settled) = settled {
match Box::pin(retry_active_turn_after_quota_exhaustion(
bound, state, context, settled,
))
.await
{
Ok(true) => continue,
Ok(false) => {}
Err(()) => {
finalize_active_turn(
bound,
state,
ResponsesWebSocketTurnOutcome::connection_admission_lost(),
)
.await;
close_bound_upstream(bound).await;
terminate_responses_websocket_for_plan_permit_loss(client_socket)
.await;
break;
}
}
}
// 重试失败。旧 attempt 已经结算，logical turn 仍停在
// Replanning，所以后面分支里的 end() / finalize_active_turn
@@ -531,27 +560,32 @@ pub(super) async fn relay_bound_connection(
let mut relay_serialization_failed = false;
match relay_directive {
Some(ResponsesWebSocketRelayDirective::ForwardOriginal) => {
let restored = parsed_upstream_frame
.as_ref()
.and_then(|frame| {
bound
.redaction_restorer
.restore_provider_frame_text(frame.event())
});
let client_frame = match restored {
Some(text) => AxumWsMessage::Text(text.into()),
None => upstream_message_to_client(upstream_message.clone()),
let client_frame = match parsed_upstream_frame.as_ref().map(|frame| {
bound
.redaction_restorer
.restore_provider_frame_text(frame.event())
}) {
Some(Ok(Some(text))) => Some(AxumWsMessage::Text(text.into())),
Some(Ok(None)) | None => {
Some(upstream_message_to_client(upstream_message.clone()))
}
Some(Err(_)) => {
relay_serialization_failed = true;
None
}
};
match send_client_message(client_socket, client_frame).await {
Ok(()) => {
if let (Some(turn), Some(frame)) = (
bound.turn_state.attempt_mut(),
parsed_upstream_frame.as_ref(),
) {
turn.capture_client_frame(frame.event());
if let Some(client_frame) = client_frame {
match send_client_message(client_socket, client_frame).await {
Ok(()) => {
if let (Some(turn), Some(frame)) = (
bound.turn_state.attempt_mut(),
parsed_upstream_frame.as_ref(),
) {
turn.capture_client_frame(frame.event());
}
}
Err(error) => relay_send_error = Some(error),
}
Err(error) => relay_send_error = Some(error),
}
}
Some(ResponsesWebSocketRelayDirective::ForwardEvents(events)) => {
@@ -560,14 +594,18 @@ pub(super) async fn relay_bound_connection(
.redaction_restorer
.restore_provider_frame_text(event)
{
Some(restored) => restored,
None => match encode_opaque_websocket_event(event) {
Ok(Some(restored)) => restored,
Ok(None) => match encode_opaque_websocket_event(event) {
Ok(encoded) => encoded,
Err(_) => {
relay_serialization_failed = true;
break;
}
},
Err(_) => {
relay_serialization_failed = true;
break;
}
};
match send_client_message(
client_socket,
apps/aether-gateway/src/handlers/proxy/websocket/responses/lifecycle.rs
+52 -3
View File
@@ -18,6 +18,7 @@ use crate::handlers::proxy::websocket::session::{
CLOSE_INTERNAL_ERROR, CLOSE_POLICY_VIOLATION, CLOSE_TRY_AGAIN, WEBSOCKET_LOG_TRANSPORT,
};
use crate::handlers::proxy::websocket::transport::send_responses_websocket_error;
use crate::plan_usage_policy::PlanUsagePolicySnapshot;
use crate::{AppState, GatewayError};
const RESPONSES_WEBSOCKET_ADAPTER_OBSERVATION_TIMEOUT: Duration = Duration::from_secs(5);
@@ -85,6 +86,7 @@ pub(super) async fn begin_responses_websocket_turn(
control_decision: &crate::control::GatewayControlDecision,
decision: crate::ai_serving::AiExecutionDecision,
client_event: &serde_json::Value,
plan_usage_policy_snapshot: Option<PlanUsagePolicySnapshot>,
) -> Result<ActiveProviderAttempt, GatewayError> {
let state = state.clone();
let trace_id = trace_id.to_string();
@@ -108,6 +110,7 @@ pub(super) async fn begin_responses_websocket_turn(
&control_decision,
decision,
&client_event,
plan_usage_policy_snapshot,
)
.await?;
Ok(ActiveProviderAttempt::new(&state, turn))
@@ -231,12 +234,29 @@ impl PreviousAttemptSettled {
pub(super) async fn settle_turn_finalization(
bound: &mut BoundResponsesConnection,
state: &AppState,
turn: ActiveProviderAttempt,
mut turn: ActiveProviderAttempt,
outcome: ResponsesWebSocketTurnOutcome,
) -> PreviousAttemptSettled {
) -> Option<PreviousAttemptSettled> {
// This awaited finalization entry is used only by transparent quota retry.
// Free the old attempt's durable cost reservation before planning can
// reserve the replacement, otherwise two per-attempt estimates briefly
// count against the same user and can manufacture a false local 429.
if let Err(error) = turn.release_plan_usage_cost_for_retry(state).await {
warn!(
event_name = "responses_websocket_plan_usage_cost_retry_release_failed",
log_type = "ops",
transport = WEBSOCKET_LOG_TRANSPORT,
websocket = true,
error = ?error,
"gateway stopped a transparent Responses WebSocket retry because the old plan cost reservation could not be released"
);
queue_turn_finalization(bound, state, turn, outcome).await;
await_pending_turn_finalization(bound).await;
return None;
}
queue_turn_finalization(bound, state, turn, outcome).await;
await_pending_turn_finalization(bound).await;
PreviousAttemptSettled(())
Some(PreviousAttemptSettled(()))
}
pub(super) fn spawn_bounded_adapter_observation(
@@ -334,6 +354,16 @@ pub(super) async fn send_responses_websocket_turn_start_error(
) {
let status_code = responses_websocket_turn_start_http_status(error);
match error {
GatewayError::PlanUsageLimited(_) => {
send_responses_websocket_error(
client_socket,
status_code,
"rate_limit_error",
"plan_usage_limit_exceeded",
"Subscription plan usage limit exceeded; retry later",
)
.await;
}
GatewayError::Client { status, message } => {
let (error_type, code) = if status.as_u16() == 429 {
("rate_limit_error", "gateway_request_capacity_exceeded")
@@ -378,6 +408,7 @@ pub(super) async fn send_responses_websocket_turn_start_error(
fn responses_websocket_turn_start_http_status(error: &GatewayError) -> u16 {
match error {
GatewayError::PlanUsageLimited(_) => StatusCode::TOO_MANY_REQUESTS.as_u16(),
GatewayError::Client { status, .. } => status.as_u16(),
GatewayError::AdmissionTimeout { .. } => StatusCode::TOO_MANY_REQUESTS.as_u16(),
GatewayError::LocalExecutionPlanningTimeout { .. } => StatusCode::GATEWAY_TIMEOUT.as_u16(),
@@ -387,6 +418,7 @@ fn responses_websocket_turn_start_http_status(error: &GatewayError) -> u16 {
pub(super) fn responses_websocket_turn_start_close(error: &GatewayError) -> (u16, &'static str) {
match error {
GatewayError::PlanUsageLimited(_) => (CLOSE_TRY_AGAIN, "plan_usage_limit_exceeded"),
GatewayError::Client { .. } => (CLOSE_POLICY_VIOLATION, "request_not_allowed"),
GatewayError::AdmissionTimeout { .. }
| GatewayError::LocalExecutionPlanningTimeout { .. } => (CLOSE_TRY_AGAIN, "gateway_busy"),
@@ -405,6 +437,7 @@ mod tests {
await_turn_finalization_handle, responses_websocket_turn_start_close,
responses_websocket_turn_start_http_status, spawn_bounded_adapter_observation_with_timeout,
};
use crate::plan_usage_policy::PlanUsagePolicyRejection;
use crate::GatewayError;
#[test]
@@ -422,6 +455,22 @@ mod tests {
);
}
#[test]
fn plan_usage_limit_uses_http_429_and_retry_later_close_code() {
let error = GatewayError::PlanUsageLimited(PlanUsagePolicyRejection {
metric: "actual_cost_usd",
limit: 10.0,
retry_after: 60,
window: "calendar_month",
});
assert_eq!(responses_websocket_turn_start_http_status(&error), 429);
assert_eq!(
responses_websocket_turn_start_close(&error),
(1013, "plan_usage_limit_exceeded")
);
}
/// C6 依赖的性质：结算是「等到落地」而不是「排进队列」。
///
/// 透明重试在这之后立刻按 health / adaptive / pool 状态规划下一个 attempt，
apps/aether-gateway/src/handlers/proxy/websocket/responses/mod.rs
+1
View File
@@ -18,6 +18,7 @@ mod frame;
mod lifecycle;
mod observation;
mod ownership;
mod plan_admission;
mod quota;
mod redaction;
mod relay_policy;
apps/aether-gateway/src/handlers/proxy/websocket/responses/ownership.rs
+3
View File
@@ -18,6 +18,7 @@ use crate::ai_serving::{
};
use crate::control::GatewayControlDecision;
use crate::orchestration::release_pool_key_lease_from_report_context;
use crate::plan_usage_policy::PlanUsagePolicySnapshot;
use crate::{AppState, GatewayError};
/// Owns a selected pool-key lease until the attempt lifecycle has taken over
@@ -149,6 +150,7 @@ pub(super) async fn begin_responses_websocket_turn_with_planned_lease(
control_decision: &GatewayControlDecision,
decision: AiExecutionDecision,
client_event: &Value,
plan_usage_policy_snapshot: Option<PlanUsagePolicySnapshot>,
mut planned_lease: PlannedPoolKeyLeaseGuard,
) -> Result<ActiveProviderAttempt, GatewayError> {
let state = state.clone();
@@ -163,6 +165,7 @@ pub(super) async fn begin_responses_websocket_turn_with_planned_lease(
&control_decision,
decision,
&client_event,
plan_usage_policy_snapshot,
)
.await?;
// ActiveProviderAttempt now owns the report context containing the
apps/aether-gateway/src/handlers/proxy/websocket/responses/plan_admission.rs
+232
View File
@@ -0,0 +1,232 @@
//! Subscription-plan request and concurrency admission for logical WS turns.
//!
//! The HTTP Upgrade is only a transport connection. Each accepted
//! `response.create` is admitted here as one logical request, while transparent
//! provider retries reuse the permit stored on that logical turn.
use axum::extract::ws::WebSocket;
use std::time::Duration;
use crate::handlers::proxy::websocket::session::{CLOSE_INTERNAL_ERROR, CLOSE_TRY_AGAIN};
use crate::handlers::proxy::websocket::transport::{
close_client_socket, send_responses_websocket_error,
};
use crate::plan_usage_policy::PlanUsageAdmissionError;
use crate::{AppState, GatewayError};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct ResponsesWebSocketPlanAdmissionError {
status: u16,
error_type: &'static str,
code: &'static str,
message: &'static str,
close_code: u16,
close_reason: &'static str,
}
pub(super) async fn acquire_responses_websocket_plan_admission(
state: &AppState,
decision: &crate::control::GatewayControlDecision,
logical_turn_id: &str,
) -> Result<crate::plan_usage_policy::PlanUsageAdmission, PlanUsageAdmissionError> {
crate::plan_usage_policy::check_and_acquire_plan_usage_policy_admission(
state,
Some(decision),
logical_turn_id,
crate::clock::current_unix_ms(),
)
.await
}
pub(super) fn responses_websocket_plan_permit_is_healthy(
permit: Option<&aether_runtime::AdmissionPermit>,
) -> bool {
permit.is_none_or(aether_runtime::AdmissionPermit::is_healthy)
}
pub(super) async fn send_responses_websocket_plan_admission_error(
client_socket: &mut WebSocket,
error: &PlanUsageAdmissionError,
) {
let mapped = map_plan_admission_error(error);
send_responses_websocket_error(
client_socket,
mapped.status,
mapped.error_type,
mapped.code,
mapped.message,
)
.await;
}
pub(super) fn responses_websocket_plan_admission_close(
error: &PlanUsageAdmissionError,
) -> (u16, &'static str) {
let mapped = map_plan_admission_error(error);
(mapped.close_code, mapped.close_reason)
}
pub(super) async fn wait_for_admission_permit_loss(
permit: Option<&aether_runtime::AdmissionPermit>,
) {
wait_for_admission_permit_loss_with_interval(permit, Duration::from_secs(1)).await;
}
async fn wait_for_admission_permit_loss_with_interval(
permit: Option<&aether_runtime::AdmissionPermit>,
health_poll_interval: Duration,
) {
let Some(permit) = permit else {
std::future::pending::<()>().await;
return;
};
let mut health = tokio::time::interval(health_poll_interval);
health.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
loop {
health.tick().await;
if !permit.is_healthy() {
return;
}
}
}
pub(super) async fn terminate_responses_websocket_for_plan_permit_loss(
client_socket: &mut WebSocket,
) {
let mapped = plan_permit_loss_error();
send_responses_websocket_error(
client_socket,
mapped.status,
mapped.error_type,
mapped.code,
mapped.message,
)
.await;
close_client_socket(client_socket, mapped.close_code, mapped.close_reason).await;
}
fn map_plan_admission_error(
error: &PlanUsageAdmissionError,
) -> ResponsesWebSocketPlanAdmissionError {
match error {
PlanUsageAdmissionError::Rejected(_) => plan_limit_exceeded_error(),
PlanUsageAdmissionError::Runtime(
aether_runtime_state::RuntimeSemaphoreError::Saturated { .. },
) => plan_limit_exceeded_error(),
PlanUsageAdmissionError::Runtime(
aether_runtime_state::RuntimeSemaphoreError::Unavailable { .. },
) => ResponsesWebSocketPlanAdmissionError {
status: 503,
error_type: "server_error",
code: "plan_usage_policy_unavailable",
message: "Gateway could not evaluate the subscription plan limit",
close_code: CLOSE_TRY_AGAIN,
close_reason: "plan_usage_policy_unavailable",
},
PlanUsageAdmissionError::Runtime(
aether_runtime_state::RuntimeSemaphoreError::InvalidConfiguration(_),
)
| PlanUsageAdmissionError::Gateway(GatewayError::Internal(_)) => {
ResponsesWebSocketPlanAdmissionError {
status: 500,
error_type: "server_error",
code: "plan_usage_policy_unavailable",
message: "Gateway could not evaluate the subscription plan limit",
close_code: CLOSE_INTERNAL_ERROR,
close_reason: "plan_usage_policy_unavailable",
}
}
PlanUsageAdmissionError::Gateway(GatewayError::Client { status, .. }) => {
ResponsesWebSocketPlanAdmissionError {
status: status.as_u16(),
error_type: "invalid_request_error",
code: "plan_usage_policy_unavailable",
message: "Gateway could not evaluate the subscription plan limit",
close_code: CLOSE_INTERNAL_ERROR,
close_reason: "plan_usage_policy_unavailable",
}
}
PlanUsageAdmissionError::Gateway(_) => ResponsesWebSocketPlanAdmissionError {
status: 500,
error_type: "server_error",
code: "plan_usage_policy_unavailable",
message: "Gateway could not evaluate the subscription plan limit",
close_code: CLOSE_INTERNAL_ERROR,
close_reason: "plan_usage_policy_unavailable",
},
}
}
const fn plan_permit_loss_error() -> ResponsesWebSocketPlanAdmissionError {
ResponsesWebSocketPlanAdmissionError {
status: 503,
error_type: "server_error",
code: "plan_usage_concurrency_unavailable",
message: "Subscription plan concurrency admission was lost; retry this response",
close_code: CLOSE_TRY_AGAIN,
close_reason: "plan_usage_concurrency_unavailable",
}
}
const fn plan_limit_exceeded_error() -> ResponsesWebSocketPlanAdmissionError {
ResponsesWebSocketPlanAdmissionError {
status: 429,
error_type: "rate_limit_error",
code: "plan_usage_limit_exceeded",
message: "Subscription plan usage limit exceeded; retry later",
close_code: CLOSE_TRY_AGAIN,
close_reason: "plan_usage_limit_exceeded",
}
}
#[cfg(test)]
mod tests {
use super::{
map_plan_admission_error, plan_permit_loss_error, responses_websocket_plan_admission_close,
};
use crate::plan_usage_policy::{PlanUsageAdmissionError, PlanUsagePolicyRejection};
#[test]
fn plan_limit_rejection_is_a_machine_readable_429() {
let error = PlanUsageAdmissionError::Rejected(PlanUsagePolicyRejection {
metric: "request_count",
limit: 10.0,
retry_after: 60,
window: "rolling",
});
let mapped = map_plan_admission_error(&error);
assert_eq!(mapped.status, 429);
assert_eq!(mapped.error_type, "rate_limit_error");
assert_eq!(mapped.code, "plan_usage_limit_exceeded");
assert_eq!(
responses_websocket_plan_admission_close(&error),
(1013, "plan_usage_limit_exceeded")
);
}
#[test]
fn saturated_plan_concurrency_uses_the_same_limit_error() {
let error = PlanUsageAdmissionError::Runtime(
aether_runtime_state::RuntimeSemaphoreError::Saturated {
gate: "plan_usage_concurrency",
limit: 2,
},
);
let mapped = map_plan_admission_error(&error);
assert_eq!(mapped.status, 429);
assert_eq!(mapped.code, "plan_usage_limit_exceeded");
}
#[test]
fn lost_plan_permit_is_a_retryable_client_visible_503() {
let mapped = plan_permit_loss_error();
assert_eq!(mapped.status, 503);
assert_eq!(mapped.error_type, "server_error");
assert_eq!(mapped.code, "plan_usage_concurrency_unavailable");
assert_eq!(mapped.close_code, 1013);
assert_eq!(mapped.close_reason, "plan_usage_concurrency_unavailable");
}
}
apps/aether-gateway/src/handlers/proxy/websocket/responses/quota.rs
+56 -15
View File
@@ -18,7 +18,9 @@ use super::request::{
};
use super::state::BoundResponsesConnection;
use super::turn::{prepare_responses_websocket_turn_decision, ResponsesWebSocketTurnOutcome};
use super::upstream::{bind_responses_upstream, close_bound_upstream};
use super::upstream::{
bind_responses_upstream, close_bound_upstream, ResponsesWebSocketUpstreamBindError,
};
use crate::clock::current_unix_secs;
use crate::handlers::proxy::websocket::ingress::WebSocketRequestContext;
use crate::handlers::proxy::websocket::session::WEBSOCKET_LOG_TRANSPORT;
@@ -103,7 +105,7 @@ pub(super) async fn retry_active_turn_after_quota_exhaustion(
state: &AppState,
context: &WebSocketRequestContext,
_previous_settled: PreviousAttemptSettled,
) -> bool {
) -> Result<bool, ()> {
// `LogicalTurn::client_event` is intentionally redacted before it is
// retained for replay. The binding, however, keeps the hash of the raw
// client-side Responses Lite tools/instructions so a later continuation
@@ -113,7 +115,7 @@ pub(super) async fn retry_active_turn_after_quota_exhaustion(
// redacted replay event.
let responses_lite_static_config = bound.responses_lite_static_config.clone();
let Some(active) = bound.turn_state.logical_mut() else {
return false;
return Ok(false);
};
if let Some(reason) = active.quota_retry_block_reason() {
debug!(
@@ -128,7 +130,7 @@ pub(super) async fn retry_active_turn_after_quota_exhaustion(
reason,
"gateway will not transparently replay an unsafe Responses WebSocket turn"
);
return false;
return Ok(false);
}
active.retry_attempted = true;
active.turn_attempt = active.turn_attempt.saturating_add(1);
@@ -142,12 +144,14 @@ pub(super) async fn retry_active_turn_after_quota_exhaustion(
trace_id = %context.trace_id,
"gateway refused to retry a WebSocket turn without its live authorization snapshot"
);
return false;
return Ok(false);
};
let turn_index = active.turn_index;
let logical_turn_id = active.logical_turn_id.clone();
let codex_fingerprint_context = active.codex_fingerprint_context.clone();
let turn_attempt = active.turn_attempt;
let plan_usage_permit = active.plan_usage_permit.clone();
let plan_usage_policy_snapshot = active.plan_usage_policy_snapshot.clone();
let retry_exclusion_until_unix_secs = bound
.pending_adapter_drain
@@ -193,7 +197,7 @@ pub(super) async fn retry_active_turn_after_quota_exhaustion(
exhausted_key_id = ?exhausted_key_id,
"gateway could not find an alternate Responses WebSocket provider after quota exhaustion"
);
return false;
return Ok(false);
}
Err(error) => {
warn!(
@@ -206,7 +210,7 @@ pub(super) async fn retry_active_turn_after_quota_exhaustion(
error = ?error,
"gateway could not plan an alternate Responses WebSocket provider after quota exhaustion"
);
return false;
return Ok(false);
}
};
let OwnedResponsesWebSocketDecision {
@@ -228,7 +232,7 @@ pub(super) async fn retry_active_turn_after_quota_exhaustion(
key_id = ?decision.key_id,
"gateway rejected an alternate Responses WebSocket plan that reused the exhausted key"
);
return false;
return Ok(false);
}
let provider_event = match planned_response_create_event(
&decision,
@@ -250,7 +254,7 @@ pub(super) async fn retry_active_turn_after_quota_exhaustion(
error_code = code,
"gateway could not rebuild a Responses response.create for transparent quota retry"
);
return false;
return Ok(false);
}
};
let replacement_provider_store = provider_event.get("store") == Some(&Value::Bool(true));
@@ -272,6 +276,7 @@ pub(super) async fn retry_active_turn_after_quota_exhaustion(
&turn_control.decision,
turn_decision,
&client_event,
plan_usage_policy_snapshot,
planned_lease,
)
.await
@@ -287,7 +292,7 @@ pub(super) async fn retry_active_turn_after_quota_exhaustion(
error = ?error,
"gateway could not start usage and audit tracking for transparent quota retry"
);
return false;
return Ok(false);
}
};
let mut replacement = match bind_responses_upstream(
@@ -295,11 +300,41 @@ pub(super) async fn retry_active_turn_after_quota_exhaustion(
normalization,
&client_event,
adapter,
plan_usage_permit.as_ref(),
|state| turn.record_upstream_request_state(state),
)
.await
{
Ok(connection) => connection,
Err(code) => {
Err(ResponsesWebSocketUpstreamBindError::PlanUsagePermitLost) => {
turn.release_plan_usage_cost_before_upstream_send(
state,
"transparent_retry_plan_usage_permit_lost",
)
.await;
queue_turn_finalization(
bound,
state,
turn,
ResponsesWebSocketTurnOutcome::connection_admission_lost(),
)
.await;
warn!(
event_name = "responses_websocket_quota_retry_plan_usage_concurrency_lost",
log_type = "ops",
transport = WEBSOCKET_LOG_TRANSPORT,
websocket = true,
trace_id = %context.trace_id,
"gateway stopped a transparent Responses WebSocket retry before its upstream send after the subscription plan concurrency lease became unhealthy"
);
return Err(());
}
Err(ResponsesWebSocketUpstreamBindError::Transport(code)) => {
turn.release_plan_usage_cost_before_upstream_send(
state,
"transparent_retry_rebind_failed",
)
.await;
queue_turn_finalization(
bound,
state,
@@ -316,11 +351,10 @@ pub(super) async fn retry_active_turn_after_quota_exhaustion(
error_code = code,
"gateway could not bind an alternate Responses WebSocket provider after quota exhaustion"
);
return false;
return Ok(false);
}
};
turn.mark_upstream_request_sent();
turn.set_provider_response_headers(replacement.upstream_response_headers.clone());
let replacement_upstream = replacement
.upstream
@@ -348,8 +382,15 @@ pub(super) async fn retry_active_turn_after_quota_exhaustion(
// drop guard 结算并让调用方走「透明重试失败」分支，不静默丢弃一条已经写了
// pending usage 行、占着 candidate 和 pool key lease 的 attempt。
if let Err(orphan) = bound.turn_state.resume(turn) {
let mut orphan = orphan;
orphan
.release_plan_usage_cost_before_upstream_send(
state,
"transparent_retry_state_handoff_failed",
)
.await;
drop(orphan);
return false;
return Ok(false);
}
if let Some(logical) = bound.turn_state.logical_mut() {
logical.provider_store = replacement_provider_store;
@@ -369,7 +410,7 @@ pub(super) async fn retry_active_turn_after_quota_exhaustion(
key_id = ?bound.decision_template.key_id,
"gateway transparently rebound a Responses WebSocket turn after quota exhaustion"
);
true
Ok(true)
}
fn responses_lite_static_config_after_rebind(
apps/aether-gateway/src/handlers/proxy/websocket/responses/redaction.rs
+67 -11
View File
@@ -15,7 +15,7 @@
//! （`privacy::restore_sync_response_body` / `privacy::StreamingResponseRestorer`），
//! WS 少了这一步，客户端就会直接看到 `<AETHER:EMAIL:...>`。
//! [`ResponsesWebSocketRedactionRestorer`] 补上这一跳，语义与 HTTP 完全一致：
//! 复用 `privacy::restore_json_strings`，只还原本连接自己 mask 出来的映射，
//! 复用 `privacy::restore_json_strings_with_budget`，只还原本连接自己 mask 出来的映射，
//! 未映射的占位符原样透传。
//!
//! ## session 为什么活在连接上而不是活在这一轮里
@@ -48,7 +48,11 @@ 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::privacy::{
restore_json_strings_with_budget, serialize_json_value_with_limit, RedactionSession,
RedactionSessionSlot, RestoreExpansionBudget, MAX_STREAM_RESTORE_EXPANSION_BYTES,
MAX_STREAM_RESTORE_OUTPUT_BYTES,
};
use crate::{AppState, GatewayError};
/// Responses WebSocket 只承载 `openai:responses`，脱敏规则按这个客户端格式选取。
@@ -191,29 +195,48 @@ impl ResponsesWebSocketRedactionRestorer {
/// 把一帧 provider 事件里的占位符换回真实值，返回要发给客户端的帧文本。
///
/// `None` 表示这一帧没有任何东西要还原，调用方必须原样转发上游字节：未启用
/// 脱敏（没有任何 session）时连 clone 都不做。
/// `Ok(None)` 表示这一帧没有任何东西要还原，调用方必须原样转发上游字节：
/// 未启用脱敏（没有任何 session）时连 clone 都不做。`Err` 表示恢复或有界
/// 序列化失败，调用方必须停止转发，不能把它降级成“未命中”而泄漏占位符。
///
/// 入参只读：审计与终态观测继续消费脱敏态的事件，还原只作用于发往客户端的
/// 那一份拷贝，和 HTTP 侧「审计存脱敏体、线上还原」保持一致。
pub(super) fn restore_provider_frame_text(&self, event: &Value) -> Option<String> {
pub(super) fn restore_provider_frame_text(
&self,
event: &Value,
) -> Result<Option<String>, GatewayError> {
if self.sessions.is_empty() {
return None;
return Ok(None);
}
let mut restored_event = event.clone();
let mut restored = false;
// Fresh per provider frame: this bounds one allocation amplification
// without imposing a cumulative byte or duration cap on the socket.
let mut budget = RestoreExpansionBudget::new(MAX_STREAM_RESTORE_EXPANSION_BYTES);
for session in &self.sessions {
// 逐 session 还原而不是合并映射：每个 session 只认自己 mask 过的
// sentinel（`RedactionSession::restore_text`），跨 session 合并会绕开
// 这条边界。同一个值在不同轮派生出的 sentinel 相同，所以顺序无关。
restored |= restore_json_strings(&mut restored_event, session);
restored |=
restore_json_strings_with_budget(&mut restored_event, session, &mut budget)?;
}
if !restored {
return None;
return Ok(None);
}
// 刚从 JSON 解析出来的 Value 再序列化不会失败；真失败时宁可让客户端看到
// 占位符，也不能丢掉这一帧——丢帧会让客户端的协议状态机卡死。
serde_json::to_string(&restored_event).ok()
// The parsed input frame is bounded at ingress, but use the same bounded
// serializer here so calculating the restored output limit cannot create
// an unchecked temporary allocation.
let original_len =
serialize_json_value_with_limit(event, MAX_STREAM_RESTORE_OUTPUT_BYTES)?.len();
let output_limit = original_len
.checked_add(MAX_STREAM_RESTORE_EXPANSION_BYTES)
.ok_or_else(|| {
GatewayError::Internal(
"WebSocket redaction restored frame length overflow".to_string(),
)
})?
.min(MAX_STREAM_RESTORE_OUTPUT_BYTES);
serialize_json_value_with_limit(&restored_event, output_limit).map(Some)
}
}
@@ -334,6 +357,7 @@ mod tests {
local_rejection: None,
allowed_models: None,
ip_rules: None,
verified_api_key_hash: None,
});
decision
}
@@ -691,6 +715,7 @@ mod tests {
let frame = provider_delta_frame(&format!("your mail is {sentinel}"));
let restored = restorer
.restore_provider_frame_text(&frame)
.expect("frame restoration must not fail")
.expect("a frame echoing this turn's sentinel must be restored");
assert!(
@@ -725,6 +750,7 @@ mod tests {
});
let restored = restorer
.restore_provider_frame_text(&frame)
.expect("frame restoration must not fail")
.expect("a batched sentinel must be restored");
assert!(restored.contains(TEST_EMAIL), "{restored}");
@@ -744,6 +770,7 @@ mod tests {
let frame = provider_delta_frame(&format!("{FOREIGN_SENTINEL} and {sentinel}"));
let restored = restorer
.restore_provider_frame_text(&frame)
.expect("frame restoration must not fail")
.expect("the mapped sentinel is still restored");
assert!(restored.contains(TEST_EMAIL), "{restored}");
@@ -764,17 +791,36 @@ mod tests {
assert!(
restorer
.restore_provider_frame_text(&provider_delta_frame("nothing to restore"))
.expect("frame restoration must not fail")
.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))
.expect("frame restoration must not fail")
.is_none(),
"a frame that only carries unmapped placeholders must not be rewritten"
);
}
#[tokio::test]
async fn frame_restore_budget_does_not_accumulate_across_frames() {
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);
for _ in 0..3 {
let restored = restorer
.restore_provider_frame_text(&provider_delta_frame(&sentinel))
.expect("frame restoration must not fail")
.expect("each provider frame gets an independent expansion budget");
assert!(restored.contains(TEST_EMAIL));
}
}
/// 未启用脱敏（或这条连接从没 mask 到东西）时，还原器必须完全不介入：
/// 连 clone 都不做，输出就是上游原字节。
#[tokio::test]
@@ -783,9 +829,11 @@ mod tests {
assert!(restorer
.restore_provider_frame_text(&provider_delta_frame(FOREIGN_SENTINEL))
.expect("frame restoration must not fail")
.is_none());
assert!(restorer
.restore_provider_frame_text(&provider_delta_frame(TEST_EMAIL))
.expect("frame restoration must not fail")
.is_none());
}
@@ -806,6 +854,7 @@ mod tests {
assert!(restorer
.restore_provider_frame_text(&provider_delta_frame(FOREIGN_SENTINEL))
.expect("frame restoration must not fail")
.is_none());
}
@@ -822,6 +871,7 @@ mod tests {
let before = frame.clone();
let _ = restorer
.restore_provider_frame_text(&frame)
.expect("frame restoration must not fail")
.expect("the frame is restored for the client");
assert_eq!(
@@ -849,6 +899,7 @@ mod tests {
let frame = provider_delta_frame(&format!("{first_sentinel} then {second_sentinel}"));
let restored = restorer
.restore_provider_frame_text(&frame)
.expect("frame restoration must not fail")
.expect("both turns' sentinels are restorable on this connection");
assert!(restored.contains(TEST_EMAIL), "{restored}");
@@ -876,11 +927,13 @@ mod tests {
assert!(
restorer
.restore_provider_frame_text(&provider_delta_frame(&prior_sentinel))
.expect("prior-chain restoration check must not fail")
.is_none(),
"an independent chain must not restore PII from its predecessor"
);
let restored = restorer
.restore_provider_frame_text(&provider_delta_frame(&current_sentinel))
.expect("the new chain's first-turn restoration must not fail")
.expect("the new chain's first-turn mapping must remain available");
assert!(restored.contains(OTHER_TEST_EMAIL), "{restored}");
assert!(!restored.contains(&current_sentinel), "{restored}");
@@ -901,6 +954,7 @@ mod tests {
assert!(!restorer.has_sessions());
assert!(restorer
.restore_provider_frame_text(&provider_delta_frame(&prior_sentinel))
.expect("cleared-chain restoration check must not fail")
.is_none());
}
@@ -927,12 +981,14 @@ mod tests {
assert!(
restorer
.restore_provider_frame_text(&provider_delta_frame(&oldest_sentinel))
.expect("frame restoration must not fail")
.is_none(),
"the evicted turn's sentinel is relayed verbatim, never mis-restored"
);
assert!(
restorer
.restore_provider_frame_text(&provider_delta_frame(&newest_sentinel))
.expect("frame restoration must not fail")
.is_some(),
"the most recent turns stay restorable"
);
apps/aether-gateway/src/handlers/proxy/websocket/responses/session.rs
+105 -31
View File
@@ -30,6 +30,11 @@ use super::ownership::{
await_owned_responses_websocket_plan, begin_responses_websocket_turn_with_planned_lease,
spawn_owned_responses_websocket_plan, OwnedResponsesWebSocketDecision,
};
use super::plan_admission::{
acquire_responses_websocket_plan_admission, responses_websocket_plan_admission_close,
send_responses_websocket_plan_admission_error,
terminate_responses_websocket_for_plan_permit_loss,
};
use super::redaction::redact_responses_websocket_client_event_with_reasoning_replay_policy;
use super::relay_policy::{fatal_relay_policy, FatalRelaySignal};
use super::request::{
@@ -41,7 +46,9 @@ use super::request::{
use super::state::BoundResponsesConnection;
use super::turn::{prepare_responses_websocket_turn_decision, ResponsesWebSocketTurnOutcome};
use super::turn_state::LogicalTurn;
use super::upstream::{bind_responses_upstream, close_bound_upstream};
use super::upstream::{
bind_responses_upstream, close_bound_upstream, ResponsesWebSocketUpstreamBindError,
};
use crate::ai_serving::ResponsesWebSocketPinnedCandidate;
use crate::handlers::proxy::websocket::ingress::{
@@ -522,6 +529,35 @@ async fn bootstrap_responses_websocket(
}
}
let plan_usage_admission = match acquire_responses_websocket_plan_admission(
&state,
&turn_control.decision,
&first_logical_turn_id,
)
.await
{
Ok(admission) => admission,
Err(error) => {
warn!(
event_name = "responses_websocket_initial_plan_usage_rejected",
log_type = "event",
transport = WEBSOCKET_LOG_TRANSPORT,
websocket = true,
trace_id = %context.trace_id,
logical_turn_id = %first_logical_turn_id,
error = ?error,
"gateway rejected the initial Responses WebSocket turn at its subscription plan limit"
);
send_responses_websocket_plan_admission_error(client_socket, &error).await;
let (close_code, close_reason) = responses_websocket_plan_admission_close(&error);
close_client_socket(client_socket, close_code, close_reason).await;
return None;
}
};
let crate::plan_usage_policy::PlanUsageAdmission {
permit: plan_usage_permit,
policy_snapshot: plan_usage_policy_snapshot,
} = plan_usage_admission;
// A cross-socket response chain must be owned by this exact live
// authenticated principal. Missing, expired, corrupt or unavailable state
// fails closed; allowing the normal scheduler to choose a provider/key
@@ -885,6 +921,7 @@ async fn bootstrap_responses_websocket(
&turn_control.decision,
first_turn_decision,
&first_event,
plan_usage_policy_snapshot.clone(),
planned_lease,
)
.await
@@ -907,36 +944,68 @@ async fn bootstrap_responses_websocket(
}
};
let mut bound =
match bind_responses_upstream(&decision, normalization, &first_event, adapter).await {
Ok(connection) => connection,
Err(code) => {
let finalizer = finalize_unbound_turn(
state.clone(),
first_turn,
ResponsesWebSocketTurnOutcome::upstream_connect_failed(code),
);
warn!(
event_name = "responses_websocket_upstream_connect_failed",
log_type = "ops",
transport = WEBSOCKET_LOG_TRANSPORT,
websocket = true,
trace_id = %context.trace_id,
error_code = code,
"gateway failed to establish Responses WebSocket upstream"
);
send_gateway_error_with_status(
client_socket,
502,
code,
"Gateway could not establish the Provider connection",
let mut bound = match bind_responses_upstream(
&decision,
normalization,
&first_event,
adapter,
plan_usage_permit.as_ref(),
|state| first_turn.record_upstream_request_state(state),
)
.await
{
Ok(connection) => connection,
Err(ResponsesWebSocketUpstreamBindError::PlanUsagePermitLost) => {
first_turn
.release_plan_usage_cost_before_upstream_send(
&state,
"initial_plan_usage_permit_lost",
)
.await;
close_client_socket(client_socket, CLOSE_TRY_AGAIN, code).await;
await_turn_finalization_handle(finalizer).await;
return None;
}
};
let finalizer = finalize_unbound_turn(
state.clone(),
first_turn,
ResponsesWebSocketTurnOutcome::connection_admission_lost(),
);
warn!(
event_name = "responses_websocket_initial_plan_usage_concurrency_lost",
log_type = "ops",
transport = WEBSOCKET_LOG_TRANSPORT,
websocket = true,
trace_id = %context.trace_id,
"gateway stopped the initial Responses WebSocket turn before its upstream send after the subscription plan concurrency lease became unhealthy"
);
terminate_responses_websocket_for_plan_permit_loss(client_socket).await;
await_turn_finalization_handle(finalizer).await;
return None;
}
Err(ResponsesWebSocketUpstreamBindError::Transport(code)) => {
let finalizer = finalize_unbound_turn(
state.clone(),
first_turn,
ResponsesWebSocketTurnOutcome::upstream_connect_failed(code),
);
warn!(
event_name = "responses_websocket_upstream_connect_failed",
log_type = "ops",
transport = WEBSOCKET_LOG_TRANSPORT,
websocket = true,
trace_id = %context.trace_id,
error_code = code,
"gateway failed to establish Responses WebSocket upstream"
);
send_gateway_error_with_status(
client_socket,
502,
code,
"Gateway could not establish the Provider connection",
)
.await;
close_client_socket(client_socket, CLOSE_TRY_AGAIN, code).await;
await_turn_finalization_handle(finalizer).await;
return None;
}
};
if bound.responses_lite_static_config.is_some() {
bound.responses_lite_static_config = continuation_record
.as_ref()
@@ -953,7 +1022,6 @@ async fn bootstrap_responses_websocket(
.continuation_response_ids
.remember_persisted(previous_response_id);
}
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 {
register_initial_redaction_session(&mut bound, session);
@@ -962,7 +1030,9 @@ async fn bootstrap_responses_websocket(
LogicalTurn::new(first_event, 1, first_logical_turn_id)
.with_codex_fingerprint_context(first_codex_fingerprint_context)
.with_provider_store(first_provider_event.get("store") == Some(&Value::Bool(true)))
.with_turn_control(turn_control),
.with_turn_control(turn_control)
.with_plan_usage_permit(plan_usage_permit)
.with_plan_usage_policy_snapshot(plan_usage_policy_snapshot),
first_turn,
);
@@ -2067,6 +2137,8 @@ mod tests {
resolve_responses_websocket_adapter(
crate::orchestration::ResponsesWebSocketAdapter::Standard,
),
None,
|_| {},
)
.await
.expect("upstream binding should succeed");
@@ -2278,6 +2350,7 @@ mod tests {
let restored = ordinary_bound
.redaction_restorer
.restore_provider_frame_text(&provider_event)
.expect("ordinary OpenAI replay policy restoration must not fail")
.expect("ordinary OpenAI replay policy should restore response text");
let restored: serde_json::Value =
serde_json::from_str(&restored).expect("restored provider event should remain JSON");
@@ -2297,6 +2370,7 @@ mod tests {
deepseek_bound
.redaction_restorer
.restore_provider_frame_text(&provider_event)
.expect("DeepSeek opaque restoration check must not fail")
.is_none(),
"the authenticated DeepSeek binding must keep opaque reasoning state byte-identical"
);
apps/aether-gateway/src/handlers/proxy/websocket/responses/settlement.rs
+24
View File
@@ -283,4 +283,28 @@ mod tests {
// 投递失败不是供应商的错误，摘要不该因此补 parser_error。
assert_eq!(facts.forced_error(), None);
}
#[test]
fn plan_permit_loss_is_gateway_cancellation_not_provider_failure() {
let facts = attempt_facts_for_outcome(
None,
AttemptClientDelivery::Complete,
ResponsesWebSocketTurnOutcome::connection_admission_lost(),
);
assert_eq!(
facts.provider,
aborted(
499,
"gateway WebSocket connection admission became unhealthy"
)
);
assert_eq!(
facts.delivery,
AttemptClientDelivery::Aborted {
reason: "gateway WebSocket connection admission became unhealthy"
}
);
assert_eq!(facts.forced_error(), None);
}
}
apps/aether-gateway/src/handlers/proxy/websocket/responses/turn.rs
+433 -14
View File
@@ -17,7 +17,8 @@ use aether_contracts::{
};
use aether_data_contracts::repository::candidates::RequestCandidateStatus;
use aether_data_contracts::repository::usage::{
UsageBodyCaptureState, WEBSOCKET_MODE_METADATA_KEY, WEBSOCKET_TRANSPORT_METADATA_KEY,
UsageBodyCaptureState, PLAN_USAGE_RESERVATION_DEFERRED_METADATA_KEY,
WEBSOCKET_MODE_METADATA_KEY, WEBSOCKET_TRANSPORT_METADATA_KEY,
};
use aether_scheduler_core::SchedulerRequestCandidateStatusUpdate;
use aether_usage_runtime::{
@@ -51,6 +52,7 @@ use crate::orchestration::{
release_local_pool_key_lease, release_pool_key_lease_from_report_context,
LocalExecutionEffectContext, LocalStreamFailureEffect,
};
use crate::plan_usage_policy::PlanUsagePolicySnapshot;
use crate::request_candidate_runtime::{
ensure_execution_request_candidate_slot, record_local_request_candidate_status,
};
@@ -146,6 +148,10 @@ impl ResponsesWebSocketTurnOutcome {
}
pub(super) const fn connection_admission_lost() -> Self {
// This gateway-owned cancellation is also used when a logical turn's
// subscription-plan permit becomes unhealthy. Keeping it Cancelled
// prevents the settlement layer from projecting a provider failure;
// the client-visible 503 is emitted separately by plan_admission.
Self::Cancelled {
reason: "gateway WebSocket connection admission became unhealthy",
}
@@ -250,16 +256,226 @@ pub(super) struct ResponsesProviderAttempt {
first_event_timeout: Duration,
terminal_timeout: Duration,
admission: Option<ResponsesWebSocketTurnAdmission>,
/// A provider attempt gets its own server-issued cost reservation. The
/// trusted token also lives in the lifecycle report context so the normal
/// terminal usage path can reconcile actual cost. We retain the control
/// snapshot here only for the quota-retry path, which must release the old
/// attempt synchronously before reserving the replacement.
plan_usage_cost_reservation: Option<ResponsesPlanUsageCostReservation>,
terminal_error_body: Option<String>,
/// 观察到的 provider 终态事实，与「为什么现在结算」这个信号分开保存。
/// 客户端投递失败不会把它擦掉。
provider_outcome: Option<AttemptProviderOutcome>,
/// 这一个 attempt 的内容是否完整交付给了客户端。与 provider 终态正交。
client_delivery: AttemptClientDelivery,
/// True only after `response.create` has been accepted by the upstream
/// socket writer. Cancellation before this point must not be projected as
/// provider failure or billed usage.
upstream_request_sent: bool,
/// Tracks the irreversible transport handoff separately from flush
/// confirmation. Once start_send succeeds, a later cancellation cannot
/// prove the provider did not receive or execute response.create.
upstream_request_state: UpstreamRequestState,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) enum UpstreamRequestState {
NotStarted,
PossiblySent,
Sent,
}
impl UpstreamRequestState {
const fn may_have_reached_provider(self) -> bool {
!matches!(self, Self::NotStarted)
}
}
struct ResponsesPlanUsageCostReservation {
control_decision: GatewayControlDecision,
reservation_token: String,
}
struct OwnedResponsesPlanUsageCostReservation {
control_decision: GatewayControlDecision,
plan: ExecutionPlan,
reservation_token: String,
}
/// Owns a successful durable reservation until an attempt lifecycle owns the
/// trusted report context. If startup is cancelled anywhere in between, Drop
/// releases the reservation instead of leaving a false 429 until its TTL.
struct ResponsesPlanUsageCostReservationGuard {
state: AppState,
reservation: Option<OwnedResponsesPlanUsageCostReservation>,
}
impl ResponsesPlanUsageCostReservationGuard {
fn new(
state: &AppState,
control_decision: GatewayControlDecision,
plan: ExecutionPlan,
reservation_token: String,
) -> Self {
Self {
state: state.clone(),
reservation: Some(OwnedResponsesPlanUsageCostReservation {
control_decision,
plan,
reservation_token,
}),
}
}
fn reservation_token(&self) -> &str {
self.reservation
.as_ref()
.expect("an armed plan cost reservation guard owns its reservation")
.reservation_token
.as_str()
}
fn disarm(mut self) -> ResponsesPlanUsageCostReservation {
let reservation = self
.reservation
.take()
.expect("an armed plan cost reservation guard owns its reservation");
ResponsesPlanUsageCostReservation {
control_decision: reservation.control_decision,
reservation_token: reservation.reservation_token,
}
}
}
impl Drop for ResponsesPlanUsageCostReservationGuard {
fn drop(&mut self) {
let Some(reservation) = self.reservation.take() else {
return;
};
let state = self.state.clone();
if let Ok(handle) = tokio::runtime::Handle::try_current() {
handle.spawn(async move {
release_responses_plan_usage_cost_best_effort(
&state,
&reservation.control_decision,
&reservation.plan,
reservation.reservation_token.as_str(),
"turn_start_cancelled",
)
.await;
});
}
}
}
enum ResponsesPlanUsageCostReservationStart {
NotRequired,
Reserved(ResponsesPlanUsageCostReservationGuard),
Rejected(crate::plan_usage_policy::PlanUsagePolicyRejection),
}
async fn reserve_responses_plan_usage_cost_owned(
state: &AppState,
control_decision: &GatewayControlDecision,
plan: &ExecutionPlan,
report_context: Option<&Value>,
plan_usage_policy_snapshot: Option<&PlanUsagePolicySnapshot>,
reservation_token: String,
) -> Result<ResponsesPlanUsageCostReservationStart, GatewayError> {
let outcome = crate::plan_usage_policy::reserve_admitted_plan_usage_policy_cost(
state,
control_decision,
plan,
report_context,
plan_usage_policy_snapshot,
reservation_token.as_str(),
)
.await?;
match outcome {
crate::plan_usage_policy::PlanUsageCostReservationOutcome::NotRequired => {
Ok(ResponsesPlanUsageCostReservationStart::NotRequired)
}
crate::plan_usage_policy::PlanUsageCostReservationOutcome::Reserved => {
Ok(ResponsesPlanUsageCostReservationStart::Reserved(
ResponsesPlanUsageCostReservationGuard::new(
state,
control_decision.clone(),
plan.clone(),
reservation_token,
),
))
}
crate::plan_usage_policy::PlanUsageCostReservationOutcome::Rejected(rejection) => {
Ok(ResponsesPlanUsageCostReservationStart::Rejected(rejection))
}
}
}
fn attach_plan_usage_reservation_token(
report_context: Option<Value>,
reservation_token: &str,
) -> Option<Value> {
let mut object = match report_context {
Some(Value::Object(object)) => object,
Some(other) => Map::from_iter([("seed".to_string(), other)]),
None => Map::new(),
};
// Always overwrite a client-derived/seed value with the server-issued
// token. Terminal reconciliation must never trust caller-controlled JSON.
object.insert(
"plan_usage_reservation_token".to_string(),
Value::String(reservation_token.to_string()),
);
// The token and its reconciliation controls are server-owned. A seed or
// client event must never be able to defer reconciliation on its own.
object.insert(
PLAN_USAGE_RESERVATION_DEFERRED_METADATA_KEY.to_string(),
Value::Bool(false),
);
Some(Value::Object(object))
}
fn websocket_plan_usage_rejection_error(
rejection: &crate::plan_usage_policy::PlanUsagePolicyRejection,
) -> GatewayError {
GatewayError::PlanUsageLimited(rejection.clone())
}
async fn release_responses_plan_usage_cost(
state: &AppState,
control_decision: &GatewayControlDecision,
plan: &ExecutionPlan,
reservation_token: &str,
_reason: &'static str,
) -> Result<(), GatewayError> {
crate::plan_usage_policy::release_plan_usage_policy_cost(
state,
control_decision,
plan,
reservation_token,
current_unix_ms() / 1_000,
)
.await
}
async fn release_responses_plan_usage_cost_best_effort(
state: &AppState,
control_decision: &GatewayControlDecision,
plan: &ExecutionPlan,
reservation_token: &str,
reason: &'static str,
) {
if let Err(error) =
release_responses_plan_usage_cost(state, control_decision, plan, reservation_token, reason)
.await
{
warn!(
event_name = "responses_websocket_plan_usage_cost_release_failed",
log_type = "ops",
request_id = %plan.request_id,
candidate_id = ?plan.candidate_id,
reservation_token,
reason,
error = ?error,
"gateway failed to release a Responses WebSocket plan cost reservation"
);
}
}
/// 组装一轮 turn 的 decision。
@@ -307,6 +523,7 @@ pub(super) async fn begin_unowned_responses_websocket_turn(
control_decision: &GatewayControlDecision,
decision: AiExecutionDecision,
client_event: &Value,
plan_usage_policy_snapshot: Option<PlanUsagePolicySnapshot>,
) -> Result<ResponsesProviderAttempt, GatewayError> {
let planned_report_context = decision.report_context.clone();
let attempt = match build_openai_responses_stream_plan_from_decision(
@@ -415,6 +632,68 @@ pub(super) async fn begin_unowned_responses_websocket_turn(
}
};
let reservation_token = uuid::Uuid::new_v4().to_string();
let cost_reservation = match reserve_responses_plan_usage_cost_owned(
state,
control_decision,
&plan,
report_context.as_ref(),
plan_usage_policy_snapshot.as_ref(),
reservation_token,
)
.await
{
Ok(ResponsesPlanUsageCostReservationStart::NotRequired) => None,
Ok(ResponsesPlanUsageCostReservationStart::Reserved(guard)) => {
report_context =
attach_plan_usage_reservation_token(report_context, guard.reservation_token());
Some(guard)
}
Ok(ResponsesPlanUsageCostReservationStart::Rejected(rejection)) => {
let error = websocket_plan_usage_rejection_error(&rejection);
admission.release().await;
release_then_record_responses_websocket_admission_failure(
release_local_pool_key_lease(
state,
LocalExecutionEffectContext {
plan: &plan,
report_context: report_context.as_ref(),
},
),
record_responses_websocket_admission_failure(
state,
&plan,
report_context.as_ref(),
candidate_started_at_unix_ms,
&error,
),
)
.await;
return Err(error);
}
Err(error) => {
admission.release().await;
release_then_record_responses_websocket_admission_failure(
release_local_pool_key_lease(
state,
LocalExecutionEffectContext {
plan: &plan,
report_context: report_context.as_ref(),
},
),
record_responses_websocket_admission_failure(
state,
&plan,
report_context.as_ref(),
candidate_started_at_unix_ms,
&error,
),
)
.await;
return Err(error);
}
};
let lifecycle = ExecutionAttemptLifecycle::begin(
state,
AttemptLifecycleSeed {
@@ -427,6 +706,7 @@ pub(super) async fn begin_unowned_responses_websocket_turn(
},
)
.await;
let plan_usage_cost_reservation = cost_reservation.map(|guard| guard.disarm());
Ok(ResponsesProviderAttempt {
lifecycle,
@@ -442,10 +722,11 @@ pub(super) async fn begin_unowned_responses_websocket_turn(
first_event_timeout,
terminal_timeout,
admission: Some(admission),
plan_usage_cost_reservation,
terminal_error_body: None,
provider_outcome: None,
client_delivery: AttemptClientDelivery::Complete,
upstream_request_sent: false,
upstream_request_state: UpstreamRequestState::NotStarted,
})
}
@@ -496,6 +777,11 @@ fn responses_websocket_admission_failure_update(
"gateway_admission_timeout",
format!("gateway admission gate {gate} timed out after {queue_budget_ms}ms"),
),
GatewayError::PlanUsageLimited(_) => (
StatusCode::TOO_MANY_REQUESTS.as_u16(),
"plan_usage_limit_exceeded",
"subscription plan usage limit exceeded".to_string(),
),
other => (
StatusCode::INTERNAL_SERVER_ERROR.as_u16(),
"gateway_admission_failed",
@@ -577,6 +863,64 @@ impl ResponsesProviderAttempt {
}
}
/// Releases this attempt's reserved cost before a transparent retry. The
/// terminal usage write that follows may reconcile the same token again;
/// repository reconciliation is idempotent, so the first Released state
/// wins and the replacement attempt can reserve independently.
pub(super) async fn release_plan_usage_cost_for_retry(
&mut self,
state: &AppState,
) -> Result<(), GatewayError> {
if self.upstream_request_state.may_have_reached_provider()
&& self.provider_outcome.is_none()
{
return Err(GatewayError::Internal(
"cannot release an unresolved plan cost reservation after the upstream request may have been sent"
.to_string(),
));
}
let Some(reservation) = self.plan_usage_cost_reservation.take() else {
return Ok(());
};
let release = release_responses_plan_usage_cost(
state,
&reservation.control_decision,
self.lifecycle.plan(),
reservation.reservation_token.as_str(),
"transparent_retry",
)
.await;
if release.is_err() {
self.plan_usage_cost_reservation = Some(reservation);
}
release
}
/// A transparent replacement failed before its request reached an
/// upstream. No terminal usage event can safely own the still-reserved
/// estimate, so release it explicitly before the orphan attempt is
/// finalized or dropped.
pub(super) async fn release_plan_usage_cost_before_upstream_send(
&mut self,
state: &AppState,
reason: &'static str,
) {
if self.upstream_request_state.may_have_reached_provider() {
return;
}
let Some(reservation) = self.plan_usage_cost_reservation.take() else {
return;
};
release_responses_plan_usage_cost_best_effort(
state,
&reservation.control_decision,
self.lifecycle.plan(),
reservation.reservation_token.as_str(),
reason,
)
.await;
}
pub(super) fn set_provider_response_headers(&mut self, headers: BTreeMap<String, String>) {
let observed_at_unix_ms = current_unix_ms();
let report_context = attach_provider_response_headers_to_report_context(
@@ -592,12 +936,33 @@ impl ResponsesProviderAttempt {
/// Starts the per-turn response deadlines only after the corresponding
/// `response.create` has been accepted by the upstream socket writer.
pub(super) fn mark_upstream_request_sent(&mut self) {
self.upstream_request_sent = true;
self.started_at = Instant::now();
self.provider_request_started_at_unix_ms = current_unix_ms();
self.provider_request_order_id = uuid::Uuid::now_v7().to_string();
self.first_event_elapsed_ms = None;
pub(super) fn record_upstream_request_state(&mut self, state: UpstreamRequestState) {
match state {
UpstreamRequestState::NotStarted => {
debug_assert!(
false,
"upstream request state cannot move back to not started"
);
}
UpstreamRequestState::PossiblySent => {
debug_assert_eq!(
self.upstream_request_state,
UpstreamRequestState::NotStarted
);
self.upstream_request_state = state;
self.started_at = Instant::now();
self.provider_request_started_at_unix_ms = current_unix_ms();
self.provider_request_order_id = uuid::Uuid::now_v7().to_string();
self.first_event_elapsed_ms = None;
}
UpstreamRequestState::Sent => {
debug_assert_eq!(
self.upstream_request_state,
UpstreamRequestState::PossiblySent
);
self.upstream_request_state = state;
}
}
}
/// Selects a cancellation-safe fallback for an attempt whose owner task
@@ -605,7 +970,9 @@ impl ResponsesProviderAttempt {
/// after the write it remains a gateway relay failure because provider
/// work may already have started.
pub(super) const fn abandonment_outcome(&self) -> ResponsesWebSocketTurnOutcome {
ResponsesWebSocketTurnOutcome::relay_task_abandonment(self.upstream_request_sent)
ResponsesWebSocketTurnOutcome::relay_task_abandonment(
self.upstream_request_state.may_have_reached_provider(),
)
}
pub(super) fn deadline(&self) -> ResponsesWebSocketTurnDeadline {
@@ -734,6 +1101,14 @@ impl ResponsesProviderAttempt {
/// 这里只提供 WS 观察到的终态事实。
async fn settle(mut self, state: &AppState, outcome: ResponsesWebSocketTurnOutcome) {
let facts = attempt_facts_for_outcome(self.provider_outcome, self.client_delivery, outcome);
if self.upstream_request_state.may_have_reached_provider()
&& !facts.provider.is_terminal()
&& self.plan_usage_cost_reservation.is_some()
{
let report_context =
attach_plan_usage_reservation_deferred(self.lifecycle.take_report_context());
self.lifecycle.set_report_context(report_context);
}
if let Some(reason) = facts.delivery.aborted_reason() {
let report_context = attach_client_delivery_to_report_context(
self.lifecycle.take_report_context(),
@@ -951,6 +1326,19 @@ fn attach_client_delivery_to_report_context(
Some(Value::Object(object))
}
fn attach_plan_usage_reservation_deferred(report_context: Option<Value>) -> Option<Value> {
let mut object = match report_context {
Some(Value::Object(object)) => object,
Some(other) => Map::from_iter([("seed".to_string(), other)]),
None => Map::new(),
};
object.insert(
PLAN_USAGE_RESERVATION_DEFERRED_METADATA_KEY.to_string(),
Value::Bool(true),
);
Some(Value::Object(object))
}
fn provider_terminal_outcome(
frame: &ParsedResponsesWebSocketFrame<'_>,
) -> Option<ResponsesWebSocketTurnOutcome> {
@@ -1016,7 +1404,8 @@ mod tests {
attempt_facts_for_outcome, settle_signal_for_client_delivery_failure,
};
use super::{
attach_client_delivery_to_report_context, prepare_websocket_report_context,
attach_client_delivery_to_report_context, attach_plan_usage_reservation_deferred,
attach_plan_usage_reservation_token, prepare_websocket_report_context,
provider_terminal_outcome, release_then_record_responses_websocket_admission_failure,
resolve_responses_websocket_turn_timeouts, responses_websocket_admission_failure_update,
websocket_event_as_sse_line, ResponsesWebSocketTurnDeadline, ResponsesWebSocketTurnOutcome,
@@ -1028,6 +1417,36 @@ mod tests {
};
use crate::GatewayError;
#[test]
fn server_plan_usage_reservation_token_overrides_seed_value() {
let context = attach_plan_usage_reservation_token(
Some(json!({
"candidate_index": 2,
"plan_usage_reservation_token": "client-controlled-token"
})),
"server-reservation-token",
)
.expect("reservation context");
assert_eq!(context["candidate_index"], 2);
assert_eq!(
context["plan_usage_reservation_token"],
"server-reservation-token"
);
assert_eq!(context["plan_usage_reservation_deferred"], false);
}
#[test]
fn gateway_can_defer_a_server_owned_reservation_after_transport_handoff() {
let context = attach_plan_usage_reservation_deferred(Some(json!({
"plan_usage_reservation_token": "server-reservation-token",
"plan_usage_reservation_deferred": false
})))
.expect("reservation context");
assert_eq!(context["plan_usage_reservation_deferred"], true);
}
#[tokio::test]
async fn admission_failure_releases_pool_lease_before_recording_candidate_terminal() {
let phase = Arc::new(AtomicU8::new(0));
apps/aether-gateway/src/handlers/proxy/websocket/responses/turn_state.rs
+28 -1
View File
@@ -11,6 +11,7 @@ use serde_json::Value;
use super::control::ResponsesWebSocketTurnControl;
use super::lifecycle::ActiveProviderAttempt;
use super::request::response_create_has_previous_response_id;
use crate::plan_usage_policy::PlanUsagePolicySnapshot;
/// 客户端一次 `response.create` 对应的 logical turn。
///
@@ -18,7 +19,7 @@ use super::request::response_create_has_previous_response_id;
/// 换一条上游连接重放同一份客户端事件，但对客户端始终是同一轮请求。
/// `client_event` 保存的必须是**已脱敏**的事件（见 `super::redaction`），
/// 因为透明重试直接重放它。
#[derive(Debug, Clone)]
#[derive(Debug)]
pub(super) struct LogicalTurn {
pub(super) client_event: Value,
/// Effective `store` after provider body rules and WebSocket framing. Only
@@ -37,6 +38,14 @@ pub(super) struct LogicalTurn {
/// this logical turn. Quota retries reuse it instead of falling back to the
/// connection's Upgrade-time authorization snapshot.
pub(super) turn_control: Option<ResponsesWebSocketTurnControl>,
/// Subscription-plan concurrency is per logical `response.create`, not per
/// provider attempt. Keeping the permit here makes transparent retries
/// retain the same slot through `Replanning` until the logical terminal.
pub(super) plan_usage_permit: Option<aether_runtime::AdmissionPermit>,
/// Cost limits are evaluated against the policy and entitlement window
/// that admitted this logical response.create. Every transparent provider
/// attempt reuses this snapshot while receiving its own reservation token.
pub(super) plan_usage_policy_snapshot: Option<PlanUsagePolicySnapshot>,
}
impl LogicalTurn {
@@ -51,6 +60,8 @@ impl LogicalTurn {
retry_attempted: false,
retry_unsafe_reason: None,
turn_control: None,
plan_usage_permit: None,
plan_usage_policy_snapshot: None,
}
}
@@ -59,6 +70,22 @@ impl LogicalTurn {
self
}
pub(super) fn with_plan_usage_permit(
mut self,
permit: Option<aether_runtime::AdmissionPermit>,
) -> Self {
self.plan_usage_permit = permit;
self
}
pub(super) fn with_plan_usage_policy_snapshot(
mut self,
snapshot: Option<PlanUsagePolicySnapshot>,
) -> Self {
self.plan_usage_policy_snapshot = snapshot;
self
}
pub(super) fn with_codex_fingerprint_context(
mut self,
context: CodexFingerprintConvergenceContext,
apps/aether-gateway/src/handlers/proxy/websocket/responses/upstream.rs
+308 -26
View File
@@ -7,18 +7,33 @@ use wreq::ws::message::Message as WreqWsMessage;
use super::adapter::ResponsesWebSocketProtocolAdapter;
use super::binding::{UpstreamBindingIdentity, UpstreamBindingIdentityError};
use super::plan_admission::responses_websocket_plan_permit_is_healthy;
use super::redaction::ResponsesWebSocketRedactionRestorer;
use super::request::{planned_request_uses_codex_responses_lite, planned_response_create_event};
use super::state::{
BoundResponsesConnection, ContinuationResponseIds, ExhaustedResponsesWebSocketExclusions,
};
use super::turn::UpstreamRequestState;
use super::turn_state::ResponsesTurnState;
use crate::ai_serving::{AiExecutionDecision, ResponsesWebSocketBodyNormalization};
use crate::handlers::proxy::websocket::session::RESPONSES_WEBSOCKET_SESSION_LIMITS;
use crate::handlers::proxy::websocket::transport::{
close_upstream_socket, connect_upstream_websocket, send_upstream_message,
close_upstream_socket, connect_upstream_websocket, feed_upstream_message,
flush_upstream_messages, WebSocketWriteError,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) enum ResponsesWebSocketUpstreamSendError {
PlanUsagePermitLost,
Transport(WebSocketWriteError),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) enum ResponsesWebSocketUpstreamBindError {
PlanUsagePermitLost,
Transport(&'static str),
}
/// 上游 WebSocket 握手的默认绝对 deadline（30 秒）。
/// 覆盖 DNS → TCP connect → TLS → HTTP 101 Upgrade → 发送首条 event 的完整链路。
/// 如果 decision 配置了更短的 first_byte_ms 或 total_ms，取其与此值的较小者。
@@ -39,59 +54,100 @@ pub(super) fn resolve_upstream_handshake_deadline(decision: &AiExecutionDecision
Duration::from_millis(deadline_ms)
}
pub(super) async fn bind_responses_upstream(
pub(super) async fn bind_responses_upstream<F>(
decision: &AiExecutionDecision,
normalization: ResponsesWebSocketBodyNormalization,
initial_event: &Value,
adapter: &'static dyn ResponsesWebSocketProtocolAdapter,
) -> Result<BoundResponsesConnection, &'static str> {
plan_usage_permit: Option<&aether_runtime::AdmissionPermit>,
record_upstream_request_state: F,
) -> Result<BoundResponsesConnection, ResponsesWebSocketUpstreamBindError>
where
F: FnMut(UpstreamRequestState),
{
// 绝对 deadline：从此刻起必须在限定时间内完成握手 + 首条事件发送，
// 防止慢 TLS / 慢 HTTP Upgrade 无限占用 connection permit。
let handshake_deadline = resolve_upstream_handshake_deadline(decision);
tokio::time::timeout(
handshake_deadline,
bind_responses_upstream_inner(decision, normalization, initial_event, adapter),
bind_responses_upstream_inner(
decision,
normalization,
initial_event,
adapter,
plan_usage_permit,
record_upstream_request_state,
),
)
.await
.map_err(|_| "responses_websocket_upstream_handshake_timeout")?
.map_err(|_| {
ResponsesWebSocketUpstreamBindError::Transport(
"responses_websocket_upstream_handshake_timeout",
)
})?
}
/// 实际执行握手 + 首条事件发送的内部函数，由外层 timeout 包裹。
async fn bind_responses_upstream_inner(
async fn bind_responses_upstream_inner<F>(
decision: &AiExecutionDecision,
normalization: ResponsesWebSocketBodyNormalization,
initial_event: &Value,
adapter: &'static dyn ResponsesWebSocketProtocolAdapter,
) -> Result<BoundResponsesConnection, &'static str> {
plan_usage_permit: Option<&aether_runtime::AdmissionPermit>,
record_upstream_request_state: F,
) -> Result<BoundResponsesConnection, ResponsesWebSocketUpstreamBindError>
where
F: FnMut(UpstreamRequestState),
{
let binding_identity =
UpstreamBindingIdentity::from_decision(adapter, decision).map_err(|error| match error {
UpstreamBindingIdentityError::MissingUpstreamUrl => {
adapter.upstream_errors().upstream_url_missing
}
UpstreamBindingIdentityError::InvalidUpstreamUrl => {
adapter.upstream_errors().upstream_url_invalid
}
UpstreamBindingIdentityError::InvalidHandshakeHeaders => {
adapter.upstream_errors().headers_invalid
}
UpstreamBindingIdentity::from_decision(adapter, decision).map_err(|error| {
ResponsesWebSocketUpstreamBindError::Transport(match error {
UpstreamBindingIdentityError::MissingUpstreamUrl => {
adapter.upstream_errors().upstream_url_missing
}
UpstreamBindingIdentityError::InvalidUpstreamUrl => {
adapter.upstream_errors().upstream_url_invalid
}
UpstreamBindingIdentityError::InvalidHandshakeHeaders => {
adapter.upstream_errors().headers_invalid
}
})
})?;
let mut upstream = connect_upstream_websocket(
decision,
RESPONSES_WEBSOCKET_SESSION_LIMITS,
adapter.upstream_errors(),
)
.await?;
let first_event = planned_response_create_event(decision, &normalization, initial_event)?;
send_upstream_message(&mut upstream.socket, WreqWsMessage::text(first_event))
.await
.map_err(|_| "responses_websocket_initial_send_failed")?;
.await
.map_err(ResponsesWebSocketUpstreamBindError::Transport)?;
let first_event = planned_response_create_event(decision, &normalization, initial_event)
.map_err(ResponsesWebSocketUpstreamBindError::Transport)?;
send_responses_websocket_upstream_message(
&mut upstream.socket,
WreqWsMessage::text(first_event),
plan_usage_permit,
record_upstream_request_state,
)
.await
.map_err(|error| match error {
ResponsesWebSocketUpstreamSendError::PlanUsagePermitLost => {
ResponsesWebSocketUpstreamBindError::PlanUsagePermitLost
}
ResponsesWebSocketUpstreamSendError::Transport(_) => {
ResponsesWebSocketUpstreamBindError::Transport(
"responses_websocket_initial_send_failed",
)
}
})?;
let client_model = initial_event
.get("model")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.ok_or("responses_websocket_model_missing")?
.ok_or(ResponsesWebSocketUpstreamBindError::Transport(
"responses_websocket_model_missing",
))?
.to_string();
let provider_model = decision
.provider_request_body
@@ -107,7 +163,9 @@ async fn bind_responses_upstream_inner(
.map(str::trim)
.filter(|value| !value.is_empty())
})
.ok_or("responses_websocket_mapped_model_missing")?
.ok_or(ResponsesWebSocketUpstreamBindError::Transport(
"responses_websocket_mapped_model_missing",
))?
.to_string();
let responses_lite_static_config =
@@ -139,6 +197,64 @@ async fn bind_responses_upstream_inner(
})
}
pub(super) async fn send_responses_websocket_upstream_message<F>(
upstream: &mut wreq::ws::WebSocket,
message: WreqWsMessage,
plan_usage_permit: Option<&aether_runtime::AdmissionPermit>,
record_upstream_request_state: F,
) -> Result<(), ResponsesWebSocketUpstreamSendError>
where
F: FnMut(UpstreamRequestState),
{
let mut record_upstream_request_state = record_upstream_request_state;
if !responses_websocket_plan_permit_is_healthy(plan_usage_permit) {
return Err(ResponsesWebSocketUpstreamSendError::PlanUsagePermitLost);
}
feed_upstream_message(upstream, message)
.await
.map_err(ResponsesWebSocketUpstreamSendError::Transport)?;
// No await between successful start_send and the state transition: an
// outer supervisor cannot cancel this attempt in the ambiguity window.
record_upstream_request_state(UpstreamRequestState::PossiblySent);
flush_upstream_messages(upstream)
.await
.map_err(ResponsesWebSocketUpstreamSendError::Transport)?;
record_upstream_request_state(UpstreamRequestState::Sent);
Ok(())
}
#[cfg(test)]
async fn complete_responses_websocket_upstream_send<Q, F, R>(
queue: Q,
flush: F,
plan_usage_permit: Option<&aether_runtime::AdmissionPermit>,
mut record_upstream_request_state: R,
) -> Result<(), ResponsesWebSocketUpstreamSendError>
where
Q: std::future::Future<Output = Result<(), WebSocketWriteError>>,
F: std::future::Future<Output = Result<(), WebSocketWriteError>>,
R: FnMut(UpstreamRequestState),
{
// Concurrency lease health is an admission condition, so check it before
// the socket writer is polled. Once polling begins, SinkExt::send may have
// completed start_send while still waiting for flush; cancelling it at
// that point cannot prove the provider did not receive response.create.
if !responses_websocket_plan_permit_is_healthy(plan_usage_permit) {
return Err(ResponsesWebSocketUpstreamSendError::PlanUsagePermitLost);
}
queue
.await
.map_err(ResponsesWebSocketUpstreamSendError::Transport)?;
// No await between successful start_send and the state transition: an
// outer supervisor cannot cancel this attempt in the ambiguity window.
record_upstream_request_state(UpstreamRequestState::PossiblySent);
flush
.await
.map_err(ResponsesWebSocketUpstreamSendError::Transport)?;
record_upstream_request_state(UpstreamRequestState::Sent);
Ok(())
}
pub(super) async fn receive_optional_upstream(
upstream: &mut Option<wreq::ws::WebSocket>,
) -> Option<Result<WreqWsMessage, ()>> {
@@ -181,13 +297,175 @@ pub(super) fn decision_bound_upstream_change_fields(
#[cfg(test)]
mod tests {
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::task::{Context, Poll};
use std::time::Duration;
use aether_contracts::ExecutionTimeouts;
use crate::ai_serving::AiExecutionDecision;
use super::{resolve_upstream_handshake_deadline, DEFAULT_UPSTREAM_HANDSHAKE_DEADLINE_MS};
use super::{
complete_responses_websocket_upstream_send, resolve_upstream_handshake_deadline,
ResponsesWebSocketUpstreamBindError, DEFAULT_UPSTREAM_HANDSHAKE_DEADLINE_MS,
};
use crate::handlers::proxy::websocket::responses::turn::UpstreamRequestState;
use crate::handlers::proxy::websocket::transport::WebSocketWriteError;
struct ReadySendThatRequestsSupervisorCancellation {
cancellation_requested: Arc<AtomicBool>,
}
struct MutablePermitHealth(Arc<AtomicBool>);
impl aether_runtime::AdmissionPermitHealth for MutablePermitHealth {
fn is_healthy(&self) -> bool {
self.0.load(Ordering::Acquire)
}
}
impl Future for ReadySendThatRequestsSupervisorCancellation {
type Output = Result<(), WebSocketWriteError>;
fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
// Model the socket writer waking an already-running outer
// supervisor in the same poll in which it reports a successful
// provider transfer. The supervisor cannot drop this future until
// control returns from this poll.
self.cancellation_requested.store(true, Ordering::Release);
Poll::Ready(Ok(()))
}
}
#[tokio::test]
async fn successful_send_marks_attempt_before_returning_to_cancelling_supervisor() {
let cancellation_requested = Arc::new(AtomicBool::new(false));
let upstream_request_sent = Arc::new(AtomicBool::new(false));
let state_by_handoff = Arc::clone(&upstream_request_sent);
let result = complete_responses_websocket_upstream_send(
ReadySendThatRequestsSupervisorCancellation {
cancellation_requested: Arc::clone(&cancellation_requested),
},
async { Ok(()) },
None,
move |state| {
if state == UpstreamRequestState::Sent {
state_by_handoff.store(true, Ordering::Release);
}
},
)
.await;
assert_eq!(result, Ok(()));
assert!(cancellation_requested.load(Ordering::Acquire));
assert!(
upstream_request_sent.load(Ordering::Acquire),
"a successful provider write must transfer lifecycle ownership before an outer supervisor can cancel the send future"
);
}
#[tokio::test]
async fn unhealthy_plan_permit_rejects_before_polling_the_upstream_send() {
let healthy = Arc::new(AtomicBool::new(false));
let permit =
aether_runtime::AdmissionPermit::from_parts(None, Some(MutablePermitHealth(healthy)))
.expect("test permit");
let send_polled = Arc::new(AtomicBool::new(false));
let send_polled_by_future = Arc::clone(&send_polled);
let result = complete_responses_websocket_upstream_send(
async move {
send_polled_by_future.store(true, Ordering::Release);
Ok(())
},
async { Ok(()) },
Some(&permit),
|_| {},
)
.await;
assert_eq!(
result,
Err(super::ResponsesWebSocketUpstreamSendError::PlanUsagePermitLost)
);
assert!(!send_polled.load(Ordering::Acquire));
}
#[tokio::test]
async fn plan_permit_loss_after_send_start_does_not_cancel_the_send() {
let healthy = Arc::new(AtomicBool::new(true));
let permit = aether_runtime::AdmissionPermit::from_parts(
None,
Some(MutablePermitHealth(Arc::clone(&healthy))),
)
.expect("test permit");
let flush_started = Arc::new(tokio::sync::Notify::new());
let allow_flush_to_finish = Arc::new(tokio::sync::Notify::new());
let flush_started_by_future = Arc::clone(&flush_started);
let allow_flush_to_finish_by_future = Arc::clone(&allow_flush_to_finish);
let marked_possibly_sent = Arc::new(AtomicBool::new(false));
let marked_possibly_sent_by_callback = Arc::clone(&marked_possibly_sent);
let marked_sent = Arc::new(AtomicBool::new(false));
let marked_sent_by_callback = Arc::clone(&marked_sent);
let send = complete_responses_websocket_upstream_send(
async { Ok(()) },
async move {
flush_started_by_future.notify_one();
allow_flush_to_finish_by_future.notified().await;
Ok(())
},
Some(&permit),
move |state| match state {
UpstreamRequestState::PossiblySent => {
marked_possibly_sent_by_callback.store(true, Ordering::Release)
}
UpstreamRequestState::Sent => {
marked_sent_by_callback.store(true, Ordering::Release)
}
UpstreamRequestState::NotStarted => {}
},
);
let revoke = async move {
flush_started.notified().await;
healthy.store(false, Ordering::Release);
allow_flush_to_finish.notify_one();
};
let (result, ()) = tokio::join!(send, revoke);
assert_eq!(result, Ok(()));
assert!(marked_possibly_sent.load(Ordering::Acquire));
assert!(marked_sent.load(Ordering::Acquire));
}
#[tokio::test]
async fn flush_failure_keeps_the_attempt_possibly_sent() {
let observed = Arc::new(std::sync::Mutex::new(Vec::new()));
let observed_by_callback = Arc::clone(&observed);
let result = complete_responses_websocket_upstream_send(
async { Ok(()) },
async { Err(WebSocketWriteError::TimedOut) },
None,
move |state| observed_by_callback.lock().expect("states").push(state),
)
.await;
assert_eq!(
result,
Err(super::ResponsesWebSocketUpstreamSendError::Transport(
WebSocketWriteError::TimedOut
))
);
assert_eq!(
*observed.lock().expect("states"),
vec![UpstreamRequestState::PossiblySent]
);
}
fn sample_decision() -> AiExecutionDecision {
AiExecutionDecision {
@@ -336,12 +614,16 @@ mod tests {
ResponsesWebSocketBodyNormalization::for_tests("test-model"),
&json!({"type": "response.create", "model": "test-model"}),
adapter,
None,
|_| {},
)
.await;
assert_eq!(
result.err().expect("bind should fail with timeout"),
"responses_websocket_upstream_handshake_timeout"
ResponsesWebSocketUpstreamBindError::Transport(
"responses_websocket_upstream_handshake_timeout"
)
);
}
}