From f1b7cdc3bd3a03930456214b7fa7c49b3ba89260 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Thu, 4 Dec 2025 08:34:09 -0800 Subject: [PATCH] Use shared check sandboxing (#7547) --- codex-rs/core/src/unified_exec/session.rs | 32 ++++--- .../core/src/unified_exec/session_manager.rs | 92 ++++++------------- 2 files changed, 47 insertions(+), 77 deletions(-) diff --git a/codex-rs/core/src/unified_exec/session.rs b/codex-rs/core/src/unified_exec/session.rs index a6e4167ad..02465538e 100644 --- a/codex-rs/core/src/unified_exec/session.rs +++ b/codex-rs/core/src/unified_exec/session.rs @@ -154,10 +154,6 @@ impl UnifiedExecSession { } pub(super) async fn check_for_sandbox_denial(&self) -> Result<(), UnifiedExecError> { - if self.sandbox_type() == SandboxType::None || !self.has_exited() { - return Ok(()); - } - let _ = tokio::time::timeout(Duration::from_millis(20), self.output_notify.notified()).await; @@ -167,28 +163,40 @@ impl UnifiedExecSession { aggregated.extend_from_slice(&chunk); } let aggregated_text = String::from_utf8_lossy(&aggregated).to_string(); - let exit_code = self.exit_code().unwrap_or(-1); + self.check_for_sandbox_denial_with_text(&aggregated_text) + .await?; + Ok(()) + } + + pub(super) async fn check_for_sandbox_denial_with_text( + &self, + text: &str, + ) -> Result<(), UnifiedExecError> { + let sandbox_type = self.sandbox_type(); + if sandbox_type == SandboxType::None || !self.has_exited() { + return Ok(()); + } + + let exit_code = self.exit_code().unwrap_or(-1); let exec_output = ExecToolCallOutput { exit_code, - stdout: StreamOutput::new(aggregated_text.clone()), - aggregated_output: StreamOutput::new(aggregated_text.clone()), + stderr: StreamOutput::new(text.to_string()), + aggregated_output: StreamOutput::new(text.to_string()), ..Default::default() }; - - if is_likely_sandbox_denied(self.sandbox_type(), &exec_output) { + if is_likely_sandbox_denied(sandbox_type, &exec_output) { let snippet = formatted_truncate_text( - &aggregated_text, + text, TruncationPolicy::Tokens(UNIFIED_EXEC_OUTPUT_MAX_TOKENS), ); let message = if snippet.is_empty() { - format!("Session creation failed with exit code {exit_code}") + format!("Session exited with code {exit_code}") } else { snippet }; return Err(UnifiedExecError::sandbox_denied(message, exec_output)); } - Ok(()) } diff --git a/codex-rs/core/src/unified_exec/session_manager.rs b/codex-rs/core/src/unified_exec/session_manager.rs index 0eab30dce..da9cb338d 100644 --- a/codex-rs/core/src/unified_exec/session_manager.rs +++ b/codex-rs/core/src/unified_exec/session_manager.rs @@ -154,10 +154,24 @@ impl UnifiedExecSessionManager { let output = formatted_truncate_text(&text, TruncationPolicy::Tokens(max_tokens)); let has_exited = session.has_exited(); let exit_code = session.exit_code(); - let sandbox_type = session.sandbox_type(); let chunk_id = generate_chunk_id(); - let process_id = if has_exited { - None + let process_id = request.process_id.clone(); + if has_exited { + let exit = exit_code.unwrap_or(-1); + Self::emit_exec_end_from_context( + context, + &request.command, + cwd, + output.clone(), + exit, + wall_time, + // We always emit the process ID in order to keep consistency between the Begin + // event and the End event. + Some(process_id), + ) + .await; + + session.check_for_sandbox_denial_with_text(&text).await?; } else { // Only store session if not exited. self.store_session( @@ -166,48 +180,29 @@ impl UnifiedExecSessionManager { &request.command, cwd.clone(), start, - request.process_id.clone(), + process_id, ) .await; - Some(request.process_id.clone()) - }; - let original_token_count = approx_token_count(&text); + Self::emit_waiting_status(&context.session, &context.turn, &request.command).await; + }; + + let original_token_count = approx_token_count(&text); let response = UnifiedExecResponse { event_call_id: context.call_id.clone(), chunk_id, wall_time, output, - process_id: process_id.clone(), + process_id: if has_exited { + None + } else { + Some(request.process_id.clone()) + }, exit_code, original_token_count: Some(original_token_count), session_command: Some(request.command.clone()), }; - if !has_exited { - Self::emit_waiting_status(&context.session, &context.turn, &request.command).await; - } - - // If the command completed during this call, emit an ExecCommandEnd via the emitter. - if has_exited { - let exit = response.exit_code.unwrap_or(-1); - Self::emit_exec_end_from_context( - context, - &request.command, - cwd, - response.output.clone(), - exit, - response.wall_time, - // We always emit the process ID in order to keep consistency between the Begin - // event and the End event. - Some(request.process_id), - ) - .await; - - // Exit code should always be Some - sandboxing::check_sandboxing(sandbox_type, &text, exit_code.unwrap_or_default())?; - } - Ok(response) } @@ -714,39 +709,6 @@ impl UnifiedExecSessionManager { } } -mod sandboxing { - use super::*; - use crate::exec::SandboxType; - use crate::exec::is_likely_sandbox_denied; - use crate::unified_exec::UNIFIED_EXEC_OUTPUT_MAX_TOKENS; - - pub(crate) fn check_sandboxing( - sandbox_type: SandboxType, - text: &str, - exit_code: i32, - ) -> Result<(), UnifiedExecError> { - let exec_output = ExecToolCallOutput { - exit_code, - stderr: StreamOutput::new(text.to_string()), - aggregated_output: StreamOutput::new(text.to_string()), - ..Default::default() - }; - if is_likely_sandbox_denied(sandbox_type, &exec_output) { - let snippet = formatted_truncate_text( - text, - TruncationPolicy::Tokens(UNIFIED_EXEC_OUTPUT_MAX_TOKENS), - ); - let message = if snippet.is_empty() { - format!("Session exited with code {exit_code}") - } else { - snippet - }; - return Err(UnifiedExecError::sandbox_denied(message, exec_output)); - } - Ok(()) - } -} - enum SessionStatus { Alive { exit_code: Option,