[codex] Reuse the tool router within a turn

codex-rs/core/src/session/turn.rs

@@ -110,6 +110,7 @@ use codex_utils_stream_parser::strip_citations;
 use futures::future::BoxFuture;
 use futures::prelude::*;
 use futures::stream::FuturesOrdered;
+use tokio::sync::OnceCell;
 use tokio_util::sync::CancellationToken;
 use tracing::Instrument;
 use tracing::error;
@@ -201,6 +202,7 @@ pub(crate) async fn run_turn(
     // 1. At the start of a turn, so the fresh turn input in `input` gets sampled first.
     // 2. After auto-compact, when model/tool continuation needs to resume before any steer.
 
+    let tool_router = OnceCell::new();
     loop {
         // Note that pending_input would be something like a message the user
         // submitted through the UI while the model was running. Though the UI
@@ -239,6 +241,7 @@ pub(crate) async fn run_turn(
             &mut client_session,
             &responses_metadata,
             sampling_request_input,
+            &tool_router,
             cancellation_token.child_token(),
         )
         .await
@@ -1041,14 +1044,17 @@ async fn run_sampling_request(
     client_session: &mut ModelClientSession,
     responses_metadata: &CodexResponsesMetadata,
     input: Vec<ResponseItem>,
+    tool_router: &OnceCell<Arc<ToolRouter>>,
     cancellation_token: CancellationToken,
 ) -> CodexResult<(SamplingRequestResult, Vec<ResponseItem>)> {
-    let router = built_tools(sess.as_ref(), turn_context.as_ref(), &cancellation_token).await?;
+    let router = tool_router
+        .get_or_try_init(|| built_tools(sess.as_ref(), turn_context.as_ref(), &cancellation_token))
+        .await?;
 
     let base_instructions = sess.get_base_instructions().await;
 
     let tool_runtime = ToolCallRuntime::new(
-        Arc::clone(&router),
+        Arc::clone(router),
         Arc::clone(&sess),
         Arc::clone(&turn_context),
         Arc::clone(&turn_diff_tracker),
@@ -1056,7 +1062,7 @@ async fn run_sampling_request(
     let _code_mode_worker = sess.services.code_mode_service.start_turn_worker(
         &sess,
         &turn_context,
-        Arc::clone(&router),
+        Arc::clone(router),
         Arc::clone(&turn_diff_tracker),
     );
     let max_retries = turn_context.provider.info().stream_max_retries();

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

@@ -30,6 +30,7 @@ use core_test_support::skip_if_no_network;
 use core_test_support::skip_if_sandbox;
 use core_test_support::test_codex::local;
 use core_test_support::test_codex::test_codex;
+use pretty_assertions::assert_eq;
 use regex_lite::Regex;
 use serde_json::Value;
 use serde_json::json;
@@ -130,6 +131,54 @@ async fn turn_environment_selection_keeps_environment_backed_tools() -> Result<(
     Ok(())
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn tool_specs_stay_stable_across_follow_up_sampling_requests() -> Result<()> {
+    skip_if_no_network!(Ok(()));
+
+    let server = start_mock_server().await;
+    let first_call_args = json!({
+        "plan": [{"step": "Inspect tools", "status": "in_progress"}],
+    })
+    .to_string();
+    let second_call_args = json!({
+        "plan": [{"step": "Inspect tools", "status": "completed"}],
+    })
+    .to_string();
+    let response_mock = mount_sse_sequence(
+        &server,
+        vec![
+            sse(vec![
+                ev_response_created("resp-1"),
+                ev_function_call("call-1", "update_plan", &first_call_args),
+                ev_completed("resp-1"),
+            ]),
+            sse(vec![
+                ev_response_created("resp-2"),
+                ev_function_call("call-2", "update_plan", &second_call_args),
+                ev_completed("resp-2"),
+            ]),
+            sse(vec![
+                ev_response_created("resp-3"),
+                ev_assistant_message("msg-1", "done"),
+                ev_completed("resp-3"),
+            ]),
+        ],
+    )
+    .await;
+    let mut builder = test_codex();
+    let test = builder.build(&server).await?;
+
+    test.submit_turn("inspect the available tools").await?;
+
+    let requests = response_mock.requests();
+    assert_eq!(requests.len(), 3);
+    let first_request_tools = requests[0].body_json()["tools"].clone();
+    assert_eq!(requests[1].body_json()["tools"], first_request_tools);
+    assert_eq!(requests[2].body_json()["tools"], first_request_tools);
+
+    Ok(())
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn custom_tool_unknown_returns_custom_output_error() -> Result<()> {
     skip_if_no_network!(Ok(()));