codex-proxy-rs/backend/crates/gateway-core/tests/engine/execution.rs

use std::{
    collections::{BTreeMap, BTreeSet},
    sync::{
        Arc, Mutex,
        atomic::{AtomicBool, AtomicUsize, Ordering},
    },
    time::{Duration, Instant, SystemTime},
};

use gateway_core::engine::budget::{ClientBudgetCharge, ClientBudgetError, ClientBudgetPort};
use gateway_core::error::GatewayError;

#[derive(Default)]
struct Admissions {
    active: Arc<AtomicBool>,
    limits: Mutex<Vec<RateLimits>>,
    release_gate: Mutex<Option<oneshot::Receiver<()>>>,
    releases: AtomicUsize,
}

impl ClientAdmissionPort for Admissions {
    fn abandon(
        &self,
        key: &gateway_core::policy::ClientApiKeyId,
        request: &gateway_core::engine::ModelRequestId,
    ) {
        let _ = futures::FutureExt::now_or_never(self.release(key, request));
    }

    fn admit(
        &self,
        request: ClientAdmissionRequest,
    ) -> BoxFuture<'_, Result<ClientAdmissionDecision, ClientAdmissionError>> {
        Box::pin(async move {
            self.limits.lock().unwrap().push(request.limits);
            assert!(!self.active.swap(true, Ordering::SeqCst));
            Ok(ClientAdmissionDecision::Granted)
        })
    }
    fn release<'a>(
        &'a self,
        _: &'a ClientApiKeyId,
        _: &'a ModelRequestId,
    ) -> BoxFuture<'a, Result<bool, ClientAdmissionError>> {
        Box::pin(async {
            self.releases.fetch_add(1, Ordering::SeqCst);
            let gate = self.release_gate.lock().unwrap().take();
            if let Some(gate) = gate {
                gate.await.expect("release gate");
            }
            Ok(self.active.swap(false, Ordering::SeqCst))
        })
    }
    fn restore(
        &self,
        _: ClientAdmissionRecovery,
    ) -> BoxFuture<'_, Result<ClientAdmissionRestoreResult, ClientAdmissionError>> {
        Box::pin(async { Ok(ClientAdmissionRestoreResult::default()) })
    }
}

#[derive(Default)]
struct Budget {
    reject: bool,
    active: Arc<AtomicBool>,
    charges: Mutex<Vec<ClientBudgetCharge>>,
    settlement_gate: Mutex<Option<oneshot::Receiver<()>>>,
    settlements: AtomicUsize,
    fail_settlement: bool,
}

impl ClientBudgetPort for Budget {
    fn admit(&self, _: ClientApiKeyId) -> BoxFuture<'_, Result<(), GatewayError>> {
        Box::pin(async {
            assert!(self.active.load(Ordering::SeqCst));
            if self.reject {
                Err(
                    GatewayError::new(GatewayErrorKind::RateLimited, "budget exhausted")
                        .with_client_code("key_daily_budget_exceeded"),
                )
            } else {
                Ok(())
            }
        })
    }
    fn settle(&self, charge: ClientBudgetCharge) -> BoxFuture<'_, Result<(), ClientBudgetError>> {
        Box::pin(async move {
            self.settlements.fetch_add(1, Ordering::SeqCst);
            assert!(
                self.active.load(Ordering::SeqCst),
                "settle before releasing the concurrent request slot"
            );
            let gate = self.settlement_gate.lock().unwrap().take();
            if let Some(gate) = gate {
                gate.await.expect("settlement gate");
            }
            assert!(self.active.load(Ordering::SeqCst));
            self.charges.lock().unwrap().push(charge);
            if self.fail_settlement {
                Err(ClientBudgetError)
            } else {
                Ok(())
            }
        })
    }
}

fn service(admissions: Arc<Admissions>, budget: Arc<Budget>) -> DefaultExecutionService {
    DefaultExecutionService::new(
        RuntimeSnapshotHandle::new(start_snapshot()),
        Arc::new(TrackingExecutionStore::default()),
        ProviderRegistry::default(),
        admissions,
        Arc::new(UnusedCircuits),
        Arc::new(UnusedContinuation),
        Arc::new(RecordingClientApiKeyUsage::default()),
    )
    .with_budget(budget)
}

fn request(service: &DefaultExecutionService, transport: ClientTransport) -> StartExecution {
    StartExecution {
        client: service.authenticate("sk_start_test").unwrap(),
        public_model: PublicModelId::new("gpt-start").unwrap(),
        operation: start_operation(),
        metadata: ExecutionRequestMetadata {
            protocol: "openai".to_owned(),
            endpoint: "/v1/responses".to_owned(),
            transport,
            stream: transport != ClientTransport::HttpJson,
            client_ip: None,
            user_agent: None,
            previous_response_id: None,
        },
    }
}

#[test]
fn budget_rejection_releases_client_concurrency_without_creating_a_charge() {
    let admissions = Arc::new(Admissions::default());
    let budget = Arc::new(Budget {
        reject: true,
        active: admissions.active.clone(),
        charges: Mutex::default(),
        ..Default::default()
    });
    let service = service(admissions.clone(), budget.clone());
    for transport in [
        ClientTransport::HttpJson,
        ClientTransport::HttpSse,
        ClientTransport::WebSocket,
    ] {
        let result = block_on(service.start(request(&service, transport)));
        assert!(
            matches!(result, Err(error) if error.client_error_code() == Some("key_daily_budget_exceeded"))
        );
        assert!(!admissions.active.load(Ordering::SeqCst));
        assert!(budget.charges.lock().unwrap().is_empty());
    }
}

#[test]
fn reused_client_uses_updated_limits_for_each_execution() {
    let snapshots = RuntimeSnapshotHandle::new(start_snapshot());
    let admissions = Arc::new(Admissions::default());
    let service = DefaultExecutionService::new(
        snapshots.clone(),
        Arc::new(TrackingExecutionStore::default()),
        ProviderRegistry::default(),
        admissions.clone(),
        Arc::new(UnusedCircuits),
        Arc::new(UnusedContinuation),
        Arc::new(RecordingClientApiKeyUsage::default()),
    );
    let client = service.authenticate("sk_start_test").unwrap();
    let limited = RateLimits {
        max_concurrency: 1,
        requests_per_minute: 1,
    };
    for (revision, limits) in [(2, limited), (3, RateLimits::unlimited())] {
        snapshots.publish(start_snapshot_with_policy(revision, true, limits, false));
        let mut next = request(&service, ClientTransport::WebSocket);
        next.client = client.clone();
        let started = block_on(service.start(next)).expect("new execution");
        block_on(started.session.detach_finalize());
    }
    assert_eq!(
        *admissions.limits.lock().unwrap(),
        vec![limited, RateLimits::unlimited()]
    );
}

#[test]
fn reused_client_cannot_start_after_key_disable_or_snapshot_suspension() {
    for suspend in [false, true] {
        let snapshots = RuntimeSnapshotHandle::new(start_snapshot());
        let admissions = Arc::new(Admissions::default());
        let service = DefaultExecutionService::new(
            snapshots.clone(),
            Arc::new(TrackingExecutionStore::default()),
            ProviderRegistry::default(),
            admissions.clone(),
            Arc::new(UnusedCircuits),
            Arc::new(UnusedContinuation),
            Arc::new(RecordingClientApiKeyUsage::default()),
        );
        let next = request(&service, ClientTransport::WebSocket);
        if suspend {
            snapshots.suspend();
        } else {
            snapshots.publish(start_snapshot_with_policy(
                2,
                false,
                RateLimits::unlimited(),
                false,
            ));
        }
        let result = block_on(service.start(next));
        let expected = if suspend {
            GatewayErrorKind::Internal
        } else {
            GatewayErrorKind::Unauthorized
        };
        assert!(matches!(result, Err(error) if error.kind() == expected));
        assert!(admissions.limits.lock().unwrap().is_empty());
    }
}

#[test]
fn cancellation_and_pre_send_failure_settle_zero_and_release_concurrency_for_all_transports() {
    let admissions = Arc::new(Admissions::default());
    let budget = Arc::new(Budget {
        reject: false,
        active: admissions.active.clone(),
        charges: Mutex::default(),
        ..Default::default()
    });
    let service = service(admissions.clone(), budget.clone());
    for transport in [
        ClientTransport::HttpJson,
        ClientTransport::HttpSse,
        ClientTransport::WebSocket,
    ] {
        for detach in [true, false] {
            let mut started = block_on(service.start(request(&service, transport))).unwrap();
            assert!(admissions.active.load(Ordering::SeqCst));
            if detach {
                block_on(started.session.detach_finalize());
            } else {
                assert!(block_on(started.session.next_event()).is_err());
                block_on(started.session.detach_finalize());
            }
            assert!(!admissions.active.load(Ordering::SeqCst));
        }
    }
    let charges = budget.charges.lock().unwrap();
    assert_eq!(
        charges.len(),
        6,
        "one settlement per request, including detached finalized sessions"
    );
    assert!(
        charges
            .iter()
            .all(|charge| charge.amount_usd == gateway_core::metering::Decimal::ZERO)
    );
}

fn early_failure_service(
    store: Arc<TrackingExecutionStore>,
    admissions: Arc<Admissions>,
    budget: Arc<Budget>,
) -> DefaultExecutionService {
    DefaultExecutionService::new(
        RuntimeSnapshotHandle::new(start_snapshot()),
        store,
        ProviderRegistry::new([Arc::new(LocalFailingProvider) as Arc<dyn Provider>]).unwrap(),
        admissions,
        Arc::new(UnusedCircuits),
        Arc::new(UnusedContinuation),
        Arc::new(RecordingClientApiKeyUsage::default()),
    )
    .with_budget(budget)
}

