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 开关
2026-10-12 20:19:51 +08:00 · 2026-08-17 14:50:33 +08:00
parent 9a0d346ff3
commit 71b54070e8
72 changed files with 10441 additions and 108 deletions
@@ -0,0 +1,141 @@
+//! Mutable state owned by one Responses WebSocket connection.
+//!
+//! The session loop is intentionally kept separate from these containers.  A
+//! connection may survive many `response.create` turns, while the turn
+//! lifecycle and upstream binding are replaced independently.
+
+use serde_json::Value;
+use std::collections::{BTreeMap, BTreeSet};
+use tokio::task::JoinHandle;
+
+use super::adapter::{ResponsesWebSocketDrainDirective, ResponsesWebSocketProtocolAdapter};
+use super::binding::UpstreamBindingIdentity;
+use super::lifecycle::ActiveResponsesWebSocketTurn;
+use super::request::response_create_has_previous_response_id;
+use crate::ai_serving::{AiExecutionDecision, ResponsesWebSocketBodyNormalization};
+
+const EXHAUSTED_KEY_EXCLUSION_FALLBACK_SECONDS: u64 = 300;
+
+/// All mutable state associated with the physical upstream connection.
+pub(super) struct BoundResponsesConnection {
+    pub(super) upstream: Option<wreq::ws::WebSocket>,
+    pub(super) adapter: &'static dyn ResponsesWebSocketProtocolAdapter,
+    pub(super) client_model: String,
+    pub(super) provider_model: String,
+    pub(super) response_in_flight: bool,
+    pub(super) decision_template: AiExecutionDecision,
+    /// Reproduces this binding's provider-body normalization for continuation
+    /// turns, which must not re-enter the planner. Replaced whenever the
+    /// binding or its decision is replaced.
+    pub(super) body_normalization: ResponsesWebSocketBodyNormalization,
+    pub(super) binding_identity: UpstreamBindingIdentity,
+    pub(super) active_turn: Option<ActiveResponsesWebSocketTurn>,
+    pub(super) active_response_create: Option<ActiveResponsesWebSocketRequest>,
+    pub(super) next_turn_index: u64,
+    pub(super) upstream_response_headers: BTreeMap<String, String>,
+    pub(super) pending_adapter_drain: Option<ResponsesWebSocketDrainDirective>,
+    pub(super) pending_adapter_observation: Option<JoinHandle<()>>,
+    pub(super) exhausted_exclusions: ExhaustedResponsesWebSocketExclusions,
+    pub(super) pending_turn_finalization: Option<JoinHandle<()>>,
+}
+
+/// Connection-local fallback in addition to the distributed account breaker.
+/// A key and its provider account are excluded until the upstream's reset
+/// deadline (or a short fallback when the terminal payload lacks one), so an
+/// unusually long-lived client socket does not keep it unavailable after the
+/// quota has recovered.
+#[derive(Debug, Default)]
+pub(super) struct ExhaustedResponsesWebSocketExclusions {
+    expires_at_by_key: BTreeMap<String, u64>,
+    expires_at_by_codex_account: BTreeMap<String, u64>,
+}
+
+impl ExhaustedResponsesWebSocketExclusions {
+    pub(super) fn exclude(
+        &mut self,
+        key_id: String,
+        codex_account_id: Option<String>,
+        reset_at_unix_secs: Option<u64>,
+        now_unix_secs: u64,
+    ) -> u64 {
+        self.prune(now_unix_secs);
+        let requested_expiry = reset_at_unix_secs
+            .filter(|reset_at| *reset_at > now_unix_secs)
+            .unwrap_or_else(|| {
+                now_unix_secs.saturating_add(EXHAUSTED_KEY_EXCLUSION_FALLBACK_SECONDS)
+            });
+        let expiry = self
+            .expires_at_by_key
+            .entry(key_id)
+            .and_modify(|existing| *existing = (*existing).max(requested_expiry))
+            .or_insert(requested_expiry);
+        if let Some(account_id) = codex_account_id {
+            self.expires_at_by_codex_account
+                .entry(account_id)
+                .and_modify(|existing| *existing = (*existing).max(requested_expiry))
+                .or_insert(requested_expiry);
+        }
+        *expiry
+    }
+
+    pub(super) fn codex_account_ids(&mut self, now_unix_secs: u64) -> BTreeSet<String> {
+        self.prune(now_unix_secs);
+        self.expires_at_by_codex_account.keys().cloned().collect()
+    }
+
+    pub(super) fn key_ids(&mut self, now_unix_secs: u64) -> BTreeSet<String> {
+        self.prune(now_unix_secs);
+        self.expires_at_by_key.keys().cloned().collect()
+    }
+
+    pub(super) fn len(&mut self, now_unix_secs: u64) -> usize {
+        self.prune(now_unix_secs);
+        self.expires_at_by_key.len() + self.expires_at_by_codex_account.len()
+    }
+
+    fn prune(&mut self, now_unix_secs: u64) {
+        self.expires_at_by_key
+            .retain(|_, expires_at| *expires_at > now_unix_secs);
+        self.expires_at_by_codex_account
+            .retain(|_, expires_at| *expires_at > now_unix_secs);
+    }
+}
+
+#[derive(Debug, Clone)]
+pub(super) struct ActiveResponsesWebSocketRequest {
+    pub(super) client_event: Value,
+    pub(super) turn_index: u64,
+    pub(super) logical_turn_id: String,
+    pub(super) turn_attempt: u32,
+    pub(super) retry_attempted: bool,
+    pub(super) retry_unsafe_reason: Option<&'static str>,
+}
+
+impl ActiveResponsesWebSocketRequest {
+    pub(super) fn new(client_event: Value, turn_index: u64, logical_turn_id: String) -> Self {
+        Self {
+            client_event,
+            turn_index,
+            logical_turn_id,
+            turn_attempt: 1,
+            retry_attempted: false,
+            retry_unsafe_reason: None,
+        }
+    }
+
+    pub(super) fn quota_retry_block_reason(&self) -> Option<&'static str> {
+        if self.retry_attempted {
+            Some("quota_retry_already_attempted")
+        } else if let Some(reason) = self.retry_unsafe_reason {
+            Some(reason)
+        } else if response_create_has_previous_response_id(&self.client_event) {
+            Some("previous_response_id")
+        } else {
+            None
+        }
+    }
+
+    pub(super) fn mark_retry_unsafe(&mut self, reason: &'static str) {
+        self.retry_unsafe_reason.get_or_insert(reason);
+    }
+}