Run environment hooks on their owning executors (#56343)

## What changed

- Include ready environments' configured hooks in lifecycle events and execute command hooks using each owner's shell, working directory, shell environment policy, and managed network policy.
- Resolve tool environments before `PreToolUse` and use those targets for tool hooks, execution, and permission review. Route MCP hooks by server ownership and file-copy hooks to both endpoints.
- Preserve environment, MCP server, and terminal session selectors across hook input rewrites, rejecting attempts to switch targets.

## Testing

Add regression coverage for permission-request selector preservation, MCP resource hook ownership, and file-copy hooks that rewrite paths while preserving source and target environments.

GitOrigin-RevId: 31a3a6830523965285dab451cfc866c620ef6c83
This commit is contained in:
andrewgu-oai
2026-10-23 04:01:39 +00:00
committed by copyberry
parent e5ea2ed50b
commit 4610b573e5
47 changed files with 1704 additions and 433 deletions
+11 -1
View File
@@ -52,6 +52,7 @@ mod head_tail_buffer;
mod oneshot;
mod process;
mod process_manager;
pub(crate) use process_manager::exec_env_policy_from_shell_policy;
mod process_state;
mod shell_snapshot;
mod stdin_approval;
@@ -125,9 +126,18 @@ pub(crate) struct ExecCommandRequest {
pub prefix_rule: Option<Vec<String>>,
}
/// The process and environment behind a terminal session ID at the time of lookup.
#[derive(Clone, Debug)]
pub(crate) struct TerminalSession {
pub(crate) process_id: i32,
process: Arc<UnifiedExecProcess>,
call_id: String,
pub(crate) environment_id: String,
}
#[derive(Debug)]
pub(crate) struct WriteStdinRequest<'a> {
pub process_id: i32,
pub terminal: &'a TerminalSession,
pub input: &'a str,
pub yield_time_ms: u64,
pub max_output_tokens: Option<usize>,
+75 -26
View File
@@ -315,9 +315,9 @@ async fn write_stdin(
input: &str,
yield_time_ms: u64,
) -> Result<ExecCommandToolOutput, UnifiedExecError> {
session
.services
.unified_exec_manager
let manager = &session.services.unified_exec_manager;
let terminal = manager.resolve_terminal(process_id).await?;
manager
.write_stdin(
&UnifiedExecContext::new(
Arc::clone(session),
@@ -326,7 +326,7 @@ async fn write_stdin(
"write".to_string(),
),
WriteStdinRequest {
process_id,
terminal: &terminal,
input,
yield_time_ms,
max_output_tokens: None,
@@ -671,8 +671,12 @@ async fn terminating_initial_exec_command_rechecks_initial_response_state() -> a
Ok(())
}
#[test_case::test_case(false; "terminal removed")]
#[test_case::test_case(true; "session id reused")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn terminating_during_stdin_poll_returns_exited_response() -> anyhow::Result<()> {
async fn terminating_during_stdin_poll_returns_exited_response(
reuse_id: bool,
) -> anyhow::Result<()> {
let (session, turn) = test_session_and_turn().await;
let manager = &session.services.unified_exec_manager;
let process_id = manager.allocate_process_id().await;
@@ -743,7 +747,20 @@ async fn terminating_during_stdin_poll_returns_exited_response() -> anyhow::Resu
.await
.expect("poll should clone process handles");
manager.release_process_id(process_id).await;
if reuse_id {
let replacement = super::process_tests::remote_process(
WriteStatus::Accepted,
/*terminate_error*/ None,
SandboxType::None,
)
.await;
let mut store = manager.process_store.lock().await;
let entry = store.processes.get_mut(&process_id).unwrap();
entry.process = Arc::new(replacement);
entry.call_id = "replacement".into();
} else {
manager.release_process_id(process_id).await;
}
allow_terminate.notify_one();
process.terminate_confirmed().await?;
@@ -751,8 +768,19 @@ async fn terminating_during_stdin_poll_returns_exited_response() -> anyhow::Resu
.await
.expect("poll should finish")
.expect("poll task should not panic")?;
assert_eq!(output.process_id, None);
assert!(manager.process_store.lock().await.processes.is_empty());
assert_eq!(
(output.process_id, output.event_call_id.as_str()),
(None, "call")
);
assert_eq!(
manager
.process_store
.lock()
.await
.processes
.contains_key(&process_id),
reuse_id
);
Ok(())
}
@@ -871,7 +899,7 @@ async fn stdin_approval_preserves_the_reviewed_terminal() -> anyhow::Result<()>
.await?;
let process_id = opened.process_id.expect("running terminal");
let cwd = PathUri::parse("file:///C:/workspace")?;
let original = {
{
let mut store = manager.process_store.lock().await;
let entry = store.processes.get_mut(&process_id).unwrap();
entry.permissions = TerminalPermissions::for_launch(
@@ -886,8 +914,9 @@ async fn stdin_approval_preserves_the_reviewed_terminal() -> anyhow::Result<()>
);
entry.environment_id = "unselected-executor".to_string();
entry.cwd = cwd.clone();
Arc::clone(&entry.process)
};
}
let original = manager.resolve_terminal(process_id).await?;
assert_eq!(original.environment_id, "unselected-executor");
// Empty polling must complete without an approval response.
// The test deadline must allow the minimum empty-poll wait.
tokio::time::timeout(
@@ -943,7 +972,13 @@ async fn stdin_approval_preserves_the_reviewed_terminal() -> anyhow::Result<()>
),
("call", Some("write"), cwd.clone().into())
);
assert!(original.interaction_lock().try_lock_owned().is_err());
assert!(
original
.process
.interaction_lock()
.try_lock_owned()
.is_err()
);
if input == "replace\n" {
let replacement = super::process_tests::remote_process(
WriteStatus::Accepted,
@@ -979,9 +1014,32 @@ async fn stdin_approval_preserves_the_reviewed_terminal() -> anyhow::Result<()>
matches!(result, Err(UnifiedExecError::UnknownProcessId { process_id: id }) if id == process_id)
),
}
assert!(original.interaction_lock().try_lock_owned().is_ok());
assert!(original.process.interaction_lock().try_lock_owned().is_ok());
}
original.terminate();
// Reusing the numeric ID must also reject a subsequent write to the original terminal.
let context = UnifiedExecContext::new(
Arc::clone(&session),
crate::session::step_context::StepContext::for_test(turn),
tokio_util::sync::CancellationToken::new(),
"stale".to_string(),
);
let stale = manager
.write_stdin(
&context,
WriteStdinRequest {
terminal: &original,
input: "",
yield_time_ms: 250,
max_output_tokens: None,
truncation_policy: TruncationPolicy::Tokens(10_000),
interaction_event: None,
},
)
.await;
assert!(
matches!(stale, Err(UnifiedExecError::UnknownProcessId { process_id: id }) if id == process_id)
);
original.process.terminate();
assert!(session.terminate_background_terminal(process_id).await);
Ok(())
}
@@ -1010,27 +1068,18 @@ async fn stdin_approval_observes_strict_review_enabled_while_queued() -> anyhow:
)
.await?;
let process_id = opened.process_id.expect("running terminal");
let process = Arc::clone(
&manager
.process_store
.lock()
.await
.processes
.get(&process_id)
.unwrap()
.process,
);
let terminal = manager.resolve_terminal(process_id).await?;
let context = UnifiedExecContext::new(
Arc::clone(&session),
StepContext::for_test(turn),
tokio_util::sync::CancellationToken::new(),
"write".to_string(),
);
let interaction = process.interaction_lock().lock_owned().await;
let interaction = terminal.process.interaction_lock().lock_owned().await;
let mut queued = Box::pin(manager.write_stdin(
&context,
WriteStdinRequest {
process_id,
terminal: &terminal,
// NUL input must be rejected when strict review becomes required.
input: "\0",
yield_time_ms: 250,
@@ -52,6 +52,7 @@ use crate::unified_exec::MIN_EMPTY_YIELD_TIME_MS;
use crate::unified_exec::MIN_YIELD_TIME_MS;
use crate::unified_exec::ProcessEntry;
use crate::unified_exec::ProcessStore;
use crate::unified_exec::TerminalSession;
use crate::unified_exec::UnifiedExecContext;
use crate::unified_exec::UnifiedExecError;
use crate::unified_exec::UnifiedExecProcessManager;
@@ -134,7 +135,7 @@ pub(super) fn apply_unified_exec_env(mut env: HashMap<String, String>) -> HashMa
env
}
pub(super) fn exec_env_policy_from_shell_policy(
pub(crate) fn exec_env_policy_from_shell_policy(
policy: &ShellEnvironmentPolicy,
) -> codex_exec_server::ExecEnvPolicy {
let mut exclude = policy
@@ -477,13 +478,26 @@ impl UnifiedExecProcessManager {
}
pub(crate) async fn release_process_id(&self, process_id: i32) {
self.release_process(process_id, /*expected_process*/ None)
.await;
}
async fn release_process(
&self,
process_id: i32,
expected_process: Option<&Arc<UnifiedExecProcess>>,
) {
let removed = {
let mut store = self.process_store.lock().await;
let entry = store.processes.remove(&process_id);
if !should_use_deterministic_process_ids() {
store.reserved_process_ids.remove(&process_id);
if let Some(expected) = expected_process
&& !store
.processes
.get(&process_id)
.is_some_and(|entry| Arc::ptr_eq(&entry.process, expected))
{
return;
}
entry
store.remove(process_id)
};
if let Some(entry) = removed {
unregister_network_approval_for_entry(&entry).await;
@@ -732,7 +746,10 @@ impl UnifiedExecProcessManager {
} else {
OutputRetention::UntilCollected
};
match self.refresh_process_state(process_id, retention).await {
match self
.refresh_process_state(process_id, &process, retention)
.await
{
ProcessStatus::Alive {
exit_code,
process_id,
@@ -871,7 +888,7 @@ impl UnifiedExecProcessManager {
turn_id = trace_id(&context.step_context.turn.sub_id),
call_id = trace_id(&context.call_id),
original_exec_call_id = tracing::field::Empty,
unified_exec_process_id = request.process_id,
unified_exec_process_id = request.terminal.process_id,
interaction = if request.input.is_empty() { "poll" } else { "write" },
outcome = tracing::field::Empty,
)
@@ -897,23 +914,12 @@ impl UnifiedExecProcessManager {
context: &UnifiedExecContext,
request: WriteStdinRequest<'_>,
) -> Result<ExecCommandToolOutput, UnifiedExecError> {
let process_id = request.process_id;
// Different terminal sessions can be polled concurrently, but reads and
// writes against one terminal must not overlap because they share a
// draining output buffer and process lifecycle.
let locked_process = {
let store = self.process_store.lock().await;
let entry = store
.processes
.get(&process_id)
.ok_or(UnifiedExecError::UnknownProcessId { process_id })?;
// Capture the original call even if this interaction is cancelled while queued.
if let Some(call_id) = trace_id(&entry.call_id) {
tracing::Span::current().record("original_exec_call_id", call_id);
}
Arc::clone(&entry.process)
};
let process_id = request.terminal.process_id;
let locked_process = &request.terminal.process;
if let Some(call_id) = trace_id(&request.terminal.call_id) {
tracing::Span::current().record("original_exec_call_id", call_id);
}
// Serialize writes and polls on the original terminal.
let _interaction_guard = locked_process.interaction_lock().lock_owned().await;
// A queued write must observe strict review enabled while it was waiting.
let strict_auto_review = context.step_context.turn.strict_auto_review_enabled();
@@ -922,7 +928,7 @@ impl UnifiedExecProcessManager {
let entry = store
.processes
.get(&process_id)
.filter(|entry| Arc::ptr_eq(&entry.process, &locked_process))
.filter(|entry| Arc::ptr_eq(&entry.process, locked_process))
.ok_or(UnifiedExecError::UnknownProcessId { process_id })?;
entry.stdin_approval(context, request.input, strict_auto_review)?
};
@@ -994,7 +1000,7 @@ impl UnifiedExecProcessManager {
tty,
..
} = self
.prepare_process_handles(process_id, &locked_process)
.prepare_process_handles(process_id, locked_process)
.await?;
context
.session
@@ -1026,7 +1032,7 @@ impl UnifiedExecProcessManager {
if !process.has_exited() {
if matches!(err, UnifiedExecError::ProcessFailed { .. }) {
process.terminate();
self.release_process_id(process_id).await;
self.release_process(process_id, Some(locked_process)).await;
}
return Err(err);
}
@@ -1065,7 +1071,7 @@ impl UnifiedExecProcessManager {
let message =
network_denial_message_for_session(session.as_ref(), network_approval.clone())
.await;
self.release_process_id(process_id).await;
self.release_process(process_id, Some(locked_process)).await;
return Err(fail_process_with_message(process.as_ref(), message));
}
if let Some(message) = process.failure_message() {
@@ -1074,7 +1080,7 @@ impl UnifiedExecProcessManager {
network_approval.clone(),
)
.await;
self.release_process_id(process_id).await;
self.release_process(process_id, Some(locked_process)).await;
if let Err(message) = finish_result {
return Err(fail_process_with_message(process.as_ref(), message));
}
@@ -1083,7 +1089,7 @@ impl UnifiedExecProcessManager {
// Only finalize remote sessions after both process exit and output collection.
let status = self
.refresh_process_state(process_id, OutputRetention::UntilCollected)
.refresh_process_state(process_id, locked_process, OutputRetention::UntilCollected)
.await;
let (process_id, exit_code, event_call_id) = match status {
ProcessStatus::Alive {
@@ -1105,9 +1111,7 @@ impl UnifiedExecProcessManager {
if process.has_exited() {
(None, process.exit_code(), call_id)
} else {
return Err(UnifiedExecError::UnknownProcessId {
process_id: request.process_id,
});
return Err(UnifiedExecError::UnknownProcessId { process_id });
}
}
};
@@ -1134,7 +1138,7 @@ impl UnifiedExecProcessManager {
call_id: response.event_call_id.clone(),
process_id: response
.process_id
.unwrap_or(request.process_id)
.unwrap_or(request.terminal.process_id)
.to_string(),
stdin: request.input.to_string(),
};
@@ -1149,10 +1153,15 @@ impl UnifiedExecProcessManager {
async fn refresh_process_state(
&self,
process_id: i32,
expected_process: &Arc<UnifiedExecProcess>,
retention: OutputRetention,
) -> ProcessStatus {
let mut store = self.process_store.lock().await;
let Some(entry) = store.processes.get_mut(&process_id) else {
let Some(entry) = store
.processes
.get_mut(&process_id)
.filter(|entry| Arc::ptr_eq(&entry.process, expected_process))
else {
return ProcessStatus::Unknown;
};
@@ -1847,14 +1856,21 @@ impl UnifiedExecProcessManager {
}
}
/// Resolves stdin's target from the host-owned terminal, not model-supplied metadata.
pub(crate) async fn environment_id_for_process(&self, process_id: i32) -> Option<String> {
self.process_store
.lock()
.await
pub(crate) async fn resolve_terminal(
&self,
process_id: i32,
) -> Result<TerminalSession, UnifiedExecError> {
let store = self.process_store.lock().await;
let entry = store
.processes
.get(&process_id)
.map(|entry| entry.environment_id.clone())
.ok_or(UnifiedExecError::UnknownProcessId { process_id })?;
Ok(TerminalSession {
process_id,
process: Arc::clone(&entry.process),
call_id: entry.call_id.clone(),
environment_id: entry.environment_id.clone(),
})
}
pub(crate) async fn list_processes(&self) -> Vec<BackgroundTerminalInfo> {