fn assert_early_failure_recorded(
    store: &TrackingExecutionStore,
    request_id: &ModelRequestId,
    transport: ClientTransport,
) {
    assert_eq!(store.creates.load(Ordering::SeqCst), 1);
    assert_eq!(store.finalizes.load(Ordering::SeqCst), 1);
    assert!(store.attempts.lock().unwrap().is_empty());
    let requests = store.requests.lock().unwrap();
    assert_eq!(requests.len(), 1);
    let request = &requests[0];
    assert_eq!(request.id, *request_id);
    assert_eq!(request.client_api_key_ref.as_str(), "key_start_test");
    assert_eq!(
        request.client_api_key_id.as_ref(),
        Some(&request.client_api_key_ref)
    );
    assert_eq!(request.protocol, "openai");
    assert_eq!(request.endpoint, "/v1/responses");
    assert_eq!(request.client_transport, transport.as_str());
    assert_eq!(request.operation, OperationKind::Generate);
    assert_eq!(
        request.requested_model.as_ref().map(PublicModelId::as_str),
        Some("gpt-start")
    );
    let finalizations = store.finalizations.lock().unwrap();
    assert_eq!(finalizations.len(), 1);
    let finalization = &finalizations[0];
    assert_eq!(finalization.request_id, *request_id);
    assert_eq!(finalization.outcome, ExecutionOutcome::Failed);
    assert_eq!(finalization.send_state, UpstreamSendState::NotSent);
    assert_eq!(finalization.attempt_count, 0);
    assert!(finalization.downstream_committed_at.is_none());
    assert!(finalization.upstream_request_id.is_none());
    assert!(finalization.upstream_response_id.is_none());
    assert!(finalization.upstream_status_code.is_none());
    assert!(finalization.upstream_transport.is_none());
    assert!(finalization.http_version.is_none());
    assert!(finalization.websocket_pool.is_none());
    assert!(finalization.provider_metadata_json.is_none());
    assert_eq!(
        finalization
            .error
            .as_ref()
            .expect("original failure")
            .kind(),
        GatewayErrorKind::Unsupported,
        "detached cancellation must not replace the original provider failure"
    );
    assert_eq!(finalization.usage, Default::default());
    assert!(finalization.cost.total().is_none());
}

fn assert_zero_cleanup_completed(
    admissions: &Admissions,
    budget: &Budget,
    store: &TrackingExecutionStore,
    request_id: &ModelRequestId,
) {
    assert!(!admissions.active.load(Ordering::SeqCst));
    assert_eq!(admissions.releases.load(Ordering::SeqCst), 1);
    assert_eq!(budget.settlements.load(Ordering::SeqCst), 1);
    let charges = budget.charges.lock().unwrap();
    assert_eq!(charges.len(), 1);
    assert_eq!(charges[0].request_id, *request_id);
    assert_eq!(charges[0].key_id.as_str(), "key_start_test");
    assert_eq!(charges[0].amount_usd, Decimal::ZERO);
    assert_eq!(
        charges[0].completed_at,
        store.finalizations.lock().unwrap()[0].completed_at,
        "settlement must retain the original failure completion time"
    );
}

#[test]
fn early_provider_failure_records_request_and_settles_zero_for_all_transports() {
    block_on(async {
        for transport in [
            ClientTransport::HttpJson,
            ClientTransport::HttpSse,
            ClientTransport::WebSocket,
        ] {
            let store = Arc::new(TrackingExecutionStore::default());
            let admissions = Arc::new(Admissions::default());
            let budget = Arc::new(Budget {
                active: admissions.active.clone(),
                ..Default::default()
            });
            let service = early_failure_service(store.clone(), admissions.clone(), budget.clone());
            let mut started = service.start(request(&service, transport)).await.unwrap();
            assert!(admissions.active.load(Ordering::SeqCst));
            assert_eq!(store.creates.load(Ordering::SeqCst), 0);
            let error = if transport == ClientTransport::HttpJson {
                started.session.collect_uncommitted().await.unwrap_err()
            } else {
                started.session.next_event().await.unwrap_err()
            };
            assert!(matches!(
                error,
                EngineError::Provider(error)
                    if error.kind() == ProviderErrorKind::Unsupported
                        && error.send_state() == UpstreamSendState::NotSent
            ));
            assert!(started.session.is_finalized());
            assert_early_failure_recorded(&store, &started.request_id, transport);
            assert_zero_cleanup_completed(&admissions, &budget, &store, &started.request_id);
            assert!(started.session.next_event().await.unwrap().is_none());
            started.session.detach_finalize().await;
            assert_early_failure_recorded(&store, &started.request_id, transport);
            assert_zero_cleanup_completed(&admissions, &budget, &store, &started.request_id);
        }
    });
}

#[test]
fn detached_early_failure_resumes_cancelled_store_write_and_settles_once_for_all_transports() {
    block_on(async {
        for transport in [
            ClientTransport::HttpJson,
            ClientTransport::HttpSse,
            ClientTransport::WebSocket,
        ] {
            for suspend_create in [true, false] {
                let (complete_write, write_gate) = oneshot::channel();
                let store = Arc::new(TrackingExecutionStore::default());
                if suspend_create {
                    *store.create_gate.lock().unwrap() = Some(write_gate);
                } else {
                    *store.finalize_gate.lock().unwrap() = Some(write_gate);
                }
                let (complete_release, release_gate) = oneshot::channel();
                let admissions = Arc::new(Admissions {
                    release_gate: Mutex::new(Some(release_gate)),
                    ..Default::default()
                });
                let (complete_settlement, settlement_gate) = oneshot::channel();
                let budget = Arc::new(Budget {
                    active: admissions.active.clone(),
                    settlement_gate: Mutex::new(Some(settlement_gate)),
                    ..Default::default()
                });
                let service =
                    early_failure_service(store.clone(), admissions.clone(), budget.clone());
                let mut started = service.start(request(&service, transport)).await.unwrap();
                let mut next = started.session.next_event();
                assert!(futures::poll!(next.as_mut()).is_pending());
                drop(next);
                assert!(!started.session.is_finalized());
                assert_eq!(store.creates.load(Ordering::SeqCst), 1);
                assert_eq!(
                    store.finalizes.load(Ordering::SeqCst),
                    usize::from(!suspend_create)
                );
                assert_eq!(
                    store.requests.lock().unwrap().len(),
                    usize::from(!suspend_create)
                );
                assert!(store.finalizations.lock().unwrap().is_empty());
                assert_eq!(budget.settlements.load(Ordering::SeqCst), 0);
                assert_eq!(admissions.releases.load(Ordering::SeqCst), 0);
                assert!(admissions.active.load(Ordering::SeqCst));

                let mut detached = started.session.detach_finalize();
                assert!(futures::poll!(detached.as_mut()).is_pending());
                assert_eq!(store.creates.load(Ordering::SeqCst), 1);
                assert_eq!(
                    store.finalizes.load(Ordering::SeqCst),
                    usize::from(!suspend_create)
                );
                // receiver 已从 Store 替身取走；只有延续原 future 才能继续接收此信号。
                complete_write
                    .send(())
                    .expect("detached cleanup retains the original store write");
                assert!(futures::poll!(detached.as_mut()).is_pending());
                assert_early_failure_recorded(&store, &started.request_id, transport);
                assert_eq!(budget.settlements.load(Ordering::SeqCst), 1);
                assert!(budget.charges.lock().unwrap().is_empty());
                assert_eq!(admissions.releases.load(Ordering::SeqCst), 0);
                assert!(admissions.active.load(Ordering::SeqCst));

                complete_settlement.send(()).expect("original settlement");
                assert!(futures::poll!(detached.as_mut()).is_pending());
                assert_eq!(budget.settlements.load(Ordering::SeqCst), 1);
                {
                    let charges = budget.charges.lock().unwrap();
                    assert_eq!(charges.len(), 1);
                    assert_eq!(charges[0].amount_usd, Decimal::ZERO);
                }
                assert_eq!(admissions.releases.load(Ordering::SeqCst), 1);
                assert!(admissions.active.load(Ordering::SeqCst));
                complete_release
                    .send(())
                    .expect("original admission release");
                detached.await;
                assert_early_failure_recorded(&store, &started.request_id, transport);
                assert_zero_cleanup_completed(&admissions, &budget, &store, &started.request_id);
            }
        }
    });
}

#[derive(Default)]
struct ChargedProvider {
    policies: Mutex<Vec<bool>>,
    fail: bool,
}

fn known_charge() -> Decimal {
    "1.25".parse().unwrap()
}

#[async_trait]
impl Provider for ChargedProvider {
    fn name(&self) -> &'static str {
        "openai"
    }

    fn catalog_generation(&self) -> ProviderCatalogGeneration {
        ProviderCatalogGeneration::default()
    }

    async fn query_model_capabilities(
        &self,
    ) -> Result<Vec<ProviderModelCapabilities>, ProviderError> {
        Ok(Vec::new())
    }

    async fn execute(
        &self,
        request: ProviderRequest,
        context: AttemptContext,
    ) -> Result<ProviderStream, ProviderError> {
        self.policies.lock().unwrap().push(context.disable_fast());
        let candidate = request.candidate();
        let metadata = ProviderCallMetadata::new(
            candidate.provider().clone(),
            candidate.upstream_model().unwrap().clone(),
            ProviderAccountId::new("acct_start").unwrap(),
            UpstreamTransport::new("websocket").unwrap(),
        );
        let response = ResponseMeta::new("resp_cleanup", "gpt-start");
        let terminal = if self.fail {
            Err(ProviderError::new(
                ProviderErrorKind::InvalidRequest,
                UpstreamSendState::Sent,
            ))
        } else {
            Ok(GatewayEvent::Completed(response.clone()).into())
        };
        let events: Vec<Result<ProviderEvent, ProviderError>> = vec![
            Ok(GatewayEvent::Started(response).into()),
            Ok(GatewayEvent::ProviderCost(
                ProviderReportedCost::from_usd_ticks(known_charge().scaled()).unwrap(),
            )
            .into()),
            terminal,
        ];
        Ok(ProviderStream::new(
            metadata,
            futures::stream::iter(events),
            (),
        ))
    }
}

