//! Responses WebSocket end-to-end coverage.
//!
//! Every test starts a protocol-aware mock upstream, seeds a throwaway SQLite
//! store, mounts the real gateway router, and drives the public
//! `/v1/responses` WebSocket the way a client would.
//!
//! The assertions deliberately reach back into the database. A turn settles its
//! billing row from a task that outlives the relay loop, so a client that saw
//! `response.completed` is not evidence that the turn was ever accounted for —
//! only the row is.

use std::path::PathBuf;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;

use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_data::repository::auth::CreateStandaloneApiKeyRecord;
use aether_data::repository::wallet::WalletLookupKey;
use aether_data::{
    DataBackends, DataLayerConfig, DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig,
};
use aether_data_contracts::repository::global_models::{
    CreateAdminGlobalModelRecord, UpsertAdminProviderModelRecord,
};
use aether_data_contracts::repository::provider_catalog::{
    StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UsageAuditListQuery};
use aether_gateway::{build_router_with_state, AppState, GatewayDataConfig, UsageRuntimeConfig};
use aether_testkit::SpawnedServer;
use axum::extract::ws::{Message as AxumWsMessage, WebSocket, WebSocketUpgrade};
use axum::extract::State;
use axum::http::HeaderMap;
use axum::response::Response;
use axum::routing::get;
use axum::Router;
use futures_util::{SinkExt, StreamExt};
use serde_json::{json, Value};
use sha2::Digest;
use tokio::net::TcpStream;
use tokio::sync::Mutex;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};

type BoxError = Box<dyn std::error::Error>;
type ClientSocket = WebSocketStream<MaybeTlsStream<TcpStream>>;

const CLIENT_API_KEY: &str = "sk-aether-responses-ws-e2e";
const PROVIDER_API_KEY: &str = "sk-upstream-responses-ws-e2e";
const PROVIDER_ID: &str = "provider-responses-ws-e2e";
const ENDPOINT_ID: &str = "endpoint-responses-ws-e2e";
const PROVIDER_KEY_ID: &str = "provider-key-responses-ws-e2e";
/// 透明重试的替代 key。只有配额重试用例会 seed 它。
const ALTERNATE_PROVIDER_KEY_ID: &str = "provider-key-responses-ws-e2e-alt";
const ALTERNATE_PROVIDER_API_KEY: &str = "sk-upstream-responses-ws-e2e-alt";
const GLOBAL_MODEL_ID: &str = "global-model-responses-ws-e2e";
const PROVIDER_MODEL_ID: &str = "provider-model-responses-ws-e2e";
const API_KEY_ID: &str = "api-key-responses-ws-e2e";
const PUBLIC_MODEL: &str = "gpt-responses-ws-e2e";
const UPSTREAM_MODEL: &str = "gpt-responses-ws-upstream";

/// 2100-01-01，保证 oauth 凭证在测试期间不会被判为过期。
const FAR_FUTURE_UNIX_SECS: u64 = 4_102_444_800;

const INPUT_TOKENS: u64 = 4;
const OUTPUT_TOKENS: u64 = 2;

/// Generous enough to absorb a loaded CI runner, short enough that a genuinely
/// lost row fails the test instead of hanging the job.
const SETTLE_TIMEOUT: Duration = Duration::from_secs(30);
const RECEIVE_TIMEOUT: Duration = Duration::from_secs(15);

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

/// The headline guarantee: a continuation stays on one physical upstream socket,
/// and both turns are billed independently.
#[tokio::test]
async fn continuation_reuses_one_upstream_connection_and_bills_both_turns() -> Result<(), BoxError>
{
    let harness = Harness::start(UpstreamBehavior::CompleteEveryTurn).await?;
    let mut client = harness.connect().await?;

    client
        .send(response_create(json!({"input": "first turn"})))
        .await?;
    let first = receive_event(&mut client, "response.completed").await?;
    assert_eq!(
        first.pointer("/response/id").and_then(Value::as_str),
        Some("resp-e2e-1")
    );

    client
        .send(response_create(json!({
            "previous_response_id": "resp-e2e-1",
            "input": "second turn"
        })))
        .await?;
    let second = receive_event(&mut client, "response.completed").await?;
    assert_eq!(
        second.pointer("/response/id").and_then(Value::as_str),
        Some("resp-e2e-2")
    );

    // The mock records each `response.create` before answering it, so both
    // turns completing means both are already on record.
    let upstream_events = harness.upstream.observed_events().await;
    assert_eq!(
        upstream_events.len(),
        2,
        "one upstream turn per client turn"
    );
    assert_eq!(
        harness.upstream.connections(),
        1,
        "the continuation must reuse the bound upstream socket"
    );
    for event in &upstream_events {
        assert_eq!(
            event.get("model").and_then(Value::as_str),
            Some(UPSTREAM_MODEL),
            "every turn is rewritten to the mapped provider model"
        );
    }
    assert_eq!(
        upstream_events[1]
            .get("previous_response_id")
            .and_then(Value::as_str),
        Some("resp-e2e-1"),
        "the continuation id survives provider body normalization"
    );
    assert_eq!(
        harness.upstream.authorization_headers().await,
        vec![Some(format!("Bearer {PROVIDER_API_KEY}"))],
        "the upstream is opened with the configured provider key"
    );

    let audits = harness
        .usage_audits_where(2, "billed turns", is_billed)
        .await?;
    assert_eq!(audits.len(), 2, "each response.create bills separately");
    for audit in &audits {
        assert!(
            audit.is_websocket(),
            "turns are recorded as WebSocket usage, metadata: {:?}",
            audit.request_metadata
        );
        assert_eq!(audit.model, PUBLIC_MODEL);
        assert_eq!(audit.input_tokens, INPUT_TOKENS);
        assert_eq!(audit.output_tokens, OUTPUT_TOKENS);
        assert_eq!(audit.total_tokens, INPUT_TOKENS + OUTPUT_TOKENS);
        assert_eq!(audit.status_code, Some(200));
    }
    assert_ne!(
        audits[0].request_id, audits[1].request_id,
        "each turn gets its own logical request identity"
    );

    Ok(())
}

