openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-28 18:58:46 +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/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,