fn charged_service(
    admissions: Arc<Admissions>,
    budget: Option<Arc<Budget>>,
    fail: bool,
) -> DefaultExecutionService {
    let service = DefaultExecutionService::new(
        RuntimeSnapshotHandle::new(start_snapshot()),
        Arc::new(TrackingExecutionStore::default()),
        ProviderRegistry::new([Arc::new(ChargedProvider {
            fail,
            ..Default::default()
        }) as Arc<dyn Provider>])
        .unwrap(),
        admissions,
        Arc::new(UnusedCircuits),
        Arc::new(UnusedContinuation),
        Arc::new(RecordingClientApiKeyUsage::default()),
    );
    match budget {
        Some(budget) => service.with_budget(budget),
        None => service,
    }
}

async fn consume_charged_prefix(session: &mut dyn ExecutionSession) {
    let first = session.next_event().await.unwrap().unwrap();
    assert_eq!(
        first.commit_requirement(),
        CommitRequirement::CommitBeforeDelivery
    );
    session.commit_downstream(None).await.unwrap();
    let cost = session.next_event().await.unwrap().unwrap();
    assert_eq!(
        cost.commit_requirement(),
        CommitRequirement::AlreadyCommitted
    );
}

fn assert_cleanup_completed(admissions: &Admissions, budget: &Budget, request_id: &ModelRequestId) {
    assert!(!admissions.active.load(Ordering::SeqCst));
    assert_eq!(admissions.releases.load(Ordering::SeqCst), 1);
    assert_eq!(budget.settlements.load(Ordering::SeqCst), 1);
    let charges = budget.charges.lock().unwrap();
    assert_eq!(charges.len(), 1);
    assert_eq!(charges[0].request_id, *request_id);
    assert_eq!(charges[0].key_id.as_str(), "key_start_test");
    assert_eq!(charges[0].amount_usd, known_charge());
}

#[test]
fn normal_completion_settles_known_charge_once_for_all_transports() {
    block_on(async {
        for transport in [
            ClientTransport::HttpJson,
            ClientTransport::HttpSse,
            ClientTransport::WebSocket,
        ] {
            let admissions = Arc::new(Admissions::default());
            let budget = Arc::new(Budget {
                active: admissions.active.clone(),
                ..Default::default()
            });
            let service = charged_service(admissions.clone(), Some(budget.clone()), false);
            let mut started = service.start(request(&service, transport)).await.unwrap();
            if transport == ClientTransport::HttpJson {
                assert_eq!(
                    started.session.collect_uncommitted().await.unwrap().len(),
                    3
                );
                assert!(!started.session.is_finalized());
                started.session.commit_downstream(Some(200)).await.unwrap();
            } else {
                consume_charged_prefix(started.session.as_mut()).await;
                assert!(started.session.next_event().await.unwrap().is_some());
            }
            assert!(started.session.next_event().await.unwrap().is_none());
            assert!(started.session.is_finalized());
            started.session.detach_finalize().await;
            assert_cleanup_completed(&admissions, &budget, &started.request_id);
        }
    });
}

#[test]
fn settlement_survives_cancelled_event_wait_and_resumes_without_restart() {
    block_on(async {
        for fail in [false, true] {
            let admissions = Arc::new(Admissions::default());
            let (release, gate) = oneshot::channel();
            let budget = Arc::new(Budget {
                active: admissions.active.clone(),
                settlement_gate: Mutex::new(Some(gate)),
                ..Default::default()
            });
            let service = charged_service(admissions.clone(), Some(budget.clone()), fail);
            let mut started = service
                .start(request(&service, ClientTransport::WebSocket))
                .await
                .unwrap();
            consume_charged_prefix(started.session.as_mut()).await;
            if !fail {
                assert!(started.session.next_event().await.unwrap().is_some());
            }
            for _ in 0..2 {
                let mut terminal = started.session.next_event();
                assert!(futures::poll!(terminal.as_mut()).is_pending());
                drop(terminal);
                assert!(!started.session.is_finalized());
                assert_eq!(budget.settlements.load(Ordering::SeqCst), 1);
                assert!(budget.charges.lock().unwrap().is_empty());
                assert_eq!(admissions.releases.load(Ordering::SeqCst), 0);
            }
            release
                .send(())
                .expect("the original settlement is still alive");
            assert!(started.session.next_event().await.unwrap().is_none());
            assert!(started.session.is_finalized());
            started.session.detach_finalize().await;
            assert_cleanup_completed(&admissions, &budget, &started.request_id);
        }
    });
}

#[test]
fn detached_cleanup_settles_cancelled_execution_and_resumes_existing_settlement() {
    block_on(async {
        for cancel_event_wait in [false, true] {
            let admissions = Arc::new(Admissions::default());
            let (release, gate) = oneshot::channel();
            let budget = Arc::new(Budget {
                active: admissions.active.clone(),
                settlement_gate: Mutex::new(Some(gate)),
                ..Default::default()
            });
            let service = charged_service(admissions.clone(), Some(budget.clone()), false);
            let mut started = service
                .start(request(&service, ClientTransport::WebSocket))
                .await
                .unwrap();
            consume_charged_prefix(started.session.as_mut()).await;
            if cancel_event_wait {
                started.session.cancel();
                let mut terminal = started.session.next_event();
                assert!(futures::poll!(terminal.as_mut()).is_pending());
                drop(terminal);
                assert!(!started.session.is_finalized());
            }
            let mut detached = started.session.detach_finalize();
            assert!(futures::poll!(detached.as_mut()).is_pending());
            assert_eq!(budget.settlements.load(Ordering::SeqCst), 1);
            assert_eq!(admissions.releases.load(Ordering::SeqCst), 0);
            release
                .send(())
                .expect("detached cleanup retains settlement");
            detached.await;
            assert_cleanup_completed(&admissions, &budget, &started.request_id);
        }
    });
}

#[test]
fn cancelled_release_resumes_without_repeating_settlement_with_or_without_budget() {
    block_on(async {
        for with_budget in [false, true] {
            let (release, gate) = oneshot::channel();
            let admissions = Arc::new(Admissions {
                release_gate: Mutex::new(Some(gate)),
                ..Default::default()
            });
            let budget = Arc::new(Budget {
                active: admissions.active.clone(),
                ..Default::default()
            });
            let service = charged_service(
                admissions.clone(),
                with_budget.then(|| budget.clone()),
                false,
            );
            let mut started = service
                .start(request(&service, ClientTransport::HttpSse))
                .await
                .unwrap();
            consume_charged_prefix(started.session.as_mut()).await;
            assert!(started.session.next_event().await.unwrap().is_some());
            let mut terminal = started.session.next_event();
            assert!(futures::poll!(terminal.as_mut()).is_pending());
            drop(terminal);
            assert!(!started.session.is_finalized());
            assert!(admissions.active.load(Ordering::SeqCst));
            assert_eq!(
                budget.settlements.load(Ordering::SeqCst),
                usize::from(with_budget)
            );
            let mut detached = started.session.detach_finalize();
            assert!(futures::poll!(detached.as_mut()).is_pending());
            assert_eq!(admissions.releases.load(Ordering::SeqCst), 1);
            release
                .send(())
                .expect("the original release is still alive");
            detached.await;
            if with_budget {
                assert_cleanup_completed(&admissions, &budget, &started.request_id);
            } else {
                assert!(!admissions.active.load(Ordering::SeqCst));
                assert_eq!(admissions.releases.load(Ordering::SeqCst), 1);
                assert!(budget.charges.lock().unwrap().is_empty());
            }
        }
    });
}

#[test]
fn cancelled_buffered_commit_keeps_settlement_for_detached_cleanup() {
    block_on(async {
        let admissions = Arc::new(Admissions::default());
        let (release, gate) = oneshot::channel();
        let budget = Arc::new(Budget {
            active: admissions.active.clone(),
            settlement_gate: Mutex::new(Some(gate)),
            ..Default::default()
        });
        let service = charged_service(admissions.clone(), Some(budget.clone()), false);
        let mut started = service
            .start(request(&service, ClientTransport::HttpJson))
            .await
            .unwrap();
        assert_eq!(
            started.session.collect_uncommitted().await.unwrap().len(),
            3
        );
        let mut commit = started.session.commit_downstream(Some(200));
        assert!(futures::poll!(commit.as_mut()).is_pending());
        drop(commit);
        assert!(!started.session.is_finalized());
        release
            .send(())
            .expect("buffered commit retains settlement");
        started.session.detach_finalize().await;
        assert_cleanup_completed(&admissions, &budget, &started.request_id);
    });
}

#[test]
fn settlement_failure_keeps_provider_error_and_releases_concurrency_once() {
    block_on(async {
        for fail_settlement in [false, true] {
            let admissions = Arc::new(Admissions::default());
            let budget = Arc::new(Budget {
                active: admissions.active.clone(),
                fail_settlement,
                ..Default::default()
            });
            let service = charged_service(admissions.clone(), Some(budget.clone()), true);
            let mut started = service
                .start(request(&service, ClientTransport::WebSocket))
                .await
                .unwrap();
            consume_charged_prefix(started.session.as_mut()).await;
            assert!(matches!(
                started.session.next_event().await,
                Err(EngineError::Provider(error)) if error.kind() == ProviderErrorKind::InvalidRequest
            ));
            assert!(started.session.is_finalized());
            started.session.detach_finalize().await;
            // Store 端口已接管精确费用后，结算错误不能改写 Provider 错误或触发第二次结算。
            assert_cleanup_completed(&admissions, &budget, &started.request_id);
        }
    });
}

