Implemented thread-level atomic elicitation counter for stopwatch pausing (#12296)

### Purpose
While trying to build out CLI-Tools for the agent to use under skills we
have found that those tools sometimes need to invoke a user elicitation.
These elicitations are handled out of band of the codex app-server but
need to indicate to the exec manager that the command running is not
going to progress on the usual timeout horizon.

### Example
Model calls universal exec:
`$ download-credit-card-history --start-date 2026-01-19 --end-date
2026-02-19 > credit_history.jsonl`

download-cred-card-history might hit a hosted/preauthenticated service
to fetch data. That service might decide that the request requires an
end user approval the access to the personal data. It should be able to
signal to the running thread that the command in question is blocked on
user elicitation. In that case we want the exec to continue, but the
timeout to not expire on the tool call, essentially freezing time until
the user approves or rejects the command at which point the tool would
signal the app-server to decrement the outstanding elicitation count.
Now timeouts would proceed as normal.

### What's Added

- New v2 RPC methods:
    - thread/increment_elicitation
    - thread/decrement_elicitation
- Protocol updates in:
    - codex-rs/app-server-protocol/src/protocol/common.rs
    - codex-rs/app-server-protocol/src/protocol/v2.rs
- App-server handlers wired in:
    - codex-rs/app-server/src/codex_message_processor.rs

### Behavior

- Counter starts at 0 per thread.
- increment atomically increases the counter.
- decrement atomically decreases the counter; decrement at 0 returns
invalid request.
- Transition rules:
- 0 -> 1: broadcast pause state, pausing all active stopwatches
immediately.
    - \>0 -> >0: remain paused.
    - 1 -> 0: broadcast unpause state, resuming stopwatches.
- Core thread/session logic:
    - codex-rs/core/src/codex_thread.rs
    - codex-rs/core/src/codex.rs
    - codex-rs/core/src/mcp_connection_manager.rs

### Exec-server stopwatch integration

- Added centralized stopwatch tracking/controller:
    - codex-rs/exec-server/src/posix/stopwatch_controller.rs
- Hooked pause/unpause broadcast handling + stopwatch registration:
    - codex-rs/exec-server/src/posix/mcp.rs
    - codex-rs/exec-server/src/posix/stopwatch.rs
    - codex-rs/exec-server/src/posix.rs
This commit is contained in:
Channing Conger
2026-03-09 22:29:26 -07:00
committed by GitHub
Unverified
parent 79307b7933
commit c6343e0649
12 changed files with 773 additions and 15 deletions
+33
View File
@@ -422,6 +422,39 @@ mod tests {
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn unified_exec_pause_blocks_yield_timeout() -> anyhow::Result<()> {
skip_if_sandbox!(Ok(()));
let (session, turn) = test_session_and_turn().await;
session.set_out_of_band_elicitation_pause_state(true);
let paused_session = Arc::clone(&session);
tokio::spawn(async move {
tokio::time::sleep(Duration::from_secs(2)).await;
paused_session.set_out_of_band_elicitation_pause_state(false);
});
let started = tokio::time::Instant::now();
let response =
exec_command(&session, &turn, "sleep 1 && echo unified-exec-done", 250).await?;
assert!(
started.elapsed() >= Duration::from_secs(2),
"pause should block the unified exec yield timeout"
);
assert!(
response.output.contains("unified-exec-done"),
"exec_command should wait for output after the pause lifts"
);
assert!(
response.process_id.is_none(),
"completed command should not leave a background process"
);
Ok(())
}
#[tokio::test]
#[ignore] // Ignored while we have a better way to test this.
async fn requests_with_large_timeout_are_capped() -> anyhow::Result<()> {
@@ -8,6 +8,7 @@ use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use tokio::sync::Notify;
use tokio::sync::mpsc;
use tokio::sync::watch;
use tokio::time::Duration;
use tokio::time::Instant;
use tokio_util::sync::CancellationToken;
@@ -95,6 +96,7 @@ struct PreparedProcessHandles {
output_closed: Arc<AtomicBool>,
output_closed_notify: Arc<Notify>,
cancellation_token: CancellationToken,
pause_state: Option<watch::Receiver<bool>>,
command: Vec<String>,
process_id: String,
tty: bool,
@@ -211,6 +213,11 @@ impl UnifiedExecProcessManager {
&output_closed,
&output_closed_notify,
&cancellation_token,
Some(
context
.session
.subscribe_out_of_band_elicitation_pause_state(),
),
deadline,
)
.await;
@@ -303,6 +310,7 @@ impl UnifiedExecProcessManager {
output_closed,
output_closed_notify,
cancellation_token,
pause_state,
command: session_command,
process_id,
tty,
@@ -337,6 +345,7 @@ impl UnifiedExecProcessManager {
&output_closed,
&output_closed_notify,
&cancellation_token,
pause_state,
deadline,
)
.await;
@@ -435,6 +444,10 @@ impl UnifiedExecProcessManager {
output_closed_notify,
cancellation_token,
} = entry.process.output_handles();
let pause_state = entry
.session
.upgrade()
.map(|session| session.subscribe_out_of_band_elicitation_pause_state());
Ok(PreparedProcessHandles {
writer_tx: entry.process.writer_sender(),
@@ -443,6 +456,7 @@ impl UnifiedExecProcessManager {
output_closed,
output_closed_notify,
cancellation_token,
pause_state,
command: entry.command.clone(),
process_id: entry.process_id.clone(),
tty: entry.tty,
@@ -621,7 +635,8 @@ impl UnifiedExecProcessManager {
output_closed: &Arc<AtomicBool>,
output_closed_notify: &Arc<Notify>,
cancellation_token: &CancellationToken,
deadline: Instant,
mut pause_state: Option<watch::Receiver<bool>>,
mut deadline: Instant,
) -> Vec<u8> {
const POST_EXIT_CLOSE_WAIT_CAP: Duration = Duration::from_millis(50);
@@ -629,6 +644,12 @@ impl UnifiedExecProcessManager {
let mut exit_signal_received = cancellation_token.is_cancelled();
let mut post_exit_deadline: Option<Instant> = None;
loop {
Self::extend_deadlines_while_paused(
&mut pause_state,
&mut deadline,
&mut post_exit_deadline,
)
.await;
let drained_chunks: Vec<Vec<u8>>;
let mut wait_for_output = None;
{
@@ -666,6 +687,7 @@ impl UnifiedExecProcessManager {
_ = &mut notified => {}
_ = &mut closed => {}
_ = tokio::time::sleep(close_wait_remaining) => break,
_ = Self::wait_for_pause_change(pause_state.as_ref()) => {}
}
continue;
}
@@ -678,6 +700,7 @@ impl UnifiedExecProcessManager {
_ = &mut notified => {}
_ = &mut exit_notified => exit_signal_received = true,
_ = tokio::time::sleep(remaining) => break,
_ = Self::wait_for_pause_change(pause_state.as_ref()) => {}
}
continue;
}
@@ -695,6 +718,42 @@ impl UnifiedExecProcessManager {
collected
}
async fn extend_deadlines_while_paused(
pause_state: &mut Option<watch::Receiver<bool>>,
deadline: &mut Instant,
post_exit_deadline: &mut Option<Instant>,
) {
let Some(receiver) = pause_state.as_mut() else {
return;
};
if !*receiver.borrow() {
return;
}
let paused_at = Instant::now();
while *receiver.borrow() {
if receiver.changed().await.is_err() {
break;
}
}
let paused_for = paused_at.elapsed();
*deadline += paused_for;
if let Some(post_exit_deadline) = post_exit_deadline.as_mut() {
*post_exit_deadline += paused_for;
}
}
async fn wait_for_pause_change(pause_state: Option<&watch::Receiver<bool>>) {
match pause_state {
Some(pause_state) => {
let mut receiver = pause_state.clone();
let _ = receiver.changed().await;
}
None => std::future::pending::<()>().await,
}
}
fn prune_processes_if_needed(store: &mut ProcessStore) -> Option<ProcessEntry> {
if store.processes.len() < MAX_UNIFIED_EXEC_PROCESSES {
return None;