120 lines
3.9 KiB
Rust
120 lines
3.9 KiB
Rust
//! Verifies that the agent retries when the SSE stream terminates before
|
|
//! delivering a `response.completed` event.
|
|
|
|
use std::time::Duration;
|
|
|
|
use codex_core::Codex;
|
|
use codex_core::CodexSpawnOk;
|
|
use codex_core::ModelProviderInfo;
|
|
use codex_core::exec::CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR;
|
|
use codex_core::protocol::EventMsg;
|
|
use codex_core::protocol::InputItem;
|
|
use codex_core::protocol::Op;
|
|
use core_test_support::load_default_config_for_test;
|
|
use core_test_support::load_sse_fixture;
|
|
use core_test_support::load_sse_fixture_with_id;
|
|
use tempfile::TempDir;
|
|
use tokio::time::timeout;
|
|
use wiremock::Mock;
|
|
use wiremock::MockServer;
|
|
use wiremock::Request;
|
|
use wiremock::Respond;
|
|
use wiremock::ResponseTemplate;
|
|
use wiremock::matchers::method;
|
|
use wiremock::matchers::path;
|
|
|
|
fn sse_incomplete() -> String {
|
|
load_sse_fixture("tests/fixtures/incomplete_sse.json")
|
|
}
|
|
|
|
fn sse_completed(id: &str) -> String {
|
|
load_sse_fixture_with_id("tests/fixtures/completed_template.json", id)
|
|
}
|
|
|
|
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
|
async fn retries_on_early_close() {
|
|
#![allow(clippy::unwrap_used)]
|
|
|
|
if std::env::var(CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR).is_ok() {
|
|
println!(
|
|
"Skipping test because it cannot execute when network is disabled in a Codex sandbox."
|
|
);
|
|
return;
|
|
}
|
|
|
|
let server = MockServer::start().await;
|
|
|
|
struct SeqResponder;
|
|
impl Respond for SeqResponder {
|
|
fn respond(&self, _: &Request) -> ResponseTemplate {
|
|
use std::sync::atomic::AtomicUsize;
|
|
use std::sync::atomic::Ordering;
|
|
static CALLS: AtomicUsize = AtomicUsize::new(0);
|
|
let n = CALLS.fetch_add(1, Ordering::SeqCst);
|
|
if n == 0 {
|
|
ResponseTemplate::new(200)
|
|
.insert_header("content-type", "text/event-stream")
|
|
.set_body_raw(sse_incomplete(), "text/event-stream")
|
|
} else {
|
|
ResponseTemplate::new(200)
|
|
.insert_header("content-type", "text/event-stream")
|
|
.set_body_raw(sse_completed("resp_ok"), "text/event-stream")
|
|
}
|
|
}
|
|
}
|
|
|
|
Mock::given(method("POST"))
|
|
.and(path("/v1/responses"))
|
|
.respond_with(SeqResponder {})
|
|
.expect(2)
|
|
.mount(&server)
|
|
.await;
|
|
|
|
// Configure retry behavior explicitly to avoid mutating process-wide
|
|
// environment variables.
|
|
|
|
let model_provider = ModelProviderInfo {
|
|
name: "openai".into(),
|
|
base_url: format!("{}/v1", server.uri()),
|
|
// Environment variable that should exist in the test environment.
|
|
// ModelClient will return an error if the environment variable for the
|
|
// provider is not set.
|
|
env_key: Some("PATH".into()),
|
|
env_key_instructions: None,
|
|
wire_api: codex_core::WireApi::Responses,
|
|
query_params: None,
|
|
http_headers: None,
|
|
env_http_headers: None,
|
|
// exercise retry path: first attempt yields incomplete stream, so allow 1 retry
|
|
request_max_retries: Some(0),
|
|
stream_max_retries: Some(1),
|
|
stream_idle_timeout_ms: Some(2000),
|
|
};
|
|
|
|
let ctrl_c = std::sync::Arc::new(tokio::sync::Notify::new());
|
|
let codex_home = TempDir::new().unwrap();
|
|
let mut config = load_default_config_for_test(&codex_home);
|
|
config.model_provider = model_provider;
|
|
let CodexSpawnOk { codex, .. } = Codex::spawn(config, ctrl_c).await.unwrap();
|
|
|
|
codex
|
|
.submit(Op::UserInput {
|
|
items: vec![InputItem::Text {
|
|
text: "hello".into(),
|
|
}],
|
|
})
|
|
.await
|
|
.unwrap();
|
|
|
|
// Wait until TaskComplete (should succeed after retry).
|
|
loop {
|
|
let ev = timeout(Duration::from_secs(10), codex.next_event())
|
|
.await
|
|
.unwrap()
|
|
.unwrap();
|
|
if matches!(ev.msg, EventMsg::TaskComplete(_)) {
|
|
break;
|
|
}
|
|
}
|
|
}
|