use async_trait::async_trait;
use bytes::Bytes;
use futures::{channel::oneshot, executor::block_on, future::BoxFuture};
use gateway_core::account::{AccountSelectionPolicy, ProviderAccountId, RotationStrategy};
use gateway_core::engine::admission::{
    ClientAdmissionDecision, ClientAdmissionError, ClientAdmissionPort, ClientAdmissionRecovery,
    ClientAdmissionRequest, ClientAdmissionRestoreResult,
};
use gateway_core::engine::continuation::{
    NativeContinuationPin, NativeContinuationPort, NativeContinuationStoreError, PreviousResponseId,
};
use gateway_core::engine::execution::{
    ClientApiKeyUsageSink, ClientKeyVerifier, ClientTransport, DefaultExecutionService,
    ExecutionRequestMetadata, ExecutionService, ExecutionSession, ProviderCircuitDecision,
    ProviderCircuitError, ProviderCircuitPort, StartExecution, StartProviderExecution,
    provider_failure_affects_circuit,
};
use gateway_core::engine::probe::{AccountProbe, AccountProbeErrorSource, AccountProbeRequest};
use gateway_core::engine::provider::{
    Provider, ProviderCallMetadata, ProviderRegistry, ProviderRequest, ProviderRequestObservation,
    ProviderStream,
};
use gateway_core::engine::{
    AttemptContext, AttemptRecord, CommitRequirement, EngineError, ExecutionOutcome,
    ExecutionStore, IntermediateFailure, ModelRequestFinalization, ModelRequestId, NewModelRequest,
    ProbeFailure, RecoveryReport,
};
use gateway_core::error::{
    ClientVisibleUpstreamResponse, GatewayErrorKind, ProviderError, ProviderErrorKind, StoreError,
    StoreErrorKind,
};
use gateway_core::event::{GatewayEvent, ProviderEvent, ResponseMeta};
use gateway_core::metering::{Decimal, ProviderReportedCost};
use gateway_core::operation::{
    GenerateRequest, ImageRequest, ImageRequestKind, Operation, OperationKind, ProtocolPayload,
    RawJsonPayload,
};
use gateway_core::policy::{ClientApiKeyId, ClientPolicy, PlaintextClientApiKey, RateLimits};
use gateway_core::routing::{
    ClientRoutingScope, ConfigRevision, FrozenAccountScope, ModelCapabilities,
    ProviderCatalogGeneration, ProviderKind, ProviderModel, ProviderModelCapabilities,
    PublicModelId, RuntimeAccount, RuntimeAccountDirectory, RuntimeSnapshot, UpstreamModelId,
};
use gateway_core::runtime::RuntimeSnapshotHandle;
use gateway_core::upstream::{UpstreamSendState, UpstreamTransport};
use serde_json::json;

#[test]
fn only_provider_attributable_failures_should_affect_circuit() {
    assert!(provider_failure_affects_circuit(ProviderErrorKind::Timeout));
    assert!(provider_failure_affects_circuit(
        ProviderErrorKind::Transport
    ));
    assert!(!provider_failure_affects_circuit(
        ProviderErrorKind::RateLimited
    ));
    assert!(!provider_failure_affects_circuit(
        ProviderErrorKind::InvalidRequest
    ));
    assert!(!provider_failure_affects_circuit(
        ProviderErrorKind::ContinuationRecoveryRequired
    ));
    assert!(!provider_failure_affects_circuit(
        ProviderErrorKind::UpstreamCapacityUnavailable
    ));
}

#[test]
fn account_probe_should_not_write_to_the_persistent_execution_store() {
    let store = Arc::new(TrackingExecutionStore::default());
    let service = DefaultExecutionService::new(
        RuntimeSnapshotHandle::new(probe_snapshot()),
        store.clone(),
        ProviderRegistry::default(),
        Arc::new(UnusedAdmissions),
        Arc::new(UnusedCircuits),
        Arc::new(UnusedContinuation),
        Arc::new(RecordingClientApiKeyUsage::default()),
    );

    let error = block_on(service.probe(AccountProbeRequest {
        account_id: ProviderAccountId::new("acct_probe").expect("account ID"),
        provider_kind: ProviderKind::new("openai").expect("provider kind"),
        upstream_model: UpstreamModelId::new("gpt-probe").expect("model ID"),
        operation: probe_operation(),
    }))
    .expect_err("empty Provider registry should stop the probe after it starts");

    assert_eq!(error.kind(), GatewayErrorKind::NoAvailableProvider);
    assert_eq!(error.source(), AccountProbeErrorSource::Gateway);
    assert_eq!(error.send_state(), None);
    assert!(!store.touched.load(Ordering::SeqCst));
}

#[test]
fn probe_failures_should_be_observable_without_a_model_request_row() {
    let store = Arc::new(TrackingExecutionStore::default());
    let providers = ProviderRegistry::new([Arc::new(FailingProvider) as Arc<dyn Provider>])
        .expect("provider registry");
    let service = DefaultExecutionService::new(
        RuntimeSnapshotHandle::new(probe_snapshot()),
        store.clone(),
        providers,
        Arc::new(UnusedAdmissions),
        Arc::new(UnusedCircuits),
        Arc::new(UnusedContinuation),
        Arc::new(RecordingClientApiKeyUsage::default()),
    );

    let error = block_on(service.probe(AccountProbeRequest {
        account_id: ProviderAccountId::new("acct_probe").expect("account ID"),
        provider_kind: ProviderKind::new("openai").expect("provider kind"),
        upstream_model: UpstreamModelId::new("gpt-probe").expect("model ID"),
        operation: probe_operation(),
    }))
    .expect_err("the provider rejects every probe");

    assert_eq!(error.kind(), GatewayErrorKind::UpstreamUnavailable);
    assert_eq!(error.source(), AccountProbeErrorSource::Upstream);
    assert_eq!(error.send_state(), Some(UpstreamSendState::NotSent));
    let upstream = error
        .upstream_response()
        .expect("probe must preserve its request-local upstream response");
    assert_eq!(upstream.status(), 502);
    assert_eq!(
        upstream.content_type(),
        Some(b"application/json".as_slice())
    );
    assert_eq!(
        upstream.body(),
        &Bytes::from_static(br#"{"error":{"message":"source upstream failure"}}"#),
    );
    assert!(!format!("{error:?}").contains("source upstream failure"));
    assert!(!store.touched.load(Ordering::SeqCst));
    assert_eq!(store.probe_failures(), vec!["transport".to_owned()]);
}

#[test]
fn provider_local_probe_failure_should_remain_distinct_from_upstream() {
    let store = Arc::new(TrackingExecutionStore::default());
    let providers = ProviderRegistry::new([Arc::new(LocalFailingProvider) as Arc<dyn Provider>])
        .expect("provider registry");
    let service = DefaultExecutionService::new(
        RuntimeSnapshotHandle::new(probe_snapshot()),
        store,
        providers,
        Arc::new(UnusedAdmissions),
        Arc::new(UnusedCircuits),
        Arc::new(UnusedContinuation),
        Arc::new(RecordingClientApiKeyUsage::default()),
    );

    let error = block_on(service.probe(AccountProbeRequest {
        account_id: ProviderAccountId::new("acct_probe").expect("account ID"),
        provider_kind: ProviderKind::new("openai").expect("provider kind"),
        upstream_model: UpstreamModelId::new("gpt-probe").expect("model ID"),
        operation: probe_operation(),
    }))
    .expect_err("the Provider rejects the probe before sending it");

    assert_eq!(error.source(), AccountProbeErrorSource::Provider);
    assert_eq!(error.send_state(), Some(UpstreamSendState::NotSent));
    assert!(error.upstream_response().is_none());
}

#[test]
fn probe_observation_store_failure_preserves_the_provider_error() {
    let store = Arc::new(TrackingExecutionStore {
        fail_probe_observation: AtomicBool::new(true),
        ..TrackingExecutionStore::default()
    });
    let providers = ProviderRegistry::new([Arc::new(FailingProvider) as Arc<dyn Provider>])
        .expect("provider registry");
    let service = DefaultExecutionService::new(
        RuntimeSnapshotHandle::new(probe_snapshot()),
        store.clone(),
        providers,
        Arc::new(UnusedAdmissions),
        Arc::new(UnusedCircuits),
        Arc::new(UnusedContinuation),
        Arc::new(RecordingClientApiKeyUsage::default()),
    );

    let error = block_on(service.probe(AccountProbeRequest {
        account_id: ProviderAccountId::new("acct_probe").expect("account ID"),
        provider_kind: ProviderKind::new("openai").expect("provider kind"),
        upstream_model: UpstreamModelId::new("gpt-probe").expect("model ID"),
        operation: probe_operation(),
    }))
    .expect_err("the provider error must survive observation failure");

    assert_eq!(error.kind(), GatewayErrorKind::UpstreamUnavailable);
    assert!(!store.touched.load(Ordering::SeqCst));
    assert!(store.probe_failures().is_empty());
}

struct FailingProvider;

struct NativeCatalogProvider {
    fail: bool,
}

#[async_trait]
impl Provider for NativeCatalogProvider {
    fn name(&self) -> &'static str {
        "openai"
    }
    fn catalog_generation(&self) -> ProviderCatalogGeneration {
        ProviderCatalogGeneration::default()
    }
    async fn query_model_capabilities(
        &self,
    ) -> Result<Vec<ProviderModelCapabilities>, ProviderError> {
        Ok(Vec::new())
    }
    async fn query_client_model_catalog(
        &self,
        scope: &FrozenAccountScope,
        protocol: &str,
        version: &str,
    ) -> Result<
        Option<Vec<gateway_core::routing::ProviderModelDescriptor>>,
        gateway_core::routing::ProviderCatalogUnavailable,
    > {
        assert!(scope.allows(&ProviderAccountId::new("acct_start").expect("account")));
        assert!(!scope.allows(&ProviderAccountId::new("acct_other").expect("account")));
        assert_eq!((protocol, version), ("codex", "0.154.0"));
        if self.fail {
            return Err(gateway_core::routing::ProviderCatalogUnavailable);
        }
        Ok(Some(
            ["gpt-new", "gpt-start"]
                .into_iter()
                .map(|id| gateway_core::routing::ProviderModelDescriptor {
                    model: UpstreamModelId::new(id).expect("id"),
                    content: gateway_core::routing::ProviderModelContent::Native(
                        RawJsonPayload::new(
                            "codex",
                            serde_json::to_vec(&json!({"slug":id,"future":null}))
                                .expect("JSON")
                                .into(),
                        )
                        .expect("payload"),
                    ),
                })
                .collect(),
        ))
    }
    async fn execute(
        &self,
        _: ProviderRequest,
        _: AttemptContext,
    ) -> Result<ProviderStream, ProviderError> {
        panic!("catalog must not use metered execution")
    }
}

