openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-10 20:26:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add per-session code-mode execution limits (#37114)

Browse Source 
## 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:
Sean Huang
2026-08-05 16:02:54 +00:00
committed by copyberry
parent 2707dfc219
commit 9d00bb01c0
18 changed files with 730 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files

45
codex-rs/code-mode-runtime/src/service.rs
View File

@@ -4,6 +4,7 @@ use std::time::Duration;
use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::CodeModeNestedToolCall;
use codex_code_mode_protocol::CodeModeSession;
use codex_code_mode_protocol::CodeModeSessionCellExecutionLimits;
use codex_code_mode_protocol::CodeModeSessionDelegate;
use codex_code_mode_protocol::CodeModeSessionResultFuture;
use codex_code_mode_protocol::CodeModeToolKind;
@@ -29,17 +30,9 @@ use crate::session_runtime::SessionRuntime;
const YIELD_GRACE_PERIOD: Duration = Duration::from_secs(1);
const MIN_YIELD_TIME_FOR_GRACE: Duration = Duration::from_secs(10);
fn yield_timeout(yield_time_ms: u64) -> Duration {
let yield_time = Duration::from_millis(yield_time_ms);
if yield_time >= MIN_YIELD_TIME_FOR_GRACE {
yield_time.saturating_add(YIELD_GRACE_PERIOD)
} else {
yield_time
}
}
pub struct InProcessCodeModeSession {
runtime: SessionRuntime<ProtocolDelegate>,
cell_execution_limits: CodeModeSessionCellExecutionLimits,
}
impl InProcessCodeModeSession {
@@ -48,20 +41,36 @@ impl InProcessCodeModeSession {
}
pub fn with_delegate(delegate: Arc<dyn CodeModeSessionDelegate>) -> Self {
Self::with_delegate_and_limits(delegate, CodeModeSessionCellExecutionLimits::default())
}
pub fn with_delegate_and_limits(
delegate: Arc<dyn CodeModeSessionDelegate>,
cell_execution_limits: CodeModeSessionCellExecutionLimits,
) -> Self {
Self {
runtime: SessionRuntime::new(Arc::new(ProtocolDelegate { delegate })),
cell_execution_limits: CodeModeSessionCellExecutionLimits {
max_heap_size_bytes: None,
..cell_execution_limits
},
}
}
pub fn with_delegate_and_task_failure_handler(
delegate: Arc<dyn CodeModeSessionDelegate>,
task_failure_handler: Arc<dyn Fn(String) + Send + Sync>,
cell_execution_limits: CodeModeSessionCellExecutionLimits,
) -> Self {
Self {
runtime: SessionRuntime::new_with_task_failure_handler(
Arc::new(ProtocolDelegate { delegate }),
Some(task_failure_handler),
),
cell_execution_limits: CodeModeSessionCellExecutionLimits {
max_heap_size_bytes: None,
..cell_execution_limits
},
}
}
@@ -71,7 +80,7 @@ impl InProcessCodeModeSession {
.runtime
.execute(
runtime_request(request),
runtime::ObserveMode::YieldAfter(yield_timeout(yield_time_ms)),
runtime::ObserveMode::YieldAfter(self.resolve_yield_timeout(yield_time_ms)),
)
.await
.map_err(|error| error.to_string())?;
@@ -126,7 +135,7 @@ impl InProcessCodeModeSession {
.runtime
.begin_observe(
&runtime_cell_id,
runtime::ObserveMode::YieldAfter(yield_timeout(yield_time_ms)),
runtime::ObserveMode::YieldAfter(self.resolve_yield_timeout(yield_time_ms)),
)
.await
{
@@ -185,6 +194,20 @@ impl InProcessCodeModeSession {
.await
.map_err(|error| error.to_string())
}
fn resolve_yield_timeout(&self, yield_time_ms: u64) -> Duration {
let yield_time = Duration::from_millis(yield_time_ms);
let timeout = if yield_time >= MIN_YIELD_TIME_FOR_GRACE {
yield_time.saturating_add(YIELD_GRACE_PERIOD)
} else {
yield_time
};
self.cell_execution_limits
.max_yield_time_ms
.map(Duration::from_millis)
.map_or(timeout, |limit| timeout.min(limit))
}
}
impl Default for InProcessCodeModeSession {

184
codex-rs/code-mode-runtime/src/service_tests.rs
View File

@@ -12,12 +12,12 @@ use super::WaitOutcome;
use super::WaitRequest;
use super::WaitToPendingOutcome;
use super::WaitToPendingRequest;
use super::yield_timeout;
use crate::CodeModeToolKind;
use crate::ExecuteRequest;
use crate::ExecuteToPendingOutcome;
use crate::FunctionCallOutputContentItem;
use crate::ToolDefinition;
use codex_code_mode_protocol::CodeModeSessionCellExecutionLimits;
use codex_code_mode_protocol::NotificationFuture;
use codex_code_mode_protocol::ToolInvocationFuture;
use codex_protocol::ToolName;
@@ -27,15 +27,36 @@ use tokio::sync::Notify;
use tokio_util::sync::CancellationToken;
#[test]
fn yield_timeout_adds_grace_only_at_ten_seconds() {
assert_eq!(
yield_timeout(/*yield_time_ms*/ 9_999),
Duration::from_millis(9_999)
);
assert_eq!(
yield_timeout(/*yield_time_ms*/ 10_000),
Duration::from_secs(11)
);
fn resolve_yield_timeout_applies_grace_before_session_limits() {
for (max_yield_time_ms, requested_yield_time_ms, expected_timeout) in [
(None, 0, Duration::ZERO),
(None, 9_999, Duration::from_millis(9_999)),
(None, 10_000, Duration::from_secs(11)),
(None, 10_001, Duration::from_millis(11_001)),
(Some(0), 0, Duration::ZERO),
(Some(0), 10_000, Duration::ZERO),
(Some(5_000), 9_999, Duration::from_secs(5)),
(Some(10_000), 10_000, Duration::from_secs(10)),
(Some(10_500), 10_000, Duration::from_millis(10_500)),
(Some(11_000), 10_000, Duration::from_secs(11)),
(Some(12_000), 10_000, Duration::from_secs(11)),
(Some(10_500), 5_000, Duration::from_secs(5)),
(Some(u64::MAX), u64::MAX, Duration::from_millis(u64::MAX)),
] {
let session = InProcessCodeModeSession::with_delegate_and_limits(
Arc::new(ReleasableToolDelegate::default()),
CodeModeSessionCellExecutionLimits {
max_yield_time_ms,
max_heap_size_bytes: None,
},
);
assert_eq!(
session.resolve_yield_timeout(requested_yield_time_ms),
expected_timeout,
"requested {requested_yield_time_ms} ms with limit {max_yield_time_ms:?}"
);
}
}
#[tokio::test(start_paused = true)]
@@ -68,6 +89,79 @@ async fn execute_waits_for_nested_tool_during_yield_grace() {
);
}
#[tokio::test(start_paused = true)]
async fn execute_and_wait_clamp_yield_grace_without_stopping_the_cell() {
let delegate = Arc::new(ReleasableToolDelegate::default());
let service = InProcessCodeModeSession::with_delegate_and_limits(
delegate.clone(),
CodeModeSessionCellExecutionLimits {
max_yield_time_ms: Some(/*value*/ 10_000),
max_heap_size_bytes: None,
},
);
let started = service
.execute(ExecuteRequest {
enabled_tools: vec![echo_tool()],
source: r#"await tools.echo({}); text("done");"#.to_string(),
yield_time_ms: None,
..execute_request("")
})
.await
.unwrap();
let initial_response = tokio::spawn(started.initial_response());
wait_until_tool_started(&delegate).await;
tokio::time::advance(Duration::from_millis(9_999)).await;
assert!(!initial_response.is_finished());
tokio::time::advance(Duration::from_millis(/*millis*/ 1)).await;
wait_until_finished(&initial_response).await;
assert_eq!(
initial_response.await.unwrap().unwrap(),
RuntimeResponse::Yielded {
cell_id: cell_id("1"),
content_items: Vec::new(),
}
);
let wait_response = service
.begin_wait(WaitRequest {
cell_id: cell_id("1"),
yield_time_ms: 10_000,
})
.await;
let wait_response = tokio::spawn(wait_response);
tokio::task::yield_now().await;
tokio::time::advance(Duration::from_secs(/*secs*/ 10)).await;
wait_until_finished(&wait_response).await;
assert_eq!(
wait_response.await.unwrap().unwrap(),
WaitOutcome::LiveCell(RuntimeResponse::Yielded {
cell_id: cell_id("1"),
content_items: Vec::new(),
})
);
delegate.release_tool();
let completion = service
.begin_wait(WaitRequest {
cell_id: cell_id("1"),
yield_time_ms: 10_000,
})
.await;
let completion = tokio::spawn(completion);
wait_until_finished(&completion).await;
assert_eq!(
completion.await.unwrap().unwrap(),
WaitOutcome::LiveCell(RuntimeResponse::Result {
cell_id: cell_id("1"),
content_items: vec![FunctionCallOutputContentItem::InputText {
text: "done".to_string(),
}],
error_text: None,
})
);
}
#[tokio::test(start_paused = true)]
async fn wait_waits_for_nested_tool_during_yield_grace() {
let delegate = Arc::new(ReleasableToolDelegate::default());
@@ -113,6 +207,76 @@ async fn wait_waits_for_nested_tool_during_yield_grace() {
);
}
#[tokio::test(start_paused = true)]
async fn zero_yield_limit_is_immediate_and_scoped_to_its_session() {
let zero_delegate = Arc::new(ReleasableToolDelegate::default());
let zero_session = InProcessCodeModeSession::with_delegate_and_limits(
zero_delegate.clone(),
CodeModeSessionCellExecutionLimits {
max_yield_time_ms: Some(/*value*/ 0),
max_heap_size_bytes: None,
},
);
let limited_delegate = Arc::new(ReleasableToolDelegate::default());
let limited_session = InProcessCodeModeSession::with_delegate_and_limits(
limited_delegate.clone(),
CodeModeSessionCellExecutionLimits {
max_yield_time_ms: Some(/*value*/ 10),
max_heap_size_bytes: None,
},
);
let request = ExecuteRequest {
enabled_tools: vec![echo_tool()],
source: "await tools.echo({});".to_string(),
yield_time_ms: Some(/*value*/ 60_000),
..execute_request("")
};
let zero_started = zero_session.execute(request.clone()).await.unwrap();
let limited_started = limited_session.execute(request).await.unwrap();
let zero_response = tokio::spawn(zero_started.initial_response());
let limited_response = tokio::spawn(limited_started.initial_response());
wait_until_tool_started(&zero_delegate).await;
wait_until_tool_started(&limited_delegate).await;
wait_until_finished(&zero_response).await;
assert!(!limited_response.is_finished());
assert_eq!(
zero_response.await.unwrap().unwrap(),
RuntimeResponse::Yielded {
cell_id: cell_id("1"),
content_items: Vec::new(),
}
);
let zero_wait = zero_session
.begin_wait(WaitRequest {
cell_id: cell_id("1"),
yield_time_ms: 60_000,
})
.await;
let zero_wait = tokio::spawn(zero_wait);
wait_until_finished(&zero_wait).await;
assert_eq!(
zero_wait.await.unwrap().unwrap(),
WaitOutcome::LiveCell(RuntimeResponse::Yielded {
cell_id: cell_id("1"),
content_items: Vec::new(),
})
);
tokio::time::advance(Duration::from_millis(/*millis*/ 10)).await;
wait_until_finished(&limited_response).await;
assert_eq!(
limited_response.await.unwrap().unwrap(),
RuntimeResponse::Yielded {
cell_id: cell_id("1"),
content_items: Vec::new(),
}
);
zero_session.shutdown().await.unwrap();
limited_session.shutdown().await.unwrap();
}
async fn wait_until_finished<T>(task: &tokio::task::JoinHandle<T>) {
for _ in 0..10_000 {
if task.is_finished() {