openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-02 19:43:24 +00:00
Code Issues Packages Projects Releases Wiki Activity

Prewarm shell snapshots for eligible turns (#42137)

Browse Source 
## Why

Shell Snapshot V2 previously captured the login-shell environment lazily when a
command started, adding that work to the command path.

## What changed

- Start one asynchronous snapshot capture after turn hooks accept the turn, so
  eligible local `exec_command` calls can reuse it without blocking the model.
- Use the turn's resolved tools, permissions, sandbox, environment, and shell
  configuration, and skip prewarming for unsupported or network-managed cases.
- Cancel outstanding prewarming during shutdown. Keep failed speculative
  captures retryable by real commands without consuming their retry budget.

## Testing

Added coverage for non-blocking capture and reuse, failure recovery, sandbox and
permission isolation, shutdown cancellation, hook-stopped turns, eligibility
gates, and concurrent capture retries.

GitOrigin-RevId: ff24a4f0328c06653e09d03fd0a3db1eccd65030
This commit is contained in:
jif
2026-09-01 21:16:12 +00:00
committed by copyberry
parent c5b5fa8089
commit 612e6491d5
13 changed files with 783 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/exec-server/src/local_process.rs
+41 -1
View File
@@ -71,6 +71,8 @@ use crate::rpc::internal_error;
use crate::rpc::invalid_params;
use crate::rpc::invalid_request;
use crate::rpc_server_requests::RpcServerRequestSender;
#[cfg(unix)]
use crate::shell_snapshot::CapturePurpose;
use crate::telemetry::ExecServerTelemetry;
use crate::telemetry::ProcessMetricGuard;
@@ -355,7 +357,12 @@ impl LocalProcess {
#[cfg(unix)]
self.inner
.shell_snapshots
.prepare(&params, &mut prepared, &self.inner.telemetry)
.prepare(
&params,
&mut prepared,
&self.inner.telemetry,
CapturePurpose::Execution,
)
.await?;
if prepared.command.is_empty() {
return Err(invalid_params("argv must not be empty".to_string()));
@@ -753,6 +760,39 @@ impl ExecBackend for LocalProcess {
fn start(&self, params: ExecParams) -> ExecBackendFuture<'_> {
Box::pin(LocalProcess::start(self, params))
}
#[cfg(unix)]
fn prewarm_shell_snapshot(&self, params: ExecParams) -> ExecProcessFuture<'_, ()> {
Box::pin(async move {
if params.enforce_managed_network
|| params.managed_network.is_some()
|| params.network_proxy.is_some()
{
return Err(ExecServerError::Protocol(
"shell snapshot prewarming does not support managed networking".to_string(),
));
}
let mut prepared = prepare_exec_request(
&params,
child_env(&params),
self.runtime_paths.as_ref(),
/*network_policy_decider*/ None,
/*network_policy_audit_observer*/ None,
)
.await
.map_err(map_handler_error)?;
self.inner
.shell_snapshots
.prepare(
&params,
&mut prepared,
&self.inner.telemetry,
CapturePurpose::Prewarm,
)
.await
.map_err(map_handler_error)
})
}
}
impl LocalExecProcess {
codex-rs/exec-server/src/process.rs
+11
View File
@@ -222,6 +222,17 @@ pub type ExecProcessFuture<'a, T> =
pub trait ExecBackend: Send + Sync {
fn start(&self, params: ExecParams) -> ExecBackendFuture<'_>;
/// Captures a local shell snapshot without starting the requested command.
/// Failures must remain retryable by real commands. Remote backends do not
/// support this operation; callers should leave them on the lazy path.
fn prewarm_shell_snapshot(&self, _params: ExecParams) -> ExecProcessFuture<'_, ()> {
Box::pin(async {
Err(ExecServerError::Protocol(
"exec backend does not support shell snapshot prewarming".to_string(),
))
})
}
/// Starts a process with an authoritative controller-side policy decider.
fn start_with_network_policy_decider(
&self,
codex-rs/exec-server/src/shell_snapshot.rs
+41 -16
View File
@@ -56,12 +56,19 @@ struct ShellSnapshot {
environment: HashMap<String, String>,
}
#[derive(Clone, Copy, PartialEq, Eq)]
pub(crate) enum CapturePurpose {
Execution,
Prewarm,
}
impl ShellSnapshotCache {
pub(crate) async fn prepare(
&self,
params: &ExecParams,
prepared: &mut PreparedExecRequest,
telemetry: &ExecServerTelemetry,
purpose: CapturePurpose,
) -> Result<(), JSONRPCErrorError> {
let Some(request) = params.shell_snapshot.as_ref() else {
return Ok(());
@@ -103,7 +110,8 @@ impl ShellSnapshotCache {
let mut entry = entries.remove(position)?;
// Share each failed attempt during backoff. After the retry
// budget is exhausted, keep falling back until eviction.
if entry.attempts < MAX_SNAPSHOT_ATTEMPTS
if purpose == CapturePurpose::Execution
&& entry.attempts < MAX_SNAPSHOT_ATTEMPTS
&& let Some(Err(retry_at)) = entry.snapshot.get()
&& Instant::now() >= *retry_at
{
@@ -134,23 +142,40 @@ impl ShellSnapshotCache {
snapshot
}
};
let Ok(snapshot) = snapshot
.get_or_init(|| async {
let started_at = std::time::Instant::now();
let result = capture_snapshot(params, prepared, shell_type).await;
telemetry.shell_snapshot_captured(
started_at.elapsed(),
result.as_ref().map(|_| ()).map_err(|_| "capture_failed"),
);
result.map_err(|err| {
tracing::warn!("failed to capture shell snapshot: {err:?}");
Instant::now() + SNAPSHOT_RETRY_BACKOFF
})
})
.await
else {
let capture = async {
let started_at = std::time::Instant::now();
let result = capture_snapshot(params, prepared, shell_type).await;
telemetry.shell_snapshot_captured(
started_at.elapsed(),
result.as_ref().map(|_| ()).map_err(|_| "capture_failed"),
);
result
};
let snapshot = match purpose {
CapturePurpose::Execution => {
snapshot
.get_or_init(|| async {
capture.await.map_err(|err| {
tracing::warn!("failed to capture shell snapshot: {err:?}");
Instant::now() + SNAPSHOT_RETRY_BACKOFF
})
})
.await
}
CapturePurpose::Prewarm => {
// Leave the cell uninitialized on failure: a waiting real command
// can capture immediately, without spending its retry budget.
snapshot
.get_or_try_init(|| async { capture.await.map(Ok) })
.await?
}
};
let Ok(snapshot) = snapshot else {
return Ok(());
};
if purpose == CapturePurpose::Prewarm {
return Ok(());
}
let request_overrides = params
.env
codex-rs/exec-server/src/shell_snapshot_tests.rs
+37 -11
View File
@@ -10,6 +10,7 @@ use opentelemetry_sdk::metrics::data::MetricData;
use pretty_assertions::assert_eq;
use test_case::test_case;
use super::CapturePurpose;
use super::MAX_SNAPSHOT_ATTEMPTS;
use super::SNAPSHOT_RETRY_BACKOFF;
use super::ShellSnapshotCache;
@@ -22,14 +23,19 @@ use crate::protocol::ShellInfo;
use crate::protocol::ShellSnapshotRequest;
use crate::telemetry::ExecServerTelemetry;
#[test_case(1; "succeeds_on_first_attempt")]
#[test_case(2; "recovers_on_second_attempt")]
#[test_case(3; "recovers_on_last_attempt")]
#[test_case(4; "stops_after_three_failures")]
#[test_case(1, CapturePurpose::Execution; "succeeds_on_first_attempt")]
#[test_case(2, CapturePurpose::Execution; "recovers_on_second_attempt")]
#[test_case(3, CapturePurpose::Execution; "recovers_on_last_attempt")]
#[test_case(4, CapturePurpose::Execution; "stops_after_three_failures")]
#[test_case(2, CapturePurpose::Prewarm; "failed_prewarm_releases_waiting_command")]
#[test_case(3, CapturePurpose::Prewarm; "failed_prewarm_preserves_last_attempt")]
#[test_case(4, CapturePurpose::Prewarm; "failed_prewarm_preserves_retry_limit")]
#[tokio::test]
async fn snapshot_failure_retries_are_bounded_and_single_flight(
recovery_attempt: usize,
initial_purpose: CapturePurpose,
) -> anyhow::Result<()> {
let prewarm_fails_first = initial_purpose == CapturePurpose::Prewarm;
let home = tempfile::TempDir::new()?;
let profile = home.path().join(".bashrc");
std::fs::write(&profile, "printf x >> \"$HOME/captures\"\nexit 7\n")?;
@@ -101,12 +107,23 @@ async fn snapshot_failure_retries_are_bounded_and_single_flight(
)
.await
.expect("prepare concurrent capture");
let prewarming = prewarm_fails_first && attempt == 1;
let purpose = if prewarming {
CapturePurpose::Prewarm
} else {
CapturePurpose::Execution
};
let (first, second) = tokio::join!(
cache.prepare(&params, &mut prepared, &telemetry),
cache.prepare(&params, &mut concurrent, &telemetry),
biased;
cache.prepare(&params, &mut prepared, &telemetry, purpose),
cache.prepare(&params, &mut concurrent, &telemetry, CapturePurpose::Execution),
);
first.expect("capture failure must preserve command fallback");
second.expect("concurrent request must share the capture attempt");
if prewarming {
first.expect_err("prewarm must report failure without caching it");
} else {
first.expect("capture failure must preserve command fallback");
}
second.expect("waiting command must complete even when prewarm fails");
assert_eq!(
(&prepared.command, &prepared.env),
(&concurrent.command, &concurrent.env)
@@ -115,7 +132,12 @@ async fn snapshot_failure_retries_are_bounded_and_single_flight(
tokio::time::pause();
if attempt < recovery_attempt || recovery_attempt > MAX_SNAPSHOT_ATTEMPTS {
cache
.prepare(&params, &mut prepared, &telemetry)
.prepare(
&params,
&mut prepared,
&telemetry,
CapturePurpose::Execution,
)
.await
.expect("capture must stay cached during backoff");
assert_eq!(
@@ -127,7 +149,10 @@ async fn snapshot_failure_retries_are_bounded_and_single_flight(
}
assert_eq!(
std::fs::read_to_string(home.path().join("captures"))?,
"x".repeat(attempt.min(recovery_attempt).min(MAX_SNAPSHOT_ATTEMPTS))
"x".repeat(
attempt.min(recovery_attempt).min(MAX_SNAPSHOT_ATTEMPTS)
+ usize::from(prewarm_fails_first)
)
);
tokio::time::advance(SNAPSHOT_RETRY_BACKOFF).await;
tokio::time::resume();
@@ -170,7 +195,8 @@ async fn snapshot_failure_retries_are_bounded_and_single_flight(
}
let mut expected_counters = BTreeMap::new();
let mut expected_durations = BTreeMap::new();
let failures = (recovery_attempt - 1).min(MAX_SNAPSHOT_ATTEMPTS) as u64;
let failures =
(recovery_attempt - 1).min(MAX_SNAPSHOT_ATTEMPTS) as u64 + u64::from(prewarm_fails_first);
for (success, count) in [
("false", failures),
("true", u64::from(recovery_attempt <= MAX_SNAPSHOT_ATTEMPTS)),