diff --git a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst index 88f160cc9e..b80ae23637 100644 Binary files a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst and b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst differ diff --git a/codex-rs/app-server-protocol/src/protocol/thread_history.rs b/codex-rs/app-server-protocol/src/protocol/thread_history.rs index 75dacd5dc8..44ebe0d62b 100644 --- a/codex-rs/app-server-protocol/src/protocol/thread_history.rs +++ b/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![ diff --git a/codex-rs/app-server/tests/suite/v2/thread_resume.rs b/codex-rs/app-server/tests/suite/v2/thread_resume.rs index 69da60d821..ba2133a2b5 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_resume.rs +++ b/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(), diff --git a/codex-rs/core/src/guardian/decision.rs b/codex-rs/core/src/guardian/decision.rs index 6709fb2e67..a1d0026317 100644 --- a/codex-rs/core/src/guardian/decision.rs +++ b/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, diff --git a/codex-rs/core/src/guardian/review.rs b/codex-rs/core/src/guardian/review.rs index 641b6cc175..ff2a94763a 100644 --- a/codex-rs/core/src/guardian/review.rs +++ b/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, turn_id: &str) { +pub(super) async fn record_guardian_non_denial(session: &Arc) { + 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, turn: &Arc, turn_id: &str) { +async fn record_guardian_denial(session: &Arc) { + 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, turn: &Arc, } #[cfg(test)] -pub(crate) async fn record_guardian_denial_for_test( - session: &Arc, - turn: &Arc, - turn_id: &str, -) { - record_guardian_denial(session, turn, turn_id).await; +pub(crate) async fn record_guardian_denial_for_test(session: &Arc) { + record_guardian_denial(session).await; } #[derive(Clone)] diff --git a/codex-rs/core/src/guardian/review_request.rs b/codex-rs/core/src/guardian/review_request.rs index 7683c8834f..5e117b2342 100644 --- a/codex-rs/core/src/guardian/review_request.rs +++ b/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 } diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index 8b57655124..50244fd34b 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/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()) - .await; + let _approval_metadata = + 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, - ) { - let Some(turn_state) = self - .input_queue - .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(), + ) -> Arc<(Option, McpToolApprovalMetadata)> { + let key = (invocation.server.clone(), call_id.to_string()); + let metadata = Arc::new(( (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, McpToolApprovalMetadata)> { - let turn_state = self - .input_queue - .turn_state_for_sub_id(&self.active_turn, sub_id) - .await?; - - turn_state.lock().await.mcp_tool_approval_metadata(call_id) + self.mcp_tool_approval_metadata + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .get(&(server.to_string(), call_id.to_string())) + .and_then(std::sync::Weak::upgrade) + .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) - .await + .mcp_tool_approval_metadata(server, id) .and_then(|(_, metadata)| metadata.link_id); let metadata = McpToolApprovalMetadata { annotations: None, diff --git a/codex-rs/core/src/mcp_tool_call_tests.rs b/codex-rs/core/src/mcp_tool_call_tests.rs index 980ed317c3..3233728975 100644 --- a/codex-rs/core/src/mcp_tool_call_tests.rs +++ b/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::>(); + 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() + ), + ], + ); +} diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 12bfee4acb..18a8d083c7 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/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 - .mcp_tool_approval_metadata(&turn_context.sub_id, call_id) - .await + && let Some((Some(invocation), _)) = + session.mcp_tool_approval_metadata(&request.server_name, call_id) && 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 - .mcp_tool_approval_metadata(&turn_context.sub_id, call_id) - .await + let Some((Some(invocation), metadata)) = + session.mcp_tool_approval_metadata(&request.server_name, call_id) else { return Ok(None); }; diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index aebfd641e5..7dd4fe066a 100644 --- a/codex-rs/core/src/session/session.rs +++ b/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, 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, /// 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, pub(super) mcp_elicitation_reviewer_handle: OnceLock, pub(super) mcp_elicitation_lifecycle_handle: OnceLock, 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(), diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index b77be1ae0d..0152b7b043 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/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, session: Arc, - ctx: Arc, + _ctx: Arc, _input: Vec, 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( - &session_for_review, - &turn_for_review, - &turn_id, - ) - .await; + crate::guardian::record_guardian_denial_for_test(&session_for_review).await; } }); }); diff --git a/codex-rs/core/src/state/turn.rs b/codex-rs/core/src/state/turn.rs index 324b39ae46..01b8756e1f 100644 --- a/codex-rs/core/src/state/turn.rs +++ b/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, pending_user_input: HashMap>, pending_elicitations: HashMap<(String, RequestId), oneshot::Sender>, - mcp_tool_approval_metadata: HashMap, McpToolApprovalMetadata)>, pending_dynamic_tools: HashMap>, 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, - 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, McpToolApprovalMetadata)> { - self.mcp_tool_approval_metadata.get(call_id).cloned() - } - pub(crate) fn insert_pending_dynamic_tool( &mut self, key: String, diff --git a/codex-rs/core/src/tools/code_mode/delegate.rs b/codex-rs/core/src/tools/code_mode/delegate.rs index ac8d0505ea..607fd086ff 100644 --- a/codex-rs/core/src/tools/code_mode/delegate.rs +++ b/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, + pub(super) step_context: Arc, +} + struct CellDispatchGate { ready: watch::Sender, // 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, step_context: Arc, tracker: SharedTurnDiffTracker, ) -> CodeModeDispatchWorker { - let track_completeness = ExecutedToolCalls::is_enabled(&exec.turn.config.features); - let tool_runtime = ToolCallRuntime::new(Arc::clone(&exec.session), step_context, tracker); - let host = Arc::new(CoreTurnHost { exec, tool_runtime }); + let tool_runtime = ToolCallRuntime::new(Arc::clone(&session), step_context, tracker); + let host = Arc::new(CoreTurnHost { + 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, dispatch_trace: Box, cancellation_token: CancellationToken, response_tx: oneshot::Sender>, @@ -395,7 +421,7 @@ impl Drop for CodeModeDispatchWorker { } struct CoreTurnHost { - exec: ExecContext, + session: Arc, tool_runtime: ToolCallRuntime, } @@ -403,11 +429,13 @@ impl CoreTurnHost { fn submit_tool( &self, invocation: CodeModeNestedToolCall, + step_context: Arc, call_id: String, cancellation_token: CancellationToken, ) -> impl std::future::Future> + 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 - .session + self.session .inject_if_running(vec![ResponseItem::CustomToolCallOutput { id: None, call_id, diff --git a/codex-rs/core/src/tools/code_mode/delegate_tests.rs b/codex-rs/core/src/tools/code_mode/delegate_tests.rs index 70340147c7..1dc6201623 100644 --- a/codex-rs/core/src/tools/code_mode/delegate_tests.rs +++ b/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); @@ -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 { - session: Arc::new(session), - turn, - }, + Arc::new(session), 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()), + ); +} diff --git a/codex-rs/core/src/tools/code_mode/execute_handler.rs b/codex-rs/core/src/tools/code_mode/execute_handler.rs index f7ceb3c7c7..d131019fd5 100644 --- a/codex-rs/core/src/tools/code_mode/execute_handler.rs +++ b/codex-rs/core/src/tools/code_mode/execute_handler.rs @@ -34,7 +34,7 @@ impl CodeModeExecuteHandler { async fn execute( &self, session: std::sync::Arc, - step_context: std::sync::Arc, + step_context: Arc, call_id: String, originating_item_id: Option, code: String, @@ -71,13 +71,16 @@ impl CodeModeExecuteHandler { .session .services .code_mode_service - .execute(codex_code_mode::ExecuteRequest { - tool_call_id: call_id.clone(), - enabled_tools, - source: args.code.clone(), - yield_time_ms: args.yield_time_ms, - max_output_tokens: args.max_output_tokens, - }) + .execute( + codex_code_mode::ExecuteRequest { + tool_call_id: call_id.clone(), + enabled_tools, + source: args.code.clone(), + 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(); diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index 355f0a195c..ab100b4d9f 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/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, ) -> Result { request .yield_time_ms .get_or_insert(self.default_exec_yield_time_ms); - self.session() - .await? - .execute(request, self.dispatch_broker.clone()) - .await + let delegate = Arc::new(CodeModeCellDelegate { + broker: Arc::clone(&self.dispatch_broker), + step_context, + }); + self.session().await?.execute(request, delegate).await } pub(crate) async fn wait( @@ -209,18 +212,13 @@ impl CodeModeService { step_context: Arc, tracker: SharedTurnDiffTracker, ) -> Option { - 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, + step_context: Arc, 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 turn_id = exec.turn.sub_id.clone(); + let thread_id = session.thread_id; + 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(), - turn_id: exec.turn.sub_id.clone(), + thread_id: session.thread_id.to_string(), + 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(), diff --git a/codex-rs/core/src/tools/parallel.rs b/codex-rs/core/src/tools/parallel.rs index e43aa5ff4d..8cbc861b0f 100644 --- a/codex-rs/core/src/tools/parallel.rs +++ b/codex-rs/core/src/tools/parallel.rs @@ -79,7 +79,9 @@ impl ToolCallRuntime { ) -> impl std::future::Future> { 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, call: ToolCall, source: ToolCallSource, cancellation_token: CancellationToken, ) -> impl std::future::Future> { - self.session.services.executed_tool_calls.record_tool_call( - &call, - &source, - &self.step_context, - ); - let router = &self.step_context.tool_router; + self.session + .services + .executed_tool_calls + .record_tool_call(&call, &source, &step_context); + let router = &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); diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index 4243ae75da..e4950170e0 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/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 - // the original exec item's identity for request-local freshness accounting. + // Nested calls retain the original exec item and turn context when later + // 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()) - .skip_while(|response_id| response_id != &"resp-3".to_owned()) + .map(|(_, response_id)| response_id.clone()) .collect::>(); nested_response_ids.dedup(); - assert_eq!( - nested_response_ids, - vec!["resp-3".to_owned(), "resp-5".to_owned()] - ); + assert_eq!(nested_response_ids, vec![Some("resp-2".to_owned())]); assert_eq!( observed .iter() @@ -6119,6 +6117,93 @@ impl<'call> ToolExecutor> 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::::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(())); diff --git a/codex-rs/core/tests/suite/guardian_mcp_elicitation.rs b/codex-rs/core/tests/suite/guardian_mcp_elicitation.rs index 538ce35d03..914acdb9fd 100644 --- a/codex-rs/core/tests/suite/guardian_mcp_elicitation.rs +++ b/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(()) +} diff --git a/codex-rs/core/tests/suite/guardian_review.rs b/codex-rs/core/tests/suite/guardian_review.rs index 245fba5db8..466bb77b7e 100644 --- a/codex-rs/core/tests/suite/guardian_review.rs +++ b/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(()) +} diff --git a/codex-rs/core/tests/suite/rmcp_client.rs b/codex-rs/core/tests/suite/rmcp_client.rs index 2e403ded07..93765c59ca 100644 --- a/codex-rs/core/tests/suite/rmcp_client.rs +++ b/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, diff --git a/codex-rs/protocol/src/legacy_events.rs b/codex-rs/protocol/src/legacy_events.rs index 2ef68b9b1b..bfea90b558 100644 --- a/codex-rs/protocol/src/legacy_events.rs +++ b/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 { + pub fn as_legacy_end_event(&self, turn_id: String) -> Option { 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 { 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() diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 0c79853a35..81d799b615 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/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"); diff --git a/codex-rs/thread-store/src/local/rollout_migration/legacy_event.rs b/codex-rs/thread-store/src/local/rollout_migration/legacy_event.rs index a6eb63d75a..27d4c885d0 100644 --- a/codex-rs/thread-store/src/local/rollout_migration/legacy_event.rs +++ b/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(( diff --git a/codex-rs/thread-store/src/local/rollout_migration_tests.rs b/codex-rs/thread-store/src/local/rollout_migration_tests.rs index 8dfad90b9e..bad11829c6 100644 --- a/codex-rs/thread-store/src/local/rollout_migration_tests.rs +++ b/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::>(); - assert_eq!(item_turn_ids, vec!["old", "current", "old", "current"]); + assert_eq!( + item_turn_ids, + vec!["old", "current", "old", "old", "current"] + ); } #[tokio::test]