diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_tests.rs index 97d6eacefa..efde4f6cdf 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_tests.rs @@ -251,11 +251,32 @@ async fn stalled_http_headers_exhaust_the_sampling_retry_budget() -> Result<()> config.provider = create_model_provider(provider, config.provider.auth_manager()); let sampler = LunaSampler::new(config); - let result = tokio::time::timeout( + tokio::time::pause(); + let sample = tokio::time::timeout( Duration::from_secs(10), sampler.sample(sample_request("stalled-headers")), - ) - .await?; + ); + tokio::pin!(sample); + let deadline = std::time::Instant::now() + Duration::from_secs(/*secs*/ 30); + for expected_requests in 1..=3 { + // Keep the paused runtime runnable until the server records this attempt. + // Otherwise auto-advance can exhaust a retry before its request arrives. + while mock.requests().len() < expected_requests { + anyhow::ensure!( + std::time::Instant::now() < deadline, + "server did not receive request {expected_requests}" + ); + tokio::select! { + result = &mut sample => { + anyhow::bail!("sampling finished before request {expected_requests}: {result:?}"); + } + () = tokio::task::yield_now() => {} + } + } + tokio::time::advance(Duration::from_millis(/*millis*/ 1_001)).await; + } + let result = sample.await?; + tokio::time::resume(); assert!(matches!( result, Err(LunaSamplerError::Api(ApiError::Transport( diff --git a/codex-rs/model-provider/src/models_endpoint_timeout_tests.rs b/codex-rs/model-provider/src/models_endpoint_timeout_tests.rs index 41ed3a0557..1cb11fba7b 100644 --- a/codex-rs/model-provider/src/models_endpoint_timeout_tests.rs +++ b/codex-rs/model-provider/src/models_endpoint_timeout_tests.rs @@ -10,15 +10,15 @@ use wiremock::matchers::path; #[tokio::test] async fn catalog_deadline_returns_request_timeout() { let server = MockServer::start().await; - Mock::given(method("GET")) + let mock = Mock::given(method("GET")) .and(path("/models")) .respond_with( ResponseTemplate::new(200) .set_body_json(serde_json::json!({ "models": [] })) - .set_delay(MODELS_REFRESH_TIMEOUT * 2), + .set_delay(Duration::from_secs(/*secs*/ 60)), ) .expect(1) - .mount(&server) + .mount_as_scoped(&server) .await; let endpoint = OpenAiModelsEndpoint::new( ModelProviderInfo::create_openai_provider(Some(server.uri())), @@ -26,13 +26,39 @@ async fn catalog_deadline_returns_request_timeout() { /*gateway_auth_manager*/ None, ); - let error = endpoint - .list_models( + tokio::time::pause(); + let request = tokio::time::timeout( + MODELS_REFRESH_TIMEOUT * 2, + endpoint.list_models( "0.0.0", HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), - ) + ), + ); + tokio::pin!(request); + let received = mock.wait_until_satisfied(); + tokio::pin!(received); + let deadline = std::time::Instant::now() + Duration::from_secs(/*secs*/ 30); + // Keep the paused runtime runnable until the request reaches the server. + // Otherwise auto-advance can expire the catalog deadline during dispatch. + loop { + assert!( + std::time::Instant::now() < deadline, + "server did not receive the catalog request" + ); + tokio::select! { + () = &mut received => break, + result = &mut request => { + panic!("catalog request finished before reaching the server: {result:?}"); + } + () = tokio::task::yield_now() => {} + } + } + tokio::time::advance(MODELS_REFRESH_TIMEOUT + Duration::from_millis(/*millis*/ 1)).await; + let error = request .await + .expect("catalog deadline should expire") .expect_err("delayed catalog request should time out"); + tokio::time::resume(); assert!( matches!(error.details(), CodexErrorDetails::RequestTimeout),