/// A client that walks away before the provider produced anything must settle
/// as a void row: nothing was produced, so nothing is billed.
///
/// This is the path with no protocol event to announce it: the relay loop owns
/// the turn, and losing the client is an exit the upstream never reports.
///
/// The mirror case — the provider *did* reach a terminal event and only the last
/// hop to the client failed — is billed instead. That one cannot be pinned here:
/// it depends on the relay loop's `select!` observing the upstream terminal frame
/// before it observes the closed client socket, which is a race by construction.
/// It is covered deterministically by the relay-level unit tests
/// `a_provider_terminal_that_reaches_a_closed_client_socket_is_still_billed` and
/// `a_closed_client_socket_before_any_terminal_still_voids_the_bill`.
#[tokio::test]
async fn client_disconnect_before_any_provider_output_settles_a_void_row() -> Result<(), BoxError> {
    let harness = Harness::start(UpstreamBehavior::StallAfterCreated).await?;
    let mut client = harness.connect().await?;

    client
        .send(response_create(json!({"input": "abandoned turn"})))
        .await?;
    // Leave only once the turn is genuinely in flight upstream, so this covers
    // an interrupted turn rather than racing turn start.
    receive_event(&mut client, "response.created").await?;
    drop(client);

    let audits = harness
        .usage_audits_where(1, "settled turns", |audit| !is_pending(audit))
        .await?;
    assert_eq!(audits.len(), 1, "the abandoned turn is still accounted for");
    let audit = &audits[0];
    assert_eq!(audit.model, PUBLIC_MODEL);
    assert!(
        !is_pending(audit),
        "an abandoned turn must not be left pending: {audit:?}"
    );
    // The provider never emitted a terminal event, so this row stays void.
    // Only a reached provider terminal survives a client delivery failure.
    assert!(
        !is_billed(audit),
        "a turn with no provider output must not be billed: {audit:?}"
    );
    assert_eq!(
        audit.status, "cancelled",
        "a client that left before any provider output settles as cancelled: {audit:?}"
    );
    assert_eq!(audit.status_code, Some(499));

    Ok(())
}

/// An upstream that dies mid-turn must surface an error and still settle.
#[tokio::test]
async fn upstream_drop_mid_turn_reports_an_error_and_settles_the_usage_row() -> Result<(), BoxError>
{
    let harness = Harness::start(UpstreamBehavior::CloseAfterCreated).await?;
    let mut client = harness.connect().await?;

    client
        .send(response_create(json!({"input": "doomed turn"})))
        .await?;
    let error = receive_error_or_close(&mut client)
        .await?
        .ok_or("gateway closed without telling the client why")?;
    assert_eq!(error.get("type").and_then(Value::as_str), Some("error"));

    let audits = harness
        .usage_audits_where(1, "settled turns", |audit| !is_pending(audit))
        .await?;
    assert_eq!(audits.len(), 1, "the failed turn is still accounted for");
    let audit = &audits[0];
    assert!(
        !is_pending(audit),
        "a failed turn must not be left pending: {audit:?}"
    );

    Ok(())
}

