diff --git a/codex-rs/exec-server/src/posix.rs b/codex-rs/exec-server/src/posix.rs index ea0fab1a7e..b4dd0fbf40 100644 --- a/codex-rs/exec-server/src/posix.rs +++ b/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. diff --git a/codex-rs/exec-server/src/posix/escalate_server.rs b/codex-rs/exec-server/src/posix/escalate_server.rs index 0ea81f32bf..6512eab852 100644 --- a/codex-rs/exec-server/src/posix/escalate_server.rs +++ b/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, workdir: PathBuf, timeout_ms: Option, + cancel_rx: Option>, ) -> anyhow::Result { 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(); diff --git a/codex-rs/exec-server/src/posix/mcp.rs b/codex-rs/exec-server/src/posix/mcp.rs index f5785dc5d0..d10ed6e210 100644 --- a/codex-rs/exec-server/src/posix/mcp.rs +++ b/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 { std::env::var(CODEX_BASH_PATH_ENV_VAR) .map(PathBuf::from) @@ -87,10 +91,15 @@ impl ExecTool { context: RequestContext, Parameters(params): Parameters, ) -> Result { + 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(¶ms.workdir), - params.timeout_ms, + process_timeout_ms, + Some(cancel_rx), ) .await .map_err(|e| McpError::internal_error(e.to_string(), None))?; diff --git a/codex-rs/exec-server/src/posix/mcp_escalation_policy.rs b/codex-rs/exec-server/src/posix/mcp_escalation_policy.rs index 069948ea06..9e059fdba5 100644 --- a/codex-rs/exec-server/src/posix/mcp_escalation_policy.rs +++ b/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, + stopwatch: Stopwatch, } impl McpEscalationPolicy { - pub(crate) fn new(policy: ExecPolicy, context: RequestContext) -> Self { - Self { policy, context } + pub(crate) fn new( + policy: ExecPolicy, + context: RequestContext, + 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)) } } diff --git a/codex-rs/exec-server/src/posix/stopwatch.rs b/codex-rs/exec-server/src/posix/stopwatch.rs new file mode 100644 index 0000000000..72219b5996 --- /dev/null +++ b/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>, + notify: Arc, +} + +#[derive(Debug)] +struct StopwatchState { + elapsed: Duration, + running_since: Option, + 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(&self, fut: F) -> T + where + F: Future, + { + 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"); + } +}