mirror of
https://github.com/openai/codex.git
synced 2026-11-09 08:40:33 +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:
1 parent
2707dfc219
commit
9d00bb01c0
18 files changed
+730
-57
No files matched your search
@@ -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);
|
||||
|
||||
|
||||
Reference in new issue
Block a user