openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-27 18:48:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add code-mode output truncation protocol contracts (#57272)

Browse Source 
## What changed

Add session options for client- or server-owned output truncation and optional `max_output_tokens` budgets for wait and terminate observations. Introduce `TerminateRequest` and update session implementations and callers to use it.

Extend the stdio and gRPC schemas with optional truncation settings and budgets, and define the `bounded-output-v1` capability. Existing sessions retain client-side truncation; providers reject server mode by default. This change defines the contracts without enabling server-side truncation.

## Testing

Extend protocol serialization coverage for `outputTruncation: "server"` while preserving existing wire shapes when optional fields are absent.

GitOrigin-RevId: 22145897ea16d385c5ba9e7549109f7a669fceb4
This commit is contained in:
Sean Huang
2026-10-26 21:36:35 +00:00
committed by copyberry
parent 7077b1edb9
commit e963cab11a
38 changed files with 329 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/code-mode-host/src/grpc/mod.rs
+1
View File
@@ -234,6 +234,7 @@ impl GrpcCodeModeHost {
let request = WaitRequest {
cell_id: CellId::new(request.cell_id),
yield_time_ms: request.yield_time_ms,
max_output_tokens: None,
};
let outcome = tokio::select! {
biased;
codex-rs/code-mode-host/src/grpc/robustness_tests.rs
+1
View File
@@ -255,6 +255,7 @@ async fn session_closure_cancels_pending_termination() {
let termination = host.terminate(Request::new(proto::TerminateRequest {
session_id,
cell_id,
max_output_tokens: None,
}));
tokio::pin!(termination);
assert!(termination.as_mut().now_or_never().is_none());
codex-rs/code-mode-host/src/grpc/service_tests.rs
+5
View File
@@ -14,6 +14,7 @@ pub(super) async fn open_session(
let mut stream = host
.open_session(Request::new(proto::OpenSessionRequest {
cell_execution_limits: None,
output_truncation: None,
}))
.await
.expect("open code-mode session")
@@ -120,6 +121,7 @@ async fn execute_stream_starts_immediately_and_wait_preserves_missing_cells() {
cell_id: cell_id.clone(),
wait_id: "wait-1".to_string(),
yield_time_ms: 60_000,
max_output_tokens: None,
}))
.await
.expect("wait for completion")
@@ -163,6 +165,7 @@ async fn execute_stream_starts_immediately_and_wait_preserves_missing_cells() {
cell_id: "missing-cell".to_string(),
wait_id: "wait-missing".to_string(),
yield_time_ms: 1,
max_output_tokens: None,
}))
.await
.expect("missing cells remain successful wait outcomes")
@@ -355,6 +358,7 @@ async fn cancellation_before_wait_admission_is_preserved() {
cell_id: cell_id.clone(),
wait_id: "pre-cancelled".to_string(),
yield_time_ms: 60_000,
max_output_tokens: None,
}))
.await
.unwrap_err();
@@ -363,6 +367,7 @@ async fn cancellation_before_wait_admission_is_preserved() {
host.terminate(Request::new(proto::TerminateRequest {
session_id,
cell_id,
max_output_tokens: None,
}))
.await
.unwrap();
codex-rs/code-mode-host/src/grpc/session.rs
+2 -1
View File
@@ -124,6 +124,7 @@ impl GrpcHostState {
.try_send(Ok(proto::SessionEvent {
event: Some(proto::session_event::Event::Opened(proto::SessionOpened {
session_id: id.to_string(),
output_truncation: None,
})),
}))
.map_err(|_| Status::internal("failed to publish the opened code-mode session"))?;
@@ -263,7 +264,7 @@ impl GrpcSession {
_ = self.closed.cancelled() => {
Err(Status::cancelled("code-mode session is closed"))
}
result = self.runtime.terminate(cell_id) => {
result = self.runtime.terminate(cell_id.into()) => {
result.map_err(Status::failed_precondition)
}
}
codex-rs/code-mode-host/src/host_tests.rs
+6 -1
View File
@@ -127,6 +127,7 @@ async fn handshake_and_multiple_session_lifecycles_are_ordered() {
request: HostRequest::OpenSession {
session_id: session_id(id),
cell_execution_limits: None,
output_truncation: None,
},
})
.await
@@ -198,6 +199,7 @@ async fn disconnect_cancels_a_backpressured_host_writer() {
request: HostRequest::OpenSession {
session_id: session_id("backpressured-session"),
cell_execution_limits: None,
output_truncation: None,
},
})
.await
@@ -339,6 +341,7 @@ async fn incompatible_or_invalid_handshake_is_rejected() {
request: HostRequest::OpenSession {
session_id: session_id("session-1"),
cell_execution_limits: None,
output_truncation: None,
},
})
.await
@@ -405,6 +408,7 @@ async fn session_id_cannot_be_reused_after_shutdown() {
HostRequest::OpenSession {
session_id: id.clone(),
cell_execution_limits: None,
output_truncation: None,
},
),
(
@@ -433,6 +437,7 @@ async fn session_id_cannot_be_reused_after_shutdown() {
request: HostRequest::OpenSession {
session_id: id,
cell_execution_limits: None,
output_truncation: None,
},
})
.await
@@ -556,7 +561,7 @@ async fn execute_request_id_remains_active_until_initial_response() {
state
.session(&session_id)
.expect("session")
.terminate(cell_id.into())
.terminate(codex_code_mode_protocol::CellId::from(cell_id).into())
.await
.expect("terminate cell");
state.disconnect().await;
codex-rs/code-mode-host/src/lib.rs
+14 -7
View File
@@ -382,6 +382,7 @@ impl HostState {
HostRequest::OpenSession {
session_id,
cell_execution_limits,
..
} => {
let result = CodeModeSessionCellExecutionLimits::try_from(
cell_execution_limits.unwrap_or_default(),
@@ -452,6 +453,7 @@ impl HostState {
HostRequest::Wait {
session_id,
request,
..
} => {
let result = match self.session(&session_id) {
Ok(session) => {
@@ -476,15 +478,20 @@ impl HostState {
HostRequest::Terminate {
session_id,
cell_id,
..
} => {
let result = match self.session(&session_id) {
Ok(session) => session.terminate(cell_id.into()).await.and_then(|outcome| {
let outcome = outcome.with_code_mode_host_duration(received_at.elapsed());
Ok(HostResponse::WaitCompleted {
outcome: WireWaitOutcome::try_from(outcome)
.map_err(|error| error.to_string())?,
})
}),
Ok(session) => session
.terminate(codex_code_mode_protocol::CellId::from(cell_id).into())
.await
.and_then(|outcome| {
let outcome =
outcome.with_code_mode_host_duration(received_at.elapsed());
Ok(HostResponse::WaitCompleted {
outcome: WireWaitOutcome::try_from(outcome)
.map_err(|error| error.to_string())?,
})
}),
Err(err) => Err(err),
};
self.respond(request_id, result);
codex-rs/code-mode-host/tests/grpc.rs
+26 -7
View File
@@ -327,6 +327,7 @@ async fn grpc_early_yield_preserves_cell_and_later_waits() -> Result<()> {
WaitRequest {
cell_id: cell_id.clone(),
yield_time_ms: 60_000,
max_output_tokens: None,
},
Some(wait_signal.clone()),
);
@@ -347,6 +348,7 @@ async fn grpc_early_yield_preserves_cell_and_later_waits() -> Result<()> {
WaitRequest {
cell_id: cell_id.clone(),
yield_time_ms: 60_000,
max_output_tokens: None,
},
/*preempt*/ None,
),
@@ -397,13 +399,17 @@ async fn shutdown_immediately_rejects_new_operations() -> Result<()> {
WaitRequest {
cell_id: cell_id("missing"),
yield_time_ms: 1,
max_output_tokens: None,
},
/*preempt*/ None,
)
.await,
Err(expected.clone())
);
assert_eq!(session.terminate(cell_id("missing")).await, Err(expected));
assert_eq!(
session.terminate(cell_id("missing").into()).await,
Err(expected)
);
shutdown.await.map_err(anyhow::Error::msg)?;
Ok(())
@@ -452,6 +458,7 @@ async fn cancelling_execution_before_admission_keeps_the_session_usable() -> Res
WaitRequest {
cell_id: abandoned_cell.clone(),
yield_time_ms: 1,
max_output_tokens: None,
},
/*preempt*/ None,
)
@@ -512,6 +519,7 @@ async fn dropping_a_started_cell_off_runtime_terminates_its_buffered_remote_exec
WaitRequest {
cell_id: abandoned_cell.clone(),
yield_time_ms: 1,
max_output_tokens: None,
},
/*preempt*/ None,
)
@@ -540,6 +548,7 @@ async fn dropping_a_started_cell_off_runtime_terminates_its_buffered_remote_exec
WaitRequest {
cell_id: abandoned_cell.clone(),
yield_time_ms: 1,
max_output_tokens: None,
},
/*preempt*/ None,
)
@@ -751,6 +760,7 @@ async fn concurrent_wait_rejects_without_displacing_the_active_observer() -> Res
WaitRequest {
cell_id: running_cell.clone(),
yield_time_ms: 100,
max_output_tokens: None,
},
)
.await?;
@@ -762,6 +772,7 @@ async fn concurrent_wait_rejects_without_displacing_the_active_observer() -> Res
WaitRequest {
cell_id: running_cell.clone(),
yield_time_ms: 60_000,
max_output_tokens: None,
},
/*preempt*/ None,
),
@@ -785,7 +796,7 @@ async fn concurrent_wait_rejects_without_displacing_the_active_observer() -> Res
})
);
let actual = session
.terminate(running_cell.clone())
.terminate(running_cell.clone().into())
.await
.map_err(anyhow::Error::msg)?;
assert_eq!(
@@ -838,6 +849,7 @@ async fn dropping_a_wait_retires_its_observer_before_the_next_wait() -> Result<(
WaitRequest {
cell_id: running_cell.clone(),
yield_time_ms: 60_000,
max_output_tokens: None,
},
)
.await?;
@@ -850,6 +862,7 @@ async fn dropping_a_wait_retires_its_observer_before_the_next_wait() -> Result<(
WaitRequest {
cell_id: running_cell.clone(),
yield_time_ms: 1,
max_output_tokens: None,
},
/*preempt*/ None,
),
@@ -866,7 +879,7 @@ async fn dropping_a_wait_retires_its_observer_before_the_next_wait() -> Result<(
})
);
let actual = session
.terminate(running_cell.clone())
.terminate(running_cell.clone().into())
.await
.map_err(anyhow::Error::msg)?;
assert_eq!(
@@ -1091,6 +1104,7 @@ async fn sessions_enforce_independent_yield_limits() -> Result<()> {
WaitRequest {
cell_id: cell_id("1"),
yield_time_ms: 60_000,
max_output_tokens: None,
},
/*preempt*/ None,
),
@@ -1117,7 +1131,7 @@ async fn sessions_enforce_independent_yield_limits() -> Result<()> {
text_response("1", "isolated", actual.code_mode_host_duration())
);
let actual = limited
.terminate(cell_id("1"))
.terminate(cell_id("1").into())
.await
.map_err(anyhow::Error::msg)?;
assert_eq!(
@@ -1156,6 +1170,7 @@ async fn dropping_a_grpc_lease_retires_its_server_session() -> Result<()> {
let mut lease = client
.open_session(grpc::OpenSessionRequest {
cell_execution_limits: None,
output_truncation: None,
})
.await
.context("open raw gRPC session")?
@@ -1236,6 +1251,7 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> {
WaitRequest {
cell_id: old_cell_id.clone(),
yield_time_ms: 60_000,
max_output_tokens: None,
},
)
.await?;
@@ -1280,7 +1296,7 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> {
.context("host loss did not retire the original generation's cell")?;
let _replacement = HostHarness::start(&listen_url).await?;
assert!(session.terminate(old_cell_id.clone()).await.is_err());
assert!(session.terminate(old_cell_id.clone().into()).await.is_err());
let mut rejected = request(r#"await tools.echo({must_not_run: true});"#);
rejected.enabled_tools = vec![tool("echo")];
@@ -1291,6 +1307,7 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> {
WaitRequest {
cell_id: old_cell_id.clone(),
yield_time_ms: 1,
max_output_tokens: None,
},
/*preempt*/ None,
);
@@ -1406,6 +1423,7 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> {
WaitRequest {
cell_id: replacement_cell_id.clone(),
yield_time_ms: 1,
max_output_tokens: None,
},
/*preempt*/ None,
)
@@ -1420,7 +1438,7 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> {
})
);
let actual = session
.terminate(replacement_cell_id.clone())
.terminate(replacement_cell_id.clone().into())
.await
.map_err(anyhow::Error::msg)?;
assert_eq!(
@@ -1437,13 +1455,14 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> {
WaitRequest {
cell_id: old_cell_id.clone(),
yield_time_ms: 1,
max_output_tokens: None,
},
/*preempt*/ None,
)
.await
.unwrap_err();
assert!(stale_wait.contains("stale code-mode host generation"));
let stale_termination = session.terminate(old_cell_id).await.unwrap_err();
let stale_termination = session.terminate(old_cell_id.into()).await.unwrap_err();
assert!(stale_termination.contains("stale code-mode host generation"));
session.shutdown().await.map_err(anyhow::Error::msg)?;
Ok(())
codex-rs/code-mode-host/tests/grpc/admission_tests.rs
+5 -1
View File
@@ -228,6 +228,7 @@ async fn routing_metadata_is_scoped_to_each_session_on_a_shared_provider() -> Re
WaitRequest {
cell_id: id.clone(),
yield_time_ms: 1,
max_output_tokens: None,
},
/*preempt*/ None,
)
@@ -237,7 +238,10 @@ async fn routing_metadata_is_scoped_to_each_session_on_a_shared_provider() -> Re
outcome,
WaitOutcome::LiveCell(RuntimeResponse::Yielded { .. })
));
session.terminate(id).await.map_err(anyhow::Error::msg)?;
session
.terminate(id.into())
.await
.map_err(anyhow::Error::msg)?;
}
for session in [&parent, &child] {
execute(
codex-rs/code-mode-host/tests/grpc/image_cancellation_tests.rs
+1 -1
View File
@@ -57,7 +57,7 @@ try { image(value); } catch { text(calls); }
.await
.context("getter did not signal startup")?;
let result = timeout(Duration::from_secs(1), session.terminate(id.clone()))
let result = timeout(Duration::from_secs(1), session.terminate(id.clone().into()))
.await
.context("workload did not stop within the cancellation deadline")?
.map_err(anyhow::Error::msg)?;
codex-rs/code-mode-host/tests/grpc_notifications.rs
+2 -1
View File
@@ -225,6 +225,7 @@ async fn completed_waits_drain_pending_notifications_before_returning() -> Resul
WaitRequest {
cell_id: cell_id("1"),
yield_time_ms: 5_000,
max_output_tokens: None,
},
/*preempt*/ None,
)
@@ -292,7 +293,7 @@ async fn termination_cancels_pending_notifications() -> Result<()> {
}
);
let actual = session
.terminate(cell_id("1"))
.terminate(cell_id("1").into())
.await
.map_err(anyhow::Error::msg)?;
assert_eq!(
codex-rs/code-mode-host/tests/grpc_tcp.rs
+1
View File
@@ -150,6 +150,7 @@ async fn tcp_listener_opens_a_grpc_session() -> Result<()> {
let response = client
.open_session(grpc::OpenSessionRequest {
cell_execution_limits: None,
output_truncation: None,
})
.await
.context("failed to open gRPC code-mode session")?;
codex-rs/code-mode-host/tests/grpc_timing.rs
+4 -1
View File
@@ -40,6 +40,7 @@ async fn execution_timing_includes_javascript_but_excludes_delayed_reads() -> Re
let mut events = client
.open_session(grpc::OpenSessionRequest {
cell_execution_limits: None,
output_truncation: None,
})
.await?
.into_inner();
@@ -229,6 +230,7 @@ async fn observation_timing_excludes_previous_requests_and_background_time() ->
WaitRequest {
cell_id: cell_id.clone(),
yield_time_ms: 50,
max_output_tokens: None,
},
/*preempt*/ None,
)
@@ -254,7 +256,7 @@ async fn observation_timing_excludes_previous_requests_and_background_time() ->
let observed = Instant::now();
let terminated = session
.terminate(cell_id.clone())
.terminate(cell_id.clone().into())
.await
.map_err(anyhow::Error::msg)?;
let duration = terminated
@@ -276,6 +278,7 @@ async fn observation_timing_excludes_previous_requests_and_background_time() ->
WaitRequest {
cell_id: cell_id.clone(),
yield_time_ms: 50,
max_output_tokens: None,
},
/*preempt*/ None,
)
codex-rs/code-mode-host/tests/settled.rs
+2
View File
@@ -109,6 +109,7 @@ async fn assert_partial_results(provider: &dyn CodeModeSessionProvider) {
WaitRequest {
cell_id: cell_id.clone(),
yield_time_ms: 10,
max_output_tokens: None,
},
/*preempt*/ None,
)
@@ -133,6 +134,7 @@ async fn assert_partial_results(provider: &dyn CodeModeSessionProvider) {
WaitRequest {
cell_id: cell_id.clone(),
yield_time_ms: 5_000,
max_output_tokens: None,
},
/*preempt*/ None,
)
codex-rs/code-mode-host/tests/stdio.rs
+14 -3
View File
@@ -206,6 +206,7 @@ async fn execute_to_terminal(
WaitRequest {
cell_id,
yield_time_ms: 60_000,
max_output_tokens: None,
},
/*preempt*/ None,
)
@@ -270,6 +271,7 @@ async fn interrupt_yields_observations_without_stopping_the_cell() {
let request = WaitRequest {
cell_id: cell_id.clone(),
yield_time_ms: 60_000,
max_output_tokens: None,
};
let response =
tokio::time::timeout(Duration::from_secs(5), session.wait(request, Some(signal)))
@@ -288,6 +290,7 @@ async fn interrupt_yields_observations_without_stopping_the_cell() {
WaitRequest {
cell_id,
yield_time_ms: 60_000,
max_output_tokens: None,
},
/*preempt*/ None,
),
@@ -350,6 +353,7 @@ async fn session_execution_limits_are_isolated_on_a_shared_process_host() {
WaitRequest {
cell_id: cell_id("1"),
yield_time_ms: 60_000,
max_output_tokens: None,
},
/*preempt*/ None,
),
@@ -367,7 +371,7 @@ async fn session_execution_limits_are_isolated_on_a_shared_process_host() {
);
limited
.terminate(cell_id("1"))
.terminate(cell_id("1").into())
.await
.expect("terminate cell");
limited.shutdown().await.expect("shutdown limited session");
@@ -468,6 +472,7 @@ text(result.value);
WaitRequest {
cell_id: cell_id("3"),
yield_time_ms: 1,
max_output_tokens: None,
},
/*preempt*/ None,
)
@@ -482,7 +487,7 @@ text(result.value);
})
);
let actual = session
.terminate(cell_id("3"))
.terminate(cell_id("3").into())
.await
.expect("terminate cell");
assert_eq!(
@@ -546,6 +551,7 @@ async fn dropping_long_wait_releases_observer_before_next_wait() {
WaitRequest {
cell_id: wait_cell_id,
yield_time_ms: 60_000,
max_output_tokens: None,
},
/*preempt*/ None,
)
@@ -561,6 +567,7 @@ async fn dropping_long_wait_releases_observer_before_next_wait() {
WaitRequest {
cell_id: running_cell_id.clone(),
yield_time_ms: 1,
max_output_tokens: None,
},
/*preempt*/ None,
),
@@ -577,7 +584,7 @@ async fn dropping_long_wait_releases_observer_before_next_wait() {
})
);
session
.terminate(running_cell_id)
.terminate(running_cell_id.into())
.await
.expect("terminate cell");
session.shutdown().await.expect("shutdown remote session");
@@ -652,6 +659,7 @@ return;
WaitRequest {
cell_id: wait_cell_id,
yield_time_ms: 60_000,
max_output_tokens: None,
},
/*preempt*/ None,
)
@@ -995,6 +1003,7 @@ async fn child_process_loss_cleans_up_and_rebuilds_the_shared_host() {
WaitRequest {
cell_id: wait_a_cell,
yield_time_ms: 60_000,
max_output_tokens: None,
},
/*preempt*/ None,
)
@@ -1008,6 +1017,7 @@ async fn child_process_loss_cleans_up_and_rebuilds_the_shared_host() {
WaitRequest {
cell_id: wait_b_cell,
yield_time_ms: 60_000,
max_output_tokens: None,
},
/*preempt*/ None,
)
@@ -1092,6 +1102,7 @@ async fn child_process_loss_cleans_up_and_rebuilds_the_shared_host() {
WaitRequest {
cell_id: cell_b.clone(),
yield_time_ms: 1,
max_output_tokens: None,
},
/*preempt*/ None,
)
codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto
+14
View File
@@ -45,6 +45,14 @@ service CodeModeHost {
message OpenSessionRequest {
optional SessionCellExecutionLimits cell_execution_limits = 1;
// Absent preserves the legacy client-side truncation behavior.
optional OutputTruncation output_truncation = 2;
}
enum OutputTruncation {
OUTPUT_TRUNCATION_UNSPECIFIED = 0;
OUTPUT_TRUNCATION_CLIENT = 1;
OUTPUT_TRUNCATION_SERVER = 2;
}
message SessionCellExecutionLimits {
@@ -64,6 +72,8 @@ message SessionEvent {
message SessionOpened {
string session_id = 1;
// Absent (including from an older host) means the client owns truncation.
optional OutputTruncation output_truncation = 2;
}
message CloseSessionRequest {
@@ -182,6 +192,8 @@ message WaitRequest {
string cell_id = 2;
string wait_id = 3;
uint64 yield_time_ms = 4;
// Only set for a session that selected server-side output truncation.
optional uint64 max_output_tokens = 5;
}
message WaitResponse {
@@ -211,6 +223,8 @@ message YieldObservationResponse {}
message TerminateRequest {
string session_id = 1;
string cell_id = 2;
// The terminating observation has its own output budget.
optional uint64 max_output_tokens = 3;
}
message ExecutionOutcome {
codex-rs/code-mode-protocol/src/host/host_tests.rs
+18
View File
@@ -27,6 +27,7 @@ use super::WireContentItem;
use super::WireExecuteRequest;
use super::WireImageDetail;
use super::WireNestedToolCall;
use super::WireOutputTruncation;
use super::WireResult;
use super::WireRuntimeResponse;
use super::WireSessionCellExecutionLimits;
@@ -249,6 +250,7 @@ fn open_session_serializes_optional_cell_execution_limits() {
max_yield_time_ms: Some(250),
max_heap_size_bytes: Some(16 * 1024 * 1024),
}),
output_truncation: None,
},
json!({
"method": "session/open",
@@ -304,9 +306,23 @@ fn client_to_host_v1_variants_are_pinned() {
HostRequest::OpenSession {
session_id: session_id(),
cell_execution_limits: None,
output_truncation: None,
},
json!({ "method": "session/open", "sessionId": "session-1" }),
),
(
request_id(/*value*/ 10),
HostRequest::OpenSession {
session_id: session_id(),
cell_execution_limits: None,
output_truncation: Some(WireOutputTruncation::Server),
},
json!({
"method": "session/open",
"sessionId": "session-1",
"outputTruncation": "server",
}),
),
(
request_id(/*value*/ 2),
HostRequest::Execute {
@@ -353,6 +369,7 @@ fn client_to_host_v1_variants_are_pinned() {
cell_id: cell_id("cell-1"),
yield_time_ms: 50,
},
max_output_tokens: None,
},
json!({
"method": "session/wait",
@@ -365,6 +382,7 @@ fn client_to_host_v1_variants_are_pinned() {
HostRequest::Terminate {
session_id: session_id(),
cell_id: cell_id("cell-1"),
max_output_tokens: None,
},
json!({
"method": "session/terminate",
codex-rs/code-mode-protocol/src/host/message.rs
+9
View File
@@ -15,6 +15,7 @@ use super::SupportedProtocolVersions;
use super::WireCellId;
use super::WireExecuteRequest;
use super::WireNestedToolCall;
use super::WireOutputTruncation;
use super::WireRuntimeResponse;
use super::WireSessionCellExecutionLimits;
use super::WireWaitOutcome;
@@ -186,6 +187,10 @@ pub enum HostRequest {
session_id: SessionId,
#[serde(default, skip_serializing_if = "Option::is_none")]
cell_execution_limits: Option<WireSessionCellExecutionLimits>,
/// Sent only when output truncation support was agreed at connection hello.
/// Omission preserves the legacy session behavior and wire shape.
#[serde(default, skip_serializing_if = "Option::is_none")]
output_truncation: Option<WireOutputTruncation>,
},
#[serde(rename = "session/execute")]
Execute {
@@ -196,11 +201,15 @@ pub enum HostRequest {
Wait {
session_id: SessionId,
request: WireWaitRequest,
#[serde(default, skip_serializing_if = "Option::is_none")]
max_output_tokens: Option<usize>,
},
#[serde(rename = "session/terminate")]
Terminate {
session_id: SessionId,
cell_id: WireCellId,
#[serde(default, skip_serializing_if = "Option::is_none")]
max_output_tokens: Option<usize>,
},
#[serde(rename = "session/shutdown")]
ShutdownSession { session_id: SessionId },
codex-rs/code-mode-protocol/src/host/mod.rs
+7
View File
@@ -12,6 +12,12 @@ mod types;
/// Maximum number of unresolved delegate callbacks allowed per host connection.
pub const MAX_PENDING_DELEGATE_CALLS: usize = 1_024;
/// Stdio connection support for the optional `session/open.outputTruncation` field.
/// Each session selects its own setting after the connection hello.
pub const BOUNDED_OUTPUT_CAPABILITY: &str = "bounded-output-v1";
/// Maximum retained bytes between server-truncated observations.
pub const BOUNDED_OUTPUT_MAX_BYTES: usize = 1024 * 1024;
/// Negotiated support for cell execution resource limits on `session/open`.
pub const SESSION_RESOURCE_LIMITS_CAPABILITY: &str = "session-cell-execution-resource-limits";
@@ -38,6 +44,7 @@ pub use payload::WireContentItem;
pub use payload::WireExecuteRequest;
pub use payload::WireImageDetail;
pub use payload::WireNestedToolCall;
pub use payload::WireOutputTruncation;
pub use payload::WireRuntimeResponse;
pub use payload::WireSessionCellExecutionLimits;
pub use payload::WireToolDefinition;
codex-rs/code-mode-protocol/src/host/payload.rs
+30 -1
View File
@@ -8,6 +8,7 @@ use serde_json::Value as JsonValue;
use crate::CellId;
use crate::CodeModeNestedToolCall;
use crate::CodeModeOutputTruncation;
use crate::CodeModeSessionCellExecutionLimits;
use crate::CodeModeToolKind;
use crate::ExecuteRequest;
@@ -51,6 +52,32 @@ impl TryFrom<WireSessionCellExecutionLimits> for CodeModeSessionCellExecutionLim
}
}
/// The output truncation setting carried by a V1 session-open request.
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "camelCase")]
pub enum WireOutputTruncation {
Client,
Server,
}
impl From<CodeModeOutputTruncation> for WireOutputTruncation {
fn from(value: CodeModeOutputTruncation) -> Self {
match value {
CodeModeOutputTruncation::Client => Self::Client,
CodeModeOutputTruncation::Server => Self::Server,
}
}
}
impl From<WireOutputTruncation> for CodeModeOutputTruncation {
fn from(value: WireOutputTruncation) -> Self {
match value {
WireOutputTruncation::Client => Self::Client,
WireOutputTruncation::Server => Self::Server,
}
}
}
/// A cell identifier with a wire representation owned by protocol V1.
#[derive(Clone, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)]
#[serde(transparent)]
@@ -211,7 +238,8 @@ impl TryFrom<WireExecuteRequest> for ExecuteRequest {
}
}
/// The complete wait request shape supported by protocol V1.
/// The nested wait payload for protocol V1. Its optional budget is carried by
/// the enclosing `HostRequest::Wait`.
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct WireWaitRequest {
@@ -233,6 +261,7 @@ impl From<WireWaitRequest> for WaitRequest {
Self {
cell_id: value.cell_id.into(),
yield_time_ms: value.yield_time_ms,
max_output_tokens: None,
}
}
}
codex-rs/code-mode-protocol/src/lib.rs
+3
View File
@@ -35,14 +35,17 @@ pub use runtime::ExecuteRequest;
pub use runtime::ExecuteToPendingOutcome;
pub use runtime::MissingCodeModeHostDuration;
pub use runtime::RuntimeResponse;
pub use runtime::TerminateRequest;
pub use runtime::WaitOutcome;
pub use runtime::WaitRequest;
pub use runtime::WaitToPendingOutcome;
pub use runtime::WaitToPendingRequest;
pub use session::CellId;
pub use session::CodeModeOutputTruncation;
pub use session::CodeModeSession;
pub use session::CodeModeSessionCellExecutionLimits;
pub use session::CodeModeSessionDelegate;
pub use session::CodeModeSessionOptions;
pub use session::CodeModeSessionProvider;
pub use session::CodeModeSessionProviderFuture;
pub use session::CodeModeSessionResultFuture;
codex-rs/code-mode-protocol/src/runtime.rs
+20
View File
@@ -29,6 +29,26 @@ pub struct ExecuteRequest {
pub struct WaitRequest {
pub cell_id: CellId,
pub yield_time_ms: u64,
/// Limit for this observation in Server mode; ignored in Client mode.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub max_output_tokens: Option<usize>,
}
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub struct TerminateRequest {
pub cell_id: CellId,
/// Limit for the final observation in Server mode; ignored in Client mode.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub max_output_tokens: Option<usize>,
}
impl From<CellId> for TerminateRequest {
fn from(cell_id: CellId) -> Self {
Self {
cell_id,
max_output_tokens: None,
}
}
}
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
codex-rs/code-mode-protocol/src/session.rs
+60 -1
View File
@@ -12,6 +12,7 @@ use tokio_util::sync::CancellationToken;
use crate::CodeModeNestedToolCall;
use crate::ExecuteRequest;
use crate::RuntimeResponse;
use crate::TerminateRequest;
use crate::WaitOutcome;
use crate::WaitRequest;
@@ -30,6 +31,33 @@ pub struct CodeModeSessionCellExecutionLimits {
pub max_heap_size_bytes: Option<usize>,
}
/// The owner of the output budget for each exec or wait observation.
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub enum CodeModeOutputTruncation {
/// Return raw output; the caller applies its own budget and formatting.
#[default]
Client,
/// The host applies the exec, wait, or terminate observation's budget before
/// transmitting output (10,000 tokens by default). Clients must not truncate
/// it again or append the script error: status and error metadata remain separate.
///
/// Between observations, the host retains at most
/// [`crate::host::BOUNDED_OUTPUT_MAX_BYTES`] of output payload. It also
/// bounds the storage for retained items separately, without charging it
/// against the payload or observation token budgets.
/// Ordinary text notices report overflow, even for a zero-token budget.
/// Media are indivisible and can be omitted if they exceed the collector.
/// Without overflow, formatting matches Core's client policy.
Server,
}
/// Settings shared by every cell in one code-mode session.
#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct CodeModeSessionOptions {
pub cell_execution_limits: CodeModeSessionCellExecutionLimits,
pub output_truncation: CodeModeOutputTruncation,
}
#[derive(Clone, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)]
pub struct CellId(String);
@@ -152,6 +180,12 @@ impl CodeModeSessionDelegate for NoopCodeModeSessionDelegate {
/// must keep those values isolated. Implementations may execute cells
/// in-process or remotely.
pub trait CodeModeSession: Send + Sync {
/// Identifies who applies output truncation for this session's responses.
/// Providers must not change modes when reconnecting to a new host.
fn output_truncation(&self) -> CodeModeOutputTruncation {
CodeModeOutputTruncation::Client
}
/// Executes a cell, yielding its foreground observation if `preempt` is signaled.
/// The cell continues running after its observation yields.
fn execute<'a>(
@@ -162,13 +196,19 @@ pub trait CodeModeSession: Send + Sync {
) -> CodeModeSessionResultFuture<'a, StartedCell>;
/// Waits for a cell, yielding the observation if `preempt` is signaled while the cell continues.
/// Server mode applies the request's output limit; Client mode ignores it.
fn wait<'a>(
&'a self,
request: WaitRequest,
preempt: Option<CancellationToken>,
) -> CodeModeSessionResultFuture<'a, WaitOutcome>;
fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome>;
/// Terminates a cell, applying the request's limit to its final observation
/// in Server mode; Client mode ignores the limit.
fn terminate<'a>(
&'a self,
request: TerminateRequest,
) -> CodeModeSessionResultFuture<'a, WaitOutcome>;
fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()>;
}
@@ -201,6 +241,25 @@ pub trait CodeModeSessionProvider: Send + Sync {
})
}
}
/// Requests output handling for one session, without changing other sessions
/// created by this provider. Server mode requires host support.
fn create_session_with_options<'a>(
&'a self,
options: CodeModeSessionOptions,
) -> CodeModeSessionProviderFuture<'a> {
match options.output_truncation {
CodeModeOutputTruncation::Client => {
self.create_session_with_limits(options.cell_execution_limits)
}
CodeModeOutputTruncation::Server => Box::pin(async {
Err(
"code-mode session provider does not support server output truncation"
.to_string(),
)
}),
}
}
}
#[cfg(test)]
codex-rs/code-mode-runtime/src/service.rs
+9 -3
View File
@@ -15,6 +15,7 @@ use codex_code_mode_protocol::FunctionCallOutputContentItem;
use codex_code_mode_protocol::ImageDetail;
use codex_code_mode_protocol::RuntimeResponse;
use codex_code_mode_protocol::StartedCell;
use codex_code_mode_protocol::TerminateRequest;
use codex_code_mode_protocol::WaitOutcome;
use codex_code_mode_protocol::WaitRequest;
use codex_code_mode_protocol::WaitToPendingOutcome;
@@ -129,6 +130,7 @@ impl InProcessCodeModeSession {
preempt: Option<CancellationToken>,
) -> CodeModeSessionResultFuture<'static, WaitOutcome> {
let WaitRequest {
max_output_tokens: _,
cell_id,
yield_time_ms,
} = request;
@@ -158,7 +160,8 @@ impl InProcessCodeModeSession {
}
}
pub async fn terminate(&self, cell_id: CellId) -> Result<WaitOutcome, String> {
pub async fn terminate(&self, request: TerminateRequest) -> Result<WaitOutcome, String> {
let cell_id = request.cell_id;
match self.runtime.terminate(&runtime_cell_id(&cell_id)).await {
Ok(event) => Ok(WaitOutcome::LiveCell(runtime_response(&cell_id, event)?)),
Err(runtime::Error::MissingCell(_) | runtime::Error::ClosedCell(_)) => {
@@ -239,8 +242,11 @@ impl CodeModeSession for InProcessCodeModeSession {
Box::pin(InProcessCodeModeSession::wait(self, request, preempt))
}
fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome> {
Box::pin(InProcessCodeModeSession::terminate(self, cell_id))
fn terminate<'a>(
&'a self,
request: TerminateRequest,
) -> CodeModeSessionResultFuture<'a, WaitOutcome> {
Box::pin(InProcessCodeModeSession::terminate(self, request))
}
fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()> {
codex-rs/code-mode-runtime/src/service_audio_tests.rs
+1
View File
@@ -149,6 +149,7 @@ async fn audio_helper_bounds_wav_data_and_preserves_outputs_after_yield() {
let actual = service
.wait(
WaitRequest {
max_output_tokens: None,
cell_id: cell_id("1"),
yield_time_ms: 60_000,
},
codex-rs/code-mode-runtime/src/service_contract_tests.rs
+10 -6
View File
@@ -223,6 +223,7 @@ async fn yielded_cells_retain_their_own_delegate_until_closed() {
let actual = service
.wait(
WaitRequest {
max_output_tokens: None,
cell_id: cell_id("1"),
yield_time_ms: 60_000,
},
@@ -253,7 +254,7 @@ async fn yielded_cells_retain_their_own_delegate_until_closed() {
);
assert!(weak_a.upgrade().is_none());
service.terminate(cell_id("2")).await.unwrap();
service.terminate(cell_id("2").into()).await.unwrap();
assert_eq!(
next_event(&mut events_b).await,
DelegateEvent::ToolCancelled
@@ -293,6 +294,7 @@ async fn yields_and_resumes() {
let actual = service
.wait(
WaitRequest {
max_output_tokens: None,
cell_id: cell_id("1"),
yield_time_ms: 60_000,
},
@@ -418,7 +420,7 @@ async fn observed_natural_completion_wins_over_termination() {
})
.await
.unwrap();
let actual = service.terminate(cell_id("1")).await.unwrap();
let actual = service.terminate(cell_id("1").into()).await.unwrap();
let result_buffered_for = actual.result_buffered_for().unwrap();
assert!(!result_buffered_for.is_zero());
assert_eq!(
@@ -461,7 +463,7 @@ async fn termination_cancels_pending_callbacks_before_responding() {
}
);
assert_eq!(
service.terminate(cell_id("1")).await.unwrap(),
service.terminate(cell_id("1").into()).await.unwrap(),
WaitOutcome::LiveCell(RuntimeResponse::Terminated {
code_mode_host_duration: None,
cell_id: cell_id("1"),
@@ -541,13 +543,13 @@ async fn repeated_termination_is_rejected_while_callback_cleanup_is_pending() {
let terminating_service = Arc::clone(&service);
let first_termination =
tokio::spawn(async move { terminating_service.terminate(cell_id("1")).await });
tokio::spawn(async move { terminating_service.terminate(cell_id("1").into()).await });
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::NotificationCancelled
);
let repeated_termination = service.terminate(cell_id("1")).await;
let repeated_termination = service.terminate(cell_id("1").into()).await;
delegate.release_notification();
assert_eq!(
@@ -594,6 +596,7 @@ async fn second_observer_is_rejected_without_displacing_the_first() {
WaitRequest {
cell_id: cell_id("1"),
yield_time_ms: 60_000,
max_output_tokens: None,
},
/*preempt*/ None,
)
@@ -604,6 +607,7 @@ async fn second_observer_is_rejected_without_displacing_the_first() {
WaitRequest {
cell_id: cell_id("1"),
yield_time_ms: 60_000,
max_output_tokens: None,
},
/*preempt*/ None,
)
@@ -618,7 +622,7 @@ async fn second_observer_is_rejected_without_displacing_the_first() {
content_items: Vec::new(),
};
assert_eq!(
service.terminate(cell_id("1")).await.unwrap(),
service.terminate(cell_id("1").into()).await.unwrap(),
WaitOutcome::LiveCell(terminated.clone())
);
assert_eq!(
codex-rs/code-mode-runtime/src/service_tests.rs
+13 -5
View File
@@ -98,6 +98,7 @@ async fn interrupt_yields_execute_and_wait_without_stopping_the_cell() {
WaitRequest {
cell_id: cell_id("1"),
yield_time_ms: 60_000,
max_output_tokens: None,
},
Some(wait_signal),
)
@@ -125,6 +126,7 @@ async fn interrupt_yields_execute_and_wait_without_stopping_the_cell() {
WaitRequest {
cell_id: cell_id("1"),
yield_time_ms: 60_000,
max_output_tokens: None,
},
/*preempt*/ None,
)
@@ -221,6 +223,7 @@ async fn execute_and_wait_clamp_yield_grace_without_stopping_the_cell() {
WaitRequest {
cell_id: cell_id("1"),
yield_time_ms: 10_000,
max_output_tokens: None,
},
/*preempt*/ None,
)
@@ -243,6 +246,7 @@ async fn execute_and_wait_clamp_yield_grace_without_stopping_the_cell() {
WaitRequest {
cell_id: cell_id("1"),
yield_time_ms: 10_000,
max_output_tokens: None,
},
/*preempt*/ None,
)
@@ -291,6 +295,7 @@ async fn wait_waits_for_nested_tool_during_yield_grace() {
WaitRequest {
cell_id: cell_id("1"),
yield_time_ms: 10_000,
max_output_tokens: None,
},
/*preempt*/ None,
)
@@ -366,6 +371,7 @@ async fn zero_yield_limit_is_immediate_and_scoped_to_its_session() {
WaitRequest {
cell_id: cell_id("1"),
yield_time_ms: 60_000,
max_output_tokens: None,
},
/*preempt*/ None,
)
@@ -441,7 +447,7 @@ async fn tool_start_acknowledgement_reports_termination_before_invocation() {
let (startup, terminated) = tokio::join!(
wait_until_tool_started(&delegate),
service.terminate(cell_id("1")),
service.terminate(cell_id("1").into()),
);
assert_eq!(startup, Err("cell closed before nested tool started"));
let response = RuntimeResponse::Terminated {
@@ -746,6 +752,7 @@ text("after exit");
let actual = service
.wait(
WaitRequest {
max_output_tokens: None,
cell_id: cell_id("1"),
yield_time_ms: 30_000,
},
@@ -1106,7 +1113,7 @@ async fn execute_to_pending_returns_once_the_runtime_is_quiescent() {
}
);
let termination = service.terminate(cell_id("1")).await.unwrap();
let termination = service.terminate(cell_id("1").into()).await.unwrap();
assert_eq!(
termination,
@@ -1150,7 +1157,7 @@ await Promise.all([
}
);
let termination = service.terminate(cell_id("1")).await.unwrap();
let termination = service.terminate(cell_id("1").into()).await.unwrap();
assert_eq!(
termination,
@@ -1218,7 +1225,7 @@ await Promise.all([
})
);
let termination = service.terminate(cell_id("1")).await.unwrap();
let termination = service.terminate(cell_id("1").into()).await.unwrap();
assert_eq!(
termination,
@@ -1285,7 +1292,7 @@ await new Promise(() => {});
})
);
let termination = service.terminate(cell_id("1")).await.unwrap();
let termination = service.terminate(cell_id("1").into()).await.unwrap();
assert_eq!(
termination,
@@ -2132,6 +2139,7 @@ async fn wait_reports_missing_cell_separately_from_runtime_results() {
WaitRequest {
cell_id: cell_id("missing"),
yield_time_ms: 1,
max_output_tokens: None,
},
/*preempt*/ None,
)
codex-rs/code-mode-runtime/tests/settled.rs
+1 -1
View File
@@ -242,7 +242,7 @@ async fn terminates_helpers_waiting_for_unsettled_inputs() {
assert_eq!(
tokio::time::timeout(
Duration::from_secs(/*secs*/ 5),
session.terminate(cell_id.clone())
session.terminate(cell_id.clone().into())
)
.await
.expect("termination timed out")
codex-rs/code-mode/src/grpc_session/mod.rs
+7 -2
View File
@@ -17,6 +17,7 @@ use codex_code_mode_protocol::CodeModeSessionProviderFuture;
use codex_code_mode_protocol::CodeModeSessionResultFuture;
use codex_code_mode_protocol::ExecuteRequest;
use codex_code_mode_protocol::StartedCell;
use codex_code_mode_protocol::TerminateRequest;
use codex_code_mode_protocol::WaitOutcome;
use codex_code_mode_protocol::WaitRequest;
use codex_code_mode_protocol::grpc;
@@ -119,6 +120,7 @@ impl GrpcCodeModeSessionProvider {
loop {
let mut request = tonic::Request::new(grpc::OpenSessionRequest {
cell_execution_limits,
output_truncation: None,
});
inject_span_traceparent(&mut request, &open_session_span);
let result = async {
@@ -269,8 +271,11 @@ impl CodeModeSession for GrpcCodeModeSession {
Box::pin(self.inner.wait(request, preempt))
}
fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome> {
Box::pin(self.inner.terminate(cell_id))
fn terminate<'a>(
&'a self,
request: TerminateRequest,
) -> CodeModeSessionResultFuture<'a, WaitOutcome> {
Box::pin(self.inner.terminate(request.cell_id))
}
fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()> {
codex-rs/code-mode/src/grpc_session/operations.rs
+2
View File
@@ -313,6 +313,7 @@ impl SessionInner {
cell_id: expected_cell_id.as_str().to_string(),
wait_id: wait_id.clone(),
yield_time_ms: request.yield_time_ms,
max_output_tokens: None,
};
let mut client = self.client();
let watcher = preempt.map(|signal| {
@@ -368,6 +369,7 @@ impl SessionInner {
client.terminate(self.request(grpc::TerminateRequest {
session_id: self.id.clone(),
cell_id: cell_id.as_str().to_string(),
max_output_tokens: None,
})),
)
.await?;
codex-rs/code-mode/src/grpc_session/reconnect.rs
+11 -4
View File
@@ -4,13 +4,13 @@ use std::sync::PoisonError;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering;
use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::CodeModeSession;
use codex_code_mode_protocol::CodeModeSessionCellExecutionLimits;
use codex_code_mode_protocol::CodeModeSessionDelegate;
use codex_code_mode_protocol::CodeModeSessionResultFuture;
use codex_code_mode_protocol::ExecuteRequest;
use codex_code_mode_protocol::StartedCell;
use codex_code_mode_protocol::TerminateRequest;
use codex_code_mode_protocol::WaitOutcome;
use codex_code_mode_protocol::WaitRequest;
use tokio::sync::Semaphore;
@@ -97,6 +97,7 @@ impl CodeModeSession for ReconnectableSession {
Box::pin(async move {
let binding = self.inner.binding_for_request(generation).await?;
let request = WaitRequest {
max_output_tokens: request.max_output_tokens,
cell_id: generation::remote_cell_id(binding.generation, &request.cell_id)?,
yield_time_ms: request.yield_time_ms,
};
@@ -105,15 +106,21 @@ impl CodeModeSession for ReconnectableSession {
})
}
fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome> {
fn terminate<'a>(
&'a self,
request: TerminateRequest,
) -> CodeModeSessionResultFuture<'a, WaitOutcome> {
Box::pin(async move {
// Cleanup must not replace the runtime and hide its loss from the next request.
let binding = self
.inner
.binding()
.ok_or_else(|| SHUTDOWN_ERROR.to_string())?;
let cell_id = generation::remote_cell_id(binding.generation, &cell_id)?;
let outcome = binding.session.terminate(cell_id).await?;
let cell_id = generation::remote_cell_id(binding.generation, &request.cell_id)?;
let outcome = binding
.session
.terminate(TerminateRequest { cell_id, ..request })
.await?;
Ok(generation::public_wait_outcome(binding.generation, outcome))
})
}
codex-rs/code-mode/src/remote_session.rs
+11 -5
View File
@@ -6,7 +6,6 @@ use std::sync::atomic::AtomicBool;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering;
use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::CodeModeSession;
use codex_code_mode_protocol::CodeModeSessionCellExecutionLimits;
use codex_code_mode_protocol::CodeModeSessionDelegate;
@@ -15,6 +14,7 @@ use codex_code_mode_protocol::CodeModeSessionProviderFuture;
use codex_code_mode_protocol::CodeModeSessionResultFuture;
use codex_code_mode_protocol::ExecuteRequest;
use codex_code_mode_protocol::StartedCell;
use codex_code_mode_protocol::TerminateRequest;
use codex_code_mode_protocol::WaitOutcome;
use codex_code_mode_protocol::WaitRequest;
use codex_code_mode_protocol::host::SessionId;
@@ -286,9 +286,12 @@ impl ProcessOwnedCodeModeSession {
.await
}
pub async fn terminate(&self, cell_id: CellId) -> Result<WaitOutcome, String> {
pub async fn terminate(&self, request: TerminateRequest) -> Result<WaitOutcome, String> {
let binding = self.connection().await?;
binding.connection.terminate(binding.remote, cell_id).await
binding
.connection
.terminate(binding.remote, request.cell_id)
.await
}
pub async fn shutdown(&self) -> Result<(), String> {
@@ -549,8 +552,11 @@ impl CodeModeSession for ProcessOwnedCodeModeSession {
Box::pin(ProcessOwnedCodeModeSession::wait(self, request, preempt))
}
fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome> {
Box::pin(ProcessOwnedCodeModeSession::terminate(self, cell_id))
fn terminate<'a>(
&'a self,
request: TerminateRequest,
) -> CodeModeSessionResultFuture<'a, WaitOutcome> {
Box::pin(ProcessOwnedCodeModeSession::terminate(self, request))
}
fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()> {
codex-rs/code-mode/src/remote_session/connection/driver/commands.rs
+3
View File
@@ -112,6 +112,7 @@ impl ConnectionDriver {
session_id: session.id.clone(),
cell_execution_limits: (limits != WireSessionCellExecutionLimits::default())
.then_some(limits),
output_truncation: None,
},
};
let frame = match EncodedFrame::encode(&message) {
@@ -251,6 +252,7 @@ impl ConnectionDriver {
HostRequest::Wait {
session_id: session.id.clone(),
request,
max_output_tokens: None,
},
PendingRequest::Wait {
session,
@@ -284,6 +286,7 @@ impl ConnectionDriver {
HostRequest::Terminate {
session_id: session.id.clone(),
cell_id,
max_output_tokens: None,
},
PendingRequest::Terminate {
session,
codex-rs/code-mode/src/remote_session/connection/driver/responses.rs
+1
View File
@@ -432,6 +432,7 @@ impl ConnectionDriver {
HostRequest::Terminate {
session_id: session.id.clone(),
cell_id: cell_id.clone(),
max_output_tokens: None,
},
PendingRequest::Terminate {
session,
codex-rs/code-mode/src/remote_session/connection/driver_tests.rs
+8
View File
@@ -237,6 +237,7 @@ async fn old_host_does_not_receive_yield_frames() {
WaitRequest {
cell_id: CellId::new("1".to_string()),
yield_time_ms: 60_000,
max_output_tokens: None,
},
Some(signal),
)
@@ -311,6 +312,7 @@ async fn open_session_includes_nondefault_cell_execution_limits() {
max_yield_time_ms: Some(250),
max_heap_size_bytes: Some(16 * 1024 * 1024),
}),
output_truncation: None,
},
}
);
@@ -1275,6 +1277,7 @@ async fn mismatched_wait_response_fails_connection() {
request: WaitRequest {
cell_id: CellId::new("1".to_string()),
yield_time_ms: 1,
max_output_tokens: None,
},
caller_cancellation: CancellationToken::new(),
yield_signal: None,
@@ -1379,6 +1382,7 @@ async fn remote_wait_accepts_durations_longer_than_five_minutes() {
request: WaitRequest {
cell_id: CellId::new("1".to_string()),
yield_time_ms: 300_001,
max_output_tokens: None,
},
caller_cancellation: CancellationToken::new(),
yield_signal: None,
@@ -1448,6 +1452,7 @@ async fn queued_remote_wait_times_out_and_invalidates_the_connection() {
WaitRequest {
cell_id: CellId::new("1".to_string()),
yield_time_ms: 1,
max_output_tokens: None,
},
/*yield_signal*/ None,
)
@@ -1534,6 +1539,7 @@ async fn cancelled_wait_is_retired_before_next_wait_is_sent() {
request: WaitRequest {
cell_id: CellId::new("1".to_string()),
yield_time_ms: 60_000,
max_output_tokens: None,
},
caller_cancellation: first_cancellation.clone(),
yield_signal: None,
@@ -1553,6 +1559,7 @@ async fn cancelled_wait_is_retired_before_next_wait_is_sent() {
request: WaitRequest {
cell_id: CellId::new("1".to_string()),
yield_time_ms: 1,
max_output_tokens: None,
},
caller_cancellation: CancellationToken::new(),
yield_signal: None,
@@ -2005,6 +2012,7 @@ async fn aborting_driver_marks_connection_dead_and_closes_cells() {
request: WaitRequest {
cell_id: CellId::new("1".to_string()),
yield_time_ms: 60_000,
max_output_tokens: None,
},
caller_cancellation: CancellationToken::new(),
yield_signal: None,
codex-rs/core/src/tools/code_mode/mod.rs
+2 -2
View File
@@ -158,7 +158,7 @@ impl CodeModeService {
&self,
cell_id: CellId,
) -> Result<codex_code_mode::WaitOutcome, String> {
self.session().await?.terminate(cell_id).await
self.session().await?.terminate(cell_id.into()).await
}
pub(crate) async fn interrupt_active_cells(&self) {
@@ -170,7 +170,7 @@ impl CodeModeService {
.active_cell_ids()
.into_iter()
.map(|cell_id| async move {
if let Err(error) = session.terminate(cell_id.clone()).await {
if let Err(error) = session.terminate(cell_id.clone().into()).await {
tracing::warn!(%cell_id, %error, "failed to terminate interrupted code-mode cell");
}
}),
codex-rs/core/src/tools/code_mode/wait_handler.rs
+1
View File
@@ -140,6 +140,7 @@ impl CodeModeWaitHandler {
codex_code_mode::WaitRequest {
cell_id: cell_id.clone(),
yield_time_ms: args.yield_time_ms,
max_output_tokens: None,
},
step_context.preempt.clone(),
)
codex-rs/core/src/tools/tool_dispatch_trace_tests.rs
+3 -2
View File
@@ -90,14 +90,15 @@ impl codex_code_mode::CodeModeSession for MissingCellCodeModeSession {
request: codex_code_mode::WaitRequest,
_preempt: Option<CancellationToken>,
) -> codex_code_mode::CodeModeSessionResultFuture<'a, codex_code_mode::WaitOutcome> {
self.terminate(request.cell_id)
self.terminate(request.cell_id.into())
}
fn terminate<'a>(
&'a self,
cell_id: codex_code_mode::CellId,
request: codex_code_mode::TerminateRequest,
) -> codex_code_mode::CodeModeSessionResultFuture<'a, codex_code_mode::WaitOutcome> {
Box::pin(async move {
let cell_id = request.cell_id;
Ok(codex_code_mode::WaitOutcome::MissingCell(
codex_code_mode::RuntimeResponse::Result {
code_mode_host_duration: None,
codex-rs/core/tests/code_mode_tracing.rs
+1
View File
@@ -176,6 +176,7 @@ async fn stdio_callbacks_preserve_context_across_overlapping_yielded_cells() {
WaitRequest {
cell_id,
yield_time_ms: 5_000,
max_output_tokens: None,
},
/*preempt*/ None,
)