openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-14 21:40:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

Capture direct-tool critical paths behind a feature flag (#57201)

Browse Source 
## What changed

Add the default-disabled `critical_path_timing` feature to capture dependency paths for root-agent turns. Connect direct tool dispatch, execution-lock handoffs, and completed tool batches to turn timing so overlapping work shares its common history without double counting.

Omit captures for unsupported dependencies, including CodeMode, agent communication, approval review, and user waits. Settle aborted tasks before freezing timing, invalidate captures on forced teardown, and apply runtime feature changes only to future turns.

## Testing

Add coverage for serial and overlapping tool dispatch, execution-lock dependency handoffs, cancellation, unsupported paths, frozen timing after completion, and configuration refresh boundaries.

GitOrigin-RevId: 71458b654b1c4a05799886928b449c948ea33109
This commit is contained in:
taohong-oai
2026-10-26 19:19:10 +00:00
committed by copyberry
parent 36021e1ad5
commit f371f9701f
31 changed files with 1436 additions and 106 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/tools/parallel.rs
+146 -6
View File
@@ -27,6 +27,9 @@ use crate::tools::registry::AnyToolResult;
use crate::tools::registry::ToolArgumentDiffConsumer;
use crate::tools::router::ToolCall;
use crate::tools::router::ToolCallSource;
use codex_direct_tool_timing::ToolAdmission;
use codex_direct_tool_timing::ToolBatchTiming;
use codex_direct_tool_timing::ToolGateTiming;
use codex_history::ResponseItemEnvelope;
use codex_protocol::error::CodexErr;
use codex_protocol::models::ResponseInputItem;
@@ -49,6 +52,25 @@ pub(crate) struct ToolCallRuntime {
step_context: Arc<StepContext>,
tracker: SharedTurnDiffTracker,
parallel_execution: Arc<RwLock<()>>,
timing: Option<DirectToolTiming>,
}
#[derive(Clone, Default)]
struct DirectToolTiming {
gate: ToolGateTiming,
batch: ToolBatchTiming,
}
// Remove this boundary when CodeMode and subagent dependencies have runtime capture.
fn direct_path_supported(tool_name: &codex_tools::ToolName, source: &ToolCallSource) -> bool {
!(matches!(source, ToolCallSource::CodeMode { .. })
|| (tool_name.is_default_namespace()
&& matches!(tool_name.name.as_str(),
crate::tools::code_mode::PUBLIC_TOOL_NAME | crate::tools::code_mode::WAIT_TOOL_NAME))
// Multi-agent tool namespaces can be overridden by the tool catalog.
|| matches!(tool_name.name.as_str(), "spawn_agent" | "send_input"
| "close_agent" | "resume_agent" | "wait_agent" | "list_agents"
| "followup_task" | "interrupt_agent" | "send_message"))
}
impl ToolCallRuntime {
@@ -57,11 +79,18 @@ impl ToolCallRuntime {
step_context: Arc<StepContext>,
tracker: SharedTurnDiffTracker,
) -> Self {
let timing = step_context
.turn
.config
.features
.enabled(codex_features::Feature::CriticalPathTiming)
.then(DirectToolTiming::default);
Self {
session,
step_context,
tracker,
parallel_execution: Arc::new(RwLock::new(())),
timing,
}
}
@@ -152,6 +181,35 @@ impl ToolCallRuntime {
let lock = Arc::clone(&self.parallel_execution);
let invocation_cancellation_token = cancellation_token.clone();
let started = Instant::now();
if turn
.config
.features
.enabled(codex_features::Feature::CriticalPathTiming)
&& !direct_path_supported(
&tool_runtime
.as_ref()
.map_or_else(|| call.tool_name.clone(), |runtime| runtime.tool_name()),
&source,
)
{
turn.turn_timing_state.invalidate_critical_path();
}
let work_timing = self.timing.and_then(|timing| match &source {
ToolCallSource::Direct | ToolCallSource::DirectPlaintextMessage => {
let work = turn
.turn_timing_state
.direct_tool_timings
.start_turn_direct_call(
&call.call_id,
&turn.sub_id,
timing.batch,
started,
turn.turn_timing_state.snapshot_path(started),
);
Some((timing.gate, work))
}
ToolCallSource::CodeMode { .. } => None,
});
let mut tool_call_timing_guard = ToolCallTimingGuard::capture(
started,
&session.thread_id,
@@ -208,7 +266,7 @@ impl ToolCallRuntime {
async move {
let _message_admission = message_admission?;
// Readiness and admission remain cancellable even when the handler must finish.
let guard = async {
let (guard, admission_started) = async {
if let Some(tool_runtime) = tool_runtime
&& let Some(readiness) = tool_runtime.wait_until_ready(&session)
{
@@ -217,7 +275,8 @@ impl ToolCallRuntime {
.await;
}
if supports_parallel {
let admission_started = Instant::now();
let guard = if supports_parallel {
Either::Left(
lock.read()
.instrument(trace_span!("tool_dispatch.parallel_lock"))
@@ -229,10 +288,24 @@ impl ToolCallRuntime {
.instrument(trace_span!("tool_dispatch.exclusive_lock"))
.await,
)
}
};
(guard, admission_started)
}
.or_cancel_tool(&invocation_cancellation_token)
.await?;
// Publish the completed path before releasing the admission lock,
// including when cancellation drops the handler future.
let _admitted_timing = work_timing.map(|(gate, work)| {
gate.admit(
work,
if supports_parallel {
ToolAdmission::Parallel
} else {
ToolAdmission::Exclusive
},
admission_started,
)
});
// Admission through the parallel-execution gate marks the end
// of dispatch waiting and the start of handler execution.
if let Some(execution_started_at) = execution_started_at {
@@ -250,6 +323,7 @@ impl ToolCallRuntime {
dispatch_call_state,
)
.await;
drop(_admitted_timing);
drop(guard);
// The sampling loop collects results in order only after its stream ends.
// Record readiness here, before either caller encodes or collects the result.
@@ -347,6 +421,10 @@ impl ToolCallRuntime {
}
impl ToolCallRuntime {
pub(crate) fn batch_timing(&self) -> Option<ToolBatchTiming> {
self.timing.as_ref().map(|timing| timing.batch.clone())
}
fn tool_task_join_error(err: JoinError) -> FunctionCallError {
FunctionCallError::Fatal(format!("tool task failed to receive: {err:?}"))
}
@@ -529,6 +607,40 @@ mod tests {
use tokio::sync::oneshot;
use tracing_test::internal::MockWriter;
#[test]
fn unsupported_handlers_remain_incomplete_after_namespace_resolution() {
let handler = crate::tools::multi_agent_tool::multi_agent_v2_handler(
crate::tools::handlers::multi_agents_v2::WaitAgentHandler::default(),
Some("custom_collaboration"),
/*description_override*/ None,
/*parameters_override*/ None,
);
let registry = ToolRegistry::from_tools([handler]);
let registered = registry
.registered_tool(&codex_tools::ToolName::namespaced(
"custom_collaboration",
"wait_agent",
))
.unwrap();
assert!(!direct_path_supported(
&registered.runtime.tool_name(),
&ToolCallSource::Direct
));
for name in [
crate::tools::code_mode::PUBLIC_TOOL_NAME,
crate::tools::code_mode::WAIT_TOOL_NAME,
] {
assert!(!direct_path_supported(
&codex_tools::ToolName::plain(name).with_default_namespace(),
&ToolCallSource::Direct
));
}
assert!(direct_path_supported(
&codex_tools::ToolName::namespaced("mcp", "exec"),
&ToolCallSource::Direct
));
}
#[test]
fn tool_call_timing_guard_ignores_code_mode_source() {
let buffer: &'static std::sync::Mutex<Vec<u8>> =
@@ -622,13 +734,34 @@ mod tests {
});
}
#[test_case::test_case(false; "ordinary_handler")]
#[test_case::test_case(true; "dynamic_handler")]
#[test_case::test_case(false, false; "ordinary_capture_disabled")]
#[test_case::test_case(false, true; "ordinary_capture_enabled")]
#[test_case::test_case(true, false; "dynamic_capture_disabled")]
#[test_case::test_case(true, true; "dynamic_capture_enabled")]
#[tokio::test]
async fn cancellation_before_dispatch_admission_logs_dispatch_only_timing(
dynamic: bool,
critical_path_timing: bool,
) -> anyhow::Result<()> {
let (session, turn_context) = crate::session::tests::make_session_and_context().await;
let (session, mut turn_context) = crate::session::tests::make_session_and_context().await;
Arc::make_mut(&mut turn_context.config)
.features
.enable(codex_features::Feature::RuntimeMetrics)
.unwrap();
Arc::make_mut(&mut turn_context.config)
.features
.set_enabled(
codex_features::Feature::CriticalPathTiming,
critical_path_timing,
)
.unwrap();
if critical_path_timing {
turn_context.turn_timing_state.enable_critical_path();
}
turn_context
.turn_timing_state
.mark_turn_started(Instant::now())
.await;
let session = Arc::new(session);
let turn_context = Arc::new(turn_context);
let tool_name = codex_tools::ToolName::plain("test_tool");
@@ -661,6 +794,8 @@ mod tests {
let step_context = step_context.with_tool_router_for_test(router);
let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new()));
let runtime = ToolCallRuntime::new(session, step_context, tracker);
let batch = runtime.batch_timing();
assert_eq!(batch.is_some(), critical_path_timing);
let execution_gate = Arc::clone(&runtime.parallel_execution);
let execution_gate_guard = execution_gate
.try_write_owned()
@@ -700,6 +835,11 @@ mod tests {
.expect("cancelled tool response task should join")
.expect("cancelled tool call should produce a response");
assert_eq!(
batch.and_then(|batch| batch.completed()).is_some(),
critical_path_timing,
"capture requires its own gate and must finish before cancellation returns"
);
let logs = String::from_utf8(
buffer
.lock()