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

Enable server-side output truncation for Code Mode sessions (#57870)

Browse Source 
## Why

Code Mode sessions always used client-side output handling, leaving server truncation settings and observation budgets unused.

## What changed

- Honor each session's `output_truncation` setting in the runtime and both host transports, and advertise the bounded-output capability.
- Apply `max_output_tokens` independently to execute, wait, and terminate responses in server mode. Keep client mode as the default when the setting is omitted.
- Use bounded output collection and rendering in server mode, preserving head and tail output with omission notices. Include script errors in the output budget while retaining failure status and a separate diagnostic capped at 64 KiB.
- Estimate audio tokens without the Tokio-dependent cache when rendering bounded output.

## Testing

Add runtime and host regression tests for session mode selection, per-request budgets, output overflow, oversized images, zero-token audio output, diagnostic preservation, and termination under output backpressure.

GitOrigin-RevId: 3d9661edcdd2dde31526bb92eb83fa0614f0fa88
This commit is contained in:
Sean Huang
2026-10-29 01:04:48 +00:00
committed by copyberry
parent 4e2dfd2230
commit edd1ae37ec
22 changed files with 1148 additions and 166 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/code-mode-host/src/grpc/conversions.rs
+16
View File
@@ -1,3 +1,4 @@
use codex_code_mode_protocol::CodeModeOutputTruncation;
use codex_code_mode_protocol::CodeModeSessionCellExecutionLimits;
use codex_code_mode_protocol::CodeModeToolKind;
use codex_code_mode_protocol::ExecuteRequest;
@@ -29,6 +30,21 @@ pub(super) fn session_limits(
})
}
pub(super) fn observation_budget(
max_output_tokens: Option<u64>,
output_truncation: CodeModeOutputTruncation,
) -> Result<Option<usize>, Status> {
if output_truncation == CodeModeOutputTruncation::Client {
return Ok(None);
}
max_output_tokens
.map(|value| {
usize::try_from(value)
.map_err(|_| Status::invalid_argument("output budget is too large"))
})
.transpose()
}
pub(super) fn execute_request(request: proto::ExecuteRequest) -> Result<ExecuteRequest, Status> {
validation::identifier(&request.tool_call_id, "tool call ID")?;
Ok(ExecuteRequest {
codex-rs/code-mode-host/src/grpc/mod.rs
+36 -3
View File
@@ -12,6 +12,10 @@ use std::sync::Arc;
use std::time::Instant;
use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::CodeModeOutputTruncation;
use codex_code_mode_protocol::CodeModeSession;
use codex_code_mode_protocol::CodeModeSessionOptions;
use codex_code_mode_protocol::TerminateRequest;
use codex_code_mode_protocol::WaitRequest;
use codex_code_mode_protocol::grpc as proto;
use codex_code_mode_protocol::grpc::code_mode_host_server::CodeModeHost;
@@ -75,7 +79,23 @@ impl GrpcCodeModeHost {
) -> Result<Response<GrpcStream<proto::SessionEvent>>, Status> {
let _permit = self.state.request_permit()?;
let limits = conversions::session_limits(request.cell_execution_limits)?;
Ok(Response::new(self.state.open_session(limits)?))
let output_truncation = match request
.output_truncation
.map(proto::OutputTruncation::try_from)
.transpose()
.map_err(|_| Status::invalid_argument("unknown code-mode output truncation mode"))?
{
None | Some(proto::OutputTruncation::Unspecified | proto::OutputTruncation::Client) => {
CodeModeOutputTruncation::Client
}
Some(proto::OutputTruncation::Server) => CodeModeOutputTruncation::Server,
};
Ok(Response::new(self.state.open_session(
CodeModeSessionOptions {
cell_execution_limits: limits,
output_truncation,
},
)?))
}
#[tracing::instrument(
@@ -242,11 +262,15 @@ impl GrpcCodeModeHost {
validation::identifier(&request.cell_id, "cell ID")?;
validation::identifier(&request.wait_id, "wait ID")?;
let _permit = self.state.request_permit()?;
let max_output_tokens = conversions::observation_budget(
request.max_output_tokens,
session.runtime.output_truncation(),
)?;
let registration = WaitRegistration::new(Arc::clone(&session), request.wait_id)?;
let request = WaitRequest {
cell_id: CellId::new(request.cell_id),
yield_time_ms: request.yield_time_ms,
max_output_tokens: None,
max_output_tokens,
};
let outcome = tokio::select! {
biased;
@@ -301,7 +325,16 @@ impl GrpcCodeModeHost {
let session = self.state.session(&request.session_id)?;
validation::identifier(&request.cell_id, "cell ID")?;
let _permit = self.state.request_permit()?;
let outcome = session.terminate(CellId::new(request.cell_id)).await?;
let max_output_tokens = conversions::observation_budget(
request.max_output_tokens,
session.runtime.output_truncation(),
)?;
let outcome = session
.terminate(TerminateRequest {
cell_id: CellId::new(request.cell_id),
max_output_tokens,
})
.await?;
let outcome = outcome.with_code_mode_host_duration(received_at.elapsed());
conversions::wait_response(outcome).map_err(|error| Status::internal(error.to_string()))
}
codex-rs/code-mode-host/src/grpc/service_tests.rs
+178 -4
View File
@@ -11,10 +11,18 @@ use tonic::Request;
pub(super) async fn open_session(
host: &GrpcCodeModeHost,
) -> (String, GrpcStream<proto::SessionEvent>) {
let (opened, stream) = open_session_with_truncation(host, /*output_truncation*/ None).await;
(opened.session_id, stream)
}
async fn open_session_with_truncation(
host: &GrpcCodeModeHost,
output_truncation: Option<i32>,
) -> (proto::SessionOpened, GrpcStream<proto::SessionEvent>) {
let mut stream = host
.open_session(Request::new(proto::OpenSessionRequest {
cell_execution_limits: None,
output_truncation: None,
output_truncation,
}))
.await
.expect("open code-mode session")
@@ -27,7 +35,7 @@ pub(super) async fn open_session(
let Some(proto::session_event::Event::Opened(opened)) = event.event else {
panic!("expected the first session event to open its lease");
};
(opened.session_id, stream)
(opened, stream)
}
pub(super) fn execute_request(
@@ -80,6 +88,18 @@ pub(super) async fn execute_events(
(started.cell_id, stream)
}
fn only_text(items: &[proto::ContentItem]) -> &str {
let [
proto::ContentItem {
item: Some(proto::content_item::Item::Text(text)),
},
] = items
else {
panic!("expected one text item, got {items:?}");
};
&text.text
}
#[tokio::test]
async fn execute_stream_starts_immediately_and_wait_preserves_missing_cells() {
let host = GrpcCodeModeHost::new();
@@ -358,7 +378,7 @@ async fn cancellation_before_wait_admission_is_preserved() {
cell_id: cell_id.clone(),
wait_id: "pre-cancelled".to_string(),
yield_time_ms: 60_000,
max_output_tokens: None,
max_output_tokens: Some(0),
}))
.await
.unwrap_err();
@@ -367,7 +387,7 @@ async fn cancellation_before_wait_admission_is_preserved() {
host.terminate(Request::new(proto::TerminateRequest {
session_id,
cell_id,
max_output_tokens: None,
max_output_tokens: Some(0),
}))
.await
.unwrap();
@@ -440,3 +460,157 @@ async fn notifications_do_not_delay_cell_completion() {
.await
.expect("legacy notification acknowledgments remain accepted");
}
#[tokio::test]
async fn output_truncation_is_a_session_setting_and_omitting_it_keeps_legacy_output() {
for (output_truncation, bounded_output) in [
(None, false),
(Some(proto::OutputTruncation::Unspecified.into()), false),
(Some(proto::OutputTruncation::Client.into()), false),
(Some(proto::OutputTruncation::Server.into()), true),
] {
let host = GrpcCodeModeHost::new();
let (opened, _lease) = open_session_with_truncation(&host, output_truncation).await;
assert_eq!(
opened.output_truncation,
bounded_output.then_some(proto::OutputTruncation::Server.into()),
);
let mut request = execute_request(
&opened.session_id,
"execution",
"text('x'.repeat(2 * 1024 * 1024));",
);
request.yield_time_ms = Some(60_000);
let (_, mut execution) = execute_events(&host, request).await;
let outcome = execution
.next()
.await
.expect("outcome event")
.expect("outcome");
let Some(proto::execute_event::Event::Outcome(outcome)) = outcome.event else {
panic!("expected execution outcome");
};
assert!(matches!(
outcome.outcome,
Some(proto::execution_outcome::Outcome::Completed(_))
));
assert_eq!(
outcome.content_items.iter().any(
|item| matches!(&item.item, Some(proto::content_item::Item::Text(text))
if text.text.starts_with("[Code Mode session output limit reached: omitted "))
),
bounded_output
);
if !bounded_output {
assert_eq!(
only_text(&outcome.content_items),
"x".repeat(2 * 1024 * 1024)
);
}
}
}
#[tokio::test]
async fn server_output_budgets_follow_exec_wait_and_terminate_over_grpc() {
for (execution_id, after_yield) in [
("failed", "throw '0123456789'.repeat(8);"),
(
"terminated",
"text('0123456789'.repeat(8)); notify('ready to terminate'); await new Promise(() => {});",
),
] {
let host = GrpcCodeModeHost::new();
let (opened, mut lease) =
open_session_with_truncation(&host, Some(proto::OutputTruncation::Server.into())).await;
let source = format!("text('0123456789'.repeat(4)); yield_control(); {after_yield}");
let mut request = execute_request(&opened.session_id, execution_id, &source);
request.yield_time_ms = Some(60_000);
request.max_output_tokens = Some(5);
let (cell_id, mut execution) = execute_events(&host, request).await;
let Some(proto::execute_event::Event::Outcome(first)) =
execution.next().await.unwrap().unwrap().event
else {
panic!("expected execution outcome");
};
assert_eq!(first.cell_id, cell_id);
assert!(matches!(
first.outcome,
Some(proto::execution_outcome::Outcome::Yielded(_))
));
assert_eq!(
only_text(&first.content_items),
"Warning: truncated output (original token count: 10)\nTotal output lines: 1\n\n0123456789…5 tokens truncated…0123456789"
);
let response = if execution_id == "failed" {
host.wait(Request::new(proto::WaitRequest {
session_id: opened.session_id,
cell_id: cell_id.clone(),
wait_id: "wait-after-yield".to_string(),
yield_time_ms: 60_000,
max_output_tokens: Some(7),
}))
.await
.expect("script failures are completed results, not RPC failures")
.into_inner()
} else {
// The event is ordered after text in the runtime FIFO, so the
// termination observes it without a sleep or racing the cell.
assert!(matches!(
lease.next().await.unwrap().unwrap().event,
Some(proto::session_event::Event::Notification(
proto::Notification { text, .. }
)) if text == "ready to terminate"
));
host.terminate(Request::new(proto::TerminateRequest {
session_id: opened.session_id,
cell_id: cell_id.clone(),
max_output_tokens: Some(7),
}))
.await
.expect("terminate with an independent output budget")
.into_inner()
};
let Some(proto::wait_response::State::LiveCell(last)) = response.state else {
panic!("expected final observation from the live cell");
};
let expected_text = if execution_id == "failed" {
let Some(proto::execution_outcome::Outcome::Completed(proto::ExecutionCompleted {
error_text: Some(diagnostic),
})) = &last.outcome
else {
panic!(
"expected completed failure with a diagnostic, got {:?}",
last.outcome
);
};
assert_eq!(diagnostic, &"0123456789".repeat(8));
"Warning: truncated output (original token count: 24)\nTotal output lines: 2\n\nScript error:\n…17 tokens truncated…67890123456789"
} else {
assert_eq!(
last.outcome,
Some(proto::execution_outcome::Outcome::Terminated(
proto::ExecutionTerminated {}
))
);
"Warning: truncated output (original token count: 20)\nTotal output lines: 1\n\n01234567890123…13 tokens truncated…67890123456789"
};
assert_eq!(last.cell_id, cell_id);
assert_eq!(only_text(&last.content_items), expected_text);
}
}
#[tokio::test]
async fn unknown_output_truncation_is_rejected_before_opening_a_session() {
let host = GrpcCodeModeHost::new();
let result = host
.open_session(Request::new(proto::OpenSessionRequest {
cell_execution_limits: None,
output_truncation: Some(42),
}))
.await;
let Err(status) = result else {
panic!("unknown truncation mode should be rejected");
};
assert_eq!(status.code(), Code::InvalidArgument);
}
codex-rs/code-mode-host/src/grpc/session.rs
+19 -11
View File
@@ -10,7 +10,10 @@ use std::sync::PoisonError;
use std::sync::Weak;
use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::CodeModeSessionCellExecutionLimits;
use codex_code_mode_protocol::CodeModeOutputTruncation;
use codex_code_mode_protocol::CodeModeSession;
use codex_code_mode_protocol::CodeModeSessionOptions;
use codex_code_mode_protocol::TerminateRequest;
use codex_code_mode_protocol::WaitOutcome;
use codex_code_mode_protocol::grpc as proto;
use codex_code_mode_protocol::host::MAX_PENDING_DELEGATE_CALLS;
@@ -107,7 +110,7 @@ impl GrpcHostState {
pub(super) fn open_session(
self: &Arc<Self>,
limits: CodeModeSessionCellExecutionLimits,
options: CodeModeSessionOptions,
) -> Result<GrpcStream<proto::SessionEvent>, Status> {
let id = Uuid::new_v4();
let (events, receiver) = mpsc::channel(OUTGOING_CHANNEL_CAPACITY);
@@ -118,13 +121,18 @@ impl GrpcHostState {
event_sender,
closed,
Arc::clone(&self.delegate_permits),
limits,
options,
);
events
.try_send(Ok(proto::SessionEvent {
event: Some(proto::session_event::Event::Opened(proto::SessionOpened {
session_id: id.to_string(),
output_truncation: None,
output_truncation: match session.runtime.output_truncation() {
CodeModeOutputTruncation::Server => {
Some(proto::OutputTruncation::Server.into())
}
CodeModeOutputTruncation::Client => None,
},
})),
}))
.map_err(|_| Status::internal("failed to publish the opened code-mode session"))?;
@@ -213,7 +221,7 @@ impl GrpcSession {
events: EventSender,
closed: CancellationToken,
delegate_permits: Arc<Semaphore>,
limits: CodeModeSessionCellExecutionLimits,
options: CodeModeSessionOptions,
) -> Arc<Self> {
Arc::new_cyclic(|weak: &Weak<Self>| {
let failure_session = weak.clone();
@@ -225,9 +233,9 @@ impl GrpcSession {
});
Self {
id,
runtime: Arc::new(InProcessCodeModeSession::with_task_failure_handler(
failure_handler,
limits,
runtime: Arc::new(InProcessCodeModeSession::with_options(
options,
Some(failure_handler),
)),
closed,
state: Mutex::new(SessionState::default()),
@@ -258,13 +266,13 @@ impl GrpcSession {
result
}
pub(super) async fn terminate(&self, cell_id: CellId) -> Result<WaitOutcome, Status> {
pub(super) async fn terminate(&self, request: TerminateRequest) -> Result<WaitOutcome, Status> {
tokio::select! {
biased;
_ = self.closed.cancelled() => {
Err(Status::cancelled("code-mode session is closed"))
}
result = self.runtime.terminate(cell_id.into()) => {
result = self.runtime.terminate(request) => {
result.map_err(Status::failed_precondition)
}
}
@@ -347,7 +355,7 @@ impl GrpcSession {
if let Some(cell_id) = cell_id {
let session = Arc::clone(self);
tokio::spawn(async move {
let _ = session.terminate(CellId::new(cell_id)).await;
let _ = session.terminate(CellId::new(cell_id).into()).await;
});
}
}
codex-rs/code-mode-host/src/host_tests.rs
+226 -50
View File
@@ -9,7 +9,9 @@ use std::task::Poll;
use std::time::Duration;
use std::time::Instant;
use codex_code_mode_protocol::CodeModeSessionCellExecutionLimits;
use codex_code_mode_protocol::CodeModeSessionOptions;
use codex_code_mode_protocol::TerminateRequest;
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;
@@ -27,9 +29,15 @@ use codex_code_mode_protocol::host::RequestId;
use codex_code_mode_protocol::host::SESSION_RESOURCE_LIMITS_CAPABILITY;
use codex_code_mode_protocol::host::SessionId;
use codex_code_mode_protocol::host::SupportedProtocolVersions;
use codex_code_mode_protocol::host::WireContentItem;
use codex_code_mode_protocol::host::WireExecuteRequest;
use codex_code_mode_protocol::host::WireOutputTruncation;
use codex_code_mode_protocol::host::WireResult;
use codex_code_mode_protocol::host::WireRuntimeResponse;
use codex_code_mode_protocol::host::WireWaitOutcome;
use codex_code_mode_protocol::host::WireWaitRequest;
use pretty_assertions::assert_eq;
use tokio::io::AsyncRead;
use tokio::io::AsyncWrite;
use tokio::sync::Semaphore;
use tokio::sync::mpsc;
@@ -96,6 +104,20 @@ fn execute_request(source: &str) -> WireExecuteRequest {
}
}
async fn next_response(reader: &mut FramedReader<impl AsyncRead + Unpin>) -> HostToClient {
loop {
match reader
.read()
.await
.expect("host response")
.expect("host connected")
{
HostToClient::CellClosed { .. } => {}
message => return message,
}
}
}
#[tokio::test]
async fn handshake_and_multiple_session_lifecycles_are_ordered() {
let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 4096);
@@ -258,52 +280,209 @@ impl AsyncWrite for BlockingWriter {
}
#[tokio::test]
async fn session_resource_limits_are_negotiated_when_optional_or_required() {
let resource_limits_capability =
Capability::new(SESSION_RESOURCE_LIMITS_CAPABILITY).expect("resource limits capability");
let resource_limits =
CapabilitySet::try_new([resource_limits_capability.clone()]).expect("host capabilities");
for (required, optional) in [
(
CapabilitySet::empty(),
CapabilitySet::try_new([resource_limits_capability]).expect("optional capabilities"),
),
(resource_limits.clone(), CapabilitySet::empty()),
async fn supported_capabilities_are_negotiated_when_optional_or_required() {
for capability in [
SESSION_RESOURCE_LIMITS_CAPABILITY,
BOUNDED_OUTPUT_CAPABILITY,
] {
let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 1024);
let (host_reader, host_writer) = tokio::io::split(host_stream);
let (client_reader, client_writer) = tokio::io::split(client_stream);
let host = tokio::spawn(run(host_reader, host_writer));
let mut reader = FramedReader::new(client_reader);
let mut writer = FramedWriter::new(client_writer);
let capability = Capability::new(capability).expect("supported capability");
let host_capabilities =
CapabilitySet::try_new([capability.clone()]).expect("host capabilities");
writer
.write(&ClientToHost::ClientHello(
ClientHello::new(
SupportedProtocolVersions::try_new([ProtocolVersion::V1])
.expect("supported versions"),
required,
optional,
)
.expect("client hello"),
))
.await
.expect("write hello");
assert_eq!(
reader.read::<HostToClient>().await.expect("host hello"),
Some(HostToClient::HostHello(HostHello::new(
ProtocolVersion::V1,
resource_limits.clone(),
)))
);
for (required, optional) in [
(
CapabilitySet::empty(),
CapabilitySet::try_new([capability]).expect("optional capabilities"),
),
(host_capabilities.clone(), CapabilitySet::empty()),
] {
let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 1024);
let (host_reader, host_writer) = tokio::io::split(host_stream);
let (client_reader, client_writer) = tokio::io::split(client_stream);
let host = tokio::spawn(run(host_reader, host_writer));
let mut reader = FramedReader::new(client_reader);
let mut writer = FramedWriter::new(client_writer);
drop(writer);
drop(reader);
host.await.expect("host task").expect("host connection");
writer
.write(&ClientToHost::ClientHello(
ClientHello::new(
SupportedProtocolVersions::try_new([ProtocolVersion::V1])
.expect("supported versions"),
required,
optional,
)
.expect("client hello"),
))
.await
.expect("write hello");
assert_eq!(
reader.read::<HostToClient>().await.expect("host hello"),
Some(HostToClient::HostHello(HostHello::new(
ProtocolVersion::V1,
host_capabilities.clone(),
)))
);
drop(writer);
drop(reader);
host.await.expect("host task").expect("host connection");
}
}
}
#[tokio::test]
async fn output_selection_is_per_session() {
let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 4096);
let (host_reader, host_writer) = tokio::io::split(host_stream);
let (client_reader, client_writer) = tokio::io::split(client_stream);
let host = tokio::spawn(run(host_reader, host_writer));
let mut reader = FramedReader::new(client_reader);
let mut writer = FramedWriter::new(client_writer);
writer
.write(&client_hello([ProtocolVersion::V1], CapabilitySet::empty()))
.await
.expect("write hello");
assert_eq!(
reader.read::<HostToClient>().await.expect("host hello"),
Some(HostToClient::HostHello(HostHello::new(
ProtocolVersion::V1,
CapabilitySet::empty(),
)))
);
for (id, name, output_truncation) in [
(1, "client", None),
(2, "server", Some(WireOutputTruncation::Server)),
] {
writer
.write(&ClientToHost::Request {
id: request_id(id),
request: HostRequest::OpenSession {
session_id: session_id(name),
cell_execution_limits: None,
output_truncation,
},
})
.await
.expect("open session");
assert_eq!(
reader.read::<HostToClient>().await.expect("session ready"),
Some(HostToClient::Response {
id: request_id(id),
result: WireResult::Ok {
value: HostResponse::SessionReady {
session_id: session_id(name),
},
},
})
);
}
for (id, name, expected_text, expected_wait) in [
(
3,
"client",
"0123456789012345678901234567890123456789",
"01234567890123456789012345678901234567890123456789012345678901234567890123456789",
),
(
4,
"server",
"Warning: truncated output (original token count: 10)\nTotal output lines: 1\n\n0123456789…5 tokens truncated…0123456789",
"Warning: truncated output (original token count: 20)\nTotal output lines: 1\n\n01234567890123…13 tokens truncated…67890123456789",
),
] {
let mut request = execute_request(
"text('0123456789'.repeat(4)); yield_control(); text('0123456789'.repeat(8));",
);
request.max_output_tokens = Some(5);
writer
.write(&ClientToHost::Request {
id: request_id(id),
request: HostRequest::Execute {
session_id: session_id(name),
request,
},
})
.await
.expect("execute");
// Admission is sent before start_cell can send its initial result.
assert!(matches!(
next_response(&mut reader).await,
HostToClient::Response {
id: actual_id,
result: WireResult::Ok { value: HostResponse::ExecutionStarted { .. } },
} if actual_id == request_id(id)
));
let HostToClient::InitialResponse {
id: actual_id,
result:
WireResult::Ok {
value:
WireRuntimeResponse::Yielded {
cell_id,
content_items,
..
},
},
} = next_response(&mut reader).await
else {
panic!("expected initial yield");
};
assert_eq!(actual_id, request_id(id));
assert_eq!(
content_items,
vec![WireContentItem::InputText {
text: expected_text.to_string()
}]
);
writer
.write(&ClientToHost::Request {
id: request_id(id + 2),
request: HostRequest::Wait {
session_id: session_id(name),
request: WireWaitRequest {
cell_id,
yield_time_ms: 60_000,
},
// A supplied budget only affects Server mode and belongs
// to this observation, not the preceding exec.
max_output_tokens: Some(7),
},
})
.await
.expect("wait");
let HostToClient::Response {
id: actual_id,
result:
WireResult::Ok {
value:
HostResponse::WaitCompleted {
outcome:
WireWaitOutcome::LiveCell(WireRuntimeResponse::Result {
content_items,
error_text,
..
}),
},
},
} = next_response(&mut reader).await
else {
panic!("expected final wait");
};
assert_eq!(actual_id, request_id(id + 2));
assert_eq!(error_text, None);
assert_eq!(
content_items,
vec![WireContentItem::InputText {
text: expected_wait.to_string()
}]
);
}
drop(writer);
drop(reader);
host.await.expect("host task").expect("host connection");
}
#[tokio::test]
async fn incompatible_or_invalid_handshake_is_rejected() {
let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 1024);
@@ -521,10 +700,7 @@ async fn execute_request_id_remains_active_until_initial_response() {
});
let session_id = session_id("session-1");
state
.open_session(
session_id.clone(),
CodeModeSessionCellExecutionLimits::default(),
)
.open_session(session_id.clone(), CodeModeSessionOptions::default())
.expect("open session");
let request_id = request_id(/*value*/ 1);
@@ -561,7 +737,10 @@ async fn execute_request_id_remains_active_until_initial_response() {
state
.session(&session_id)
.expect("session")
.terminate(codex_code_mode_protocol::CellId::from(cell_id).into())
.terminate(TerminateRequest {
cell_id: cell_id.into(),
max_output_tokens: None,
})
.await
.expect("terminate cell");
state.disconnect().await;
@@ -585,10 +764,7 @@ async fn active_cell_limit_rejects_execute_without_disconnecting() {
};
let session_id = session_id("session-1");
state
.open_session(
session_id.clone(),
CodeModeSessionCellExecutionLimits::default(),
)
.open_session(session_id.clone(), CodeModeSessionOptions::default())
.expect("open session");
let request_id = request_id(/*value*/ 1);
codex-rs/code-mode-host/src/lib.rs
+32 -10
View File
@@ -11,7 +11,12 @@ use std::time::Instant;
use anyhow::Context;
use anyhow::Result;
use codex_code_mode_protocol::CodeModeOutputTruncation;
use codex_code_mode_protocol::CodeModeSessionCellExecutionLimits;
use codex_code_mode_protocol::CodeModeSessionOptions;
use codex_code_mode_protocol::TerminateRequest;
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;
@@ -278,6 +283,7 @@ async fn negotiate<R: AsyncRead + Unpin, W: AsyncWrite + Unpin>(
[
SESSION_RESOURCE_LIMITS_CAPABILITY,
YIELD_OBSERVATION_CAPABILITY,
BOUNDED_OUTPUT_CAPABILITY,
]
.into_iter()
.map(Capability::new)
@@ -398,13 +404,23 @@ impl HostState {
HostRequest::OpenSession {
session_id,
cell_execution_limits,
..
output_truncation,
} => {
let result = CodeModeSessionCellExecutionLimits::try_from(
cell_execution_limits.unwrap_or_default(),
)
.map_err(|error| format!("invalid code-mode session execution limits: {error}"))
.and_then(|limits| self.open_session(session_id.clone(), limits))
.and_then(|limits| {
self.open_session(
session_id.clone(),
CodeModeSessionOptions {
cell_execution_limits: limits,
output_truncation: output_truncation
.map(CodeModeOutputTruncation::from)
.unwrap_or_default(),
},
)
})
.map(|()| HostResponse::SessionReady { session_id });
self.respond(request_id, result);
}
@@ -469,7 +485,7 @@ impl HostState {
HostRequest::Wait {
session_id,
request,
..
max_output_tokens,
} => {
let result = match self.session(&session_id) {
Ok(session) => {
@@ -478,7 +494,10 @@ impl HostState {
_ = cancellation.cancelled() => {
Err("code-mode request cancelled".to_string())
}
result = session.wait(request.into(), Some(yield_signal)) => result.and_then(|outcome| {
result = session.wait(WaitRequest {
max_output_tokens,
..request.into()
}, Some(yield_signal)) => result.and_then(|outcome| {
let outcome = outcome.with_code_mode_host_duration(received_at.elapsed());
Ok(HostResponse::WaitCompleted {
outcome: WireWaitOutcome::try_from(outcome)
@@ -494,11 +513,14 @@ impl HostState {
HostRequest::Terminate {
session_id,
cell_id,
..
max_output_tokens,
} => {
let result = match self.session(&session_id) {
Ok(session) => session
.terminate(codex_code_mode_protocol::CellId::from(cell_id).into())
.terminate(TerminateRequest {
cell_id: cell_id.into(),
max_output_tokens,
})
.await
.and_then(|outcome| {
let outcome =
@@ -536,7 +558,7 @@ impl HostState {
fn open_session(
&self,
session_id: SessionId,
cell_execution_limits: CodeModeSessionCellExecutionLimits,
options: CodeModeSessionOptions,
) -> Result<(), String> {
let mut sessions = self.sessions.lock().unwrap_or_else(PoisonError::into_inner);
if sessions.contains_key(&session_id) {
@@ -563,9 +585,9 @@ impl HostState {
});
sessions.insert(
session_id,
Arc::new(InProcessCodeModeSession::with_task_failure_handler(
task_failure_handler,
cell_execution_limits,
Arc::new(InProcessCodeModeSession::with_options(
options,
Some(task_failure_handler),
)),
);
Ok(())
codex-rs/code-mode-host/tests/grpc_timing.rs
+4 -5
View File
@@ -165,11 +165,10 @@ async fn stdio_execution_timing_includes_javascript_but_excludes_delayed_reads()
if suffix.is_empty() {
assert_eq!(error_text, None);
} else {
assert!(
error_text
.context("expected JS failure")?
.contains("timed failure")
);
let Some(error_text) = error_text else {
anyhow::bail!("expected JS failure");
};
assert!(error_text.contains("timed failure"));
}
}
session.shutdown().await.map_err(anyhow::Error::msg)?;