From 35eab50501746159d89fafd8cf1f048b00cf5174 Mon Sep 17 00:00:00 2001 From: jif Date: Fri, 31 Jul 2026 15:59:04 +0000 Subject: [PATCH] Pass sessions directly to session tasks (#36354) ## What changed - Replace `SessionTaskContext` with `Arc` in the session task run and abort interfaces. - Read extension data from `TurnContext` during regular turns instead of storing and forwarding a duplicate reference through the running task. GitOrigin-RevId: 04bc0d324e5ba781f35ac9c61e811888208b74ae --- codex-rs/core/src/session/tests.rs | 10 ++-- codex-rs/core/src/session/turn.rs | 3 +- codex-rs/core/src/state/turn.rs | 2 - codex-rs/core/src/tasks/compact.rs | 5 +- codex-rs/core/src/tasks/mod.rs | 72 ++++----------------------- codex-rs/core/src/tasks/regular.rs | 7 +-- codex-rs/core/src/tasks/review.rs | 38 ++++++-------- codex-rs/core/src/tasks/user_shell.rs | 5 +- 8 files changed, 36 insertions(+), 106 deletions(-) diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index ad6c05ed41..8ee16e14f4 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -84,7 +84,6 @@ use crate::rollout::recorder::RolloutRecorder; use crate::state::ActiveTurn; use crate::state::TaskKind; use crate::tasks::SessionTask; -use crate::tasks::SessionTaskContext; use crate::tasks::SessionTaskResult; use crate::tasks::UserShellCommandMode; use crate::tasks::execute_user_shell_command; @@ -7253,7 +7252,7 @@ async fn spawn_task_turn_span_inherits_dispatch_trace_context() { async fn run( self: Arc, - _session: Arc, + _session: Arc, _ctx: Arc, _input: Vec, _cancellation_token: CancellationToken, @@ -9829,7 +9828,7 @@ impl SessionTask for CompletingTask { async fn run( self: Arc, - _session: Arc, + _session: Arc, _ctx: Arc, _input: Vec, _cancellation_token: CancellationToken, @@ -9977,7 +9976,7 @@ impl SessionTask for NeverEndingTask { async fn run( self: Arc, - _session: Arc, + _session: Arc, _ctx: Arc, _input: Vec, cancellation_token: CancellationToken, @@ -10006,12 +10005,11 @@ impl SessionTask for GuardianDeniedApprovalTask { async fn run( self: Arc, - session: Arc, + session: Arc, ctx: Arc, _input: Vec, cancellation_token: CancellationToken, ) -> SessionTaskResult { - let session = session.clone_session(); for _ in 0..3 { crate::guardian::record_guardian_denial_for_test(&session, &ctx, &ctx.sub_id).await; } diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 09358b82a8..ec93b4db6b 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -150,7 +150,6 @@ const POST_SAMPLING_TOKEN_ESTIMATE_TARGET: &str = "codex_core::post_sampling_tok pub(crate) async fn run_turn( sess: Arc, turn_context: Arc, - turn_extension_data: Arc, input: Vec, prewarmed_client_session: Option, cancellation_token: CancellationToken, @@ -345,7 +344,7 @@ pub(crate) async fn run_turn( run_sampling_request( Arc::clone(&sess), Arc::clone(&step_context), - Arc::clone(&turn_extension_data), + Arc::clone(&turn_context.extension_data), Arc::clone(&turn_diff_tracker), &mut client_session, &responses_metadata, diff --git a/codex-rs/core/src/state/turn.rs b/codex-rs/core/src/state/turn.rs index b74005f208..2a18f72d02 100644 --- a/codex-rs/core/src/state/turn.rs +++ b/codex-rs/core/src/state/turn.rs @@ -7,7 +7,6 @@ use tokio::sync::Notify; use tokio_util::sync::CancellationToken; use tokio_util::task::AbortOnDropHandle; -use codex_extension_api::ExtensionData; use codex_protocol::dynamic_tools::DynamicToolResponse; use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::request_permissions::RequestPermissionProfile; @@ -77,7 +76,6 @@ pub(crate) struct RunningTask { pub(crate) cancellation_token: CancellationToken, pub(crate) handle: AbortOnDropHandle<()>, pub(crate) turn_context: Arc, - pub(crate) turn_extension_data: Arc, pub(crate) _agent_execution_guard: Option, // Timer recorded when the task drops to capture the full turn duration. pub(crate) _timer: Option, diff --git a/codex-rs/core/src/tasks/compact.rs b/codex-rs/core/src/tasks/compact.rs index def729b2ac..18c9a27a5b 100644 --- a/codex-rs/core/src/tasks/compact.rs +++ b/codex-rs/core/src/tasks/compact.rs @@ -1,10 +1,10 @@ use std::sync::Arc; use super::SessionTask; -use super::SessionTaskContext; use super::SessionTaskResult; use super::emit_compact_metric; use crate::session::TurnInput; +use crate::session::session::Session; use crate::session::turn_context::TurnContext; use crate::state::TaskKind; use codex_features::Feature; @@ -26,12 +26,11 @@ impl SessionTask for CompactTask { async fn run( self: Arc, - session: Arc, + session: Arc, ctx: Arc, _input: Vec, _cancellation_token: CancellationToken, ) -> SessionTaskResult { - let session = session.clone_session(); let _profile_guard = ctx.turn_timing_state.begin_compaction(); if ctx.config.features.enabled(Feature::TokenBudget) { crate::compact_token_budget::run_manual_compact_task(session, ctx).await?; diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 83ca4512a4..e065686b46 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -8,7 +8,6 @@ use std::sync::Arc; use std::time::Duration; use std::time::Instant; -use codex_extension_api::ExtensionData; use futures::future::BoxFuture; use tokio::select; use tokio::sync::Notify; @@ -36,8 +35,6 @@ use crate::state::RunningTask; use crate::state::TaskKind; use codex_analytics::TurnProfileFact; use codex_analytics::TurnTokenUsageFact; -use codex_login::AuthManager; -use codex_models_manager::manager::SharedModelsManager; use codex_otel::SessionTelemetry; use codex_otel::TURN_E2E_DURATION_METRIC; use codex_otel::TURN_MEMORY_METRIC; @@ -176,38 +173,6 @@ fn bool_tag(value: bool) -> &'static str { if value { "true" } else { "false" } } -/// Thin wrapper that exposes the parts of [`Session`] task runners need. -#[derive(Clone)] -pub(crate) struct SessionTaskContext { - session: Arc, - turn_extension_data: Arc, -} - -impl SessionTaskContext { - pub(crate) fn new(session: Arc, turn_extension_data: Arc) -> Self { - Self { - session, - turn_extension_data, - } - } - - pub(crate) fn clone_session(&self) -> Arc { - Arc::clone(&self.session) - } - - pub(crate) fn turn_extension_data(&self) -> Arc { - Arc::clone(&self.turn_extension_data) - } - - pub(crate) fn auth_manager(&self) -> Arc { - Arc::clone(&self.session.services.auth_manager) - } - - pub(crate) fn models_manager(&self) -> SharedModelsManager { - Arc::clone(&self.session.services.models_manager) - } -} - /// Async task that drives a [`Session`] turn. /// /// Implementations encapsulate a specific Codex workflow (regular chat, @@ -236,7 +201,7 @@ pub(crate) trait SessionTask: Send + Sync + 'static { /// lifecycle instead. fn run( self: Arc, - session: Arc, + session: Arc, ctx: Arc, input: Vec, cancellation_token: CancellationToken, @@ -249,7 +214,7 @@ pub(crate) trait SessionTask: Send + Sync + 'static { /// [`Session::abort_all_tasks`] cancels the task. fn abort( &self, - session: Arc, + session: Arc, ctx: Arc, ) -> impl std::future::Future + Send { async move { @@ -265,17 +230,13 @@ pub(crate) trait AnySessionTask: Send + Sync + 'static { fn run( self: Arc, - session: Arc, + session: Arc, ctx: Arc, input: Vec, cancellation_token: CancellationToken, ) -> BoxFuture<'static, SessionTaskResult>; - fn abort<'a>( - &'a self, - session: Arc, - ctx: Arc, - ) -> BoxFuture<'a, ()>; + fn abort<'a>(&'a self, session: Arc, ctx: Arc) -> BoxFuture<'a, ()>; } impl AnySessionTask for T @@ -292,7 +253,7 @@ where fn run( self: Arc, - session: Arc, + session: Arc, ctx: Arc, input: Vec, cancellation_token: CancellationToken, @@ -306,11 +267,7 @@ where )) } - fn abort<'a>( - &'a self, - session: Arc, - ctx: Arc, - ) -> BoxFuture<'a, ()> { + fn abort<'a>(&'a self, session: Arc, ctx: Arc) -> BoxFuture<'a, ()> { Box::pin(SessionTask::abort(self, session, ctx)) } } @@ -377,7 +334,6 @@ impl Session { self.emit_turn_start_lifecycle(turn_context.as_ref(), &token_usage_at_turn_start) .await; - let turn_extension_data = Arc::clone(&turn_context.extension_data); let mut active = self.active_turn.lock().await; let turn = active.get_or_insert_with(ActiveTurn::default); debug_assert!(turn.task.is_none()); @@ -386,10 +342,7 @@ impl Session { &turn_context.session_source, ); let done_clone = Arc::clone(&done); - let session_ctx = Arc::new(SessionTaskContext::new( - Arc::clone(self), - Arc::clone(&turn_extension_data), - )); + let session = Arc::clone(self); let ctx = Arc::clone(&turn_context); let task_for_run = Arc::clone(&task); let task_input = input; @@ -417,14 +370,14 @@ impl Session { let ctx_for_finish = Arc::clone(&ctx); let task_result = task_for_run .run( - Arc::clone(&session_ctx), + Arc::clone(&session), ctx, task_input, task_cancellation_token.child_token(), ) .instrument(trace_span!("session_task.run")) .await; - let sess = session_ctx.clone_session(); + let sess = Arc::clone(&session); if let Err(err) = sess.flush_rollout().await { warn!("failed to flush rollout before completing turn: {err}"); sess.send_event( @@ -457,7 +410,6 @@ impl Session { task, cancellation_token, turn_context: Arc::clone(&turn_context), - turn_extension_data, _agent_execution_guard: agent_execution_guard, _timer: timer, }; @@ -917,12 +869,8 @@ impl Session { task.handle.abort(); - let session_ctx = Arc::new(SessionTaskContext::new( - Arc::clone(self), - Arc::clone(&task.turn_extension_data), - )); session_task - .abort(session_ctx, Arc::clone(&task.turn_context)) + .abort(Arc::clone(self), Arc::clone(&task.turn_context)) .await; if reason == TurnAbortReason::Interrupted diff --git a/codex-rs/core/src/tasks/regular.rs b/codex-rs/core/src/tasks/regular.rs index 3aa39bee90..b88bd51b03 100644 --- a/codex-rs/core/src/tasks/regular.rs +++ b/codex-rs/core/src/tasks/regular.rs @@ -3,6 +3,7 @@ use std::sync::Arc; use tokio_util::sync::CancellationToken; use crate::session::TurnInput; +use crate::session::session::Session; use crate::session::turn::run_hooks_and_record_inputs; use crate::session::turn::run_turn; use crate::session::turn_context::TurnContext; @@ -14,7 +15,6 @@ use tracing::Instrument; use tracing::trace_span; use super::SessionTask; -use super::SessionTaskContext; use super::SessionTaskResult; #[derive(Default)] @@ -37,13 +37,11 @@ impl SessionTask for RegularTask { async fn run( self: Arc, - session: Arc, + sess: Arc, ctx: Arc, input: Vec, cancellation_token: CancellationToken, ) -> SessionTaskResult { - let sess = session.clone_session(); - let turn_extension_data = session.turn_extension_data(); let run_turn_span = trace_span!("run_turn"); // Regular turns emit `TurnStarted` inline so first-turn lifecycle does // not wait on startup prewarm resolution. @@ -78,7 +76,6 @@ impl SessionTask for RegularTask { let last_agent_message = run_turn( Arc::clone(&sess), Arc::clone(&ctx), - Arc::clone(&turn_extension_data), next_input, prewarmed_client_session.take(), cancellation_token.child_token(), diff --git a/codex-rs/core/src/tasks/review.rs b/codex-rs/core/src/tasks/review.rs index 0e2e2712e0..4a0c1c6855 100644 --- a/codex-rs/core/src/tasks/review.rs +++ b/codex-rs/core/src/tasks/review.rs @@ -29,7 +29,6 @@ use codex_features::Feature; use codex_protocol::user_input::UserInput; use super::SessionTask; -use super::SessionTaskContext; use super::SessionTaskResult; #[derive(Clone, Copy)] @@ -52,16 +51,15 @@ impl SessionTask for ReviewTask { async fn run( self: Arc, - session: Arc, + session: Arc, ctx: Arc, input: Vec, cancellation_token: CancellationToken, ) -> SessionTaskResult { - session.session.services.session_telemetry.counter( - "codex.task.review", - /*inc*/ 1, - &[], - ); + session + .services + .session_telemetry + .counter("codex.task.review", /*inc*/ 1, &[]); let mut user_input = Vec::new(); for item in input { @@ -84,18 +82,18 @@ impl SessionTask for ReviewTask { None => None, }; if !cancellation_token.is_cancelled() { - exit_review_mode(session.clone_session(), output.clone(), ctx.clone()).await; + exit_review_mode(Arc::clone(&session), output.clone(), ctx.clone()).await; } Ok(None) } - async fn abort(&self, session: Arc, ctx: Arc) { - exit_review_mode(session.clone_session(), /*review_output*/ None, ctx).await; + async fn abort(&self, session: Arc, ctx: Arc) { + exit_review_mode(session, /*review_output*/ None, ctx).await; } } async fn start_review_conversation( - session: Arc, + session: Arc, ctx: Arc, input: Vec, cancellation_token: CancellationToken, @@ -124,10 +122,10 @@ async fn start_review_conversation( sub_agent_config.model = Some(model); (run_codex_thread_one_shot( sub_agent_config, - session.auth_manager(), - session.models_manager(), + Arc::clone(&session.services.auth_manager), + Arc::clone(&session.services.models_manager), input, - session.clone_session(), + Arc::clone(&session), ctx.clone(), cancellation_token, SubAgentSource::Review, @@ -140,7 +138,7 @@ async fn start_review_conversation( } async fn process_review_events( - session: Arc, + session: Arc, ctx: Arc, receiver: async_channel::Receiver, ) -> Option { @@ -149,10 +147,7 @@ async fn process_review_events( match event.clone().msg { EventMsg::AgentMessage(_) => { if let Some(prev) = prev_agent_message.take() { - session - .clone_session() - .send_event(ctx.as_ref(), prev.msg) - .await; + session.send_event(ctx.as_ref(), prev.msg).await; } prev_agent_message = Some(event); } @@ -177,10 +172,7 @@ async fn process_review_events( return None; } other => { - session - .clone_session() - .send_event(ctx.as_ref(), other) - .await; + session.send_event(ctx.as_ref(), other).await; } } } diff --git a/codex-rs/core/src/tasks/user_shell.rs b/codex-rs/core/src/tasks/user_shell.rs index 3b5f352b0a..c8bdd5aa5b 100644 --- a/codex-rs/core/src/tasks/user_shell.rs +++ b/codex-rs/core/src/tasks/user_shell.rs @@ -39,7 +39,6 @@ use codex_sandboxing::SandboxType; use codex_shell_command::parse_command::parse_command; use super::SessionTask; -use super::SessionTaskContext; use super::SessionTaskResult; use crate::session::session::Session; use codex_protocol::models::PermissionProfile; @@ -78,13 +77,13 @@ impl SessionTask for UserShellCommandTask { async fn run( self: Arc, - session: Arc, + session: Arc, turn_context: Arc, _input: Vec, cancellation_token: CancellationToken, ) -> SessionTaskResult { execute_user_shell_command( - session.clone_session(), + session, turn_context, self.command.clone(), cancellation_token,