Send warmup request (#11258)

Send a request with `generate: falls` but a full set of tools and
instructions to pre-warm inference.

---------

Co-authored-by: Codex <noreply@openai.com>

codex-rs/core/src/tasks/regular.rs

@@ -3,77 +3,61 @@ use std::sync::Mutex;
 
 use crate::client::ModelClient;
 use crate::client::ModelClientSession;
+use crate::client_common::Prompt;
 use crate::codex::TurnContext;
 use crate::codex::run_turn;
+use crate::error::Result as CodexResult;
 use crate::state::TaskKind;
 use async_trait::async_trait;
-use codex_otel::OtelManager;
-use codex_protocol::openai_models::ModelInfo;
 use codex_protocol::user_input::UserInput;
-use tokio::task::JoinHandle;
 use tokio_util::sync::CancellationToken;
 use tracing::Instrument;
 use tracing::trace_span;
-use tracing::warn;
 
 use super::SessionTask;
 use super::SessionTaskContext;
 
-type PrewarmedSessionTask = JoinHandle<Option<ModelClientSession>>;
-
 pub(crate) struct RegularTask {
-    prewarmed_session_task: Mutex<Option<PrewarmedSessionTask>>,
+    prewarmed_session: Mutex<Option<ModelClientSession>>,
 }
 
 impl Default for RegularTask {
     fn default() -> Self {
         Self {
-            prewarmed_session_task: Mutex::new(None),
+            prewarmed_session: Mutex::new(None),
         }
     }
 }
 
 impl RegularTask {
-    pub(crate) fn with_startup_prewarm(
+    pub(crate) async fn with_startup_prewarm(
         model_client: ModelClient,
-        otel_manager: OtelManager,
-        model_info: ModelInfo,
-    ) -> Self {
-        let prewarmed_session_task = tokio::spawn(async move {
-            let mut client_session = model_client.new_session();
-            match client_session
-                .prewarm_websocket(&otel_manager, &model_info)
-                .await
-            {
-                Ok(()) => Some(client_session),
-                Err(err) => {
-                    warn!("startup websocket prewarm task failed: {err}");
-                    None
-                }
-            }
-        });
+        prompt: Prompt,
+        turn_context: Arc<TurnContext>,
+        turn_metadata_header: Option<String>,
+    ) -> CodexResult<Self> {
+        let mut client_session = model_client.new_session();
+        client_session
+            .prewarm_websocket(
+                &prompt,
+                &turn_context.model_info,
+                &turn_context.otel_manager,
+                turn_context.reasoning_effort,
+                turn_context.reasoning_summary,
+                turn_metadata_header.as_deref(),
+            )
+            .await?;
 
-        Self {
-            prewarmed_session_task: Mutex::new(Some(prewarmed_session_task)),
-        }
+        Ok(Self {
+            prewarmed_session: Mutex::new(Some(client_session)),
+        })
     }
 
     async fn take_prewarmed_session(&self) -> Option<ModelClientSession> {
-        let prewarmed_session_task = self
-            .prewarmed_session_task
+        self.prewarmed_session
             .lock()
             .unwrap_or_else(std::sync::PoisonError::into_inner)
-            .take();
-        match prewarmed_session_task {
-            Some(task) => match task.await {
-                Ok(client_session) => client_session,
-                Err(err) => {
-                    warn!("startup websocket prewarm task join failed: {err}");
-                    None
-                }
-            },
-            None => None,
-        }
+            .take()
     }
 }