openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-04 15:08:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

bound unified exec output collection

Browse Source
This commit is contained in:
jif-oai
2026-06-15 14:50:53 +02:00
parent 127224cacc
commit 2c8b819469
2 changed files with 89 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

63
codex-rs/core/src/unified_exec/mod_tests.rs
View File

@@ -21,7 +21,6 @@ use codex_exec_server::WriteResponse;
use codex_exec_server::WriteStatus;
use codex_sandboxing::SandboxType;
use codex_utils_output_truncation::TruncationPolicy;
use codex_utils_output_truncation::approx_token_count;
use core_test_support::get_remote_test_env;
use core_test_support::skip_if_sandbox;
use core_test_support::test_codex::test_env as remote_test_env;
@@ -163,7 +162,6 @@ async fn exec_command_with_tty(
)
.await;
let wall_time = Instant::now().saturating_duration_since(started_at);
let text = String::from_utf8_lossy(&collected).to_string();
let has_exited = process.has_exited();
let exit_code = process.exit_code();
let response_process_id = if process_started_alive && !has_exited {
@@ -189,12 +187,12 @@ async fn exec_command_with_tty(
event_call_id: context.call_id,
chunk_id: generate_chunk_id(),
wall_time,
raw_output: collected,
raw_output: collected.raw_output,
truncation_policy: turn.truncation_policy,
max_output_tokens: None,
process_id: response_process_id,
exit_code,
original_token_count: Some(approx_token_count(&text)),
original_token_count: Some(collected.original_token_count),
hook_command: Some(cmd.to_string()),
})
}
@@ -351,6 +349,61 @@ fn head_tail_buffer_default_preserves_prefix_and_suffix() {
assert!(rendered.ends_with(b"bc"));
}
#[tokio::test]
async fn collect_output_bounds_multiple_drains_and_reports_original_size() {
let output_buffer = Arc::new(tokio::sync::Mutex::new(HeadTailBuffer::default()));
let output_notify = Arc::new(Notify::new());
let output_closed = Arc::new(std::sync::atomic::AtomicBool::new(false));
let output_closed_notify = Arc::new(Notify::new());
let cancellation_token = tokio_util::sync::CancellationToken::new();
let chunk_len = UNIFIED_EXEC_OUTPUT_MAX_BYTES * 3 / 4;
output_buffer.lock().await.push_chunk(vec![b'a'; chunk_len]);
let collector = tokio::spawn({
let output_buffer = Arc::clone(&output_buffer);
let output_notify = Arc::clone(&output_notify);
let output_closed = Arc::clone(&output_closed);
let output_closed_notify = Arc::clone(&output_closed_notify);
let cancellation_token = cancellation_token.clone();
async move {
UnifiedExecProcessManager::collect_output_until_deadline(
&output_buffer,
&output_notify,
&output_closed,
&output_closed_notify,
&cancellation_token,
/*pause_state*/ None,
Instant::now() + Duration::from_secs(1),
)
.await
}
});
tokio::time::timeout(Duration::from_secs(1), async {
loop {
if output_buffer.lock().await.retained_bytes() == 0 {
break;
}
tokio::task::yield_now().await;
}
})
.await
.expect("collector should drain the first chunk");
output_buffer.lock().await.push_chunk(vec![b'b'; chunk_len]);
output_closed.store(true, std::sync::atomic::Ordering::Release);
output_closed_notify.notify_waiters();
output_notify.notify_waiters();
cancellation_token.cancel();
let collected = collector.await.expect("collector task should succeed");
let half = UNIFIED_EXEC_OUTPUT_MAX_BYTES / 2;
let expected_output = [vec![b'a'; half], vec![b'b'; half]].concat();
assert_eq!(collected.raw_output, expected_output);
assert_eq!(collected.original_token_count, chunk_len * 2 / 4);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn unified_exec_persists_across_requests() -> anyhow::Result<()> {
skip_if_sandbox!(Ok(()));
@@ -881,7 +934,7 @@ async fn unified_exec_uses_remote_exec_server_when_configured() -> anyhow::Resul
)
.await;
assert!(String::from_utf8_lossy(&collected).contains("remote-unified-exec"));
assert!(String::from_utf8_lossy(&collected.raw_output).contains("remote-unified-exec"));
Ok(())
}

41
codex-rs/core/src/unified_exec/process_manager.rs
View File

