Move warmup to the task level (#11216)

Instead of storing a special connection on the client level make the
regular task responsible for establishing a normal client session and
open a connection on it.

Then when the turn is started we pass in a pre-established session.
This commit is contained in:
pakrym-oai
2026-02-09 10:57:52 -08:00
committed by GitHub
parent 9346d321d2
commit ccd17374cb
6 changed files with 177 additions and 221 deletions

View File

@@ -35,7 +35,6 @@ use core_test_support::responses::start_websocket_server_with_headers;
use core_test_support::skip_if_no_network;
use core_test_support::test_codex::test_codex;
use core_test_support::wait_for_event;
use futures::FutureExt;
use futures::StreamExt;
use opentelemetry_sdk::metrics::InMemoryMetricExporter;
use pretty_assertions::assert_eq;
@@ -103,11 +102,11 @@ async fn responses_websocket_preconnect_reuses_connection() {
.await;
let harness = websocket_harness(&server).await;
harness
.client
.pre_establish_connection(harness.otel_manager.clone(), async { None }.boxed());
let mut client_session = harness.client.new_session();
client_session
.prewarm_websocket(&harness.otel_manager, None)
.await
.expect("websocket prewarm failed");
let prompt = prompt_with_input(vec![message_item("hello")]);
stream_until_complete(&mut client_session, &harness, &prompt).await;
@@ -128,11 +127,11 @@ async fn responses_websocket_preconnect_is_reused_even_with_header_changes() {
.await;
let harness = websocket_harness(&server).await;
harness
.client
.pre_establish_connection(harness.otel_manager.clone(), async { None }.boxed());
let mut client_session = harness.client.new_session();
client_session
.prewarm_websocket(&harness.otel_manager, None)
.await
.expect("websocket prewarm failed");
let prompt = prompt_with_input(vec![message_item("hello")]);
let mut stream = client_session
.stream(