#[test]
fn client_catalog_forwards_scope_maps_whole_objects_and_omits_unroutable_models() {
    for fail in [false, true] {
        let snapshot = start_snapshot().with_model_mappings(BTreeMap::from([
            ("alias".to_owned(), "gpt-start".to_owned()),
            ("missing-alias".to_owned(), "unavailable-model".to_owned()),
        ]));
        let service = DefaultExecutionService::new(
            RuntimeSnapshotHandle::new(snapshot),
            Arc::new(TrackingExecutionStore::default()),
            ProviderRegistry::new([Arc::new(NativeCatalogProvider { fail }) as Arc<dyn Provider>])
                .expect("registry"),
            Arc::new(UnusedAdmissions),
            Arc::new(UnusedCircuits),
            Arc::new(UnusedContinuation),
            Arc::new(RecordingClientApiKeyUsage::default()),
        );
        let client = service.authenticate("sk_start_test").expect("authenticate");
        let result = block_on(service.client_model_catalog(&client, "codex", "0.154.0"));
        if fail {
            assert!(
                result.is_err(),
                "native failure cannot use the global synthesized catalog"
            );
            continue;
        }
        let models = result.expect("models");
        let pairs = models
            .iter()
            .map(|entry| {
                let gateway_core::routing::PublicModelDescriptor::Native { model, payload } = entry
                else {
                    panic!("native required")
                };
                (model.as_str(), payload.body())
            })
            .collect::<Vec<_>>();
        assert_eq!(
            pairs.iter().map(|(id, _)| *id).collect::<Vec<_>>(),
            ["gpt-start", "alias"]
        );
        assert_eq!(pairs[0].1, pairs[1].1, "Core preserves alias payload bytes");
    }
}

#[async_trait]
impl Provider for FailingProvider {
    fn name(&self) -> &'static str {
        "openai"
    }

    fn catalog_generation(&self) -> ProviderCatalogGeneration {
        ProviderCatalogGeneration::default()
    }

    async fn query_model_capabilities(
        &self,
    ) -> Result<Vec<ProviderModelCapabilities>, ProviderError> {
        Ok(Vec::new())
    }

    async fn execute(
        &self,
        _request: ProviderRequest,
        _context: AttemptContext,
    ) -> Result<ProviderStream, ProviderError> {
        Err(
            ProviderError::new(ProviderErrorKind::Transport, UpstreamSendState::NotSent)
                .with_client_visible_upstream_response(ClientVisibleUpstreamResponse::new(
                    502,
                    Some(b"application/json".to_vec()),
                    Bytes::from_static(br#"{"error":{"message":"source upstream failure"}}"#),
                )),
        )
    }
}

struct LocalFailingProvider;

#[async_trait]
impl Provider for LocalFailingProvider {
    fn name(&self) -> &'static str {
        "openai"
    }

    fn catalog_generation(&self) -> ProviderCatalogGeneration {
        ProviderCatalogGeneration::default()
    }

    async fn query_model_capabilities(
        &self,
    ) -> Result<Vec<ProviderModelCapabilities>, ProviderError> {
        Ok(Vec::new())
    }

    async fn execute(
        &self,
        _: ProviderRequest,
        _: AttemptContext,
    ) -> Result<ProviderStream, ProviderError> {
        Err(ProviderError::new(
            ProviderErrorKind::Unsupported,
            UpstreamSendState::NotSent,
        ))
    }
}

struct ColdFailingProvider {
    requested_model: Option<PublicModelId>,
}

#[async_trait]
impl Provider for ColdFailingProvider {
    fn name(&self) -> &'static str {
        "openai"
    }

    fn catalog_generation(&self) -> ProviderCatalogGeneration {
        ProviderCatalogGeneration::default()
    }

    fn request_observation(&self, _: &Operation, _: &ClientApiKeyId) -> ProviderRequestObservation {
        ProviderRequestObservation {
            requested_model: self.requested_model.clone(),
            ..Default::default()
        }
    }

    async fn query_model_capabilities(
        &self,
    ) -> Result<Vec<ProviderModelCapabilities>, ProviderError> {
        Ok(Vec::new())
    }

    async fn execute(
        &self,
        request: ProviderRequest,
        _: AttemptContext,
    ) -> Result<ProviderStream, ProviderError> {
        let metadata = ProviderCallMetadata::for_provider_endpoint(
            request.candidate().provider().clone(),
            ProviderAccountId::new("acct_usage").expect("account"),
            UpstreamTransport::new("http_json").expect("transport"),
        );
        Ok(ProviderStream::new(
            metadata,
            futures::stream::once(async {
                Err(ProviderError::new(
                    ProviderErrorKind::Transport,
                    UpstreamSendState::NotSent,
                ))
            }),
            (),
        ))
    }
}

#[test]
fn successful_authentication_should_record_client_key_usage() {
    let usage = Arc::new(RecordingClientApiKeyUsage::default());
    let service = DefaultExecutionService::new(
        RuntimeSnapshotHandle::new(client_snapshot()),
        Arc::new(TrackingExecutionStore::default()),
        ProviderRegistry::default(),
        Arc::new(UnusedAdmissions),
        Arc::new(UnusedCircuits),
        Arc::new(UnusedContinuation),
        usage.clone(),
    );

    service
        .authenticate("sk_usage_test")
        .expect("successful authentication");

    assert_eq!(usage.recorded(), vec!["key_usage_test".to_owned()]);
}

#[test]
fn client_key_verification_should_not_record_client_key_usage() {
    let usage = Arc::new(RecordingClientApiKeyUsage::default());
    let service = DefaultExecutionService::new(
        RuntimeSnapshotHandle::new(client_snapshot()),
        Arc::new(TrackingExecutionStore::default()),
        ProviderRegistry::default(),
        Arc::new(UnusedAdmissions),
        Arc::new(UnusedCircuits),
        Arc::new(UnusedContinuation),
        usage.clone(),
    );

    let key_id = service
        .verify_client_key("sk_usage_test")
        .expect("client key verification");

    assert_eq!(key_id.as_str(), "key_usage_test");
    assert!(usage.recorded().is_empty());
}

#[test]
fn provider_endpoint_should_persist_its_real_v1_endpoint() {
    assert_provider_endpoint_observation(None);
    assert_provider_endpoint_observation(Some("gpt-image-2"));
}

