Aether/apps/aether-gateway/src/handlers/proxy/websocket/responses/connection.rs

//! Connection-level Responses WebSocket FSM.

use std::time::Duration;

use axum::extract::ws::{Message as AxumWsMessage, WebSocket};
use futures_util::{SinkExt, StreamExt};
use serde_json::Value;
use wreq::ws::message::Message as WreqWsMessage;

use super::adapter::ResponsesWebSocketRelayDirective;
use super::client::{adapter_drain_ready, forward_client_message, RelayDisposition};
use super::continuation::{
    ResponsesWebSocketContinuationRecord, ResponsesWebSocketContinuationRegistry,
};
use super::frame::{encode_opaque_websocket_event, ParsedResponsesWebSocketFrame};
use super::lifecycle::{
    await_pending_adapter_observation, finalize_active_turn, queue_turn_finalization,
    settle_turn_finalization, spawn_bounded_adapter_observation, PreviousAttemptSettled,
};
use super::plan_admission::terminate_responses_websocket_for_plan_permit_loss;
use super::quota::{
    detach_exhausted_upstream, is_usage_limit_error_event, mark_active_response_retry_unsafe,
    observe_active_response_rebind_safety, retry_active_turn_after_quota_exhaustion,
};
use super::relay_policy::{
    classify_quota_relay, fatal_relay_policy, FatalRelaySignal, QuotaRelayAction, QuotaRelayFacts,
};
use super::settlement::settle_signal_for_client_delivery_failure;
use super::state::BoundResponsesConnection;
use super::turn::{
    ResponsesProviderAttempt, ResponsesWebSocketTurnObservation, ResponsesWebSocketTurnOutcome,
};
use super::turn_state::LogicalTurn;
use super::upstream::{close_bound_upstream, receive_optional_upstream};
use crate::handlers::proxy::websocket::ingress::WebSocketRequestContext;
use crate::handlers::proxy::websocket::session::{
    wait_for_optional_deadline, CLOSE_INTERNAL_ERROR, CLOSE_TRY_AGAIN, WEBSOCKET_LOG_TRANSPORT,
};
use crate::handlers::proxy::websocket::transport::{
    close_client_socket, send_client_message, send_gateway_error_with_status,
    send_responses_websocket_error, upstream_message_to_client,
};
use crate::AppState;

const LOG_TARGET: &str = "aether_gateway::handlers::proxy::responses_ws";
const CONTINUATION_REGISTRATION_TIMEOUT: Duration = Duration::from_millis(500);

/// 写客户端 socket 失败时记录的投递失败原因。刻意不说「客户端在终态前断开」：
/// 供应商的终态可能已经到达，只是最后一跳没送出去。
const CLIENT_DELIVERY_FAILED_REASON: &str =
    "gateway could not relay the provider event to the client";

macro_rules! debug {
    ($($arg:tt)*) => {
        tracing::debug!(target: LOG_TARGET, $($arg)*)
    };
}

macro_rules! warn {
    ($($arg:tt)*) => {
        tracing::warn!(target: LOG_TARGET, $($arg)*)
    };
}

