fix(ws): harden Responses connection lifecycle

Revalidate control policy per turn, isolate downstream credentials, and make planner/turn ownership cancellation-safe. Preserve opaque protocol events, align configurable timeout semantics, and extend end-to-end security and settlement coverage.
2026-10-12 12:09:50 +08:00 · 2026-08-17 18:50:29 +08:00
parent 4a0775c4ea
commit c8118edf36
42 changed files with 4017 additions and 1276 deletions
@@ -0,0 +1,245 @@
+//! Cancellation-safe ownership handoff for WebSocket planning leases.
+//!
+//! The relay races every turn against connection and response deadlines. A
+//! planner future therefore cannot directly own a distributed pool-key lease:
+//! losing the race would drop the future between scheduler selection and turn
+//! startup, leaving that key unavailable until the lease TTL elapsed.
+
+use std::collections::BTreeSet;
+use std::time::Duration;
+
+use serde_json::Value;
+use tokio::task::JoinHandle;
+
+use super::lifecycle::{begin_responses_websocket_turn, ActiveProviderAttempt};
+use crate::ai_serving::{
+    maybe_build_responses_websocket_decision, AiExecutionDecision, GatewayAuthApiKeySnapshot,
+    ResponsesWebSocketDecision, ResponsesWebSocketPinnedCandidate,
+};
+use crate::control::GatewayControlDecision;
+use crate::orchestration::release_pool_key_lease_from_report_context;
+use crate::{AppState, GatewayError};
+
+/// Owns a selected pool-key lease until the attempt lifecycle has taken over
+/// the decision report context.
+pub(super) struct PlannedPoolKeyLeaseGuard {
+    state: AppState,
+    report_context: Option<Value>,
+}
+
+/// Planner output coupled to both its request parts and lease guard.
+pub(super) struct OwnedResponsesWebSocketDecision {
+    pub(super) planned: ResponsesWebSocketDecision,
+    pub(super) planning_parts: http::request::Parts,
+    pub(super) planned_lease: PlannedPoolKeyLeaseGuard,
+}
+
+/// Runs planning in an owner task. Dropping the caller's waiter detaches this
+/// task; an unobserved successful output drops its guard and releases the
+/// selected pool-key lease.
+#[allow(clippy::too_many_arguments)]
+pub(super) fn spawn_owned_responses_websocket_plan(
+    state: AppState,
+    parts: http::request::Parts,
+    trace_id: String,
+    control_decision: GatewayControlDecision,
+    auth_snapshot: Option<GatewayAuthApiKeySnapshot>,
+    client_event: Value,
+    excluded_key_ids: Option<BTreeSet<String>>,
+    excluded_codex_account_ids: Option<BTreeSet<String>>,
+    pinned_candidate: Option<ResponsesWebSocketPinnedCandidate>,
+) -> JoinHandle<Result<Option<OwnedResponsesWebSocketDecision>, GatewayError>> {
+    let owner_timeout = state
+        .frontdoor_runtime_guards
+        .local_execution_planning_timeout;
+    tokio::spawn(async move {
+        let planned = await_owned_planning_deadline(
+            maybe_build_responses_websocket_decision(
+                &state,
+                &parts,
+                &trace_id,
+                &control_decision,
+                auth_snapshot.as_ref(),
+                &client_event,
+                excluded_key_ids.as_ref(),
+                excluded_codex_account_ids.as_ref(),
+                pinned_candidate.as_ref(),
+            ),
+            owner_timeout,
+        )
+        .await
+        .map_err(|_| GatewayError::LocalExecutionPlanningTimeout {
+            trace_id: trace_id.clone(),
+            phase: "responses_websocket_plan_owner",
+            timeout_ms: owner_timeout.as_millis() as u64,
+        })??;
+
+        Ok(planned.map(|planned| {
+            let planned_lease =
+                PlannedPoolKeyLeaseGuard::new(&state, planned.execution.report_context.as_ref());
+            OwnedResponsesWebSocketDecision {
+                planned,
+                planning_parts: parts,
+                planned_lease,
+            }
+        }))
+    })
+}
+
+async fn await_owned_planning_deadline<F, T>(
+    planning: F,
+    deadline: Duration,
+) -> Result<T, tokio::time::error::Elapsed>
+where
+    F: std::future::Future<Output = T>,
+{
+    tokio::time::timeout(deadline, planning).await
+}
+
+pub(super) async fn await_owned_responses_websocket_plan(
+    handle: JoinHandle<Result<Option<OwnedResponsesWebSocketDecision>, GatewayError>>,
+) -> Result<Option<OwnedResponsesWebSocketDecision>, GatewayError> {
+    handle.await.map_err(|error| {
+        GatewayError::Internal(format!(
+            "Responses WebSocket planning task failed before ownership transfer: {error}"
+        ))
+    })?
+}
+
+impl PlannedPoolKeyLeaseGuard {
+    fn new(state: &AppState, report_context: Option<&Value>) -> Self {
+        Self {
+            state: state.clone(),
+            report_context: report_context.cloned(),
+        }
+    }
+
+    pub(super) async fn release(mut self) {
+        release_pool_key_lease_from_report_context(&self.state, self.report_context.as_ref()).await;
+        self.report_context = None;
+    }
+
+    fn disarm(&mut self) {
+        self.report_context = None;
+    }
+}
+
+impl Drop for PlannedPoolKeyLeaseGuard {
+    fn drop(&mut self) {
+        let Some(report_context) = self.report_context.take() else {
+            return;
+        };
+        let state = self.state.clone();
+        if let Ok(handle) = tokio::runtime::Handle::try_current() {
+            handle.spawn(async move {
+                release_pool_key_lease_from_report_context(&state, Some(&report_context)).await;
+            });
+        }
+    }
+}
+
+/// Keeps the planning guard in the same detached owner task as lifecycle
+/// startup. If the relay loses a deadline race while awaiting startup, the
+/// task completes the handoff (or releases the lease on failure) without a
+/// cancellation gap.
+pub(super) async fn begin_responses_websocket_turn_with_planned_lease(
+    state: &AppState,
+    trace_id: &str,
+    parts: http::request::Parts,
+    control_decision: &GatewayControlDecision,
+    decision: AiExecutionDecision,
+    client_event: &Value,
+    mut planned_lease: PlannedPoolKeyLeaseGuard,
+) -> Result<ActiveProviderAttempt, GatewayError> {
+    let state = state.clone();
+    let trace_id = trace_id.to_string();
+    let control_decision = control_decision.clone();
+    let client_event = client_event.clone();
+    tokio::spawn(async move {
+        let turn = begin_responses_websocket_turn(
+            &state,
+            &trace_id,
+            parts,
+            &control_decision,
+            decision,
+            &client_event,
+        )
+        .await?;
+        // ActiveProviderAttempt now owns the report context containing the
+        // lease. No await occurs between that handoff and disarming the guard.
+        planned_lease.disarm();
+        Ok(turn)
+    })
+    .await
+    .map_err(|error| {
+        GatewayError::Internal(format!(
+            "Responses WebSocket guarded turn startup task failed: {error}"
+        ))
+    })?
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::atomic::{AtomicUsize, Ordering};
+    use std::sync::Arc;
+    use std::time::Duration;
+
+    struct DropProbe(Arc<AtomicUsize>);
+
+    impl Drop for DropProbe {
+        fn drop(&mut self) {
+            self.0.fetch_add(1, Ordering::SeqCst);
+        }
+    }
+
+    #[tokio::test]
+    async fn dropping_a_planning_waiter_detaches_the_owner_and_drops_its_output() {
+        let started = Arc::new(tokio::sync::Notify::new());
+        let release = Arc::new(tokio::sync::Notify::new());
+        let dropped = Arc::new(AtomicUsize::new(0));
+        let task_started = Arc::clone(&started);
+        let task_release = Arc::clone(&release);
+        let task_dropped = Arc::clone(&dropped);
+        let owner = tokio::spawn(async move {
+            task_started.notify_one();
+            task_release.notified().await;
+            DropProbe(task_dropped)
+        });
+        started.notified().await;
+
+        let waiter = tokio::spawn(async move {
+            let _ = owner.await;
+        });
+        waiter.abort();
+        let _ = waiter.await;
+        release.notify_one();
+
+        tokio::time::timeout(Duration::from_secs(1), async {
+            while dropped.load(Ordering::SeqCst) == 0 {
+                tokio::task::yield_now().await;
+            }
+        })
+        .await
+        .expect("detached owner output should be dropped after it finishes");
+        assert_eq!(dropped.load(Ordering::SeqCst), 1);
+    }
+
+    #[tokio::test]
+    async fn planning_owner_deadline_drops_stalled_work_and_its_guards() {
+        let dropped = Arc::new(AtomicUsize::new(0));
+        let task_dropped = Arc::clone(&dropped);
+        let planning = async move {
+            let _probe = DropProbe(task_dropped);
+            std::future::pending::<()>().await;
+        };
+
+        let result =
+            super::await_owned_planning_deadline(planning, Duration::from_millis(20)).await;
+
+        assert!(
+            result.is_err(),
+            "stalled planning must hit its owner deadline"
+        );
+        assert_eq!(dropped.load(Ordering::SeqCst), 1);
+    }
+}