fn assert_provider_endpoint_observation(model: Option<&str>) {
    let store = Arc::new(TrackingExecutionStore::default());
    let service = DefaultExecutionService::new(
        RuntimeSnapshotHandle::new(client_snapshot()),
        store.clone(),
        ProviderRegistry::new([Arc::new(ColdFailingProvider {
            requested_model: model.map(|model| PublicModelId::new(model).expect("model")),
        }) as Arc<dyn Provider>])
        .expect("provider registry"),
        Arc::new(UnusedAdmissions),
        Arc::new(UnusedCircuits),
        Arc::new(UnusedContinuation),
        Arc::new(RecordingClientApiKeyUsage::default()),
    );
    let client = service
        .authenticate("sk_usage_test")
        .expect("authenticated client");
    let operation = Operation::GenerateImage(ImageRequest::from_raw_json(
        ImageRequestKind::Generation,
        RawJsonPayload::new(
            "openai",
            Bytes::from_static(br#"{"model":"gpt-image-2","prompt":"hello"}"#),
        )
        .expect("image payload"),
    ));

    let mut started = block_on(service.start_provider_endpoint(StartProviderExecution {
        client,
        provider: ProviderKind::new("openai").expect("provider"),
        operation,
        metadata: ExecutionRequestMetadata {
            protocol: "openai".to_owned(),
            endpoint: "/v1/images/generations".to_owned(),
            transport: ClientTransport::HttpJson,
            stream: false,
            client_ip: None,
            user_agent: None,
            previous_response_id: None,
        },
    }))
    .expect("provider endpoint request should start without a text catalog entry");

    let error = block_on(started.session.collect_uncommitted())
        .expect_err("the cold provider stops execution after persistence");
    assert_eq!(
        store.model_request_endpoints(),
        vec!["/v1/images/generations".to_owned()],
        "execution error: {error:?}"
    );
    assert_eq!(store.requested_models(), vec![model.map(str::to_owned)]);
    let upstream_models = store.upstream_models();
    assert!(!upstream_models.is_empty());
    assert!(upstream_models.iter().all(Option::is_none));
}

#[test]
fn circuit_store_failure_should_fail_open_during_request_start() {
    let service = DefaultExecutionService::new(
        RuntimeSnapshotHandle::new(start_snapshot()),
        Arc::new(TrackingExecutionStore::default()),
        ProviderRegistry::default(),
        Arc::new(UnusedAdmissions),
        Arc::new(FailingDecisionCircuits),
        Arc::new(UnusedContinuation),
        Arc::new(RecordingClientApiKeyUsage::default()),
    );
    let client = service
        .authenticate("sk_start_test")
        .expect("authenticated client");

    let started = block_on(service.start(StartExecution {
        client,
        public_model: PublicModelId::new("gpt-start").expect("public model"),
        operation: start_operation(),
        metadata: ExecutionRequestMetadata {
            protocol: "openai".to_owned(),
            endpoint: "/v1/responses".to_owned(),
            transport: ClientTransport::HttpJson,
            stream: false,
            client_ip: None,
            user_agent: None,
            previous_response_id: None,
        },
    }))
    .expect("recoverable circuit state must not reject the request");

    assert!(!started.session.is_finalized());
}

#[test]
fn slow_circuit_store_should_time_out_and_fail_open_during_request_start() {
    let service = DefaultExecutionService::new(
        RuntimeSnapshotHandle::new(start_snapshot()),
        Arc::new(TrackingExecutionStore::default()),
        ProviderRegistry::default(),
        Arc::new(UnusedAdmissions),
        Arc::new(PendingDecisionCircuits),
        Arc::new(UnusedContinuation),
        Arc::new(RecordingClientApiKeyUsage::default()),
    );
    let client = service
        .authenticate("sk_start_test")
        .expect("authenticated client");
    let started_at = Instant::now();

    let started = block_on(service.start(StartExecution {
        client,
        public_model: PublicModelId::new("gpt-start").expect("public model"),
        operation: start_operation(),
        metadata: ExecutionRequestMetadata {
            protocol: "openai".to_owned(),
            endpoint: "/v1/responses".to_owned(),
            transport: ClientTransport::HttpJson,
            stream: false,
            client_ip: None,
            user_agent: None,
            previous_response_id: None,
        },
    }))
    .expect("slow recoverable circuit state must not reject the request");

    assert!(started_at.elapsed() < Duration::from_secs(2));
    assert!(!started.session.is_finalized());
}

#[test]
fn known_catalog_should_reject_a_model_that_the_provider_did_not_publish() {
    let service = DefaultExecutionService::new(
        RuntimeSnapshotHandle::new(start_snapshot()),
        Arc::new(TrackingExecutionStore::default()),
        ProviderRegistry::default(),
        Arc::new(UnusedAdmissions),
        Arc::new(UnusedCircuits),
        Arc::new(UnusedContinuation),
        Arc::new(RecordingClientApiKeyUsage::default()),
    );
    let client = service
        .authenticate("sk_start_test")
        .expect("authenticated client");
    let model = "gpt-future-not-in-catalog";

    let result = block_on(service.start(StartExecution {
        client,
        public_model: PublicModelId::from_client_wire(model).expect("client model"),
        operation: start_operation_for_model(model),
        metadata: ExecutionRequestMetadata {
            protocol: "openai".to_owned(),
            endpoint: "/v1/responses".to_owned(),
            transport: ClientTransport::HttpJson,
            stream: false,
            client_ip: None,
            user_agent: None,
            previous_response_id: None,
        },
    }));
    let Err(error) = result else {
        panic!("a known provider catalog must be authoritative for model availability");
    };

    assert_eq!(
        (error.kind(), error.client_message()),
        (
            GatewayErrorKind::ModelNotFound,
            "the requested model was not found in the provider catalogs available to this API key; check the model name",
        )
    );
}

#[test]
fn continuation_owned_by_another_client_api_key_should_fail_closed() {
    let service = DefaultExecutionService::new(
        RuntimeSnapshotHandle::new(start_snapshot()),
        Arc::new(TrackingExecutionStore::default()),
        ProviderRegistry::default(),
        Arc::new(UnusedAdmissions),
        Arc::new(UnusedCircuits),
        Arc::new(RejectedContinuation::OwnershipMismatch),
        Arc::new(RecordingClientApiKeyUsage::default()),
    );
    let client = service
        .authenticate("sk_start_test")
        .expect("authenticated client");

    let result = block_on(service.start(StartExecution {
        client,
        public_model: PublicModelId::new("gpt-start").expect("public model"),
        operation: start_operation(),
        metadata: execution_metadata_with_continuation(),
    }));
    let Err(error) = result else {
        panic!("cross-client continuation must be rejected");
    };

    assert_eq!(error.kind(), GatewayErrorKind::PolicyDenied);
}

#[test]
fn invalid_continuation_record_should_not_be_forwarded_as_an_external_handle() {
    let service = DefaultExecutionService::new(
        RuntimeSnapshotHandle::new(start_snapshot()),
        Arc::new(TrackingExecutionStore::default()),
        ProviderRegistry::default(),
        Arc::new(UnusedAdmissions),
        Arc::new(UnusedCircuits),
        Arc::new(RejectedContinuation::InvalidData),
        Arc::new(RecordingClientApiKeyUsage::default()),
    );
    let client = service
        .authenticate("sk_start_test")
        .expect("authenticated client");

    let result = block_on(service.start(StartExecution {
        client,
        public_model: PublicModelId::new("gpt-start").expect("public model"),
        operation: start_operation(),
        metadata: execution_metadata_with_continuation(),
    }));
    let Err(error) = result else {
        panic!("invalid continuation state must be rejected");
    };

    assert_eq!(error.kind(), GatewayErrorKind::Internal);
}

fn execution_metadata_with_continuation() -> ExecutionRequestMetadata {
    ExecutionRequestMetadata {
        protocol: "openai".to_owned(),
        endpoint: "/v1/responses".to_owned(),
        transport: ClientTransport::HttpJson,
        stream: false,
        client_ip: None,
        user_agent: None,
        previous_response_id: Some(PreviousResponseId::new("response-private")),
    }
}

#[derive(Default)]
struct RecordingClientApiKeyUsage {
    ids: Mutex<Vec<String>>,
}

impl RecordingClientApiKeyUsage {
    fn recorded(&self) -> Vec<String> {
        self.ids.lock().expect("recorded API Key IDs").clone()
    }
}

impl ClientApiKeyUsageSink for RecordingClientApiKeyUsage {
    fn record_used(&self, key_id: &ClientApiKeyId) {
        self.ids
            .lock()
            .expect("recorded API Key IDs")
            .push(key_id.as_str().to_owned());
    }
}

#[derive(Default)]
struct TrackingExecutionStore {
    touched: AtomicBool,
    probe_failures: Mutex<Vec<String>>,
    requests: Mutex<Vec<NewModelRequest>>,
    attempts: Mutex<Vec<AttemptRecord>>,
    finalizations: Mutex<Vec<ModelRequestFinalization>>,
    create_gate: Mutex<Option<oneshot::Receiver<()>>>,
    finalize_gate: Mutex<Option<oneshot::Receiver<()>>>,
    creates: AtomicUsize,
    finalizes: AtomicUsize,
    fail_probe_observation: AtomicBool,
}

impl TrackingExecutionStore {
    fn touch(&self) {
        self.touched.store(true, Ordering::SeqCst);
    }

    fn probe_failures(&self) -> Vec<String> {
        self.probe_failures
            .lock()
            .expect("probe failures lock")
            .clone()
    }

    fn model_request_endpoints(&self) -> Vec<String> {
        self.requests
            .lock()
            .expect("model requests lock")
            .iter()
            .map(|request| request.endpoint.clone())
            .collect()
    }

    fn requested_models(&self) -> Vec<Option<String>> {
        self.requests
            .lock()
            .expect("model requests lock")
            .iter()
            .map(|request| {
                request
                    .requested_model
                    .as_ref()
                    .map(|model| model.as_str().to_owned())
            })
            .collect()
    }

    fn upstream_models(&self) -> Vec<Option<String>> {
        self.attempts
            .lock()
            .expect("attempts lock")
            .iter()
            .map(|attempt| {
                attempt
                    .upstream_model_id
                    .as_ref()
                    .map(|model| model.as_str().to_owned())
            })
            .collect()
    }
}

#[async_trait]
impl ExecutionStore for TrackingExecutionStore {
    async fn create_model_request(&self, request: NewModelRequest) -> Result<(), StoreError> {
        self.touch();
        self.creates.fetch_add(1, Ordering::SeqCst);
        let gate = self.create_gate.lock().unwrap().take();
        if let Some(gate) = gate {
            gate.await.expect("create model request gate");
        }
        self.requests
            .lock()
            .expect("model requests lock")
            .push(request);
        Ok(())
    }

    async fn record_attempt(&self, attempt: AttemptRecord) -> Result<(), StoreError> {
        self.touch();
        self.attempts.lock().expect("attempts lock").push(attempt);
        Ok(())
    }

    async fn mark_send_state(
        &self,
        _: &ModelRequestId,
        _: UpstreamSendState,
    ) -> Result<(), StoreError> {
        self.touch();
        Ok(())
    }

    async fn mark_downstream_committed(
        &self,
        _: &ModelRequestId,
        _: SystemTime,
        _: Option<u16>,
    ) -> Result<(), StoreError> {
        self.touch();
        Ok(())
    }

    async fn record_client_status(&self, _: &ModelRequestId, _: u16) -> Result<(), StoreError> {
        self.touch();
        Ok(())
    }

    async fn record_intermediate_failure(&self, _: IntermediateFailure) -> Result<(), StoreError> {
        self.touch();
        Ok(())
    }

    async fn record_probe_failure(&self, failure: ProbeFailure) -> Result<(), StoreError> {
        if self.fail_probe_observation.load(Ordering::SeqCst) {
            return Err(StoreError::new(StoreErrorKind::Unavailable));
        }
        self.probe_failures
            .lock()
            .expect("probe failures lock")
            .push(failure.error.kind().as_str().to_owned());
        Ok(())
    }

    async fn finalize_model_request(
        &self,
        finalization: ModelRequestFinalization,
    ) -> Result<(), StoreError> {
        self.touch();
        self.finalizes.fetch_add(1, Ordering::SeqCst);
        let gate = self.finalize_gate.lock().unwrap().take();
        if let Some(gate) = gate {
            gate.await.expect("finalize model request gate");
        }
        self.finalizations
            .lock()
            .expect("finalizations lock")
            .push(finalization);
        Ok(())
    }

    async fn recover_expired(&self, _: SystemTime) -> Result<RecoveryReport, StoreError> {
        self.touch();
        Ok(RecoveryReport::default())
    }
}

struct UnusedAdmissions;

impl ClientAdmissionPort for UnusedAdmissions {
    fn abandon(
        &self,
        key: &gateway_core::policy::ClientApiKeyId,
        request: &gateway_core::engine::ModelRequestId,
    ) {
        let _ = futures::FutureExt::now_or_never(self.release(key, request));
    }

    fn admit(
        &self,
        _: ClientAdmissionRequest,
    ) -> BoxFuture<'_, Result<ClientAdmissionDecision, ClientAdmissionError>> {
        Box::pin(async { Ok(ClientAdmissionDecision::Granted) })
    }

    fn release<'a>(
        &'a self,
        _: &'a ClientApiKeyId,
        _: &'a ModelRequestId,
    ) -> BoxFuture<'a, Result<bool, ClientAdmissionError>> {
        Box::pin(async { Ok(true) })
    }

    fn restore(
        &self,
        _: ClientAdmissionRecovery,
    ) -> BoxFuture<'_, Result<ClientAdmissionRestoreResult, ClientAdmissionError>> {
        Box::pin(async { Ok(ClientAdmissionRestoreResult::default()) })
    }
}