pub(super) async fn relay_bound_connection(
    client_socket: &mut WebSocket,
    bound: &mut BoundResponsesConnection,
    state: &AppState,
    context: &WebSocketRequestContext,
) {
    let mut client_connected = true;
    loop {
        if !client_connected && !bound.turn_state.response_in_flight() {
            close_bound_upstream(bound).await;
            break;
        }
        let active_turn_deadline = bound.turn_state.attempt().map(|turn| turn.deadline());
        tokio::select! {
            _ = wait_for_optional_deadline(active_turn_deadline.map(|deadline| deadline.deadline)) => {
                let Some(turn_deadline) = active_turn_deadline else {
                    continue;
                };
                warn!(
                    event_name = "responses_websocket_turn_timeout",
                    log_type = "ops",
                    transport = WEBSOCKET_LOG_TRANSPORT,
                    websocket = true,
                    trace_id = %context.trace_id,
                    timeout_phase = ?turn_deadline.phase,
                    timeout_ms = turn_deadline.timeout.as_millis() as u64,
                    "Responses WebSocket response did not reach its configured deadline"
                );
                finalize_active_turn(bound, state, turn_deadline.phase.outcome()).await;
                send_gateway_error_with_status(
                    client_socket,
                    504,
                    turn_deadline.phase.error_code(),
                    turn_deadline.phase.client_message(),
                ).await;
                close_bound_upstream(bound).await;
                close_client_socket(
                    client_socket,
                    CLOSE_TRY_AGAIN,
                    turn_deadline.phase.error_code(),
                ).await;
                break;
            }
            client_message = client_socket.next(), if client_connected => {
                let Some(client_message) = client_message else {
                    if retain_disconnected_turn(bound) {
                        client_connected = false;
                        continue;
                    }
                    finalize_active_turn(
                        bound,
                        state,
                        ResponsesWebSocketTurnOutcome::client_disconnected(),
                    ).await;
                    close_bound_upstream(bound).await;
                    break;
                };
                let Ok(client_message) = client_message else {
                    if retain_disconnected_turn(bound) {
                        client_connected = false;
                        continue;
                    }
                    warn!(
                        event_name = "responses_websocket_client_receive_failed",
                        log_type = "ops",
                        transport = WEBSOCKET_LOG_TRANSPORT,
                        websocket = true,
                        trace_id = %context.trace_id,
                        "client WebSocket receive failed"
                    );
                    finalize_active_turn(
                        bound,
                        state,
                        ResponsesWebSocketTurnOutcome::client_disconnected(),
                    ).await;
                    close_bound_upstream(bound).await;
                    break;
                };
                if matches!(client_message, AxumWsMessage::Close(_))
                    && retain_disconnected_turn(bound)
                {
                    client_connected = false;
                    continue;
                }
                match Box::pin(forward_client_message(
                    client_message,
                    bound,
                    client_socket,
                    state,
                    context,
                ))
                .await
                {
                    RelayDisposition::Continue => {}
                    RelayDisposition::Close => {
                        finalize_active_turn(
                            bound,
                            state,
                            ResponsesWebSocketTurnOutcome::client_disconnected(),
                        ).await;
                        break;
                    }
                    RelayDisposition::PlanUsagePermitLost => {
                        warn!(
                            event_name = "responses_websocket_plan_usage_concurrency_lost_before_send",
                            log_type = "ops",
                            transport = WEBSOCKET_LOG_TRANSPORT,
                            websocket = true,
                            trace_id = %context.trace_id,
                            "gateway stopped a Responses WebSocket turn before its upstream send after the subscription plan concurrency lease became unhealthy"
                        );
                        close_bound_upstream(bound).await;
                        terminate_responses_websocket_for_plan_permit_loss(client_socket).await;
                        break;
                    }
                    RelayDisposition::UpstreamError(code) => {
                        warn!(
                            event_name = "responses_websocket_upstream_send_failed",
                            log_type = "ops",
                            transport = WEBSOCKET_LOG_TRANSPORT,
                            websocket = true,
                            trace_id = %context.trace_id,
                            error_code = code,
                            "Upstream WebSocket send failed"
                        );
                        finalize_active_turn(
                            bound,
                            state,
                            ResponsesWebSocketTurnOutcome::upstream_send_failed(),
                        ).await;
                        send_gateway_error_with_status(
                            client_socket,
                            502,
                            code,
                            "Gateway could not forward the WebSocket event upstream",
                        ).await;
                        close_bound_upstream(bound).await;
                        close_client_socket(client_socket, CLOSE_INTERNAL_ERROR, code).await;
                        break;
                    }
                }
            }
            upstream_message = receive_optional_upstream(&mut bound.upstream) => {
                let Some(upstream_message) = upstream_message else {
                    finalize_active_turn(
                        bound,
                        state,
                        ResponsesWebSocketTurnOutcome::upstream_closed(),
                    ).await;
                    bound.upstream = None;
                    close_client_socket(client_socket, 1000, "upstream_closed").await;
                    break;
                };
                let Ok(upstream_message) = upstream_message else {
                    warn!(
                        event_name = "responses_websocket_upstream_receive_failed",
                        log_type = "ops",
                        transport = WEBSOCKET_LOG_TRANSPORT,
                        websocket = true,
                        trace_id = %context.trace_id,
                        "Upstream WebSocket receive failed"
                    );
                    finalize_active_turn(
                        bound,
                        state,
                        ResponsesWebSocketTurnOutcome::upstream_receive_failed(),
                    ).await;
                    send_gateway_error_with_status(
                        client_socket,
                        502,
                        "responses_websocket_receive_failed",
                        "Provider connection closed unexpectedly",
                    ).await;
                    bound.upstream = None;
                    close_client_socket(client_socket, CLOSE_INTERNAL_ERROR, "upstream_receive_failed").await;
                    break;
                };
                let parsed_upstream_frame = match &upstream_message {
                    WreqWsMessage::Text(text) => {
                        ParsedResponsesWebSocketFrame::parse(text.as_str()).ok()
                    }
                    _ => None,
                };
                if parsed_upstream_frame
                    .as_ref()
                    .is_some_and(ParsedResponsesWebSocketFrame::carries_stream_id)
                {
                    // The session currently owns only the implicit default
                    // lane. Reject a provider-side named identity before the
                    // adapter, usage observer, continuation cache, PII
                    // restorer, or logical-turn state can attribute an
                    // interleaved event to the sole default-lane attempt.
                    let policy = fatal_relay_policy(
                        FatalRelaySignal::UnexpectedUpstreamStreamId,
                    );
                    let frame = parsed_upstream_frame
                        .as_ref()
                        .expect("the stream-id guard requires a parsed frame");
                    warn!(
                        event_name = "responses_websocket_unexpected_upstream_stream_id",
                        log_type = "ops",
                        transport = WEBSOCKET_LOG_TRANSPORT,
                        websocket = true,
                        trace_id = %context.trace_id,
                        event_type = %frame.event_type_for_log(),
                        frame_bytes = frame.raw_text().len(),
                        chunked = frame.is_chunked(),
                        "gateway rejected a named-lane provider event on a default-lane Responses WebSocket"
                    );
                    finalize_active_turn(
                        bound,
                        state,
                        ResponsesWebSocketTurnOutcome::upstream_receive_failed(),
                    )
                    .await;
                    send_responses_websocket_error(
                        client_socket,
                        policy.status_code,
                        "server_error",
                        policy.error_code,
                        policy.client_message,
                    )
                    .await;
                    close_bound_upstream(bound).await;
                    close_client_socket(
                        client_socket,
                        policy.close_code,
                        policy.close_reason,
                    )
                    .await;
                    break;
                }
                let parsed_upstream_event = parsed_upstream_frame
                    .as_ref()
                    .map(ParsedResponsesWebSocketFrame::event);
                if let WreqWsMessage::Text(text) = &upstream_message {
                    debug!(
                        event_name = "responses_websocket_upstream_event",
                        log_type = "event",
                        transport = WEBSOCKET_LOG_TRANSPORT,
                        websocket = true,
                        trace_id = %context.trace_id,
                        event_type = %parsed_upstream_frame
                            .as_ref()
                            .map(ParsedResponsesWebSocketFrame::event_type_for_log)
                            .unwrap_or_else(|| "invalid_json".to_string()),
                        frame_bytes = text.len(),
                        chunked = parsed_upstream_frame
                            .as_ref()
                            .is_some_and(ParsedResponsesWebSocketFrame::is_chunked),
                        active_turn = bound.turn_state.response_in_flight(),
                        "gateway received Responses WebSocket event"
                    );
                }
                if matches!(&upstream_message, WreqWsMessage::Binary(_)) {
                    mark_active_response_retry_unsafe(bound, "upstream_binary_frame");
                } else if matches!(&upstream_message, WreqWsMessage::Text(_))
                    && parsed_upstream_event.is_none()
                {
                    mark_active_response_retry_unsafe(bound, "invalid_upstream_event");
                }
                if let Some(event) = parsed_upstream_event {
                    observe_active_response_rebind_safety(bound, event);
                    if bound.pending_adapter_drain.is_none()
                        && bound.adapter.observes_upstream_events()
                    {
                        let adapter = bound.adapter;
                        if let Some(observation) = adapter.observe_upstream_event(event) {
                            let directive = observation.drain;
                            await_pending_adapter_observation(bound).await;
                            let state_for_observation = state.clone();
                            let trace_id = context.trace_id.clone();
                            let report_context = bound.decision_template.report_context.clone();
                            bound.pending_adapter_observation = Some(spawn_bounded_adapter_observation(async move {
                                adapter
                                    .persist_upstream_observation(
                                        &state_for_observation,
                                        &trace_id,
                                        report_context.as_ref(),
                                        observation,
                                    )
                                    .await;
                            }));
                            if let Some(directive) = directive {
                                bound.pending_adapter_drain = Some(directive);
                                // A definitive quota signal must be visible to
                                // the next planner before a transparent retry.
                                await_pending_adapter_observation(bound).await;
                            }
                        }
                    }
                }
                let observation = match &upstream_message {
                    WreqWsMessage::Text(text) => {
                        let adapter = bound.adapter;
                        match parsed_upstream_frame.as_ref() {
                            Some(frame) => bound
                                .turn_state
                                .attempt_mut()
                                .and_then(|turn| turn.observe_upstream_frame(frame, adapter)),
                            None => {
                                if let Some(turn) = bound.turn_state.attempt_mut() {
                                    turn.observe_invalid_upstream_text(text.as_str())
                                }
                                else {
                                    None
                                }
                            }
                        }
                    }
                    _ => None,
                };
                if matches!(
                    observation,
                    Some(ResponsesWebSocketTurnObservation::Started)
                        | Some(ResponsesWebSocketTurnObservation::Terminal(_))
                ) {
                    if let Some(turn) = bound.turn_state.attempt_mut() {
                        turn.mark_stream_started(state).await;
                    }
                }
                let terminal_outcome = match observation {
                    Some(ResponsesWebSocketTurnObservation::Terminal(outcome)) => Some(outcome),
                    _ => None,
                };
                if let Some(frame) = parsed_upstream_frame.as_ref() {
                    if let Some(response_id) = evicted_default_lane_continuation_response_id(
                        bound.turn_state.logical(),
                        frame,
                    ) {
                        // A 4xx/5xx continuation terminal evicts the referenced
                        // ID from the provider's implicit default-lane cache.
                        // Do not keep claiming local ownership and replay it.
                        bound
                            .continuation_response_ids
                            .forget_connection_local(response_id);
                    }
                    if let Some(response_id) = connection_local_terminal_response_id(
                        bound.turn_state.logical(),
                        frame,
                    ) {
                        // Remember every successful response on this physical
                        // socket, including store=false responses that exist
                        // only in the provider's connection-local cache.
                        bound
                            .continuation_response_ids
                            .remember_connection_local(response_id);
                    }
                    if let Some(registration) =
                        prepare_persisted_continuation_registration(bound, context, frame)
                    {
                        if let Some(response_id) =
                            register_persisted_continuation_before_terminal_delivery(
                            state,
                            context,
                            registration,
                        )
                            .await
                        {
                            bound
                                .continuation_response_ids
                                .remember_persisted(response_id.as_str());
                        }
                    }
                }
                if matches!(&upstream_message, WreqWsMessage::Text(_))
                    && parsed_upstream_frame.is_none()
                {
                    let policy = fatal_relay_policy(FatalRelaySignal::InvalidUpstreamText);
                    finalize_active_turn(
                        bound,
                        state,
                        terminal_outcome.unwrap_or_else(
                            ResponsesWebSocketTurnOutcome::upstream_receive_failed,
                        ),
                    )
                    .await;
                    send_responses_websocket_error(
                        client_socket,
                        policy.status_code,
                        "server_error",
                        policy.error_code,
                        policy.client_message,
                    )
                    .await;
                    close_bound_upstream(bound).await;
                    close_client_socket(
                        client_socket,
                        policy.close_code,
                        policy.close_reason,
                    )
                    .await;
                    break;
                }
                let is_close = matches!(upstream_message, WreqWsMessage::Close(_));
                let drain_for_adapter = adapter_drain_ready(
                    bound.pending_adapter_drain,
                    bound.turn_state.response_in_flight(),
                    observation,
                    is_close,
                );
                let quota_facts = QuotaRelayFacts {
                    drain_ready: drain_for_adapter,
                    retry_current_turn: bound
                        .pending_adapter_drain
                        .is_some_and(|directive| directive.retry_current_turn)
                        && bound
                            .turn_state
                            .logical()
                            .is_some_and(|turn| turn.quota_retry_block_reason().is_none()),
                    transparent_retry_failed: false,
                    usage_limit_error: parsed_upstream_event.is_some_and(is_usage_limit_error_event),
                    upstream_closed: is_close,
                };
                let mut quota_relay_action = classify_quota_relay(quota_facts);
                if matches!(quota_relay_action, QuotaRelayAction::AttemptTransparentRetry) {
                    // detach_attempt 保留 logical turn：重试是同一轮请求的下一个 attempt。
                    let retry_turn = bound.turn_state.detach_attempt();
                    // 先结算旧 attempt 并等它落地，再规划下一个 attempt。两个理由：
                    //
                    // 1. 规划要读 health / adaptive / pool 状态，而这些正是旧
                    //    attempt 结算时才投射的。普通的新 turn 早就在 client.rs 里
                    //    用 await_pending_turn_finalization 挡住了「基于陈旧状态
                    //    规划」，透明重试这条路径原先漏了这一步。
                    // 2. 旧 attempt 还占着自己的 pool key lease。不先释放，重试就
                    //    可能因为「这把 key 仍被占用」而挑不到本该可用的替代 key，
                    //    或者干脆判成无可用供应商。
                    let settled = match retry_turn {
                        Some(mut turn) => {
                            turn.release_admission().await;
                            settle_turn_finalization(
                                bound,
                                state,
                                turn,
                                terminal_outcome.unwrap_or_else(
                                    ResponsesWebSocketTurnOutcome::upstream_closed,
                                ),
                            )
                            .await
                        }
                        None => Some(PreviousAttemptSettled::nothing_to_settle()),
                    };
                    // Planning and binding a replacement carries the complete
                    // scheduler/provider state machine. Keep that large future
                    // off the relay task's stack; the default Tokio/test worker
                    // stack is otherwise easy to exhaust on this rare branch.
                    if let Some(settled) = settled {
                        match Box::pin(retry_active_turn_after_quota_exhaustion(
                            bound, state, context, settled,
                        ))
                        .await
                        {
                            Ok(true) => continue,
                            Ok(false) => {}
                            Err(()) => {
                                finalize_active_turn(
                                    bound,
                                    state,
                                    ResponsesWebSocketTurnOutcome::connection_admission_lost(),
                                )
                                .await;
                                close_bound_upstream(bound).await;
                                terminate_responses_websocket_for_plan_permit_loss(client_socket)
                                    .await;
                                break;
                            }
                        }
                    }
                    // 重试失败。旧 attempt 已经结算，logical turn 仍停在
                    // Replanning，所以后面分支里的 end() / finalize_active_turn
                    // 只会清掉 logical turn 而不会交出 attempt——不存在重复结算。
                    quota_relay_action = classify_quota_relay(QuotaRelayFacts {
                        retry_current_turn: false,
                        transparent_retry_failed: true,
                        ..quota_facts
                    });
                }
                let detach_after_forward =
                    matches!(quota_relay_action, QuotaRelayAction::ForwardQuotaAndDetach);
                if detach_after_forward && is_close {
                    let directive = bound
                        .pending_adapter_drain
                        .expect("adapter drain state should be present");
                    finalize_active_turn(
                        bound,
                        state,
                        terminal_outcome
                            .unwrap_or_else(ResponsesWebSocketTurnOutcome::provider_quota_exhausted),
                    )
                    .await;
                    send_gateway_error_with_status(
                        client_socket,
                        429,
                        directive.error_code,
                        "Provider connection closed after reporting exhausted quota; send a new response.create to select another Provider connection",
                    )
                    .await;
                    detach_exhausted_upstream(bound, directive, &context.trace_id).await;
                    continue;
                }
                // Standard Responses frames cross the gateway byte-for-byte unless PII
                // restoration has something to replace. Codex may wrap public events with
                // provider-private side-channel chunks; only that explicit envelope is
                // peeled, and each retained event is serialized as a complete opaque Value.
                // Observation and capture continue to consume the redacted event, while the
                // final client hop receives restored text.
                let relay_directive = parsed_upstream_frame
                    .as_ref()
                    .map(|frame| {
                        bound
                            .adapter
                            .relay_directive_for_upstream_event(frame.event())
                    });
                let mut relay_send_error = None;
                let mut relay_serialization_failed = false;
                match relay_directive {
                    _ if !client_connected => {}
                    Some(ResponsesWebSocketRelayDirective::ForwardOriginal) => {
                        let client_frame = match parsed_upstream_frame.as_ref().map(|frame| {
                            bound
                                .redaction_restorer
                                .restore_provider_frame_text(frame.event())
                        }) {
                            Some(Ok(Some(text))) => Some(AxumWsMessage::Text(text.into())),
                            Some(Ok(None)) | None => {
                                Some(upstream_message_to_client(upstream_message.clone()))
                            }
                            Some(Err(_)) => {
                                relay_serialization_failed = true;
                                None
                            }
                        };
                        if let Some(client_frame) = client_frame {
                            match send_client_message(client_socket, client_frame).await {
                                Ok(()) => {
                                    if let (Some(turn), Some(frame)) = (
                                        bound.turn_state.attempt_mut(),
                                        parsed_upstream_frame.as_ref(),
                                    ) {
                                        turn.capture_client_frame(frame.event());
                                    }
                                }
                                Err(error) => relay_send_error = Some(error),
                            }
                        }
                    }
                    Some(ResponsesWebSocketRelayDirective::ForwardEvents(events)) => {
                        for event in events {
                            let event = event.as_ref();
                            let text = match bound
                                .redaction_restorer
                                .restore_provider_frame_text(event)
                            {
                                Ok(Some(restored)) => restored,
                                Ok(None) => match encode_opaque_websocket_event(event) {
                                    Ok(encoded) => encoded,
                                    Err(_) => {
                                        relay_serialization_failed = true;
                                        break;
                                    }
                                },
                                Err(_) => {
                                    relay_serialization_failed = true;
                                    break;
                                }
                            };
                            match send_client_message(
                                client_socket,
                                AxumWsMessage::Text(text.into()),
                            )
                            .await
                            {
                                Ok(()) => {
                                    if let Some(turn) = bound.turn_state.attempt_mut() {
                                        turn.capture_client_frame(event);
                                    }
                                }
                                Err(error) => {
                                    relay_send_error = Some(error);
                                    break;
                                }
                            }
                        }
                    }
                    Some(ResponsesWebSocketRelayDirective::SuppressProviderPrivate) => {}
                    None => {
                        if let Err(error) = send_client_message(
                            client_socket,
                            upstream_message_to_client(upstream_message.clone()),
                        )
                        .await
                        {
                            relay_send_error = Some(error);
                        }
                    }
                }
                if relay_serialization_failed {
                    warn!(
                        event_name = "responses_websocket_provider_event_serialization_failed",
                        log_type = "ops",
                        transport = WEBSOCKET_LOG_TRANSPORT,
                        websocket = true,
                        trace_id = %context.trace_id,
                        provider_terminal_reached = terminal_outcome.is_some(),
                        "gateway could not serialize an opaque provider event"
                    );
                    bound
                        .turn_state
                        .record_client_delivery_aborted(CLIENT_DELIVERY_FAILED_REASON);
                    finalize_active_turn(
                        bound,
                        state,
                        settle_signal_for_client_delivery_failure(terminal_outcome),
                    )
                    .await;
                    send_gateway_error_with_status(
                        client_socket,
                        502,
                        "responses_websocket_event_serialization_failed",
                        "Gateway could not relay the provider event",
                    )
                    .await;
                    close_bound_upstream(bound).await;
                    close_client_socket(
                        client_socket,
                        CLOSE_INTERNAL_ERROR,
                        "provider_event_serialization_failed",
                    )
                    .await;
                    break;
                }
                if let Some(error) = relay_send_error {
                    if terminal_outcome.is_none() && retain_disconnected_turn(bound) {
                        client_connected = false;
                        continue;
                    }
                    warn!(
                        event_name = "responses_websocket_client_send_failed",
                        log_type = "ops",
                        transport = WEBSOCKET_LOG_TRANSPORT,
                        websocket = true,
                        trace_id = %context.trace_id,
                        error_code = error.as_str(),
                        provider_terminal_reached = terminal_outcome.is_some(),
                        "gateway could not relay a provider event to the client"
                    );
                    // 投递失败是独立事实，不能覆盖已经到达的 provider 终态：
                    // 供应商已经完成推理并消耗 token，账单按它的终态计。
                    bound
                        .turn_state
                        .record_client_delivery_aborted(CLIENT_DELIVERY_FAILED_REASON);
                    finalize_active_turn(
                        bound,
                        state,
                        settle_signal_for_client_delivery_failure(terminal_outcome),
                    ).await;
                    close_bound_upstream(bound).await;
                    break;
                }
                if let Some(outcome) = terminal_outcome {
                    finalize_active_turn(bound, state, outcome).await;
                } else if is_close {
                    finalize_active_turn(
                        bound,
                        state,
                        ResponsesWebSocketTurnOutcome::upstream_closed(),
                    )
                    .await;
                }
                if detach_after_forward {
                    let directive = bound
                        .pending_adapter_drain
                        .expect("adapter drain state should be present");
                    if bound.turn_state.response_in_flight() {
                        finalize_active_turn(
                            bound,
                            state,
                            ResponsesWebSocketTurnOutcome::provider_quota_exhausted(),
                        )
                        .await;
                    }
                    detach_exhausted_upstream(bound, directive, &context.trace_id).await;
                    continue;
                }
                if drain_for_adapter {
                    let directive = bound
                        .pending_adapter_drain
                        .expect("adapter drain state should be present");
                    detach_exhausted_upstream(bound, directive, &context.trace_id).await;
                    continue;
                }
                if is_close {
                    bound.upstream = None;
                    break;
                }
            }
        }
    }
}

