code-mode: make all approvals trigger elicitation pause (#31650)

### summary

We want to pause code-mode from yielding back to the model when a
subcommand triggers an approval prompt. This means that all of these
previously inline blocking requests should also take out a
ElicitationService registration.

This also does some plumbing refactoring to request patch approval to
make it match the other `request_*_approval` methods in that it blocks
on the approval in the function instead of returning the oneshot
channel, this affords our ability to encapsulate the ElicitationService
registration via RAII.

Adds tests to confirm the blocking behavior for code_mode both in suite
tests and that the session holds them.
This commit is contained in:
Channing Conger
2026-07-08 15:27:04 -07:00
committed by GitHub
parent b6f9aee16d
commit c55cb4b363
6 changed files with 461 additions and 11 deletions
+3 -4
View File
@@ -656,11 +656,10 @@ async fn handle_patch_approval(
let decision = if let Some(decision) = guardian_decision {
decision
} else {
let decision_rx = parent_session
.request_patch_approval(parent_ctx, call_id, changes, reason, grant_root)
.await;
let decision =
parent_session.request_patch_approval(parent_ctx, call_id, changes, reason, grant_root);
await_approval_with_cancel(
async move { decision_rx.await.unwrap_or_default() },
decision,
parent_session,
&approval_id,
cancel_token,
@@ -0,0 +1,182 @@
use std::collections::HashMap;
use codex_protocol::protocol::ReviewDecision;
use codex_protocol::request_permissions::PermissionGrantScope;
use codex_protocol::request_permissions::RequestPermissionProfile;
use codex_protocol::request_permissions::RequestPermissionsArgs;
use codex_protocol::request_permissions::RequestPermissionsResponse;
use codex_protocol::request_user_input::RequestUserInputArgs;
use codex_protocol::request_user_input::RequestUserInputResponse;
use tokio::sync::watch;
use tokio_util::sync::CancellationToken;
use super::tests::make_session_and_context_with_rx;
use crate::state::ActiveTurn;
async fn wait_until_held(pause_state: &mut watch::Receiver<bool>) {
pause_state
.wait_for(|paused| *paused)
.await
.expect("elicitation service should remain available");
}
async fn wait_until_released(pause_state: &mut watch::Receiver<bool>) {
pause_state
.wait_for(|paused| !*paused)
.await
.expect("elicitation service should remain available");
}
#[tokio::test]
async fn command_approval_holds_an_elicitation_until_response() {
let (session, turn_context, events) = make_session_and_context_with_rx().await;
*session.active_turn.lock().await = Some(ActiveTurn::default());
let mut pause_state = session.subscribe_elicitation_pause_state();
#[allow(deprecated)]
let cwd = turn_context.cwd.clone();
let request = tokio::spawn({
let session = session.clone();
let turn_context = turn_context.clone();
async move {
session
.request_command_approval(
turn_context.as_ref(),
"call-1".to_string(),
/*approval_id*/ None,
/*environment_id*/ None,
vec!["echo".to_string()],
cwd,
/*reason*/ None,
/*network_approval_context*/ None,
/*proposed_execpolicy_amendment*/ None,
/*additional_permissions*/ None,
/*available_decisions*/ None,
)
.await
}
});
events.recv().await.expect("approval event");
wait_until_held(&mut pause_state).await;
session
.notify_approval("call-1", ReviewDecision::Approved)
.await;
request.await.expect("approval task");
wait_until_released(&mut pause_state).await;
}
#[tokio::test]
async fn patch_approval_holds_an_elicitation_until_response() {
let (session, turn_context, events) = make_session_and_context_with_rx().await;
*session.active_turn.lock().await = Some(ActiveTurn::default());
let mut pause_state = session.subscribe_elicitation_pause_state();
let request = tokio::spawn({
let session = session.clone();
let turn_context = turn_context.clone();
async move {
session
.request_patch_approval(
turn_context.as_ref(),
"call-1".to_string(),
HashMap::new(),
/*reason*/ None,
/*grant_root*/ None,
)
.await
}
});
events.recv().await.expect("approval event");
wait_until_held(&mut pause_state).await;
session
.notify_approval("call-1", ReviewDecision::Approved)
.await;
request.await.expect("approval task");
wait_until_released(&mut pause_state).await;
}
#[tokio::test]
async fn permission_request_holds_an_elicitation_until_response() {
let (session, turn_context, events) = make_session_and_context_with_rx().await;
*session.active_turn.lock().await = Some(ActiveTurn::default());
let mut pause_state = session.subscribe_elicitation_pause_state();
let request = tokio::spawn({
let session = session.clone();
let turn_context = turn_context.clone();
async move {
let environment = turn_context
.environments
.primary()
.expect("primary environment")
.selection();
session
.request_permissions_for_environment(
&turn_context,
"call-1".to_string(),
RequestPermissionsArgs {
environment_id: None,
reason: None,
permissions: RequestPermissionProfile::default(),
},
environment,
CancellationToken::new(),
)
.await
}
});
events.recv().await.expect("permission request event");
wait_until_held(&mut pause_state).await;
session
.notify_request_permissions_response(
"call-1",
RequestPermissionsResponse {
permissions: RequestPermissionProfile::default(),
scope: PermissionGrantScope::Turn,
strict_auto_review: false,
},
)
.await;
request.await.expect("permission request task");
wait_until_released(&mut pause_state).await;
}
#[tokio::test]
async fn request_user_input_holds_an_elicitation_until_response() {
let (session, turn_context, events) = make_session_and_context_with_rx().await;
*session.active_turn.lock().await = Some(ActiveTurn::default());
let mut pause_state = session.subscribe_elicitation_pause_state();
let request = tokio::spawn({
let session = session.clone();
let turn_context = turn_context.clone();
async move {
session
.request_user_input(
turn_context.as_ref(),
"call-1".to_string(),
RequestUserInputArgs {
questions: Vec::new(),
auto_resolution_ms: None,
},
)
.await
}
});
events.recv().await.expect("request user input event");
wait_until_held(&mut pause_state).await;
let response = RequestUserInputResponse {
answers: HashMap::new(),
};
session
.notify_user_input_response(&turn_context.sub_id, response)
.await;
request.await.expect("request user input task");
wait_until_released(&mut pause_state).await;
}
+10 -2
View File
@@ -2158,6 +2158,7 @@ impl Session {
additional_permissions: Option<AdditionalPermissionProfile>,
available_decisions: Option<Vec<ReviewDecision>>,
) -> ReviewDecision {
let _elicitation = self.services.elicitations.register();
// command-level approvals use `call_id`.
// `approval_id` is only present for subcommand callbacks (execve intercept)
let effective_approval_id = approval_id.clone().unwrap_or_else(|| call_id.clone());
@@ -2229,7 +2230,8 @@ impl Session {
changes: HashMap<PathBuf, FileChange>,
reason: Option<String>,
grant_root: Option<PathBuf>,
) -> oneshot::Receiver<ReviewDecision> {
) -> ReviewDecision {
let _elicitation = self.services.elicitations.register();
// Add the tx_approve callback to the map before sending the request.
let (tx_approve, rx_approve) = oneshot::channel();
let approval_id = call_id.clone();
@@ -2256,7 +2258,7 @@ impl Session {
grant_root,
});
self.send_event(turn_context, event).await;
rx_approve
rx_approve.await.unwrap_or(ReviewDecision::Abort)
}
#[expect(
@@ -2384,6 +2386,7 @@ impl Session {
return Some(response);
}
let _elicitation = self.services.elicitations.register();
let (tx_response, rx_response) = oneshot::channel();
let prev_entry = {
let mut active = self.active_turn.lock().await;
@@ -2475,6 +2478,7 @@ impl Session {
call_id: String,
args: RequestUserInputArgs,
) -> Option<RequestUserInputResponse> {
let _elicitation = self.services.elicitations.register();
let sub_id = turn_context.sub_id.clone();
let (tx_response, rx_response) = oneshot::channel();
let event_id = sub_id.clone();
@@ -4046,5 +4050,9 @@ async fn build_hooks_for_config(
})
}
#[cfg(test)]
#[path = "elicitation_holders_tests.rs"]
mod elicitation_holders_tests;
#[cfg(test)]
pub(crate) mod tests;
@@ -156,7 +156,7 @@ impl Approvable<ApplyPatchRequest> for ApplyPatchRuntime {
return ReviewDecision::Approved;
}
if let Some(reason) = retry_reason {
let rx_approve = session
return session
.request_patch_approval(
turn,
call_id,
@@ -165,7 +165,6 @@ impl Approvable<ApplyPatchRequest> for ApplyPatchRuntime {
/*grant_root*/ None,
)
.await;
return rx_approve.await.unwrap_or_default();
}
with_cached_approval(
@@ -173,12 +172,11 @@ impl Approvable<ApplyPatchRequest> for ApplyPatchRuntime {
"apply_patch",
approval_keys,
|| async move {
let rx_approve = session
session
.request_patch_approval(
turn, call_id, changes, /*reason*/ None, /*grant_root*/ None,
)
.await;
rx_approve.await.unwrap_or_default()
.await
},
)
.await
@@ -0,0 +1,262 @@
#![allow(clippy::unwrap_used)]
use std::time::Duration;
use anyhow::Result;
use codex_core::config::Config;
use codex_features::Feature;
use codex_protocol::models::PermissionProfile;
use codex_protocol::protocol::AskForApproval;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::Op;
use codex_protocol::protocol::ReviewDecision;
use codex_protocol::request_permissions::PermissionGrantScope;
use codex_protocol::request_permissions::RequestPermissionsResponse;
use codex_protocol::user_input::UserInput;
use core_test_support::responses;
use core_test_support::responses::ResponseMock;
use core_test_support::responses::ev_assistant_message;
use core_test_support::responses::ev_completed;
use core_test_support::responses::ev_custom_tool_call;
use core_test_support::responses::ev_response_created;
use core_test_support::responses::sse;
use core_test_support::skip_if_no_network;
use core_test_support::skip_if_wine_exec;
use core_test_support::test_codex::TestCodex;
use core_test_support::test_codex::test_codex;
use core_test_support::test_codex::turn_permission_fields;
use core_test_support::wait_for_event;
use core_test_support::wait_for_event_match;
use core_test_support::wait_for_event_with_timeout;
use wiremock::MockServer;
const YIELD_TIME_MS: u64 = 1_000;
const TURN_COMPLETE_TIMEOUT: Duration = Duration::from_secs(30);
struct CodeModeElicitationHarness {
_server: MockServer,
test: TestCodex,
follow_up: ResponseMock,
turn_id: String,
}
impl CodeModeElicitationHarness {
async fn start(
code: &str,
permission_profile: PermissionProfile,
configure: impl FnOnce(&mut Config) + Send + 'static,
) -> Result<Self> {
let server = responses::start_mock_server().await;
let mut builder =
test_codex()
.with_model("test-gpt-5.1-codex")
.with_config(move |config| {
let _ = config.features.enable(Feature::CodeMode);
configure(config);
});
let test = builder.build_with_auto_env(&server).await?;
let follow_up = mount_code_mode_responses(&server, code).await;
let turn_id = submit_turn(&test, permission_profile).await?;
Ok(Self {
_server: server,
test,
follow_up,
turn_id,
})
}
async fn assert_result_held(&self) {
tokio::time::sleep(Duration::from_millis(YIELD_TIME_MS + 250)).await;
assert!(
self.follow_up.requests().is_empty(),
"captured exec result should not return during a user elicitation"
);
}
async fn finish(self) {
wait_for_event_with_timeout(
&self.test.codex,
|event| match event {
EventMsg::TurnComplete(event) => event.turn_id == self.turn_id,
_ => false,
},
TURN_COMPLETE_TIMEOUT,
)
.await;
self.follow_up.single_request();
}
}
async fn mount_code_mode_responses(server: &MockServer, code: &str) -> ResponseMock {
responses::mount_sse_once(
server,
sse(vec![
ev_response_created("resp-1"),
ev_custom_tool_call("call-1", "exec", code),
ev_completed("resp-1"),
]),
)
.await;
responses::mount_sse_once(
server,
sse(vec![
ev_assistant_message("msg-1", "done"),
ev_completed("resp-2"),
]),
)
.await
}
async fn submit_turn(test: &TestCodex, permission_profile: PermissionProfile) -> Result<String> {
let (sandbox_policy, permission_profile) =
turn_permission_fields(permission_profile, test.config.cwd.as_path());
test.codex
.submit(Op::UserInput {
items: vec![UserInput::Text {
text: "run a code-mode tool that needs user input".into(),
text_elements: Vec::new(),
}],
final_output_json_schema: None,
responsesapi_client_metadata: None,
additional_context: Default::default(),
thread_settings: codex_protocol::protocol::ThreadSettingsOverrides {
approval_policy: Some(AskForApproval::OnRequest),
sandbox_policy: Some(sandbox_policy),
permission_profile,
collaboration_mode: Some(codex_protocol::config_types::CollaborationMode {
mode: codex_protocol::config_types::ModeKind::Default,
settings: codex_protocol::config_types::Settings {
model: test.session_configured.model.clone(),
reasoning_effort: None,
developer_instructions: None,
},
}),
..Default::default()
},
})
.await?;
Ok(wait_for_event_match(&test.codex, |event| match event {
EventMsg::TurnStarted(event) => Some(event.turn_id.clone()),
_ => None,
})
.await)
}
#[cfg_attr(windows, ignore = "no exec_command on Windows")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn code_mode_holds_yielded_result_during_command_approval() -> Result<()> {
skip_if_wine_exec!(
Ok(()),
"command approvals currently require a host-native cwd"
);
skip_if_no_network!(Ok(()));
let harness = CodeModeElicitationHarness::start(
r#"// @exec: {"yield_time_ms": 1000}
await tools.exec_command({
cmd: "printf code_mode_approval_marker",
sandbox_permissions: "require_escalated",
justification: "test command approval",
});"#,
PermissionProfile::read_only(),
|_| {},
)
.await?;
let approval = wait_for_event_match(&harness.test.codex, |event| match event {
EventMsg::ExecApprovalRequest(approval) => Some(approval.clone()),
_ => None,
})
.await;
harness.assert_result_held().await;
harness
.test
.codex
.submit(Op::ExecApproval {
id: approval.effective_approval_id(),
turn_id: Some(harness.turn_id.clone()),
decision: ReviewDecision::Approved,
})
.await?;
harness.finish().await;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn code_mode_holds_yielded_result_during_patch_approval() -> Result<()> {
skip_if_no_network!(Ok(()));
let harness = CodeModeElicitationHarness::start(
r#"// @exec: {"yield_time_ms": 1000}
await tools.apply_patch("*** Begin Patch\n*** Add File: code_mode_patch_approval.txt\n+held\n*** End Patch\n");"#,
PermissionProfile::read_only(),
|_| {},
)
.await?;
let approval = wait_for_event_match(&harness.test.codex, |event| match event {
EventMsg::ApplyPatchApprovalRequest(approval) => Some(approval.clone()),
_ => None,
})
.await;
harness.assert_result_held().await;
harness
.test
.codex
.submit(Op::PatchApproval {
id: approval.call_id,
decision: ReviewDecision::Approved,
})
.await?;
harness.finish().await;
Ok(())
}
#[cfg_attr(
target_os = "linux",
ignore = "request_permissions tool integration is not supported on Linux"
)]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn code_mode_holds_yielded_result_during_permission_request() -> Result<()> {
skip_if_no_network!(Ok(()));
let harness = CodeModeElicitationHarness::start(
r#"// @exec: {"yield_time_ms": 1000}
await tools.request_permissions({
reason: "test permission request",
permissions: { network: { enabled: true } },
});"#,
PermissionProfile::read_only(),
|config| {
let _ = config.features.enable(Feature::RequestPermissionsTool);
},
)
.await?;
let request = wait_for_event(&harness.test.codex, |event| {
matches!(
event,
EventMsg::RequestPermissions(_) | EventMsg::TurnComplete(_) | EventMsg::Error(_)
)
})
.await;
let EventMsg::RequestPermissions(request) = request else {
panic!("expected request_permissions before turn completion, got {request:?}");
};
harness.assert_result_held().await;
harness
.test
.codex
.submit(Op::RequestPermissionsResponse {
id: request.call_id,
response: RequestPermissionsResponse {
permissions: Default::default(),
scope: PermissionGrantScope::Turn,
strict_auto_review: false,
},
})
.await?;
harness.finish().await;
Ok(())
}
+1
View File
@@ -42,6 +42,7 @@ mod cli_stream;
mod client;
mod client_websockets;
mod code_mode;
mod code_mode_elicitation;
mod codex_delegate;
mod collaboration_instructions;
mod compact;