Support negotiated server-side truncation in Code Mode clients (#58125)

## Why

Host-side truncation can bound Code Mode output before it crosses the transport, including script errors that would otherwise exceed the frame limit.

## What changed

- Negotiate server-side output truncation for stdio and gRPC sessions, with client-side fallback for older hosts.
- Forward output budgets for execution, waiting, and termination when server-side truncation is selected. Avoid truncating host-formatted output or appending script errors a second time.
- Preserve the negotiated mode across reconnects and reject hosts that cannot maintain it.

## Testing

Add transport and integration coverage for output budgets, oversized script errors, reconnect compatibility, and preservation of output-limit notices with a zero-token budget.

GitOrigin-RevId: 2075101e4845eca64f45491c27a3b326788645c8
This commit is contained in:
Sean Huang
2026-10-29 16:57:12 +00:00
committed by copyberry
parent 3062f7aa4a
commit 09db085f94
21 changed files with 751 additions and 154 deletions
@@ -5,17 +5,21 @@ use std::path::Path;
use std::path::PathBuf;
use std::process::Stdio;
use std::sync::Arc;
use std::sync::OnceLock;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::time::Duration;
use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::CodeModeOutputTruncation;
use codex_code_mode_protocol::CodeModeSessionCellExecutionLimits;
use codex_code_mode_protocol::CodeModeSessionDelegate;
use codex_code_mode_protocol::CodeModeSessionOptions;
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::BOUNDED_OUTPUT_CAPABILITY;
use codex_code_mode_protocol::host::Capability;
use codex_code_mode_protocol::host::CapabilitySet;
use codex_code_mode_protocol::host::ClientHello;
@@ -209,9 +213,14 @@ impl Connection {
.map_err(|error| error.to_string())?;
let yield_capability =
Capability::new(YIELD_OBSERVATION_CAPABILITY).map_err(|error| error.to_string())?;
let optional_capabilities =
CapabilitySet::try_new([session_limits_capability, yield_capability])
.map_err(|error| error.to_string())?;
let output_capability =
Capability::new(BOUNDED_OUTPUT_CAPABILITY).map_err(|error| error.to_string())?;
let optional_capabilities = CapabilitySet::try_new([
session_limits_capability,
yield_capability,
output_capability,
])
.map_err(|error| error.to_string())?;
let hello = ClientHello::new(
SupportedProtocolVersions::try_new([ProtocolVersion::V1])
.map_err(|err| err.to_string())?,
@@ -326,27 +335,42 @@ impl Connection {
self.alive.load(Ordering::Acquire)
}
fn supports(&self, capability: &str) -> bool {
self.capabilities
.iter()
.any(|value| value.as_str() == capability)
}
pub(super) async fn open_session(
&self,
session: RemoteSession,
limits: CodeModeSessionCellExecutionLimits,
mut options: CodeModeSessionOptions,
pinned: &OnceLock<CodeModeOutputTruncation>,
) -> Result<SessionCleanup, String> {
if limits != CodeModeSessionCellExecutionLimits::default()
&& !self
.capabilities
.iter()
.any(|capability| capability.as_str() == SESSION_RESOURCE_LIMITS_CAPABILITY)
if options.cell_execution_limits != CodeModeSessionCellExecutionLimits::default()
&& !self.supports(SESSION_RESOURCE_LIMITS_CAPABILITY)
{
return Err(format!(
"code-mode host does not support session resource limits: missing `{SESSION_RESOURCE_LIMITS_CAPABILITY}` capability"
));
}
let pinned_mode = pinned.get().copied();
let output_truncation = if self.supports(BOUNDED_OUTPUT_CAPABILITY) {
pinned_mode.unwrap_or(options.output_truncation)
} else if pinned_mode == Some(CodeModeOutputTruncation::Server) {
return Err(
"code-mode host cannot preserve server output truncation after reconnect".into(),
);
} else {
CodeModeOutputTruncation::Client
};
options.output_truncation = output_truncation;
let cleanup = SessionCleanup::new();
let cancellation = CallerCancellation::new();
let (response_tx, response_rx) = oneshot::channel();
self.send(DriverCommand::OpenSession {
session,
limits,
options,
cleanup: cleanup.clone(),
caller_cancellation: cancellation.token(),
response_tx,
@@ -355,6 +379,7 @@ impl Connection {
let result = self.receive(response_rx).await;
cancellation.disarm();
result?;
pinned.get_or_init(|| output_truncation);
Ok(cleanup)
}
@@ -424,9 +449,7 @@ impl Connection {
&self,
signal: Option<CancellationToken>,
) -> Option<CancellationToken> {
self.capabilities
.iter()
.any(|capability| capability.as_str() == YIELD_OBSERVATION_CAPABILITY)
self.supports(YIELD_OBSERVATION_CAPABILITY)
.then_some(signal)
.flatten()
}
@@ -434,13 +457,13 @@ impl Connection {
pub(super) async fn terminate(
&self,
session: RemoteSession,
cell_id: CellId,
request: TerminateRequest,
) -> Result<WaitOutcome, String> {
let (response_tx, response_rx) = oneshot::channel();
self.with_transport_deadline(Duration::ZERO, "terminate", async {
self.send(DriverCommand::Terminate {
session,
cell_id,
request,
response_tx,
})
.await?;
@@ -1,14 +1,16 @@
use std::sync::Arc;
use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::CodeModeSessionCellExecutionLimits;
use codex_code_mode_protocol::CodeModeOutputTruncation;
use codex_code_mode_protocol::CodeModeSessionDelegate;
use codex_code_mode_protocol::CodeModeSessionOptions;
use codex_code_mode_protocol::ExecuteRequest;
use codex_code_mode_protocol::TerminateRequest;
use codex_code_mode_protocol::WaitOutcome;
use codex_code_mode_protocol::WaitRequest;
use codex_code_mode_protocol::host::ClientToHost;
use codex_code_mode_protocol::host::EncodedFrame;
use codex_code_mode_protocol::host::HostRequest;
use codex_code_mode_protocol::host::WireOutputTruncation;
use codex_code_mode_protocol::host::WireSessionCellExecutionLimits;
use codex_code_mode_protocol::host::WireWaitRequest;
use tokio::sync::oneshot;
@@ -30,11 +32,11 @@ impl ConnectionDriver {
match command {
DriverCommand::OpenSession {
session,
limits,
options,
cleanup,
caller_cancellation,
response_tx,
} => self.open_session(session, limits, cleanup, caller_cancellation, response_tx),
} => self.open_session(session, options, cleanup, caller_cancellation, response_tx),
DriverCommand::Execute {
session,
request,
@@ -65,9 +67,9 @@ impl ConnectionDriver {
),
DriverCommand::Terminate {
session,
cell_id,
request,
response_tx,
} => self.terminate(session, cell_id, response_tx),
} => self.terminate(session, request, response_tx),
DriverCommand::ShutdownSession {
session,
response_tx,
@@ -78,7 +80,7 @@ impl ConnectionDriver {
fn open_session(
&mut self,
session: RemoteSession,
limits: CodeModeSessionCellExecutionLimits,
options: CodeModeSessionOptions,
cleanup: super::cleanup::SessionCleanup,
caller_cancellation: CancellationToken,
response_tx: oneshot::Sender<Result<(), String>>,
@@ -90,7 +92,7 @@ impl ConnectionDriver {
)));
return true;
}
let limits = match WireSessionCellExecutionLimits::try_from(limits) {
let limits = match WireSessionCellExecutionLimits::try_from(options.cell_execution_limits) {
Ok(limits) => limits,
Err(error) => {
let _ = response_tx.send(Err(format!(
@@ -112,7 +114,8 @@ impl ConnectionDriver {
session_id: session.id.clone(),
cell_execution_limits: (limits != WireSessionCellExecutionLimits::default())
.then_some(limits),
output_truncation: None,
output_truncation: (options.output_truncation == CodeModeOutputTruncation::Server)
.then_some(WireOutputTruncation::Server),
},
};
let frame = match EncodedFrame::encode(&message) {
@@ -213,6 +216,7 @@ impl ConnectionDriver {
let _ = response_tx.send(Err(err));
return true;
}
let max_output_tokens = request.max_output_tokens;
let request = match remote_wait_request(&session, request) {
Ok(request) => request,
Err(err) => {
@@ -224,6 +228,7 @@ impl ConnectionDriver {
self.requests.push_deferred_wait(DeferredWait {
session,
request,
max_output_tokens,
caller_cancellation,
yield_signal,
response_tx,
@@ -233,6 +238,7 @@ impl ConnectionDriver {
self.start_wait(
session,
request,
max_output_tokens,
caller_cancellation,
yield_signal,
response_tx,
@@ -243,6 +249,7 @@ impl ConnectionDriver {
&mut self,
session: RemoteSession,
request: WireWaitRequest,
max_output_tokens: Option<usize>,
caller_cancellation: CancellationToken,
yield_signal: Option<CancellationToken>,
response_tx: oneshot::Sender<Result<WaitOutcome, String>>,
@@ -252,7 +259,7 @@ impl ConnectionDriver {
HostRequest::Wait {
session_id: session.id.clone(),
request,
max_output_tokens: None,
max_output_tokens,
},
PendingRequest::Wait {
session,
@@ -267,14 +274,14 @@ impl ConnectionDriver {
fn terminate(
&mut self,
session: RemoteSession,
cell_id: CellId,
request: TerminateRequest,
response_tx: oneshot::Sender<Result<WaitOutcome, String>>,
) -> bool {
if let Err(err) = self.sessions.require_ready(&session) {
let _ = response_tx.send(Err(err));
return true;
}
let cell_id = match remote_cell_id(&session, &cell_id) {
let cell_id = match remote_cell_id(&session, &request.cell_id) {
Ok(cell_id) => cell_id,
Err(err) => {
let _ = response_tx.send(Err(err));
@@ -286,7 +293,7 @@ impl ConnectionDriver {
HostRequest::Terminate {
session_id: session.id.clone(),
cell_id,
max_output_tokens: None,
max_output_tokens: request.max_output_tokens,
},
PendingRequest::Terminate {
session,
@@ -41,6 +41,7 @@ impl ConnectionDriver {
if !self.start_wait(
wait.session,
wait.request,
wait.max_output_tokens,
wait.caller_cancellation,
wait.yield_signal,
wait.response_tx,
@@ -1,11 +1,11 @@
use std::sync::Arc;
use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::CodeModeSessionCellExecutionLimits;
use codex_code_mode_protocol::CodeModeSessionDelegate;
use codex_code_mode_protocol::CodeModeSessionOptions;
use codex_code_mode_protocol::ExecuteRequest;
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::host::DelegateRequestId;
@@ -31,7 +31,7 @@ pub(in crate::remote_session) struct RemoteSession {
pub(in crate::remote_session::connection) enum DriverCommand {
OpenSession {
session: RemoteSession,
limits: CodeModeSessionCellExecutionLimits,
options: CodeModeSessionOptions,
cleanup: SessionCleanup,
caller_cancellation: CancellationToken,
response_tx: oneshot::Sender<Result<(), String>>,
@@ -53,7 +53,7 @@ pub(in crate::remote_session::connection) enum DriverCommand {
},
Terminate {
session: RemoteSession,
cell_id: CellId,
request: TerminateRequest,
response_tx: oneshot::Sender<Result<WaitOutcome, String>>,
},
ShutdownSession {
@@ -200,6 +200,7 @@ pub(super) enum PendingRequest {
pub(super) struct DeferredWait {
pub(super) session: RemoteSession,
pub(super) request: WireWaitRequest,
pub(super) max_output_tokens: Option<usize>,
pub(super) caller_cancellation: CancellationToken,
pub(super) yield_signal: Option<CancellationToken>,
pub(super) response_tx: oneshot::Sender<Result<WaitOutcome, String>>,
@@ -7,12 +7,18 @@ use std::time::Duration;
use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::CodeModeNestedToolCall;
use codex_code_mode_protocol::CodeModeOutputTruncation;
use codex_code_mode_protocol::CodeModeSession;
use codex_code_mode_protocol::CodeModeSessionCellExecutionLimits;
use codex_code_mode_protocol::CodeModeSessionDelegate;
use codex_code_mode_protocol::CodeModeSessionOptions;
use codex_code_mode_protocol::ExecuteRequest;
use codex_code_mode_protocol::NotificationFuture;
use codex_code_mode_protocol::TerminateRequest;
use codex_code_mode_protocol::ToolInvocationFuture;
use codex_code_mode_protocol::WaitRequest;
use codex_code_mode_protocol::host::BOUNDED_OUTPUT_CAPABILITY;
use codex_code_mode_protocol::host::Capability;
use codex_code_mode_protocol::host::CapabilitySet;
use codex_code_mode_protocol::host::ClientToHost;
use codex_code_mode_protocol::host::DelegateRequest;
@@ -31,6 +37,7 @@ use codex_code_mode_protocol::host::WireSessionCellExecutionLimits;
use codex_code_mode_protocol::host::WireWaitOutcome;
use codex_protocol::ToolName;
use pretty_assertions::assert_eq;
use serde_json::json;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tokio_util::sync::CancellationToken;
@@ -43,6 +50,8 @@ use super::DriverEvent;
use super::DriverLifecycle;
use super::RemoteSession;
use super::SessionCleanup;
use crate::remote_session::OwnedCodeModeHost;
use crate::remote_session::ProcessOwnedCodeModeSession;
struct DriverHarness {
command_tx: mpsc::Sender<DriverCommand>,
@@ -93,26 +102,15 @@ impl DriverHarness {
self.command_tx
.send(DriverCommand::OpenSession {
session: session.clone(),
limits: Default::default(),
options: CodeModeSessionOptions::default(),
cleanup: cleanup.clone(),
caller_cancellation: CancellationToken::new(),
response_tx,
})
.await
.expect("open command");
self.outgoing_rx.recv().await.expect("open frame");
self.event_tx
.send(DriverEvent::HostMessage(HostToClient::Response {
id: RequestId::new(/*value*/ 1),
result: WireResult::Ok {
value: HostResponse::SessionReady {
session_id: session.id,
},
},
}))
.await
.expect("open response");
self.accept_open(session.id, CodeModeOutputTruncation::Client)
.await;
response_rx
.await
.expect("open reply")
@@ -120,6 +118,30 @@ impl DriverHarness {
cleanup
}
async fn accept_open(&mut self, session_id: SessionId, mode: CodeModeOutputTruncation) {
let mut expected = json!({ "method": "session/open", "sessionId": session_id });
if mode == CodeModeOutputTruncation::Server {
expected["outputTruncation"] = json!("server");
}
assert_eq!(self.next_request().await, expected);
self.event_tx
.send(DriverEvent::HostMessage(HostToClient::Response {
id: RequestId::new(/*value*/ 1),
result: WireResult::Ok {
value: HostResponse::SessionReady { session_id },
},
}))
.await
.expect("open response");
}
async fn next_request(&mut self) -> serde_json::Value {
let frame = self.outgoing_rx.recv().await.expect("request");
let mut message =
EncodedFrame::decode_framed::<serde_json::Value>(&frame.into_framed_bytes()).unwrap();
message["request"].take()
}
async fn start_cell(
&mut self,
session: RemoteSession,
@@ -167,6 +189,17 @@ impl DriverHarness {
delivered.started
}
fn connection(&self, capabilities: CapabilitySet) -> Connection {
Connection {
command_tx: self.command_tx.clone(),
execute_claim_tx: self.execute_claim_tx.clone(),
alive: Arc::clone(&self.alive),
failure: Arc::clone(&self.failure),
cancellation: self.cancellation.clone(),
capabilities,
}
}
async fn start_tool_delegate(&self, session: &RemoteSession, id: DelegateRequestId) {
self.event_tx
.send(DriverEvent::HostMessage(HostToClient::DelegateRequest {
@@ -220,14 +253,7 @@ async fn old_host_does_not_receive_yield_frames() {
}))
.await
.unwrap();
let connection = Connection {
command_tx: harness.command_tx.clone(),
execute_claim_tx: harness.execute_claim_tx.clone(),
alive: Arc::clone(&harness.alive),
failure: Arc::clone(&harness.failure),
cancellation: harness.cancellation.clone(),
capabilities: CapabilitySet::empty(),
};
let connection = harness.connection(CapabilitySet::empty());
let signal = CancellationToken::new();
signal.cancel();
let wait = tokio::spawn(async move {
@@ -287,8 +313,10 @@ async fn open_session_includes_nondefault_cell_execution_limits() {
.command_tx
.send(DriverCommand::OpenSession {
session: session.clone(),
limits,
options: CodeModeSessionOptions {
cell_execution_limits: limits,
..CodeModeSessionOptions::default()
},
cleanup: SessionCleanup::new(),
caller_cancellation: CancellationToken::new(),
response_tx,
@@ -318,6 +346,80 @@ async fn open_session_includes_nondefault_cell_execution_limits() {
);
}
#[tokio::test]
async fn stdio_logical_session_pins_output_mode_across_reconnect() {
use CodeModeOutputTruncation::Client;
use CodeModeOutputTruncation::Server;
let legacy = CapabilitySet::empty();
let bounded =
CapabilitySet::try_new([Capability::new(BOUNDED_OUTPUT_CAPABILITY).unwrap()]).unwrap();
for (requested, hosts, expected_mode) in [
(Client, [&bounded, &bounded].as_slice(), Client),
(Server, [&legacy, &bounded].as_slice(), Client),
(Server, [&bounded, &bounded, &legacy].as_slice(), Server),
] {
let host = Arc::new(OwnedCodeModeHost::new("unused".into()));
let session = ProcessOwnedCodeModeSession::with_host(
Arc::clone(&host),
CodeModeSessionOptions {
output_truncation: requested,
..CodeModeSessionOptions::default()
},
);
for (index, &capabilities) in hosts.iter().enumerate() {
let mut harness = DriverHarness::start();
*host.connection.lock().unwrap() =
Some(Arc::new(harness.connection(capabilities.clone())));
if expected_mode == Server && capabilities == &legacy {
tokio::select! {
result = session.connection() => assert!(matches!(result, Err(error) if error.contains("cannot preserve"))),
_ = harness.outgoing_rx.recv() => panic!("opened a server session on an incompatible host"),
}
assert_eq!(session.output_truncation(), Server);
break;
}
let session_id = SessionId::new(format!("session-{}", index + 1)).unwrap();
tokio::try_join!(session.connection(), async {
harness.accept_open(session_id, expected_mode).await;
Ok(())
})
.unwrap();
assert_eq!(session.output_truncation(), expected_mode);
if capabilities == &legacy {
let cell_id = CellId::new("1".into());
let wait = WaitRequest {
cell_id: cell_id.clone(),
yield_time_ms: 1,
max_output_tokens: Some(5),
};
let terminate = TerminateRequest {
cell_id,
max_output_tokens: Some(5),
};
let legacy_host = async {
for method in ["session/wait", "session/terminate"] {
let request = harness.next_request().await;
assert_eq!(request["method"], method);
// A null value would also be rejected by a legacy host.
assert!(request.get("maxOutputTokens").is_none());
}
harness.cancellation.cancel();
};
let (wait, terminate, ()) = tokio::join!(biased;
session.wait(wait, /*preempt*/ None), session.terminate(terminate), legacy_host
);
assert!(wait.is_err() && terminate.is_err());
}
harness.cancellation.cancel();
(&mut harness.driver_task).await.unwrap();
}
}
}
#[derive(Default)]
struct RecordingDelegate {
closed_cells: StdMutex<Vec<CellId>>,
@@ -626,8 +728,7 @@ async fn dropped_open_waiter_shuts_down_committed_session() {
.command_tx
.send(DriverCommand::OpenSession {
session: session.clone(),
limits: Default::default(),
options: CodeModeSessionOptions::default(),
cleanup,
caller_cancellation: CancellationToken::new(),
response_tx: open_tx,
@@ -857,7 +958,7 @@ async fn terminate_closes_cell_without_waiting_for_delegate_cleanup() {
.command_tx
.send(DriverCommand::Terminate {
session: session.clone(),
cell_id: CellId::new("1".to_string()),
request: CellId::new("1".to_string()).into(),
response_tx,
})
.await
@@ -1330,7 +1431,7 @@ async fn mismatched_terminate_response_fails_connection() {
.command_tx
.send(DriverCommand::Terminate {
session,
cell_id: CellId::new("1".to_string()),
request: CellId::new("1".to_string()).into(),
response_tx,
})
.await
@@ -1498,7 +1599,7 @@ async fn queued_remote_termination_times_out_and_invalidates_the_connection() {
};
let response = tokio::spawn(async move {
let result = connection
.terminate(session, CellId::new("1".to_string()))
.terminate(session, CellId::new("1".to_string()).into())
.await;
(connection, result)
});
@@ -1539,7 +1640,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,
max_output_tokens: Some(19),
},
caller_cancellation: first_cancellation.clone(),
yield_signal: None,
@@ -1559,7 +1660,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,
max_output_tokens: Some(37),
},
caller_cancellation: CancellationToken::new(),
yield_signal: None,
@@ -1587,7 +1688,7 @@ async fn cancelled_wait_is_retired_before_next_wait_is_sent() {
}))
.await
.expect("cancelled wait response");
harness.outgoing_rx.recv().await.expect("second wait frame");
assert_eq!(harness.next_request().await["maxOutputTokens"], 37);
harness
.event_tx
.send(DriverEvent::HostMessage(HostToClient::Response {