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
@@ -30,11 +30,13 @@ use super::WireImageDetail;
use super::WireNestedToolCall;
use super::WireResult;
use super::WireRuntimeResponse;
use super::WireSessionCellExecutionLimits;
use super::WireToolDefinition;
use super::WireToolKind;
use super::WireToolName;
use super::WireWaitOutcome;
use super::WireWaitRequest;
use crate::CodeModeSessionCellExecutionLimits;
use crate::ExecuteRequest;
fn session_id() -> SessionId {
@@ -363,6 +365,61 @@ fn handshake_v1_variants_are_pinned() {
}
}
#[test]
fn open_session_serializes_optional_cell_execution_limits() {
assert_wire_round_trip(
HostRequest::OpenSession {
session_id: session_id(),
cell_execution_limits: Some(WireSessionCellExecutionLimits {
max_yield_time_ms: Some(250),
max_heap_size_bytes: Some(16 * 1024 * 1024),
}),
},
json!({
"method": "session/open",
"sessionId": "session-1",
"cellExecutionLimits": {
"maxYieldTimeMs": 250,
"maxHeapSizeBytes": 16 * 1024 * 1024,
},
}),
);
}
#[test]
fn session_cell_execution_limits_convert_between_domain_and_wire() {
let domain_limits = CodeModeSessionCellExecutionLimits {
max_yield_time_ms: Some(250),
max_heap_size_bytes: Some(16_usize * 1024 * 1024),
};
let wire_limits = WireSessionCellExecutionLimits {
max_yield_time_ms: Some(250),
max_heap_size_bytes: Some(16_u64 * 1024 * 1024),
};
assert_eq!(
WireSessionCellExecutionLimits::try_from(domain_limits.clone())
.expect("domain limits convert to wire limits"),
wire_limits
);
assert_eq!(
CodeModeSessionCellExecutionLimits::try_from(wire_limits)
.expect("wire limits convert to domain limits"),
domain_limits
);
}
#[cfg(target_pointer_width = "32")]
#[test]
fn session_cell_execution_limits_reject_heap_sizes_that_exceed_usize() {
let wire_limits = WireSessionCellExecutionLimits {
max_yield_time_ms: None,
max_heap_size_bytes: Some(u64::from(u32::MAX) + 1),
};
assert!(CodeModeSessionCellExecutionLimits::try_from(wire_limits).is_err());
}
#[test]
fn client_to_host_v1_variants_are_pinned() {
let execute_request = execute_request();
@@ -371,6 +428,7 @@ fn client_to_host_v1_variants_are_pinned() {
request_id(/*value*/ 1),
HostRequest::OpenSession {
session_id: session_id(),
cell_execution_limits: None,
},
json!({ "method": "session/open", "sessionId": "session-1" }),
),
@@ -810,6 +868,17 @@ fn every_nested_v1_object_rejects_unknown_fields() {
}))
.is_err()
);
assert!(
serde_json::from_value::<HostRequest>(json!({
"method": "session/open",
"sessionId": "session-1",
"cellExecutionLimits": {
"maxYieldTimeMs": 250,
"unexpected": true,
},
}))
.is_err()
);
assert!(
serde_json::from_value::<WireExecuteRequest>(json!({
"tool_call_id": "call-1",