mirror of
https://github.com/openai/codex.git
synced 2026-10-18 17:18:52 +00:00
Add per-session code-mode execution limits (#37114)
## What changed - Add `create_session_with_limits` and session-scoped cell execution limits. - Clamp execute and wait yield times to the session's `max_yield_time_ms` without terminating the running cell. - Negotiate support with remote code-mode hosts and include non-default limits in `session/open`, while keeping unlimited sessions compatible with hosts and providers that do not support limits. ## Testing - Cover yield-time clamping, zero-timeout behavior, and isolation between sessions. - Cover wire serialization, capability negotiation, unsupported hosts, and shared process-host execution. GitOrigin-RevId: 9517321cd605bb87f93eeaa6ba331cc2e346e582
This commit is contained in:
@@ -12,6 +12,7 @@ use std::os::unix::fs::PermissionsExt;
|
||||
use codex_code_mode::CellId;
|
||||
use codex_code_mode::CodeModeNestedToolCall;
|
||||
use codex_code_mode::CodeModeSession;
|
||||
use codex_code_mode::CodeModeSessionCellExecutionLimits;
|
||||
use codex_code_mode::CodeModeSessionDelegate;
|
||||
use codex_code_mode::CodeModeSessionProvider;
|
||||
use codex_code_mode::CodeModeToolKind;
|
||||
@@ -260,6 +261,74 @@ async fn next_callback_event(
|
||||
.expect("callback event stream closed")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn session_execution_limits_are_isolated_on_a_shared_process_host() {
|
||||
let provider: Arc<dyn CodeModeSessionProvider> =
|
||||
Arc::new(ProcessOwnedCodeModeSessionProvider::with_host_program(
|
||||
codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").expect("host binary"),
|
||||
));
|
||||
let limited = provider
|
||||
.create_session_with_limits(
|
||||
Arc::new(RecordingDelegate::default()),
|
||||
CodeModeSessionCellExecutionLimits {
|
||||
max_yield_time_ms: Some(1),
|
||||
max_heap_size_bytes: None,
|
||||
},
|
||||
)
|
||||
.await
|
||||
.expect("create limited session");
|
||||
let other = provider
|
||||
.create_session_with_limits(
|
||||
Arc::new(RecordingDelegate::default()),
|
||||
CodeModeSessionCellExecutionLimits {
|
||||
max_yield_time_ms: Some(1_000),
|
||||
max_heap_size_bytes: None,
|
||||
},
|
||||
)
|
||||
.await
|
||||
.expect("create independently limited session");
|
||||
|
||||
let response = tokio::time::timeout(
|
||||
Duration::from_secs(5),
|
||||
execute(&limited, execute_request("await new Promise(() => {});")),
|
||||
)
|
||||
.await
|
||||
.expect("session limit should bound the default execution wait");
|
||||
assert_eq!(
|
||||
response,
|
||||
RuntimeResponse::Yielded {
|
||||
cell_id: cell_id("1"),
|
||||
content_items: Vec::new(),
|
||||
}
|
||||
);
|
||||
assert_eq!(
|
||||
tokio::time::timeout(
|
||||
Duration::from_secs(5),
|
||||
limited.wait(WaitRequest {
|
||||
cell_id: cell_id("1"),
|
||||
yield_time_ms: 60_000,
|
||||
}),
|
||||
)
|
||||
.await
|
||||
.expect("session limit should bound explicit waits")
|
||||
.expect("wait for yielded cell"),
|
||||
WaitOutcome::LiveCell(RuntimeResponse::Yielded {
|
||||
cell_id: cell_id("1"),
|
||||
content_items: Vec::new(),
|
||||
})
|
||||
);
|
||||
|
||||
limited
|
||||
.terminate(cell_id("1"))
|
||||
.await
|
||||
.expect("terminate cell");
|
||||
limited.shutdown().await.expect("shutdown limited session");
|
||||
other
|
||||
.shutdown()
|
||||
.await
|
||||
.expect("shutdown independent session");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn remote_session_persists_values_forwards_delegates_and_controls_cells() {
|
||||
let provider = ProcessOwnedCodeModeSessionProvider::with_host_program(
|
||||
|
||||
@@ -6,6 +6,7 @@ use anyhow::Context;
|
||||
use anyhow::Result;
|
||||
use codex_code_mode::CellId;
|
||||
use codex_code_mode::CodeModeNestedToolCall;
|
||||
use codex_code_mode::CodeModeSessionCellExecutionLimits;
|
||||
use codex_code_mode::CodeModeSessionDelegate;
|
||||
use codex_code_mode::CodeModeSessionProvider;
|
||||
use codex_code_mode::CodeModeToolKind;
|
||||
@@ -32,6 +33,7 @@ use codex_code_mode_protocol::host::HostToClient;
|
||||
use codex_code_mode_protocol::host::MAX_FRAME_BYTES;
|
||||
use codex_code_mode_protocol::host::ProtocolVersion;
|
||||
use codex_code_mode_protocol::host::RequestId;
|
||||
use codex_code_mode_protocol::host::SESSION_RESOURCE_LIMITS_CAPABILITY;
|
||||
use codex_code_mode_protocol::host::SessionId;
|
||||
use codex_code_mode_protocol::host::SupportedProtocolVersions;
|
||||
use codex_code_mode_protocol::host::WireContentItem;
|
||||
@@ -225,16 +227,18 @@ impl HostClient {
|
||||
|
||||
async fn negotiate_dual(&mut self, websocket_url: &str) -> Result<HostClient> {
|
||||
let capability = Capability::new(DUAL_WEBSOCKET_CAPABILITY)?;
|
||||
let resource_limits = Capability::new(SESSION_RESOURCE_LIMITS_CAPABILITY)?;
|
||||
let hello = ClientHello::new(
|
||||
SupportedProtocolVersions::try_new([ProtocolVersion::V1])?,
|
||||
CapabilitySet::empty(),
|
||||
CapabilitySet::try_new([capability.clone()])?,
|
||||
CapabilitySet::try_new([capability.clone(), resource_limits.clone()])?,
|
||||
)?;
|
||||
self.send(&ClientToHost::ClientHello(hello)).await?;
|
||||
let HostToClient::HostHello(hello) = self.read().await? else {
|
||||
anyhow::bail!("expected code-mode host hello");
|
||||
};
|
||||
assert!(hello.capabilities().contains(&capability));
|
||||
assert!(hello.capabilities().contains(&resource_limits));
|
||||
let token = hello
|
||||
.bulk_connection_token()
|
||||
.context("dual websocket handshake omitted its pairing token")?;
|
||||
@@ -257,6 +261,7 @@ impl HostClient {
|
||||
id,
|
||||
request: HostRequest::OpenSession {
|
||||
session_id: session_id.clone(),
|
||||
cell_execution_limits: None,
|
||||
},
|
||||
})
|
||||
.await?;
|
||||
@@ -416,7 +421,13 @@ async fn production_websocket_client_runs_nested_tools_while_other_sessions_prog
|
||||
.await
|
||||
.map_err(anyhow::Error::msg)?;
|
||||
let fast_session = provider
|
||||
.create_session(Arc::new(NoopCodeModeSessionDelegate))
|
||||
.create_session_with_limits(
|
||||
Arc::new(NoopCodeModeSessionDelegate),
|
||||
CodeModeSessionCellExecutionLimits {
|
||||
max_yield_time_ms: Some(5_000),
|
||||
max_heap_size_bytes: None,
|
||||
},
|
||||
)
|
||||
.await
|
||||
.map_err(anyhow::Error::msg)?;
|
||||
|
||||
|
||||
Reference in New Issue
Block a user