/// 供应商配额耗尽后的透明重试：客户端不该看到 429，两个 attempt 都要结算。
///
/// 第一个 attempt 拿到 Codex 的 `usage_limit_reached`，网关换到第二把 key 重开一条
/// 上游连接重放同一个 `response.create`。C6 之前，重试的规划发生在旧 attempt 结算
/// 之前：规划读到的是旧 attempt 还没投射的 health / adaptive / pool 状态，而且旧
/// attempt 的 pool key lease 还被它自己占着。
///
/// 顺序本身在这里无法确定性断言（结算与规划都在同一个任务里、DB 里看不到先后），
/// 由 lifecycle 的单测确定性覆盖；这个用例保证整条路径真的能跑通，并且两个
/// attempt 都留下了终态记账行。
#[tokio::test]
async fn provider_quota_exhaustion_transparently_retries_onto_another_key() -> Result<(), BoxError>
{
    let harness = Harness::start_with_fixture(
        UpstreamBehavior::QuotaExhaustedThenComplete,
        ProviderFixture::CodexKeyPair,
    )
    .await?;
    let mut client = harness.connect().await?;

    client
        .send(response_create(json!({"input": "retry after quota"})))
        .await?;

    // 客户端只应该看到重试之后那次成功的响应，看不到 429。
    let completed = receive_event(&mut client, "response.completed").await?;
    assert_eq!(
        completed
            .pointer("/response/status")
            .and_then(Value::as_str),
        Some("completed")
    );

    // 上游被连了两次：配额耗尽的那条 + 重试用的那条。
    assert_eq!(
        harness.upstream.connections(),
        2,
        "the transparent retry must open a second upstream connection"
    );
    let observed = harness.upstream.observed_events().await;
    assert_eq!(
        observed.len(),
        2,
        "the same response.create must be replayed once"
    );

    // 两把不同的 key 被用过：重试不能落回那把已经耗尽的 key。
    let authorizations = harness.upstream.authorization_headers().await;
    assert_eq!(authorizations.len(), 2);
    assert_ne!(
        authorizations[0], authorizations[1],
        "the retry must not reuse the exhausted key: {authorizations:?}"
    );

    // 两个 attempt 各自留下一条终态行：配额失败的那条 + 成功计费的那条。
    let audits = harness
        .usage_audits_where(2, "settled attempts", |audit| !is_pending(audit))
        .await?;
    let settled = audits
        .iter()
        .filter(|audit| !is_pending(audit))
        .collect::<Vec<_>>();
    assert_eq!(
        settled.len(),
        2,
        "both attempts must reach a terminal accounting row: {:?}",
        audits
            .iter()
            .map(|audit| (audit.status.clone(), audit.status_code, audit.total_tokens))
            .collect::<Vec<_>>()
    );
    assert!(
        settled.iter().any(|audit| audit.status_code == Some(429)),
        "the exhausted attempt keeps its own 429 row: {:?}",
        settled
            .iter()
            .map(|audit| (audit.status.clone(), audit.status_code))
            .collect::<Vec<_>>()
    );

    let billed = harness
        .usage_audits_where(1, "the billed retry attempt", is_billed)
        .await?;
    let retry = billed
        .iter()
        .find(|audit| is_billed(audit))
        .ok_or("the successful retry attempt must be billed")?;
    assert_eq!(
        retry.total_tokens,
        INPUT_TOKENS + OUTPUT_TOKENS,
        "the retry attempt is billed for what it actually consumed"
    );

    client.close(None).await?;
    Ok(())
}

/// 脱敏的另一半：请求侧把真实 PII 换成占位符发给上游，响应侧必须在推给客户端之前
/// 换回真实值。
///
/// 上游把收到的 `input` 原样回显，所以它回来的就是占位符——这一条同时钉住了两个
/// 方向：上游不能看到原文，客户端不能看到占位符。
#[tokio::test]
async fn redacted_pii_is_restored_before_the_client_sees_a_provider_frame() -> Result<(), BoxError>
{
    const CLIENT_EMAIL: &str = "responses.ws.pii@example.com";

    let harness = Harness::start_with_pii_redaction(UpstreamBehavior::EchoInputBack).await?;
    let mut client = harness.connect().await?;

    client
        .send(response_create(
            json!({"input": format!("my mail is {CLIENT_EMAIL}")}),
        ))
        .await?;

    let delta = receive_event(&mut client, "response.output_text.delta").await?;
    let delta_text = delta
        .get("delta")
        .and_then(Value::as_str)
        .ok_or("the provider delta must carry text")?;
    assert!(
        delta_text.contains(CLIENT_EMAIL),
        "the client must receive the restored value: {delta_text}"
    );
    assert!(
        !delta_text.contains("<AETHER:"),
        "no redaction placeholder may reach the client: {delta_text}"
    );

    // 请求侧仍然成立：上游只看到占位符，看不到原文。
    let upstream_events = harness.upstream.observed_events().await;
    assert_eq!(upstream_events.len(), 1);
    let upstream_input = upstream_events[0]
        .get("input")
        .and_then(Value::as_str)
        .ok_or("the upstream request must carry input")?;
    assert!(
        !upstream_input.contains(CLIENT_EMAIL),
        "the upstream must never see the raw PII: {upstream_input}"
    );
    assert!(
        upstream_input.contains("<AETHER:EMAIL:"),
        "the upstream must see the placeholder: {upstream_input}"
    );

    // 还原只发生在最后一跳：终态照常到达，计费不受影响。
    let completed = receive_event(&mut client, "response.completed").await?;
    assert_eq!(
        completed
            .pointer("/response/status")
            .and_then(Value::as_str),
        Some("completed")
    );
    let audits = harness
        .usage_audits_where(1, "the billed redacted turn", is_billed)
        .await?;
    assert_eq!(audits.len(), 1);
    assert_eq!(audits[0].total_tokens, INPUT_TOKENS + OUTPUT_TOKENS);

    client.close(None).await?;
    Ok(())
}

fn is_pending(audit: &StoredRequestUsageAudit) -> bool {
    audit.status.eq_ignore_ascii_case("pending")
}

/// A turn that finished accounting: settled, and carrying what it consumed.
fn is_billed(audit: &StoredRequestUsageAudit) -> bool {
    !is_pending(audit) && audit.total_tokens > 0
}

// ---------------------------------------------------------------------------
// Harness
// ---------------------------------------------------------------------------

/// A live gateway wired to a mock Responses WebSocket upstream over a throwaway
/// SQLite store.
struct Harness {
    database: TemporarySqlite,
    upstream: Arc<MockUpstreamState>,
    websocket_url: String,
    _upstream_server: SpawnedServer,
    _gateway_server: SpawnedServer,
}

