From 7a3eec6fdb356bd71f80582119eb829179ff0da1 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Thu, 2 Apr 2026 16:39:56 -0700 Subject: [PATCH 1/2] core: cut codex-core compile time 48% with native async SessionTask (#16631) ## Why This continues the compile-time cleanup from #16630. `SessionTask` implementations are monomorphized, but `Session` stores the task behind a `dyn` boundary so it can drive and abort heterogenous turn tasks uniformly. That means we can move the `#[async_trait]` expansion off the implementation trait, keep a small boxed adapter only at the storage boundary, and preserve the existing task lifecycle semantics while reducing the amount of generated async-trait glue in `codex-core`. One measurement caveat showed up while exploring this: a warm incremental benchmark based on `touch core/src/tasks/mod.rs && cargo check -p codex-core --lib` was basically flat, but that was the wrong benchmark for this change. Using package-clean `codex-core` rebuilds, like #16630, shows the real win. Relevant pre-change code: - [`SessionTask` with `#[async_trait]`](https://github.com/openai/codex/blob/3c7f013f9735e67796c70d95f75f436b7f97e3ec/codex-rs/core/src/tasks/mod.rs#L129-L182) - [`RunningTask` storing `Arc`](https://github.com/openai/codex/blob/3c7f013f9735e67796c70d95f75f436b7f97e3ec/codex-rs/core/src/state/turn.rs#L69-L77) ## What changed - Switched `SessionTask::{run, abort}` to native RPITIT futures with explicit `Send` bounds. - Added a private `AnySessionTask` adapter that boxes those futures only at the `Arc` storage boundary. - Updated `RunningTask` to store `Arc` and removed `#[async_trait]` from the concrete task impls plus test-only `SessionTask` impls. ## Timing Benchmarked package-clean `codex-core` rebuilds with dependencies left warm: ```shell cargo check -p codex-core --lib >/dev/null cargo clean -p codex-core >/dev/null /usr/bin/time -p cargo +nightly rustc -p codex-core --lib -- \ -Z time-passes \ -Z time-passes-format=json >/dev/null ``` | revision | rustc `total` | process `real` | `generate_crate_metadata` | `MIR_borrow_checking` | `monomorphization_collector_graph_walk` | | --- | ---: | ---: | ---: | ---: | ---: | | parent `3c7f013f9735` | 67.21s | 67.71s | 24.61s | 23.43s | 22.43s | | this PR `2cafd783ac22` | 35.08s | 35.60s | 8.01s | 7.25s | 7.15s | | delta | -47.8% | -47.4% | -67.5% | -69.1% | -68.1% | For completeness, the warm touched-file benchmark stayed flat (`1.96s` parent vs `1.97s` this PR), which is why that benchmark should not be used to evaluate this refactor. ## Verification - Ran `cargo test -p codex-core`; this change compiled and task-related tests passed before hitting the same unrelated 5 `config::tests::*guardian*` failures already present on the parent stack. --- codex-rs/core/src/codex_tests.rs | 2 - codex-rs/core/src/state/turn.rs | 4 +- codex-rs/core/src/tasks/compact.rs | 6 +- codex-rs/core/src/tasks/ghost_snapshot.rs | 2 - codex-rs/core/src/tasks/mod.rs | 76 +++++++++++++++++-- codex-rs/core/src/tasks/regular.rs | 2 - codex-rs/core/src/tasks/review.rs | 4 +- codex-rs/core/src/tasks/undo.rs | 12 ++- codex-rs/core/src/tasks/user_shell.rs | 2 - .../src/tools/handlers/multi_agents_tests.rs | 1 - 10 files changed, 79 insertions(+), 32 deletions(-) diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index 4ff31e8bd0..c939120739 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/codex-rs/core/src/codex_tests.rs @@ -3119,7 +3119,6 @@ async fn spawn_task_turn_span_inherits_dispatch_trace_context() { captured_trace: Arc>>, } - #[async_trait::async_trait] impl SessionTask for TraceCaptureTask { fn kind(&self) -> TaskKind { TaskKind::Regular @@ -4375,7 +4374,6 @@ struct NeverEndingTask { listen_to_cancellation_token: bool, } -#[async_trait::async_trait] impl SessionTask for NeverEndingTask { fn kind(&self) -> TaskKind { self.kind diff --git a/codex-rs/core/src/state/turn.rs b/codex-rs/core/src/state/turn.rs index a8e3e167b5..214fd8be15 100644 --- a/codex-rs/core/src/state/turn.rs +++ b/codex-rs/core/src/state/turn.rs @@ -18,7 +18,7 @@ use rmcp::model::RequestId; use tokio::sync::oneshot; use crate::codex::TurnContext; -use crate::tasks::SessionTask; +use crate::tasks::AnySessionTask; use codex_protocol::models::PermissionProfile; use codex_protocol::protocol::ReviewDecision; use codex_protocol::protocol::TokenUsage; @@ -69,7 +69,7 @@ pub(crate) enum TaskKind { pub(crate) struct RunningTask { pub(crate) done: Arc, pub(crate) kind: TaskKind, - pub(crate) task: Arc, + pub(crate) task: Arc, pub(crate) cancellation_token: CancellationToken, pub(crate) handle: Arc>, pub(crate) turn_context: Arc, diff --git a/codex-rs/core/src/tasks/compact.rs b/codex-rs/core/src/tasks/compact.rs index a2d94bdc0a..8c7998d853 100644 --- a/codex-rs/core/src/tasks/compact.rs +++ b/codex-rs/core/src/tasks/compact.rs @@ -4,14 +4,12 @@ use super::SessionTask; use super::SessionTaskContext; use crate::codex::TurnContext; use crate::state::TaskKind; -use async_trait::async_trait; use codex_protocol::user_input::UserInput; use tokio_util::sync::CancellationToken; #[derive(Clone, Copy, Default)] pub(crate) struct CompactTask; -#[async_trait] impl SessionTask for CompactTask { fn kind(&self) -> TaskKind { TaskKind::Compact @@ -30,14 +28,14 @@ impl SessionTask for CompactTask { ) -> Option { let session = session.clone_session(); let _ = if crate::compact::should_use_remote_compact_task(&ctx.provider) { - let _ = session.services.session_telemetry.counter( + session.services.session_telemetry.counter( "codex.task.compact", /*inc*/ 1, &[("type", "remote")], ); crate::compact_remote::run_remote_compact_task(session.clone(), ctx).await } else { - let _ = session.services.session_telemetry.counter( + session.services.session_telemetry.counter( "codex.task.compact", /*inc*/ 1, &[("type", "local")], diff --git a/codex-rs/core/src/tasks/ghost_snapshot.rs b/codex-rs/core/src/tasks/ghost_snapshot.rs index 7b848098a4..a9cf19cb4d 100644 --- a/codex-rs/core/src/tasks/ghost_snapshot.rs +++ b/codex-rs/core/src/tasks/ghost_snapshot.rs @@ -2,7 +2,6 @@ use crate::codex::TurnContext; use crate::state::TaskKind; use crate::tasks::SessionTask; use crate::tasks::SessionTaskContext; -use async_trait::async_trait; use codex_git_utils::CreateGhostCommitOptions; use codex_git_utils::GhostSnapshotReport; use codex_git_utils::GitToolingError; @@ -26,7 +25,6 @@ pub(crate) struct GhostSnapshotTask { const SNAPSHOT_WARNING_THRESHOLD: Duration = Duration::from_secs(240); -#[async_trait] impl SessionTask for GhostSnapshotTask { fn kind(&self) -> TaskKind { TaskKind::Regular diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 74b2e471cd..c0b5ad4d91 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -9,7 +9,7 @@ use std::sync::Arc; use std::time::Duration; use std::time::Instant; -use async_trait::async_trait; +use futures::future::BoxFuture; use tokio::select; use tokio::sync::Notify; use tokio_util::sync::CancellationToken; @@ -126,7 +126,6 @@ impl SessionTaskContext { /// intentionally small: implementers identify themselves via /// [`SessionTask::kind`], perform their work in [`SessionTask::run`], and may /// release resources in [`SessionTask::abort`]. -#[async_trait] pub(crate) trait SessionTask: Send + Sync + 'static { /// Describes the type of work the task performs so the session can /// surface it in telemetry and UI. @@ -143,21 +142,84 @@ pub(crate) trait SessionTask: Send + Sync + 'static { /// abort; implementers should watch for it and terminate quickly once it /// fires. Returning [`Some`] yields a final message that /// [`Session::on_task_finished`] will emit to the client. - async fn run( + fn run( self: Arc, session: Arc, ctx: Arc, input: Vec, cancellation_token: CancellationToken, - ) -> Option; + ) -> impl std::future::Future> + Send; /// Gives the task a chance to perform cleanup after an abort. /// /// The default implementation is a no-op; override this if additional /// teardown or notifications are required once /// [`Session::abort_all_tasks`] cancels the task. - async fn abort(&self, session: Arc, ctx: Arc) { - let _ = (session, ctx); + fn abort( + &self, + session: Arc, + ctx: Arc, + ) -> impl std::future::Future + Send { + async move { + let _ = (session, ctx); + } + } +} + +pub(crate) trait AnySessionTask: Send + Sync + 'static { + fn kind(&self) -> TaskKind; + + fn span_name(&self) -> &'static str; + + fn run( + self: Arc, + session: Arc, + ctx: Arc, + input: Vec, + cancellation_token: CancellationToken, + ) -> BoxFuture<'static, Option>; + + fn abort<'a>( + &'a self, + session: Arc, + ctx: Arc, + ) -> BoxFuture<'a, ()>; +} + +impl AnySessionTask for T +where + T: SessionTask, +{ + fn kind(&self) -> TaskKind { + SessionTask::kind(self) + } + + fn span_name(&self) -> &'static str { + SessionTask::span_name(self) + } + + fn run( + self: Arc, + session: Arc, + ctx: Arc, + input: Vec, + cancellation_token: CancellationToken, + ) -> BoxFuture<'static, Option> { + Box::pin(SessionTask::run( + self, + session, + ctx, + input, + cancellation_token, + )) + } + + fn abort<'a>( + &'a self, + session: Arc, + ctx: Arc, + ) -> BoxFuture<'a, ()> { + Box::pin(SessionTask::abort(self, session, ctx)) } } @@ -179,7 +241,7 @@ impl Session { input: Vec, task: T, ) { - let task: Arc = Arc::new(task); + let task: Arc = Arc::new(task); let task_kind = task.kind(); let span_name = task.span_name(); let started_at = Instant::now(); diff --git a/codex-rs/core/src/tasks/regular.rs b/codex-rs/core/src/tasks/regular.rs index 7a274d534f..f2a29ee7ab 100644 --- a/codex-rs/core/src/tasks/regular.rs +++ b/codex-rs/core/src/tasks/regular.rs @@ -1,6 +1,5 @@ use std::sync::Arc; -use async_trait::async_trait; use tokio_util::sync::CancellationToken; use crate::codex::TurnContext; @@ -25,7 +24,6 @@ impl RegularTask { } } -#[async_trait] impl SessionTask for RegularTask { fn kind(&self) -> TaskKind { TaskKind::Regular diff --git a/codex-rs/core/src/tasks/review.rs b/codex-rs/core/src/tasks/review.rs index e0c6033483..a1cc071108 100644 --- a/codex-rs/core/src/tasks/review.rs +++ b/codex-rs/core/src/tasks/review.rs @@ -1,7 +1,6 @@ use std::borrow::Cow; use std::sync::Arc; -use async_trait::async_trait; use codex_protocol::config_types::WebSearchMode; use codex_protocol::items::TurnItem; use codex_protocol::models::ContentItem; @@ -48,7 +47,6 @@ impl ReviewTask { } } -#[async_trait] impl SessionTask for ReviewTask { fn kind(&self) -> TaskKind { TaskKind::Review @@ -65,7 +63,7 @@ impl SessionTask for ReviewTask { input: Vec, cancellation_token: CancellationToken, ) -> Option { - let _ = session.session.services.session_telemetry.counter( + session.session.services.session_telemetry.counter( "codex.task.review", /*inc*/ 1, &[], diff --git a/codex-rs/core/src/tasks/undo.rs b/codex-rs/core/src/tasks/undo.rs index 48cdf11aab..dd65530092 100644 --- a/codex-rs/core/src/tasks/undo.rs +++ b/codex-rs/core/src/tasks/undo.rs @@ -4,7 +4,6 @@ use crate::codex::TurnContext; use crate::state::TaskKind; use crate::tasks::SessionTask; use crate::tasks::SessionTaskContext; -use async_trait::async_trait; use codex_git_utils::RestoreGhostCommitOptions; use codex_git_utils::restore_ghost_commit_with_options; use codex_protocol::models::ResponseItem; @@ -25,7 +24,6 @@ impl UndoTask { } } -#[async_trait] impl SessionTask for UndoTask { fn kind(&self) -> TaskKind { TaskKind::Regular @@ -42,11 +40,11 @@ impl SessionTask for UndoTask { _input: Vec, cancellation_token: CancellationToken, ) -> Option { - let _ = session.session.services.session_telemetry.counter( - "codex.task.undo", - /*inc*/ 1, - &[], - ); + session + .session + .services + .session_telemetry + .counter("codex.task.undo", /*inc*/ 1, &[]); let sess = session.clone_session(); sess.send_event( ctx.as_ref(), diff --git a/codex-rs/core/src/tasks/user_shell.rs b/codex-rs/core/src/tasks/user_shell.rs index 4d4ff383ac..449db837fd 100644 --- a/codex-rs/core/src/tasks/user_shell.rs +++ b/codex-rs/core/src/tasks/user_shell.rs @@ -1,7 +1,6 @@ use std::sync::Arc; use std::time::Duration; -use async_trait::async_trait; use codex_async_utils::CancelErr; use codex_async_utils::OrCancelExt; use codex_protocol::user_input::UserInput; @@ -62,7 +61,6 @@ impl UserShellCommandTask { } } -#[async_trait] impl SessionTask for UserShellCommandTask { fn kind(&self) -> TaskKind { TaskKind::Regular diff --git a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs index ee81b961e3..fa231e66b8 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -115,7 +115,6 @@ fn history_contains_inter_agent_communication( #[derive(Clone, Copy)] struct NeverEndingTask; -#[async_trait::async_trait] impl SessionTask for NeverEndingTask { fn kind(&self) -> TaskKind { TaskKind::Regular From 083b908813a98df92040d23c179f9cd5dd4baafa Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Thu, 2 Apr 2026 17:15:41 -0700 Subject: [PATCH 2/2] app-server: make thread shell command tests shell-aware --- .../tests/suite/v2/thread_shell_command.rs | 34 +++++++++++++++---- 1 file changed, 27 insertions(+), 7 deletions(-) diff --git a/codex-rs/app-server/tests/suite/v2/thread_shell_command.rs b/codex-rs/app-server/tests/suite/v2/thread_shell_command.rs index 4e8a40fca2..a2f4b3a2bf 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_shell_command.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_shell_command.rs @@ -26,6 +26,7 @@ use codex_app_server_protocol::TurnCompletedNotification; use codex_app_server_protocol::TurnStartParams; use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::UserInput as V2UserInput; +use codex_core::shell::default_user_shell; use codex_features::FEATURES; use codex_features::Feature; use pretty_assertions::assert_eq; @@ -67,11 +68,12 @@ async fn thread_shell_command_runs_as_standalone_turn_and_persists_history() -> ) .await??; let ThreadStartResponse { thread, .. } = to_response::(start_resp)?; + let (shell_command, expected_output) = current_shell_output_command("hello from bang"); let shell_id = mcp .send_thread_shell_command_request(ThreadShellCommandParams { thread_id: thread.id.clone(), - command: "printf 'hello from bang\\n'".to_string(), + command: shell_command, }) .await?; let shell_resp: JSONRPCResponse = timeout( @@ -93,7 +95,7 @@ async fn thread_shell_command_runs_as_standalone_turn_and_persists_history() -> assert_eq!(status, &CommandExecutionStatus::InProgress); let delta = wait_for_command_execution_output_delta(&mut mcp, &command_id).await?; - assert_eq!(delta.delta, "hello from bang\n"); + assert_eq!(delta.delta, expected_output); let completed = wait_for_command_execution_completed(&mut mcp, Some(&command_id)).await?; let ThreadItem::CommandExecution { @@ -110,7 +112,7 @@ async fn thread_shell_command_runs_as_standalone_turn_and_persists_history() -> assert_eq!(id, &command_id); assert_eq!(source, &CommandExecutionSource::UserShell); assert_eq!(status, &CommandExecutionStatus::Completed); - assert_eq!(aggregated_output.as_deref(), Some("hello from bang\n")); + assert_eq!(aggregated_output.as_deref(), Some(expected_output.as_str())); assert_eq!(*exit_code, Some(0)); timeout( @@ -147,7 +149,7 @@ async fn thread_shell_command_runs_as_standalone_turn_and_persists_history() -> }; assert_eq!(source, &CommandExecutionSource::UserShell); assert_eq!(status, &CommandExecutionStatus::Completed); - assert_eq!(aggregated_output.as_deref(), Some("hello from bang\n")); + assert_eq!(aggregated_output.as_deref(), Some(expected_output.as_str())); Ok(()) } @@ -196,6 +198,7 @@ async fn thread_shell_command_uses_existing_active_turn() -> Result<()> { ) .await??; let ThreadStartResponse { thread, .. } = to_response::(start_resp)?; + let (shell_command, expected_output) = current_shell_output_command("active turn bang"); let turn_id = mcp .send_turn_start_request(TurnStartParams { @@ -240,7 +243,7 @@ async fn thread_shell_command_uses_existing_active_turn() -> Result<()> { let shell_id = mcp .send_thread_shell_command_request(ThreadShellCommandParams { thread_id: thread.id.clone(), - command: "printf 'active turn bang\\n'".to_string(), + command: shell_command, }) .await?; let shell_resp: JSONRPCResponse = timeout( @@ -269,7 +272,7 @@ async fn thread_shell_command_uses_existing_active_turn() -> Result<()> { unreachable!("helper returns command execution item"); }; assert_eq!(source, &CommandExecutionSource::UserShell); - assert_eq!(aggregated_output.as_deref(), Some("active turn bang\n")); + assert_eq!(aggregated_output.as_deref(), Some(expected_output.as_str())); mcp.send_response( request_id, @@ -309,7 +312,7 @@ async fn thread_shell_command_uses_existing_active_turn() -> Result<()> { source: CommandExecutionSource::UserShell, aggregated_output, .. - } if aggregated_output.as_deref() == Some("active turn bang\n") + } if aggregated_output.as_deref() == Some(expected_output.as_str()) ) }), "expected active-turn shell command to be persisted on the existing turn" @@ -318,6 +321,23 @@ async fn thread_shell_command_uses_existing_active_turn() -> Result<()> { Ok(()) } +fn current_shell_output_command(text: &str) -> (String, String) { + match default_user_shell().name() { + "powershell" => { + let escaped_text = text.replace('\'', "''"); + ( + format!("Write-Output '{escaped_text}'"), + format!("{text}\r\n"), + ) + } + "cmd" => (format!("echo {text}"), format!("{text}\r\n")), + _ => { + let quoted_text = shlex::try_quote(text).expect("test output text should be quotable"); + (format!("printf '%s\\n' {quoted_text}"), format!("{text}\n")) + } + } +} + async fn wait_for_command_execution_started( mcp: &mut McpProcess, expected_id: Option<&str>,