mirror of
https://github.com/openai/codex.git
synced 2026-05-03 10:56:37 +00:00
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:
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user