openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-18 17:18:52 +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:09:01 +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
codex-rs/code-mode-host/src/host_tests.rs
+68 -3
View File
@@ -8,6 +8,7 @@ use std::task::Context;
use std::task::Poll;
use std::time::Duration;
use codex_code_mode_protocol::CodeModeSessionCellExecutionLimits;
use codex_code_mode_protocol::host::Capability;
use codex_code_mode_protocol::host::CapabilitySet;
use codex_code_mode_protocol::host::ClientHello;
@@ -23,6 +24,7 @@ use codex_code_mode_protocol::host::HostResponse;
use codex_code_mode_protocol::host::HostToClient;
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::WireExecuteRequest;
@@ -129,6 +131,7 @@ async fn handshake_and_multiple_session_lifecycles_are_ordered() {
id: request_id,
request: HostRequest::OpenSession {
session_id: session_id(id),
cell_execution_limits: None,
},
})
.await
@@ -199,6 +202,7 @@ async fn disconnect_cancels_a_backpressured_host_writer() {
id: request_id(/*value*/ 1),
request: HostRequest::OpenSession {
session_id: session_id("backpressured-session"),
cell_execution_limits: None,
},
})
.await
@@ -357,6 +361,56 @@ async fn optional_dual_websocket_capability_falls_back_to_a_single_connection()
host.await.expect("host task").expect("host connection");
}
#[tokio::test]
async fn session_resource_limits_are_negotiated_when_optional_or_required() {
let dual_capability =
Capability::new(DUAL_WEBSOCKET_CAPABILITY).expect("dual websocket capability");
let resource_limits_capability =
Capability::new(SESSION_RESOURCE_LIMITS_CAPABILITY).expect("resource limits capability");
let resource_limits =
CapabilitySet::try_new([resource_limits_capability.clone()]).expect("host capabilities");
for (required, optional) in [
(
CapabilitySet::empty(),
CapabilitySet::try_new([dual_capability, resource_limits_capability])
.expect("optional capabilities"),
),
(resource_limits.clone(), CapabilitySet::empty()),
] {
let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 1024);
let (host_reader, host_writer) = tokio::io::split(host_stream);
let (client_reader, client_writer) = tokio::io::split(client_stream);
let host = tokio::spawn(run(host_reader, host_writer));
let mut reader = FramedReader::new(client_reader);
let mut writer = FramedWriter::new(client_writer);
writer
.write(&ClientToHost::ClientHello(
ClientHello::new(
SupportedProtocolVersions::try_new([ProtocolVersion::V1])
.expect("supported versions"),
required,
optional,
)
.expect("client hello"),
))
.await
.expect("write hello");
assert_eq!(
reader.read::<HostToClient>().await.expect("host hello"),
Some(HostToClient::HostHello(HostHello::new(
ProtocolVersion::V1,
resource_limits.clone(),
)))
);
drop(writer);
drop(reader);
host.await.expect("host task").expect("host connection");
}
}
#[tokio::test]
async fn incompatible_or_invalid_handshake_is_rejected() {
let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 1024);
@@ -393,6 +447,7 @@ async fn incompatible_or_invalid_handshake_is_rejected() {
id: request_id(/*value*/ 1),
request: HostRequest::OpenSession {
session_id: session_id("session-1"),
cell_execution_limits: None,
},
})
.await
@@ -458,6 +513,7 @@ async fn session_id_cannot_be_reused_after_shutdown() {
request_id(/*value*/ 1),
HostRequest::OpenSession {
session_id: id.clone(),
cell_execution_limits: None,
},
),
(
@@ -483,7 +539,10 @@ async fn session_id_cannot_be_reused_after_shutdown() {
writer
.write(&ClientToHost::Request {
id: request_id(/*value*/ 3),
request: HostRequest::OpenSession { session_id: id },
request: HostRequest::OpenSession {
session_id: id,
cell_execution_limits: None,
},
})
.await
.expect("reuse session ID");
@@ -568,7 +627,10 @@ async fn execute_request_id_remains_active_until_initial_response() {
});
let session_id = session_id("session-1");
state
.open_session(session_id.clone())
.open_session(
session_id.clone(),
CodeModeSessionCellExecutionLimits::default(),
)
.expect("open session");
let request_id = request_id(/*value*/ 1);
@@ -627,7 +689,10 @@ async fn active_cell_limit_rejects_execute_without_disconnecting() {
};
let session_id = session_id("session-1");
state
.open_session(session_id.clone())
.open_session(
session_id.clone(),
CodeModeSessionCellExecutionLimits::default(),
)
.expect("open session");
let request_id = request_id(/*value*/ 1);
codex-rs/code-mode-host/src/lib.rs
+33 -10
View File
@@ -10,6 +10,7 @@ use std::time::Duration;
use anyhow::Context;
use anyhow::Result;
use codex_code_mode_protocol::CodeModeSessionCellExecutionLimits;
use codex_code_mode_protocol::host::Capability;
use codex_code_mode_protocol::host::CapabilitySet;
use codex_code_mode_protocol::host::ClientToHost;
@@ -23,6 +24,7 @@ use codex_code_mode_protocol::host::HostToClient;
use codex_code_mode_protocol::host::MAX_PENDING_DELEGATE_CALLS;
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::TransportLane;
@@ -314,11 +316,21 @@ async fn negotiate(
} else {
None
};
let host_capabilities = if registration.is_some() {
CapabilitySet::try_new([dual_capability])?
} else {
CapabilitySet::empty()
};
let resource_limits_capability = Capability::new(SESSION_RESOURCE_LIMITS_CAPABILITY)?;
let resource_limits_requested = client_hello
.required_capabilities()
.contains(&resource_limits_capability)
|| client_hello
.optional_capabilities()
.contains(&resource_limits_capability);
let host_capabilities = CapabilitySet::try_new(
[
registration.is_some().then_some(dual_capability),
resource_limits_requested.then_some(resource_limits_capability),
]
.into_iter()
.flatten(),
)?;
if let Some(capability) = client_hello
.required_capabilities()
.iter()
@@ -419,10 +431,16 @@ impl HostState {
return;
}
match request {
HostRequest::OpenSession { session_id } => {
let result = self
.open_session(session_id.clone())
.map(|()| HostResponse::SessionReady { session_id });
HostRequest::OpenSession {
session_id,
cell_execution_limits,
} => {
let result = CodeModeSessionCellExecutionLimits::try_from(
cell_execution_limits.unwrap_or_default(),
)
.map_err(|error| format!("invalid code-mode session execution limits: {error}"))
.and_then(|limits| self.open_session(session_id.clone(), limits))
.map(|()| HostResponse::SessionReady { session_id });
self.respond(request_id, result);
}
HostRequest::Execute {
@@ -537,7 +555,11 @@ impl HostState {
}
}
fn open_session(&self, session_id: SessionId) -> Result<(), String> {
fn open_session(
&self,
session_id: SessionId,
cell_execution_limits: CodeModeSessionCellExecutionLimits,
) -> Result<(), String> {
let mut sessions = self.sessions.lock().unwrap_or_else(PoisonError::into_inner);
if sessions.contains_key(&session_id) {
return Err(format!(
@@ -571,6 +593,7 @@ impl HostState {
InProcessCodeModeSession::with_delegate_and_task_failure_handler(
delegate,
task_failure_handler,
cell_execution_limits,
),
),
);