/// 供应商夹具形态。
///
/// 透明配额重试只有 Codex adapter 会开启（`retry_current_turn: true` 只从
/// codex.rs 出），而且重试要有第二把 key 可挑，否则规划直接判无可用供应商。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ProviderFixture {
    /// 单个 openai 类型供应商、单把 key。
    SingleOpenAiKey,
    /// codex 类型供应商 + 两把 key：第一把配额耗尽后重试落到第二把。
    CodexKeyPair,
}

impl ProviderFixture {
    const fn provider_type(self) -> &'static str {
        match self {
            Self::SingleOpenAiKey => "openai",
            Self::CodexKeyPair => "codex",
        }
    }

    const fn has_alternate_key(self) -> bool {
        matches!(self, Self::CodexKeyPair)
    }
}

/// 这条用例要不要打开 chat PII 脱敏模块。
///
/// 默认关闭：其余用例都靠原文 body 断言上游看到了什么，打开脱敏会把断言目标换成
/// 占位符。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum PiiRedaction {
    Disabled,
    Enabled,
}

impl PiiRedaction {
    const fn is_enabled(self) -> bool {
        matches!(self, Self::Enabled)
    }
}

impl Harness {
    async fn start(behavior: UpstreamBehavior) -> Result<Self, BoxError> {
        Self::start_with(
            behavior,
            ProviderFixture::SingleOpenAiKey,
            PiiRedaction::Disabled,
        )
        .await
    }

    async fn start_with_fixture(
        behavior: UpstreamBehavior,
        fixture: ProviderFixture,
    ) -> Result<Self, BoxError> {
        Self::start_with(behavior, fixture, PiiRedaction::Disabled).await
    }

    async fn start_with_pii_redaction(behavior: UpstreamBehavior) -> Result<Self, BoxError> {
        Self::start_with(
            behavior,
            ProviderFixture::SingleOpenAiKey,
            PiiRedaction::Enabled,
        )
        .await
    }

    async fn start_with(
        behavior: UpstreamBehavior,
        fixture: ProviderFixture,
        redaction: PiiRedaction,
    ) -> Result<Self, BoxError> {
        let upstream = Arc::new(MockUpstreamState::new(behavior));
        let upstream_server =
            SpawnedServer::start(mock_upstream_router(Arc::clone(&upstream))).await?;

        let database = TemporarySqlite::new();
        prepare_and_seed_database(
            &database.config,
            upstream_server.base_url(),
            fixture,
            redaction,
        )
        .await?;

        let data_config = GatewayDataConfig::from_database_config(database.config.clone())
            .with_encryption_key(DEVELOPMENT_ENCRYPTION_KEY);
        let state = AppState::new()?
            .with_data_config_and_background_isolation(data_config, false)?
            // The usage runtime defaults to disabled, which silently turns every
            // terminal usage write into a no-op. Without this the suite could
            // not observe billing at all. Queueing stays off so the terminal
            // write lands through the in-process path instead of Redis.
            .with_usage_runtime_config(UsageRuntimeConfig {
                enabled: true,
                ..UsageRuntimeConfig::default()
            })?;
        let gateway_server = SpawnedServer::start(build_router_with_state(state)).await?;
        let websocket_url = format!(
            "{}/v1/responses",
            gateway_server.base_url().replacen("http://", "ws://", 1)
        );

        Ok(Self {
            database,
            upstream,
            websocket_url,
            _upstream_server: upstream_server,
            _gateway_server: gateway_server,
        })
    }

    async fn connect(&self) -> Result<ClientSocket, BoxError> {
        let mut request = self.websocket_url.clone().into_client_request()?;
        request.headers_mut().insert(
            "authorization",
            http::HeaderValue::from_str(&format!("Bearer {CLIENT_API_KEY}"))?,
        );
        let (socket, response) =
            tokio::time::timeout(RECEIVE_TIMEOUT, tokio_tungstenite::connect_async(request))
                .await
                .map_err(|_| "timed out connecting to the gateway WebSocket")??;
        if response.status() != http::StatusCode::SWITCHING_PROTOCOLS {
            return Err(
                format!("unexpected gateway handshake status: {}", response.status()).into(),
            );
        }
        Ok(socket)
    }

    /// Waits until `expected` usage rows satisfy `settled`.
    ///
    /// A row is created `Pending` at turn start and reaches its final shape
    /// through several independent writes, so "no longer pending" does not imply
    /// "finished": a row can briefly read as completed with zero tokens and no
    /// WebSocket metadata before the terminal write lands. Each caller waits for
    /// the specific end state it is about to assert.
    async fn usage_audits_where(
        &self,
        expected: usize,
        what: &str,
        settled: impl Fn(&StoredRequestUsageAudit) -> bool,
    ) -> Result<Vec<StoredRequestUsageAudit>, BoxError> {
        let deadline = tokio::time::Instant::now() + SETTLE_TIMEOUT;
        loop {
            let audits = self.usage_audits().await?;
            if audits.iter().filter(|audit| settled(audit)).count() >= expected {
                return Ok(audits);
            }
            if tokio::time::Instant::now() >= deadline {
                let observed = audits
                    .iter()
                    .map(|audit| {
                        format!(
                            "{} status={} code={:?} tokens={} websocket={}",
                            audit.request_id,
                            audit.status,
                            audit.status_code,
                            audit.total_tokens,
                            audit.is_websocket()
                        )
                    })
                    .collect::<Vec<_>>();
                return Err(format!(
                    "timed out waiting for {expected} {what}; observed {}: {observed:?}",
                    audits.len()
                )
                .into());
            }
            tokio::time::sleep(Duration::from_millis(25)).await;
        }
    }