fn retain_disconnected_turn(bound: &mut BoundResponsesConnection) -> bool {
    if bound
        .turn_state
        .attempt()
        .is_none_or(|attempt| attempt.cancel_on_client_disconnect())
    {
        return false;
    }
    bound
        .turn_state
        .record_client_delivery_aborted(CLIENT_DELIVERY_FAILED_REASON);
    true
}

struct PendingContinuationRegistration {
    user_id: String,
    api_key_id: String,
    response_id: String,
    record: ResponsesWebSocketContinuationRecord,
}

fn prepare_persisted_continuation_registration(
    bound: &BoundResponsesConnection,
    context: &WebSocketRequestContext,
    frame: &ParsedResponsesWebSocketFrame<'_>,
) -> Option<PendingContinuationRegistration> {
    let logical = bound.turn_state.logical();
    let Some(response_id) = persistable_terminal_response_id(logical, frame).map(str::to_string)
    else {
        return None;
    };
    let logical = logical.expect("a persistable terminal requires an active logical turn");
    let Some(auth_context) = logical
        .turn_control
        .as_ref()
        .and_then(|control| control.decision.auth_context.as_ref())
    else {
        warn!(
            event_name = "responses_websocket_continuation_registration_skipped",
            log_type = "ops",
            transport = WEBSOCKET_LOG_TRANSPORT,
            websocket = true,
            trace_id = %context.trace_id,
            reason = "missing_live_auth_context",
            "gateway did not register a persisted Responses continuation"
        );
        return None;
    };
    let Some(pinned_candidate) =
        crate::ai_serving::ResponsesWebSocketPinnedCandidate::from_decision(
            &bound.decision_template,
        )
    else {
        warn!(
            event_name = "responses_websocket_continuation_registration_skipped",
            log_type = "ops",
            transport = WEBSOCKET_LOG_TRANSPORT,
            websocket = true,
            trace_id = %context.trace_id,
            reason = "missing_binding_identity",
            "gateway did not register a persisted Responses continuation"
        );
        return None;
    };
    let record = match ResponsesWebSocketContinuationRecord::from_binding(
        pinned_candidate,
        bound.client_model.as_str(),
        bound.provider_model.as_str(),
        &bound.binding_identity,
        &bound.body_normalization,
        bound.redaction_restorer.has_sessions(),
        bound.responses_lite_static_config.clone(),
    ) {
        Ok(record) => record,
        Err(error) => {
            warn!(
                event_name = "responses_websocket_continuation_registration_skipped",
                log_type = "ops",
                transport = WEBSOCKET_LOG_TRANSPORT,
                websocket = true,
                trace_id = %context.trace_id,
                provider_id = %bound.decision_template.provider_id.as_deref().unwrap_or("-"),
                endpoint_id = %bound.decision_template.endpoint_id.as_deref().unwrap_or("-"),
                key_id = %bound.decision_template.key_id.as_deref().unwrap_or("-"),
                reason = error.kind(),
                "gateway could not build a persisted Responses continuation record"
            );
            return None;
        }
    };
    Some(PendingContinuationRegistration {
        user_id: auth_context.user_id.clone(),
        api_key_id: auth_context.api_key_id.clone(),
        response_id,
        record,
    })
}