struct UnusedCircuits;

impl ProviderCircuitPort for UnusedCircuits {
    fn decision<'a>(
        &'a self,
        _: &'a ProviderKind,
    ) -> BoxFuture<'a, Result<ProviderCircuitDecision, ProviderCircuitError>> {
        Box::pin(async { Ok(ProviderCircuitDecision::Allow) })
    }

    fn observe_failure<'a>(
        &'a self,
        _: &'a ProviderKind,
    ) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
        Box::pin(async { Ok(()) })
    }

    fn observe_success<'a>(
        &'a self,
        _: &'a ProviderKind,
    ) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
        Box::pin(async { Ok(()) })
    }
}

struct FailingDecisionCircuits;

impl ProviderCircuitPort for FailingDecisionCircuits {
    fn decision<'a>(
        &'a self,
        _: &'a ProviderKind,
    ) -> BoxFuture<'a, Result<ProviderCircuitDecision, ProviderCircuitError>> {
        Box::pin(async { Err(ProviderCircuitError) })
    }

    fn observe_failure<'a>(
        &'a self,
        _: &'a ProviderKind,
    ) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
        Box::pin(async { Ok(()) })
    }

    fn observe_success<'a>(
        &'a self,
        _: &'a ProviderKind,
    ) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
        Box::pin(async { Ok(()) })
    }
}

struct PendingDecisionCircuits;

impl ProviderCircuitPort for PendingDecisionCircuits {
    fn decision<'a>(
        &'a self,
        _: &'a ProviderKind,
    ) -> BoxFuture<'a, Result<ProviderCircuitDecision, ProviderCircuitError>> {
        Box::pin(futures::future::pending())
    }

    fn observe_failure<'a>(
        &'a self,
        _: &'a ProviderKind,
    ) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
        Box::pin(async { Ok(()) })
    }

    fn observe_success<'a>(
        &'a self,
        _: &'a ProviderKind,
    ) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
        Box::pin(async { Ok(()) })
    }
}

struct UnusedContinuation;

impl NativeContinuationPort for UnusedContinuation {
    fn resolve<'a>(
        &'a self,
        _: &'a ClientApiKeyId,
        _: &'a PreviousResponseId,
    ) -> BoxFuture<'a, Result<Option<NativeContinuationPin>, NativeContinuationStoreError>> {
        Box::pin(async { Ok(None) })
    }

    fn record<'a>(
        &'a self,
        _: NativeContinuationPin,
    ) -> BoxFuture<'a, Result<(), NativeContinuationStoreError>> {
        Box::pin(async { Ok(()) })
    }
}

enum RejectedContinuation {
    OwnershipMismatch,
    InvalidData,
}

impl NativeContinuationPort for RejectedContinuation {
    fn resolve<'a>(
        &'a self,
        _: &'a ClientApiKeyId,
        _: &'a PreviousResponseId,
    ) -> BoxFuture<'a, Result<Option<NativeContinuationPin>, NativeContinuationStoreError>> {
        Box::pin(async move {
            Err(match self {
                Self::OwnershipMismatch => NativeContinuationStoreError::ownership_mismatch(),
                Self::InvalidData => NativeContinuationStoreError::invalid_data("invalid record"),
            })
        })
    }

    fn record<'a>(
        &'a self,
        _: NativeContinuationPin,
    ) -> BoxFuture<'a, Result<(), NativeContinuationStoreError>> {
        Box::pin(async { Ok(()) })
    }
}

fn account_scope(provider: &ProviderKind, account_id: &str) -> Arc<FrozenAccountScope> {
    Arc::new(FrozenAccountScope::new(
        Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([(
            ProviderAccountId::new(account_id).expect("account ID"),
            RuntimeAccount::new(provider.clone(), BTreeSet::new()),
        )]))),
        ClientRoutingScope::all_accounts(),
    ))
}

fn probe_snapshot() -> RuntimeSnapshot {
    let provider = ProviderKind::new("openai").expect("provider kind");
    let capabilities =
        ModelCapabilities::new(BTreeSet::from([OperationKind::Generate]), Some(16_000));
    RuntimeSnapshot::new(
        ConfigRevision::new(1).expect("config revision"),
        AccountSelectionPolicy::new(
            RotationStrategy::Smart,
            std::num::NonZeroU32::new(1).expect("concurrency"),
            Duration::from_millis(1),
        ),
        vec![provider.clone()],
        vec![ProviderModel::new(
            provider,
            UpstreamModelId::new("gpt-probe").expect("model ID"),
            capabilities,
        )],
        Vec::new(),
    )
    .expect("probe snapshot")
}

fn client_snapshot() -> RuntimeSnapshot {
    let provider = ProviderKind::new("openai").expect("provider kind");
    RuntimeSnapshot::new(
        ConfigRevision::new(1).expect("config revision"),
        AccountSelectionPolicy::new(
            RotationStrategy::Smart,
            std::num::NonZeroU32::new(1).expect("concurrency"),
            Duration::from_millis(1),
        ),
        vec![provider.clone()],
        Vec::new(),
        vec![ClientPolicy::new(
            ClientApiKeyId::new("key_usage_test").expect("client API key ID"),
            PlaintextClientApiKey::new("sk_usage_test").expect("plaintext client API key"),
            account_scope(&provider, "acct_usage"),
            true,
            RateLimits::unlimited(),
        )],
    )
    .expect("client snapshot")
}

fn start_snapshot() -> RuntimeSnapshot {
    start_snapshot_with_policy(1, true, RateLimits::unlimited(), false)
}

fn start_snapshot_with_policy(
    revision: u64,
    enabled: bool,
    limits: RateLimits,
    disable_fast: bool,
) -> RuntimeSnapshot {
    let provider = ProviderKind::new("openai").expect("provider kind");
    let capabilities =
        ModelCapabilities::new(BTreeSet::from([OperationKind::Generate]), Some(16_000));
    RuntimeSnapshot::new(
        ConfigRevision::new(revision).expect("config revision"),
        AccountSelectionPolicy::new(
            RotationStrategy::Smart,
            std::num::NonZeroU32::new(1).expect("concurrency"),
            Duration::from_millis(1),
        ),
        vec![provider.clone()],
        vec![ProviderModel::new(
            provider.clone(),
            UpstreamModelId::new("gpt-start").expect("model ID"),
            capabilities,
        )],
        vec![ClientPolicy::new(
            ClientApiKeyId::new("key_start_test").expect("client API key ID"),
            PlaintextClientApiKey::new("sk_start_test").expect("plaintext client API key"),
            Arc::new(
                account_scope(&provider, "acct_start")
                    .as_ref()
                    .clone()
                    .with_disable_fast(disable_fast),
            ),
            enabled,
            limits,
        )],
    )
    .expect("start snapshot")
}

fn probe_operation() -> Operation {
    let body = json!({
        "model": "gpt-probe",
        "input": [{"type": "message", "role": "user", "content": "ping"}],
    });
    Operation::Generate(GenerateRequest::from_protocol_payload(
        ProtocolPayload::json_object("openai", body.as_object().expect("request object").clone())
            .expect("OpenAI payload"),
    ))
}

fn start_operation() -> Operation {
    start_operation_for_model("gpt-start")
}

fn start_operation_for_model(model: &str) -> Operation {
    let body = json!({
        "model": model,
        "input": [{"type": "message", "role": "user", "content": "ping"}],
    });
    Operation::Generate(GenerateRequest::from_protocol_payload(
        ProtocolPayload::json_object("openai", body.as_object().expect("request object").clone())
            .expect("OpenAI payload"),
    ))
}

#[derive(Default)]
struct BoundedAdmissions {
    active: Mutex<BTreeSet<ModelRequestId>>,
    granted: AtomicUsize,
    rpm_exhausted: AtomicBool,
    hold_acquisition: AtomicBool,
}

#[derive(Default)]
struct QueuedAccountProvider {
    waiting: gateway_core::concurrency::ConcurrencyWaitQueue<&'static str>,
}

#[async_trait]
impl Provider for QueuedAccountProvider {
    fn name(&self) -> &'static str {
        "openai"
    }

    fn catalog_generation(&self) -> ProviderCatalogGeneration {
        ProviderCatalogGeneration::default()
    }

    async fn query_model_capabilities(
        &self,
    ) -> Result<Vec<ProviderModelCapabilities>, ProviderError> {
        Ok(Vec::new())
    }

    async fn execute(
        &self,
        _: ProviderRequest,
        context: AttemptContext,
    ) -> Result<ProviderStream, ProviderError> {
        let mut waiting = gateway_core::concurrency::CapacityWait::new(
            &self.waiting,
            gateway_core::concurrency::ConcurrencyQueuePolicy {
                max_waiting: 1,
                timeout: Duration::from_secs(5),
            },
            context.deadline(),
            context.concurrency_wait_budget(),
        );
        loop {
            waiting.wait(&["account"]).await.map_err(|error| {
                ProviderError::new(error.provider_kind(), UpstreamSendState::NotSent)
            })?;
        }
    }
}

