Reuse tool runtime for code mode worker (#14496)

## Summary
- create the turn-scoped `ToolCallRuntime` before starting the code mode
worker so the worker reuses the same runtime and router
- thread the shared runtime through the code mode service/worker path
and use it for nested tool calls
- model aborted tool calls as a concrete `ToolOutput` so aborted
responses still produce valid tool output shapes

## Testing
- `just fmt`
- `cargo test -p codex-core` (still running locally)
This commit is contained in:
pakrym-oai
2026-03-12 12:48:32 -07:00
committed by GitHub
parent d3e6680531
commit 09ba6b47ae
8 changed files with 112 additions and 81 deletions

View File

@@ -5551,11 +5551,6 @@ pub(crate) async fn run_turn(
// Although from the perspective of codex.rs, TurnDiffTracker has the lifecycle of a Task which contains
// many turns, from the perspective of the user, it is a single turn.
let turn_diff_tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new()));
let _code_mode_worker = sess
.services
.code_mode_service
.start_turn_worker(&sess, &turn_context, &turn_diff_tracker)
.await;
let mut server_model_warning_emitted_for_turn = false;
// `ModelClientSession` is turn-scoped and caches WebSocket + sticky routing state, so we reuse
@@ -6161,10 +6156,26 @@ async fn run_sampling_request(
turn_context.as_ref(),
base_instructions,
);
let tool_runtime = ToolCallRuntime::new(
Arc::clone(&router),
Arc::clone(&sess),
Arc::clone(&turn_context),
Arc::clone(&turn_diff_tracker),
);
let _code_mode_worker = sess
.services
.code_mode_service
.start_turn_worker(
&sess,
&turn_context,
Arc::clone(&router),
Arc::clone(&turn_diff_tracker),
)
.await;
let mut retries = 0;
loop {
let err = match try_run_sampling_request(
Arc::clone(&router),
tool_runtime.clone(),
Arc::clone(&sess),
Arc::clone(&turn_context),
client_session,
@@ -6919,7 +6930,7 @@ async fn drain_in_flight(
)
)]
async fn try_run_sampling_request(
router: Arc<ToolRouter>,
tool_runtime: ToolCallRuntime,
sess: Arc<Session>,
turn_context: Arc<TurnContext>,
client_session: &mut ModelClientSession,
@@ -6950,13 +6961,6 @@ async fn try_run_sampling_request(
.instrument(trace_span!("stream_request"))
.or_cancel(&cancellation_token)
.await??;
let tool_runtime = ToolCallRuntime::new(
Arc::clone(&router),
Arc::clone(&sess),
Arc::clone(&turn_context),
Arc::clone(&turn_diff_tracker),
);
let mut in_flight: FuturesOrdered<BoxFuture<'static, CodexResult<ResponseInputItem>>> =
FuturesOrdered::new();
let mut needs_follow_up = false;