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

codex-rs/code-mode-host/src/host_tests.rs

@@ -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

@@ -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,
                 ),
             ),
         );

codex-rs/code-mode-host/tests/stdio.rs

@@ -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(

codex-rs/code-mode-host/tests/websocket.rs

@@ -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)?;