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:
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
+69
View File
@@ -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(