feat(gateway): Codex/OpenAI Responses WebSocket 代理模式

在 /v1/responses 上支持 WebSocket 升级,把客户端帧中继到上游 Codex /
OpenAI Responses WebSocket 端点,同时保持既有的路由、鉴权、配额与用量
语义:

- 路由与准入:control/route/ai.rs 识别 WebSocket 升级请求;
  websocket/ingress.rs 复用 API Key 鉴权、IP 规则与并发许可,并引入
  独立的 WebSocket 连接许可
- 中继:websocket/responses/* 按 connection / session / turn 分层,
  帧解析归一化、socket 写入有界、continuation 保持调度亲和性
- 配额:orchestration/codex_quota_breaker.rs 在账号配额耗尽时熔断并
  自动恢复,不再直接断开客户端连接
- 用量:每个 turn 的终态用量落库,request_metadata 记录
  websocket_mode / websocket_transport,管理端与 usage 视图暴露
  is_websocket
- 管理端:provider 可配置 Responses WebSocket 开关
This commit is contained in:
AAEE86
2026-08-17 14:50:33 +08:00
committed by ZheFox
parent 9a0d346ff3
commit 71b54070e8
72 changed files with 10441 additions and 108 deletions
@@ -80,8 +80,9 @@ pub(crate) use self::standard::{
build_local_stream_plan_and_reports as build_standard_family_stream_plan_and_reports,
build_local_sync_attempt_source as build_standard_family_sync_attempt_source,
build_local_sync_plan_and_reports as build_standard_family_sync_plan_and_reports,
codex_model_capabilities_for_transport, set_local_openai_chat_execution_exhausted_diagnostic,
validate_final_openai_provider_request,
codex_model_capabilities_for_transport, maybe_build_responses_websocket_decision,
set_local_openai_chat_execution_exhausted_diagnostic, validate_final_openai_provider_request,
ResponsesWebSocketBodyNormalization, ResponsesWebSocketDecision,
};
pub(crate) use self::state::{
GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth,
@@ -42,13 +42,14 @@ pub(crate) use self::openai::{
build_local_openai_responses_sync_attempt_source_for_kind,
build_local_openai_responses_sync_plan_and_reports_for_kind, copy_request_number_field,
copy_request_number_field_as, map_openai_reasoning_effort_to_claude_output,
map_openai_reasoning_effort_to_gemini_budget, maybe_build_stream_local_decision_payload,
map_openai_reasoning_effort_to_gemini_budget, maybe_build_responses_websocket_decision,
maybe_build_stream_local_decision_payload,
maybe_build_stream_local_openai_responses_decision_payload,
maybe_build_sync_local_decision_payload,
maybe_build_sync_local_openai_embedding_decision_payload,
maybe_build_sync_local_openai_responses_decision_payload, parse_openai_stop_sequences,
resolve_openai_chat_max_tokens, set_local_openai_chat_execution_exhausted_diagnostic,
value_as_u64,
value_as_u64, ResponsesWebSocketBodyNormalization, ResponsesWebSocketDecision,
};
pub(crate) use crate::ai_serving::normalize_standard_request_to_openai_chat_request;
pub(crate) use crate::ai_serving::{
@@ -23,6 +23,8 @@ pub(crate) use responses::{
build_local_openai_responses_stream_plan_and_reports_for_kind,
build_local_openai_responses_sync_attempt_source_for_kind,
build_local_openai_responses_sync_plan_and_reports_for_kind,
maybe_build_responses_websocket_decision,
maybe_build_stream_local_openai_responses_decision_payload,
maybe_build_sync_local_openai_responses_decision_payload,
maybe_build_sync_local_openai_responses_decision_payload, ResponsesWebSocketBodyNormalization,
ResponsesWebSocketDecision,
};
@@ -1,6 +1,16 @@
use crate::ai_serving::planner::common::endpoint_config_forces_body_stream_field;
use crate::ai_serving::planner::plan_builders::{AiStreamAttempt, AiSyncAttempt};
use crate::ai_serving::planner::spec_metadata::local_openai_responses_spec_metadata;
use crate::ai_serving::planner::standard::codex::codex_model_capabilities_for_transport;
use crate::ai_serving::planner::standard::normalize::build_local_openai_responses_request_body_with_codex_model_capabilities;
use crate::ai_serving::GatewayControlDecision;
use crate::orchestration::{
codex_quota_breaker_blocks_candidate, log_codex_quota_breaker_check_failure,
responses_websocket_adapter, ResponsesWebSocketAdapter,
};
use crate::{AiExecutionDecision, AppState, GatewayError};
use aether_runtime_state::RuntimeLockLease;
use std::collections::BTreeSet;
mod decision;
mod plans;
@@ -165,3 +175,314 @@ pub(crate) async fn maybe_build_stream_local_openai_responses_decision_payload(
Ok(None)
}
/// One eligible upstream plus the adapter that is allowed to speak to it.
///
/// The adapter is selected from the provider-scoped capability before the
/// decision leaves the planner. This prevents a public Responses socket from
/// choosing an arbitrary provider protocol after scheduling has completed.
pub(crate) struct ResponsesWebSocketDecision {
pub(crate) execution: AiExecutionDecision,
pub(crate) adapter: ResponsesWebSocketAdapter,
pub(crate) normalization: ResponsesWebSocketBodyNormalization,
}
/// Everything needed to re-run provider-body normalization for the candidate a
/// socket is already bound to.
///
/// A continuation turn (`previous_response_id` on the bound upstream) cannot
/// re-enter the planner, because planning selects a candidate and a different
/// key would break the response chain. Without this, such turns reached the
/// provider with only their `model` rewritten — skipping model directives,
/// endpoint body rules, and the Codex body contract that turn 1 received.
///
/// This value holds cloned scalars and JSON only: no candidate, no pool key
/// lease, no `AppState`. It cannot influence selection.
#[derive(Debug, Clone)]
pub(crate) struct ResponsesWebSocketBodyNormalization {
provider_type: String,
provider_api_format: String,
client_api_format: String,
mapped_model: String,
requested_model: String,
upstream_is_stream: bool,
force_body_stream_field: bool,
body_rules: Option<serde_json::Value>,
request_headers: http::HeaderMap,
codex_model_capabilities: Option<crate::ai_serving::CodexResponsesModelCapabilities>,
model_directive_patch: Option<serde_json::Value>,
}
impl ResponsesWebSocketBodyNormalization {
/// Builds a normalizer for a plain `openai:responses` upstream with no
/// endpoint body rules, directives or Codex capabilities, so relay tests can
/// construct a bound connection without standing up a provider snapshot.
#[cfg(test)]
pub(crate) fn for_tests(mapped_model: &str) -> Self {
Self {
provider_type: "openai".to_string(),
provider_api_format: "openai:responses".to_string(),
client_api_format: "openai:responses".to_string(),
mapped_model: mapped_model.to_string(),
requested_model: mapped_model.to_string(),
upstream_is_stream: true,
force_body_stream_field: false,
body_rules: None,
request_headers: http::HeaderMap::new(),
codex_model_capabilities: None,
model_directive_patch: None,
}
}
#[cfg(test)]
pub(crate) fn with_provider_type_for_tests(mut self, provider_type: &str) -> Self {
self.provider_type = provider_type.to_string();
self
}
#[cfg(test)]
pub(crate) fn with_model_directive_patch_for_tests(mut self, patch: serde_json::Value) -> Self {
self.model_directive_patch = Some(patch);
self
}
/// Applies the same body transformations the planner applied on the turn
/// that bound this upstream.
///
/// Mirrors the same-format branch of
/// `resolve_local_openai_responses_candidate_payload_parts`. The
/// cross-format, Kiro, Windsurf and Antigravity branches are unreachable
/// here: the WebSocket planner only returns candidates whose provider API
/// format is `openai:responses`.
///
/// Returns `None` when normalization fails, leaving the caller to fall back
/// to the unnormalized event — a continuation cannot re-select a candidate,
/// so failing the turn outright would be worse than sending it as-is.
pub(crate) fn normalize_response_create(
&self,
client_event: &serde_json::Value,
) -> Option<serde_json::Value> {
use crate::ai_serving::planner::common::{
enforce_provider_body_stream_policy, request_requires_body_stream_field,
};
let source_model = client_event
.get("model")
.and_then(serde_json::Value::as_str)
.unwrap_or(self.requested_model.as_str());
let require_body_stream_field =
request_requires_body_stream_field(client_event, self.force_body_stream_field);
let mut body = build_local_openai_responses_request_body_with_codex_model_capabilities(
client_event,
&self.mapped_model,
self.upstream_is_stream,
self.force_body_stream_field,
self.provider_type.as_str(),
self.provider_api_format.as_str(),
self.body_rules.as_ref(),
&self.request_headers,
self.codex_model_capabilities.as_ref(),
false,
)?;
if let Some(patch) = self.model_directive_patch.as_ref() {
crate::ai_serving::apply_model_directive_mapping_patch(&mut body, patch);
// The patch is a deep merge and may reintroduce `stream`.
enforce_provider_body_stream_policy(
&mut body,
self.provider_api_format.as_str(),
self.upstream_is_stream,
require_body_stream_field,
);
}
crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities(
&mut body,
crate::ai_serving::OpenAiProviderRequestFinalization {
source_api_format: self.client_api_format.as_str(),
provider_api_format: self.provider_api_format.as_str(),
provider_type: self.provider_type.as_str(),
provider_model: self.mapped_model.as_str(),
source_model,
body_rules: self.body_rules.as_ref(),
upstream_is_stream: self.upstream_is_stream,
require_body_stream_field,
},
self.codex_model_capabilities.as_ref(),
)
.ok()?;
Some(body)
}
}
/// Builds one upstream decision for a Responses WebSocket turn. The session
/// reuses this decision for same-model turns and invokes the planner again when
/// a later `response.create` changes the public model.
pub(crate) async fn maybe_build_responses_websocket_decision(
state: &AppState,
parts: &http::request::Parts,
trace_id: &str,
decision: &GatewayControlDecision,
body_json: &serde_json::Value,
excluded_key_ids: Option<&BTreeSet<String>>,
excluded_codex_account_ids: Option<&BTreeSet<String>>,
) -> Result<Option<ResponsesWebSocketDecision>, GatewayError> {
let Some(spec) = resolve_stream_spec(crate::ai_serving::OPENAI_RESPONSES_STREAM_PLAN_KIND)
else {
return Ok(None);
};
let Some(input) = resolve_local_openai_responses_decision_input(
state,
parts,
trace_id,
decision,
body_json,
spec.decision_kind,
)
.await?
else {
return Ok(None);
};
let body_json = input.effective_body_json(body_json);
let (mut source, _) = build_local_openai_responses_candidate_attempt_source(
state, trace_id, &input, body_json, spec,
)
.await?;
while let Some(attempt) = source.next_attempt().await? {
let pool_key_lease = attempt.eligible.orchestration.pool_key_lease.clone();
if excluded_key_ids
.is_some_and(|key_ids| key_ids.contains(attempt.eligible.candidate.key_id.as_str()))
{
release_responses_websocket_planning_lease(state, pool_key_lease.as_ref()).await;
continue;
}
let Some(adapter) = responses_websocket_adapter(
&attempt.eligible.transport.provider.provider_type,
attempt.eligible.transport.provider.config.as_ref(),
) else {
release_responses_websocket_planning_lease(state, pool_key_lease.as_ref()).await;
continue;
};
// Captured before `attempt` is consumed so a later continuation turn can
// reproduce this candidate's body normalization without re-planning.
let transport = std::sync::Arc::clone(&attempt.eligible.transport);
let candidate_provider_api_format = attempt.eligible.provider_api_format.clone();
let payload = match maybe_build_local_openai_responses_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)
.await
{
Ok(Some(payload)) => payload,
Ok(None) => {
release_responses_websocket_planning_lease(state, pool_key_lease.as_ref()).await;
continue;
}
Err(error) => {
release_responses_websocket_planning_lease(state, pool_key_lease.as_ref()).await;
return Err(error);
}
};
if payload
.provider_type
.as_deref()
.is_some_and(|value| value.trim().eq_ignore_ascii_case("codex"))
&& crate::orchestration::codex_account_id_from_headers(
&payload.provider_request_headers,
)
.is_some_and(|account_id| {
excluded_codex_account_ids
.is_some_and(|account_ids| account_ids.contains(account_id))
})
{
release_responses_websocket_planning_lease(state, pool_key_lease.as_ref()).await;
continue;
}
match codex_quota_breaker_blocks_candidate(
state,
payload.provider_type.as_deref(),
payload.key_id.as_deref(),
&payload.provider_request_headers,
)
.await
{
Ok(true) => {
release_responses_websocket_planning_lease(state, pool_key_lease.as_ref()).await;
continue;
}
Ok(false) => {}
Err(error) => log_codex_quota_breaker_check_failure(&error),
}
if payload
.provider_type
.as_deref()
.is_some_and(|value| adapter.supports_provider_type(value))
&& payload.provider_api_format.as_deref().is_some_and(|value| {
crate::ai_serving::normalize_api_format_alias(value) == "openai:responses"
})
{
let mapped_model = payload.mapped_model.clone().unwrap_or_default();
let source_model = body_json
.get("model")
.and_then(serde_json::Value::as_str)
.unwrap_or(input.requested_model.as_str());
let normalization = ResponsesWebSocketBodyNormalization {
provider_type: transport.provider.provider_type.clone(),
provider_api_format: candidate_provider_api_format.clone(),
client_api_format: local_openai_responses_spec_metadata(spec)
.api_format
.to_string(),
requested_model: input.requested_model.clone(),
upstream_is_stream: payload.upstream_is_stream,
force_body_stream_field: endpoint_config_forces_body_stream_field(
transport.endpoint.config.as_ref(),
),
body_rules: transport.endpoint.body_rules.clone(),
request_headers: input.effective_headers(&parts.headers).clone(),
codex_model_capabilities: codex_model_capabilities_for_transport(
&transport,
candidate_provider_api_format.as_str(),
mapped_model.as_str(),
source_model,
),
model_directive_patch: input
.model_directive_policy
.resolve_reasoning(
candidate_provider_api_format.as_str(),
Some(&input.requested_model),
)
.mapping_patch_for_mapped_model(mapped_model.as_str())
.ok()
.flatten(),
mapped_model,
};
return Ok(Some(ResponsesWebSocketDecision {
execution: payload,
adapter,
normalization,
}));
}
release_responses_websocket_planning_lease(state, pool_key_lease.as_ref()).await;
}
Ok(None)
}
async fn release_responses_websocket_planning_lease(
state: &AppState,
lease: Option<&RuntimeLockLease>,
) {
let Some(lease) = lease else {
return;
};
if let Err(error) =
crate::handlers::shared::provider_pool::release_admin_provider_pool_key_lease(
state.runtime_state.as_ref(),
lease,
)
.await
{
tracing::warn!(
error = ?error,
"gateway Responses WebSocket planner failed to release an unused pool key lease"
);
}
}