feat: waiting for an elicitation should not count against a shell tool timeout

codex-rs/exec-server/src/posix.rs

@@ -71,6 +71,7 @@ mod escalation_policy;
 mod mcp;
 mod mcp_escalation_policy;
 mod socket;
+mod stopwatch;
 
 /// Default value of --execve option relative to the current executable.
 /// Note this must match the name of the binary as specified in Cargo.toml.

codex-rs/exec-server/src/posix/escalate_server.rs

@@ -13,6 +13,7 @@ use codex_core::exec::process_exec_tool_call;
 use codex_core::get_platform_sandbox;
 use codex_core::protocol::SandboxPolicy;
 use tokio::process::Command;
+use tokio::sync::oneshot;
 
 use crate::posix::escalate_protocol::BASH_EXEC_WRAPPER_ENV_VAR;
 use crate::posix::escalate_protocol::ESCALATE_SOCKET_ENV_VAR;
@@ -49,6 +50,7 @@ impl EscalateServer {
         env: HashMap<String, String>,
         workdir: PathBuf,
         timeout_ms: Option<u64>,
+        cancel_rx: Option<oneshot::Receiver<()>>,
     ) -> anyhow::Result<ExecResult> {
         let (escalate_server, escalate_client) = AsyncDatagramSocket::pair()?;
         let client_socket = escalate_client.into_inner();
@@ -90,7 +92,7 @@ impl EscalateServer {
             &sandbox_cwd,
             &None,
             None,
-            None,
+            cancel_rx,
         )
         .await?;
         escalate_task.abort();

codex-rs/exec-server/src/posix/mcp.rs

@@ -22,10 +22,14 @@ use crate::posix::escalate_server::EscalateServer;
 use crate::posix::escalate_server::{self};
 use crate::posix::mcp_escalation_policy::ExecPolicy;
 use crate::posix::mcp_escalation_policy::McpEscalationPolicy;
+use crate::posix::stopwatch::Stopwatch;
 
 /// Path to our patched bash.
 const CODEX_BASH_PATH_ENV_VAR: &str = "CODEX_BASH_PATH";
 
+/// Default timeout for the shell tool.
+const DEFAULT_TIMEOUT_MS: u64 = 10_000;
+
 pub(crate) fn get_bash_path() -> Result<PathBuf> {
     std::env::var(CODEX_BASH_PATH_ENV_VAR)
         .map(PathBuf::from)
@@ -87,10 +91,15 @@ impl ExecTool {
         context: RequestContext<RoleServer>,
         Parameters(params): Parameters<ExecParams>,
     ) -> Result<CallToolResult, McpError> {
+        let effective_timeout =
+            Duration::from_millis(params.timeout_ms.unwrap_or(DEFAULT_TIMEOUT_MS));
+        let stopwatch = Stopwatch::new(effective_timeout);
+        let cancel_rx = stopwatch.cancellation_receiver();
+        let process_timeout_ms = Some(u64::MAX);
         let escalate_server = EscalateServer::new(
             self.bash_path.clone(),
             self.execve_wrapper.clone(),
-            McpEscalationPolicy::new(self.policy, context),
+            McpEscalationPolicy::new(self.policy, context, stopwatch.clone()),
         );
         let result = escalate_server
             .exec(
@@ -98,7 +107,8 @@ impl ExecTool {
                 // TODO: use ShellEnvironmentPolicy
                 std::env::vars().collect(),
                 PathBuf::from(&params.workdir),
-                params.timeout_ms,
+                process_timeout_ms,
+                Some(cancel_rx),
             )
             .await
             .map_err(|e| McpError::internal_error(e.to_string(), None))?;

codex-rs/exec-server/src/posix/mcp_escalation_policy.rs

@@ -10,6 +10,7 @@ use rmcp::service::RequestContext;
 
 use crate::posix::escalate_protocol::EscalateAction;
 use crate::posix::escalation_policy::EscalationPolicy;
+use crate::posix::stopwatch::Stopwatch;
 
 /// This is the policy which decides how to handle an exec() call.
 ///
@@ -34,11 +35,20 @@ pub(crate) enum ExecPolicyOutcome {
 pub(crate) struct McpEscalationPolicy {
     policy: ExecPolicy,
     context: RequestContext<RoleServer>,
+    stopwatch: Stopwatch,
 }
 
 impl McpEscalationPolicy {
-    pub(crate) fn new(policy: ExecPolicy, context: RequestContext<RoleServer>) -> Self {
-        Self { policy, context }
+    pub(crate) fn new(
+        policy: ExecPolicy,
+        context: RequestContext<RoleServer>,
+        stopwatch: Stopwatch,
+    ) -> Self {
+        Self {
+            policy,
+            context,
+            stopwatch,
+        }
     }
 
     async fn prompt(
@@ -54,25 +64,34 @@ impl McpEscalationPolicy {
         } else {
             format!("{} {}", file.display(), args)
         };
-        context
-            .peer
-            .create_elicitation(CreateElicitationRequestParam {
-                message: format!("Allow agent to run `{command}` in `{}`?", workdir.display()),
-                requested_schema: ElicitationSchema::builder()
-                    .title("Execution Permission Request")
-                    .optional_string_with("reason", |schema| {
-                        schema.description("Optional reason for allowing or denying execution")
+        self.stopwatch
+            .pause_for(async {
+                context
+                    .peer
+                    .create_elicitation(CreateElicitationRequestParam {
+                        message: format!(
+                            "Allow agent to run `{command}` in `{}`?",
+                            workdir.display()
+                        ),
+                        requested_schema: ElicitationSchema::builder()
+                            .title("Execution Permission Request")
+                            .optional_string_with("reason", |schema| {
+                                schema.description(
+                                    "Optional reason for allowing or denying execution",
+                                )
+                            })
+                            .build()
+                            .map_err(|e| {
+                                McpError::internal_error(
+                                    format!("failed to build elicitation schema: {e}"),
+                                    None,
+                                )
+                            })?,
                     })
-                    .build()
-                    .map_err(|e| {
-                        McpError::internal_error(
-                            format!("failed to build elicitation schema: {e}"),
-                            None,
-                        )
-                    })?,
+                    .await
+                    .map_err(|e| McpError::internal_error(e.to_string(), None))
             })
             .await
-            .map_err(|e| McpError::internal_error(e.to_string(), None))
     }
 }
 

codex-rs/exec-server/src/posix/stopwatch.rs

@@ -0,0 +1,159 @@
+use std::future::Future;
+use std::sync::Arc;
+use std::time::Duration;
+use std::time::Instant;
+
+use tokio::sync::Mutex;
+use tokio::sync::Notify;
+use tokio::sync::oneshot;
+
+#[derive(Clone, Debug)]
+pub(crate) struct Stopwatch {
+    limit: Duration,
+    inner: Arc<Mutex<StopwatchState>>,
+    notify: Arc<Notify>,
+}
+
+#[derive(Debug)]
+struct StopwatchState {
+    elapsed: Duration,
+    running_since: Option<Instant>,
+    active_pauses: u32,
+}
+
+impl Stopwatch {
+    pub(crate) fn new(limit: Duration) -> Self {
+        Self {
+            inner: Arc::new(Mutex::new(StopwatchState {
+                elapsed: Duration::ZERO,
+                running_since: Some(Instant::now()),
+                active_pauses: 0,
+            })),
+            notify: Arc::new(Notify::new()),
+            limit,
+        }
+    }
+
+    pub(crate) fn cancellation_receiver(&self) -> oneshot::Receiver<()> {
+        let limit = self.limit;
+        let (tx, rx) = oneshot::channel();
+        let inner = Arc::clone(&self.inner);
+        let notify = Arc::clone(&self.notify);
+        tokio::spawn(async move {
+            loop {
+                let (remaining, running) = {
+                    let guard = inner.lock().await;
+                    let elapsed = guard.elapsed
+                        + guard
+                            .running_since
+                            .map(|since| since.elapsed())
+                            .unwrap_or_default();
+                    if elapsed >= limit {
+                        break;
+                    }
+                    (limit - elapsed, guard.running_since.is_some())
+                };
+
+                if remaining.is_zero() {
+                    break;
+                }
+
+                if !running {
+                    notify.notified().await;
+                    continue;
+                }
+
+                let sleep = tokio::time::sleep(remaining);
+                tokio::pin!(sleep);
+                tokio::select! {
+                    _ = &mut sleep => {
+                        break;
+                    }
+                    _ = notify.notified() => {
+                        continue;
+                    }
+                }
+            }
+            let _ = tx.send(());
+        });
+        rx
+    }
+
+    /// Runs `fut`, pausing the stopwatch while the future is pending. The clock resumes
+    /// automatically when the future completes. Nested calls are reference-counted so the
+    /// stopwatch only resumes when every pause is lifted.
+    pub(crate) async fn pause_for<F, T>(&self, fut: F) -> T
+    where
+        F: Future<Output = T>,
+    {
+        self.pause().await;
+        let result = fut.await;
+        self.resume().await;
+        result
+    }
+
+    async fn pause(&self) {
+        let mut guard = self.inner.lock().await;
+        guard.active_pauses += 1;
+        if guard.active_pauses == 1
+            && let Some(since) = guard.running_since.take()
+        {
+            guard.elapsed += since.elapsed();
+            self.notify.notify_waiters();
+        }
+    }
+
+    async fn resume(&self) {
+        let mut guard = self.inner.lock().await;
+        if guard.active_pauses == 0 {
+            return;
+        }
+        guard.active_pauses -= 1;
+        if guard.active_pauses == 0 && guard.running_since.is_none() {
+            guard.running_since = Some(Instant::now());
+            self.notify.notify_waiters();
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::Stopwatch;
+    use tokio::time::Duration;
+    use tokio::time::Instant;
+    use tokio::time::sleep;
+    use tokio::time::timeout;
+
+    #[tokio::test]
+    async fn cancellation_receiver_fires_after_limit() {
+        let stopwatch = Stopwatch::new(Duration::from_millis(50));
+        let rx = stopwatch.cancellation_receiver();
+        let start = Instant::now();
+        rx.await.expect("cancellation should fire");
+        assert!(start.elapsed() >= Duration::from_millis(50));
+    }
+
+    #[tokio::test]
+    async fn pause_prevents_timeout_until_resumed() {
+        let stopwatch = Stopwatch::new(Duration::from_millis(50));
+        let mut rx = stopwatch.cancellation_receiver();
+
+        let pause_handle = tokio::spawn({
+            let stopwatch = stopwatch.clone();
+            async move {
+                stopwatch
+                    .pause_for(async {
+                        sleep(Duration::from_millis(100)).await;
+                    })
+                    .await;
+            }
+        });
+
+        assert!(timeout(Duration::from_millis(30), &mut rx).await.is_err());
+
+        pause_handle.await.expect("pause task should finish");
+
+        rx.await
+            .expect("cancellation should eventually fire after resume");
+    }
+}