fn persistable_terminal_response_id<'a>(
    logical: Option<&LogicalTurn>,
    frame: &'a ParsedResponsesWebSocketFrame<'_>,
) -> Option<&'a str> {
    let logical = logical?;
    // `provider_store` is derived from the final framed provider event. False
    // or absent is ZDR/connection-local and must never create a 24-hour KV
    // record, even if the provider happens to return an ID.
    if !logical.provider_store {
        return None;
    }
    connection_local_terminal_response_id(Some(logical), frame)
}

fn connection_local_terminal_response_id<'a>(
    logical: Option<&LogicalTurn>,
    frame: &'a ParsedResponsesWebSocketFrame<'_>,
) -> Option<&'a str> {
    logical?;
    frame.continuation_response_id()
}

fn evicted_default_lane_continuation_response_id<'a>(
    logical: Option<&'a LogicalTurn>,
    frame: &ParsedResponsesWebSocketFrame<'_>,
) -> Option<&'a str> {
    let logical = logical?;
    let terminal = frame.terminal()?;
    if terminal.status_code < 400 || terminal.cancelled {
        return None;
    }
    logical
        .client_event
        .get("previous_response_id")
        .and_then(Value::as_str)
        .filter(|response_id| !response_id.trim().is_empty())
}

async fn register_persisted_continuation_before_terminal_delivery(
    state: &AppState,
    context: &WebSocketRequestContext,
    registration: PendingContinuationRegistration,
) -> Option<String> {
    let PendingContinuationRegistration {
        user_id,
        api_key_id,
        response_id,
        record,
    } = registration;
    let registry = ResponsesWebSocketContinuationRegistry::new(state.runtime_state.as_ref());
    match tokio::time::timeout(
        CONTINUATION_REGISTRATION_TIMEOUT,
        registry.register(
            user_id.as_str(),
            api_key_id.as_str(),
            response_id.as_str(),
            &record,
        ),
    )
    .await
    {
        Ok(Ok(())) => {
            debug!(
                event_name = "responses_websocket_continuation_registered",
                log_type = "event",
                transport = WEBSOCKET_LOG_TRANSPORT,
                websocket = true,
                trace_id = %context.trace_id,
                user_id = %user_id,
                api_key_id = %api_key_id,
                provider_id = %record.pinned_candidate().provider_id(),
                endpoint_id = %record.pinned_candidate().endpoint_id(),
                key_id = %record.pinned_candidate().key_id(),
                client_model = %record.client_model(),
                provider_model = %record.provider_model(),
                "gateway registered a persisted Responses continuation before terminal delivery"
            );
            Some(response_id)
        }
        Ok(Err(error)) => {
            warn!(
                event_name = "responses_websocket_continuation_registration_failed",
                log_type = "ops",
                transport = WEBSOCKET_LOG_TRANSPORT,
                websocket = true,
                trace_id = %context.trace_id,
                provider_id = %record.pinned_candidate().provider_id(),
                endpoint_id = %record.pinned_candidate().endpoint_id(),
                key_id = %record.pinned_candidate().key_id(),
                reason = error.kind(),
                "gateway failed to register a persisted Responses continuation"
            );
            None
        }
        Err(_) => {
            warn!(
                event_name = "responses_websocket_continuation_registration_failed",
                log_type = "ops",
                transport = WEBSOCKET_LOG_TRANSPORT,
                websocket = true,
                trace_id = %context.trace_id,
                provider_id = %record.pinned_candidate().provider_id(),
                endpoint_id = %record.pinned_candidate().endpoint_id(),
                key_id = %record.pinned_candidate().key_id(),
                reason = "timeout",
                timeout_ms = CONTINUATION_REGISTRATION_TIMEOUT.as_millis() as u64,
                "gateway timed out registering a persisted Responses continuation"
            );
            None
        }
    }
}

