mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
committed by
GitHub
Unverified
parent
124a09e577
commit
4cef89a122
@@ -8,7 +8,7 @@ use tokio::time::Instant;
|
||||
use tokio::time::Sleep;
|
||||
|
||||
use super::UnifiedExecContext;
|
||||
use super::session::UnifiedExecSession;
|
||||
use super::process::UnifiedExecProcess;
|
||||
use crate::codex::Session;
|
||||
use crate::codex::TurnContext;
|
||||
use crate::exec::ExecToolCallOutput;
|
||||
@@ -37,13 +37,13 @@ const UNIFIED_EXEC_OUTPUT_DELTA_MAX_BYTES: usize = 8192;
|
||||
/// shared transcript, and emits ExecCommandOutputDelta events on UTF‑8
|
||||
/// boundaries.
|
||||
pub(crate) fn start_streaming_output(
|
||||
session: &UnifiedExecSession,
|
||||
process: &UnifiedExecProcess,
|
||||
context: &UnifiedExecContext,
|
||||
transcript: Arc<Mutex<HeadTailBuffer>>,
|
||||
) {
|
||||
let mut receiver = session.output_receiver();
|
||||
let output_drained = session.output_drained_notify();
|
||||
let exit_token = session.cancellation_token();
|
||||
let mut receiver = process.output_receiver();
|
||||
let output_drained = process.output_drained_notify();
|
||||
let exit_token = process.cancellation_token();
|
||||
|
||||
let session_ref = Arc::clone(&context.session);
|
||||
let turn_ref = Arc::clone(&context.turn);
|
||||
@@ -104,7 +104,7 @@ pub(crate) fn start_streaming_output(
|
||||
/// single ExecCommandEnd event with the aggregated transcript.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) fn spawn_exit_watcher(
|
||||
session: Arc<UnifiedExecSession>,
|
||||
process: Arc<UnifiedExecProcess>,
|
||||
session_ref: Arc<Session>,
|
||||
turn_ref: Arc<TurnContext>,
|
||||
call_id: String,
|
||||
@@ -114,14 +114,14 @@ pub(crate) fn spawn_exit_watcher(
|
||||
transcript: Arc<Mutex<HeadTailBuffer>>,
|
||||
started_at: Instant,
|
||||
) {
|
||||
let exit_token = session.cancellation_token();
|
||||
let output_drained = session.output_drained_notify();
|
||||
let exit_token = process.cancellation_token();
|
||||
let output_drained = process.output_drained_notify();
|
||||
|
||||
tokio::spawn(async move {
|
||||
exit_token.cancelled().await;
|
||||
output_drained.notified().await;
|
||||
|
||||
let exit_code = session.exit_code().unwrap_or(-1);
|
||||
let exit_code = process.exit_code().unwrap_or(-1);
|
||||
let duration = Instant::now().saturating_duration_since(started_at);
|
||||
emit_exec_end_for_unified_exec(
|
||||
session_ref,
|
||||
|
||||
@@ -3,11 +3,11 @@ use thiserror::Error;
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub(crate) enum UnifiedExecError {
|
||||
#[error("Failed to create unified exec session: {message}")]
|
||||
CreateSession { message: String },
|
||||
// Called "session" in the model's training.
|
||||
#[error("Unknown session id {process_id}")]
|
||||
UnknownSessionId { process_id: String },
|
||||
#[error("Failed to create unified exec process: {message}")]
|
||||
CreateProcess { message: String },
|
||||
// The model is trained on `session_id`, but internally we track a `process_id`.
|
||||
#[error("Unknown process id {process_id}")]
|
||||
UnknownProcessId { process_id: String },
|
||||
#[error("failed to write to stdin")]
|
||||
WriteToStdin,
|
||||
#[error("missing command line for unified exec request")]
|
||||
@@ -20,8 +20,8 @@ pub(crate) enum UnifiedExecError {
|
||||
}
|
||||
|
||||
impl UnifiedExecError {
|
||||
pub(crate) fn create_session(message: String) -> Self {
|
||||
Self::CreateSession { message }
|
||||
pub(crate) fn create_process(message: String) -> Self {
|
||||
Self::CreateProcess { message }
|
||||
}
|
||||
|
||||
pub(crate) fn sandbox_denied(message: String, output: ExecToolCallOutput) -> Self {
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
//! Unified Exec: interactive PTY execution orchestrated with approvals + sandboxing.
|
||||
//! Unified Exec: interactive process execution orchestrated with approvals + sandboxing.
|
||||
//!
|
||||
//! Responsibilities
|
||||
//! - Manages interactive PTY sessions (create, reuse, buffer output with caps).
|
||||
//! - 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 `ExecEnv`; on sandbox denial,
|
||||
@@ -9,17 +9,17 @@
|
||||
//! - Uses the shared `is_likely_sandbox_denied` heuristic to keep denial messages
|
||||
//! consistent with other exec paths.
|
||||
//!
|
||||
//! Flow at a glance (open session)
|
||||
//! 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 `CommandSpec` → `ExecEnv` → spawn PTY.
|
||||
//! 4) If denial, orchestrator retries with `SandboxType::None`.
|
||||
//! 5) Session is returned with streaming output + metadata.
|
||||
//! 5) Process handle is returned with streaming output + metadata.
|
||||
//!
|
||||
//! This keeps policy logic and user interaction centralized while the PTY/session
|
||||
//! This keeps policy logic and user interaction centralized while the PTY/process
|
||||
//! concerns remain isolated here. The implementation is split between:
|
||||
//! - `session.rs`: PTY session lifecycle + output buffering.
|
||||
//! - `session_manager.rs`: orchestration (approvals, sandboxing, reuse) and request handling.
|
||||
//! - `process.rs`: PTY process lifecycle + output buffering.
|
||||
//! - `process_manager.rs`: orchestration (approvals, sandboxing, reuse) and request handling.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::collections::HashSet;
|
||||
@@ -38,21 +38,21 @@ use crate::sandboxing::SandboxPermissions;
|
||||
mod async_watcher;
|
||||
mod errors;
|
||||
mod head_tail_buffer;
|
||||
mod session;
|
||||
mod session_manager;
|
||||
mod process;
|
||||
mod process_manager;
|
||||
|
||||
pub(crate) use errors::UnifiedExecError;
|
||||
pub(crate) use session::UnifiedExecSession;
|
||||
pub(crate) use process::UnifiedExecProcess;
|
||||
|
||||
pub(crate) const MIN_YIELD_TIME_MS: u64 = 250;
|
||||
pub(crate) const MAX_YIELD_TIME_MS: u64 = 30_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_SESSIONS: usize = 64;
|
||||
pub(crate) const MAX_UNIFIED_EXEC_PROCESSES: usize = 64;
|
||||
|
||||
// Send a warning message to the models when it reaches this number of sessions.
|
||||
pub(crate) const WARNING_UNIFIED_EXEC_SESSIONS: usize = 60;
|
||||
// Send a warning message to the models when it reaches this number of processes.
|
||||
pub(crate) const WARNING_UNIFIED_EXEC_PROCESSES: usize = 60;
|
||||
|
||||
pub(crate) struct UnifiedExecContext {
|
||||
pub session: Arc<Session>,
|
||||
@@ -104,32 +104,32 @@ pub(crate) struct UnifiedExecResponse {
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub(crate) struct SessionStore {
|
||||
sessions: HashMap<String, SessionEntry>,
|
||||
reserved_sessions_id: HashSet<String>,
|
||||
pub(crate) struct ProcessStore {
|
||||
processes: HashMap<String, ProcessEntry>,
|
||||
reserved_process_ids: HashSet<String>,
|
||||
}
|
||||
|
||||
impl SessionStore {
|
||||
fn remove(&mut self, session_id: &str) -> Option<SessionEntry> {
|
||||
self.reserved_sessions_id.remove(session_id);
|
||||
self.sessions.remove(session_id)
|
||||
impl ProcessStore {
|
||||
fn remove(&mut self, process_id: &str) -> Option<ProcessEntry> {
|
||||
self.reserved_process_ids.remove(process_id);
|
||||
self.processes.remove(process_id)
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) struct UnifiedExecSessionManager {
|
||||
session_store: Mutex<SessionStore>,
|
||||
pub(crate) struct UnifiedExecProcessManager {
|
||||
process_store: Mutex<ProcessStore>,
|
||||
}
|
||||
|
||||
impl Default for UnifiedExecSessionManager {
|
||||
impl Default for UnifiedExecProcessManager {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
session_store: Mutex::new(SessionStore::default()),
|
||||
process_store: Mutex::new(ProcessStore::default()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct SessionEntry {
|
||||
session: Arc<UnifiedExecSession>,
|
||||
struct ProcessEntry {
|
||||
process: Arc<UnifiedExecProcess>,
|
||||
session_ref: Arc<Session>,
|
||||
turn_ref: Arc<TurnContext>,
|
||||
call_id: String,
|
||||
@@ -421,10 +421,10 @@ mod tests {
|
||||
session
|
||||
.services
|
||||
.unified_exec_manager
|
||||
.session_store
|
||||
.process_store
|
||||
.lock()
|
||||
.await
|
||||
.sessions
|
||||
.processes
|
||||
.is_empty()
|
||||
);
|
||||
|
||||
@@ -432,7 +432,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn reusing_completed_session_returns_unknown_session() -> anyhow::Result<()> {
|
||||
async fn reusing_completed_process_returns_unknown_process() -> anyhow::Result<()> {
|
||||
skip_if_sandbox!(Ok(()));
|
||||
|
||||
let (session, turn) = test_session_and_turn().await;
|
||||
@@ -450,23 +450,23 @@ mod tests {
|
||||
|
||||
let err = write_stdin(&session, process_id, "", 100)
|
||||
.await
|
||||
.expect_err("expected unknown session error");
|
||||
.expect_err("expected unknown process error");
|
||||
|
||||
match err {
|
||||
UnifiedExecError::UnknownSessionId { process_id: err_id } => {
|
||||
UnifiedExecError::UnknownProcessId { process_id: err_id } => {
|
||||
assert_eq!(err_id, process_id, "process id should match request");
|
||||
}
|
||||
other => panic!("expected UnknownSessionId, got {other:?}"),
|
||||
other => panic!("expected UnknownProcessId, got {other:?}"),
|
||||
}
|
||||
|
||||
assert!(
|
||||
session
|
||||
.services
|
||||
.unified_exec_manager
|
||||
.session_store
|
||||
.process_store
|
||||
.lock()
|
||||
.await
|
||||
.sessions
|
||||
.processes
|
||||
.is_empty()
|
||||
);
|
||||
|
||||
|
||||
+14
-14
@@ -30,8 +30,8 @@ pub(crate) struct OutputHandles {
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct UnifiedExecSession {
|
||||
session: ExecCommandSession,
|
||||
pub(crate) struct UnifiedExecProcess {
|
||||
process_handle: ExecCommandSession,
|
||||
output_buffer: OutputBuffer,
|
||||
output_notify: Arc<Notify>,
|
||||
cancellation_token: CancellationToken,
|
||||
@@ -40,9 +40,9 @@ pub(crate) struct UnifiedExecSession {
|
||||
sandbox_type: SandboxType,
|
||||
}
|
||||
|
||||
impl UnifiedExecSession {
|
||||
impl UnifiedExecProcess {
|
||||
pub(super) fn new(
|
||||
session: ExecCommandSession,
|
||||
process_handle: ExecCommandSession,
|
||||
initial_output_rx: tokio::sync::broadcast::Receiver<Vec<u8>>,
|
||||
sandbox_type: SandboxType,
|
||||
) -> Self {
|
||||
@@ -69,7 +69,7 @@ impl UnifiedExecSession {
|
||||
});
|
||||
|
||||
Self {
|
||||
session,
|
||||
process_handle,
|
||||
output_buffer,
|
||||
output_notify,
|
||||
cancellation_token,
|
||||
@@ -80,7 +80,7 @@ impl UnifiedExecSession {
|
||||
}
|
||||
|
||||
pub(super) fn writer_sender(&self) -> mpsc::Sender<Vec<u8>> {
|
||||
self.session.writer_sender()
|
||||
self.process_handle.writer_sender()
|
||||
}
|
||||
|
||||
pub(super) fn output_handles(&self) -> OutputHandles {
|
||||
@@ -92,7 +92,7 @@ impl UnifiedExecSession {
|
||||
}
|
||||
|
||||
pub(super) fn output_receiver(&self) -> tokio::sync::broadcast::Receiver<Vec<u8>> {
|
||||
self.session.output_receiver()
|
||||
self.process_handle.output_receiver()
|
||||
}
|
||||
|
||||
pub(super) fn cancellation_token(&self) -> CancellationToken {
|
||||
@@ -104,15 +104,15 @@ impl UnifiedExecSession {
|
||||
}
|
||||
|
||||
pub(super) fn has_exited(&self) -> bool {
|
||||
self.session.has_exited()
|
||||
self.process_handle.has_exited()
|
||||
}
|
||||
|
||||
pub(super) fn exit_code(&self) -> Option<i32> {
|
||||
self.session.exit_code()
|
||||
self.process_handle.exit_code()
|
||||
}
|
||||
|
||||
pub(super) fn terminate(&self) {
|
||||
self.session.terminate();
|
||||
self.process_handle.terminate();
|
||||
self.cancellation_token.cancel();
|
||||
self.output_task.abort();
|
||||
}
|
||||
@@ -164,7 +164,7 @@ impl UnifiedExecSession {
|
||||
TruncationPolicy::Tokens(UNIFIED_EXEC_OUTPUT_MAX_TOKENS),
|
||||
);
|
||||
let message = if snippet.is_empty() {
|
||||
format!("Session exited with code {exit_code}")
|
||||
format!("Process exited with code {exit_code}")
|
||||
} else {
|
||||
snippet
|
||||
};
|
||||
@@ -178,11 +178,11 @@ impl UnifiedExecSession {
|
||||
sandbox_type: SandboxType,
|
||||
) -> Result<Self, UnifiedExecError> {
|
||||
let SpawnedPty {
|
||||
session,
|
||||
session: process_handle,
|
||||
output_rx,
|
||||
mut exit_rx,
|
||||
} = spawned;
|
||||
let managed = Self::new(session, output_rx, sandbox_type);
|
||||
let managed = Self::new(process_handle, output_rx, sandbox_type);
|
||||
|
||||
let exit_ready = matches!(exit_rx.try_recv(), Ok(_) | Err(TryRecvError::Closed));
|
||||
|
||||
@@ -217,7 +217,7 @@ impl UnifiedExecSession {
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for UnifiedExecSession {
|
||||
impl Drop for UnifiedExecProcess {
|
||||
fn drop(&mut self) {
|
||||
self.terminate();
|
||||
}
|
||||
+93
-90
@@ -30,14 +30,14 @@ use crate::truncate::TruncationPolicy;
|
||||
use crate::truncate::approx_token_count;
|
||||
use crate::truncate::formatted_truncate_text;
|
||||
use crate::unified_exec::ExecCommandRequest;
|
||||
use crate::unified_exec::MAX_UNIFIED_EXEC_SESSIONS;
|
||||
use crate::unified_exec::SessionEntry;
|
||||
use crate::unified_exec::SessionStore;
|
||||
use crate::unified_exec::MAX_UNIFIED_EXEC_PROCESSES;
|
||||
use crate::unified_exec::ProcessEntry;
|
||||
use crate::unified_exec::ProcessStore;
|
||||
use crate::unified_exec::UnifiedExecContext;
|
||||
use crate::unified_exec::UnifiedExecError;
|
||||
use crate::unified_exec::UnifiedExecProcessManager;
|
||||
use crate::unified_exec::UnifiedExecResponse;
|
||||
use crate::unified_exec::UnifiedExecSessionManager;
|
||||
use crate::unified_exec::WARNING_UNIFIED_EXEC_SESSIONS;
|
||||
use crate::unified_exec::WARNING_UNIFIED_EXEC_PROCESSES;
|
||||
use crate::unified_exec::WriteStdinRequest;
|
||||
use crate::unified_exec::async_watcher::emit_exec_end_for_unified_exec;
|
||||
use crate::unified_exec::async_watcher::spawn_exit_watcher;
|
||||
@@ -45,10 +45,10 @@ use crate::unified_exec::async_watcher::start_streaming_output;
|
||||
use crate::unified_exec::clamp_yield_time;
|
||||
use crate::unified_exec::generate_chunk_id;
|
||||
use crate::unified_exec::head_tail_buffer::HeadTailBuffer;
|
||||
use crate::unified_exec::process::OutputBuffer;
|
||||
use crate::unified_exec::process::OutputHandles;
|
||||
use crate::unified_exec::process::UnifiedExecProcess;
|
||||
use crate::unified_exec::resolve_max_tokens;
|
||||
use crate::unified_exec::session::OutputBuffer;
|
||||
use crate::unified_exec::session::OutputHandles;
|
||||
use crate::unified_exec::session::UnifiedExecSession;
|
||||
|
||||
const UNIFIED_EXEC_ENV: [(&str, &str); 9] = [
|
||||
("NO_COLOR", "1"),
|
||||
@@ -69,7 +69,7 @@ fn apply_unified_exec_env(mut env: HashMap<String, String>) -> HashMap<String, S
|
||||
env
|
||||
}
|
||||
|
||||
struct PreparedSessionHandles {
|
||||
struct PreparedProcessHandles {
|
||||
writer_tx: mpsc::Sender<Vec<u8>>,
|
||||
output_buffer: OutputBuffer,
|
||||
output_notify: Arc<Notify>,
|
||||
@@ -80,10 +80,10 @@ struct PreparedSessionHandles {
|
||||
process_id: String,
|
||||
}
|
||||
|
||||
impl UnifiedExecSessionManager {
|
||||
impl UnifiedExecProcessManager {
|
||||
pub(crate) async fn allocate_process_id(&self) -> String {
|
||||
loop {
|
||||
let mut store = self.session_store.lock().await;
|
||||
let mut store = self.process_store.lock().await;
|
||||
|
||||
let process_id = if !cfg!(test) && !cfg!(feature = "deterministic_process_ids") {
|
||||
// production mode → random
|
||||
@@ -91,7 +91,7 @@ impl UnifiedExecSessionManager {
|
||||
} else {
|
||||
// test or deterministic mode
|
||||
let next = store
|
||||
.reserved_sessions_id
|
||||
.reserved_process_ids
|
||||
.iter()
|
||||
.filter_map(|s| s.parse::<i32>().ok())
|
||||
.max()
|
||||
@@ -101,17 +101,17 @@ impl UnifiedExecSessionManager {
|
||||
next.to_string()
|
||||
};
|
||||
|
||||
if store.reserved_sessions_id.contains(&process_id) {
|
||||
if store.reserved_process_ids.contains(&process_id) {
|
||||
continue;
|
||||
}
|
||||
|
||||
store.reserved_sessions_id.insert(process_id.clone());
|
||||
store.reserved_process_ids.insert(process_id.clone());
|
||||
return process_id;
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn release_process_id(&self, process_id: &str) {
|
||||
let mut store = self.session_store.lock().await;
|
||||
let mut store = self.process_store.lock().await;
|
||||
store.remove(process_id);
|
||||
}
|
||||
|
||||
@@ -125,7 +125,7 @@ impl UnifiedExecSessionManager {
|
||||
.clone()
|
||||
.unwrap_or_else(|| context.turn.cwd.clone());
|
||||
|
||||
let session = self
|
||||
let process = self
|
||||
.open_session_with_sandbox(
|
||||
&request.command,
|
||||
cwd.clone(),
|
||||
@@ -135,8 +135,8 @@ impl UnifiedExecSessionManager {
|
||||
)
|
||||
.await;
|
||||
|
||||
let session = match session {
|
||||
Ok(session) => Arc::new(session),
|
||||
let process = match process {
|
||||
Ok(process) => Arc::new(process),
|
||||
Err(err) => {
|
||||
self.release_process_id(&request.process_id).await;
|
||||
return Err(err);
|
||||
@@ -158,7 +158,7 @@ impl UnifiedExecSessionManager {
|
||||
);
|
||||
emitter.emit(event_ctx, ToolEventStage::Begin).await;
|
||||
|
||||
start_streaming_output(&session, context, Arc::clone(&transcript));
|
||||
start_streaming_output(&process, context, Arc::clone(&transcript));
|
||||
|
||||
let max_tokens = resolve_max_tokens(request.max_output_tokens);
|
||||
let yield_time_ms = clamp_yield_time(request.yield_time_ms);
|
||||
@@ -171,7 +171,7 @@ impl UnifiedExecSessionManager {
|
||||
output_buffer,
|
||||
output_notify,
|
||||
cancellation_token,
|
||||
} = session.output_handles();
|
||||
} = process.output_handles();
|
||||
let deadline = start + Duration::from_millis(yield_time_ms);
|
||||
let collected = Self::collect_output_until_deadline(
|
||||
&output_buffer,
|
||||
@@ -184,8 +184,8 @@ impl UnifiedExecSessionManager {
|
||||
|
||||
let text = String::from_utf8_lossy(&collected).to_string();
|
||||
let output = formatted_truncate_text(&text, TruncationPolicy::Tokens(max_tokens));
|
||||
let exit_code = session.exit_code();
|
||||
let has_exited = session.has_exited() || exit_code.is_some();
|
||||
let exit_code = process.exit_code();
|
||||
let has_exited = process.has_exited() || exit_code.is_some();
|
||||
let chunk_id = generate_chunk_id();
|
||||
let process_id = request.process_id.clone();
|
||||
if has_exited {
|
||||
@@ -208,14 +208,14 @@ impl UnifiedExecSessionManager {
|
||||
.await;
|
||||
|
||||
self.release_process_id(&request.process_id).await;
|
||||
session.check_for_sandbox_denial_with_text(&text).await?;
|
||||
process.check_for_sandbox_denial_with_text(&text).await?;
|
||||
} else {
|
||||
// Long‑lived command: persist the session so write_stdin can reuse
|
||||
// Long‑lived command: persist the process so write_stdin can reuse
|
||||
// it, and register a background watcher that will emit
|
||||
// ExecCommandEnd when the PTY eventually exits (even if no further
|
||||
// tool calls are made).
|
||||
self.store_session(
|
||||
Arc::clone(&session),
|
||||
self.store_process(
|
||||
Arc::clone(&process),
|
||||
context,
|
||||
&request.command,
|
||||
cwd.clone(),
|
||||
@@ -254,7 +254,7 @@ impl UnifiedExecSessionManager {
|
||||
) -> Result<UnifiedExecResponse, UnifiedExecError> {
|
||||
let process_id = request.process_id.to_string();
|
||||
|
||||
let PreparedSessionHandles {
|
||||
let PreparedProcessHandles {
|
||||
writer_tx,
|
||||
output_buffer,
|
||||
output_notify,
|
||||
@@ -264,7 +264,7 @@ impl UnifiedExecSessionManager {
|
||||
command: session_command,
|
||||
process_id,
|
||||
..
|
||||
} = self.prepare_session_handles(process_id.as_str()).await?;
|
||||
} = self.prepare_process_handles(process_id.as_str()).await?;
|
||||
|
||||
if !request.input.is_empty() {
|
||||
Self::send_input(&writer_tx, request.input.as_bytes()).await?;
|
||||
@@ -291,23 +291,23 @@ impl UnifiedExecSessionManager {
|
||||
let original_token_count = approx_token_count(&text);
|
||||
let chunk_id = generate_chunk_id();
|
||||
|
||||
// After polling, refresh_session_state tells us whether the PTY is
|
||||
// After polling, refresh_process_state tells us whether the PTY is
|
||||
// still alive or has exited and been removed from the store; we thread
|
||||
// that through so the handler can tag TerminalInteraction with an
|
||||
// appropriate process_id and exit_code.
|
||||
let status = self.refresh_session_state(process_id.as_str()).await;
|
||||
let status = self.refresh_process_state(process_id.as_str()).await;
|
||||
let (process_id, exit_code, event_call_id) = match status {
|
||||
SessionStatus::Alive {
|
||||
ProcessStatus::Alive {
|
||||
exit_code,
|
||||
call_id,
|
||||
process_id,
|
||||
} => (Some(process_id), exit_code, call_id),
|
||||
SessionStatus::Exited { exit_code, entry } => {
|
||||
ProcessStatus::Exited { exit_code, entry } => {
|
||||
let call_id = entry.call_id.clone();
|
||||
(None, exit_code, call_id)
|
||||
}
|
||||
SessionStatus::Unknown => {
|
||||
return Err(UnifiedExecError::UnknownSessionId {
|
||||
ProcessStatus::Unknown => {
|
||||
return Err(UnifiedExecError::UnknownProcessId {
|
||||
process_id: request.process_id.to_string(),
|
||||
});
|
||||
}
|
||||
@@ -332,25 +332,25 @@ impl UnifiedExecSessionManager {
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
async fn refresh_session_state(&self, process_id: &str) -> SessionStatus {
|
||||
let mut store = self.session_store.lock().await;
|
||||
let Some(entry) = store.sessions.get(process_id) else {
|
||||
return SessionStatus::Unknown;
|
||||
async fn refresh_process_state(&self, process_id: &str) -> ProcessStatus {
|
||||
let mut store = self.process_store.lock().await;
|
||||
let Some(entry) = store.processes.get(process_id) else {
|
||||
return ProcessStatus::Unknown;
|
||||
};
|
||||
|
||||
let exit_code = entry.session.exit_code();
|
||||
let exit_code = entry.process.exit_code();
|
||||
let process_id = entry.process_id.clone();
|
||||
|
||||
if entry.session.has_exited() {
|
||||
if entry.process.has_exited() {
|
||||
let Some(entry) = store.remove(&process_id) else {
|
||||
return SessionStatus::Unknown;
|
||||
return ProcessStatus::Unknown;
|
||||
};
|
||||
SessionStatus::Exited {
|
||||
ProcessStatus::Exited {
|
||||
exit_code,
|
||||
entry: Box::new(entry),
|
||||
}
|
||||
} else {
|
||||
SessionStatus::Alive {
|
||||
ProcessStatus::Alive {
|
||||
exit_code,
|
||||
call_id: entry.call_id.clone(),
|
||||
process_id,
|
||||
@@ -358,16 +358,16 @@ impl UnifiedExecSessionManager {
|
||||
}
|
||||
}
|
||||
|
||||
async fn prepare_session_handles(
|
||||
async fn prepare_process_handles(
|
||||
&self,
|
||||
process_id: &str,
|
||||
) -> Result<PreparedSessionHandles, UnifiedExecError> {
|
||||
let mut store = self.session_store.lock().await;
|
||||
) -> Result<PreparedProcessHandles, UnifiedExecError> {
|
||||
let mut store = self.process_store.lock().await;
|
||||
let entry =
|
||||
store
|
||||
.sessions
|
||||
.processes
|
||||
.get_mut(process_id)
|
||||
.ok_or(UnifiedExecError::UnknownSessionId {
|
||||
.ok_or(UnifiedExecError::UnknownProcessId {
|
||||
process_id: process_id.to_string(),
|
||||
})?;
|
||||
entry.last_used = Instant::now();
|
||||
@@ -375,10 +375,10 @@ impl UnifiedExecSessionManager {
|
||||
output_buffer,
|
||||
output_notify,
|
||||
cancellation_token,
|
||||
} = entry.session.output_handles();
|
||||
} = entry.process.output_handles();
|
||||
|
||||
Ok(PreparedSessionHandles {
|
||||
writer_tx: entry.session.writer_sender(),
|
||||
Ok(PreparedProcessHandles {
|
||||
writer_tx: entry.process.writer_sender(),
|
||||
output_buffer,
|
||||
output_notify,
|
||||
cancellation_token,
|
||||
@@ -400,9 +400,9 @@ impl UnifiedExecSessionManager {
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
async fn store_session(
|
||||
async fn store_process(
|
||||
&self,
|
||||
session: Arc<UnifiedExecSession>,
|
||||
process: Arc<UnifiedExecProcess>,
|
||||
context: &UnifiedExecContext,
|
||||
command: &[String],
|
||||
cwd: PathBuf,
|
||||
@@ -410,8 +410,8 @@ impl UnifiedExecSessionManager {
|
||||
process_id: String,
|
||||
transcript: Arc<tokio::sync::Mutex<HeadTailBuffer>>,
|
||||
) {
|
||||
let entry = SessionEntry {
|
||||
session: Arc::clone(&session),
|
||||
let entry = ProcessEntry {
|
||||
process: Arc::clone(&process),
|
||||
session_ref: Arc::clone(&context.session),
|
||||
turn_ref: Arc::clone(&context.turn),
|
||||
call_id: context.call_id.clone(),
|
||||
@@ -419,25 +419,25 @@ impl UnifiedExecSessionManager {
|
||||
command: command.to_vec(),
|
||||
last_used: started_at,
|
||||
};
|
||||
let number_sessions = {
|
||||
let mut store = self.session_store.lock().await;
|
||||
Self::prune_sessions_if_needed(&mut store);
|
||||
store.sessions.insert(process_id.clone(), entry);
|
||||
store.sessions.len()
|
||||
let number_processes = {
|
||||
let mut store = self.process_store.lock().await;
|
||||
Self::prune_processes_if_needed(&mut store);
|
||||
store.processes.insert(process_id.clone(), entry);
|
||||
store.processes.len()
|
||||
};
|
||||
|
||||
if number_sessions >= WARNING_UNIFIED_EXEC_SESSIONS {
|
||||
if number_processes >= WARNING_UNIFIED_EXEC_PROCESSES {
|
||||
context
|
||||
.session
|
||||
.record_model_warning(
|
||||
format!("The maximum number of unified exec sessions you can keep open is {WARNING_UNIFIED_EXEC_SESSIONS} and you currently have {number_sessions} sessions open. Reuse older sessions or close them to prevent automatic pruning of old session"),
|
||||
format!("The maximum number of unified exec processes you can keep open is {WARNING_UNIFIED_EXEC_PROCESSES} and you currently have {number_processes} processes open. Reuse older processes or close them to prevent automatic pruning of old processes"),
|
||||
&context.turn
|
||||
)
|
||||
.await;
|
||||
};
|
||||
|
||||
spawn_exit_watcher(
|
||||
Arc::clone(&session),
|
||||
Arc::clone(&process),
|
||||
Arc::clone(&context.session),
|
||||
Arc::clone(&context.turn),
|
||||
context.call_id.clone(),
|
||||
@@ -471,7 +471,7 @@ impl UnifiedExecSessionManager {
|
||||
pub(crate) async fn open_session_with_exec_env(
|
||||
&self,
|
||||
env: &ExecEnv,
|
||||
) -> Result<UnifiedExecSession, UnifiedExecError> {
|
||||
) -> Result<UnifiedExecProcess, UnifiedExecError> {
|
||||
let (program, args) = env
|
||||
.command
|
||||
.split_first()
|
||||
@@ -485,8 +485,8 @@ impl UnifiedExecSessionManager {
|
||||
&env.arg0,
|
||||
)
|
||||
.await
|
||||
.map_err(|err| UnifiedExecError::create_session(err.to_string()))?;
|
||||
UnifiedExecSession::from_spawned(spawned, env.sandbox).await
|
||||
.map_err(|err| UnifiedExecError::create_process(err.to_string()))?;
|
||||
UnifiedExecProcess::from_spawned(spawned, env.sandbox).await
|
||||
}
|
||||
|
||||
pub(super) async fn open_session_with_sandbox(
|
||||
@@ -496,7 +496,7 @@ impl UnifiedExecSessionManager {
|
||||
sandbox_permissions: SandboxPermissions,
|
||||
justification: Option<String>,
|
||||
context: &UnifiedExecContext,
|
||||
) -> Result<UnifiedExecSession, UnifiedExecError> {
|
||||
) -> Result<UnifiedExecProcess, UnifiedExecError> {
|
||||
let env = apply_unified_exec_env(create_env(&context.turn.shell_environment_policy));
|
||||
let features = context.session.features();
|
||||
let mut orchestrator = ToolOrchestrator::new();
|
||||
@@ -536,7 +536,7 @@ impl UnifiedExecSessionManager {
|
||||
context.turn.approval_policy,
|
||||
)
|
||||
.await
|
||||
.map_err(|e| UnifiedExecError::create_session(format!("{e:?}")))
|
||||
.map_err(|e| UnifiedExecError::create_process(format!("{e:?}")))
|
||||
}
|
||||
|
||||
pub(super) async fn collect_output_until_deadline(
|
||||
@@ -600,20 +600,20 @@ impl UnifiedExecSessionManager {
|
||||
collected
|
||||
}
|
||||
|
||||
fn prune_sessions_if_needed(store: &mut SessionStore) -> bool {
|
||||
if store.sessions.len() < MAX_UNIFIED_EXEC_SESSIONS {
|
||||
fn prune_processes_if_needed(store: &mut ProcessStore) -> bool {
|
||||
if store.processes.len() < MAX_UNIFIED_EXEC_PROCESSES {
|
||||
return false;
|
||||
}
|
||||
|
||||
let meta: Vec<(String, Instant, bool)> = store
|
||||
.sessions
|
||||
.processes
|
||||
.iter()
|
||||
.map(|(id, entry)| (id.clone(), entry.last_used, entry.session.has_exited()))
|
||||
.map(|(id, entry)| (id.clone(), entry.last_used, entry.process.has_exited()))
|
||||
.collect();
|
||||
|
||||
if let Some(session_id) = Self::session_id_to_prune_from_meta(&meta) {
|
||||
if let Some(entry) = store.remove(&session_id) {
|
||||
entry.session.terminate();
|
||||
if let Some(process_id) = Self::process_id_to_prune_from_meta(&meta) {
|
||||
if let Some(entry) = store.remove(&process_id) {
|
||||
entry.process.terminate();
|
||||
}
|
||||
return true;
|
||||
}
|
||||
@@ -622,7 +622,7 @@ impl UnifiedExecSessionManager {
|
||||
}
|
||||
|
||||
// Centralized pruning policy so we can easily swap strategies later.
|
||||
fn session_id_to_prune_from_meta(meta: &[(String, Instant, bool)]) -> Option<String> {
|
||||
fn process_id_to_prune_from_meta(meta: &[(String, Instant, bool)]) -> Option<String> {
|
||||
if meta.is_empty() {
|
||||
return None;
|
||||
}
|
||||
@@ -650,22 +650,25 @@ impl UnifiedExecSessionManager {
|
||||
.map(|(process_id, _, _)| process_id)
|
||||
}
|
||||
|
||||
pub(crate) async fn terminate_all_sessions(&self) {
|
||||
let entries: Vec<SessionEntry> = {
|
||||
let mut sessions = self.session_store.lock().await;
|
||||
let entries: Vec<SessionEntry> =
|
||||
sessions.sessions.drain().map(|(_, entry)| entry).collect();
|
||||
sessions.reserved_sessions_id.clear();
|
||||
pub(crate) async fn terminate_all_processes(&self) {
|
||||
let entries: Vec<ProcessEntry> = {
|
||||
let mut processes = self.process_store.lock().await;
|
||||
let entries: Vec<ProcessEntry> = processes
|
||||
.processes
|
||||
.drain()
|
||||
.map(|(_, entry)| entry)
|
||||
.collect();
|
||||
processes.reserved_process_ids.clear();
|
||||
entries
|
||||
};
|
||||
|
||||
for entry in entries {
|
||||
entry.session.terminate();
|
||||
entry.process.terminate();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
enum SessionStatus {
|
||||
enum ProcessStatus {
|
||||
Alive {
|
||||
exit_code: Option<i32>,
|
||||
call_id: String,
|
||||
@@ -673,7 +676,7 @@ enum SessionStatus {
|
||||
},
|
||||
Exited {
|
||||
exit_code: Option<i32>,
|
||||
entry: Box<SessionEntry>,
|
||||
entry: Box<ProcessEntry>,
|
||||
},
|
||||
Unknown,
|
||||
}
|
||||
@@ -716,7 +719,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pruning_prefers_exited_sessions_outside_recently_used() {
|
||||
fn pruning_prefers_exited_processes_outside_recently_used() {
|
||||
let now = Instant::now();
|
||||
let id = |n: i32| n.to_string();
|
||||
let meta = vec![
|
||||
@@ -732,7 +735,7 @@ mod tests {
|
||||
(id(10), now - Duration::from_secs(13), false),
|
||||
];
|
||||
|
||||
let candidate = UnifiedExecSessionManager::session_id_to_prune_from_meta(&meta);
|
||||
let candidate = UnifiedExecProcessManager::process_id_to_prune_from_meta(&meta);
|
||||
|
||||
assert_eq!(candidate, Some(id(2)));
|
||||
}
|
||||
@@ -754,13 +757,13 @@ mod tests {
|
||||
(id(10), now - Duration::from_secs(13), false),
|
||||
];
|
||||
|
||||
let candidate = UnifiedExecSessionManager::session_id_to_prune_from_meta(&meta);
|
||||
let candidate = UnifiedExecProcessManager::process_id_to_prune_from_meta(&meta);
|
||||
|
||||
assert_eq!(candidate, Some(id(1)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pruning_protects_recent_sessions_even_if_exited() {
|
||||
fn pruning_protects_recent_processes_even_if_exited() {
|
||||
let now = Instant::now();
|
||||
let id = |n: i32| n.to_string();
|
||||
let meta = vec![
|
||||
@@ -776,7 +779,7 @@ mod tests {
|
||||
(id(10), now - Duration::from_secs(13), true),
|
||||
];
|
||||
|
||||
let candidate = UnifiedExecSessionManager::session_id_to_prune_from_meta(&meta);
|
||||
let candidate = UnifiedExecProcessManager::process_id_to_prune_from_meta(&meta);
|
||||
|
||||
// (10) is exited but among the last 8; we should drop the LRU outside that set.
|
||||
assert_eq!(candidate, Some(id(1)));
|
||||
Reference in New Issue
Block a user