Reuse websocket connection (#9127)

Reuses the connection but still sends full requests.

codex-rs/core/tests/suite/client.rs

@@ -1171,7 +1171,7 @@ async fn azure_responses_request_includes_store_and_reasoning_ids() {
         SessionSource::Exec,
     );
 
-    let client = ModelClient::new(
+    let mut client = ModelClient::new(
         Arc::clone(&config),
         None,
         model_info,

codex-rs/core/tests/suite/websocket.rs

@@ -1,7 +1,9 @@
+#![allow(clippy::expect_used, clippy::unwrap_used)]
 use codex_core::AuthManager;
 use codex_core::CodexAuth;
 use codex_core::ContentItem;
 use codex_core::ModelClient;
+use codex_core::ModelClientSession;
 use codex_core::ModelProviderInfo;
 use codex_core::Prompt;
 use codex_core::ResponseEvent;
@@ -11,23 +13,97 @@ use codex_core::models_manager::manager::ModelsManager;
 use codex_core::protocol::SessionSource;
 use codex_otel::OtelManager;
 use codex_protocol::ThreadId;
+use codex_protocol::config_types::ReasoningSummary;
 use core_test_support::load_default_config_for_test;
+use core_test_support::responses::WebSocketTestServer;
 use core_test_support::responses::ev_completed;
 use core_test_support::responses::ev_response_created;
 use core_test_support::responses::start_websocket_server;
+use core_test_support::skip_if_no_network;
 use futures::StreamExt;
+use pretty_assertions::assert_eq;
 use std::sync::Arc;
 use tempfile::TempDir;
 
+const MODEL: &str = "gpt-5.2-codex";
+
+struct WebsocketTestHarness {
+    _codex_home: TempDir,
+    client: ModelClient,
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn responses_websocket_streams_request() {
+    skip_if_no_network!();
+
     let server = start_websocket_server(vec![vec![vec![
         ev_response_created("resp-1"),
         ev_completed("resp-1"),
     ]]])
     .await;
 
-    let provider = ModelProviderInfo {
+    let harness = websocket_harness(&server).await;
+    let mut session = harness.client.new_session();
+    let mut prompt = Prompt::default();
+    prompt.input = vec![ResponseItem::Message {
+        id: None,
+        role: "user".into(),
+        content: vec![ContentItem::InputText {
+            text: "hello".into(),
+        }],
+    }];
+
+    stream_until_complete(&mut session, &prompt).await;
+
+    let connection = server.single_connection();
+    assert_eq!(connection.len(), 1);
+    let body = connection.first().expect("missing request").body_json();
+
+    assert_eq!(body["model"].as_str(), Some(MODEL));
+    assert_eq!(body["stream"], serde_json::Value::Bool(true));
+    assert_eq!(body["input"].as_array().map(Vec::len), Some(1));
+
+    server.shutdown().await;
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn responses_websocket_reuses_connection() {
+    skip_if_no_network!();
+
+    let server = start_websocket_server(vec![vec![
+        vec![ev_response_created("resp-1"), ev_completed("resp-1")],
+        vec![ev_response_created("resp-2"), ev_completed("resp-2")],
+    ]])
+    .await;
+
+    let harness = websocket_harness(&server).await;
+    let mut session = harness.client.new_session();
+    let mut prompt = Prompt::default();
+    prompt.input = vec![ResponseItem::Message {
+        id: None,
+        role: "user".into(),
+        content: vec![ContentItem::InputText {
+            text: "hello".into(),
+        }],
+    }];
+
+    for _ in 0..2 {
+        stream_until_complete(&mut session, &prompt).await;
+    }
+
+    let connection = server.single_connection();
+    assert_eq!(connection.len(), 2);
+    let body = connection.first().expect("missing request").body_json();
+
+    assert_eq!(body["model"].as_str(), Some(MODEL));
+    assert_eq!(body["stream"], serde_json::Value::Bool(true));
+    assert_eq!(body["input"].as_array().map(Vec::len), Some(1));
+
+    server.shutdown().await;
+}
+
+fn websocket_provider(server: &WebSocketTestServer) -> ModelProviderInfo {
+    ModelProviderInfo {
         name: "mock-ws".into(),
         base_url: Some(format!("{}/v1", server.uri())),
         env_key: None,
@@ -41,23 +117,21 @@ async fn responses_websocket_streams_request() {
         stream_max_retries: Some(0),
         stream_idle_timeout_ms: Some(5_000),
         requires_openai_auth: false,
-    };
+    }
+}
 
+async fn websocket_harness(server: &WebSocketTestServer) -> WebsocketTestHarness {
+    let provider = websocket_provider(server);
     let codex_home = TempDir::new().unwrap();
     let mut config = load_default_config_for_test(&codex_home).await;
-    config.model_provider_id = provider.name.clone();
-    config.model_provider = provider.clone();
-    let effort = config.model_reasoning_effort;
-    let summary = config.model_reasoning_summary;
-    let model = ModelsManager::get_model_offline(config.model.as_deref());
-    config.model = Some(model.clone());
+    config.model = Some(MODEL.to_string());
     let config = Arc::new(config);
-    let model_info = ModelsManager::construct_model_info_offline(model.as_str(), &config);
+    let model_info = ModelsManager::construct_model_info_offline(MODEL, &config);
     let conversation_id = ThreadId::new();
     let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("Test API Key"));
     let otel_manager = OtelManager::new(
         conversation_id,
-        model.as_str(),
+        MODEL,
         model_info.slug.as_str(),
         None,
         Some("test@test.com".to_string()),
@@ -66,31 +140,27 @@ async fn responses_websocket_streams_request() {
         "test".to_string(),
         SessionSource::Exec,
     );
-
     let client = ModelClient::new(
         Arc::clone(&config),
         None,
         model_info,
         otel_manager,
-        provider,
-        effort,
-        summary,
+        provider.clone(),
+        None,
+        ReasoningSummary::Auto,
         conversation_id,
         SessionSource::Exec,
-    )
-    .new_session();
+    );
 
-    let mut prompt = Prompt::default();
-    prompt.input = vec![ResponseItem::Message {
-        id: None,
-        role: "user".into(),
-        content: vec![ContentItem::InputText {
-            text: "hello".into(),
-        }],
-    }];
+    WebsocketTestHarness {
+        _codex_home: codex_home,
+        client,
+    }
+}
 
-    let mut stream = client
-        .stream(&prompt)
+async fn stream_until_complete(session: &mut ModelClientSession, prompt: &Prompt) {
+    let mut stream = session
+        .stream(prompt)
         .await
         .expect("websocket stream failed");
 
@@ -99,14 +169,4 @@ async fn responses_websocket_streams_request() {
             break;
         }
     }
-
-    let connection = server.single_connection();
-    assert_eq!(connection.len(), 1);
-    let request = connection.first().cloned().unwrap();
-    let body = request.body_json();
-    assert_eq!(body["model"].as_str(), Some(model.as_str()));
-    assert_eq!(body["stream"], serde_json::Value::Bool(true));
-    assert_eq!(body["input"].as_array().map(Vec::len), Some(1));
-
-    server.shutdown().await;
 }