    /// Reads the persisted audit rows, oldest first.
    ///
    /// Opens its own handle per call rather than holding one for the lifetime of
    /// the harness: the gateway keeps its own pool on the same SQLite file for
    /// the whole test, and an idle second pool only adds contention.
    async fn usage_audits(&self) -> Result<Vec<StoredRequestUsageAudit>, BoxError> {
        let backends = DataBackends::from_config(DataLayerConfig::from_database(
            self.database.config.clone(),
        ))?;
        let audits = backends
            .read()
            .usage()
            .ok_or("usage reader unavailable")?
            .list_usage_audits(&UsageAuditListQuery {
                limit: Some(50),
                newest_first: false,
                ..UsageAuditListQuery::default()
            })
            .await?;
        drop(backends);
        Ok(audits)
    }
}

// ---------------------------------------------------------------------------
// Client protocol helpers
// ---------------------------------------------------------------------------

/// Builds a `response.create` frame for the seeded public model.
fn response_create(fields: Value) -> Message {
    let mut event = json!({"type": "response.create", "model": PUBLIC_MODEL});
    let object = event
        .as_object_mut()
        .expect("the literal above is an object");
    for (key, value) in fields
        .as_object()
        .expect("response.create fields must be an object")
    {
        object.insert(key.clone(), value.clone());
    }
    Message::Text(event.to_string().into())
}

/// Reads frames until `expected_type` arrives, failing fast on a gateway error.
async fn receive_event<S>(
    socket: &mut WebSocketStream<S>,
    expected_type: &str,
) -> Result<Value, BoxError>
where
    S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
{
    tokio::time::timeout(RECEIVE_TIMEOUT, async {
        loop {
            let message = socket
                .next()
                .await
                .ok_or("gateway WebSocket closed before the expected event")??;
            match message {
                Message::Text(text) => {
                    let event: Value = serde_json::from_str(text.as_ref())?;
                    match event.get("type").and_then(Value::as_str) {
                        Some("error") => {
                            return Err(format!("gateway returned an error event: {event}").into())
                        }
                        Some(event_type) if event_type == expected_type => return Ok(event),
                        _ => {}
                    }
                }
                Message::Ping(payload) => socket.send(Message::Pong(payload)).await?,
                Message::Close(frame) => {
                    return Err(format!("gateway closed before {expected_type}: {frame:?}").into())
                }
                _ => {}
            }
        }
    })
    .await
    .map_err(|_| format!("timed out waiting for {expected_type}"))?
}

/// Drains the socket until the gateway reports an error or hangs up.
///
/// Returns the error event when one arrives, `None` when the gateway closed
/// without explaining itself.
async fn receive_error_or_close<S>(
    socket: &mut WebSocketStream<S>,
) -> Result<Option<Value>, BoxError>
where
    S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
{
    tokio::time::timeout(RECEIVE_TIMEOUT, async {
        loop {
            let Some(message) = socket.next().await else {
                return Ok(None);
            };
            match message? {
                Message::Text(text) => {
                    let event: Value = serde_json::from_str(text.as_ref())?;
                    if event.get("type").and_then(Value::as_str) == Some("error") {
                        return Ok(Some(event));
                    }
                }
                Message::Ping(payload) => socket.send(Message::Pong(payload)).await?,
                Message::Close(_) => return Ok(None),
                _ => {}
            }
        }
    })
    .await
    .map_err(|_| "timed out waiting for a gateway error or close")?
}

// ---------------------------------------------------------------------------
// Mock upstream
// ---------------------------------------------------------------------------

/// How the mock upstream answers a `response.create`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum UpstreamBehavior {
    /// Announce, stream one delta, and complete — the ordinary turn.
    CompleteEveryTurn,
    /// Announce the response and then go quiet, leaving the turn in flight.
    StallAfterCreated,
    /// Announce the response and then hang up mid-turn.
    CloseAfterCreated,
    /// 第一轮只回一个 Codex 配额耗尽错误，之后的每一轮正常完成。
    ///
    /// 第一轮刻意不发 `response.created`：任何标准 `response.*` 事件都会让
    /// codex adapter 把这一轮判成 replay-unsafe，透明重试就不会发生。
    QuotaExhaustedThenComplete,
    /// 把收到的 `input` 原样回显成一个 delta，再正常完成。
    ///
    /// 上游看到的是脱敏后的 body，所以回显出来的就是占位符——正是响应侧还原要处理
    /// 的形状。
    EchoInputBack,
}

#[derive(Debug)]
struct MockUpstreamState {
    behavior: UpstreamBehavior,
    connections: AtomicUsize,
    events: Mutex<Vec<Value>>,
    authorization_headers: Mutex<Vec<Option<String>>>,
}

impl MockUpstreamState {
    fn new(behavior: UpstreamBehavior) -> Self {
        Self {
            behavior,
            connections: AtomicUsize::new(0),
            events: Mutex::new(Vec::new()),
            authorization_headers: Mutex::new(Vec::new()),
        }
    }

    fn connections(&self) -> usize {
        self.connections.load(Ordering::Acquire)
    }

