From e963cab11a8663b2d7e1746ea11d0b556cb5eee1 Mon Sep 17 00:00:00 2001 From: Sean Huang Date: Mon, 26 Oct 2026 19:03:12 +0000 Subject: [PATCH] Add code-mode output truncation protocol contracts (#57272) ## What changed Add session options for client- or server-owned output truncation and optional `max_output_tokens` budgets for wait and terminate observations. Introduce `TerminateRequest` and update session implementations and callers to use it. Extend the stdio and gRPC schemas with optional truncation settings and budgets, and define the `bounded-output-v1` capability. Existing sessions retain client-side truncation; providers reject server mode by default. This change defines the contracts without enabling server-side truncation. ## Testing Extend protocol serialization coverage for `outputTruncation: "server"` while preserving existing wire shapes when optional fields are absent. GitOrigin-RevId: 22145897ea16d385c5ba9e7549109f7a669fceb4 --- codex-rs/code-mode-host/src/grpc/mod.rs | 1 + .../src/grpc/robustness_tests.rs | 1 + .../code-mode-host/src/grpc/service_tests.rs | 5 ++ codex-rs/code-mode-host/src/grpc/session.rs | 3 +- codex-rs/code-mode-host/src/host_tests.rs | 7 ++- codex-rs/code-mode-host/src/lib.rs | 21 ++++--- codex-rs/code-mode-host/tests/grpc.rs | 33 +++++++--- .../tests/grpc/admission_tests.rs | 6 +- .../tests/grpc/image_cancellation_tests.rs | 2 +- .../tests/grpc_notifications.rs | 3 +- codex-rs/code-mode-host/tests/grpc_tcp.rs | 1 + codex-rs/code-mode-host/tests/grpc_timing.rs | 5 +- codex-rs/code-mode-host/tests/settled.rs | 2 + codex-rs/code-mode-host/tests/stdio.rs | 17 +++++- .../src/grpc/codex.code_mode.v1.proto | 14 +++++ .../code-mode-protocol/src/host/host_tests.rs | 18 ++++++ .../code-mode-protocol/src/host/message.rs | 9 +++ codex-rs/code-mode-protocol/src/host/mod.rs | 7 +++ .../code-mode-protocol/src/host/payload.rs | 31 +++++++++- codex-rs/code-mode-protocol/src/lib.rs | 3 + codex-rs/code-mode-protocol/src/runtime.rs | 20 ++++++ codex-rs/code-mode-protocol/src/session.rs | 61 ++++++++++++++++++- codex-rs/code-mode-runtime/src/service.rs | 12 +++- .../src/service_audio_tests.rs | 1 + .../src/service_contract_tests.rs | 16 +++-- .../code-mode-runtime/src/service_tests.rs | 18 ++++-- codex-rs/code-mode-runtime/tests/settled.rs | 2 +- codex-rs/code-mode/src/grpc_session/mod.rs | 9 ++- .../code-mode/src/grpc_session/operations.rs | 2 + .../code-mode/src/grpc_session/reconnect.rs | 15 +++-- codex-rs/code-mode/src/remote_session.rs | 16 +++-- .../connection/driver/commands.rs | 3 + .../connection/driver/responses.rs | 1 + .../remote_session/connection/driver_tests.rs | 8 +++ codex-rs/core/src/tools/code_mode/mod.rs | 4 +- .../core/src/tools/code_mode/wait_handler.rs | 1 + .../src/tools/tool_dispatch_trace_tests.rs | 5 +- codex-rs/core/tests/code_mode_tracing.rs | 1 + 38 files changed, 329 insertions(+), 55 deletions(-) diff --git a/codex-rs/code-mode-host/src/grpc/mod.rs b/codex-rs/code-mode-host/src/grpc/mod.rs index 088077e33c..13fb48d3c5 100644 --- a/codex-rs/code-mode-host/src/grpc/mod.rs +++ b/codex-rs/code-mode-host/src/grpc/mod.rs @@ -234,6 +234,7 @@ impl GrpcCodeModeHost { let request = WaitRequest { cell_id: CellId::new(request.cell_id), yield_time_ms: request.yield_time_ms, + max_output_tokens: None, }; let outcome = tokio::select! { biased; diff --git a/codex-rs/code-mode-host/src/grpc/robustness_tests.rs b/codex-rs/code-mode-host/src/grpc/robustness_tests.rs index 45aac60556..fe6cdf56b0 100644 --- a/codex-rs/code-mode-host/src/grpc/robustness_tests.rs +++ b/codex-rs/code-mode-host/src/grpc/robustness_tests.rs @@ -255,6 +255,7 @@ async fn session_closure_cancels_pending_termination() { let termination = host.terminate(Request::new(proto::TerminateRequest { session_id, cell_id, + max_output_tokens: None, })); tokio::pin!(termination); assert!(termination.as_mut().now_or_never().is_none()); diff --git a/codex-rs/code-mode-host/src/grpc/service_tests.rs b/codex-rs/code-mode-host/src/grpc/service_tests.rs index cf30586ad2..c847145221 100644 --- a/codex-rs/code-mode-host/src/grpc/service_tests.rs +++ b/codex-rs/code-mode-host/src/grpc/service_tests.rs @@ -14,6 +14,7 @@ pub(super) async fn open_session( let mut stream = host .open_session(Request::new(proto::OpenSessionRequest { cell_execution_limits: None, + output_truncation: None, })) .await .expect("open code-mode session") @@ -120,6 +121,7 @@ async fn execute_stream_starts_immediately_and_wait_preserves_missing_cells() { cell_id: cell_id.clone(), wait_id: "wait-1".to_string(), yield_time_ms: 60_000, + max_output_tokens: None, })) .await .expect("wait for completion") @@ -163,6 +165,7 @@ async fn execute_stream_starts_immediately_and_wait_preserves_missing_cells() { cell_id: "missing-cell".to_string(), wait_id: "wait-missing".to_string(), yield_time_ms: 1, + max_output_tokens: None, })) .await .expect("missing cells remain successful wait outcomes") @@ -355,6 +358,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, })) .await .unwrap_err(); @@ -363,6 +367,7 @@ async fn cancellation_before_wait_admission_is_preserved() { host.terminate(Request::new(proto::TerminateRequest { session_id, cell_id, + max_output_tokens: None, })) .await .unwrap(); diff --git a/codex-rs/code-mode-host/src/grpc/session.rs b/codex-rs/code-mode-host/src/grpc/session.rs index dbfc96ec95..f95e008b21 100644 --- a/codex-rs/code-mode-host/src/grpc/session.rs +++ b/codex-rs/code-mode-host/src/grpc/session.rs @@ -124,6 +124,7 @@ impl GrpcHostState { .try_send(Ok(proto::SessionEvent { event: Some(proto::session_event::Event::Opened(proto::SessionOpened { session_id: id.to_string(), + output_truncation: None, })), })) .map_err(|_| Status::internal("failed to publish the opened code-mode session"))?; @@ -263,7 +264,7 @@ impl GrpcSession { _ = self.closed.cancelled() => { Err(Status::cancelled("code-mode session is closed")) } - result = self.runtime.terminate(cell_id) => { + result = self.runtime.terminate(cell_id.into()) => { result.map_err(Status::failed_precondition) } } diff --git a/codex-rs/code-mode-host/src/host_tests.rs b/codex-rs/code-mode-host/src/host_tests.rs index 0d091ac4ef..368a7ee6cb 100644 --- a/codex-rs/code-mode-host/src/host_tests.rs +++ b/codex-rs/code-mode-host/src/host_tests.rs @@ -127,6 +127,7 @@ async fn handshake_and_multiple_session_lifecycles_are_ordered() { request: HostRequest::OpenSession { session_id: session_id(id), cell_execution_limits: None, + output_truncation: None, }, }) .await @@ -198,6 +199,7 @@ async fn disconnect_cancels_a_backpressured_host_writer() { request: HostRequest::OpenSession { session_id: session_id("backpressured-session"), cell_execution_limits: None, + output_truncation: None, }, }) .await @@ -339,6 +341,7 @@ async fn incompatible_or_invalid_handshake_is_rejected() { request: HostRequest::OpenSession { session_id: session_id("session-1"), cell_execution_limits: None, + output_truncation: None, }, }) .await @@ -405,6 +408,7 @@ async fn session_id_cannot_be_reused_after_shutdown() { HostRequest::OpenSession { session_id: id.clone(), cell_execution_limits: None, + output_truncation: None, }, ), ( @@ -433,6 +437,7 @@ async fn session_id_cannot_be_reused_after_shutdown() { request: HostRequest::OpenSession { session_id: id, cell_execution_limits: None, + output_truncation: None, }, }) .await @@ -556,7 +561,7 @@ async fn execute_request_id_remains_active_until_initial_response() { state .session(&session_id) .expect("session") - .terminate(cell_id.into()) + .terminate(codex_code_mode_protocol::CellId::from(cell_id).into()) .await .expect("terminate cell"); state.disconnect().await; diff --git a/codex-rs/code-mode-host/src/lib.rs b/codex-rs/code-mode-host/src/lib.rs index 7c73ae23df..77c80e80b2 100644 --- a/codex-rs/code-mode-host/src/lib.rs +++ b/codex-rs/code-mode-host/src/lib.rs @@ -382,6 +382,7 @@ impl HostState { HostRequest::OpenSession { session_id, cell_execution_limits, + .. } => { let result = CodeModeSessionCellExecutionLimits::try_from( cell_execution_limits.unwrap_or_default(), @@ -452,6 +453,7 @@ impl HostState { HostRequest::Wait { session_id, request, + .. } => { let result = match self.session(&session_id) { Ok(session) => { @@ -476,15 +478,20 @@ impl HostState { HostRequest::Terminate { session_id, cell_id, + .. } => { let result = match self.session(&session_id) { - Ok(session) => session.terminate(cell_id.into()).await.and_then(|outcome| { - let outcome = outcome.with_code_mode_host_duration(received_at.elapsed()); - Ok(HostResponse::WaitCompleted { - outcome: WireWaitOutcome::try_from(outcome) - .map_err(|error| error.to_string())?, - }) - }), + Ok(session) => session + .terminate(codex_code_mode_protocol::CellId::from(cell_id).into()) + .await + .and_then(|outcome| { + let outcome = + outcome.with_code_mode_host_duration(received_at.elapsed()); + Ok(HostResponse::WaitCompleted { + outcome: WireWaitOutcome::try_from(outcome) + .map_err(|error| error.to_string())?, + }) + }), Err(err) => Err(err), }; self.respond(request_id, result); diff --git a/codex-rs/code-mode-host/tests/grpc.rs b/codex-rs/code-mode-host/tests/grpc.rs index acbf3af2f2..d83c30f69c 100644 --- a/codex-rs/code-mode-host/tests/grpc.rs +++ b/codex-rs/code-mode-host/tests/grpc.rs @@ -327,6 +327,7 @@ async fn grpc_early_yield_preserves_cell_and_later_waits() -> Result<()> { WaitRequest { cell_id: cell_id.clone(), yield_time_ms: 60_000, + max_output_tokens: None, }, Some(wait_signal.clone()), ); @@ -347,6 +348,7 @@ async fn grpc_early_yield_preserves_cell_and_later_waits() -> Result<()> { WaitRequest { cell_id: cell_id.clone(), yield_time_ms: 60_000, + max_output_tokens: None, }, /*preempt*/ None, ), @@ -397,13 +399,17 @@ async fn shutdown_immediately_rejects_new_operations() -> Result<()> { WaitRequest { cell_id: cell_id("missing"), yield_time_ms: 1, + max_output_tokens: None, }, /*preempt*/ None, ) .await, Err(expected.clone()) ); - assert_eq!(session.terminate(cell_id("missing")).await, Err(expected)); + assert_eq!( + session.terminate(cell_id("missing").into()).await, + Err(expected) + ); shutdown.await.map_err(anyhow::Error::msg)?; Ok(()) @@ -452,6 +458,7 @@ async fn cancelling_execution_before_admission_keeps_the_session_usable() -> Res WaitRequest { cell_id: abandoned_cell.clone(), yield_time_ms: 1, + max_output_tokens: None, }, /*preempt*/ None, ) @@ -512,6 +519,7 @@ async fn dropping_a_started_cell_off_runtime_terminates_its_buffered_remote_exec WaitRequest { cell_id: abandoned_cell.clone(), yield_time_ms: 1, + max_output_tokens: None, }, /*preempt*/ None, ) @@ -540,6 +548,7 @@ async fn dropping_a_started_cell_off_runtime_terminates_its_buffered_remote_exec WaitRequest { cell_id: abandoned_cell.clone(), yield_time_ms: 1, + max_output_tokens: None, }, /*preempt*/ None, ) @@ -751,6 +760,7 @@ async fn concurrent_wait_rejects_without_displacing_the_active_observer() -> Res WaitRequest { cell_id: running_cell.clone(), yield_time_ms: 100, + max_output_tokens: None, }, ) .await?; @@ -762,6 +772,7 @@ async fn concurrent_wait_rejects_without_displacing_the_active_observer() -> Res WaitRequest { cell_id: running_cell.clone(), yield_time_ms: 60_000, + max_output_tokens: None, }, /*preempt*/ None, ), @@ -785,7 +796,7 @@ async fn concurrent_wait_rejects_without_displacing_the_active_observer() -> Res }) ); let actual = session - .terminate(running_cell.clone()) + .terminate(running_cell.clone().into()) .await .map_err(anyhow::Error::msg)?; assert_eq!( @@ -838,6 +849,7 @@ async fn dropping_a_wait_retires_its_observer_before_the_next_wait() -> Result<( WaitRequest { cell_id: running_cell.clone(), yield_time_ms: 60_000, + max_output_tokens: None, }, ) .await?; @@ -850,6 +862,7 @@ async fn dropping_a_wait_retires_its_observer_before_the_next_wait() -> Result<( WaitRequest { cell_id: running_cell.clone(), yield_time_ms: 1, + max_output_tokens: None, }, /*preempt*/ None, ), @@ -866,7 +879,7 @@ async fn dropping_a_wait_retires_its_observer_before_the_next_wait() -> Result<( }) ); let actual = session - .terminate(running_cell.clone()) + .terminate(running_cell.clone().into()) .await .map_err(anyhow::Error::msg)?; assert_eq!( @@ -1091,6 +1104,7 @@ async fn sessions_enforce_independent_yield_limits() -> Result<()> { WaitRequest { cell_id: cell_id("1"), yield_time_ms: 60_000, + max_output_tokens: None, }, /*preempt*/ None, ), @@ -1117,7 +1131,7 @@ async fn sessions_enforce_independent_yield_limits() -> Result<()> { text_response("1", "isolated", actual.code_mode_host_duration()) ); let actual = limited - .terminate(cell_id("1")) + .terminate(cell_id("1").into()) .await .map_err(anyhow::Error::msg)?; assert_eq!( @@ -1156,6 +1170,7 @@ async fn dropping_a_grpc_lease_retires_its_server_session() -> Result<()> { let mut lease = client .open_session(grpc::OpenSessionRequest { cell_execution_limits: None, + output_truncation: None, }) .await .context("open raw gRPC session")? @@ -1236,6 +1251,7 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> { WaitRequest { cell_id: old_cell_id.clone(), yield_time_ms: 60_000, + max_output_tokens: None, }, ) .await?; @@ -1280,7 +1296,7 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> { .context("host loss did not retire the original generation's cell")?; let _replacement = HostHarness::start(&listen_url).await?; - assert!(session.terminate(old_cell_id.clone()).await.is_err()); + assert!(session.terminate(old_cell_id.clone().into()).await.is_err()); let mut rejected = request(r#"await tools.echo({must_not_run: true});"#); rejected.enabled_tools = vec![tool("echo")]; @@ -1291,6 +1307,7 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> { WaitRequest { cell_id: old_cell_id.clone(), yield_time_ms: 1, + max_output_tokens: None, }, /*preempt*/ None, ); @@ -1406,6 +1423,7 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> { WaitRequest { cell_id: replacement_cell_id.clone(), yield_time_ms: 1, + max_output_tokens: None, }, /*preempt*/ None, ) @@ -1420,7 +1438,7 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> { }) ); let actual = session - .terminate(replacement_cell_id.clone()) + .terminate(replacement_cell_id.clone().into()) .await .map_err(anyhow::Error::msg)?; assert_eq!( @@ -1437,13 +1455,14 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> { WaitRequest { cell_id: old_cell_id.clone(), yield_time_ms: 1, + max_output_tokens: None, }, /*preempt*/ None, ) .await .unwrap_err(); assert!(stale_wait.contains("stale code-mode host generation")); - let stale_termination = session.terminate(old_cell_id).await.unwrap_err(); + let stale_termination = session.terminate(old_cell_id.into()).await.unwrap_err(); assert!(stale_termination.contains("stale code-mode host generation")); session.shutdown().await.map_err(anyhow::Error::msg)?; Ok(()) diff --git a/codex-rs/code-mode-host/tests/grpc/admission_tests.rs b/codex-rs/code-mode-host/tests/grpc/admission_tests.rs index 6847bf8174..b9bc41646e 100644 --- a/codex-rs/code-mode-host/tests/grpc/admission_tests.rs +++ b/codex-rs/code-mode-host/tests/grpc/admission_tests.rs @@ -228,6 +228,7 @@ async fn routing_metadata_is_scoped_to_each_session_on_a_shared_provider() -> Re WaitRequest { cell_id: id.clone(), yield_time_ms: 1, + max_output_tokens: None, }, /*preempt*/ None, ) @@ -237,7 +238,10 @@ async fn routing_metadata_is_scoped_to_each_session_on_a_shared_provider() -> Re outcome, WaitOutcome::LiveCell(RuntimeResponse::Yielded { .. }) )); - session.terminate(id).await.map_err(anyhow::Error::msg)?; + session + .terminate(id.into()) + .await + .map_err(anyhow::Error::msg)?; } for session in [&parent, &child] { execute( diff --git a/codex-rs/code-mode-host/tests/grpc/image_cancellation_tests.rs b/codex-rs/code-mode-host/tests/grpc/image_cancellation_tests.rs index 28a160a466..863c11a0f1 100644 --- a/codex-rs/code-mode-host/tests/grpc/image_cancellation_tests.rs +++ b/codex-rs/code-mode-host/tests/grpc/image_cancellation_tests.rs @@ -57,7 +57,7 @@ try { image(value); } catch { text(calls); } .await .context("getter did not signal startup")?; - let result = timeout(Duration::from_secs(1), session.terminate(id.clone())) + let result = timeout(Duration::from_secs(1), session.terminate(id.clone().into())) .await .context("workload did not stop within the cancellation deadline")? .map_err(anyhow::Error::msg)?; diff --git a/codex-rs/code-mode-host/tests/grpc_notifications.rs b/codex-rs/code-mode-host/tests/grpc_notifications.rs index 42fac5a4bb..3cddef1310 100644 --- a/codex-rs/code-mode-host/tests/grpc_notifications.rs +++ b/codex-rs/code-mode-host/tests/grpc_notifications.rs @@ -225,6 +225,7 @@ async fn completed_waits_drain_pending_notifications_before_returning() -> Resul WaitRequest { cell_id: cell_id("1"), yield_time_ms: 5_000, + max_output_tokens: None, }, /*preempt*/ None, ) @@ -292,7 +293,7 @@ async fn termination_cancels_pending_notifications() -> Result<()> { } ); let actual = session - .terminate(cell_id("1")) + .terminate(cell_id("1").into()) .await .map_err(anyhow::Error::msg)?; assert_eq!( diff --git a/codex-rs/code-mode-host/tests/grpc_tcp.rs b/codex-rs/code-mode-host/tests/grpc_tcp.rs index a995119548..4daceb2d79 100644 --- a/codex-rs/code-mode-host/tests/grpc_tcp.rs +++ b/codex-rs/code-mode-host/tests/grpc_tcp.rs @@ -150,6 +150,7 @@ async fn tcp_listener_opens_a_grpc_session() -> Result<()> { let response = client .open_session(grpc::OpenSessionRequest { cell_execution_limits: None, + output_truncation: None, }) .await .context("failed to open gRPC code-mode session")?; diff --git a/codex-rs/code-mode-host/tests/grpc_timing.rs b/codex-rs/code-mode-host/tests/grpc_timing.rs index c59c17812f..835635bf6c 100644 --- a/codex-rs/code-mode-host/tests/grpc_timing.rs +++ b/codex-rs/code-mode-host/tests/grpc_timing.rs @@ -40,6 +40,7 @@ async fn execution_timing_includes_javascript_but_excludes_delayed_reads() -> Re let mut events = client .open_session(grpc::OpenSessionRequest { cell_execution_limits: None, + output_truncation: None, }) .await? .into_inner(); @@ -229,6 +230,7 @@ async fn observation_timing_excludes_previous_requests_and_background_time() -> WaitRequest { cell_id: cell_id.clone(), yield_time_ms: 50, + max_output_tokens: None, }, /*preempt*/ None, ) @@ -254,7 +256,7 @@ async fn observation_timing_excludes_previous_requests_and_background_time() -> let observed = Instant::now(); let terminated = session - .terminate(cell_id.clone()) + .terminate(cell_id.clone().into()) .await .map_err(anyhow::Error::msg)?; let duration = terminated @@ -276,6 +278,7 @@ async fn observation_timing_excludes_previous_requests_and_background_time() -> WaitRequest { cell_id: cell_id.clone(), yield_time_ms: 50, + max_output_tokens: None, }, /*preempt*/ None, ) diff --git a/codex-rs/code-mode-host/tests/settled.rs b/codex-rs/code-mode-host/tests/settled.rs index a724d9ec11..4aef9b3957 100644 --- a/codex-rs/code-mode-host/tests/settled.rs +++ b/codex-rs/code-mode-host/tests/settled.rs @@ -109,6 +109,7 @@ async fn assert_partial_results(provider: &dyn CodeModeSessionProvider) { WaitRequest { cell_id: cell_id.clone(), yield_time_ms: 10, + max_output_tokens: None, }, /*preempt*/ None, ) @@ -133,6 +134,7 @@ async fn assert_partial_results(provider: &dyn CodeModeSessionProvider) { WaitRequest { cell_id: cell_id.clone(), yield_time_ms: 5_000, + max_output_tokens: None, }, /*preempt*/ None, ) diff --git a/codex-rs/code-mode-host/tests/stdio.rs b/codex-rs/code-mode-host/tests/stdio.rs index ce53e6b9ad..80b2ed9018 100644 --- a/codex-rs/code-mode-host/tests/stdio.rs +++ b/codex-rs/code-mode-host/tests/stdio.rs @@ -206,6 +206,7 @@ async fn execute_to_terminal( WaitRequest { cell_id, yield_time_ms: 60_000, + max_output_tokens: None, }, /*preempt*/ None, ) @@ -270,6 +271,7 @@ async fn interrupt_yields_observations_without_stopping_the_cell() { let request = WaitRequest { cell_id: cell_id.clone(), yield_time_ms: 60_000, + max_output_tokens: None, }; let response = tokio::time::timeout(Duration::from_secs(5), session.wait(request, Some(signal))) @@ -288,6 +290,7 @@ async fn interrupt_yields_observations_without_stopping_the_cell() { WaitRequest { cell_id, yield_time_ms: 60_000, + max_output_tokens: None, }, /*preempt*/ None, ), @@ -350,6 +353,7 @@ async fn session_execution_limits_are_isolated_on_a_shared_process_host() { WaitRequest { cell_id: cell_id("1"), yield_time_ms: 60_000, + max_output_tokens: None, }, /*preempt*/ None, ), @@ -367,7 +371,7 @@ async fn session_execution_limits_are_isolated_on_a_shared_process_host() { ); limited - .terminate(cell_id("1")) + .terminate(cell_id("1").into()) .await .expect("terminate cell"); limited.shutdown().await.expect("shutdown limited session"); @@ -468,6 +472,7 @@ text(result.value); WaitRequest { cell_id: cell_id("3"), yield_time_ms: 1, + max_output_tokens: None, }, /*preempt*/ None, ) @@ -482,7 +487,7 @@ text(result.value); }) ); let actual = session - .terminate(cell_id("3")) + .terminate(cell_id("3").into()) .await .expect("terminate cell"); assert_eq!( @@ -546,6 +551,7 @@ async fn dropping_long_wait_releases_observer_before_next_wait() { WaitRequest { cell_id: wait_cell_id, yield_time_ms: 60_000, + max_output_tokens: None, }, /*preempt*/ None, ) @@ -561,6 +567,7 @@ async fn dropping_long_wait_releases_observer_before_next_wait() { WaitRequest { cell_id: running_cell_id.clone(), yield_time_ms: 1, + max_output_tokens: None, }, /*preempt*/ None, ), @@ -577,7 +584,7 @@ async fn dropping_long_wait_releases_observer_before_next_wait() { }) ); session - .terminate(running_cell_id) + .terminate(running_cell_id.into()) .await .expect("terminate cell"); session.shutdown().await.expect("shutdown remote session"); @@ -652,6 +659,7 @@ return; WaitRequest { cell_id: wait_cell_id, yield_time_ms: 60_000, + max_output_tokens: None, }, /*preempt*/ None, ) @@ -995,6 +1003,7 @@ async fn child_process_loss_cleans_up_and_rebuilds_the_shared_host() { WaitRequest { cell_id: wait_a_cell, yield_time_ms: 60_000, + max_output_tokens: None, }, /*preempt*/ None, ) @@ -1008,6 +1017,7 @@ async fn child_process_loss_cleans_up_and_rebuilds_the_shared_host() { WaitRequest { cell_id: wait_b_cell, yield_time_ms: 60_000, + max_output_tokens: None, }, /*preempt*/ None, ) @@ -1092,6 +1102,7 @@ async fn child_process_loss_cleans_up_and_rebuilds_the_shared_host() { WaitRequest { cell_id: cell_b.clone(), yield_time_ms: 1, + max_output_tokens: None, }, /*preempt*/ None, ) diff --git a/codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto b/codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto index 4e4134c21c..31db1f631b 100644 --- a/codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto +++ b/codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto @@ -45,6 +45,14 @@ service CodeModeHost { message OpenSessionRequest { optional SessionCellExecutionLimits cell_execution_limits = 1; + // Absent preserves the legacy client-side truncation behavior. + optional OutputTruncation output_truncation = 2; +} + +enum OutputTruncation { + OUTPUT_TRUNCATION_UNSPECIFIED = 0; + OUTPUT_TRUNCATION_CLIENT = 1; + OUTPUT_TRUNCATION_SERVER = 2; } message SessionCellExecutionLimits { @@ -64,6 +72,8 @@ message SessionEvent { message SessionOpened { string session_id = 1; + // Absent (including from an older host) means the client owns truncation. + optional OutputTruncation output_truncation = 2; } message CloseSessionRequest { @@ -182,6 +192,8 @@ message WaitRequest { string cell_id = 2; string wait_id = 3; uint64 yield_time_ms = 4; + // Only set for a session that selected server-side output truncation. + optional uint64 max_output_tokens = 5; } message WaitResponse { @@ -211,6 +223,8 @@ message YieldObservationResponse {} message TerminateRequest { string session_id = 1; string cell_id = 2; + // The terminating observation has its own output budget. + optional uint64 max_output_tokens = 3; } message ExecutionOutcome { diff --git a/codex-rs/code-mode-protocol/src/host/host_tests.rs b/codex-rs/code-mode-protocol/src/host/host_tests.rs index 01eedf3b28..47ae5907d1 100644 --- a/codex-rs/code-mode-protocol/src/host/host_tests.rs +++ b/codex-rs/code-mode-protocol/src/host/host_tests.rs @@ -27,6 +27,7 @@ use super::WireContentItem; use super::WireExecuteRequest; use super::WireImageDetail; use super::WireNestedToolCall; +use super::WireOutputTruncation; use super::WireResult; use super::WireRuntimeResponse; use super::WireSessionCellExecutionLimits; @@ -249,6 +250,7 @@ fn open_session_serializes_optional_cell_execution_limits() { max_yield_time_ms: Some(250), max_heap_size_bytes: Some(16 * 1024 * 1024), }), + output_truncation: None, }, json!({ "method": "session/open", @@ -304,9 +306,23 @@ fn client_to_host_v1_variants_are_pinned() { HostRequest::OpenSession { session_id: session_id(), cell_execution_limits: None, + output_truncation: None, }, json!({ "method": "session/open", "sessionId": "session-1" }), ), + ( + request_id(/*value*/ 10), + HostRequest::OpenSession { + session_id: session_id(), + cell_execution_limits: None, + output_truncation: Some(WireOutputTruncation::Server), + }, + json!({ + "method": "session/open", + "sessionId": "session-1", + "outputTruncation": "server", + }), + ), ( request_id(/*value*/ 2), HostRequest::Execute { @@ -353,6 +369,7 @@ fn client_to_host_v1_variants_are_pinned() { cell_id: cell_id("cell-1"), yield_time_ms: 50, }, + max_output_tokens: None, }, json!({ "method": "session/wait", @@ -365,6 +382,7 @@ fn client_to_host_v1_variants_are_pinned() { HostRequest::Terminate { session_id: session_id(), cell_id: cell_id("cell-1"), + max_output_tokens: None, }, json!({ "method": "session/terminate", diff --git a/codex-rs/code-mode-protocol/src/host/message.rs b/codex-rs/code-mode-protocol/src/host/message.rs index 39dd61315a..e6a4d8d28e 100644 --- a/codex-rs/code-mode-protocol/src/host/message.rs +++ b/codex-rs/code-mode-protocol/src/host/message.rs @@ -15,6 +15,7 @@ use super::SupportedProtocolVersions; use super::WireCellId; use super::WireExecuteRequest; use super::WireNestedToolCall; +use super::WireOutputTruncation; use super::WireRuntimeResponse; use super::WireSessionCellExecutionLimits; use super::WireWaitOutcome; @@ -186,6 +187,10 @@ pub enum HostRequest { session_id: SessionId, #[serde(default, skip_serializing_if = "Option::is_none")] cell_execution_limits: Option, + /// Sent only when output truncation support was agreed at connection hello. + /// Omission preserves the legacy session behavior and wire shape. + #[serde(default, skip_serializing_if = "Option::is_none")] + output_truncation: Option, }, #[serde(rename = "session/execute")] Execute { @@ -196,11 +201,15 @@ pub enum HostRequest { Wait { session_id: SessionId, request: WireWaitRequest, + #[serde(default, skip_serializing_if = "Option::is_none")] + max_output_tokens: Option, }, #[serde(rename = "session/terminate")] Terminate { session_id: SessionId, cell_id: WireCellId, + #[serde(default, skip_serializing_if = "Option::is_none")] + max_output_tokens: Option, }, #[serde(rename = "session/shutdown")] ShutdownSession { session_id: SessionId }, diff --git a/codex-rs/code-mode-protocol/src/host/mod.rs b/codex-rs/code-mode-protocol/src/host/mod.rs index af3f960fdd..b064cfc7f8 100644 --- a/codex-rs/code-mode-protocol/src/host/mod.rs +++ b/codex-rs/code-mode-protocol/src/host/mod.rs @@ -12,6 +12,12 @@ mod types; /// Maximum number of unresolved delegate callbacks allowed per host connection. pub const MAX_PENDING_DELEGATE_CALLS: usize = 1_024; +/// Stdio connection support for the optional `session/open.outputTruncation` field. +/// Each session selects its own setting after the connection hello. +pub const BOUNDED_OUTPUT_CAPABILITY: &str = "bounded-output-v1"; +/// Maximum retained bytes between server-truncated observations. +pub const BOUNDED_OUTPUT_MAX_BYTES: usize = 1024 * 1024; + /// Negotiated support for cell execution resource limits on `session/open`. pub const SESSION_RESOURCE_LIMITS_CAPABILITY: &str = "session-cell-execution-resource-limits"; @@ -38,6 +44,7 @@ pub use payload::WireContentItem; pub use payload::WireExecuteRequest; pub use payload::WireImageDetail; pub use payload::WireNestedToolCall; +pub use payload::WireOutputTruncation; pub use payload::WireRuntimeResponse; pub use payload::WireSessionCellExecutionLimits; pub use payload::WireToolDefinition; diff --git a/codex-rs/code-mode-protocol/src/host/payload.rs b/codex-rs/code-mode-protocol/src/host/payload.rs index ea24adb15b..53bbabb2be 100644 --- a/codex-rs/code-mode-protocol/src/host/payload.rs +++ b/codex-rs/code-mode-protocol/src/host/payload.rs @@ -8,6 +8,7 @@ use serde_json::Value as JsonValue; use crate::CellId; use crate::CodeModeNestedToolCall; +use crate::CodeModeOutputTruncation; use crate::CodeModeSessionCellExecutionLimits; use crate::CodeModeToolKind; use crate::ExecuteRequest; @@ -51,6 +52,32 @@ impl TryFrom for CodeModeSessionCellExecutionLim } } +/// The output truncation setting carried by a V1 session-open request. +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(rename_all = "camelCase")] +pub enum WireOutputTruncation { + Client, + Server, +} + +impl From for WireOutputTruncation { + fn from(value: CodeModeOutputTruncation) -> Self { + match value { + CodeModeOutputTruncation::Client => Self::Client, + CodeModeOutputTruncation::Server => Self::Server, + } + } +} + +impl From for CodeModeOutputTruncation { + fn from(value: WireOutputTruncation) -> Self { + match value { + WireOutputTruncation::Client => Self::Client, + WireOutputTruncation::Server => Self::Server, + } + } +} + /// A cell identifier with a wire representation owned by protocol V1. #[derive(Clone, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)] #[serde(transparent)] @@ -211,7 +238,8 @@ impl TryFrom for ExecuteRequest { } } -/// The complete wait request shape supported by protocol V1. +/// The nested wait payload for protocol V1. Its optional budget is carried by +/// the enclosing `HostRequest::Wait`. #[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] #[serde(deny_unknown_fields)] pub struct WireWaitRequest { @@ -233,6 +261,7 @@ impl From for WaitRequest { Self { cell_id: value.cell_id.into(), yield_time_ms: value.yield_time_ms, + max_output_tokens: None, } } } diff --git a/codex-rs/code-mode-protocol/src/lib.rs b/codex-rs/code-mode-protocol/src/lib.rs index 608b871b34..6ab492f39c 100644 --- a/codex-rs/code-mode-protocol/src/lib.rs +++ b/codex-rs/code-mode-protocol/src/lib.rs @@ -35,14 +35,17 @@ pub use runtime::ExecuteRequest; pub use runtime::ExecuteToPendingOutcome; pub use runtime::MissingCodeModeHostDuration; pub use runtime::RuntimeResponse; +pub use runtime::TerminateRequest; pub use runtime::WaitOutcome; pub use runtime::WaitRequest; pub use runtime::WaitToPendingOutcome; pub use runtime::WaitToPendingRequest; pub use session::CellId; +pub use session::CodeModeOutputTruncation; pub use session::CodeModeSession; pub use session::CodeModeSessionCellExecutionLimits; pub use session::CodeModeSessionDelegate; +pub use session::CodeModeSessionOptions; pub use session::CodeModeSessionProvider; pub use session::CodeModeSessionProviderFuture; pub use session::CodeModeSessionResultFuture; diff --git a/codex-rs/code-mode-protocol/src/runtime.rs b/codex-rs/code-mode-protocol/src/runtime.rs index 66334cfeb9..5996ac41b5 100644 --- a/codex-rs/code-mode-protocol/src/runtime.rs +++ b/codex-rs/code-mode-protocol/src/runtime.rs @@ -29,6 +29,26 @@ pub struct ExecuteRequest { pub struct WaitRequest { pub cell_id: CellId, pub yield_time_ms: u64, + /// Limit for this observation in Server mode; ignored in Client mode. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub max_output_tokens: Option, +} + +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +pub struct TerminateRequest { + pub cell_id: CellId, + /// Limit for the final observation in Server mode; ignored in Client mode. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub max_output_tokens: Option, +} + +impl From for TerminateRequest { + fn from(cell_id: CellId) -> Self { + Self { + cell_id, + max_output_tokens: None, + } + } } #[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] diff --git a/codex-rs/code-mode-protocol/src/session.rs b/codex-rs/code-mode-protocol/src/session.rs index 1c87cb8ea1..38655bc41e 100644 --- a/codex-rs/code-mode-protocol/src/session.rs +++ b/codex-rs/code-mode-protocol/src/session.rs @@ -12,6 +12,7 @@ use tokio_util::sync::CancellationToken; use crate::CodeModeNestedToolCall; use crate::ExecuteRequest; use crate::RuntimeResponse; +use crate::TerminateRequest; use crate::WaitOutcome; use crate::WaitRequest; @@ -30,6 +31,33 @@ pub struct CodeModeSessionCellExecutionLimits { pub max_heap_size_bytes: Option, } +/// The owner of the output budget for each exec or wait observation. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub enum CodeModeOutputTruncation { + /// Return raw output; the caller applies its own budget and formatting. + #[default] + Client, + /// The host applies the exec, wait, or terminate observation's budget before + /// transmitting output (10,000 tokens by default). Clients must not truncate + /// it again or append the script error: status and error metadata remain separate. + /// + /// Between observations, the host retains at most + /// [`crate::host::BOUNDED_OUTPUT_MAX_BYTES`] of output payload. It also + /// bounds the storage for retained items separately, without charging it + /// against the payload or observation token budgets. + /// Ordinary text notices report overflow, even for a zero-token budget. + /// Media are indivisible and can be omitted if they exceed the collector. + /// Without overflow, formatting matches Core's client policy. + Server, +} + +/// Settings shared by every cell in one code-mode session. +#[derive(Clone, Debug, Default, Eq, PartialEq)] +pub struct CodeModeSessionOptions { + pub cell_execution_limits: CodeModeSessionCellExecutionLimits, + pub output_truncation: CodeModeOutputTruncation, +} + #[derive(Clone, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)] pub struct CellId(String); @@ -152,6 +180,12 @@ impl CodeModeSessionDelegate for NoopCodeModeSessionDelegate { /// must keep those values isolated. Implementations may execute cells /// in-process or remotely. pub trait CodeModeSession: Send + Sync { + /// Identifies who applies output truncation for this session's responses. + /// Providers must not change modes when reconnecting to a new host. + fn output_truncation(&self) -> CodeModeOutputTruncation { + CodeModeOutputTruncation::Client + } + /// Executes a cell, yielding its foreground observation if `preempt` is signaled. /// The cell continues running after its observation yields. fn execute<'a>( @@ -162,13 +196,19 @@ pub trait CodeModeSession: Send + Sync { ) -> CodeModeSessionResultFuture<'a, StartedCell>; /// Waits for a cell, yielding the observation if `preempt` is signaled while the cell continues. + /// Server mode applies the request's output limit; Client mode ignores it. fn wait<'a>( &'a self, request: WaitRequest, preempt: Option, ) -> CodeModeSessionResultFuture<'a, WaitOutcome>; - fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome>; + /// Terminates a cell, applying the request's limit to its final observation + /// in Server mode; Client mode ignores the limit. + fn terminate<'a>( + &'a self, + request: TerminateRequest, + ) -> CodeModeSessionResultFuture<'a, WaitOutcome>; fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()>; } @@ -201,6 +241,25 @@ pub trait CodeModeSessionProvider: Send + Sync { }) } } + + /// Requests output handling for one session, without changing other sessions + /// created by this provider. Server mode requires host support. + fn create_session_with_options<'a>( + &'a self, + options: CodeModeSessionOptions, + ) -> CodeModeSessionProviderFuture<'a> { + match options.output_truncation { + CodeModeOutputTruncation::Client => { + self.create_session_with_limits(options.cell_execution_limits) + } + CodeModeOutputTruncation::Server => Box::pin(async { + Err( + "code-mode session provider does not support server output truncation" + .to_string(), + ) + }), + } + } } #[cfg(test)] diff --git a/codex-rs/code-mode-runtime/src/service.rs b/codex-rs/code-mode-runtime/src/service.rs index 9e7265562f..6cd5663dc0 100644 --- a/codex-rs/code-mode-runtime/src/service.rs +++ b/codex-rs/code-mode-runtime/src/service.rs @@ -15,6 +15,7 @@ use codex_code_mode_protocol::FunctionCallOutputContentItem; use codex_code_mode_protocol::ImageDetail; 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::WaitToPendingOutcome; @@ -129,6 +130,7 @@ impl InProcessCodeModeSession { preempt: Option, ) -> CodeModeSessionResultFuture<'static, WaitOutcome> { let WaitRequest { + max_output_tokens: _, cell_id, yield_time_ms, } = request; @@ -158,7 +160,8 @@ impl InProcessCodeModeSession { } } - pub async fn terminate(&self, cell_id: CellId) -> Result { + pub async fn terminate(&self, request: TerminateRequest) -> Result { + let cell_id = request.cell_id; match self.runtime.terminate(&runtime_cell_id(&cell_id)).await { Ok(event) => Ok(WaitOutcome::LiveCell(runtime_response(&cell_id, event)?)), Err(runtime::Error::MissingCell(_) | runtime::Error::ClosedCell(_)) => { @@ -239,8 +242,11 @@ impl CodeModeSession for InProcessCodeModeSession { Box::pin(InProcessCodeModeSession::wait(self, request, preempt)) } - fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome> { - Box::pin(InProcessCodeModeSession::terminate(self, cell_id)) + fn terminate<'a>( + &'a self, + request: TerminateRequest, + ) -> CodeModeSessionResultFuture<'a, WaitOutcome> { + Box::pin(InProcessCodeModeSession::terminate(self, request)) } fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()> { diff --git a/codex-rs/code-mode-runtime/src/service_audio_tests.rs b/codex-rs/code-mode-runtime/src/service_audio_tests.rs index aa1579c3b1..313656782a 100644 --- a/codex-rs/code-mode-runtime/src/service_audio_tests.rs +++ b/codex-rs/code-mode-runtime/src/service_audio_tests.rs @@ -149,6 +149,7 @@ async fn audio_helper_bounds_wav_data_and_preserves_outputs_after_yield() { let actual = service .wait( WaitRequest { + max_output_tokens: None, cell_id: cell_id("1"), yield_time_ms: 60_000, }, diff --git a/codex-rs/code-mode-runtime/src/service_contract_tests.rs b/codex-rs/code-mode-runtime/src/service_contract_tests.rs index 224a8893e1..de19e9c831 100644 --- a/codex-rs/code-mode-runtime/src/service_contract_tests.rs +++ b/codex-rs/code-mode-runtime/src/service_contract_tests.rs @@ -223,6 +223,7 @@ async fn yielded_cells_retain_their_own_delegate_until_closed() { let actual = service .wait( WaitRequest { + max_output_tokens: None, cell_id: cell_id("1"), yield_time_ms: 60_000, }, @@ -253,7 +254,7 @@ async fn yielded_cells_retain_their_own_delegate_until_closed() { ); assert!(weak_a.upgrade().is_none()); - service.terminate(cell_id("2")).await.unwrap(); + service.terminate(cell_id("2").into()).await.unwrap(); assert_eq!( next_event(&mut events_b).await, DelegateEvent::ToolCancelled @@ -293,6 +294,7 @@ async fn yields_and_resumes() { let actual = service .wait( WaitRequest { + max_output_tokens: None, cell_id: cell_id("1"), yield_time_ms: 60_000, }, @@ -418,7 +420,7 @@ async fn observed_natural_completion_wins_over_termination() { }) .await .unwrap(); - let actual = service.terminate(cell_id("1")).await.unwrap(); + let actual = service.terminate(cell_id("1").into()).await.unwrap(); let result_buffered_for = actual.result_buffered_for().unwrap(); assert!(!result_buffered_for.is_zero()); assert_eq!( @@ -461,7 +463,7 @@ async fn termination_cancels_pending_callbacks_before_responding() { } ); assert_eq!( - service.terminate(cell_id("1")).await.unwrap(), + service.terminate(cell_id("1").into()).await.unwrap(), WaitOutcome::LiveCell(RuntimeResponse::Terminated { code_mode_host_duration: None, cell_id: cell_id("1"), @@ -541,13 +543,13 @@ async fn repeated_termination_is_rejected_while_callback_cleanup_is_pending() { let terminating_service = Arc::clone(&service); let first_termination = - tokio::spawn(async move { terminating_service.terminate(cell_id("1")).await }); + tokio::spawn(async move { terminating_service.terminate(cell_id("1").into()).await }); assert_eq!( next_event(&mut events_rx).await, DelegateEvent::NotificationCancelled ); - let repeated_termination = service.terminate(cell_id("1")).await; + let repeated_termination = service.terminate(cell_id("1").into()).await; delegate.release_notification(); assert_eq!( @@ -594,6 +596,7 @@ async fn second_observer_is_rejected_without_displacing_the_first() { WaitRequest { cell_id: cell_id("1"), yield_time_ms: 60_000, + max_output_tokens: None, }, /*preempt*/ None, ) @@ -604,6 +607,7 @@ async fn second_observer_is_rejected_without_displacing_the_first() { WaitRequest { cell_id: cell_id("1"), yield_time_ms: 60_000, + max_output_tokens: None, }, /*preempt*/ None, ) @@ -618,7 +622,7 @@ async fn second_observer_is_rejected_without_displacing_the_first() { content_items: Vec::new(), }; assert_eq!( - service.terminate(cell_id("1")).await.unwrap(), + service.terminate(cell_id("1").into()).await.unwrap(), WaitOutcome::LiveCell(terminated.clone()) ); assert_eq!( diff --git a/codex-rs/code-mode-runtime/src/service_tests.rs b/codex-rs/code-mode-runtime/src/service_tests.rs index fd78ad7866..e239875c3b 100644 --- a/codex-rs/code-mode-runtime/src/service_tests.rs +++ b/codex-rs/code-mode-runtime/src/service_tests.rs @@ -98,6 +98,7 @@ async fn interrupt_yields_execute_and_wait_without_stopping_the_cell() { WaitRequest { cell_id: cell_id("1"), yield_time_ms: 60_000, + max_output_tokens: None, }, Some(wait_signal), ) @@ -125,6 +126,7 @@ async fn interrupt_yields_execute_and_wait_without_stopping_the_cell() { WaitRequest { cell_id: cell_id("1"), yield_time_ms: 60_000, + max_output_tokens: None, }, /*preempt*/ None, ) @@ -221,6 +223,7 @@ async fn execute_and_wait_clamp_yield_grace_without_stopping_the_cell() { WaitRequest { cell_id: cell_id("1"), yield_time_ms: 10_000, + max_output_tokens: None, }, /*preempt*/ None, ) @@ -243,6 +246,7 @@ async fn execute_and_wait_clamp_yield_grace_without_stopping_the_cell() { WaitRequest { cell_id: cell_id("1"), yield_time_ms: 10_000, + max_output_tokens: None, }, /*preempt*/ None, ) @@ -291,6 +295,7 @@ async fn wait_waits_for_nested_tool_during_yield_grace() { WaitRequest { cell_id: cell_id("1"), yield_time_ms: 10_000, + max_output_tokens: None, }, /*preempt*/ None, ) @@ -366,6 +371,7 @@ async fn zero_yield_limit_is_immediate_and_scoped_to_its_session() { WaitRequest { cell_id: cell_id("1"), yield_time_ms: 60_000, + max_output_tokens: None, }, /*preempt*/ None, ) @@ -441,7 +447,7 @@ async fn tool_start_acknowledgement_reports_termination_before_invocation() { let (startup, terminated) = tokio::join!( wait_until_tool_started(&delegate), - service.terminate(cell_id("1")), + service.terminate(cell_id("1").into()), ); assert_eq!(startup, Err("cell closed before nested tool started")); let response = RuntimeResponse::Terminated { @@ -746,6 +752,7 @@ text("after exit"); let actual = service .wait( WaitRequest { + max_output_tokens: None, cell_id: cell_id("1"), yield_time_ms: 30_000, }, @@ -1106,7 +1113,7 @@ async fn execute_to_pending_returns_once_the_runtime_is_quiescent() { } ); - let termination = service.terminate(cell_id("1")).await.unwrap(); + let termination = service.terminate(cell_id("1").into()).await.unwrap(); assert_eq!( termination, @@ -1150,7 +1157,7 @@ await Promise.all([ } ); - let termination = service.terminate(cell_id("1")).await.unwrap(); + let termination = service.terminate(cell_id("1").into()).await.unwrap(); assert_eq!( termination, @@ -1218,7 +1225,7 @@ await Promise.all([ }) ); - let termination = service.terminate(cell_id("1")).await.unwrap(); + let termination = service.terminate(cell_id("1").into()).await.unwrap(); assert_eq!( termination, @@ -1285,7 +1292,7 @@ await new Promise(() => {}); }) ); - let termination = service.terminate(cell_id("1")).await.unwrap(); + let termination = service.terminate(cell_id("1").into()).await.unwrap(); assert_eq!( termination, @@ -2132,6 +2139,7 @@ async fn wait_reports_missing_cell_separately_from_runtime_results() { WaitRequest { cell_id: cell_id("missing"), yield_time_ms: 1, + max_output_tokens: None, }, /*preempt*/ None, ) diff --git a/codex-rs/code-mode-runtime/tests/settled.rs b/codex-rs/code-mode-runtime/tests/settled.rs index d9ed6a5c7c..493a709277 100644 --- a/codex-rs/code-mode-runtime/tests/settled.rs +++ b/codex-rs/code-mode-runtime/tests/settled.rs @@ -242,7 +242,7 @@ async fn terminates_helpers_waiting_for_unsettled_inputs() { assert_eq!( tokio::time::timeout( Duration::from_secs(/*secs*/ 5), - session.terminate(cell_id.clone()) + session.terminate(cell_id.clone().into()) ) .await .expect("termination timed out") diff --git a/codex-rs/code-mode/src/grpc_session/mod.rs b/codex-rs/code-mode/src/grpc_session/mod.rs index 825e68cc64..c031e927b9 100644 --- a/codex-rs/code-mode/src/grpc_session/mod.rs +++ b/codex-rs/code-mode/src/grpc_session/mod.rs @@ -17,6 +17,7 @@ use codex_code_mode_protocol::CodeModeSessionProviderFuture; use codex_code_mode_protocol::CodeModeSessionResultFuture; use codex_code_mode_protocol::ExecuteRequest; use codex_code_mode_protocol::StartedCell; +use codex_code_mode_protocol::TerminateRequest; use codex_code_mode_protocol::WaitOutcome; use codex_code_mode_protocol::WaitRequest; use codex_code_mode_protocol::grpc; @@ -119,6 +120,7 @@ impl GrpcCodeModeSessionProvider { loop { let mut request = tonic::Request::new(grpc::OpenSessionRequest { cell_execution_limits, + output_truncation: None, }); inject_span_traceparent(&mut request, &open_session_span); let result = async { @@ -269,8 +271,11 @@ impl CodeModeSession for GrpcCodeModeSession { Box::pin(self.inner.wait(request, preempt)) } - fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome> { - Box::pin(self.inner.terminate(cell_id)) + fn terminate<'a>( + &'a self, + request: TerminateRequest, + ) -> CodeModeSessionResultFuture<'a, WaitOutcome> { + Box::pin(self.inner.terminate(request.cell_id)) } fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()> { diff --git a/codex-rs/code-mode/src/grpc_session/operations.rs b/codex-rs/code-mode/src/grpc_session/operations.rs index 56c4c48ef0..bd02e961d3 100644 --- a/codex-rs/code-mode/src/grpc_session/operations.rs +++ b/codex-rs/code-mode/src/grpc_session/operations.rs @@ -313,6 +313,7 @@ impl SessionInner { cell_id: expected_cell_id.as_str().to_string(), wait_id: wait_id.clone(), yield_time_ms: request.yield_time_ms, + max_output_tokens: None, }; let mut client = self.client(); let watcher = preempt.map(|signal| { @@ -368,6 +369,7 @@ impl SessionInner { client.terminate(self.request(grpc::TerminateRequest { session_id: self.id.clone(), cell_id: cell_id.as_str().to_string(), + max_output_tokens: None, })), ) .await?; diff --git a/codex-rs/code-mode/src/grpc_session/reconnect.rs b/codex-rs/code-mode/src/grpc_session/reconnect.rs index 5338a4f3c0..0b4dae97be 100644 --- a/codex-rs/code-mode/src/grpc_session/reconnect.rs +++ b/codex-rs/code-mode/src/grpc_session/reconnect.rs @@ -4,13 +4,13 @@ use std::sync::PoisonError; use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; -use codex_code_mode_protocol::CellId; use codex_code_mode_protocol::CodeModeSession; use codex_code_mode_protocol::CodeModeSessionCellExecutionLimits; use codex_code_mode_protocol::CodeModeSessionDelegate; use codex_code_mode_protocol::CodeModeSessionResultFuture; use codex_code_mode_protocol::ExecuteRequest; use codex_code_mode_protocol::StartedCell; +use codex_code_mode_protocol::TerminateRequest; use codex_code_mode_protocol::WaitOutcome; use codex_code_mode_protocol::WaitRequest; use tokio::sync::Semaphore; @@ -97,6 +97,7 @@ impl CodeModeSession for ReconnectableSession { Box::pin(async move { let binding = self.inner.binding_for_request(generation).await?; let request = WaitRequest { + max_output_tokens: request.max_output_tokens, cell_id: generation::remote_cell_id(binding.generation, &request.cell_id)?, yield_time_ms: request.yield_time_ms, }; @@ -105,15 +106,21 @@ impl CodeModeSession for ReconnectableSession { }) } - fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome> { + fn terminate<'a>( + &'a self, + request: TerminateRequest, + ) -> CodeModeSessionResultFuture<'a, WaitOutcome> { Box::pin(async move { // Cleanup must not replace the runtime and hide its loss from the next request. let binding = self .inner .binding() .ok_or_else(|| SHUTDOWN_ERROR.to_string())?; - let cell_id = generation::remote_cell_id(binding.generation, &cell_id)?; - let outcome = binding.session.terminate(cell_id).await?; + let cell_id = generation::remote_cell_id(binding.generation, &request.cell_id)?; + let outcome = binding + .session + .terminate(TerminateRequest { cell_id, ..request }) + .await?; Ok(generation::public_wait_outcome(binding.generation, outcome)) }) } diff --git a/codex-rs/code-mode/src/remote_session.rs b/codex-rs/code-mode/src/remote_session.rs index d58b083974..593417c505 100644 --- a/codex-rs/code-mode/src/remote_session.rs +++ b/codex-rs/code-mode/src/remote_session.rs @@ -6,7 +6,6 @@ use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; -use codex_code_mode_protocol::CellId; use codex_code_mode_protocol::CodeModeSession; use codex_code_mode_protocol::CodeModeSessionCellExecutionLimits; use codex_code_mode_protocol::CodeModeSessionDelegate; @@ -15,6 +14,7 @@ use codex_code_mode_protocol::CodeModeSessionProviderFuture; use codex_code_mode_protocol::CodeModeSessionResultFuture; use codex_code_mode_protocol::ExecuteRequest; use codex_code_mode_protocol::StartedCell; +use codex_code_mode_protocol::TerminateRequest; use codex_code_mode_protocol::WaitOutcome; use codex_code_mode_protocol::WaitRequest; use codex_code_mode_protocol::host::SessionId; @@ -286,9 +286,12 @@ impl ProcessOwnedCodeModeSession { .await } - pub async fn terminate(&self, cell_id: CellId) -> Result { + pub async fn terminate(&self, request: TerminateRequest) -> Result { let binding = self.connection().await?; - binding.connection.terminate(binding.remote, cell_id).await + binding + .connection + .terminate(binding.remote, request.cell_id) + .await } pub async fn shutdown(&self) -> Result<(), String> { @@ -549,8 +552,11 @@ impl CodeModeSession for ProcessOwnedCodeModeSession { Box::pin(ProcessOwnedCodeModeSession::wait(self, request, preempt)) } - fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome> { - Box::pin(ProcessOwnedCodeModeSession::terminate(self, cell_id)) + fn terminate<'a>( + &'a self, + request: TerminateRequest, + ) -> CodeModeSessionResultFuture<'a, WaitOutcome> { + Box::pin(ProcessOwnedCodeModeSession::terminate(self, request)) } fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()> { diff --git a/codex-rs/code-mode/src/remote_session/connection/driver/commands.rs b/codex-rs/code-mode/src/remote_session/connection/driver/commands.rs index bdbde44c1e..3c4e050418 100644 --- a/codex-rs/code-mode/src/remote_session/connection/driver/commands.rs +++ b/codex-rs/code-mode/src/remote_session/connection/driver/commands.rs @@ -112,6 +112,7 @@ impl ConnectionDriver { session_id: session.id.clone(), cell_execution_limits: (limits != WireSessionCellExecutionLimits::default()) .then_some(limits), + output_truncation: None, }, }; let frame = match EncodedFrame::encode(&message) { @@ -251,6 +252,7 @@ impl ConnectionDriver { HostRequest::Wait { session_id: session.id.clone(), request, + max_output_tokens: None, }, PendingRequest::Wait { session, @@ -284,6 +286,7 @@ impl ConnectionDriver { HostRequest::Terminate { session_id: session.id.clone(), cell_id, + max_output_tokens: None, }, PendingRequest::Terminate { session, diff --git a/codex-rs/code-mode/src/remote_session/connection/driver/responses.rs b/codex-rs/code-mode/src/remote_session/connection/driver/responses.rs index 6c95cc6e5d..f3bd836c12 100644 --- a/codex-rs/code-mode/src/remote_session/connection/driver/responses.rs +++ b/codex-rs/code-mode/src/remote_session/connection/driver/responses.rs @@ -432,6 +432,7 @@ impl ConnectionDriver { HostRequest::Terminate { session_id: session.id.clone(), cell_id: cell_id.clone(), + max_output_tokens: None, }, PendingRequest::Terminate { session, diff --git a/codex-rs/code-mode/src/remote_session/connection/driver_tests.rs b/codex-rs/code-mode/src/remote_session/connection/driver_tests.rs index 4442fa62af..99e65368d8 100644 --- a/codex-rs/code-mode/src/remote_session/connection/driver_tests.rs +++ b/codex-rs/code-mode/src/remote_session/connection/driver_tests.rs @@ -237,6 +237,7 @@ async fn old_host_does_not_receive_yield_frames() { WaitRequest { cell_id: CellId::new("1".to_string()), yield_time_ms: 60_000, + max_output_tokens: None, }, Some(signal), ) @@ -311,6 +312,7 @@ async fn open_session_includes_nondefault_cell_execution_limits() { max_yield_time_ms: Some(250), max_heap_size_bytes: Some(16 * 1024 * 1024), }), + output_truncation: None, }, } ); @@ -1275,6 +1277,7 @@ async fn mismatched_wait_response_fails_connection() { request: WaitRequest { cell_id: CellId::new("1".to_string()), yield_time_ms: 1, + max_output_tokens: None, }, caller_cancellation: CancellationToken::new(), yield_signal: None, @@ -1379,6 +1382,7 @@ async fn remote_wait_accepts_durations_longer_than_five_minutes() { request: WaitRequest { cell_id: CellId::new("1".to_string()), yield_time_ms: 300_001, + max_output_tokens: None, }, caller_cancellation: CancellationToken::new(), yield_signal: None, @@ -1448,6 +1452,7 @@ async fn queued_remote_wait_times_out_and_invalidates_the_connection() { WaitRequest { cell_id: CellId::new("1".to_string()), yield_time_ms: 1, + max_output_tokens: None, }, /*yield_signal*/ None, ) @@ -1534,6 +1539,7 @@ async fn cancelled_wait_is_retired_before_next_wait_is_sent() { request: WaitRequest { cell_id: CellId::new("1".to_string()), yield_time_ms: 60_000, + max_output_tokens: None, }, caller_cancellation: first_cancellation.clone(), yield_signal: None, @@ -1553,6 +1559,7 @@ async fn cancelled_wait_is_retired_before_next_wait_is_sent() { request: WaitRequest { cell_id: CellId::new("1".to_string()), yield_time_ms: 1, + max_output_tokens: None, }, caller_cancellation: CancellationToken::new(), yield_signal: None, @@ -2005,6 +2012,7 @@ async fn aborting_driver_marks_connection_dead_and_closes_cells() { request: WaitRequest { cell_id: CellId::new("1".to_string()), yield_time_ms: 60_000, + max_output_tokens: None, }, caller_cancellation: CancellationToken::new(), yield_signal: None, diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index 7c4172d9e5..78960de43f 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -158,7 +158,7 @@ impl CodeModeService { &self, cell_id: CellId, ) -> Result { - self.session().await?.terminate(cell_id).await + self.session().await?.terminate(cell_id.into()).await } pub(crate) async fn interrupt_active_cells(&self) { @@ -170,7 +170,7 @@ impl CodeModeService { .active_cell_ids() .into_iter() .map(|cell_id| async move { - if let Err(error) = session.terminate(cell_id.clone()).await { + if let Err(error) = session.terminate(cell_id.clone().into()).await { tracing::warn!(%cell_id, %error, "failed to terminate interrupted code-mode cell"); } }), diff --git a/codex-rs/core/src/tools/code_mode/wait_handler.rs b/codex-rs/core/src/tools/code_mode/wait_handler.rs index 6ad9265ff9..1623d4e427 100644 --- a/codex-rs/core/src/tools/code_mode/wait_handler.rs +++ b/codex-rs/core/src/tools/code_mode/wait_handler.rs @@ -140,6 +140,7 @@ impl CodeModeWaitHandler { codex_code_mode::WaitRequest { cell_id: cell_id.clone(), yield_time_ms: args.yield_time_ms, + max_output_tokens: None, }, step_context.preempt.clone(), ) diff --git a/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs b/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs index 4365be3df2..e5cb4a2a1a 100644 --- a/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs +++ b/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs @@ -90,14 +90,15 @@ impl codex_code_mode::CodeModeSession for MissingCellCodeModeSession { request: codex_code_mode::WaitRequest, _preempt: Option, ) -> codex_code_mode::CodeModeSessionResultFuture<'a, codex_code_mode::WaitOutcome> { - self.terminate(request.cell_id) + self.terminate(request.cell_id.into()) } fn terminate<'a>( &'a self, - cell_id: codex_code_mode::CellId, + request: codex_code_mode::TerminateRequest, ) -> codex_code_mode::CodeModeSessionResultFuture<'a, codex_code_mode::WaitOutcome> { Box::pin(async move { + let cell_id = request.cell_id; Ok(codex_code_mode::WaitOutcome::MissingCell( codex_code_mode::RuntimeResponse::Result { code_mode_host_duration: None, diff --git a/codex-rs/core/tests/code_mode_tracing.rs b/codex-rs/core/tests/code_mode_tracing.rs index 3e8a4a2b12..a6442c0ee7 100644 --- a/codex-rs/core/tests/code_mode_tracing.rs +++ b/codex-rs/core/tests/code_mode_tracing.rs @@ -176,6 +176,7 @@ async fn stdio_callbacks_preserve_context_across_overlapping_yielded_cells() { WaitRequest { cell_id, yield_time_ms: 5_000, + max_output_tokens: None, }, /*preempt*/ None, )