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-09 10:57:03 +08:00 · 2026-08-17 14:50:33 +08:00
parent 9a0d346ff3
commit 71b54070e8
72 changed files with 10441 additions and 108 deletions
@@ -377,11 +377,13 @@ pub struct AppState {
    pub(crate) frontdoor_runtime_guards: Arc<FrontdoorRuntimeGuardConfig>,
    pub(crate) request_body_buffer_budget: Arc<Semaphore>,
    pub(crate) request_gate: Option<Arc<ConcurrencyGate>>,
+    pub(crate) websocket_connection_gate: Option<Arc<ConcurrencyGate>>,
    pub(crate) auth_snapshot_load_gate: Option<Arc<ConcurrencyGate>>,
    pub(crate) candidate_planning_gate: Option<Arc<ConcurrencyGate>>,
    pub(crate) upstream_execution_gate: Option<Arc<ConcurrencyGate>>,
    pub(crate) upstream_target_admission: Arc<crate::upstream_admission::UpstreamTargetAdmission>,
    pub(crate) distributed_request_gate: Option<Arc<RuntimeSemaphore>>,
+    pub(crate) distributed_websocket_connection_gate: Option<Arc<RuntimeSemaphore>>,
    pub(crate) client: reqwest::Client,
    pub(crate) owner_forward_client: reqwest::Client,
    pub(crate) auth_context_cache: Arc<AuthContextCache>,
@@ -325,6 +325,7 @@ impl AppState {
                frontdoor_runtime_guards.request_body_buffer_budget_permits,
            )),
            request_gate: None,
+            websocket_connection_gate: None,
            auth_snapshot_load_gate: frontdoor_runtime_guards
                .auth_snapshot_load_gate_limit
                .map(|limit| Arc::new(ConcurrencyGate::new("gateway_auth_snapshot_load", limit))),
@@ -341,6 +342,7 @@ impl AppState {
                ),
            ),
            distributed_request_gate: None,
+            distributed_websocket_connection_gate: None,
            client,
            owner_forward_client,
            auth_context_cache: Arc::new(AuthContextCache::default()),