    async fn observed_events(&self) -> Vec<Value> {
        self.events.lock().await.clone()
    }

    async fn authorization_headers(&self) -> Vec<Option<String>> {
        self.authorization_headers.lock().await.clone()
    }
}

fn mock_upstream_router(state: Arc<MockUpstreamState>) -> Router {
    Router::new()
        .route("/v1/responses", get(mock_responses_websocket))
        .with_state(state)
}

async fn mock_responses_websocket(
    State(state): State<Arc<MockUpstreamState>>,
    headers: HeaderMap,
    ws: WebSocketUpgrade,
) -> Response {
    let authorization = headers
        .get("authorization")
        .and_then(|value| value.to_str().ok())
        .map(str::to_string);
    ws.on_upgrade(move |socket| run_mock_upstream(socket, state, authorization))
}

async fn run_mock_upstream(
    mut socket: WebSocket,
    state: Arc<MockUpstreamState>,
    authorization: Option<String>,
) {
    state.connections.fetch_add(1, Ordering::AcqRel);
    state.authorization_headers.lock().await.push(authorization);
    while let Some(message) = socket.recv().await {
        let Ok(message) = message else {
            break;
        };
        match message {
            AxumWsMessage::Text(text) => {
                let Ok(event) = serde_json::from_str::<Value>(text.as_str()) else {
                    break;
                };
                if event.get("type").and_then(Value::as_str) != Some("response.create") {
                    continue;
                }
                let echoed_input = event
                    .get("input")
                    .and_then(Value::as_str)
                    .unwrap_or_default()
                    .to_string();
                let turn = {
                    let mut events = state.events.lock().await;
                    events.push(event);
                    events.len()
                };
                let response_id = format!("resp-e2e-{turn}");
                match state.behavior {
                    UpstreamBehavior::CompleteEveryTurn => {
                        if send_mock_turn(&mut socket, &response_id).await.is_err() {
                            break;
                        }
                    }
                    UpstreamBehavior::StallAfterCreated => {
                        if send_mock_created(&mut socket, &response_id).await.is_err() {
                            break;
                        }
                    }
                    UpstreamBehavior::CloseAfterCreated => {
                        let _ = send_mock_created(&mut socket, &response_id).await;
                        break;
                    }
                    UpstreamBehavior::QuotaExhaustedThenComplete => {
                        if turn == 1 {
                            let _ =
                                send_mock_event(&mut socket, codex_quota_exhausted_error()).await;
                            break;
                        }
                        if send_mock_turn(&mut socket, &response_id).await.is_err() {
                            break;
                        }
                    }
                    UpstreamBehavior::EchoInputBack => {
                        if send_mock_turn_with_delta(&mut socket, &response_id, &echoed_input)
                            .await
                            .is_err()
                        {
                            break;
                        }
                    }
                }
            }
            AxumWsMessage::Ping(payload) => {
                if socket.send(AxumWsMessage::Pong(payload)).await.is_err() {
                    break;
                }
            }
            AxumWsMessage::Close(_) => break,
            _ => {}
        }
    }
}

/// Codex 的账户级配额耗尽信号。
///
/// `status_code: 429` + `error.type: usage_limit_reached` 是 adapter 识别
/// 「配额耗尽、可透明重试」的最小载荷：解析出的元数据被强制标上
/// `limit_reached: true`，于是 drain 指令带着 `retry_current_turn: true` 下来。
fn codex_quota_exhausted_error() -> Value {
    json!({
        "type": "error",
        "status_code": 429,
        "error": {
            "type": "usage_limit_reached",
            "message": "You have hit your usage limit",
            "plan_type": "plus",
            "resets_in_seconds": 3_600
        }
    })
}

async fn send_mock_created(socket: &mut WebSocket, response_id: &str) -> Result<(), axum::Error> {
    send_mock_event(
        socket,
        json!({
            "type": "response.created",
            "response": {
                "id": response_id,
                "status": "in_progress",
                "model": UPSTREAM_MODEL
            }
        }),
    )
    .await
}

async fn send_mock_turn(socket: &mut WebSocket, response_id: &str) -> Result<(), axum::Error> {
    send_mock_turn_with_delta(socket, response_id, "hello").await
}

async fn send_mock_turn_with_delta(
    socket: &mut WebSocket,
    response_id: &str,
    delta: &str,
) -> Result<(), axum::Error> {
    send_mock_created(socket, response_id).await?;
    send_mock_event(
        socket,
        json!({
            "type": "response.output_text.delta",
            "response_id": response_id,
            "delta": delta
        }),
    )
    .await?;
    send_mock_event(
        socket,
        json!({
            "type": "response.completed",
            "response": {
                "id": response_id,
                "status": "completed",
                "model": UPSTREAM_MODEL,
                "output": [],
                "usage": {
                    "input_tokens": INPUT_TOKENS,
                    "output_tokens": OUTPUT_TOKENS,
                    "total_tokens": INPUT_TOKENS + OUTPUT_TOKENS
                }
            }
        }),
    )
    .await
}

async fn send_mock_event(socket: &mut WebSocket, event: Value) -> Result<(), axum::Error> {
    socket
        .send(AxumWsMessage::Text(event.to_string().into()))
        .await
}

// ---------------------------------------------------------------------------
// Seeded data store
// ---------------------------------------------------------------------------

