mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
## Why Code mode can use nested unified exec calls as data sources. When those calls omit `max_output_tokens`, code mode should receive raw command output so the script can parse or summarize it itself. When code mode does provide `max_output_tokens`, that explicit nested budget should be respected, including values above the default unified exec limit, rather than being capped before code mode sees the result. ## What - Preserve direct unified exec truncation behavior, while letting code-mode exec/write_stdin keep `max_output_tokens` as `None` unless explicitly supplied. - Make code-mode tool results use raw output when no explicit limit is present, and use the explicit nested limit directly when one is specified. - Refactor unified exec output formatting so `truncated_output` takes the caller-selected token budget. - Add e2e integration coverage for explicit nested exec limits, omitted nested exec limits, outer exec limit propagation, omitted-limit outputs that exceed both the default and a small truncation policy, explicit nested limits above those caps, and high explicit limits that still compact larger command output. - Reuse the code-mode turn setup helper while directly asserting the exact exec output item in each test. ## Testing - `just fmt` - `git diff --check` - Not run locally per repo guidance; CI should validate the e2e integration tests.
186 lines
5.9 KiB
Rust
186 lines
5.9 KiB
Rust
//! Unified Exec: interactive process execution orchestrated with approvals + sandboxing.
|
||
//!
|
||
//! Responsibilities
|
||
//! - Manages interactive processes (create, reuse, buffer output with caps).
|
||
//! - Uses the shared ToolOrchestrator to handle approval, sandbox selection, and
|
||
//! retry semantics in a single, descriptive flow.
|
||
//! - Spawns the PTY from a sandbox-transformed `ExecRequest`; on sandbox denial,
|
||
//! retries without sandbox when policy allows (no re‑prompt thanks to caching).
|
||
//! - Uses the shared `is_likely_sandbox_denied` heuristic to keep denial messages
|
||
//! consistent with other exec paths.
|
||
//!
|
||
//! Flow at a glance (open process)
|
||
//! 1) Build a small request `{ command, cwd }`.
|
||
//! 2) Orchestrator: approval (bypass/cache/prompt) → select sandbox → run.
|
||
//! 3) Runtime: transform `SandboxTransformRequest` -> `ExecRequest` -> spawn PTY.
|
||
//! 4) If denial, orchestrator retries with `SandboxType::None`.
|
||
//! 5) Process handle is returned with streaming output + metadata.
|
||
//!
|
||
//! This keeps policy logic and user interaction centralized while the PTY/process
|
||
//! concerns remain isolated here. The implementation is split between:
|
||
//! - `process.rs`: PTY process lifecycle + output buffering.
|
||
//! - `process_state.rs`: shared exit/failure state for local and remote processes.
|
||
//! - `process_manager.rs`: orchestration (approvals, sandboxing, reuse) and request handling.
|
||
|
||
use std::collections::HashMap;
|
||
use std::collections::HashSet;
|
||
use std::sync::Arc;
|
||
use std::sync::Weak;
|
||
|
||
use codex_exec_server::Environment;
|
||
use codex_network_proxy::NetworkProxy;
|
||
use codex_protocol::models::AdditionalPermissionProfile;
|
||
use codex_utils_absolute_path::AbsolutePathBuf;
|
||
use codex_utils_output_truncation::TruncationPolicy;
|
||
use rand::Rng;
|
||
use rand::rng;
|
||
use tokio::sync::Mutex;
|
||
|
||
use crate::sandboxing::SandboxPermissions;
|
||
use crate::session::session::Session;
|
||
use crate::session::turn_context::TurnContext;
|
||
use crate::shell::ShellType;
|
||
use crate::tools::network_approval::DeferredNetworkApproval;
|
||
|
||
mod async_watcher;
|
||
mod errors;
|
||
mod head_tail_buffer;
|
||
mod process;
|
||
mod process_manager;
|
||
mod process_state;
|
||
|
||
pub(crate) fn set_deterministic_process_ids_for_tests(enabled: bool) {
|
||
process_manager::set_deterministic_process_ids_for_tests(enabled);
|
||
}
|
||
|
||
pub(crate) use errors::UnifiedExecError;
|
||
pub(crate) use process::NoopSpawnLifecycle;
|
||
#[cfg(unix)]
|
||
pub(crate) use process::SpawnLifecycle;
|
||
pub(crate) use process::SpawnLifecycleHandle;
|
||
pub(crate) use process::UnifiedExecProcess;
|
||
|
||
pub(crate) const MIN_YIELD_TIME_MS: u64 = 250;
|
||
// Minimum yield time for an empty `write_stdin`.
|
||
pub(crate) const MIN_EMPTY_YIELD_TIME_MS: u64 = 5_000;
|
||
pub(crate) const MAX_YIELD_TIME_MS: u64 = 30_000;
|
||
pub(crate) const DEFAULT_MAX_BACKGROUND_TERMINAL_TIMEOUT_MS: u64 = 300_000;
|
||
pub(crate) const DEFAULT_MAX_OUTPUT_TOKENS: usize = 10_000;
|
||
pub(crate) const UNIFIED_EXEC_OUTPUT_MAX_BYTES: usize = 1024 * 1024; // 1 MiB
|
||
pub(crate) const UNIFIED_EXEC_OUTPUT_MAX_TOKENS: usize = UNIFIED_EXEC_OUTPUT_MAX_BYTES / 4;
|
||
pub(crate) const MAX_UNIFIED_EXEC_PROCESSES: usize = 64;
|
||
|
||
pub(crate) struct UnifiedExecContext {
|
||
pub session: Arc<Session>,
|
||
pub turn: Arc<TurnContext>,
|
||
pub call_id: String,
|
||
}
|
||
|
||
impl UnifiedExecContext {
|
||
pub fn new(session: Arc<Session>, turn: Arc<TurnContext>, call_id: String) -> Self {
|
||
Self {
|
||
session,
|
||
turn,
|
||
call_id,
|
||
}
|
||
}
|
||
}
|
||
|
||
#[derive(Debug)]
|
||
pub(crate) struct ExecCommandRequest {
|
||
pub command: Vec<String>,
|
||
pub shell_type: ShellType,
|
||
pub hook_command: String,
|
||
pub process_id: i32,
|
||
pub yield_time_ms: u64,
|
||
pub max_output_tokens: Option<usize>,
|
||
pub cwd: AbsolutePathBuf,
|
||
pub sandbox_cwd: AbsolutePathBuf,
|
||
pub environment: Arc<Environment>,
|
||
pub network: Option<NetworkProxy>,
|
||
pub tty: bool,
|
||
pub sandbox_permissions: SandboxPermissions,
|
||
pub additional_permissions: Option<AdditionalPermissionProfile>,
|
||
pub additional_permissions_preapproved: bool,
|
||
pub justification: Option<String>,
|
||
pub prefix_rule: Option<Vec<String>>,
|
||
}
|
||
|
||
#[derive(Debug)]
|
||
pub(crate) struct WriteStdinRequest<'a> {
|
||
pub process_id: i32,
|
||
pub input: &'a str,
|
||
pub yield_time_ms: u64,
|
||
pub max_output_tokens: Option<usize>,
|
||
pub truncation_policy: TruncationPolicy,
|
||
}
|
||
|
||
#[derive(Default)]
|
||
pub(crate) struct ProcessStore {
|
||
processes: HashMap<i32, ProcessEntry>,
|
||
reserved_process_ids: HashSet<i32>,
|
||
}
|
||
|
||
impl ProcessStore {
|
||
fn remove(&mut self, process_id: i32) -> Option<ProcessEntry> {
|
||
self.reserved_process_ids.remove(&process_id);
|
||
self.processes.remove(&process_id)
|
||
}
|
||
}
|
||
|
||
pub(crate) struct UnifiedExecProcessManager {
|
||
process_store: Mutex<ProcessStore>,
|
||
max_write_stdin_yield_time_ms: u64,
|
||
}
|
||
|
||
impl UnifiedExecProcessManager {
|
||
pub(crate) fn new(max_write_stdin_yield_time_ms: u64) -> Self {
|
||
Self {
|
||
process_store: Mutex::new(ProcessStore::default()),
|
||
max_write_stdin_yield_time_ms: max_write_stdin_yield_time_ms
|
||
.max(MIN_EMPTY_YIELD_TIME_MS),
|
||
}
|
||
}
|
||
}
|
||
|
||
impl Default for UnifiedExecProcessManager {
|
||
fn default() -> Self {
|
||
Self::new(DEFAULT_MAX_BACKGROUND_TERMINAL_TIMEOUT_MS)
|
||
}
|
||
}
|
||
|
||
struct ProcessEntry {
|
||
process: Arc<UnifiedExecProcess>,
|
||
call_id: String,
|
||
process_id: i32,
|
||
hook_command: String,
|
||
tty: bool,
|
||
network_approval: Option<DeferredNetworkApproval>,
|
||
session: Weak<Session>,
|
||
last_used: tokio::time::Instant,
|
||
}
|
||
|
||
pub(crate) fn clamp_yield_time(yield_time_ms: u64) -> u64 {
|
||
yield_time_ms.clamp(MIN_YIELD_TIME_MS, MAX_YIELD_TIME_MS)
|
||
}
|
||
|
||
pub(crate) fn resolve_max_tokens(max_tokens: Option<usize>) -> usize {
|
||
max_tokens.unwrap_or(DEFAULT_MAX_OUTPUT_TOKENS)
|
||
}
|
||
|
||
pub(crate) fn generate_chunk_id() -> String {
|
||
let mut rng = rng();
|
||
(0..6)
|
||
.map(|_| format!("{:x}", rng.random_range(0..16)))
|
||
.collect()
|
||
}
|
||
|
||
#[cfg(test)]
|
||
#[cfg(unix)]
|
||
#[path = "process_tests.rs"]
|
||
mod process_tests;
|
||
#[cfg(test)]
|
||
#[cfg(unix)]
|
||
#[path = "mod_tests.rs"]
|
||
mod tests;
|