diff --git a/codex-rs/app-server/src/bespoke_event_handling.rs b/codex-rs/app-server/src/bespoke_event_handling.rs index 9ef4657a36..da4d0e1050 100644 --- a/codex-rs/app-server/src/bespoke_event_handling.rs +++ b/codex-rs/app-server/src/bespoke_event_handling.rs @@ -3,6 +3,7 @@ use crate::outgoing_message::ClientRequestResult; use crate::outgoing_message::ThreadScopedOutgoingMessageSender; use crate::request_processors::thread_settings_from_config_snapshot; use crate::server_request_error::is_turn_transition_server_request_error; +use crate::thread_state::PendingNetworkApproval; use crate::thread_state::ThreadState; use crate::thread_state::TurnSummary; use crate::thread_state::resolve_server_request_on_thread_listener; @@ -779,6 +780,7 @@ pub(crate) async fn apply_bespoke_event_handling( let additional_permissions = additional_permissions.map(V2AdditionalPermissionProfile::from); + let is_network_approval = network_approval_context.is_some(); let params = CommandExecutionRequestApprovalParams { kind: kind.into(), thread_id: conversation_id.to_string(), @@ -802,6 +804,18 @@ pub(crate) async fn apply_bespoke_event_handling( params, )) .await; + let permission_guard = if is_network_approval { + thread_state.lock().await.network_approvals.insert( + approval_id.as_ref().unwrap_or(&call_id).clone(), + PendingNetworkApproval { + request_id: pending_request_id.clone(), + permission_guard, + }, + ); + None + } else { + Some(permission_guard) + }; tokio::spawn(async move { on_command_execution_request_approval_response( event_turn_id, @@ -813,12 +827,23 @@ pub(crate) async fn apply_bespoke_event_handling( rx, conversation, outgoing, - thread_state.clone(), + thread_state, permission_guard, ) .await; }); } + EventMsg::NetworkApprovalResolved { approval_id } => { + let pending = thread_state + .lock() + .await + .network_approvals + .remove(&approval_id); + if let Some(pending) = pending { + outgoing.resolve_server_request(pending.request_id).await; + drop(pending.permission_guard); + } + } EventMsg::RequestUserInput(request) => { let user_input_guard = thread_watch_manager .note_user_input_requested(&conversation_id.to_string()) @@ -1305,6 +1330,12 @@ pub(crate) async fn apply_bespoke_event_handling( .await; } EventMsg::ShutdownComplete => { + // The proxy waiter may send its resolution after this last core event. + let approvals = std::mem::take(&mut thread_state.lock().await.network_approvals); + for pending in approvals.into_values() { + outgoing.resolve_server_request(pending.request_id).await; + drop(pending.permission_guard); + } thread_watch_manager .note_thread_shutdown(&conversation_id.to_string()) .await; @@ -2034,11 +2065,22 @@ async fn on_command_execution_request_approval_response( conversation: Arc, outgoing: ThreadScopedOutgoingMessageSender, thread_state: Arc>, - permission_guard: ThreadWatchActiveGuard, + permission_guard: Option, ) { let response = receiver.await; - resolve_server_request_on_thread_listener(&thread_state, pending_request_id).await; - drop(permission_guard); + // Core resolves network prompts in event order, including when the execution disappears. + if let Some(permission_guard) = permission_guard { + resolve_server_request_on_thread_listener(&thread_state, pending_request_id).await; + drop(permission_guard); + } else if response.is_err() + && !thread_state + .lock() + .await + .network_approvals + .contains_key(approval_id.as_ref().unwrap_or(&item_id)) + { + return; + } let (decision, completion_status) = match response { Ok(Ok(value)) => { match serde_json::from_value::(value) { diff --git a/codex-rs/app-server/src/outgoing_message.rs b/codex-rs/app-server/src/outgoing_message.rs index d07a14830e..e1a0151eac 100644 --- a/codex-rs/app-server/src/outgoing_message.rs +++ b/codex-rs/app-server/src/outgoing_message.rs @@ -16,6 +16,7 @@ use codex_app_server_protocol::ServerNotification; use codex_app_server_protocol::ServerNotificationEnvelope; use codex_app_server_protocol::ServerRequest; use codex_app_server_protocol::ServerRequestPayload; +use codex_app_server_protocol::ServerRequestResolvedNotification; use codex_app_server_protocol::ServerResponse; use codex_diagnostics::Gauge; use codex_diagnostics::GaugeGuard; @@ -228,6 +229,17 @@ impl ThreadScopedOutgoingMessageSender { self.outgoing.send_server_notification(notification).await; } + pub(crate) async fn resolve_server_request(&self, request_id: RequestId) { + self.outgoing.cancel_request(&request_id).await; + self.send_server_notification(ServerNotification::ServerRequestResolved( + ServerRequestResolvedNotification { + thread_id: self.thread_id.to_string(), + request_id, + }, + )) + .await; + } + pub(crate) async fn abort_pending_server_requests(&self) { self.outgoing .cancel_requests_for_thread( diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index 8162ca9f7c..a14a448fae 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -179,7 +179,6 @@ use codex_app_server_protocol::SandboxMode; use codex_app_server_protocol::SendAddCreditsNudgeEmailParams; use codex_app_server_protocol::SendAddCreditsNudgeEmailResponse; use codex_app_server_protocol::ServerNotification; -use codex_app_server_protocol::ServerRequestResolvedNotification; use codex_app_server_protocol::SkillSummary; use codex_app_server_protocol::SkillsConfigWriteParams; use codex_app_server_protocol::SkillsConfigWriteResponse; diff --git a/codex-rs/app-server/src/request_processors/thread_lifecycle.rs b/codex-rs/app-server/src/request_processors/thread_lifecycle.rs index a6d8250d60..555671fb78 100644 --- a/codex-rs/app-server/src/request_processors/thread_lifecycle.rs +++ b/codex-rs/app-server/src/request_processors/thread_lifecycle.rs @@ -1317,7 +1317,6 @@ pub(super) async fn resolve_pending_server_request( outgoing: &Arc, request_id: RequestId, ) { - let thread_id = conversation_id.to_string(); let subscribed_connection_ids = thread_state_manager .subscribed_connection_ids(conversation_id) .await; @@ -1326,14 +1325,7 @@ pub(super) async fn resolve_pending_server_request( subscribed_connection_ids, conversation_id, ); - outgoing - .send_server_notification(ServerNotification::ServerRequestResolved( - ServerRequestResolvedNotification { - thread_id, - request_id, - }, - )) - .await; + outgoing.resolve_server_request(request_id).await; } pub(super) fn merge_turn_history_with_active_turn(turns: &mut Vec, active_turn: Turn) { diff --git a/codex-rs/app-server/src/request_processors/thread_listener_shutdown_tests.rs b/codex-rs/app-server/src/request_processors/thread_listener_shutdown_tests.rs index 9b7fd97c38..bbb872b35d 100644 --- a/codex-rs/app-server/src/request_processors/thread_listener_shutdown_tests.rs +++ b/codex-rs/app-server/src/request_processors/thread_listener_shutdown_tests.rs @@ -11,6 +11,7 @@ use crate::outgoing_message::OutgoingEnvelope; use crate::outgoing_message::OutgoingMessage; use crate::thread_sample::SampleResponseRegistry; use crate::thread_state::ConnectionCapabilities; +use crate::thread_state::PendingNetworkApproval; use crate::thread_state::resolve_server_request_on_thread_listener; use anyhow::Context; use anyhow::Result; @@ -83,7 +84,7 @@ async fn listener_shutdown_releases_runtime_and_preserves_other_subscribers( test.thread_store.clone(), Default::default(), thread_state_manager.clone(), - thread_watch_manager, + thread_watch_manager.clone(), Arc::new(Semaphore::new(/*permits*/ 1)), thread_goal_processor, /*state_db*/ None, @@ -106,10 +107,28 @@ async fn listener_shutdown_releases_runtime_and_preserves_other_subscribers( .connection_initialized(connection_id, ConnectionCapabilities::default()) .await; } + let network_request_id = RequestId::Integer(4); if attach_before_shutdown { processor .try_attach_thread_listener(thread_id, Arc::clone(&thread), connection_ids.clone()) .await; + // Leave the prompt outstanding so shutdown runs before the proxy waiter resolves it. + let permission_guard = thread_watch_manager + .note_permission_requested(&thread_id.to_string()) + .await; + thread_state_manager + .thread_state(thread_id) + .await + .lock() + .await + .network_approvals + .insert( + "network-approval".to_string(), + PendingNetworkApproval { + request_id: network_request_id.clone(), + permission_guard, + }, + ); } thread.shutdown_and_wait().await?; let shutdown_result = tokio::time::timeout(Duration::from_secs(/*secs*/ 10), async { @@ -135,8 +154,23 @@ async fn listener_shutdown_releases_runtime_and_preserves_other_subscribers( .try_attach_thread_listener(thread_id, Arc::clone(&thread), connection_ids) .await; } + let mut network_resolved = false; loop { match outgoing_rx.recv().await.expect("outgoing channel closed") { + OutgoingEnvelope::ToConnection { + connection_id: recipient, + message: OutgoingMessage::AppServerNotification(notification), + .. + } if recipient == connection_id + && matches!( + ¬ification.notification, + ServerNotification::ServerRequestResolved(resolved) + if resolved.request_id == network_request_id + ) => + { + assert!(!network_resolved, "duplicate network prompt resolution"); + network_resolved = true; + } OutgoingEnvelope::Broadcast { message: OutgoingMessage::AppServerNotification(notification), } if matches!( @@ -150,6 +184,7 @@ async fn listener_shutdown_releases_runtime_and_preserves_other_subscribers( _ => {} } } + assert_eq!(network_resolved, attach_before_shutdown); drop(thread); while runtime.upgrade().is_some() { tokio::task::yield_now().await; diff --git a/codex-rs/app-server/src/thread_state.rs b/codex-rs/app-server/src/thread_state.rs index 6189a05c47..e47c9ed81c 100644 --- a/codex-rs/app-server/src/thread_state.rs +++ b/codex-rs/app-server/src/thread_state.rs @@ -94,6 +94,12 @@ pub(crate) enum ThreadListenerCommand { }, } +/// Retained until core resolves or withdraws the prompt, even if the client replies first. +pub(crate) struct PendingNetworkApproval { + pub(crate) request_id: RequestId, + pub(crate) permission_guard: crate::thread_status::ThreadWatchActiveGuard, +} + /// Per-conversation accumulation of the latest states e.g. error message while a turn runs. #[derive(Default, Clone)] pub(crate) struct TurnSummary { @@ -120,6 +126,7 @@ impl TurnSummary { pub(crate) struct ThreadState { goal_resume_lock: Arc>, pub(crate) pending_interrupts: PendingInterruptQueue, + pub(crate) network_approvals: HashMap, pub(crate) turn_summary: TurnSummary, pub(crate) last_terminal_turn_id: Option, /// Lets an internal runtime replacement wait until the old listener has processed Core's diff --git a/codex-rs/app-server/tests/suite/v2/thread_resume_network_tests.rs b/codex-rs/app-server/tests/suite/v2/thread_resume_network_tests.rs index d17f9bb7cb..29f3941b14 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_resume_network_tests.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_resume_network_tests.rs @@ -17,6 +17,7 @@ use codex_app_server_protocol::NetworkPolicyRuleAction; use codex_app_server_protocol::SandboxMode; use codex_app_server_protocol::SandboxPolicy; use codex_app_server_protocol::ServerRequest; +use codex_app_server_protocol::ServerRequestResolvedNotification; use codex_app_server_protocol::ThreadResumeParams; use codex_app_server_protocol::ThreadResumeResponse; use codex_app_server_protocol::ThreadStartParams; @@ -32,6 +33,7 @@ use core_test_support::skip_if_remote; use pretty_assertions::assert_eq; use serde_json::json; use std::time::Duration; +use tempfile::NamedTempFile; use tempfile::TempDir; use tokio::time::timeout; @@ -150,6 +152,8 @@ allow_local_binding = true .to_str() .context("test executable must be UTF-8")?; let probe_timeout = Duration::from_secs(/* secs */ 30); + let disconnect_signal = NamedTempFile::new()?; + disconnect_signal.as_file().lock()?; for (call_id, host, expected_approvals) in [ ("remembered-domain", "127.0.0.2", 0), ("new-domain", "127.0.0.3", 1), @@ -177,17 +181,34 @@ allow_local_binding = true "--nocapture", "--test-threads=1", ])?; - responses::mount_sse_once( - &server, - responses::sse(vec![ - responses::ev_exec_command_call_with_args( - call_id, - &json!({"cmd": command, "shell": "/bin/sh", "login": false, "yield_time_ms": probe_timeout.as_millis()}), - ), - responses::ev_completed(call_id), - ]), - ) - .await; + let mut probes = Vec::new(); + if call_id == "new-domain" { + let signal = format!( + "CODEX_TEST_DISCONNECT_NETWORK_PROBE={}", + disconnect_signal.path().display() + ); + probes.push(( + "withdrawn-probe", + format!("env {} {command}", shlex::try_quote(&signal)?), + )); + } + probes.push((call_id, command)); + for (probe_id, command) in probes { + let args = json!({ + "cmd": command, + "shell": "/bin/sh", + "login": false, + "yield_time_ms": probe_timeout.as_millis(), + }); + responses::mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_exec_command_call_with_args(probe_id, &args), + responses::ev_completed(probe_id), + ]), + ) + .await; + } let final_response = responses::mount_sse_once( &server, create_final_assistant_message_sse_response("Network probe complete")?, @@ -205,6 +226,8 @@ allow_local_binding = true .await?; let started: TurnStartResponse = app.read_response(turn_id).await?; let mut approvals = 0; + let mut withdrawn_request = None; + let mut sent_late_cancel = false; loop { match timeout(probe_timeout, app.read_next_message()).await?? { JSONRPCMessage::Request(request) => { @@ -214,11 +237,6 @@ allow_local_binding = true else { anyhow::bail!("unexpected approval request: {request:?}"); }; - approvals += 1; - assert_eq!( - approvals, expected_approvals, - "unexpected approval for {call_id}" - ); assert_eq!( params.network_approval_context, Some(NetworkApprovalContext { @@ -226,6 +244,19 @@ allow_local_binding = true protocol: NetworkApprovalProtocol::Http }), ); + if call_id == "new-domain" { + if withdrawn_request.is_none() { + withdrawn_request = Some(request_id); + disconnect_signal.as_file().unlock()?; + continue; + } + assert!(sent_late_cancel, "old approval still visible"); + } + approvals += 1; + assert_eq!( + approvals, expected_approvals, + "unexpected approval for {call_id}" + ); app.send_response( request_id, serde_json::to_value(CommandExecutionRequestApprovalResponse { @@ -240,6 +271,26 @@ allow_local_binding = true ) .await?; } + JSONRPCMessage::Notification(notification) + if notification.method == "serverRequest/resolved" => + { + let resolved: ServerRequestResolvedNotification = serde_json::from_value( + notification + .params + .context("serverRequest/resolved params")?, + )?; + if withdrawn_request.as_ref() == Some(&resolved.request_id) { + assert_eq!(resolved.thread_id, thread_id); + app.send_response( + resolved.request_id, + serde_json::to_value(CommandExecutionRequestApprovalResponse { + decision: CommandExecutionApprovalDecision::Cancel, + })?, + ) + .await?; + sent_late_cancel = true; + } + } JSONRPCMessage::Notification(notification) if notification.method == "turn/completed" => { diff --git a/codex-rs/app-server/tests/suite/v2/turn_start.rs b/codex-rs/app-server/tests/suite/v2/turn_start.rs index 70531070d1..154f1529a0 100644 --- a/codex-rs/app-server/tests/suite/v2/turn_start.rs +++ b/codex-rs/app-server/tests/suite/v2/turn_start.rs @@ -2432,6 +2432,11 @@ fn managed_network_sandbox_probe() -> Result<()> { connection, "GET http://{target}/sandbox HTTP/1.0\r\nHost: {target}\r\nConnection: close\r\n\r\n" )?; + if let Some(signal) = std::env::var_os("CODEX_TEST_DISCONNECT_NETWORK_PROBE") { + // The parent releases its lock after receiving the approval, or when the test exits. + std::fs::File::open(signal)?.lock_shared()?; + return Ok(()); + } let mut response = String::new(); connection .take(/*limit*/ 4096) diff --git a/codex-rs/core/src/session/elicitation_holders_tests.rs b/codex-rs/core/src/session/elicitation_holders_tests.rs index 303f7ed6fe..2fe1b77a1e 100644 --- a/codex-rs/core/src/session/elicitation_holders_tests.rs +++ b/codex-rs/core/src/session/elicitation_holders_tests.rs @@ -100,7 +100,7 @@ async fn command_approval_holds_an_elicitation_until_response(capture: bool, pro vec!["echo".to_string()], cwd.into(), /*reason*/ None, - /*network_approval_context*/ None, + /*network_approval*/ None, /*proposed_execpolicy_amendment*/ None, /*additional_permissions*/ None, /*available_decisions*/ None, diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 2564c5663a..f4f6dcceb5 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -27,6 +27,7 @@ use crate::tasks::UserShellCommandTask; use crate::tasks::execute_user_shell_command; use crate::thread_manager::AgentTreeShutdownFailure; use crate::thread_manager::thread_store_error_kind; +use codex_protocol::approvals::ExecApprovalKind; use codex_protocol::error::CodexErr; use codex_protocol::protocol::CodexErrorInfo; use codex_protocol::protocol::ErrorEvent; @@ -175,68 +176,51 @@ pub async fn resolve_elicitation( } } -/// Propagate a user's exec approval decision to the session. -/// Also optionally applies an execpolicy amendment. -#[expect( - clippy::await_holding_invalid_type, - reason = "bind a human decision to its active turn before resolving it" -)] -pub async fn exec_approval( - sess: &Arc, - approval_id: String, - turn_id: Option, - decision: ReviewDecision, -) { - let human_approval = { - let active = sess.active_turn.lock().await; - let Some(turn) = active.as_ref() else { +impl Session { + /// Claim the prompt before applying a decision that can outlive the requesting execution. + pub(crate) async fn resolve_exec_approval( + self: &Arc, + approval_id: String, + turn_id: Option, + decision: ReviewDecision, + ) { + let Some((pending, active_turn_id)) = self + .take_pending_approval(&approval_id, turn_id.as_deref()) + .await + else { return; }; - let mut state = turn.turn_state.lock().await; - let Some(pending) = state.pending_approval(&approval_id) else { - return; - }; - if pending.kind == codex_protocol::approvals::ExecApprovalKind::HumanCommand { - state.remove_pending_approval(&approval_id) - } else { - None - } - }; - if let Some(pending) = human_approval { let decision = pending.validate(decision); if decision == ReviewDecision::Abort { - sess.abort_turn_if_active( - &pending.turn_id, - TurnAbortReason::Interrupted, - /*error*/ None, - ) + let turn_id = if pending.kind == ExecApprovalKind::HumanCommand { + Some(&pending.turn_id) + } else { + active_turn_id.as_ref() + }; + if let Some(turn_id) = turn_id { + self.abort_turn_if_active( + turn_id, + TurnAbortReason::Interrupted, + /*error*/ None, + ) + .await; + } + } else if let ReviewDecision::ApprovedExecpolicyAmendment { + proposed_execpolicy_amendment, + } = &decision + && let Err(err) = self + .persist_execpolicy_amendment(proposed_execpolicy_amendment) + .await + { + let message = format!("Failed to apply execpolicy amendment: {err}"); + tracing::warn!("{message}"); + self.send_event_raw(Event { + id: turn_id.unwrap_or(approval_id), + msg: EventMsg::Warning(WarningEvent { message }), + }) .await; } - pending.sender.send(decision).ok(); - return; - } - let event_turn_id = turn_id.unwrap_or_else(|| approval_id.clone()); - if let ReviewDecision::ApprovedExecpolicyAmendment { - proposed_execpolicy_amendment, - } = &decision - && let Err(err) = sess - .persist_execpolicy_amendment(proposed_execpolicy_amendment) - .await - { - let message = format!("Failed to apply execpolicy amendment: {err}"); - tracing::warn!("{message}"); - let warning = EventMsg::Warning(WarningEvent { message }); - sess.send_event_raw(Event { - id: event_turn_id.clone(), - msg: warning, - }) - .await; - } - match decision { - ReviewDecision::Abort => { - sess.interrupt_task().await; - } - other => sess.notify_approval(&approval_id, other).await, + let _ = pending.sender.send(decision); } } @@ -703,7 +687,8 @@ pub(super) async fn submission_loop( turn_id, decision, } => { - exec_approval(&sess, approval_id, turn_id, decision).await; + sess.resolve_exec_approval(approval_id, turn_id, decision) + .await; false } Op::PatchApproval { id, decision } => { diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 9f4a4e3917..2ae6f480c2 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -322,6 +322,7 @@ use crate::shell; use crate::state::AcceptedUserInputResponse; use crate::state::AutoCompactWindowIds; use crate::state::AutoCompactWindowSnapshot; +use crate::state::PendingApproval; use crate::state::PendingRequestPermissions; use crate::state::ReasoningEffortPin; use crate::state::SessionServices; @@ -421,6 +422,13 @@ pub(crate) enum TokenUsageContext { SpendingOnly, } +/// Which active turn can service a network approval. +#[derive(Clone, Copy, PartialEq, Eq)] +pub(crate) enum NetworkApprovalTurn { + Originating, + AnyActive, +} + /// Queue and lifecycle endpoints for a running [`Session`]. /// /// Runtime state lives on `Session`; keeping these endpoints separate lets all @@ -3089,12 +3097,13 @@ impl Session { command: Vec, cwd: PathUri, reason: Option, - network_approval_context: Option, + network_approval: Option<(NetworkApprovalContext, NetworkApprovalTurn)>, proposed_execpolicy_amendment: Option, additional_permissions: Option, available_decisions: Option>, plugin_attribution_override: Option, ) -> ReviewDecision { + let (network_approval_context, turn_binding) = network_approval.unzip(); let user_wait_owner = self.user_wait_timing_owner(turn_context, timing_call_id); let _elicitation = self.services.elicitations.register(); // command-level approvals use `call_id`. @@ -3102,19 +3111,40 @@ impl Session { let effective_approval_id = approval_id.clone().unwrap_or_else(|| call_id.clone()); // Add the tx_approve callback to the map before sending the request. let (tx_approve, rx_approve) = oneshot::channel(); - let prev_entry = { + let (prev_entry, _network_prompt) = { let mut active = self.active_turn.lock().await; + // The proxy reviews with the active turn's settings. A sandbox retry + // keeps its originating step's settings even after a background cell yields. + if let Some(binding) = turn_binding + && active + .as_ref() + .and_then(|turn| turn.task.as_ref()) + .is_none_or(|task| { + binding == NetworkApprovalTurn::Originating + && task.turn_context.sub_id != turn_context.sub_id + }) + { + return ReviewDecision::Abort; + } match active.as_mut() { Some(at) => { let mut ts = at.turn_state.lock().await; - ts.insert_pending_approval( + let prev = ts.insert_pending_approval( effective_approval_id.clone(), tx_approve, kind, turn_context.sub_id.clone(), - ) + ); + let prompt = network_approval_context.as_ref().map(|_| { + ts.track_network_prompt( + &self.tx_event, + &turn_context.sub_id, + &effective_approval_id, + ) + }); + (prev, prompt) } - None => None, + None => (None, None), } }; if prev_entry.is_some() { @@ -3622,25 +3652,32 @@ impl Session { clippy::await_holding_invalid_type, reason = "active turn checks and turn state updates must remain atomic" )] + async fn take_pending_approval( + &self, + approval_id: &str, + expected_turn_id: Option<&str>, + ) -> Option<(PendingApproval, Option)> { + let active = self.active_turn.lock().await; + let turn = active.as_ref()?; + // A live Cancel interrupts the turn servicing the prompt, which can differ from its origin. + let turn_id = turn.task.as_ref().map(|task| &task.turn_context.sub_id); + let pending = turn + .turn_state + .lock() + .await + .remove_pending_approval(approval_id, expected_turn_id)?; + (!pending.sender.is_closed()).then(|| (pending, turn_id.cloned())) + } + pub async fn notify_approval(&self, approval_id: &str, decision: ReviewDecision) { - let entry = { - let mut active = self.active_turn.lock().await; - match active.as_mut() { - Some(at) => { - let mut ts = at.turn_state.lock().await; - ts.remove_pending_approval(approval_id) - } - None => None, - } - }; - match entry { - Some(pending) => { - let decision = pending.validate(decision); - pending.sender.send(decision).ok(); - } - None => { - warn!("No pending approval found for call_id: {approval_id}"); - } + if let Some((pending, _)) = self + .take_pending_approval(approval_id, /*expected_turn_id*/ None) + .await + { + let decision = pending.validate(decision); + let _ = pending.sender.send(decision); + } else { + warn!("No pending approval found for call_id: {approval_id}"); } } diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index df0c4a3c57..7ac50acc55 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -33,7 +33,10 @@ use codex_execpolicy::Policy; use codex_execpolicy::RuleMatch; use codex_features::Feature; use codex_model_provider::create_model_provider; +use codex_network_proxy::EnvironmentNetworkPolicy; use codex_network_proxy::NetworkDecision; +use codex_network_proxy::NetworkDecisionSource; +use codex_network_proxy::NetworkPolicyDecision; use codex_network_proxy::NetworkPolicyRequest; use codex_network_proxy::NetworkProtocol; use codex_protocol::config_types::ApprovalsReviewer; @@ -43,6 +46,7 @@ use codex_protocol::models::ContentItemKind; use codex_protocol::models::ImageReference; use codex_protocol::models::NetworkPermissions; use codex_protocol::models::ResponseInputItem; +use codex_protocol::network_policy::NetworkPolicyDecisionPayload; use codex_protocol::protocol::AskForApproval; use codex_protocol::request_permissions::PermissionGrantScope; use codex_protocol::request_permissions::RequestPermissionProfile; @@ -984,7 +988,14 @@ async fn sandbox_denied_retry_uses_the_action_policy_and_reviewer() { exit_code: 1, ..Default::default() }), - network_policy_decision: None, + network_policy_decision: Some(NetworkPolicyDecisionPayload { + decision: NetworkPolicyDecision::Ask, + source: NetworkDecisionSource::Decider, + protocol: Some(NetworkApprovalProtocol::Https), + host: Some("example.com".to_string()), + port: Some(443), + reason: None, + }), }))); } Ok("sandbox-retry-succeeded".to_string()) @@ -1009,11 +1020,19 @@ async fn sandbox_denied_retry_uses_the_action_policy_and_reviewer() { call_id: call_id.to_string(), tool_name: codex_tools::ToolName::plain("exec_command"), }; - let environment = context + let mut environment = context .step_context .environments .primary() - .expect("primary environment"); + .expect("primary environment") + .clone(); + environment.config_mut().network_policy = Some(EnvironmentNetworkPolicy::from_config( + &NetworkProxyConfig { + enabled: true, + ..Default::default() + }, + /*managed_allowed_domains_only*/ false, + )); let mut orchestrator = ToolOrchestrator::new(); let mut runtime = DeniedOnceRuntime::default(); let approve = async { @@ -1021,16 +1040,30 @@ async fn sandbox_denied_retry_uses_the_action_policy_and_reviewer() { assert_eq!(approval.call_id, call_id); assert_eq!(approval.turn_id, action_turn.sub_id); assert_eq!( - approval.reason.as_deref(), - Some("command failed; retry without sandbox?") + approval.network_approval_context, + Some(NetworkApprovalContext { + host: "example.com".to_string(), + protocol: NetworkApprovalProtocol::Https, + }) ); session - .notify_approval(call_id, ReviewDecision::Approved) + .resolve_exec_approval( + approval.call_id.clone(), + Some(current_turn.sub_id.clone()), + ReviewDecision::Abort, + ) + .await; + session + .resolve_exec_approval( + approval.call_id, + Some(approval.turn_id), + ReviewDecision::Approved, + ) .await; }; let (output, ()) = tokio::join!( - orchestrator.run(&mut runtime, environment, &context), + orchestrator.run(&mut runtime, &environment, &context), approve ); assert_eq!( diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 65e7592a67..c575824c46 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -2289,6 +2289,7 @@ pub(super) fn realtime_text_for_event(msg: &EventMsg) -> Option, + pending_approvals: Arc, pending_request_permissions: HashMap, pending_user_input: HashMap>, pending_elicitations: HashMap<(String, RequestId), oneshot::Sender>, @@ -107,6 +112,7 @@ pub(crate) struct TurnState { pub(crate) struct PendingApproval { pub(crate) sender: oneshot::Sender, pub(crate) kind: codex_protocol::approvals::ExecApprovalKind, + // A yielded cell can emit a prompt for its original turn while a later turn services it. pub(crate) turn_id: String, } @@ -140,7 +146,49 @@ pub(crate) struct PendingRequestPermissions { pub(crate) turn_context: Arc, } +type PendingApprovals = SyncMutex>; + +/// The originating turn and its unique prompt ID stay valid even if the active turn changes. +pub(crate) struct NetworkApprovalPrompt<'a> { + approvals: Weak, + events: &'a async_channel::Sender, + turn_id: &'a str, + approval_id: &'a str, +} + +impl Drop for NetworkApprovalPrompt<'_> { + fn drop(&mut self) { + if let Some(approvals) = self.approvals.upgrade() { + approvals + .lock() + .unwrap_or_else(PoisonError::into_inner) + .remove(self.approval_id); + } + // The session channel is unbounded. Enqueue resolution before a same-host waiter can prompt. + let _ = self.events.try_send(Event { + id: self.turn_id.to_owned(), + msg: EventMsg::NetworkApprovalResolved { + approval_id: self.approval_id.to_owned(), + }, + }); + } +} + impl TurnState { + pub(crate) fn track_network_prompt<'a>( + &self, + events: &'a async_channel::Sender, + turn_id: &'a str, + approval_id: &'a str, + ) -> NetworkApprovalPrompt<'a> { + NetworkApprovalPrompt { + approvals: Arc::downgrade(&self.pending_approvals), + events, + turn_id, + approval_id, + } + } + pub(crate) fn insert_pending_approval( &mut self, key: String, @@ -148,26 +196,41 @@ impl TurnState { kind: codex_protocol::approvals::ExecApprovalKind, turn_id: String, ) -> Option { - self.pending_approvals.insert( - key, - PendingApproval { - sender: tx, - kind, - turn_id, - }, - ) + self.pending_approvals + .lock() + .unwrap_or_else(PoisonError::into_inner) + .insert( + key, + PendingApproval { + sender: tx, + kind, + turn_id, + }, + ) } - pub(crate) fn pending_approval(&self, key: &str) -> Option<&PendingApproval> { - self.pending_approvals.get(key) - } - - pub(crate) fn remove_pending_approval(&mut self, key: &str) -> Option { - self.pending_approvals.remove(key) + pub(crate) fn remove_pending_approval( + &mut self, + key: &str, + expected_turn_id: Option<&str>, + ) -> Option { + let mut approvals = self + .pending_approvals + .lock() + .unwrap_or_else(PoisonError::into_inner); + if let Some(expected) = expected_turn_id + && approvals.get(key)?.turn_id != expected + { + return None; + } + approvals.remove(key) } pub(crate) fn clear_pending_waiters(&mut self) { - self.pending_approvals.clear(); + self.pending_approvals + .lock() + .unwrap_or_else(PoisonError::into_inner) + .clear(); self.pending_request_permissions.clear(); self.pending_user_input.clear(); self.pending_elicitations.clear(); diff --git a/codex-rs/core/src/tools/approvals.rs b/codex-rs/core/src/tools/approvals.rs index e8461eea8a..9e1cd00013 100644 --- a/codex-rs/core/src/tools/approvals.rs +++ b/codex-rs/core/src/tools/approvals.rs @@ -12,6 +12,7 @@ use crate::guardian::spawn_approval_decision; use crate::hook_runtime::run_permission_request_hooks; use crate::mcp_tool_call::request_mcp_tool_user_approval; use crate::sandboxing::SandboxPermissions; +use crate::session::NetworkApprovalTurn; use crate::session::session::Session; use crate::tools::hook_names::HookToolName; use crate::tools::runtimes::apply_patch::ApplyPatchApprovalKey; @@ -793,7 +794,9 @@ impl Session { if human_once { None } else { - ctx.network_approval_context.clone() + ctx.network_approval_context + .clone() + .map(|context| (context, NetworkApprovalTurn::AnyActive)) }, if human_once { None @@ -857,7 +860,7 @@ impl Session { ], cwd.clone(), ctx.approval_reason.clone(), - /*network_approval_context*/ None, + /*network_approval*/ None, /*proposed_execpolicy_amendment*/ None, additional_permissions.clone(), Some(vec![ReviewDecision::Approved, ReviewDecision::Abort]), @@ -931,7 +934,9 @@ impl Session { command.clone(), cwd.clone().into(), ctx.approval_reason.clone(), - ctx.network_approval_context.clone(), + ctx.network_approval_context + .clone() + .map(|context| (context, NetworkApprovalTurn::Originating)), /*proposed_execpolicy_amendment*/ None, /*additional_permissions*/ None, /*available_decisions*/ None, diff --git a/codex-rs/core/src/tools/network_approval_tests.rs b/codex-rs/core/src/tools/network_approval_tests.rs index 139cff008e..350ac5ed38 100644 --- a/codex-rs/core/src/tools/network_approval_tests.rs +++ b/codex-rs/core/src/tools/network_approval_tests.rs @@ -17,10 +17,12 @@ use tokio_util::sync::CancellationToken; enum ExecutionCancellationTiming { PendingReview, AcceptedReview, + TurnFinished, } #[test_case(ExecutionCancellationTiming::PendingReview; "pending_review")] #[test_case(ExecutionCancellationTiming::AcceptedReview; "accepted_review")] +#[test_case(ExecutionCancellationTiming::TurnFinished; "turn_finishes_during_review")] #[tokio::test] #[allow( clippy::await_holding_invalid_type, @@ -32,7 +34,7 @@ async fn execution_cancellation_respects_network_approval_boundary( use crate::session::tests::make_session_and_context_with_rx; use crate::tasks::SessionTask; - struct PendingTask; + struct PendingTask(CancellationToken); impl SessionTask for PendingTask { fn kind(&self) -> crate::state::TaskKind { crate::state::TaskKind::Regular @@ -49,7 +51,10 @@ async fn execution_cancellation_respects_network_approval_boundary( _input: Vec, cancellation_token: CancellationToken, ) -> crate::tasks::SessionTaskResult { - cancellation_token.cancelled().await; + tokio::select! { + () = self.0.cancelled() => {} + () = cancellation_token.cancelled() => {} + } Ok(None) } } @@ -70,7 +75,10 @@ async fn execution_cancellation_respects_network_approval_boundary( ) .await .unwrap(); - session.start_task(turn, Vec::new(), PendingTask).await; + let finish_turn = CancellationToken::new(); + session + .start_task(turn, Vec::new(), PendingTask(finish_turn.clone())) + .await; let service = &session.services.network_approval; register_call_with_default_shell_trigger(service, "execution-1").await; let request = NetworkPolicyRequest { @@ -103,6 +111,17 @@ async fn execution_cancellation_respects_network_approval_boundary( .await .unwrap(); let expected_decision = match timing { + ExecutionCancellationTiming::TurnFinished => { + finish_turn.cancel(); + assert_eq!( + timeout(Duration::from_secs(5), decision) + .await + .expect("turn completion should release the waiting process"), + NetworkDecision::deny("not_allowed"), + ); + assert!(!*session.subscribe_elicitation_pause_state().borrow()); + return; + } ExecutionCancellationTiming::PendingReview => { service.record_call_outcome("execution-1", "original execution denial".to_string()); NetworkDecision::deny("not_allowed") @@ -148,12 +167,26 @@ async fn execution_cancellation_respects_network_approval_boundary( Some("original execution denial".to_string()) ); assert!(service.pending_host_approvals.lock().unwrap().is_empty()); + let mut resolved = Vec::new(); while let Ok(event) = events.try_recv() { - assert!(!matches!( - event.msg, - EventMsg::ExecApprovalRequest(_) | EventMsg::GuardianAssessment(_) - )); + match event.msg { + EventMsg::NetworkApprovalResolved { approval_id } => resolved.push(approval_id), + EventMsg::ExecApprovalRequest(_) | EventMsg::GuardianAssessment(_) => { + panic!("unexpected review after cancellation") + } + _ => {} + } } + assert_eq!(resolved, [approval.effective_approval_id()]); + // The app-server may already have received Cancel when it learns this prompt was withdrawn. + session + .resolve_exec_approval( + approval.effective_approval_id(), + Some(approval.turn_id), + ReviewDecision::Abort, + ) + .await; + assert!(session.active_turn.lock().await.is_some()); session .abort_all_tasks(codex_protocol::protocol::TurnAbortReason::Interrupted) .await; diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index e9cb691bc0..57499b1f43 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -67,6 +67,7 @@ use codex_protocol::openai_models::ToolMode; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; +use codex_protocol::protocol::ReviewDecision; use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::ThreadSettingsOverrides; use codex_protocol::user_input::UserInput; @@ -6087,20 +6088,29 @@ async fn code_mode_background_keeps_running_on_later_turn_without_wait() -> Resu skip_if_no_network!(Ok(())); let server = responses::start_mock_server().await; - let mut builder = test_codex().with_config(move |config| { - let _ = config.features.enable(Feature::CodeMode); - }); + let mut builder = test_codex() + .with_model("test-gpt-5.1-codex") + .with_config(|config| { + let _ = config.features.enable(Feature::CodeMode); + }); let test = builder.build(&server).await?; let resumed_file = test.workspace_path("code-mode-yield-resumed.txt"); let resumed_file_quoted = shlex::try_join([resumed_file.to_string_lossy().as_ref()])?; let write_file_command = format!("printf resumed > {resumed_file_quoted}"); let wait_for_file_command = format!("while [ ! -f {resumed_file_quoted} ]; do sleep 0.01; done; printf ready"); + let barrier = serde_json::json!({"barrier": { + "id": "background-command-approval", "participants": 2, "timeout_ms": 60_000, + }}); let code = format!( r#" text("before yield"); yield_control(); -await tools.exec_command({{ cmd: {write_file_command:?} }}); +await tools.test_sync_tool({barrier}); +await tools.exec_command({{ + cmd: {write_file_command:?}, sandbox_permissions: "require_escalated", + justification: "write the resumed file", +}}); text("after yield"); "# ); @@ -6123,7 +6133,12 @@ text("after yield"); ) .await; - test.submit_turn("start yielded exec").await?; + test.submit_turn_with_approval_and_permission_profile( + "start yielded exec", + AskForApproval::OnRequest, + PermissionProfile::read_only(), + ) + .await?; let first_request = first_completion.single_request(); let first_items = custom_tool_output_items(&first_request, "call-1"); @@ -6141,6 +6156,7 @@ text("after yield"); &server, sse(vec![ ev_response_created("resp-3"), + responses::ev_function_call("release-cell", "test_sync_tool", &barrier.to_string()), responses::ev_function_call( "call-2", "exec_command", @@ -6161,7 +6177,34 @@ text("after yield"); ) .await; - test.submit_turn("wait for resumed file").await?; + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "wait for resumed file".to_string(), + text_elements: Vec::new(), + }])) + .await?; + let servicing_turn = wait_for_event_match(&test.codex, |event| match event { + EventMsg::TurnStarted(turn) => Some(turn.turn_id.clone()), + _ => None, + }) + .await; + let approval = wait_for_event_match(&test.codex, |event| match event { + EventMsg::ExecApprovalRequest(request) => Some(request.clone()), + _ => None, + }) + .await; + assert_ne!(approval.turn_id, servicing_turn); + test.codex + .submit(Op::ExecApproval { + id: approval.effective_approval_id(), + turn_id: Some(approval.turn_id), + decision: ReviewDecision::Approved, + }) + .await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; let second_request = second_completion.single_request(); assert!( diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 62d01f556d..1c1282946b 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -1644,6 +1644,14 @@ pub enum EventMsg { ExecApprovalRequest(ExecApprovalRequestEvent), + /// Ends an in-process network prompt, including when its proxy request is dropped. + #[serde(skip)] + #[schemars(skip)] + #[ts(skip)] + NetworkApprovalResolved { + approval_id: String, + }, + RequestPermissions(RequestPermissionsEvent), RequestUserInput(RequestUserInputEvent), diff --git a/codex-rs/rollout-trace/src/protocol_event.rs b/codex-rs/rollout-trace/src/protocol_event.rs index ed49f8a943..a911b2b5f3 100644 --- a/codex-rs/rollout-trace/src/protocol_event.rs +++ b/codex-rs/rollout-trace/src/protocol_event.rs @@ -416,6 +416,7 @@ pub(crate) fn tool_runtime_trace_event(event: &EventMsg) -> Option Option<&'static s | EventMsg::DynamicToolCallResponse(_) | EventMsg::ElicitationRequest(_) | EventMsg::ElicitationAbandoned(_) + | EventMsg::NetworkApprovalResolved { .. } | EventMsg::ApplyPatchApprovalRequest(_) | EventMsg::GuardianAssessment(_) | EventMsg::DeprecationNotice(_) diff --git a/codex-rs/rollout/src/policy.rs b/codex-rs/rollout/src/policy.rs index b49049f6e9..e62c423204 100644 --- a/codex-rs/rollout/src/policy.rs +++ b/codex-rs/rollout/src/policy.rs @@ -199,6 +199,7 @@ fn persisted_event_msg( | EventMsg::RequestUserInput(_) | EventMsg::ElicitationRequest(_) | EventMsg::ElicitationAbandoned(_) + | EventMsg::NetworkApprovalResolved { .. } | EventMsg::ApplyPatchApprovalRequest(_) | EventMsg::StreamError(_) | EventMsg::PatchApplyBegin(_)