Preserve originating context for yielded code-mode tool calls (#44866)

## Why

A code-mode cell can outlive its turn. Nested tool calls dispatched during a later turn need the original step's tools and context, while approval handling must still work across that turn boundary.

## What changed

- Retain each cell's originating `StepContext` for nested tool dispatch and cancel queued calls whose context has been released.
- Keep MCP approval metadata available for the lifetime of the running invocation, keyed by server and call ID.
- Apply Guardian denial accounting and interruption to the active turn, including reviews from an earlier cell.
- Carry `turn_id` in legacy MCP begin/end events so history reconstruction and rollout migration assign late completions to their original turn. Preserve the current-turn fallback for older records without it.

## Testing

Add regression coverage for originating tool selection across turns, callback and metadata cleanup, delayed MCP elicitation review, Guardian interruption of the active turn, and late MCP completion attribution with legacy compatibility.

GitOrigin-RevId: 4533379aae6e4ecb4aabe5a25e4aca5d65385d76

codex-rs/app-server-protocol/src/protocol/thread_history.rs

@@ -788,7 +788,11 @@ impl ThreadHistoryBuilder {
             error: None,
             duration_ms: None,
         };
-        self.upsert_item_in_current_turn(item);
+        if payload.turn_id.is_empty() {
+            self.upsert_item_in_current_turn(item);
+        } else {
+            self.upsert_item_in_turn_id(&payload.turn_id, item);
+        }
     }
 
     fn handle_mcp_tool_call_end(&mut self, payload: &McpToolCallEndEvent) {
@@ -841,7 +845,11 @@ impl ThreadHistoryBuilder {
             error,
             duration_ms,
         };
-        self.upsert_item_in_current_turn(item);
+        if payload.turn_id.is_empty() {
+            self.upsert_item_in_current_turn(item);
+        } else {
+            self.upsert_item_in_turn_id(&payload.turn_id, item);
+        }
     }
 
     fn handle_view_image_tool_call(&mut self, payload: &ViewImageToolCallEvent) {
@@ -2998,6 +3006,7 @@ mod tests {
                 status: CoreExecCommandStatus::Completed,
             }),
             EventMsg::McpToolCallEnd(McpToolCallEndEvent {
+                turn_id: String::new(),
                 call_id: "mcp-1".into(),
                 invocation: McpInvocation {
                     server: "docs".into(),
@@ -3091,6 +3100,7 @@ mod tests {
                 collaboration_mode_kind: Default::default(),
             }),
             EventMsg::McpToolCallEnd(McpToolCallEndEvent {
+                turn_id: String::new(),
                 call_id: "mcp-1".into(),
                 invocation: McpInvocation {
                     server: "docs".into(),
@@ -3508,6 +3518,99 @@ mod tests {
         );
     }
 
+    #[test]
+    fn assigns_late_legacy_mcp_completion_to_original_turn() {
+        use codex_protocol::items::McpToolCallItem;
+        use codex_protocol::items::McpToolCallStatus as CoreMcpToolCallStatus;
+        use codex_protocol::protocol::HasLegacyEvent;
+
+        let completion = ItemCompletedEvent {
+            thread_id: ThreadId::new(),
+            turn_id: "turn-a".into(),
+            started_at_ms: None,
+            completed_at_ms: 0,
+            item: CoreTurnItem::McpToolCall(McpToolCallItem {
+                id: "mcp-1".into(),
+                server: "slack".into(),
+                tool: "search".into(),
+                arguments: serde_json::json!({"query": "incident"}),
+                connector_id: None,
+                mcp_app_resource_uri: None,
+                link_id: None,
+                app_name: None,
+                action_name: None,
+                plugin_id: None,
+                read_only_hint: Some(true),
+                status: CoreMcpToolCallStatus::Completed,
+                result: Some(CallToolResult {
+                    content: vec![serde_json::json!({"type": "text", "text": "Found incident"})],
+                    structured_content: None,
+                    is_error: None,
+                    meta: None,
+                }),
+                error: None,
+                duration: Some(Duration::from_millis(8)),
+            }),
+        };
+        let events = vec![
+            EventMsg::TurnStarted(TurnStartedEvent {
+                turn_id: "turn-a".into(),
+                root_turn_id: None,
+                trace_id: None,
+                started_at: None,
+                model_context_window: None,
+                collaboration_mode_kind: Default::default(),
+            }),
+            EventMsg::UserMessage(UserMessageEvent {
+                message: "Find related incident discussions".into(),
+                ..Default::default()
+            }),
+            EventMsg::TurnComplete(TurnCompleteEvent {
+                turn_id: "turn-a".into(),
+                started_at: None,
+                last_agent_message: None,
+                error: None,
+                completed_at: None,
+                duration_ms: None,
+                time_to_first_token_ms: None,
+            }),
+            EventMsg::TurnStarted(TurnStartedEvent {
+                turn_id: "turn-b".into(),
+                root_turn_id: None,
+                trace_id: None,
+                started_at: None,
+                model_context_window: None,
+                collaboration_mode_kind: Default::default(),
+            }),
+            EventMsg::UserMessage(UserMessageEvent {
+                message: "Explain retry logic while that runs".into(),
+                ..Default::default()
+            }),
+        ];
+        let mut items: Vec<_> = events.into_iter().map(RolloutItem::EventMsg).collect();
+        let mut expected = build_turns_from_rollout_items(&items);
+        expected[0]
+            .items
+            .push(ThreadItem::from(completion.item.clone()));
+
+        let legacy = completion.as_legacy_events(/*show_raw_agent_reasoning*/ false);
+        let serialized = serde_json::to_string(&legacy[0]).unwrap();
+        items.push(RolloutItem::EventMsg(
+            serde_json::from_str(&serialized).unwrap(),
+        ));
+        assert_eq!(build_turns_from_rollout_items(&items), expected);
+
+        // Old records have no owner, so retain their current-turn fallback.
+        let mut old_record: serde_json::Value = serde_json::from_str(&serialized).unwrap();
+        old_record.as_object_mut().unwrap().remove("turn_id");
+        *items.last_mut().unwrap() =
+            RolloutItem::EventMsg(serde_json::from_value(old_record).unwrap());
+        let mut expected_old = expected;
+        let mcp_item = expected_old[0].items.pop().unwrap();
+        expected_old[1].items.push(mcp_item);
+        assert_eq!(build_turns_from_rollout_items(&items), expected_old);
+    }
+
     #[test]
     fn assigns_late_exec_completion_to_original_turn() {
         let events = vec![

codex-rs/app-server/tests/suite/v2/thread_resume.rs

@@ -2678,6 +2678,7 @@ fn append_resume_redaction_history(
     let persisted_rollout = std::fs::read_to_string(&rollout_file_path)?;
     let appended_rollout = [
         EventMsg::McpToolCallEnd(McpToolCallEndEvent {
+            turn_id: String::new(),
             call_id: "mcp-1".to_string(),
             invocation: McpInvocation {
                 server: "docs".to_string(),

codex-rs/core/src/guardian/decision.rs

@@ -177,7 +177,7 @@ pub(crate) async fn decide_approval(
                         },
                     });
             }
-            record_guardian_non_denial(&session, turn_id).await;
+            record_guardian_non_denial(&session).await;
             Some(ReviewDecision::Approved)
         }
         Some(ApprovalDecision::Allow) => runtime.review(GuardianReviewReason::FreshRequired).await,

codex-rs/core/src/guardian/review.rs

@@ -136,13 +136,28 @@ pub(crate) fn is_basic_session_source(session_source: &SessionSource) -> bool {
     }
 }
 
-pub(super) async fn record_guardian_non_denial(session: &Arc<Session>, turn_id: &str) {
+pub(super) async fn record_guardian_non_denial(session: &Arc<Session>) {
+    let turn_id = {
+        let active = session.active_turn.lock().await;
+        let Some(task) = active.as_ref().and_then(|active| active.task.as_ref()) else {
+            return;
+        };
+        task.turn_context.sub_id.clone()
+    };
     codex_guardian_reviewer::ReviewDenials::for_thread(&session.services.thread_extension_data)
-        .record_non_denial(turn_id)
+        .record_non_denial(&turn_id)
         .await;
 }
 
-async fn record_guardian_denial(session: &Arc<Session>, turn: &Arc<TurnContext>, turn_id: &str) {
+async fn record_guardian_denial(session: &Arc<Session>) {
+    let turn = {
+        let active = session.active_turn.lock().await;
+        let Some(task) = active.as_ref().and_then(|active| active.task.as_ref()) else {
+            return;
+        };
+        Arc::clone(&task.turn_context)
+    };
+    let turn_id = &turn.sub_id;
     let Some(message) =
         codex_guardian_reviewer::ReviewDenials::for_thread(&session.services.thread_extension_data)
             .record_denial(turn_id, turn.model_info())
@@ -180,12 +195,8 @@ async fn record_guardian_denial(session: &Arc<Session>, turn: &Arc<TurnContext>,
 }
 
 #[cfg(test)]
-pub(crate) async fn record_guardian_denial_for_test(
+pub(crate) async fn record_guardian_denial_for_test(session: &Arc<Session>) {
-    session: &Arc<Session>,
+    record_guardian_denial(session).await;
-    turn: &Arc<TurnContext>,
-    turn_id: &str,
-) {
-    record_guardian_denial(session, turn, turn_id).await;
 }
 
 #[derive(Clone)]

codex-rs/core/src/guardian/review_request.rs

@@ -129,7 +129,7 @@ impl ReviewHost for super::super::runtime::ReviewRuntime {
             session
                 .send_event(turn.as_ref(), EventMsg::GuardianAssessment(completed.event))
                 .await;
-            record_guardian_non_denial(&session, report.turn_id()).await;
+            record_guardian_non_denial(&session).await;
             return Err(ReviewDecision::Abort);
         }
 
@@ -259,9 +259,9 @@ impl ReviewHost for super::super::runtime::ReviewRuntime {
             .send_event(turn.as_ref(), EventMsg::GuardianAssessment(completed.event))
             .await;
         if completed.assessment_outcome == Some(GuardianAssessmentOutcome::Deny) {
-            record_guardian_denial(&session, &turn, report.turn_id()).await;
+            record_guardian_denial(&session).await;
         } else {
-            record_guardian_non_denial(&session, report.turn_id()).await;
+            record_guardian_non_denial(&session).await;
         }
         completed.decision
     }

codex-rs/core/src/mcp_tool_call.rs

@@ -253,8 +253,8 @@ pub(crate) async fn handle_mcp_tool_call(
                 .unwrap_or_else(|| JsonValue::Object(serde_json::Map::new())),
         };
     }
-    sess.register_mcp_tool_approval_metadata(turn_context, &call_id, &invocation, metadata.clone())
+    let _approval_metadata =
-        .await;
+        sess.register_mcp_tool_approval_metadata(&call_id, &invocation, metadata.clone());
     notify_mcp_tool_call_started(
         sess.as_ref(),
         turn_context.as_ref(),
@@ -1139,40 +1139,41 @@ pub(crate) struct McpToolApprovalMetadata {
 }
 
 impl Session {
-    async fn register_mcp_tool_approval_metadata(
+    fn register_mcp_tool_approval_metadata(
         &self,
-        turn_context: &TurnContext,
         call_id: &str,
         invocation: &McpInvocation,
         metadata: McpToolApprovalMetadata,
-    ) {
+    ) -> Arc<(Option<McpInvocation>, McpToolApprovalMetadata)> {
-        let Some(turn_state) = self
+        let key = (invocation.server.clone(), call_id.to_string());
-            .input_queue
+        let metadata = Arc::new((
-            .turn_state_for_sub_id(&self.active_turn, &turn_context.sub_id)
-            .await
-        else {
-            return;
-        };
-        turn_state.lock().await.insert_mcp_tool_approval_metadata(
-            call_id.to_string(),
             (invocation.server == CODEX_APPS_MCP_SERVER_NAME
                 || is_node_repl_backed_server(&invocation.server))
             .then(|| invocation.clone()),
             metadata,
-        );
+        ));
+        let mut registry = self
+            .mcp_tool_approval_metadata
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        // Keep approval details available while a tool is running, even across turns.
+        // When a new call starts, clean up entries left by finished calls.
+        registry.retain(|_, entry| entry.strong_count() != 0);
+        registry.insert(key, Arc::downgrade(&metadata));
+        metadata
     }
 
-    pub(crate) async fn mcp_tool_approval_metadata(
+    pub(crate) fn mcp_tool_approval_metadata(
         &self,
-        sub_id: &str,
+        server: &str,
         call_id: &str,
     ) -> Option<(Option<McpInvocation>, McpToolApprovalMetadata)> {
-        let turn_state = self
+        self.mcp_tool_approval_metadata
-            .input_queue
+            .lock()
-            .turn_state_for_sub_id(&self.active_turn, sub_id)
+            .unwrap_or_else(std::sync::PoisonError::into_inner)
-            .await?;
+            .get(&(server.to_string(), call_id.to_string()))
-
+            .and_then(std::sync::Weak::upgrade)
-        turn_state.lock().await.mcp_tool_approval_metadata(call_id)
+            .map(|metadata| (*metadata).clone())
     }
 }
 
@@ -1571,8 +1572,7 @@ pub(crate) async fn request_mcp_tool_user_approval(
     );
     if tool_call_mcp_elicitation_enabled {
         let link_id = sess
-            .mcp_tool_approval_metadata(&turn_context.sub_id, id)
+            .mcp_tool_approval_metadata(server, id)
-            .await
             .and_then(|(_, metadata)| metadata.link_id);
         let metadata = McpToolApprovalMetadata {
             annotations: None,

codex-rs/core/src/mcp_tool_call_tests.rs

@@ -3309,3 +3309,63 @@ async fn approve_mode_skips_guardian_in_every_permission_mode() {
         assert_eq!(decision, None);
     }
 }
+
+#[tokio::test]
+async fn approval_metadata_is_released_when_the_invocation_future_is_dropped() {
+    let (session, _) = crate::session::tests::make_session_and_context().await;
+    let invocation = McpInvocation {
+        server: CODEX_APPS_MCP_SERVER_NAME.to_string(),
+        tool: "write_record".to_string(),
+        arguments: Some(serde_json::json!({"value": 42})),
+    };
+    let metadata = approval_metadata(
+        Some("connector"),
+        /*connector_name*/ None,
+        /*connector_description*/ None,
+        /*tool_title*/ None,
+        /*tool_description*/ None,
+    );
+    let mut call = Box::pin(async {
+        let _approval_metadata =
+            session.register_mcp_tool_approval_metadata("call", &invocation, metadata.clone());
+        std::future::pending::<()>().await;
+    });
+    assert!(futures::poll!(call.as_mut()).is_pending());
+    let _other_metadata =
+        session.register_mcp_tool_approval_metadata("other-call", &invocation, metadata.clone());
+    assert_eq!(
+        session
+            .mcp_tool_approval_metadata(CODEX_APPS_MCP_SERVER_NAME, "call")
+            .map(|(invocation, metadata)| (invocation, metadata.connector_id)),
+        Some((Some(invocation.clone()), Some("connector".to_string()))),
+    );
+    assert!(
+        session
+            .mcp_tool_approval_metadata("another-server", "call")
+            .is_none()
+    );
+    drop(call);
+    assert!(
+        session
+            .mcp_tool_approval_metadata(CODEX_APPS_MCP_SERVER_NAME, "call")
+            .is_none()
+    );
+    let _next_metadata =
+        session.register_mcp_tool_approval_metadata("next-call", &invocation, metadata);
+    let registry = session.mcp_tool_approval_metadata.lock().unwrap();
+    let mut keys = registry.keys().cloned().collect::<Vec<_>>();
+    keys.sort();
+    assert_eq!(
+        keys,
+        vec![
+            (
+                CODEX_APPS_MCP_SERVER_NAME.to_string(),
+                "next-call".to_string()
+            ),
+            (
+                CODEX_APPS_MCP_SERVER_NAME.to_string(),
+                "other-call".to_string()
+            ),
+        ],
+    );
+}

codex-rs/core/src/session/mcp.rs

@@ -789,9 +789,8 @@ async fn review_guardian_mcp_elicitation(
             .meta()
             .and_then(|meta| meta.get("callId"))
             .and_then(Value::as_str)
-        && let Some((Some(invocation), _)) = session
+        && let Some((Some(invocation), _)) =
-            .mcp_tool_approval_metadata(&turn_context.sub_id, call_id)
+            session.mcp_tool_approval_metadata(&request.server_name, call_id)
-            .await
         && invocation.server == request.server_name
     {
         Some(call_id)
@@ -828,9 +827,8 @@ async fn review_guardian_mcp_elicitation(
             else {
                 return Ok(None);
             };
-            let Some((Some(invocation), metadata)) = session
+            let Some((Some(invocation), metadata)) =
-                .mcp_tool_approval_metadata(&turn_context.sub_id, call_id)
+                session.mcp_tool_approval_metadata(&request.server_name, call_id)
-                .await
             else {
                 return Ok(None);
             };

codex-rs/core/src/session/session.rs

@@ -12,6 +12,7 @@ use crate::context::GuardianContextMode;
 use crate::environment_selection::ThreadEnvironments;
 use crate::environment_selection::TurnEnvironmentSnapshot;
 use crate::hook_mcp_executor::CoreHookMcpExecutor;
+use crate::mcp_tool_call::McpToolApprovalMetadata;
 use crate::responses_metadata::CodexResponsesMetadata;
 use crate::responses_metadata::CodexResponsesRequestKind;
 use crate::shell_snapshot::ShellSnapshot;
@@ -30,6 +31,7 @@ use codex_protocol::permissions::FileSystemPath;
 use codex_protocol::permissions::FileSystemSpecialPath;
 use codex_protocol::protocol::EnvironmentConfig;
 use codex_protocol::protocol::HookCompletedEvent;
+use codex_protocol::protocol::McpInvocation;
 use codex_protocol::protocol::MultiAgentVersion;
 use codex_protocol::protocol::ThreadHistoryMode;
 use codex_protocol::protocol::ThreadSource;
@@ -40,6 +42,9 @@ use codex_utils_path::replace_path_and_deduplicate;
 use std::sync::OnceLock;
 use tokio::sync::Semaphore;
 
+type McpToolApprovalMetadataMap =
+    HashMap<(String, String), std::sync::Weak<(Option<McpInvocation>, McpToolApprovalMetadata)>>;
+
 /// Context for an initialized model agent
 ///
 /// A session has at most 1 running task at a time, and can be interrupted by user input.
@@ -65,6 +70,8 @@ pub(crate) struct Session {
     pub(super) multi_agent_version: OnceLock<MultiAgentVersion>,
     /// Owns invalidation and serializes refreshes without blocking captured calls.
     pub(super) mcp_refresh: McpRefresh,
+    /// Non-owning lookup for approval data retained by running MCP invocations.
+    pub(crate) mcp_tool_approval_metadata: std::sync::Mutex<McpToolApprovalMetadataMap>,
     pub(super) mcp_elicitation_reviewer_handle: OnceLock<codex_mcp::ElicitationReviewerHandle>,
     pub(super) mcp_elicitation_lifecycle_handle: OnceLock<codex_mcp::ElicitationLifecycle>,
     pub(super) mcp_prewarm_tx: async_channel::Sender<()>,
@@ -1625,7 +1632,8 @@ impl Session {
                 windows_sandbox_proxy_settings_mode,
                 multi_agent_version,
                 mcp_refresh: McpRefresh::new(),
-                mcp_elicitation_reviewer_handle: OnceLock::new(),
+                mcp_tool_approval_metadata: Default::default(),
+        mcp_elicitation_reviewer_handle: OnceLock::new(),
                 mcp_elicitation_lifecycle_handle: OnceLock::new(),
                 mcp_prewarm_tx,
                 mcp_prewarm_shutdown: CancellationToken::new(),

codex-rs/core/src/session/tests.rs

@@ -6743,6 +6743,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
             codex_sandboxing::WindowsSandboxProxySettingsMode::Reconcile,
         multi_agent_version: OnceLock::from(config.multi_agent_version_from_features()),
         mcp_refresh: McpRefresh::new(),
+        mcp_tool_approval_metadata: Default::default(),
         mcp_elicitation_reviewer_handle: OnceLock::new(),
         mcp_elicitation_lifecycle_handle: OnceLock::new(),
         mcp_prewarm_tx: async_channel::bounded(1).0,
@@ -8917,6 +8918,7 @@ where
             codex_sandboxing::WindowsSandboxProxySettingsMode::Reconcile,
         multi_agent_version: OnceLock::from(config.multi_agent_version_from_features()),
         mcp_refresh: McpRefresh::new(),
+        mcp_tool_approval_metadata: Default::default(),
         mcp_elicitation_reviewer_handle: OnceLock::new(),
         mcp_elicitation_lifecycle_handle: OnceLock::new(),
         mcp_prewarm_tx: async_channel::bounded(1).0,
@@ -11245,12 +11247,12 @@ impl SessionTask for GuardianDeniedApprovalTask {
     async fn run(
         self: Arc<Self>,
         session: Arc<Session>,
-        ctx: Arc<TurnContext>,
+        _ctx: Arc<TurnContext>,
         _input: Vec<TurnInput>,
         cancellation_token: CancellationToken,
     ) -> SessionTaskResult {
         for _ in 0..3 {
-            crate::guardian::record_guardian_denial_for_test(&session, &ctx, &ctx.sub_id).await;
+            crate::guardian::record_guardian_denial_for_test(&session).await;
         }
 
         cancellation_token.cancelled().await;
@@ -11621,8 +11623,6 @@ async fn guardian_helper_review_interrupts_after_three_consecutive_denials() {
     .await;
 
     let session_for_review = Arc::clone(&sess);
-    let turn_for_review = Arc::clone(&tc);
-    let turn_id = tc.sub_id.clone();
     let review_thread = std::thread::spawn(move || {
         let runtime = tokio::runtime::Builder::new_current_thread()
             .enable_all()
@@ -11630,12 +11630,7 @@ async fn guardian_helper_review_interrupts_after_three_consecutive_denials() {
             .expect("helper review runtime");
         runtime.block_on(async move {
             for _ in 0..3 {
-                crate::guardian::record_guardian_denial_for_test(
+                crate::guardian::record_guardian_denial_for_test(&session_for_review).await;
-                    &session_for_review,
-                    &turn_for_review,
-                    &turn_id,
-                )
-                .await;
             }
         });
     });

codex-rs/core/src/state/turn.rs

@@ -19,14 +19,12 @@ use tokio::sync::oneshot;
 
 use super::TurnTokenUsage;
 use crate::agent::control::AgentExecutionGuard;
-use crate::mcp_tool_call::McpToolApprovalMetadata;
 use crate::session::TurnInputQueue;
 use crate::session::step_context::StepContext;
 use crate::session::turn_context::TurnContext;
 use crate::session::turn_context::TurnEnvironment;
 use crate::tasks::AnySessionTask;
 use codex_protocol::models::AdditionalPermissionProfile;
-use codex_protocol::protocol::McpInvocation;
 use codex_protocol::protocol::ReviewDecision;
 use codex_protocol::protocol::TokenUsage;
 
@@ -93,7 +91,6 @@ pub(crate) struct TurnState {
     pending_request_permissions: HashMap<String, PendingRequestPermissions>,
     pending_user_input: HashMap<String, oneshot::Sender<AcceptedUserInputResponse>>,
     pending_elicitations: HashMap<(String, RequestId), oneshot::Sender<ElicitationResponse>>,
-    mcp_tool_approval_metadata: HashMap<String, (Option<McpInvocation>, McpToolApprovalMetadata)>,
     pending_dynamic_tools: HashMap<String, oneshot::Sender<DynamicToolResponse>>,
     pub(crate) pending_input: TurnInputQueue,
     mailbox_delivery_phase: MailboxDeliveryPhase,
@@ -142,7 +139,6 @@ impl TurnState {
         self.pending_request_permissions.clear();
         self.pending_user_input.clear();
         self.pending_elicitations.clear();
-        self.mcp_tool_approval_metadata.clear();
         self.pending_dynamic_tools.clear();
     }
 
@@ -196,23 +192,6 @@ impl TurnState {
             .remove(&(server_name.to_string(), request_id.clone()))
     }
 
-    pub(crate) fn insert_mcp_tool_approval_metadata(
-        &mut self,
-        call_id: String,
-        invocation: Option<McpInvocation>,
-        metadata: McpToolApprovalMetadata,
-    ) {
-        self.mcp_tool_approval_metadata
-            .insert(call_id, (invocation, metadata));
-    }
-
-    pub(crate) fn mcp_tool_approval_metadata(
-        &self,
-        call_id: &str,
-    ) -> Option<(Option<McpInvocation>, McpToolApprovalMetadata)> {
-        self.mcp_tool_approval_metadata.get(call_id).cloned()
-    }
-
     pub(crate) fn insert_pending_dynamic_tool(
         &mut self,
         key: String,

codex-rs/core/src/tools/code_mode/delegate.rs

@@ -1,6 +1,7 @@
 use std::collections::HashMap;
 use std::sync::Arc;
 use std::sync::Mutex;
+use std::sync::Weak;
 
 use codex_code_mode::CellId;
 use codex_code_mode::CodeModeNestedToolCall;
@@ -17,12 +18,12 @@ use tokio::sync::watch;
 use tokio_util::sync::CancellationToken;
 use tracing::Instrument;
 
-use super::ExecContext;
 use super::PUBLIC_TOOL_NAME;
 use super::submit_nested_tool;
 use super::telemetry::DispatchInterruption;
 use super::telemetry::NestedToolDispatchTrace;
 use super::telemetry::trace_id;
+use crate::session::session::Session;
 use crate::session::step_context::StepContext;
 use crate::tools::ExecutedToolCalls;
 use crate::tools::call_trace;
@@ -37,6 +38,12 @@ pub(super) struct CodeModeDispatchBroker {
     executed_tool_calls: ExecutedToolCalls,
 }
 
+/// Retains the step advertised to one execution, including callbacks after it yields.
+pub(super) struct CodeModeCellDelegate {
+    pub(super) broker: Arc<CodeModeDispatchBroker>,
+    pub(super) step_context: Arc<StepContext>,
+}
+
 struct CellDispatchGate {
     ready: watch::Sender<bool>,
     // Keep the original exec item when later waits resume this cell.
@@ -105,13 +112,15 @@ impl CodeModeDispatchBroker {
 
     pub(super) fn start_turn_worker(
         &self,
-        exec: ExecContext,
+        session: Arc<Session>,
         step_context: Arc<StepContext>,
         tracker: SharedTurnDiffTracker,
     ) -> CodeModeDispatchWorker {
-        let track_completeness = ExecutedToolCalls::is_enabled(&exec.turn.config.features);
+        let tool_runtime = ToolCallRuntime::new(Arc::clone(&session), step_context, tracker);
-        let tool_runtime = ToolCallRuntime::new(Arc::clone(&exec.session), step_context, tracker);
+        let host = Arc::new(CoreTurnHost {
-        let host = Arc::new(CoreTurnHost { exec, tool_runtime });
+            session,
+            tool_runtime,
+        });
         let dispatch_rx = self.dispatch_rx.clone();
         let dispatch_gates = Arc::clone(&self.dispatch_gates);
         let (shutdown_tx, mut shutdown_rx) = oneshot::channel();
@@ -148,6 +157,7 @@ impl CodeModeDispatchBroker {
                     }
                     DispatchMessage::InvokeTool {
                         invocation,
+                        step_context,
                         mut dispatch_trace,
                         cancellation_token,
                         response_tx,
@@ -169,6 +179,16 @@ impl CodeModeDispatchBroker {
                             remove_dispatch_gate(&dispatch_gates, &cell_id);
                             continue;
                         }
+                        let Some(step_context) = step_context
+                            .upgrade()
+                            .filter(|_| !cancellation_token.is_cancelled())
+                        else {
+                            let _ = response_tx
+                                .send(Err("code mode nested tool call cancelled".to_string()));
+                            continue;
+                        };
+                        let track_completeness =
+                            ExecutedToolCalls::is_enabled(&step_context.turn.config.features);
                         let host = Arc::clone(&host);
                         let dispatch_gates = Arc::clone(&dispatch_gates);
                         let span = dispatch_trace.span.clone();
@@ -196,6 +216,7 @@ impl CodeModeDispatchBroker {
                                     dispatch_trace.interruption = None;
                                     host.submit_tool(
                                         invocation,
+                                        step_context,
                                         dispatch_trace.call_id.clone(),
                                         cancellation_token.clone(),
                                     )
@@ -273,13 +294,13 @@ async fn wait_until_cell_ready_for_dispatch(
     }
 }
 
-impl CodeModeSessionDelegate for CodeModeDispatchBroker {
+impl CodeModeSessionDelegate for CodeModeCellDelegate {
     #[tracing::instrument(
         name = "code_mode.broker.invoke_tool",
         level = "info",
         skip_all,
         fields(
-            conversation.id = %self.thread_id,
+            conversation.id = %self.broker.thread_id,
             cell.id = %invocation.cell_id,
             runtime_tool_call_id = invocation.runtime_tool_call_id.as_str(),
             tool_name = invocation.tool_name.name.as_str(),
@@ -294,7 +315,7 @@ impl CodeModeSessionDelegate for CodeModeDispatchBroker {
         Box::pin(async move {
             let call_id = format!("{PUBLIC_TOOL_NAME}-{}", uuid::Uuid::new_v4());
             call_trace::received(
-                self.thread_id,
+                self.broker.thread_id,
                 &invocation.tool_name,
                 &call_id,
                 call_trace::Receipt::CodeModeBroker {
@@ -303,16 +324,18 @@ impl CodeModeSessionDelegate for CodeModeDispatchBroker {
                 },
             );
             let mut dispatch_trace =
-                Box::new(NestedToolDispatchTrace::new(self.thread_id, call_id));
+                Box::new(NestedToolDispatchTrace::new(self.broker.thread_id, call_id));
             if cancellation_token.is_cancelled() {
                 dispatch_trace.interruption = Some(DispatchInterruption::Cancelled);
                 return Err("code mode nested tool call cancelled".to_string());
             }
             let (response_tx, response_rx) = oneshot::channel();
             // Only the worker can tell whether dispatch beats cancellation once the call is queued.
-            self.dispatch_tx
+            self.broker
+                .dispatch_tx
                 .send(DispatchMessage::InvokeTool {
                     invocation,
+                    step_context: Arc::downgrade(&self.step_context),
                     dispatch_trace,
                     cancellation_token: cancellation_token.clone(),
                     response_tx,
@@ -341,7 +364,8 @@ impl CodeModeSessionDelegate for CodeModeDispatchBroker {
                 return Err("code mode notification cancelled".to_string());
             }
             let (response_tx, response_rx) = oneshot::channel();
-            self.dispatch_tx
+            self.broker
+                .dispatch_tx
                 .send(DispatchMessage::Notify {
                     call_id,
                     cell_id,
@@ -362,13 +386,15 @@ impl CodeModeSessionDelegate for CodeModeDispatchBroker {
     }
 
     fn cell_closed(&self, cell_id: &CellId) {
-        self.close_cell(cell_id);
+        self.broker.close_cell(cell_id);
     }
 }
 
 enum DispatchMessage {
     InvokeTool {
         invocation: CodeModeNestedToolCall,
+        // The delegate owns the step while the callback is live; stale queued work must not.
+        step_context: Weak<StepContext>,
         dispatch_trace: Box<NestedToolDispatchTrace>,
         cancellation_token: CancellationToken,
         response_tx: oneshot::Sender<Result<JsonValue, String>>,
@@ -395,7 +421,7 @@ impl Drop for CodeModeDispatchWorker {
 }
 
 struct CoreTurnHost {
-    exec: ExecContext,
+    session: Arc<Session>,
     tool_runtime: ToolCallRuntime,
 }
 
@@ -403,11 +429,13 @@ impl CoreTurnHost {
     fn submit_tool(
         &self,
         invocation: CodeModeNestedToolCall,
+        step_context: Arc<StepContext>,
         call_id: String,
         cancellation_token: CancellationToken,
     ) -> impl std::future::Future<Output = Result<JsonValue, String>> + Send + 'static {
         let invocation = submit_nested_tool(
-            self.exec.clone(),
+            Arc::clone(&self.session),
+            step_context,
             self.tool_runtime.clone(),
             invocation,
             call_id,
@@ -421,8 +449,7 @@ impl CoreTurnHost {
         if text.trim().is_empty() {
             return Ok(());
         }
-        self.exec
+        self.session
-            .session
             .inject_if_running(vec![ResponseItem::CustomToolCallOutput {
                 id: None,
                 call_id,

codex-rs/core/src/tools/code_mode/delegate_tests.rs

@@ -1,4 +1,4 @@
-//! Tracing identity coverage for code-mode callbacks outside local turn spans.
+//! Callback tracing identity and originating-step lifetime coverage.
 
 use std::collections::BTreeMap;
 use std::sync::Arc;
@@ -20,8 +20,9 @@ use tracing_subscriber::layer::Context;
 use tracing_subscriber::prelude::*;
 use tracing_subscriber::registry::Registry;
 
+use super::CodeModeCellDelegate;
 use super::CodeModeDispatchBroker;
-use super::ExecContext;
+use super::DispatchMessage;
 use crate::session::step_context::StepContext;
 use crate::session::tests::make_session_and_context;
 use crate::tools::context::FunctionToolOutput;
@@ -31,6 +32,7 @@ use crate::tools::registry::ToolExecutor;
 use crate::tools::registry::ToolRegistry;
 use crate::tools::router::ToolRouter;
 use crate::turn_diff_tracker::TurnDiffTracker;
+use tokio::sync::oneshot;
 
 #[derive(Default)]
 struct DispatchFields(BTreeMap<String, String>);
@@ -125,17 +127,18 @@ async fn detached_code_mode_callback_keeps_thread_id_on_dispatch_span() -> anyho
     let subscriber = tracing_subscriber::registry().with(DispatchCapture(Arc::clone(&records)));
     let _untraced = tracing::Dispatch::new(tracing::subscriber::NoSubscriber::default());
     let _subscriber = tracing::subscriber::set_default(subscriber);
+    let delegate = CodeModeCellDelegate {
+        broker: Arc::clone(&broker),
+        step_context: Arc::clone(&step),
+    };
     let _worker = broker.start_turn_worker(
-        ExecContext {
+        Arc::new(session),
-            session: Arc::new(session),
-            turn,
-        },
         step,
         Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())),
     );
     // Process-owned code-mode callbacks run on a fresh task without a turn span.
     let result = tokio::spawn(async move {
-        broker
+        delegate
             .invoke_tool(
                 CodeModeNestedToolCall {
                     cell_id,
@@ -170,3 +173,63 @@ async fn detached_code_mode_callback_keeps_thread_id_on_dispatch_span() -> anyho
     uuid::Uuid::parse_str(uuid)?;
     Ok(())
 }
+
+#[tokio::test]
+async fn dropped_tool_callbacks_release_the_origin_before_dispatch() {
+    let (session, turn) = make_session_and_context().await;
+    let session = Arc::new(session);
+    let turn = Arc::new(turn);
+    let broker = Arc::new(CodeModeDispatchBroker::new(
+        session.thread_id,
+        /*executed_tool_calls*/ Default::default(),
+    ));
+    let delegate = CodeModeCellDelegate {
+        broker: Arc::clone(&broker),
+        step_context: StepContext::for_test(Arc::clone(&turn)),
+    };
+    let origin = Arc::downgrade(&delegate.step_context);
+    let cell_id = CellId::new("queued-cell".to_string());
+    let cancellation_token = CancellationToken::new();
+    let mut callback = delegate.invoke_tool(
+        CodeModeNestedToolCall {
+            cell_id: cell_id.clone(),
+            runtime_tool_call_id: "call-a".to_string(),
+            tool_name: codex_tools::ToolName::plain("test-tool"),
+            tool_kind: codex_code_mode::CodeModeToolKind::Function,
+            input: None,
+        },
+        cancellation_token.clone(),
+    );
+    assert!(futures::poll!(callback.as_mut()).is_pending());
+    assert_eq!(broker.dispatch_rx.len(), 1);
+    drop(callback);
+    drop(delegate);
+    assert!(
+        origin.upgrade().is_none(),
+        "queued work must not own the step"
+    );
+
+    // Observe the queued message's result after its original receiver was dropped.
+    let mut message = broker.dispatch_rx.try_recv().unwrap();
+    let (response_tx, response_rx) = oneshot::channel();
+    let DispatchMessage::InvokeTool {
+        response_tx: queued_response,
+        ..
+    } = &mut message
+    else {
+        panic!("expected tool invocation");
+    };
+    *queued_response = response_tx;
+    broker.dispatch_tx.send(message).await.unwrap();
+    broker.mark_cell_ready_for_dispatch(&cell_id, /*originating_item_id*/ None);
+    let _worker = broker.start_turn_worker(
+        Arc::clone(&session),
+        StepContext::for_test(Arc::clone(&turn)),
+        Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())),
+    );
+    let response = response_rx.await;
+    assert_eq!(
+        response.unwrap(),
+        Err("code mode nested tool call cancelled".to_string()),
+    );
+}

codex-rs/core/src/tools/code_mode/execute_handler.rs

@@ -34,7 +34,7 @@ impl CodeModeExecuteHandler {
     async fn execute(
         &self,
         session: std::sync::Arc<crate::session::session::Session>,
-        step_context: std::sync::Arc<crate::session::step_context::StepContext>,
+        step_context: Arc<crate::session::step_context::StepContext>,
         call_id: String,
         originating_item_id: Option<codex_protocol::ResponseItemId>,
         code: String,
@@ -71,13 +71,16 @@ impl CodeModeExecuteHandler {
             .session
             .services
             .code_mode_service
-            .execute(codex_code_mode::ExecuteRequest {
+            .execute(
-                tool_call_id: call_id.clone(),
+                codex_code_mode::ExecuteRequest {
-                enabled_tools,
+                    tool_call_id: call_id.clone(),
-                source: args.code.clone(),
+                    enabled_tools,
-                yield_time_ms: args.yield_time_ms,
+                    source: args.code.clone(),
-                max_output_tokens: args.max_output_tokens,
+                    yield_time_ms: args.yield_time_ms,
-            })
+                    max_output_tokens: args.max_output_tokens,
+                },
+                Arc::clone(&step_context),
+            )
             .await
             .map_err(FunctionCallError::RespondToModel)?;
         let cell_id = started_cell.cell_id.clone();

codex-rs/core/src/tools/code_mode/mod.rs

@@ -47,6 +47,7 @@ use codex_utils_output_truncation::TruncationPolicy;
 use codex_utils_output_truncation::formatted_truncate_text_content_items_with_policy;
 use codex_utils_output_truncation::truncate_function_output_items_with_policy;
 
+use delegate::CodeModeCellDelegate;
 use delegate::CodeModeDispatchBroker;
 use delegate::CodeModeDispatchWorker;
 pub(crate) use execute_handler::CodeModeExecuteHandler;
@@ -125,14 +126,16 @@ impl CodeModeService {
     pub(crate) async fn execute(
         &self,
         mut request: codex_code_mode::ExecuteRequest,
+        step_context: Arc<StepContext>,
     ) -> Result<codex_code_mode::StartedCell, String> {
         request
             .yield_time_ms
             .get_or_insert(self.default_exec_yield_time_ms);
-        self.session()
+        let delegate = Arc::new(CodeModeCellDelegate {
-            .await?
+            broker: Arc::clone(&self.dispatch_broker),
-            .execute(request, self.dispatch_broker.clone())
+            step_context,
-            .await
+        });
+        self.session().await?.execute(request, delegate).await
     }
 
     pub(crate) async fn wait(
@@ -209,18 +212,13 @@ impl CodeModeService {
         step_context: Arc<StepContext>,
         tracker: SharedTurnDiffTracker,
     ) -> Option<CodeModeDispatchWorker> {
-        let turn = &step_context.turn;
         if !step_context.tool_router.requires_code_mode_worker() {
             return None;
         }
 
-        let exec = ExecContext {
-            session: Arc::clone(session),
-            turn: Arc::clone(turn),
-        };
         Some(
             self.dispatch_broker
-                .start_turn_worker(exec, step_context, tracker),
+                .start_turn_worker(Arc::clone(session), step_context, tracker),
         )
     }
 
@@ -346,7 +344,8 @@ fn truncate_code_mode_result(
 
 // Submit synchronously so the recorder sees the call before the cell's dispatch gate closes.
 fn submit_nested_tool(
-    exec: ExecContext,
+    session: Arc<Session>,
+    step_context: Arc<StepContext>,
     tool_runtime: ToolCallRuntime,
     invocation: CodeModeNestedToolCall,
     call_id: String,
@@ -362,8 +361,8 @@ fn submit_nested_tool(
         tool_kind,
         input,
     } = invocation;
-    let thread_id = exec.session.thread_id;
+    let thread_id = session.thread_id;
-    let turn_id = exec.turn.sub_id.clone();
+    let turn_id = step_context.turn.sub_id.clone();
     let tool_name = tool_name.with_default_namespace();
     // A cell can outlive a turn; the broker records arrival before a dispatching turn is known.
     tracing::event!(
@@ -402,16 +401,17 @@ fn submit_nested_tool(
         payload,
         encrypted_function_args: None,
     };
-    exec.session
+    session
         .services
         .analytics_events_client
         .track_code_mode_tool_call(codex_analytics::CodeModeToolCallFact::ChildStarted {
-            thread_id: exec.session.thread_id.to_string(),
+            thread_id: session.thread_id.to_string(),
-            turn_id: exec.turn.sub_id.clone(),
+            turn_id: step_context.turn.sub_id.clone(),
             call_id: call.call_id.clone(),
             cell_id: cell_id.to_string(),
         });
     let result = tool_runtime.handle_tool_call_with_source(
+        step_context,
         call,
         ToolCallSource::CodeMode {
             cell_id: cell_id.to_string(),

codex-rs/core/src/tools/parallel.rs

@@ -79,7 +79,9 @@ impl ToolCallRuntime {
     ) -> impl std::future::Future<Output = Result<ResponseItemEnvelope, CodexErr>> {
         let error_call = call.clone();
         let source = call.direct_source();
-        let future = self.handle_tool_call_with_source(call, source, cancellation_token);
+        let step_context = Arc::clone(&self.step_context);
+        let future =
+            self.handle_tool_call_with_source(step_context, call, source, cancellation_token);
         async move {
             match future.await {
                 Ok(response) => Ok(response.into_response()),
@@ -94,21 +96,20 @@ impl ToolCallRuntime {
     #[instrument(level = "trace", skip_all)]
     pub(crate) fn handle_tool_call_with_source(
         self,
+        step_context: Arc<StepContext>,
         call: ToolCall,
         source: ToolCallSource,
         cancellation_token: CancellationToken,
     ) -> impl std::future::Future<Output = Result<AnyToolResult, FunctionCallError>> {
-        self.session.services.executed_tool_calls.record_tool_call(
+        self.session
-            &call,
+            .services
-            &source,
+            .executed_tool_calls
-            &self.step_context,
+            .record_tool_call(&call, &source, &step_context);
-        );
+        let router = &step_context.tool_router;
-        let router = &self.step_context.tool_router;
         let supports_parallel = router.tool_supports_parallel(&call);
         let tool_runtime = router.tool_runtime(&call.tool_name);
         let router = Arc::clone(router);
         let session = Arc::clone(&self.session);
-        let step_context = Arc::clone(&self.step_context);
         let turn = Arc::clone(&step_context.turn);
         let tracker = Arc::clone(&self.tracker);
         let lock = Arc::clone(&self.parallel_execution);

codex-rs/core/tests/suite/code_mode.rs

@@ -3342,6 +3342,7 @@ text((await tools.exec_command({{cmd: "printf 'phase 3'"}})).output);
     let first_completion = responses::mount_sse_once(
         &server,
         sse(vec![
+            responses::ev_response_created("resp-2"),
             ev_assistant_message("msg-1", "waiting"),
             ev_completed("resp-2"),
         ]),
@@ -3453,8 +3454,8 @@ text((await tools.exec_command({{cmd: "printf 'phase 3'"}})).output);
     );
     assert_eq!(text_item(&third_items, /*index*/ 1), "phase 3");
 
-    // Resuming the cell changes the response id, but its nested calls retain
+    // Nested calls retain the original exec item and turn context when later
-    // the original exec item's identity for request-local freshness accounting.
+    // turns resume the cell with wait.
     let originating_items = observer.originating_items.lock().unwrap();
     let wrapper_item = originating_items[0]
         .1
@@ -3470,15 +3471,12 @@ text((await tools.exec_command({{cmd: "printf 'phase 3'"}})).output);
     let observed = observer.response_ids.lock().unwrap();
     let mut nested_response_ids = observed
         .iter()
+        .skip_while(|(tool, _)| tool != "wait")
         .filter(|(tool, _)| tool == "exec_command")
-        .filter_map(|(_, response_id)| response_id.clone())
+        .map(|(_, response_id)| response_id.clone())
-        .skip_while(|response_id| response_id != &"resp-3".to_owned())
         .collect::<Vec<_>>();
     nested_response_ids.dedup();
-    assert_eq!(
+    assert_eq!(nested_response_ids, vec![Some("resp-2".to_owned())]);
-        nested_response_ids,
-        vec!["resp-3".to_owned(), "resp-5".to_owned()]
-    );
     assert_eq!(
         observed
             .iter()
@@ -6119,6 +6117,93 @@ impl<'call> ToolExecutor<ToolCall<'call>> for NamespacedCustomTool {
     }
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn yielded_code_mode_tool_callbacks_keep_their_originating_step() -> Result<()> {
+    skip_if_no_network!(Ok(()));
+    let server = responses::start_mock_server().await;
+    let mut extensions = ExtensionRegistryBuilder::<Config>::new();
+    extensions.tool_contributor(Arc::new(NamespacedCustomTool {
+        generation: 0,
+        generations: Arc::new(AtomicUsize::new(0)),
+    }));
+    let mut builder = test_codex()
+        .with_extensions(Arc::new(extensions.build()))
+        .with_model_info_override("gpt-5.4", |model| {
+            model.tool_mode = Some(ToolMode::CodeMode);
+            model.experimental_supported_tools = vec!["test_sync_tool".to_string()];
+        })
+        .with_config(|config| {
+            config.features.enable(Feature::CodeMode).unwrap();
+        });
+    let test = builder.build_with_auto_env(&server).await?;
+    // A cannot invoke its editor until B is executing in the next turn.
+    let barrier = r#"await tools.test_sync_tool({barrier: {
+        id: "code-mode-origin-turns", participants: 2, timeout_ms: 60000
+    }});"#;
+    responses::mount_sse_once(&server, sse(vec![
+        ev_response_created("resp-a"),
+        ev_custom_tool_call("call-a", "exec", &format!(
+            "// @exec: {{\"yield_time_ms\": 1}}\n{barrier}\ntext('A_ORIGIN:' + (await tools.editor__apply_patch('A')).generation);"
+        )),
+        ev_completed("resp-a"),
+    ])).await;
+    let yielded = responses::mount_sse_once(
+        &server,
+        sse(vec![
+            ev_assistant_message("msg-a", "waiting for the next turn"),
+            ev_completed("resp-a-done"),
+        ]),
+    )
+    .await;
+    test.submit_turn("start A and leave it running").await?;
+    let cell_id = extract_running_cell_id(text_item(
+        &custom_tool_output_items(&yielded.single_request(), "call-a"),
+        /*index*/ 0,
+    ));
+
+    responses::mount_sse_once(&server, sse(vec![
+        ev_response_created("resp-b"),
+        ev_custom_tool_call("call-b", "exec", &format!(
+            "// @exec: {{\"yield_time_ms\": 60000}}\n{barrier}\ntext('B_ORIGIN:' + (await tools.editor__apply_patch('B')).generation);"
+        )),
+        ev_completed("resp-b"),
+    ])).await;
+    let b_finished = responses::mount_sse_once(
+        &server,
+        sse(vec![
+            ev_response_created("resp-wait"),
+            responses::ev_function_call(
+                "call-wait-a",
+                "wait",
+                &serde_json::json!({
+                    "cell_id": cell_id, "yield_time_ms": 60000,
+                })
+                .to_string(),
+            ),
+            ev_completed("resp-wait"),
+        ]),
+    )
+    .await;
+    let completed = responses::mount_sse_once(
+        &server,
+        sse(vec![
+            ev_assistant_message("msg-b", "both cells finished"),
+            ev_completed("resp-b-done"),
+        ]),
+    )
+    .await;
+    test.submit_turn("run B, then wait for A").await?;
+    assert_eq!(
+        custom_tool_output_body_and_success(&b_finished.single_request(), "call-b").0,
+        "B_ORIGIN:3"
+    );
+    let live = completed.single_request();
+    let a_output = function_tool_output_items(&live, "call-wait-a");
+    assert!(text_item(&a_output, /*index*/ 0).starts_with("Script completed"));
+    assert_eq!(text_item(&a_output, /*index*/ 1), "A_ORIGIN:1");
+    Ok(())
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn code_mode_exposes_and_dispatches_namespaced_custom_tools() -> Result<()> {
     skip_if_no_network!(Ok(()));

codex-rs/core/tests/suite/guardian_mcp_elicitation.rs

@@ -773,3 +773,153 @@ async fn user_review_preserves_unsupported_guardian_elicitations(mode: &str) ->
     test.codex.shutdown_and_wait().await?;
     Ok(())
 }
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn yielded_code_mode_elicitation_keeps_live_invocation_metadata() -> Result<()> {
+    skip_if_no_network!(Ok(()));
+    skip_if_wine_exec!(Ok(()), "the MCP fixture requires a host Python interpreter");
+
+    let server = responses::start_mock_server().await;
+    let meta = json!([{
+        "codex_strict_auto_review": true,
+        "codex_sensitive_action": true,
+        "tool_name": "write_record",
+        "tool_params": {"value": 42},
+        "connector_id": "inner-connector",
+    }]);
+    let mcp_servers = serde_json::from_value(json!({
+        "node_repl": {
+            "command": if cfg!(windows) { "python" } else { "python3" },
+            "args": ["-u", "-c", ELICITATION_SERVER, meta.to_string(), "forward"],
+            "default_tools_approval_mode": "approve",
+        }
+    }))?;
+    let mut builder = test_codex()
+        .with_model_info_override("gpt-5.4", |model| {
+            model.tool_mode = Some(codex_protocol::openai_models::ToolMode::CodeMode);
+            model.experimental_supported_tools = vec!["test_sync_tool".to_string()];
+        })
+        .with_config(move |config| {
+            config
+                .features
+                .enable(codex_features::Feature::CodeMode)
+                .unwrap();
+            config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest);
+            config.approvals_reviewer = ApprovalsReviewer::User;
+            config.mcp_servers.set(mcp_servers).unwrap();
+        });
+    let test = builder.build_with_auto_env(&server).await?;
+    wait_for_mcp_server(&test.codex, "node_repl").await?;
+    let barrier = r#"await tools.test_sync_tool({barrier: {
+        id: "elicitation-origin", participants: 2, timeout_ms: 60000
+    }});"#;
+    responses::mount_sse_once(&server, responses::sse(vec![
+        responses::ev_response_created("origin-a"),
+        responses::ev_custom_tool_call("cell-a", "exec", &format!(
+            "// @exec: {{\"yield_time_ms\": 1}}\n{barrier}\ntext(await tools.mcp__node_repl__js({{code: 'write_record()'}}));"
+        )),
+        responses::ev_completed("origin-a"),
+    ])).await;
+    let yielded = responses::mount_sse_once(
+        &server,
+        responses::sse(vec![responses::ev_completed("a-finished")]),
+    )
+    .await;
+    test.submit_text_turn("Start A and leave its cell running.")
+        .await?;
+    let output = yielded.single_request().custom_tool_call_output("cell-a");
+    let text = output["output"]
+        .as_str()
+        .or_else(|| output["output"][0]["text"].as_str())
+        .context("A should return text")?;
+    let cell_id = text
+        .strip_prefix("Script running with cell ID ")
+        .and_then(|text| text.lines().next())
+        .context("A should yield a running cell")?
+        .to_string();
+
+    responses::mount_sse_once(
+        &server,
+        responses::sse(vec![
+            responses::ev_custom_tool_call("cell-b", "exec", barrier),
+            responses::ev_completed("b-started"),
+        ]),
+    )
+    .await;
+    let guardian = responses::mount_sse_once_match(
+        &server,
+        body_partial_json(json!({"client_metadata": {"x-openai-subagent": "guardian"}})),
+        responses::sse(vec![
+            responses::ev_assistant_message(
+                "review-result",
+                &json!({
+                    "risk_level": "low", "user_authorization": "high", "outcome": "allow",
+                    "rationale": "The user requested the action.",
+                })
+                .to_string(),
+            ),
+            responses::ev_completed("review-done"),
+        ]),
+    )
+    .await;
+    responses::mount_sse_once(
+        &server,
+        responses::sse(vec![
+            responses::ev_function_call(
+                "wait-a",
+                "wait",
+                &json!({
+                    "cell_id": cell_id, "yield_time_ms": 60000,
+                })
+                .to_string(),
+            ),
+            responses::ev_completed("b-wait"),
+        ]),
+    )
+    .await;
+    responses::mount_sse_once(
+        &server,
+        responses::sse(vec![responses::ev_completed("b-finished")]),
+    )
+    .await;
+    test.codex
+        .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text {
+            text: "Run B and wait for A.".into(),
+            text_elements: Vec::new(),
+        }]))
+        .await?;
+    let mut invocation_id = None;
+    let mut assessment_target = None;
+    loop {
+        let event =
+            tokio::time::timeout(Duration::from_secs(60), test.codex.next_event()).await??;
+        match event.msg {
+            EventMsg::ItemCompleted(event) => {
+                if let TurnItem::McpToolCall(item) = event.item {
+                    invocation_id = Some(item.id);
+                }
+            }
+            EventMsg::GuardianAssessment(assessment)
+                if assessment.status == GuardianAssessmentStatus::Approved =>
+            {
+                assessment_target = assessment.target_item_id;
+            }
+            EventMsg::ElicitationRequest(_) => panic!("strict review must not prompt the user"),
+            EventMsg::TurnComplete(_) => break,
+            _ => {}
+        }
+    }
+    assert!(invocation_id.is_some(), "A's tool must complete during B");
+    assert_eq!(assessment_target, invocation_id);
+    assert_eq!(
+        guardian
+            .requests()
+            .iter()
+            .filter(
+                |request| request.body_json()["client_metadata"]["x-openai-subagent"] == "guardian"
+            )
+            .count(),
+        1
+    );
+    Ok(())
+}

codex-rs/core/tests/suite/guardian_review.rs

@@ -2208,3 +2208,170 @@ printf '%s\n' "${@: -1}" >> "${payload_path}""#,
 
     Ok(())
 }
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn yielded_code_mode_denials_interrupt_the_servicing_turn() -> Result<()> {
+    skip_if_no_network!(Ok(()));
+    skip_if_sandbox!(Ok(()));
+    skip_if_wine_exec!(
+        Ok(()),
+        "Guardian approval actions require host-native paths"
+    );
+
+    let server = start_mock_server().await;
+    let mut builder = test_codex()
+        .with_model_info_override("gpt-5.4", |model| {
+            model.tool_mode = Some(codex_protocol::openai_models::ToolMode::CodeMode);
+            model.experimental_supported_tools = vec!["test_sync_tool".to_string()];
+        })
+        .with_config(|config| {
+            config.features.enable(Feature::CodeMode).unwrap();
+            config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest);
+            config.approvals_reviewer = ApprovalsReviewer::AutoReview;
+        });
+    let test = builder.build_with_auto_env(&server).await?;
+    let barrier = r#"await tools.test_sync_tool({barrier: {
+        id: "guardian-origin", participants: 2, timeout_ms: 60000
+    }});"#;
+    let code = format!(
+        r#"// @exec: {{"yield_time_ms": 1}}
+{barrier}
+for (let index = 0; index < 6; index++) {{
+    try {{
+        text(await tools.exec_command({{
+            cmd: `echo review-${{index}}`,
+            sandbox_permissions: "require_escalated",
+            justification: "Exercise originating-cell Guardian reviews."
+        }}));
+    }} catch (error) {{
+        text(String(error));
+    }}
+}}
+"#
+    );
+    core_test_support::responses::mount_sse_once(
+        &server,
+        sse(vec![
+            ev_response_created("origin-a"),
+            core_test_support::responses::ev_custom_tool_call("cell-a", "exec", &code),
+            ev_completed("origin-a"),
+        ]),
+    )
+    .await;
+    let yielded = core_test_support::responses::mount_sse_once(
+        &server,
+        sse(vec![ev_completed("a-finished")]),
+    )
+    .await;
+    test.submit_text_turn("Start A and leave its cell running.")
+        .await?;
+    let output = yielded.single_request().custom_tool_call_output("cell-a");
+    let text = output["output"]
+        .as_str()
+        .or_else(|| output["output"][0]["text"].as_str())
+        .context("A should return text")?;
+    let cell_id = text
+        .strip_prefix("Script running with cell ID ")
+        .and_then(|text| text.lines().next())
+        .context("A should yield a running cell")?
+        .to_string();
+
+    core_test_support::responses::mount_sse_once(
+        &server,
+        sse(vec![
+            core_test_support::responses::ev_custom_tool_call("cell-b", "exec", barrier),
+            ev_completed("b-started"),
+        ]),
+    )
+    .await;
+    let mut reviews = Vec::new();
+    // An approval resets B's consecutive count; only the final three denials interrupt it.
+    for outcome in ["deny", "deny", "allow", "deny", "deny", "deny"] {
+        reviews.push(
+            core_test_support::responses::mount_sse_once_match(
+                &server,
+                wiremock::matchers::body_partial_json(json!({
+                    "client_metadata": {"x-openai-subagent": "guardian"}
+                })),
+                sse(vec![
+                    ev_assistant_message(
+                        "review",
+                        &json!({
+                            "risk_level": if outcome == "allow" { "low" } else { "high" },
+                            "user_authorization": "high", "outcome": outcome,
+                            "rationale": "Test decision for the delayed cell.",
+                        })
+                        .to_string(),
+                    ),
+                    ev_completed("review-done"),
+                ]),
+            )
+            .await,
+        );
+    }
+    core_test_support::responses::mount_sse_once(
+        &server,
+        sse(vec![
+            ev_function_call(
+                "wait-a",
+                "wait",
+                &json!({
+                    "cell_id": cell_id, "yield_time_ms": 60000,
+                })
+                .to_string(),
+            ),
+            ev_completed("b-wait"),
+        ]),
+    )
+    .await;
+    test.codex
+        .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text {
+            text: "Run B and wait for A.".into(),
+            text_elements: Vec::new(),
+        }]))
+        .await?;
+    let mut active_id = None;
+    let mut warning_id = None;
+    loop {
+        let event =
+            tokio::time::timeout(Duration::from_secs(60), test.codex.next_event()).await??;
+        match event.msg {
+            EventMsg::TurnStarted(_) => active_id = Some(event.id),
+            EventMsg::GuardianWarning(warning)
+                if warning.message.contains("too many approval requests") =>
+            {
+                assert!(
+                    warning
+                        .message
+                        .contains("3 consecutive, 5 in the last 50 reviews"),
+                    "{}",
+                    warning.message
+                );
+                warning_id = Some(event.id);
+            }
+            EventMsg::TurnAborted(aborted) => {
+                assert_eq!(aborted.reason, TurnAbortReason::Interrupted);
+                assert_eq!(Some(event.id), active_id);
+                break;
+            }
+            EventMsg::TurnComplete(_) => panic!("B must be interrupted by A's denials"),
+            _ => {}
+        }
+    }
+    assert!(active_id.is_some());
+    assert_eq!(warning_id, active_id);
+    for review in reviews {
+        assert_eq!(
+            review
+                .requests()
+                .iter()
+                .filter(
+                    |request| request.body_json()["client_metadata"]["x-openai-subagent"]
+                        == "guardian"
+                )
+                .count(),
+            1
+        );
+    }
+    Ok(())
+}

codex-rs/core/tests/suite/rmcp_client.rs

@@ -2498,6 +2498,12 @@ async fn stdio_image_responses_round_trip() -> anyhow::Result<()> {
         .start_or_steer_turn(read_only_user_turn(&fixture, "call the rmcp image tool"))
         .await?;
 
+    let turn_id = core_test_support::wait_for_event_match(&fixture.codex, |event| match event {
+        EventMsg::TurnStarted(started) => Some(started.turn_id.clone()),
+        _ => None,
+    })
+    .await;
+
     // Wait for tool begin/end and final completion.
     let begin_event = wait_for_event(&fixture.codex, |ev| {
         matches!(ev, EventMsg::McpToolCallBegin(_))
@@ -2509,6 +2515,7 @@ async fn stdio_image_responses_round_trip() -> anyhow::Result<()> {
     assert_eq!(
         begin,
         McpToolCallBeginEvent {
+            turn_id: turn_id.clone(),
             call_id: call_id.to_string(),
             invocation: McpInvocation {
                 server: server_name.to_string(),
@@ -2532,6 +2539,7 @@ async fn stdio_image_responses_round_trip() -> anyhow::Result<()> {
     let EventMsg::McpToolCallEnd(end) = end_event else {
         unreachable!("end");
     };
+    assert_eq!(end.turn_id, turn_id);
     assert_eq!(end.call_id, call_id);
     assert_eq!(
         end.invocation,

codex-rs/protocol/src/legacy_events.rs

@@ -472,9 +472,10 @@ impl FileChangeItem {
 }
 
 impl McpToolCallItem {
-    pub fn as_legacy_begin_event(&self) -> EventMsg {
+    pub fn as_legacy_begin_event(&self, turn_id: String) -> EventMsg {
         EventMsg::McpToolCallBegin(McpToolCallBeginEvent {
             call_id: self.id.clone(),
+            turn_id,
             invocation: McpInvocation {
                 server: self.server.clone(),
                 tool: self.tool.clone(),
@@ -490,7 +491,7 @@ impl McpToolCallItem {
         })
     }
 
-    pub fn as_legacy_end_event(&self) -> Option<EventMsg> {
+    pub fn as_legacy_end_event(&self, turn_id: String) -> Option<EventMsg> {
         let result = match (&self.result, &self.error) {
             (Some(result), _) => Ok(result.clone()),
             (None, Some(error)) => Err(error.message.clone()),
@@ -499,6 +500,7 @@ impl McpToolCallItem {
 
         Some(EventMsg::McpToolCallEnd(McpToolCallEndEvent {
             call_id: self.id.clone(),
+            turn_id,
             invocation: McpInvocation {
                 server: self.server.clone(),
                 tool: self.tool.clone(),
@@ -542,7 +544,10 @@ impl TurnItem {
                 .as_legacy_end_event(String::new())
                 .into_iter()
                 .collect(),
-            TurnItem::McpToolCall(item) => item.as_legacy_end_event().into_iter().collect(),
+            TurnItem::McpToolCall(item) => item
+                .as_legacy_end_event(String::new())
+                .into_iter()
+                .collect(),
             TurnItem::Reasoning(item) => item.as_legacy_events(show_raw_agent_reasoning),
             TurnItem::ContextCompaction(item) => vec![item.as_legacy_event()],
         }
@@ -562,7 +567,7 @@ impl HasLegacyEvent for ItemStartedEvent {
                 })]
             }
             TurnItem::FileChange(item) => vec![item.as_legacy_begin_event(self.turn_id.clone())],
-            TurnItem::McpToolCall(item) => vec![item.as_legacy_begin_event()],
+            TurnItem::McpToolCall(item) => vec![item.as_legacy_begin_event(self.turn_id.clone())],
             TurnItem::CommandExecution(item) => {
                 vec![item.as_legacy_begin_event(self.turn_id.clone(), self.started_at_ms)]
             }
@@ -581,6 +586,10 @@ impl HasLegacyEvent for ItemStartedEvent {
 impl HasLegacyEvent for ItemCompletedEvent {
     fn as_legacy_events(&self, show_raw_agent_reasoning: bool) -> Vec<EventMsg> {
         match &self.item {
+            TurnItem::McpToolCall(item) => item
+                .as_legacy_end_event(self.turn_id.clone())
+                .into_iter()
+                .collect(),
             TurnItem::FileChange(item) => item
                 .as_legacy_end_event(self.turn_id.clone())
                 .into_iter()

codex-rs/protocol/src/protocol.rs

@@ -2617,6 +2617,9 @@ pub struct McpInvocation {
 pub struct McpToolCallBeginEvent {
     /// Identifier so this can be paired with the McpToolCallEnd event.
     pub call_id: String,
+    /// Originating turn; absent in older rollout records.
+    #[serde(default)]
+    pub turn_id: String,
     pub invocation: McpInvocation,
     #[serde(default, skip_serializing_if = "Option::is_none")]
     #[ts(optional)]
@@ -2646,6 +2649,9 @@ pub struct McpToolCallBeginEvent {
 pub struct McpToolCallEndEvent {
     /// Identifier for the corresponding McpToolCallBegin that finished.
     pub call_id: String,
+    /// Originating turn; absent in older rollout records.
+    #[serde(default)]
+    pub turn_id: String,
     pub invocation: McpInvocation,
     #[serde(default, skip_serializing_if = "Option::is_none")]
     #[ts(optional)]
@@ -5361,6 +5367,7 @@ mod tests {
         assert_eq!(legacy_events.len(), 1);
         match &legacy_events[0] {
             EventMsg::McpToolCallBegin(event) => {
+                assert_eq!(event.turn_id, "turn-1");
                 assert_eq!(event.call_id, "mcp-1");
                 assert_eq!(event.invocation.server, "server");
                 assert_eq!(event.invocation.tool, "tool");
@@ -5486,6 +5493,7 @@ mod tests {
         assert_eq!(legacy_events.len(), 1);
         match &legacy_events[0] {
             EventMsg::McpToolCallEnd(event) => {
+                assert_eq!(event.turn_id, "turn-1");
                 assert_eq!(event.call_id, "mcp-1");
                 assert_eq!(event.invocation.server, "server");
                 assert_eq!(event.invocation.tool, "tool");

codex-rs/thread-store/src/local/rollout_migration/legacy_event.rs

@@ -146,7 +146,7 @@ pub(super) fn completed_item(
                     error,
                     duration: Some(event.duration),
                 }),
-                None,
+                (!event.turn_id.is_empty()).then(|| event.turn_id.clone()),
             ))
         }
         EventMsg::WebSearchEnd(event) => Some((

codex-rs/thread-store/src/local/rollout_migration_tests.rs

@@ -507,6 +507,18 @@ async fn migration_keeps_late_completions_in_their_original_turn() {
             started("current"),
             user_message("current question"),
             exec_completion("old", "call-old"),
+            serde_json::from_value(json!({
+                "type": "event_msg",
+                "payload": {
+                    "type": "mcp_tool_call_end",
+                    "call_id": "mcp-old",
+                    "turn_id": "old",
+                    "invocation": {"server": "slack", "tool": "search", "arguments": {}},
+                    "duration": {"secs": 0, "nanos": 0},
+                    "result": {"Ok": {"content": []}}
+                }
+            }))
+            .expect("build late legacy MCP completion"),
             agent_message("current answer"),
             completed("current"),
         ],
@@ -545,7 +557,10 @@ async fn migration_keeps_late_completions_in_their_original_turn() {
         .iter()
         .map(|item| item.turn_id.as_str())
         .collect::<Vec<_>>();
-    assert_eq!(item_turn_ids, vec!["old", "current", "old", "current"]);
+    assert_eq!(
+        item_turn_ids,
+        vec!["old", "current", "old", "old", "current"]
+    );
 }
 
 #[tokio::test]