pub(super) async fn wait_for_connection_permit_loss(
    permit: Option<&aether_runtime::AdmissionPermit>,
) {
    let Some(permit) = permit else {
        std::future::pending::<()>().await;
        return;
    };
    let mut health = tokio::time::interval(Duration::from_secs(1));
    health.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
    loop {
        health.tick().await;
        if !permit.is_healthy() {
            return;
        }
    }
}

#[cfg(test)]
mod tests {
    use serde_json::json;

    use super::{
        connection_local_terminal_response_id, evicted_default_lane_continuation_response_id,
        persistable_terminal_response_id, LogicalTurn, ParsedResponsesWebSocketFrame,
    };

    #[test]
    fn cross_connection_registration_requires_explicit_store_true_and_a_success_terminal() {
        let completed = ParsedResponsesWebSocketFrame::parse(
            r#"{"type":"response.completed","response":{"id":"resp_persisted"}}"#,
        )
        .expect("valid completed event");
        let failed = ParsedResponsesWebSocketFrame::parse(
            r#"{"type":"response.failed","response":{"id":"resp_failed"}}"#,
        )
        .expect("valid failed event");

        let omitted_or_false = LogicalTurn::new(
            json!({"type": "response.create", "store": false}),
            1,
            "logical-local".to_string(),
        );
        assert_eq!(
            persistable_terminal_response_id(Some(&omitted_or_false), &completed),
            None,
            "store=false or an omitted provider-side store must remain connection-local"
        );
        assert_eq!(persistable_terminal_response_id(None, &completed), None);
        assert_eq!(
            connection_local_terminal_response_id(Some(&omitted_or_false), &completed),
            Some("resp_persisted"),
            "store=false continuations remain valid on the same physical socket"
        );
        assert_eq!(
            connection_local_terminal_response_id(None, &completed),
            None
        );

        let persisted = LogicalTurn::new(
            json!({"type": "response.create", "store": true}),
            1,
            "logical-persisted".to_string(),
        )
        .with_provider_store(true);
        assert_eq!(
            persistable_terminal_response_id(Some(&persisted), &completed),
            Some("resp_persisted")
        );
        assert_eq!(
            persistable_terminal_response_id(Some(&persisted), &failed),
            None,
            "a failed terminal must never establish cross-connection ownership"
        );

        let continuation = LogicalTurn::new(
            json!({
                "type": "response.create",
                "previous_response_id": "resp_parent"
            }),
            2,
            "logical-continuation".to_string(),
        );
        assert_eq!(
            evicted_default_lane_continuation_response_id(Some(&continuation), &failed),
            Some("resp_parent")
        );
        assert_eq!(
            evicted_default_lane_continuation_response_id(Some(&continuation), &completed),
            None
        );
    }
}