@@ -56,7 +56,7 @@ use codex_protocol::error::SandboxErr;
use codex_protocol::protocol::ExecCommandSource;
use codex_tools::ToolName;
use codex_utils_absolute_path::AbsolutePathBuf;
use codex_utils_output_truncation::approx_token_count;
use codex_utils_output_truncation::approx_tokens_from_byte_count;
use codex_utils_path_uri::PathUri;
const UNIFIED_EXEC_ENV: [(&str, &str); 10] = [
@@ -76,6 +76,11 @@ const NETWORK_ACCESS_DENIED_MESSAGE: &str =
const LATE_NETWORK_DENIAL_GRACE_PERIOD: Duration = Duration::from_millis(100);
const INTERRUPT: &str = "\u{3}";
pub(super) struct CollectedOutput {
pub(super) raw_output: Vec<u8>,
pub(super) original_token_count: usize,
}
/// Test-only override for deterministic unified exec process IDs.
///
/// In production builds this value should remain at its default (`false`) and
@@ -462,7 +467,10 @@ impl UnifiedExecProcessManager {
cancellation_token,
} = process.output_handles();
let deadline = start + Duration::from_millis(yield_time_ms);
let collected = Self::collect_output_until_deadline(
let CollectedOutput {
raw_output: collected,
original_token_count,
} = Self::collect_output_until_deadline(
&output_buffer,
&output_notify,
&output_closed,
@@ -596,7 +604,6 @@ impl UnifiedExecProcessManager {
(None, exit_code)
};
let original_token_count = approx_token_count(&text);
let response = ExecCommandToolOutput {
event_call_id: context.call_id.clone(),
chunk_id,
@@ -679,7 +686,10 @@ impl UnifiedExecProcessManager {
};
let start = Instant::now();
let deadline = start + Duration::from_millis(yield_time_ms);
let collected = Self::collect_output_until_deadline(
let CollectedOutput {
raw_output: collected,
original_token_count,
} = Self::collect_output_until_deadline(
&output_buffer,
&output_notify,
&output_closed,
@@ -691,8 +701,6 @@ impl UnifiedExecProcessManager {
.await;
let wall_time = Instant::now().saturating_duration_since(start);
let text = String::from_utf8_lossy(&collected).to_string();
let original_token_count = approx_token_count(&text);
let chunk_id = generate_chunk_id();
if network_approval
.as_ref()
@@ -1120,10 +1128,11 @@ impl UnifiedExecProcessManager {
cancellation_token: &CancellationToken,
mut pause_state: Option<watch::Receiver<bool>>,
mut deadline: Instant,
) -> Vec<u8> {
) -> CollectedOutput {
const POST_EXIT_CLOSE_WAIT_CAP: Duration = Duration::from_millis(50);
let mut collected: Vec<u8> = Vec::with_capacity(4096);
let mut collected = HeadTailBuffer::default();
let mut original_output_bytes = 0usize;
let mut exit_signal_received = cancellation_token.is_cancelled();
let mut post_exit_deadline: Option<Instant> = None;
loop {
@@ -1134,14 +1143,20 @@ impl UnifiedExecProcessManager {
)
.await;
let drained_chunks: Vec<Vec<u8>>;
let omitted_bytes: usize;
let mut wait_for_output = None;
{
let mut guard = output_buffer.lock().await;
omitted_bytes = guard.omitted_bytes();
drained_chunks = guard.drain_chunks();
if drained_chunks.is_empty() {
wait_for_output = Some(output_notify.notified());
}
}
original_output_bytes = drained_chunks.iter().fold(
original_output_bytes.saturating_add(omitted_bytes),
|total, chunk| total.saturating_add(chunk.len()),
);
if drained_chunks.is_empty() {
exit_signal_received |= cancellation_token.is_cancelled();
@@ -1189,7 +1204,7 @@ impl UnifiedExecProcessManager {
}
for chunk in drained_chunks {
collected.extend_from_slice(&chunk);
collected.push_chunk(chunk);
}
exit_signal_received |= cancellation_token.is_cancelled();
@@ -1198,7 +1213,13 @@ impl UnifiedExecProcessManager {
}
}
collected
CollectedOutput {
raw_output: collected.to_bytes(),
original_token_count: usize::try_from(approx_tokens_from_byte_count(
original_output_bytes,
))
.unwrap_or(usize::MAX),
}
}
async fn extend_deadlines_while_paused(