@@ -574,8 +576,20 @@ impl AppState {
    }

    pub fn with_request_concurrency_limit(mut self, limit: usize) -> Self {
-        self.request_gate = Some(Arc::new(ConcurrencyGate::new(
-            "gateway_requests",
+        let limit = limit.max(1);
+        self.request_gate = Some(Arc::new(ConcurrencyGate::new("gateway_requests", limit)));
+        if self.websocket_connection_gate.is_none() {
+            self.websocket_connection_gate = Some(Arc::new(ConcurrencyGate::new(
+                "gateway_websocket_connections",
+                limit,
+            )));
+        }
+        self
+    }
+
+    pub fn with_websocket_connection_limit(mut self, limit: usize) -> Self {
+        self.websocket_connection_gate = Some(Arc::new(ConcurrencyGate::new(
+            "gateway_websocket_connections",
            limit.max(1),
        )));
        self
@@ -646,6 +660,11 @@ impl AppState {
        self
    }

+    pub fn with_distributed_websocket_connection_gate(mut self, gate: RuntimeSemaphore) -> Self {
+        self.distributed_websocket_connection_gate = Some(Arc::new(gate));
+        self
+    }
+
    pub fn with_frontdoor_cors_config(mut self, config: FrontdoorCorsConfig) -> Self {
        self.frontdoor_cors = Some(Arc::new(config));
        self
@@ -1224,6 +1243,12 @@ impl AppState {
        self.request_gate.as_ref().map(|gate| gate.snapshot())
    }

+    pub(crate) fn websocket_connection_concurrency_snapshot(&self) -> Option<ConcurrencySnapshot> {
+        self.websocket_connection_gate
+            .as_ref()
+            .map(|gate| gate.snapshot())
+    }
+
    pub(crate) fn auth_snapshot_load_concurrency_snapshot(&self) -> Option<ConcurrencySnapshot> {
        self.auth_snapshot_load_gate
            .as_ref()
@@ -1251,6 +1276,15 @@ impl AppState {
        }
    }

+    pub(crate) async fn distributed_websocket_connection_concurrency_snapshot(
+        &self,
+    ) -> Result<Option<RuntimeSemaphoreSnapshot>, RuntimeSemaphoreError> {
+        match self.distributed_websocket_connection_gate.as_ref() {
+            Some(gate) => gate.snapshot().await.map(Some),
+            None => Ok(None),
+        }
+    }
+
    pub(crate) async fn metric_samples(&self) -> Vec<MetricSample> {
        let now = std::time::Instant::now();
        let snapshot = self.metric_snapshot.read().await.clone();
@@ -1557,6 +1591,9 @@ impl AppState {
        if let Some(snapshot) = self.request_concurrency_snapshot() {
            samples.extend(snapshot.to_metric_samples("gateway_requests"));
        }
+        if let Some(snapshot) = self.websocket_connection_concurrency_snapshot() {
+            samples.extend(snapshot.to_metric_samples("gateway_websocket_connections"));
+        }
        if let Some(snapshot) = self.auth_snapshot_load_concurrency_snapshot() {
            samples.extend(snapshot.to_metric_samples("gateway_auth_snapshot_load"));
        }
@@ -1600,6 +1637,28 @@ impl AppState {
                )])],
            }
        };
+        let distributed_websocket_connection_metrics = async {
+            let Some(gate) = self.distributed_websocket_connection_gate.as_ref() else {
+                return Vec::new();
+            };
+            match tokio::time::timeout(DISTRIBUTED_CONCURRENCY_METRICS_TIMEOUT, gate.snapshot())
+                .await
+            {
+                Ok(Ok(snapshot)) => {
+                    snapshot.to_metric_samples("gateway_websocket_connections_distributed")
+                }
+                Ok(Err(_)) | Err(_) => vec![MetricSample::new(
+                    "concurrency_unavailable",
+                    "Whether the distributed concurrency gate is currently unavailable.",
+                    MetricKind::Gauge,
+                    1,
+                )
+                .with_labels(vec![MetricLabel::new(
+                    "gate",
+                    "gateway_websocket_connections_distributed",
+                )])],
+            }
+        };
        let postgres_observability_metrics = async {
            match tokio::time::timeout(
                POSTGRES_OBSERVABILITY_METRICS_TIMEOUT,
@@ -1649,6 +1708,7 @@ impl AppState {
            );
        let (
            distributed_request_metrics,
+            distributed_websocket_connection_metrics,
            postgres_observability_metrics,
            postgres_activity_group_metrics,
            redis_runtime_metrics,
@@ -1656,6 +1716,7 @@ impl AppState {
            usage_counter_pending_health_metrics,
        ) = tokio::join!(
            distributed_request_metrics,
+            distributed_websocket_connection_metrics,
            postgres_observability_metrics,
            postgres_activity_group_metrics,
            redis_runtime_metrics,
@@ -1663,6 +1724,7 @@ impl AppState {
            usage_counter_pending_health_metrics,
        );
        samples.extend(distributed_request_metrics);
+        samples.extend(distributed_websocket_connection_metrics);
        samples.extend(postgres_observability_metrics);
        samples.extend(postgres_activity_group_metrics);
        samples.extend(redis_runtime_metrics);
@@ -1783,6 +1845,26 @@ impl AppState {
        Ok(AdmissionPermit::from_parts(local, distributed))
    }

+    pub(crate) async fn try_acquire_websocket_connection_permit(
+        &self,
+    ) -> Result<Option<AdmissionPermit>, RequestAdmissionError> {
+        let local = self
+            .websocket_connection_gate
+            .as_ref()
+            .map(|gate| gate.try_acquire())
+            .transpose()
+            .map_err(RequestAdmissionError::Local)?;
+        let distributed = match self.distributed_websocket_connection_gate.as_ref() {
+            Some(gate) => Some(
+                gate.try_acquire()
+                    .await
+                    .map_err(RequestAdmissionError::Distributed)?,
+            ),
+            None => None,
+        };
+        Ok(AdmissionPermit::from_parts(local, distributed))
+    }
+
    pub fn has_auth_api_key_data_reader(&self) -> bool {
        self.data.has_auth_api_key_reader()
    }