#[test]
fn account_wait_inherits_the_budget_spent_during_client_admission() {
    use futures::FutureExt;
    block_on(async {
        let admissions = Arc::new(BoundedAdmissions::default());
        let snapshot = start_snapshot_with_policy(
            1,
            true,
            RateLimits {
                max_concurrency: 1,
                requests_per_minute: 0,
            },
            false,
        )
        .with_client_queue_policy(gateway_core::concurrency::ConcurrencyQueuePolicy {
            max_waiting: 1,
            timeout: Duration::from_secs(1),
        });
        let service = DefaultExecutionService::new(
            RuntimeSnapshotHandle::new(snapshot),
            Arc::new(TrackingExecutionStore::default()),
            ProviderRegistry::new(
                [Arc::new(QueuedAccountProvider::default()) as Arc<dyn Provider>],
            )
            .unwrap(),
            admissions.clone(),
            Arc::new(UnusedCircuits),
            Arc::new(UnusedContinuation),
            Arc::new(RecordingClientApiKeyUsage::default()),
        );
        let running = service
            .start(request(&service, ClientTransport::HttpSse))
            .await
            .unwrap();
        let started_at = Instant::now();
        let mut queued = service.start(request(&service, ClientTransport::HttpSse));
        assert!(queued.as_mut().now_or_never().is_none());
        futures_timer::Delay::new(Duration::from_millis(200)).await;
        drop(running);
        let mut started = queued.await.unwrap();
        let error = started.session.collect_uncommitted().await.unwrap_err();
        assert!(
            matches!(
                error,
                EngineError::Provider(ref error)
                    if error.kind() == ProviderErrorKind::ConcurrencyQueueTimeout
            ),
            "{error:?}"
        );
        assert!(
            started_at.elapsed() < Duration::from_secs(3),
            "账号层不能重新获得自己的 5 秒预算，实际耗时 {:?}",
            started_at.elapsed()
        );
        assert!(admissions.active.lock().unwrap().is_empty());
    });
}

impl ClientAdmissionPort for BoundedAdmissions {
    fn abandon(&self, _: &ClientApiKeyId, id: &ModelRequestId) {
        self.active.lock().unwrap().remove(id);
    }

    fn admit(
        &self,
        request: ClientAdmissionRequest,
    ) -> BoxFuture<'_, Result<ClientAdmissionDecision, ClientAdmissionError>> {
        Box::pin(async move {
            use gateway_core::engine::admission::ClientAdmissionRejection;
            if self.rpm_exhausted.load(Ordering::SeqCst) {
                return Ok(ClientAdmissionDecision::Rejected(
                    ClientAdmissionRejection::RateLimited,
                ));
            }
            {
                let mut active = self.active.lock().unwrap();
                if !request.allow_concurrency_acquire
                    || (request.limits.max_concurrency > 0
                        && active.len() as u64 >= request.limits.max_concurrency)
                {
                    return Ok(ClientAdmissionDecision::Rejected(
                        ClientAdmissionRejection::ConcurrencyLimited,
                    ));
                }
                active.insert(request.model_request_id);
            }
            self.granted.fetch_add(1, Ordering::SeqCst);
            if self.hold_acquisition.load(Ordering::SeqCst) {
                futures::future::pending::<()>().await;
            }
            Ok(ClientAdmissionDecision::Granted)
        })
    }

    fn release<'a>(
        &'a self,
        _: &'a ClientApiKeyId,
        id: &'a ModelRequestId,
    ) -> BoxFuture<'a, Result<bool, ClientAdmissionError>> {
        Box::pin(async move { Ok(self.active.lock().unwrap().remove(id)) })
    }

    fn restore(
        &self,
        _: ClientAdmissionRecovery,
    ) -> BoxFuture<'_, Result<ClientAdmissionRestoreResult, ClientAdmissionError>> {
        Box::pin(async { Ok(ClientAdmissionRestoreResult::default()) })
    }
}

fn queue_service(
    max_concurrency: u64,
    max_waiting: u32,
    timeout: Duration,
) -> (DefaultExecutionService, Arc<BoundedAdmissions>) {
    let admissions = Arc::new(BoundedAdmissions::default());
    let snapshot = start_snapshot_with_policy(
        1,
        true,
        RateLimits {
            max_concurrency,
            requests_per_minute: 0,
        },
        false,
    )
    .with_client_queue_policy(gateway_core::concurrency::ConcurrencyQueuePolicy {
        max_waiting,
        timeout,
    });
    let service = DefaultExecutionService::new(
        RuntimeSnapshotHandle::new(snapshot),
        Arc::new(TrackingExecutionStore::default()),
        ProviderRegistry::default(),
        admissions.clone(),
        Arc::new(UnusedCircuits),
        Arc::new(UnusedContinuation),
        Arc::new(RecordingClientApiKeyUsage::default()),
    );
    (service, admissions)
}

#[test]
fn client_queue_allows_five_running_five_waiting_and_rejects_the_eleventh() {
    use futures::FutureExt;
    block_on(async {
        let (service, admissions) = queue_service(5, 5, Duration::from_secs(2));
        let mut running = Vec::new();
        for _ in 0..5 {
            running.push(
                service
                    .start(request(&service, ClientTransport::HttpSse))
                    .await
                    .unwrap(),
            );
        }
        let mut waiting = Vec::new();
        for _ in 0..5 {
            let mut future = service.start(request(&service, ClientTransport::HttpSse));
            assert!(future.as_mut().now_or_never().is_none());
            waiting.push(future);
        }
        let rejected = service
            .start(request(&service, ClientTransport::HttpSse))
            .await
            .err()
            .unwrap();
        assert_eq!(rejected.kind(), GatewayErrorKind::ConcurrencyQueueFull);
        assert_eq!(
            admissions.granted.load(Ordering::SeqCst),
            5,
            "waiting must not consume admission/RPM"
        );
        drop(running.remove(0));
        let first = waiting.remove(0).await.unwrap();
        assert_eq!(admissions.active.lock().unwrap().len(), 5);
        assert_eq!(admissions.granted.load(Ordering::SeqCst), 6);
        assert!(waiting[0].as_mut().now_or_never().is_none());
        drop((first, waiting, running));
        assert!(admissions.active.lock().unwrap().is_empty());
    });
}

#[test]
fn cancelled_waiter_releases_its_place_and_new_requests_do_not_overtake_fifo() {
    use futures::FutureExt;
    block_on(async {
        let (service, admissions) = queue_service(1, 2, Duration::from_secs(2));
        let running = service
            .start(request(&service, ClientTransport::WebSocket))
            .await
            .unwrap();
        let mut first = service.start(request(&service, ClientTransport::WebSocket));
        let mut cancelled = service.start(request(&service, ClientTransport::WebSocket));
        assert!(first.as_mut().now_or_never().is_none());
        assert!(cancelled.as_mut().now_or_never().is_none());
        drop(cancelled);
        drop(running);
        let mut later = service.start(request(&service, ClientTransport::WebSocket));
        assert!(
            later.as_mut().now_or_never().is_none(),
            "a newly freed slot belongs to the existing queue head"
        );
        let first = first.await.unwrap();
        assert_eq!(admissions.granted.load(Ordering::SeqCst), 2);
        drop(first);
        let later = later.await.unwrap();
        drop(later);
        assert!(admissions.active.lock().unwrap().is_empty());
    });
}

#[test]
fn queue_timeout_and_rpm_rejection_leave_no_new_admission() {
    block_on(async {
        let (service, admissions) = queue_service(1, 1, Duration::from_millis(20));
        let running = service
            .start(request(&service, ClientTransport::HttpJson))
            .await
            .unwrap();
        let timeout = service
            .start(request(&service, ClientTransport::HttpJson))
            .await
            .err()
            .unwrap();
        assert_eq!(timeout.kind(), GatewayErrorKind::ConcurrencyQueueTimeout);
        admissions.rpm_exhausted.store(true, Ordering::SeqCst);
        let limited = service
            .start(request(&service, ClientTransport::HttpJson))
            .await
            .err()
            .unwrap();
        assert_eq!(limited.kind(), GatewayErrorKind::RateLimited);
        assert_eq!(admissions.granted.load(Ordering::SeqCst), 1);
        drop(running);
    });
}

#[test]
fn cancelling_pending_admission_cleans_up_a_slot_acquired_before_the_reply() {
    use futures::FutureExt;
    let (service, admissions) = queue_service(1, 1, Duration::from_secs(2));
    admissions.hold_acquisition.store(true, Ordering::SeqCst);
    let mut pending = service.start(request(&service, ClientTransport::HttpSse));
    assert!(pending.as_mut().now_or_never().is_none());
    assert_eq!(admissions.active.lock().unwrap().len(), 1);
    drop(pending);
    assert!(admissions.active.lock().unwrap().is_empty());
}

#[test]
fn reused_websocket_client_gets_group_fast_policy_from_each_new_request_snapshot() {
    let snapshots = RuntimeSnapshotHandle::new(start_snapshot());
    let admissions = Arc::new(Admissions::default());
    let provider = Arc::new(ChargedProvider::default());
    let service = DefaultExecutionService::new(
        snapshots.clone(),
        Arc::new(TrackingExecutionStore::default()),
        ProviderRegistry::new([provider.clone() as Arc<dyn Provider>]).unwrap(),
        admissions,
        Arc::new(UnusedCircuits),
        Arc::new(UnusedContinuation),
        Arc::new(RecordingClientApiKeyUsage::default()),
    );
    let client = service.authenticate("sk_start_test").unwrap();
    for (revision, disable_fast) in [(2, true), (3, false)] {
        snapshots.publish(start_snapshot_with_policy(
            revision,
            true,
            RateLimits::unlimited(),
            disable_fast,
        ));
        let mut next = request(&service, ClientTransport::WebSocket);
        next.client = client.clone();
        let mut started = block_on(service.start(next)).unwrap();
        block_on(started.session.collect_uncommitted()).unwrap();
        block_on(started.session.detach_finalize());
    }
    assert_eq!(*provider.policies.lock().unwrap(), vec![true, false]);
}