struct TemporarySqlite {
    directory: PathBuf,
    config: SqlDatabaseConfig,
}

impl TemporarySqlite {
    fn new() -> Self {
        let directory = std::env::temp_dir().join(format!(
            "aether-responses-ws-e2e-{}-{}",
            std::process::id(),
            uuid::Uuid::new_v4()
        ));
        let database_path = directory.join("aether.db");
        Self {
            directory,
            config: SqlDatabaseConfig {
                driver: DatabaseDriver::Sqlite,
                url: format!("sqlite://{}", database_path.display()),
                pool: SqlPoolConfig {
                    min_connections: 1,
                    max_connections: 4,
                    acquire_timeout_ms: 5_000,
                    idle_timeout_ms: 30_000,
                    max_lifetime_ms: 300_000,
                    statement_cache_capacity: 64,
                    require_ssl: false,
                },
            },
        }
    }
}

impl Drop for TemporarySqlite {
    fn drop(&mut self) {
        let _ = std::fs::remove_dir_all(&self.directory);
    }
}

async fn prepare_and_seed_database(
    database: &SqlDatabaseConfig,
    upstream_base_url: &str,
    fixture: ProviderFixture,
    redaction: PiiRedaction,
) -> Result<(), BoxError> {
    let backends = DataBackends::from_config(DataLayerConfig::from_database(database.clone()))?;
    let pending = backends
        .prepare_database_for_startup()
        .await?
        .unwrap_or_default();
    if !pending.is_empty() {
        backends.run_database_migrations().await?;
    }

    seed_provider_catalog(&backends, upstream_base_url, fixture).await?;
    seed_models(&backends).await?;
    let user_id = seed_user(&backends).await?;
    seed_client_api_key(&backends, &user_id).await?;
    if redaction.is_enabled() {
        seed_chat_pii_redaction(&backends).await?;
    }

    let candidates = backends
        .read()
        .minimal_candidate_selection()
        .ok_or("candidate selection reader unavailable")?
        .list_for_exact_api_format_and_requested_model("openai:responses", PUBLIC_MODEL)
        .await?;
    if !candidates.iter().any(|candidate| {
        candidate.provider_id == PROVIDER_ID
            && candidate.endpoint_id == ENDPOINT_ID
            && candidate.key_id == PROVIDER_KEY_ID
    }) {
        return Err("seeded Responses WebSocket candidate is not visible".into());
    }
    drop(backends);
    Ok(())
}

async fn seed_provider_catalog(
    backends: &DataBackends,
    upstream_base_url: &str,
    fixture: ProviderFixture,
) -> Result<(), BoxError> {
    let writer = backends
        .write()
        .provider_catalog()
        .ok_or("provider catalog writer unavailable")?;
    writer
        .create_provider(
            &StoredProviderCatalogProvider::new(
                PROVIDER_ID.to_string(),
                "Responses WebSocket E2E".to_string(),
                None,
                fixture.provider_type().to_string(),
            )?
            .with_transport_fields(
                true,
                false,
                false,
                None,
                Some(0),
                None,
                Some(30.0),
                Some(10.0),
                Some(json!({"responses_websocket": {"enabled": true}})),
            ),
            None,
        )
        .await?;
    writer
        .create_endpoint(
            &StoredProviderCatalogEndpoint::new(
                ENDPOINT_ID.to_string(),
                PROVIDER_ID.to_string(),
                "openai:responses".to_string(),
                Some("openai".to_string()),
                Some("responses".to_string()),
                true,
            )?
            .with_transport_fields(
                upstream_base_url.trim_end_matches('/').to_string(),
                None,
                None,
                Some(0),
                Some("/v1/responses".to_string()),
                None,
                None,
                None,
            )?,
        )
        .await?;
    writer
        .create_key(&catalog_key(PROVIDER_KEY_ID, PROVIDER_API_KEY, fixture)?)
        .await?;
    if fixture.has_alternate_key() {
        writer
            .create_key(&catalog_key(
                ALTERNATE_PROVIDER_KEY_ID,
                ALTERNATE_PROVIDER_API_KEY,
                fixture,
            )?)
            .await?;
    }
    Ok(())
}

/// 一把健康、可服务本用例模型的 key。
///
/// codex 类型的候选要求 `auth_type = oauth`（见 candidate_selection 的
/// provider_type 约束），所以配额重试夹具走 oauth，凭证是一份未过期的
/// access_token。
fn catalog_key(
    key_id: &str,
    secret: &str,
    fixture: ProviderFixture,
) -> Result<StoredProviderCatalogKey, BoxError> {
    let oauth = fixture.has_alternate_key();
    let auth_type = if oauth { "oauth" } else { "api_key" };
    let auth_config = if oauth {
        Some(encrypt_python_fernet_plaintext(
            DEVELOPMENT_ENCRYPTION_KEY,
            &json!({
                "access_token": secret,
                "refresh_token": format!("{secret}-refresh"),
                "account_id": format!("{key_id}-account"),
                "expires_at": FAR_FUTURE_UNIX_SECS,
            })
            .to_string(),
        )?)
    } else {
        None
    };
    Ok(StoredProviderCatalogKey::new(
        key_id.to_string(),
        PROVIDER_ID.to_string(),
        "Responses WebSocket E2E".to_string(),
        auth_type.to_string(),
        Some(json!({"streaming": true})),
        true,
    )?
    .with_transport_fields(
        Some(json!(["openai:responses"])),
        encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, secret)?,
        auth_config,
        None,
        Some(json!({"openai:responses": 1})),
        Some(json!([PUBLIC_MODEL, UPSTREAM_MODEL])),
        None,
        None,
        None,
    )?
    .with_health_fields(
        Some(json!({"openai:responses": {"status": "healthy"}})),
        Some(json!({"openai:responses": {"state": "closed"}})),
    ))
}

