openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-30 19:18: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-host/tests/grpc/session_recovery_tests.rs
+139
View File
@@ -6,9 +6,13 @@ use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use anyhow::Result;
use codex_code_mode::CodeModeOutputTruncation;
use codex_code_mode::CodeModeSessionOptions;
use codex_code_mode::CodeModeSessionProvider;
use codex_code_mode::FunctionCallOutputContentItem;
use codex_code_mode::GrpcCodeModeSessionProvider;
use codex_code_mode::NoopCodeModeSessionDelegate;
use codex_code_mode::RuntimeResponse;
use codex_code_mode_host::GrpcCodeModeHost;
use codex_code_mode_protocol::grpc::code_mode_host_server::CodeModeHostServer;
use pretty_assertions::assert_eq;
@@ -56,6 +60,10 @@ async fn only_missing_sessions_reopen_without_replaying_execution() -> Result<()
.create_session()
.await
.map_err(anyhow::Error::msg)?;
assert_eq!(
session.output_truncation(),
CodeModeOutputTruncation::Client
);
execute(
&session,
request(r#"store("before", 1);"#),
@@ -121,3 +129,134 @@ async fn only_missing_sessions_reopen_without_replaying_execution() -> Result<()
}
Ok(())
}
#[tokio::test]
async fn grpc_session_pins_output_mode_across_recovery_and_shutdown() -> Result<()> {
use CodeModeOutputTruncation::Client;
use CodeModeOutputTruncation::Server;
use axum::body::Body;
use axum::extract::Request;
use axum::middleware::Next;
// A Server-capable host honors either choice. A Client-only host ignores the new field.
for (initial_host, replacement_host) in [(Server, Server), (Server, Client), (Client, Server)] {
let opens = Arc::new(AtomicUsize::new(0));
let subscriptions = Arc::new(AtomicUsize::new(0));
let subscription_count = Arc::clone(&subscriptions);
let app = axum::Router::new()
.fallback_service(CodeModeHostServer::new(GrpcCodeModeHost::new()))
.layer(axum::middleware::from_fn(
move |mut rpc: Request, next: Next| {
let opens = Arc::clone(&opens);
let subscriptions = Arc::clone(&subscription_count);
async move {
if rpc.uri().path().ends_with("/OpenSession") {
let host = if opens.fetch_add(1, Ordering::SeqCst) == 0 {
initial_host
} else {
replacement_host
};
if host == Client {
// No resource limits were requested, so ignoring the output mode
// leaves an empty OpenSessionRequest in one uncompressed gRPC frame.
rpc.headers_mut().remove("content-length");
*rpc.body_mut() = Body::from(&[0_u8; 5][..]);
}
} else if rpc.uri().path().ends_with("/Execute")
&& opens.load(Ordering::SeqCst) == 1
{
return Status::not_found("session was lost").into_http::<Body>();
} else if rpc.uri().path().ends_with("/SubscribeToToolCalls") {
subscriptions.fetch_add(1, Ordering::SeqCst);
}
next.run(rpc).await
}
},
));
let listener = TcpListener::bind("127.0.0.1:0").await?;
let endpoint = format!("http://{}", listener.local_addr()?);
let _server = AbortOnDropHandle::new(tokio::spawn(async move {
axum::serve(listener, app)
.await
.expect("recovery test host");
}));
let session = GrpcCodeModeSessionProvider::new(endpoint)
.create_session_with_options(CodeModeSessionOptions {
output_truncation: Server,
..CodeModeSessionOptions::default()
})
.await
.map_err(anyhow::Error::msg)?;
assert_eq!(session.output_truncation(), initial_host);
let delegate = Arc::new(NoopCodeModeSessionDelegate);
let failed = execute(&session, request("text('must not run')"), delegate.clone())
.await
.unwrap_err();
assert!(failed.to_string().contains("session was lost"));
let reset = execute(
&session,
request("text('must not replay')"),
delegate.clone(),
)
.await
.unwrap_err();
if initial_host == Server && replacement_host == Client {
assert_eq!(
reset.to_string(),
"gRPC code-mode host changed the session output truncation mode on reconnect"
);
assert_eq!(
subscriptions.load(Ordering::SeqCst),
1,
"the downgraded binding must not subscribe"
);
} else {
assert!(
reset
.to_string()
.contains("Stored values and running cells were lost")
);
let source = match initial_host {
Client => "text('x'.repeat(2 * 1024 * 1024))",
Server => "throw 'begin:' + 'X'.repeat(65 * 1024 * 1024) + ':end'",
};
let mut observation = request(source);
observation.max_output_tokens = Some(5);
observation.yield_time_ms = Some(60_000);
let response = execute(&session, observation, delegate).await?;
match initial_host {
Client => assert_eq!(
response,
text_response(
"g2:1",
&"x".repeat(2 * 1024 * 1024),
response.code_mode_host_duration()
)
),
Server => {
let RuntimeResponse::Result {
cell_id,
content_items,
error_text: Some(diagnostic),
..
} = response
else {
panic!("expected a script failure");
};
assert_eq!(cell_id.as_str(), "g2:1");
assert_eq!(diagnostic.len(), 64 * 1024);
assert!(diagnostic.starts_with("begin:") && diagnostic.ends_with(":end"));
assert_eq!(content_items, [
"Warning: truncated output (original token count: 262144)\nTotal output lines: 2\n\nScript err…262139 tokens truncated…",
"[Code Mode session output limit reached: omitted 67108888 text bytes; discarded 0 text segments, 0 images, 0 audio items]",
"XXXXXX:end",
].map(|text| FunctionCallOutputContentItem::InputText { text: text.into() }));
}
}
}
session.shutdown().await.map_err(anyhow::Error::msg)?;
assert_eq!(session.output_truncation(), initial_host);
}
Ok(())
}
codex-rs/code-mode-host/tests/stdio.rs
+36 -2
View File
@@ -10,9 +10,11 @@ use std::os::unix::fs::PermissionsExt;
use codex_code_mode::CellId;
use codex_code_mode::CodeModeNestedToolCall;
use codex_code_mode::CodeModeOutputTruncation;
use codex_code_mode::CodeModeSession;
use codex_code_mode::CodeModeSessionCellExecutionLimits;
use codex_code_mode::CodeModeSessionDelegate;
use codex_code_mode::CodeModeSessionOptions;
use codex_code_mode::CodeModeSessionProvider;
use codex_code_mode::CodeModeToolKind;
use codex_code_mode::ExecuteRequest;
@@ -908,9 +910,14 @@ async fn oversized_initial_response_does_not_close_the_shared_host() {
"unexpected error: {error}"
);
// The default session still returns raw output after the oversized frame failed.
let actual = execute(
&session,
execute_request(r#"text("still alive");"#),
ExecuteRequest {
max_output_tokens: Some(5),
yield_time_ms: Some(60_000),
..execute_request("text('x'.repeat(2 * 1024 * 1024));")
},
Arc::new(RecordingDelegate::default()),
)
.await;
@@ -920,13 +927,40 @@ async fn oversized_initial_response_does_not_close_the_shared_host() {
code_mode_host_duration: actual.code_mode_host_duration(),
cell_id: cell_id("2"),
content_items: vec![FunctionCallOutputContentItem::InputText {
text: "still alive".to_string(),
text: "x".repeat(2 * 1024 * 1024),
}],
result_buffered_for: Duration::ZERO,
error_text: None,
}
);
session.shutdown().await.expect("shutdown remote session");
// With Server truncation negotiated, a script error larger than the frame limit still returns.
let session = provider
.create_session_with_options(CodeModeSessionOptions {
output_truncation: CodeModeOutputTruncation::Server,
..Default::default()
})
.await
.expect("open server-truncated session");
let mut request = execute_request("throw 'begin:' + 'X'.repeat(65 * 1024 * 1024) + ':end';");
request.max_output_tokens = Some(5);
request.yield_time_ms = Some(60_000);
let response = execute(&session, request, Arc::new(RecordingDelegate::default())).await;
let RuntimeResponse::Result {
content_items,
error_text: Some(error),
..
} = response
else {
panic!("expected completed error within the frame limit");
};
assert!(error.starts_with("begin:") && error.ends_with(":end"));
assert!(error.len() > 20 && error.len() <= 64 * 1024);
assert!(
matches!(content_items.last(), Some(FunctionCallOutputContentItem::InputText { text }) if text.ends_with(":end"))
);
session.shutdown().await.expect("shutdown remote session");
}
#[cfg(unix)]
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 {
codex-rs/core/src/code_mode_host.rs
+15 -4
View File
@@ -19,6 +19,7 @@ use std::time::Duration;
use codex_code_mode::CodeModeProcessMetricsSample;
use codex_code_mode::CodeModeSessionCellExecutionLimits;
use codex_code_mode::CodeModeSessionOptions;
use codex_code_mode::CodeModeSessionProvider;
use codex_code_mode::CodeModeSessionProviderFuture;
use codex_code_mode::CodeModeSessionResultFuture;
@@ -43,14 +44,14 @@ use tonic::transport::Endpoint;
use tower::ServiceExt;
use tower::service_fn;
pub(crate) struct ProcessOwnedGrpcCodeModeSessionProvider {
pub struct ProcessOwnedGrpcCodeModeSessionProvider {
host_program: PathBuf,
provider: OnceLock<GrpcCodeModeSessionProvider>,
current: Arc<Mutex<Option<HostChannel>>>,
}
impl ProcessOwnedGrpcCodeModeSessionProvider {
pub(crate) fn with_host_program(host_program: PathBuf) -> Self {
pub fn with_host_program(host_program: PathBuf) -> Self {
Self {
host_program,
provider: OnceLock::new(),
@@ -86,12 +87,22 @@ impl CodeModeSessionProvider for ProcessOwnedGrpcCodeModeSessionProvider {
}
fn create_session(&self) -> CodeModeSessionProviderFuture<'_> {
self.create_session_with_limits(CodeModeSessionCellExecutionLimits::default())
self.create_session_with_options(CodeModeSessionOptions::default())
}
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 {
self.availability()?;
@@ -105,7 +116,7 @@ impl CodeModeSessionProvider for ProcessOwnedGrpcCodeModeSessionProvider {
async move { transport.send(request).await }
}))
});
provider.create_session_with_limits(limits).await
provider.create_session_with_options(options).await
})
}
}
codex-rs/core/src/test_support.rs
+1
View File
@@ -32,6 +32,7 @@ use codex_protocol::protocol::SessionSource;
use once_cell::sync::Lazy;
use crate::ThreadManager;
pub use crate::code_mode_host::ProcessOwnedGrpcCodeModeSessionProvider;
use crate::config::Config;
use crate::responses_metadata::CodexResponsesMetadata;
use crate::responses_metadata::CodexResponsesRequestKind;
codex-rs/core/src/tools/code_mode/execute_handler.rs
+1
View File
@@ -165,6 +165,7 @@ impl CodeModeExecuteHandler {
&step_context.settings.model_info,
response,
args.max_output_tokens,
exec.session.services.code_mode_service.output_truncation(),
wall_time,
exec.turn.config.code_mode.experimental_show_cell_overhead,
))
codex-rs/core/src/tools/code_mode/mod.rs
+18 -7
View File
@@ -15,6 +15,7 @@ use std::time::Duration;
use codex_code_mode::CellId;
use codex_code_mode::CodeModeNestedToolCall;
use codex_code_mode::CodeModeOutputTruncation;
use codex_code_mode::CodeModeSession;
use codex_code_mode::CodeModeSessionProvider;
use codex_code_mode::CodeModeToolKind;
@@ -126,6 +127,13 @@ impl CodeModeService {
Arc::clone(&self.session_provider)
}
fn output_truncation(&self) -> CodeModeOutputTruncation {
self.session
.get()
.map(|session| session.output_truncation())
.unwrap_or_default()
}
pub(crate) async fn execute(
&self,
mut request: codex_code_mode::ExecuteRequest,
@@ -155,9 +163,9 @@ impl CodeModeService {
pub(crate) async fn terminate(
&self,
cell_id: CellId,
request: codex_code_mode::TerminateRequest,
) -> Result<codex_code_mode::WaitOutcome, String> {
self.session().await?.terminate(cell_id.into()).await
self.session().await?.terminate(request).await
}
pub(crate) async fn interrupt_active_cells(&self) {
@@ -263,6 +271,7 @@ fn handle_runtime_response(
model_info: &codex_protocol::openai_models::ModelInfo,
response: RuntimeResponse,
max_output_tokens: Option<usize>,
output_truncation: CodeModeOutputTruncation,
wall_time: Duration,
experimental_show_cell_overhead: bool,
) -> CodeModeToolOutput {
@@ -284,12 +293,14 @@ fn handle_runtime_response(
let mut content_items = into_function_call_output_content_items(content_items);
sanitize_image_detail_items(supports_original, &mut content_items);
let success = error_text.is_none();
if let Some(error_text) = error_text {
content_items.push(FunctionCallOutputContentItem::InputText {
text: format!("Script error:\n{error_text}"),
});
if output_truncation == CodeModeOutputTruncation::Client {
if let Some(error_text) = error_text {
content_items.push(FunctionCallOutputContentItem::InputText {
text: format!("Script error:\n{error_text}"),
});
}
content_items = truncate_code_mode_result(content_items, max_output_tokens);
}
content_items = truncate_code_mode_result(content_items, max_output_tokens);
CodeModeToolOutput::new(
FunctionToolOutput::from_content(content_items, Some(success)),
script_status,
codex-rs/core/src/tools/code_mode/wait_handler.rs
+6 -2
View File
@@ -130,7 +130,10 @@ impl CodeModeWaitHandler {
exec.session
.services
.code_mode_service
.terminate(cell_id.clone())
.terminate(codex_code_mode::TerminateRequest {
cell_id: cell_id.clone(),
max_output_tokens: args.max_tokens,
})
.await
} else {
exec.session
@@ -140,7 +143,7 @@ impl CodeModeWaitHandler {
codex_code_mode::WaitRequest {
cell_id: cell_id.clone(),
yield_time_ms: args.yield_time_ms,
max_output_tokens: None,
max_output_tokens: args.max_tokens,
},
step_context.preempt.clone(),
)
@@ -208,6 +211,7 @@ impl CodeModeWaitHandler {
&step_context.settings.model_info,
wait_response.into(),
args.max_tokens,
exec.session.services.code_mode_service.output_truncation(),
wall_time,
exec.turn.config.code_mode.experimental_show_cell_overhead,
)))
codex-rs/core/tests/suite/code_mode.rs
+129 -25
View File
@@ -12,6 +12,11 @@ use anyhow::Result;
use base64::Engine;
use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
use codex_analytics::AnalyticsEventsClient;
use codex_code_mode::CodeModeOutputTruncation;
use codex_code_mode::CodeModeSessionOptions;
use codex_code_mode::CodeModeSessionProvider;
use codex_code_mode::CodeModeSessionProviderFuture;
use codex_code_mode::ProcessOwnedCodeModeSessionProvider;
use codex_config::types::McpServerConfig;
use codex_config::types::McpServerTransportConfig;
use codex_core::StartThreadOptions;
@@ -20,6 +25,7 @@ use codex_core::config::Config;
use codex_core::config::Constrained;
use codex_core::config::CurrentTimeReminderConfig;
use codex_core::context::NodeReplReviewEvidence;
use codex_core::test_support::ProcessOwnedGrpcCodeModeSessionProvider;
use codex_extension_api::ExtensionData;
use codex_extension_api::ExtensionFuture;
use codex_extension_api::ExtensionRegistryBuilder;
@@ -308,6 +314,74 @@ async fn run_code_mode_turn_with_builder(
Ok((test, second_mock))
}
// Request server truncation from the real host and verify what it negotiated.
struct ServerOutputProvider(Arc<dyn CodeModeSessionProvider>);
impl CodeModeSessionProvider for ServerOutputProvider {
fn availability(&self) -> std::result::Result<(), String> {
self.0.availability()
}
fn create_session(&self) -> CodeModeSessionProviderFuture<'_> {
Box::pin(async {
let session = self
.0
.create_session_with_options(CodeModeSessionOptions {
output_truncation: CodeModeOutputTruncation::Server,
..CodeModeSessionOptions::default()
})
.await?;
assert_eq!(
session.output_truncation(),
CodeModeOutputTruncation::Server
);
Ok(session)
})
}
}
pub(super) enum OutputHost {
DefaultClient,
ServerStdio,
ServerGrpc,
}
impl OutputHost {
pub(super) fn test_codex(self) -> TestCodexBuilder {
let builder = test_codex()
.with_model("test-gpt-5.1-codex")
.with_config(|config| {
config
.features
.enable(Feature::CodeMode)
.expect("enable code mode");
});
let host = || {
codex_utils_cargo_bin::cargo_bin("codex-code-mode-host")
.expect("locate the code-mode host")
};
let provider: Arc<dyn CodeModeSessionProvider> = match self {
Self::DefaultClient => return builder,
Self::ServerStdio => Arc::new(ProcessOwnedCodeModeSessionProvider::with_host_program(
host(),
)),
Self::ServerGrpc => Arc::new(
ProcessOwnedGrpcCodeModeSessionProvider::with_host_program(host()),
),
};
builder
.with_config(|config| {
config
.features
.disable(Feature::CodeModeHost)
.expect("use the injected provider");
})
.with_thread_manager(move |manager| {
manager.with_code_mode_session_provider(Arc::new(ServerOutputProvider(provider)))
})
}
}
async fn run_unavailable_code_mode_turn(
server: &MockServer,
test: &TestCodex,
@@ -4570,12 +4644,15 @@ text(`Variable truncated: ${resultVariableWasTruncated ? "True" : "False"}. Vari
// The outer directive limits output after JavaScript emits it; it does not
// limit `result.output` returned by the nested command.
#[cfg_attr(windows, ignore = "no exec_command on Windows")]
#[test_case(OutputHost::DefaultClient; "client_stdio")]
#[test_case(OutputHost::ServerStdio; "server_stdio")]
#[test_case(OutputHost::ServerGrpc; "server_grpc")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn code_mode_exec_outer_limit_truncates_emitted_output() -> Result<()> {
async fn code_mode_exec_outer_limit_truncates_emitted_output(host: OutputHost) -> Result<()> {
skip_if_no_network!(Ok(()));
let server = responses::start_mock_server().await;
let (_test, second_mock) = run_code_mode_turn(
let (_test, second_mock) = run_code_mode_turn_with_builder(
&server,
"use exec_command from code mode",
r#"// @exec: {"max_output_tokens": 5}
@@ -4584,6 +4661,7 @@ const result = await tools.exec_command({
});
text(result.output);
"#,
host.test_codex(),
)
.await?;
@@ -5546,23 +5624,49 @@ yield_control();
Ok(())
}
// The cell emits its termination output before starting the nested command. Keep
// the first turn's dispatch worker alive until that command has been dispatched.
#[derive(Default)]
struct TerminationOutputReady(tokio::sync::Notify);
impl ToolLifecycleContributor for TerminationOutputReady {
fn on_tool_finish<'a>(&'a self, input: ToolFinishInput<'a>) -> ToolLifecycleFuture<'a> {
Box::pin(async move {
if input.tool_name.name == "exec_command" {
self.0.notify_one();
} else if input.tool_name.name == "exec" && input.call_id == "call-1" {
assert_eq!(input.outcome, ToolCallOutcome::Completed { success: true });
tokio::time::timeout(Duration::from_secs(10), self.0.notified())
.await
.expect("the yielded cell should reach its nested command");
}
})
}
}
#[cfg_attr(windows, ignore = "no exec_command on Windows")]
#[test_case(OutputHost::DefaultClient; "client_stdio")]
#[test_case(OutputHost::ServerStdio; "server_stdio")]
#[test_case(OutputHost::ServerGrpc; "server_grpc")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn code_mode_wait_can_terminate_and_continue() -> Result<()> {
async fn code_mode_wait_can_terminate_and_continue(host: OutputHost) -> Result<()> {
skip_if_no_network!(Ok(()));
let server = responses::start_mock_server().await;
let mut builder = test_codex().with_config(move |config| {
let _ = config.features.enable(Feature::CodeMode);
});
let mut extensions = ExtensionRegistryBuilder::<Config>::new();
extensions.tool_lifecycle_contributor(Arc::new(TerminationOutputReady::default()));
let mut builder = host
.test_codex()
.with_extensions(Arc::new(extensions.build()));
let test = builder.build(&server).await?;
let termination_gate = test.workspace_path("code-mode-terminate.ready");
let termination_wait = wait_for_file_source(&termination_gate)?;
let code = format!(
r#"
r#"// @exec: {{"max_output_tokens": 100}}
text("phase 1");
yield_control();
text("0123456789012345678901234567890123456789");
{termination_wait}
text("phase 2");
"#
@@ -5604,6 +5708,7 @@ text("phase 2");
&serde_json::to_string(&serde_json::json!({
"cell_id": cell_id.clone(),
"terminate": true,
"max_tokens": 5,
}))?,
),
ev_completed("resp-3"),
@@ -5623,7 +5728,7 @@ text("phase 2");
let second_request = second_completion.single_request();
let second_items = function_tool_output_items(&second_request, "call-2");
assert_eq!(second_items.len(), 1);
assert_eq!(second_items.len(), 2);
assert_regex_match(
concat!(
r"(?s)\A",
@@ -5631,6 +5736,10 @@ text("phase 2");
),
text_item(&second_items, /*index*/ 0),
);
assert_eq!(
text_item(&second_items, /*index*/ 1),
"Warning: truncated output (original token count: 10)\nTotal output lines: 1\n\n0123456789…5 tokens truncated…0123456789"
);
responses::mount_sse_once(
&server,
@@ -6224,19 +6333,19 @@ async fn code_mode_interrupt_terminates_active_cells_and_nested_tools() -> Resul
Ok(())
}
#[cfg_attr(windows, ignore = "no exec_command on Windows")]
#[test_case(OutputHost::DefaultClient; "client_stdio")]
#[test_case(OutputHost::ServerStdio; "server_stdio")]
#[test_case(OutputHost::ServerGrpc; "server_grpc")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn code_mode_wait_uses_its_own_max_tokens_budget() -> Result<()> {
async fn code_mode_wait_uses_its_own_max_tokens_budget(host: OutputHost) -> Result<()> {
skip_if_no_network!(Ok(()));
let server = responses::start_mock_server().await;
let mut builder = test_codex().with_config(move |config| {
let _ = config.features.enable(Feature::CodeMode);
});
let test = builder.build(&server).await?;
let mut builder = host.test_codex();
let test = builder.build_with_auto_env(&server).await?;
// An explicit yield separates the output batches without a nested shell call.
let code = r#"// @exec: {"max_output_tokens": 100}
// Both transports accept exec budgets above i32::MAX. The wait uses its own budget.
let code = r#"// @exec: {"max_output_tokens": 2147483648}
text("phase 1");
yield_control();
text("token one token two token three token four token five token six token seven");
@@ -6306,15 +6415,10 @@ text("token one token two token three token four token five token six token seve
),
text_item(&second_items, /*index*/ 0),
);
let expected_pattern = r#"(?sx)
\A
Warning:\ truncated\ output\ \(original\ token\ count:\ \d+\)\n
Total\ output\ lines:\ 1\n
\n
.*…\d+\ tokens\ truncated….*
\z
"#;
assert_regex_match(expected_pattern, text_item(&second_items, /*index*/ 1));
assert_eq!(
text_item(&second_items, /*index*/ 1),
"Warning: truncated output (original token count: 19)\nTotal output lines: 1\n\ntoken one to…13 tokens truncated… token seven"
);
Ok(())
}
codex-rs/core/tests/suite/scenarios.rs
+30 -14
View File
@@ -2198,25 +2198,18 @@ text(`MCP: ${ping.structuredContent?.echo ?? "missing"}`);"#,
async fn astra_reads_code_mode_call_timing() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let test = test_codex()
let test = super::code_mode::OutputHost::ServerGrpc
.test_codex()
.with_model("gpt-6-astra")
.with_config(|config| {
configure_scenario_catalog(config);
// Use the selected cwd as the workspace root on local and remote executors.
config.workspace_roots = vec![config.cwd.clone()];
config.code_mode.experimental_show_cell_overhead = true;
config
.features
.enable(Feature::CodeMode)
.expect("enable code mode");
config
.features
.enable(Feature::CodeModeOnly)
.expect("enable code-mode-only tools");
config
.features
.enable(Feature::CodeModeHost)
.expect("enable the code-mode host");
})
.build_with_auto_env(&server)
.await?;
@@ -2225,7 +2218,13 @@ async fn astra_reads_code_mode_call_timing() -> Result<()> {
vec![
sse(vec![
ev_response_created("exec-response"),
ev_custom_tool_call("exec-call", "exec", "text('ready');"),
ev_custom_tool_call(
"exec-call",
"exec",
r#"// @exec: {"max_output_tokens": 0}
text('ready');
text('x'.repeat(2 * 1024 * 1024));"#,
),
ev_completed("exec-response"),
]),
sse(vec![
@@ -2239,18 +2238,35 @@ async fn astra_reads_code_mode_call_timing() -> Result<()> {
ev_completed("wait-response"),
]),
sse(vec![
ev_assistant_message("final", "The call completed; the missing-cell wait failed."),
ev_response_created("overflow-response"),
ev_custom_tool_call(
"overflow-call",
"exec",
r#"// @exec: {"max_output_tokens": 32}
text('before');
text('x'.repeat(2 * 1024 * 1024));
throw 'boom';"#,
),
ev_completed("overflow-response"),
]),
sse(vec![
ev_assistant_message(
"final",
"The first call completed. The wait and overflow script failed.",
),
ev_completed("final-response"),
]),
],
)
.await;
test.submit_turn("Run a code cell, then inspect its timing and a failed wait.")
.await?;
test.submit_turn(
"Run a code cell, inspect its timing and a failed wait, then check the diagnostic script.",
)
.await?;
insta::assert_snapshot!(
"astra_code_mode_call_timing",
context_snapshot::format_request_history_snapshot(
"Astra receives host and handler timings on completed and failed code-mode calls.",
"Astra reads code-mode timings and the collector notice even at zero tokens, then a failed wait and a truncated script error.",
&mock.requests(),
&ContextSnapshotOptions::default(),
)
codex-rs/core/tests/suite/snapshots/all__suite__scenarios__astra_code_mode_call_timing.snap
+25 -5
View File
@@ -1,8 +1,8 @@
---
source: core/tests/suite/scenarios.rs
expression: "context_snapshot::format_request_history_snapshot(\"Astra receives host and handler timings on completed and failed code-mode calls.\",\n&mock.requests(), &ContextSnapshotOptions::default(),)"
expression: "context_snapshot::format_request_history_snapshot(\"Astra reads code-mode timings and the collector notice even at zero tokens, then a failed wait and a truncated script error.\",\n&mock.requests(), &ContextSnapshotOptions::default(),)"
---
Scenario: Astra receives host and handler timings on completed and failed code-mode calls.
Scenario: Astra reads code-mode timings and the collector notice even at zero tokens, then a failed wait and a truncated script error.
## Window 1
-- request 1 (turn) --
@@ -89,14 +89,19 @@ Scenario: Astra receives host and handler timings on completed and failed code-m
<filesystem><workspace_roots><root><CWD></root></workspace_roots><permission_profile type="disabled"><...e="unrestricted" /></permission_profile></filesystem> [hash=d40b2f45cb028888]
</environment_context>
06:message/user:
Run a code cell, then inspect its timing and a failed wait.
Run a code cell, inspect its timing and a failed wait, then check the diagnostic script.
07:configuration_update:reasoning={"effort":"low"}
-- request 2 (turn) --
08:custom_tool_call/exec:text('ready');
08:custom_tool_call/exec:// @exec: {"max_output_tokens": 0}
text('ready');
text('x'.repeat(2 * 1024 * 1024));
09:custom_tool_call_output:Script completed
Wall time <DURATION> seconds (code-mode <DURATION> seconds; overhead <DURATION> seconds)
Output:
| ready
| Warning: truncated output (original token count: 262145)
Total output lines: 2
…262145 tokens truncated… | [Code Mode session output limit reached: omitted 1048581 text bytes; discarded 0 text segments, 0 images, 0 audio items]
-- request 3 (turn) --
10:function_call/functions.wait:{"cell_id":"missing"}
11:function_call_output:Script failed
@@ -104,3 +109,18 @@ Scenario: Astra receives host and handler timings on completed and failed code-m
Output:
| Script error:
exec cell missing not found
-- request 4 (turn) --
12:custom_tool_call/exec:// @exec: {"max_output_tokens": 32}
text('before');
text('x'.repeat(2 * 1024 * 1024));
throw 'boom';
13:custom_tool_call_output:Script failed
Wall time <DURATION> seconds (code-mode <DURATION> seconds; overhead <DURATION> seconds)
Output:
| Warning: truncated output (original token count: 262145)
Total output lines: 4
before
xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx…262113 tokens truncated… | [Code Mode session output limit reached: omitted 1048600 text bytes; discarded 0 text segments, 0 images, 0 audio items] | xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
Script error:
boom