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.

apps/aether-gateway/src/ai_serving/mod.rs

@@ -64,7 +64,7 @@ pub(crate) use self::planner::{
     GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot, LocalExecutionAttemptSource,
     LocalExecutionCandidateKind, LocalResolvedOAuthRequestAuth, PlannerAppState,
     ResponsesWebSocketBodyNormalization, ResponsesWebSocketDecision,
-    SkippedLocalExecutionCandidate,
+    ResponsesWebSocketPinnedCandidate, SkippedLocalExecutionCandidate,
 };
 pub(crate) use self::pure::*;
 pub(crate) use self::response_history::{

apps/aether-gateway/src/ai_serving/planner/decision_input.rs

@@ -351,6 +351,7 @@ fn apply_codex_oauth_fingerprint_convergence_to_decision(
 struct GatewayAuthenticatedDecisionInputPort<'a> {
     state: PlannerAppState<'a>,
     now_unix_secs: u64,
+    auth_snapshot_override: Option<GatewayAuthApiKeySnapshot>,
     model_directive_policy: &'a crate::system_features::ModelDirectivePolicySnapshot,
     model_directive_base_model: Option<String>,
 }
@@ -367,6 +368,17 @@ impl AiAuthenticatedDecisionInputPort for GatewayAuthenticatedDecisionInputPort<
         &self,
         auth_context: &Self::AuthContext,
     ) -> Result<Option<Self::AuthSnapshot>, Self::Error> {
+        if let Some(snapshot) = self.auth_snapshot_override.as_ref() {
+            if snapshot.user_id != auth_context.user_id
+                || snapshot.api_key_id != auth_context.api_key_id
+            {
+                return Err(GatewayError::Internal(
+                    "WebSocket auth snapshot identity does not match its control decision"
+                        .to_string(),
+                ));
+            }
+            return Ok(Some(snapshot.clone()));
+        }
         self.state
             .read_auth_api_key_snapshot(
                 &auth_context.user_id,
@@ -750,6 +762,27 @@ pub(crate) async fn resolve_local_authenticated_decision_input(
     requested_model_api_format: Option<&str>,
     explicit_required_capabilities: Option<&serde_json::Value>,
     model_directive_policy: &crate::system_features::ModelDirectivePolicySnapshot,
+) -> Result<Option<ResolvedLocalDecisionAuthInput>, GatewayError> {
+    resolve_local_authenticated_decision_input_with_snapshot(
+        state,
+        auth_context,
+        None,
+        requested_model,
+        requested_model_api_format,
+        explicit_required_capabilities,
+        model_directive_policy,
+    )
+    .await
+}
+
+pub(crate) async fn resolve_local_authenticated_decision_input_with_snapshot(
+    state: &AppState,
+    auth_context: ExecutionRuntimeAuthContext,
+    auth_snapshot_override: Option<GatewayAuthApiKeySnapshot>,
+    requested_model: Option<&str>,
+    requested_model_api_format: Option<&str>,
+    explicit_required_capabilities: Option<&serde_json::Value>,
+    model_directive_policy: &crate::system_features::ModelDirectivePolicySnapshot,
 ) -> Result<Option<ResolvedLocalDecisionAuthInput>, GatewayError> {
     let model_directive_base_model = match (requested_model, requested_model_api_format) {
         (Some(model), Some(api_format)) => model_directive_policy
@@ -761,6 +794,7 @@ pub(crate) async fn resolve_local_authenticated_decision_input(
     let port = GatewayAuthenticatedDecisionInputPort {
         state: PlannerAppState::new(state),
         now_unix_secs: current_unix_secs(),
+        auth_snapshot_override,
         model_directive_policy,
         model_directive_base_model,
     };
@@ -1066,6 +1100,52 @@ mod tests {
         }
     }
 
+    #[tokio::test]
+    async fn explicit_auth_snapshot_override_does_not_fall_back_to_the_planner_cache() {
+        // AppState::new has no auth snapshot repository. Without the explicit
+        // override this resolver returns None; a WebSocket strong snapshot must
+        // therefore be the exact value used to build the planner input.
+        let state = AppState::new().expect("test state should build");
+        let mut strong_snapshot = sample_auth_snapshot();
+        strong_snapshot.api_key_allowed_models = Some(vec!["gpt-live-only".to_string()]);
+
+        let resolved = resolve_local_authenticated_decision_input_with_snapshot(
+            &state,
+            sample_auth_context(),
+            Some(strong_snapshot.clone()),
+            Some("gpt-live-only"),
+            Some("openai:responses"),
+            None,
+            &Default::default(),
+        )
+        .await
+        .expect("snapshot override should resolve")
+        .expect("the explicit snapshot should replace the missing cached value");
+
+        assert_eq!(resolved.auth_snapshot, strong_snapshot);
+    }
+
+    #[tokio::test]
+    async fn explicit_auth_snapshot_override_rejects_an_identity_mismatch() {
+        let state = AppState::new().expect("test state should build");
+        let mut wrong_snapshot = sample_auth_snapshot();
+        wrong_snapshot.api_key_id = "another-key".to_string();
+
+        let error = resolve_local_authenticated_decision_input_with_snapshot(
+            &state,
+            sample_auth_context(),
+            Some(wrong_snapshot),
+            Some("gpt-live-only"),
+            Some("openai:responses"),
+            None,
+            &Default::default(),
+        )
+        .await
+        .expect_err("a snapshot for another API key must never be injected");
+
+        assert!(matches!(error, GatewayError::Internal(message) if message.contains("identity")));
+    }
+
     #[tokio::test]
     async fn explicit_routing_attachment_authorizes_and_caches_per_principal() {
         let repository = Arc::new(InMemoryRoutingGroupRepository::default());

apps/aether-gateway/src/ai_serving/planner/mod.rs

@@ -84,6 +84,7 @@ pub(crate) use self::standard::{
     codex_model_capabilities_for_transport, maybe_build_responses_websocket_decision,
     set_local_openai_chat_execution_exhausted_diagnostic, validate_final_openai_provider_request,
     ResponsesWebSocketBodyNormalization, ResponsesWebSocketDecision,
+    ResponsesWebSocketPinnedCandidate,
 };
 pub(crate) use self::state::{
     GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth,

apps/aether-gateway/src/ai_serving/planner/standard/mod.rs

@@ -50,6 +50,7 @@ pub(crate) use self::openai::{
     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, ResponsesWebSocketBodyNormalization, ResponsesWebSocketDecision,
+    ResponsesWebSocketPinnedCandidate,
 };
 pub(crate) use crate::ai_serving::normalize_standard_request_to_openai_chat_request;
 pub(crate) use crate::ai_serving::{

apps/aether-gateway/src/ai_serving/planner/standard/openai/mod.rs

@@ -26,5 +26,5 @@ pub(crate) use responses::{
     maybe_build_responses_websocket_decision,
     maybe_build_stream_local_openai_responses_decision_payload,
     maybe_build_sync_local_openai_responses_decision_payload, ResponsesWebSocketBodyNormalization,
-    ResponsesWebSocketDecision,
+    ResponsesWebSocketDecision, ResponsesWebSocketPinnedCandidate,
 };

apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision.rs

@@ -9,7 +9,9 @@ pub(super) use self::payload::maybe_build_local_openai_responses_decision_payloa
 pub(super) use self::support::{
     build_local_openai_responses_candidate_attempt_source,
     materialize_local_openai_responses_candidate_attempts,
-    resolve_local_openai_responses_decision_input, LocalOpenAiResponsesCandidateAttempt,
-    LocalOpenAiResponsesCandidateAttemptSource, LocalOpenAiResponsesDecisionInput,
+    resolve_local_openai_responses_decision_input,
+    resolve_local_openai_responses_decision_input_with_snapshot,
+    LocalOpenAiResponsesCandidateAttempt, LocalOpenAiResponsesCandidateAttemptSource,
+    LocalOpenAiResponsesDecisionInput,
 };
 pub(super) use crate::ai_serving::LocalOpenAiResponsesSpec;

apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/support.rs

@@ -22,6 +22,7 @@ use crate::ai_serving::planner::common::extract_standard_requested_model;
 use crate::ai_serving::planner::decision_input::{
     attach_routing_policy_to_local_requested_model_input,
     build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
+    resolve_local_authenticated_decision_input_with_snapshot,
 };
 use crate::ai_serving::planner::materialization_policy::{
     build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
@@ -32,7 +33,8 @@ use crate::ai_serving::planner::CandidateFailureDiagnostic;
 use crate::ai_serving::{
     ai_local_execution_contract_for_formats, extract_pool_sticky_session_token,
     openai_responses_request_operation, resolve_local_decision_execution_runtime_auth_context,
-    ExecutionRuntimeAuthContext, GatewayControlDecision, PlannerAppState,
+    ExecutionRuntimeAuthContext, GatewayAuthApiKeySnapshot, GatewayControlDecision,
+    PlannerAppState,
 };
 use crate::client_session_affinity::client_session_affinity_from_parts;
 use crate::{AppState, GatewayError};
@@ -51,6 +53,21 @@ pub(crate) async fn resolve_local_openai_responses_decision_input(
     decision: &GatewayControlDecision,
     body_json: &serde_json::Value,
     plan_kind: &str,
+) -> Result<Option<LocalOpenAiResponsesDecisionInput>, GatewayError> {
+    resolve_local_openai_responses_decision_input_with_snapshot(
+        state, parts, trace_id, decision, body_json, plan_kind, None,
+    )
+    .await
+}
+
+pub(crate) async fn resolve_local_openai_responses_decision_input_with_snapshot(
+    state: &AppState,
+    parts: &http::request::Parts,
+    trace_id: &str,
+    decision: &GatewayControlDecision,
+    body_json: &serde_json::Value,
+    plan_kind: &str,
+    auth_snapshot_override: Option<&GatewayAuthApiKeySnapshot>,
 ) -> Result<Option<LocalOpenAiResponsesDecisionInput>, GatewayError> {
     let Some(auth_context) = resolve_local_decision_execution_runtime_auth_context(decision) else {
         warn!(
@@ -87,16 +104,28 @@ pub(crate) async fn resolve_local_openai_responses_decision_input(
         return Ok(None);
     };
 
-    let resolved_input = match resolve_local_authenticated_decision_input(
-        state,
-        auth_context.clone(),
-        Some(requested_model.as_str()),
-        decision.auth_endpoint_signature.as_deref(),
-        None,
-        &decision.model_directive_policy,
-    )
-    .await
-    {
+    let resolved_input = match if let Some(auth_snapshot) = auth_snapshot_override {
+        resolve_local_authenticated_decision_input_with_snapshot(
+            state,
+            auth_context.clone(),
+            Some(auth_snapshot.clone()),
+            Some(requested_model.as_str()),
+            decision.auth_endpoint_signature.as_deref(),
+            None,
+            &decision.model_directive_policy,
+        )
+        .await
+    } else {
+        resolve_local_authenticated_decision_input(
+            state,
+            auth_context.clone(),
+            Some(requested_model.as_str()),
+            decision.auth_endpoint_signature.as_deref(),
+            None,
+            &decision.model_directive_policy,
+        )
+        .await
+    } {
         Ok(Some(resolved_input)) => resolved_input,
         Ok(None) => {
             warn!(

apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/mod.rs

@@ -12,6 +12,49 @@ use crate::{AiExecutionDecision, AppState, GatewayError};
 use aether_runtime_state::RuntimeLockLease;
 use std::collections::BTreeSet;
 
+/// Releases a scheduler pool-key lease if WebSocket planning is cancelled
+/// after candidate selection but before ownership reaches the turn lifecycle.
+struct ResponsesWebSocketPlanningLeaseGuard {
+    state: AppState,
+    lease: Option<RuntimeLockLease>,
+}
+
+impl ResponsesWebSocketPlanningLeaseGuard {
+    fn new(state: &AppState, lease: Option<&RuntimeLockLease>) -> Self {
+        Self {
+            state: state.clone(),
+            lease: lease.cloned(),
+        }
+    }
+
+    async fn release(mut self) {
+        // Keep the lease armed across the await. If the owner task is aborted
+        // or reaches its hard deadline while the runtime backend is stalled,
+        // Drop can still hand cleanup to a detached owner.
+        if release_responses_websocket_planning_lease(&self.state, self.lease.as_ref()).await {
+            self.lease = None;
+        }
+    }
+
+    fn disarm(&mut self) {
+        self.lease = None;
+    }
+}
+
+impl Drop for ResponsesWebSocketPlanningLeaseGuard {
+    fn drop(&mut self) {
+        let Some(lease) = self.lease.take() else {
+            return;
+        };
+        let state = self.state.clone();
+        if let Ok(handle) = tokio::runtime::Handle::try_current() {
+            handle.spawn(async move {
+                let _ = release_responses_websocket_planning_lease(&state, Some(&lease)).await;
+            });
+        }
+    }
+}
+
 mod decision;
 mod plans;
 
@@ -19,6 +62,7 @@ use self::decision::{
     build_local_openai_responses_candidate_attempt_source,
     maybe_build_local_openai_responses_decision_payload_for_candidate,
     resolve_local_openai_responses_decision_input,
+    resolve_local_openai_responses_decision_input_with_snapshot,
 };
 use self::plans::{
     build_local_stream_attempt_source, build_local_stream_plan_and_reports,
@@ -187,6 +231,45 @@ pub(crate) struct ResponsesWebSocketDecision {
     pub(crate) normalization: ResponsesWebSocketBodyNormalization,
 }
 
+/// The scheduler identity a continuation is allowed to reuse.
+///
+/// A `previous_response_id` chain cannot move to another provider connection,
+/// but it still has to pass the current scheduler runtime checks on every
+/// turn. The planner uses this identity as a filter rather than selecting an
+/// arbitrary eligible replacement.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub(crate) struct ResponsesWebSocketPinnedCandidate {
+    provider_id: String,
+    endpoint_id: String,
+    key_id: String,
+}
+
+impl ResponsesWebSocketPinnedCandidate {
+    pub(crate) fn from_decision(decision: &AiExecutionDecision) -> Option<Self> {
+        Some(Self {
+            provider_id: non_empty_decision_identity(decision.provider_id.as_deref())?,
+            endpoint_id: non_empty_decision_identity(decision.endpoint_id.as_deref())?,
+            key_id: non_empty_decision_identity(decision.key_id.as_deref())?,
+        })
+    }
+
+    fn matches(
+        &self,
+        candidate: &aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate,
+    ) -> bool {
+        candidate.provider_id == self.provider_id
+            && candidate.endpoint_id == self.endpoint_id
+            && candidate.key_id == self.key_id
+    }
+}
+
+fn non_empty_decision_identity(value: Option<&str>) -> Option<String> {
+    value
+        .map(str::trim)
+        .filter(|value| !value.is_empty())
+        .map(str::to_string)
+}
+
 /// Everything needed to re-run provider-body normalization for the candidate a
 /// socket is already bound to.
 ///
@@ -321,21 +404,24 @@ pub(crate) async fn maybe_build_responses_websocket_decision(
     parts: &http::request::Parts,
     trace_id: &str,
     decision: &GatewayControlDecision,
+    auth_snapshot: Option<&crate::ai_serving::GatewayAuthApiKeySnapshot>,
     body_json: &serde_json::Value,
     excluded_key_ids: Option<&BTreeSet<String>>,
     excluded_codex_account_ids: Option<&BTreeSet<String>>,
+    pinned_candidate: Option<&ResponsesWebSocketPinnedCandidate>,
 ) -> 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(
+    let Some(input) = resolve_local_openai_responses_decision_input_with_snapshot(
         state,
         parts,
         trace_id,
         decision,
         body_json,
         spec.decision_kind,
+        auth_snapshot,
     )
     .await?
     else {
@@ -348,18 +434,28 @@ pub(crate) async fn maybe_build_responses_websocket_decision(
     .await?;
 
     while let Some(attempt) = source.next_attempt().await? {
-        let pool_key_lease = attempt.eligible.orchestration.pool_key_lease.clone();
+        // `next_attempt` may return with a distributed pool-key lease. Arm a
+        // guard before the first await so owner-task timeout/cancellation
+        // cannot strand that lease until its server-side TTL expires.
+        let mut planning_lease = ResponsesWebSocketPlanningLeaseGuard::new(
+            state,
+            attempt.eligible.orchestration.pool_key_lease.as_ref(),
+        );
+        if pinned_candidate.is_some_and(|pinned| !pinned.matches(&attempt.eligible.candidate)) {
+            planning_lease.release().await;
+            continue;
+        }
         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;
+            planning_lease.release().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;
+            planning_lease.release().await;
             continue;
         };
         // Captured before `attempt` is consumed so a later continuation turn can
@@ -373,11 +469,11 @@ pub(crate) async fn maybe_build_responses_websocket_decision(
         {
             Ok(Some(payload)) => payload,
             Ok(None) => {
-                release_responses_websocket_planning_lease(state, pool_key_lease.as_ref()).await;
+                planning_lease.release().await;
                 continue;
             }
             Err(error) => {
-                release_responses_websocket_planning_lease(state, pool_key_lease.as_ref()).await;
+                planning_lease.release().await;
                 return Err(error);
             }
         };
@@ -393,7 +489,7 @@ pub(crate) async fn maybe_build_responses_websocket_decision(
                     .is_some_and(|account_ids| account_ids.contains(account_id))
             })
         {
-            release_responses_websocket_planning_lease(state, pool_key_lease.as_ref()).await;
+            planning_lease.release().await;
             continue;
         }
         match codex_quota_breaker_blocks_candidate(
@@ -405,7 +501,7 @@ pub(crate) async fn maybe_build_responses_websocket_decision(
         .await
         {
             Ok(true) => {
-                release_responses_websocket_planning_lease(state, pool_key_lease.as_ref()).await;
+                planning_lease.release().await;
                 continue;
             }
             Ok(false) => {}
@@ -454,13 +550,17 @@ pub(crate) async fn maybe_build_responses_websocket_decision(
                     .flatten(),
                 mapped_model,
             };
-            return Ok(Some(ResponsesWebSocketDecision {
+            let decision = ResponsesWebSocketDecision {
                 execution: payload,
                 adapter,
                 normalization,
-            }));
+            };
+            // The decision report context now carries the lease identity. The
+            // WebSocket ownership layer takes over before any further await.
+            planning_lease.disarm();
+            return Ok(Some(decision));
         }
-        release_responses_websocket_planning_lease(state, pool_key_lease.as_ref()).await;
+        planning_lease.release().await;
     }
 
     Ok(None)
@@ -469,20 +569,23 @@ pub(crate) async fn maybe_build_responses_websocket_decision(
 async fn release_responses_websocket_planning_lease(
     state: &AppState,
     lease: Option<&RuntimeLockLease>,
-) {
+) -> bool {
     let Some(lease) = lease else {
-        return;
+        return true;
     };
-    if let Err(error) =
-        crate::handlers::shared::provider_pool::release_admin_provider_pool_key_lease(
-            state.runtime_state.as_ref(),
-            lease,
-        )
-        .await
+    match 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"
-        );
+        Ok(_) => true,
+        Err(error) => {
+            tracing::warn!(
+                error = ?error,
+                "gateway Responses WebSocket planner failed to release an unused pool key lease"
+            );
+            false
+        }
     }
 }