async fn seed_models(backends: &DataBackends) -> Result<(), BoxError> {
    let writer = backends
        .write()
        .global_models()
        .ok_or("global model writer unavailable")?;
    writer
        .create_admin_global_model(&CreateAdminGlobalModelRecord::new(
            GLOBAL_MODEL_ID.to_string(),
            PUBLIC_MODEL.to_string(),
            PUBLIC_MODEL.to_string(),
            true,
            Some(0.0),
            None,
            Some(json!({"streaming": true, "chat": true})),
            Some(json!({"model_mappings": [UPSTREAM_MODEL]})),
        )?)
        .await?;
    writer
        .create_admin_provider_model(&UpsertAdminProviderModelRecord::new(
            PROVIDER_MODEL_ID.to_string(),
            PROVIDER_ID.to_string(),
            GLOBAL_MODEL_ID.to_string(),
            UPSTREAM_MODEL.to_string(),
            Some(json!([{
                "name": UPSTREAM_MODEL,
                "priority": 0,
                "api_formats": ["openai:responses"],
                "endpoint_ids": [ENDPOINT_ID]
            }])),
            Some(0.0),
            None,
            Some(false),
            Some(false),
            Some(true),
            Some(false),
            Some(false),
            true,
            true,
            Some(json!({"responses_websocket_e2e": true})),
        )?)
        .await?;
    Ok(())
}

async fn seed_user(backends: &DataBackends) -> Result<String, BoxError> {
    let users = backends.read().users().ok_or("user reader unavailable")?;
    let user = users
        .create_local_auth_user_with_settings(
            Some("responses-ws-e2e@example.test".to_string()),
            true,
            "responses-ws-e2e".to_string(),
            "disabled-password".to_string(),
            "user".to_string(),
            Some(vec![PROVIDER_ID.to_string()]),
            Some(vec!["openai:responses".to_string()]),
            Some(vec![PUBLIC_MODEL.to_string()]),
            None,
        )
        .await?
        .ok_or("failed to create E2E user")?;
    let wallets = backends
        .read()
        .wallets()
        .ok_or("wallet reader unavailable")?;
    wallets
        .initialize_auth_user_wallet(&user.id, 0.0, true)
        .await?;
    Ok(user.id)
}

async fn seed_client_api_key(backends: &DataBackends, user_id: &str) -> Result<(), BoxError> {
    backends
        .write()
        .auth_api_keys()
        .ok_or("auth API key writer unavailable")?
        .create_standalone_api_key(CreateStandaloneApiKeyRecord {
            user_id: user_id.to_string(),
            api_key_id: API_KEY_ID.to_string(),
            key_hash: sha256_hex(CLIENT_API_KEY),
            key_encrypted: Some(CLIENT_API_KEY.to_string()),
            name: Some("Responses WebSocket E2E".to_string()),
            allowed_providers: Some(vec![PROVIDER_ID.to_string()]),
            allowed_api_formats: Some(vec!["openai:responses".to_string()]),
            allowed_models: Some(vec![PUBLIC_MODEL.to_string()]),
            ip_rules: None,
            rate_limit: Some(0),
            concurrent_limit: None,
            force_capabilities: None,
            is_active: true,
            expires_at_unix_secs: None,
            auto_delete_on_expiry: false,
            total_requests: 0,
            total_tokens: 0,
            total_cost_usd: 0.0,
        })
        .await?;
    backends
        .read()
        .wallets()
        .ok_or("wallet reader unavailable")?
        .initialize_auth_api_key_wallet(API_KEY_ID, 0.0, true)
        .await?;
    if backends
        .read()
        .wallets()
        .ok_or("wallet reader unavailable")?
        .find(WalletLookupKey::ApiKeyId(API_KEY_ID))
        .await?
        .is_none()
    {
        return Err("failed to initialize E2E API key wallet".into());
    }
    Ok(())
}

/// 打开 chat PII 脱敏：系统模块开关 + 这把 client key 的 feature 开关。
///
/// 规则集刻意不写：缺省即内置规则（含 email 规则），和生产上「只打开开关」的最小
/// 配置一致。
async fn seed_chat_pii_redaction(backends: &DataBackends) -> Result<(), BoxError> {
    backends
        .upsert_system_config_entry("module.chat_pii_redaction.enabled", &json!(true), None)
        .await?;
    backends
        .write()
        .auth_api_keys()
        .ok_or("auth API key writer unavailable")?
        .set_standalone_api_key_feature_settings(
            API_KEY_ID,
            Some(json!({"chat_pii_redaction": {"enabled": true}})),
        )
        .await?
        .ok_or("failed to enable chat PII redaction on the E2E API key")?;
    Ok(())
}

fn sha256_hex(value: &str) -> String {
    let mut hasher = sha2::Sha256::new();
    hasher.update(value.as_bytes());
    format!("{:x}", hasher.finalize())
}