openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-31 19:28:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
## 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
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/code-mode/src/grpc_session/callbacks_tests.rs
+2
View File
@@ -6,6 +6,7 @@ use std::sync::Mutex;
use std::sync::atomic::AtomicBool;
use std::time::Duration;
use codex_code_mode_protocol::CodeModeOutputTruncation;
use http_body_util::Full;
use pretty_assertions::assert_eq;
use tokio::sync::Notify;
@@ -48,6 +49,7 @@ async fn missing_invocation_completion_racing_cancellation_preserves_session() {
fn session(client: GrpcClient) -> Arc<SessionInner> {
Arc::new(SessionInner {
id: "session".to_string(),
output_truncation: CodeModeOutputTruncation::Client,
client: client.clone(),
route: None,
runtime: tokio::runtime::Handle::current(),
codex-rs/code-mode/src/grpc_session/conversion.rs
+14
View File
@@ -2,6 +2,7 @@ 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::CodeModeToolKind;
use codex_code_mode_protocol::ExecuteRequest;
use codex_code_mode_protocol::FunctionCallOutputContentItem;
@@ -12,6 +13,19 @@ use codex_code_mode_protocol::WaitOutcome;
use codex_code_mode_protocol::grpc;
use codex_protocol::ToolName;
pub(super) fn output_budget(
mode: CodeModeOutputTruncation,
max_output_tokens: Option<usize>,
) -> Result<Option<u64>, String> {
match mode {
CodeModeOutputTruncation::Client => Ok(None),
CodeModeOutputTruncation::Server => max_output_tokens
.map(u64::try_from)
.transpose()
.map_err(|error| format!("invalid code-mode output token limit: {error}")),
}
}
pub(super) fn execute_request(
session_id: &str,
execution_id: String,
codex-rs/code-mode/src/grpc_session/mod.rs
+53 -4
View File
@@ -10,9 +10,11 @@ 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::CodeModeSession;
use codex_code_mode_protocol::CodeModeSessionCellExecutionLimits;
use codex_code_mode_protocol::CodeModeSessionDelegate;
use codex_code_mode_protocol::CodeModeSessionOptions;
use codex_code_mode_protocol::CodeModeSessionProvider;
use codex_code_mode_protocol::CodeModeSessionProviderFuture;
use codex_code_mode_protocol::CodeModeSessionResultFuture;
@@ -129,9 +131,12 @@ impl GrpcCodeModeSessionProvider {
#[tracing::instrument(name = "code_mode.grpc.open_binding", level = "info", skip_all)]
async fn open_binding(
&self,
limits: CodeModeSessionCellExecutionLimits,
options: &CodeModeSessionOptions,
pinned_mode: Option<CodeModeOutputTruncation>,
) -> Result<Arc<GrpcCodeModeSession>, String> {
let mut client = deadline::startup("transport connection", self.transport.client()).await?;
let requested_mode = pinned_mode.unwrap_or(options.output_truncation);
let limits = &options.cell_execution_limits;
let limits = grpc::SessionCellExecutionLimits {
max_yield_time_ms: limits.max_yield_time_ms,
max_heap_size_bytes: limits
@@ -149,7 +154,12 @@ impl GrpcCodeModeSessionProvider {
loop {
let mut request = tonic::Request::new(grpc::OpenSessionRequest {
cell_execution_limits,
output_truncation: None,
output_truncation: match requested_mode {
CodeModeOutputTruncation::Client => None,
CodeModeOutputTruncation::Server => {
Some(grpc::OutputTruncation::Server.into())
}
},
});
inject_span_traceparent(&mut request, &open_session_span);
let result = async {
@@ -197,9 +207,33 @@ impl GrpcCodeModeSessionProvider {
return Err("gRPC code-mode session lease omitted its opening event".to_string());
};
validate_identifier(&opened.session_id, "session ID")?;
let host_mode =
grpc::OutputTruncation::try_from(opened.output_truncation.unwrap_or_default())
.map_err(|_| "gRPC code-mode host returned an unknown output truncation mode")?;
let output_truncation = match host_mode {
grpc::OutputTruncation::Unspecified | grpc::OutputTruncation::Client => {
CodeModeOutputTruncation::Client
}
grpc::OutputTruncation::Server => CodeModeOutputTruncation::Server,
};
if output_truncation == CodeModeOutputTruncation::Server
&& requested_mode == CodeModeOutputTruncation::Client
{
return Err(format!(
"gRPC code-mode host selected {output_truncation:?} output truncation, requested {requested_mode:?}"
));
}
if pinned_mode.is_some_and(|previous| previous != output_truncation) {
// Dropping the opening lease releases this incompatible session before subscribing.
return Err(
"gRPC code-mode host changed the session output truncation mode on reconnect"
.to_string(),
);
}
let inner = Arc::new(SessionInner {
id: opened.session_id,
output_truncation,
client,
route,
runtime: tokio::runtime::Handle::current(),
@@ -259,9 +293,19 @@ impl CodeModeSessionProvider for GrpcCodeModeSessionProvider {
fn create_session_with_limits<'a>(
&'a self,
limits: CodeModeSessionCellExecutionLimits,
) -> CodeModeSessionProviderFuture<'a> {
self.create_session_with_options(CodeModeSessionOptions {
cell_execution_limits: limits,
..CodeModeSessionOptions::default()
})
}
fn create_session_with_options<'a>(
&'a self,
options: CodeModeSessionOptions,
) -> CodeModeSessionProviderFuture<'a> {
Box::pin(async move {
let session = Arc::new(reconnect::ReconnectableSession::new(self.clone(), limits));
let session = Arc::new(reconnect::ReconnectableSession::new(self.clone(), options));
session.initialize().await?;
Ok(session as _)
})
@@ -290,6 +334,10 @@ impl Drop for OpeningSession {
}
impl CodeModeSession for GrpcCodeModeSession {
fn output_truncation(&self) -> CodeModeOutputTruncation {
self.inner.output_truncation
}
fn execute<'a>(
&'a self,
request: ExecuteRequest,
@@ -311,7 +359,7 @@ impl CodeModeSession for GrpcCodeModeSession {
&'a self,
request: TerminateRequest,
) -> CodeModeSessionResultFuture<'a, WaitOutcome> {
Box::pin(self.inner.terminate(request.cell_id))
Box::pin(self.inner.terminate(request))
}
fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()> {
@@ -327,6 +375,7 @@ impl Drop for GrpcCodeModeSession {
pub(super) struct SessionInner {
pub(super) id: String,
output_truncation: CodeModeOutputTruncation,
pub(super) client: GrpcClient,
// Affinity belongs to the session: parents and children may share a transport
// while being admitted to different hosts. The token is opaque to Codex.
codex-rs/code-mode/src/grpc_session/operations.rs
+13 -6
View File
@@ -10,6 +10,7 @@ use codex_code_mode_protocol::DEFAULT_EXEC_YIELD_TIME_MS;
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::grpc;
@@ -58,7 +59,7 @@ impl Drop for ExecutionOwnership {
}
let session = Arc::clone(&self.session);
self.session.runtime.spawn(async move {
if let Err(error) = session.terminate(cell_id).await
if let Err(error) = session.terminate(cell_id.into()).await
&& !session.stopped.is_cancelled()
{
debug!("abandoned code-mode execution termination raced closure: {error}");
@@ -313,7 +314,10 @@ 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,
max_output_tokens: conversion::output_budget(
self.output_truncation,
request.max_output_tokens,
)?,
};
let mut client = self.client();
let watcher = preempt.map(|signal| {
@@ -359,7 +363,7 @@ impl SessionInner {
}))
}
pub(super) async fn terminate(&self, cell_id: CellId) -> Result<WaitOutcome, String> {
pub(super) async fn terminate(&self, request: TerminateRequest) -> Result<WaitOutcome, String> {
self.require_open()?;
let mut client = self.client();
let response = deadline::session_request(
@@ -368,13 +372,16 @@ impl SessionInner {
Duration::ZERO,
client.terminate(self.request(grpc::TerminateRequest {
session_id: self.id.clone(),
cell_id: cell_id.as_str().to_string(),
max_output_tokens: None,
cell_id: request.cell_id.as_str().to_string(),
max_output_tokens: conversion::output_budget(
self.output_truncation,
request.max_output_tokens,
)?,
})),
)
.await?;
let outcome = conversion::wait_outcome(response)?;
self.validate_wait_cell(&cell_id, outcome).await
self.validate_wait_cell(&request.cell_id, outcome).await
}
async fn validate_wait_cell(
codex-rs/code-mode/src/grpc_session/reconnect.rs
+21 -5
View File
@@ -1,11 +1,13 @@
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::OnceLock;
use std::sync::PoisonError;
use std::sync::atomic::Ordering;
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::CodeModeSessionResultFuture;
use codex_code_mode_protocol::ExecuteRequest;
use codex_code_mode_protocol::StartedCell;
@@ -30,8 +32,11 @@ pub(super) struct ReconnectableSession {
struct ReconnectInner {
provider: GrpcCodeModeSessionProvider,
limits: CodeModeSessionCellExecutionLimits,
options: CodeModeSessionOptions,
binding: Mutex<Option<SessionBinding>>,
// The negotiated mode belongs to the logical session, including after
// shutdown releases the binding. Never pin a mode from an unpublished open.
output_truncation: OnceLock<CodeModeOutputTruncation>,
opening_permit: Semaphore,
shutdown_requested: CancellationToken,
shutdown_result: Mutex<Option<ShutdownResultReceiver>>,
@@ -46,13 +51,14 @@ struct SessionBinding {
impl ReconnectableSession {
pub(super) fn new(
provider: GrpcCodeModeSessionProvider,
limits: CodeModeSessionCellExecutionLimits,
options: CodeModeSessionOptions,
) -> Self {
Self {
inner: Arc::new(ReconnectInner {
provider,
limits,
options,
binding: Mutex::new(None),
output_truncation: OnceLock::new(),
opening_permit: Semaphore::new(/*permits*/ 1),
shutdown_requested: CancellationToken::new(),
shutdown_result: Mutex::new(None),
@@ -66,6 +72,14 @@ impl ReconnectableSession {
}
impl CodeModeSession for ReconnectableSession {
fn output_truncation(&self) -> CodeModeOutputTruncation {
self.inner
.output_truncation
.get()
.copied()
.unwrap_or_default()
}
fn execute<'a>(
&'a self,
request: ExecuteRequest,
@@ -184,7 +198,7 @@ impl ReconnectInner {
_ = self.shutdown_requested.cancelled() => {
return Err(SHUTDOWN_ERROR.to_string());
}
session = self.provider.open_binding(self.limits.clone()) => session?,
session = self.provider.open_binding(&self.options, self.output_truncation.get().copied()) => session?,
};
let binding = SessionBinding {
session,
@@ -196,6 +210,8 @@ impl ReconnectInner {
false
} else {
*current = Some(binding.clone());
self.output_truncation
.get_or_init(|| binding.session.output_truncation());
true
}
};
codex-rs/code-mode/src/remote_session.rs
+50 -15
View File
@@ -2,13 +2,16 @@ use std::io;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::Mutex as StdMutex;
use std::sync::OnceLock;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering;
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::CodeModeSessionProvider;
use codex_code_mode_protocol::CodeModeSessionProviderFuture;
use codex_code_mode_protocol::CodeModeSessionResultFuture;
@@ -83,7 +86,17 @@ impl CodeModeSessionProvider for ProcessOwnedCodeModeSessionProvider {
&'a self,
limits: CodeModeSessionCellExecutionLimits,
) -> CodeModeSessionProviderFuture<'a> {
Box::pin(create_host_session(self.process_host(), limits))
self.create_session_with_options(CodeModeSessionOptions {
cell_execution_limits: limits,
..CodeModeSessionOptions::default()
})
}
fn create_session_with_options<'a>(
&'a self,
options: CodeModeSessionOptions,
) -> CodeModeSessionProviderFuture<'a> {
Box::pin(create_host_session(self.process_host(), options))
}
}
@@ -106,9 +119,9 @@ impl CodeModeSessionProvider for DisabledCodeModeSessionProvider {
async fn create_host_session(
host: Arc<OwnedCodeModeHost>,
limits: CodeModeSessionCellExecutionLimits,
options: CodeModeSessionOptions,
) -> Result<Arc<dyn CodeModeSession>, String> {
let session = ProcessOwnedCodeModeSession::with_host(host, limits);
let session = ProcessOwnedCodeModeSession::with_host(host, options);
session.connection().await?;
Ok(Arc::new(session))
}
@@ -218,7 +231,9 @@ struct SessionBinding {
struct SessionInner {
host: Arc<OwnedCodeModeHost>,
limits: CodeModeSessionCellExecutionLimits,
options: CodeModeSessionOptions,
// The first successful open fixes the output policy for all generations of this session.
output_truncation: OnceLock<CodeModeOutputTruncation>,
state: StdMutex<SessionState>,
shutdown_requested: AtomicBool,
shutdown_result: StdMutex<Option<ShutdownResultReceiver>>,
@@ -236,15 +251,16 @@ impl ProcessOwnedCodeModeSession {
Arc::new(OwnedCodeModeHost::new(
InstallContext::current().code_mode_host_program(),
)),
CodeModeSessionCellExecutionLimits::default(),
CodeModeSessionOptions::default(),
)
}
fn with_host(host: Arc<OwnedCodeModeHost>, limits: CodeModeSessionCellExecutionLimits) -> Self {
fn with_host(host: Arc<OwnedCodeModeHost>, options: CodeModeSessionOptions) -> Self {
Self {
inner: Arc::new(SessionInner {
host,
limits,
options,
output_truncation: OnceLock::new(),
state: StdMutex::new(SessionState::New),
shutdown_requested: AtomicBool::new(false),
shutdown_result: StdMutex::new(None),
@@ -257,9 +273,13 @@ impl ProcessOwnedCodeModeSession {
self.inner.connection().await
}
fn server_output_budget(&self, budget: Option<usize>) -> Option<usize> {
budget.filter(|_| self.output_truncation() == CodeModeOutputTruncation::Server)
}
pub async fn execute(
&self,
request: ExecuteRequest,
mut request: ExecuteRequest,
delegate: Arc<dyn CodeModeSessionDelegate>,
preempt: Option<CancellationToken>,
) -> Result<StartedCell, String> {
@@ -268,6 +288,10 @@ impl ProcessOwnedCodeModeSession {
delegate,
trace: codex_otel::current_span_w3c_trace_context(),
});
// V1 encodes the exec budget as i32. Its maximum already exceeds the host collector limit.
request.max_output_tokens = self
.server_output_budget(request.max_output_tokens)
.map(|tokens| tokens.min(i32::MAX as usize));
binding
.connection
.execute(binding.remote, request, delegate, preempt)
@@ -276,22 +300,21 @@ impl ProcessOwnedCodeModeSession {
pub async fn wait(
&self,
request: WaitRequest,
mut request: WaitRequest,
preempt: Option<CancellationToken>,
) -> Result<WaitOutcome, String> {
let binding = self.connection().await?;
request.max_output_tokens = self.server_output_budget(request.max_output_tokens);
binding
.connection
.wait(binding.remote, request, preempt)
.await
}
pub async fn terminate(&self, request: TerminateRequest) -> Result<WaitOutcome, String> {
pub async fn terminate(&self, mut request: TerminateRequest) -> Result<WaitOutcome, String> {
let binding = self.connection().await?;
binding
.connection
.terminate(binding.remote, request.cell_id)
.await
request.max_output_tokens = self.server_output_budget(request.max_output_tokens);
binding.connection.terminate(binding.remote, request).await
}
pub async fn shutdown(&self) -> Result<(), String> {
@@ -352,7 +375,11 @@ impl SessionInner {
let result = match self.host.connection().await {
Ok(connection) => {
let cleanup = connection
.open_session(remote.clone(), self.limits.clone())
.open_session(
remote.clone(),
self.options.clone(),
&self.output_truncation,
)
.await;
cleanup.map(|cleanup| SessionBinding {
connection,
@@ -529,6 +556,14 @@ impl Default for ProcessOwnedCodeModeSession {
}
impl CodeModeSession for ProcessOwnedCodeModeSession {
fn output_truncation(&self) -> CodeModeOutputTruncation {
self.inner
.output_truncation
.get()
.copied()
.unwrap_or_default()
}
fn execute<'a>(
&'a self,
request: ExecuteRequest,
codex-rs/code-mode/src/remote_session/connection.rs
+39 -16
View File
@@ -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?;
codex-rs/code-mode/src/remote_session/connection/driver/commands.rs
+20 -13
View File
@@ -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,
codex-rs/code-mode/src/remote_session/connection/driver/responses.rs
+1
View File
@@ -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,
codex-rs/code-mode/src/remote_session/connection/driver/types.rs
+5 -4
View File
@@ -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>>,
codex-rs/code-mode/src/remote_session/connection/driver_tests.rs
+133 -32
View File
@@ -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 {