From 87f3c3284e1669b90ad9f4925039e9e3bc769b61 Mon Sep 17 00:00:00 2001 From: Ruslan Nigmatullin Date: Tue, 20 Oct 2026 18:47:10 +0000 Subject: [PATCH] Track code-mode cell lifecycle and result buffering time (#55653) ## What changed - Add `CodeModeCellTiming` notifications for yielded executions, terminal waits, and unavailable cells, with thread, call, and cell identifiers. - Measure how long completed results remain buffered before retrieval and carry `result_buffered_for` through the runtime, stdio protocol, and gRPC response metadata. Immediate delivery reports zero; waits without a completed result report no buffering measurement. - Default missing buffering data to zero for compatibility with older hosts. ## Testing Extend runtime tests to verify buffering duration with a paused clock and retrieval through termination. Add gRPC conversion tests for nanosecond precision, optional metadata, and the distinction between immediate delivery and no completed result. GitOrigin-RevId: 60c0c5a3affa69ac3779bf1812fca50f5748af76 --- .../code-mode-host/src/grpc/conversions.rs | 14 +++- .../src/grpc/conversions_tests.rs | 45 ++++++++-- codex-rs/code-mode-host/src/grpc/mod.rs | 8 +- codex-rs/code-mode-host/src/peer_tests.rs | 1 + codex-rs/code-mode-host/tests/grpc.rs | 5 ++ .../tests/grpc_notifications.rs | 2 + codex-rs/code-mode-host/tests/grpc_timing.rs | 1 + codex-rs/code-mode-host/tests/settled.rs | 3 + codex-rs/code-mode-host/tests/stdio.rs | 9 ++ codex-rs/code-mode-protocol/src/grpc/mod.rs | 4 + .../code-mode-protocol/src/host/host_tests.rs | 2 + .../code-mode-protocol/src/host/payload.rs | 7 ++ codex-rs/code-mode-protocol/src/runtime.rs | 15 ++++ .../code-mode-protocol/src/runtime_tests.rs | 1 + .../code-mode-runtime/src/cell_actor/mod.rs | 2 + .../code-mode-runtime/src/cell_actor/tests.rs | 28 +++++-- .../code-mode-runtime/src/cell_actor/types.rs | 84 ++++++++++--------- codex-rs/code-mode-runtime/src/service.rs | 3 + .../src/service_audio_tests.rs | 23 ++--- .../src/service_contract_tests.rs | 52 +++++++----- .../code-mode-runtime/src/service_tests.rs | 55 +++++++++--- .../src/session_runtime/tests.rs | 3 + .../src/session_runtime/types.rs | 2 + .../code-mode-runtime/tests/array_sort.rs | 1 + codex-rs/code-mode-runtime/tests/jit.rs | 1 + .../code-mode/src/grpc_session/conversion.rs | 28 +++++-- .../src/grpc_session/conversion_tests.rs | 37 +++++++- .../code-mode/src/grpc_session/generation.rs | 2 + .../src/grpc_session/generation_tests.rs | 2 + .../code-mode/src/grpc_session/operations.rs | 5 +- .../connection/driver/cell_ids.rs | 2 + .../src/tools/code_mode/execute_handler.rs | 7 +- .../core/src/tools/code_mode/telemetry.rs | 15 ++++ .../core/src/tools/code_mode/wait_handler.rs | 24 ++++-- .../src/tools/tool_dispatch_trace_tests.rs | 1 + .../ext/extension-api/src/contributors.rs | 13 +++ .../src/contributors/tool_lifecycle.rs | 12 +++ codex-rs/ext/extension-api/src/lib.rs | 1 + 38 files changed, 397 insertions(+), 123 deletions(-) diff --git a/codex-rs/code-mode-host/src/grpc/conversions.rs b/codex-rs/code-mode-host/src/grpc/conversions.rs index 45272f62d4..6dbcb6fd9d 100644 --- a/codex-rs/code-mode-host/src/grpc/conversions.rs +++ b/codex-rs/code-mode-host/src/grpc/conversions.rs @@ -10,6 +10,7 @@ use codex_code_mode_protocol::WaitOutcome; use codex_code_mode_protocol::grpc as proto; use codex_protocol::ToolName; use serde_json::Value as JsonValue; +use tonic::Response; use tonic::Status; use super::validation; @@ -111,6 +112,7 @@ pub(super) fn execution_outcome( content_items, error_text, code_mode_host_duration, + .. } => ( cell_id, content_items, @@ -130,7 +132,8 @@ pub(super) fn execution_outcome( pub(super) fn wait_response( outcome: WaitOutcome, -) -> Result { +) -> Result, MissingCodeModeHostDuration> { + let buffered = outcome.result_buffered_for(); let state = match outcome { WaitOutcome::LiveCell(response) => { proto::wait_response::State::LiveCell(execution_outcome(response)?) @@ -139,7 +142,14 @@ pub(super) fn wait_response( proto::wait_response::State::MissingCell(execution_outcome(response)?) } }; - Ok(proto::WaitResponse { state: Some(state) }) + let mut response = Response::new(proto::WaitResponse { state: Some(state) }); + if let Some(buffered) = buffered { + let nanos = u64::try_from(buffered.as_nanos()).unwrap_or(u64::MAX); + response + .metadata_mut() + .insert(proto::RESULT_BUFFERED_FOR_NS_METADATA, nanos.into()); + } + Ok(response) } fn content_item(item: FunctionCallOutputContentItem) -> proto::ContentItem { diff --git a/codex-rs/code-mode-host/src/grpc/conversions_tests.rs b/codex-rs/code-mode-host/src/grpc/conversions_tests.rs index c8175d0a5a..0a0441ec17 100644 --- a/codex-rs/code-mode-host/src/grpc/conversions_tests.rs +++ b/codex-rs/code-mode-host/src/grpc/conversions_tests.rs @@ -69,7 +69,7 @@ fn rejects_missing_names_unknown_tool_kinds_and_invalid_json_schemas() { #[test] fn maps_text_image_audio_and_terminal_error_without_losing_details() { - let outcome = execution_outcome(RuntimeResponse::Result { + let response = wait_response(WaitOutcome::LiveCell(RuntimeResponse::Result { code_mode_host_duration: Some(Duration::from_nanos(/*nanos*/ 123_456_789)), cell_id: CellId::new("cell".to_string()), content_items: vec![ @@ -84,12 +84,22 @@ fn maps_text_image_audio_and_terminal_error_without_losing_details() { audio_url: "data:audio/wav;base64,YQ==".to_string(), }, ], + result_buffered_for: Duration::from_nanos(42), error_text: Some("failed".to_string()), - }); + })) + .unwrap(); assert_eq!( - outcome, - Ok(proto::ExecutionOutcome { + response + .metadata() + .get(proto::RESULT_BUFFERED_FOR_NS_METADATA) + .unwrap(), + "42" + ); + let state = response.into_inner().state.expect("wait outcome"); + assert_eq!( + state, + proto::wait_response::State::LiveCell(proto::ExecutionOutcome { code_mode_host_duration_ns: 123_456_789, cell_id: "cell".to_string(), content_items: vec![ @@ -119,6 +129,28 @@ fn maps_text_image_audio_and_terminal_error_without_losing_details() { ); } +#[test] +fn wait_metadata_distinguishes_an_unbuffered_result_from_no_result() { + let result = RuntimeResponse::Result { + cell_id: CellId::new("cell".to_string()), + content_items: Vec::new(), + error_text: None, + result_buffered_for: Duration::ZERO, + code_mode_host_duration: Some(Duration::ZERO), + }; + for (outcome, expected) in [ + (WaitOutcome::LiveCell(result.clone()), Some("0")), + (WaitOutcome::MissingCell(result), None), + ] { + let response = wait_response(outcome).unwrap(); + let buffered = response + .metadata() + .get(proto::RESULT_BUFFERED_FOR_NS_METADATA) + .map(|value| value.to_str().unwrap()); + assert_eq!(buffered, expected); + } +} + /// Encoding must not turn a missing request measurement into measured zero, /// including when no live cell remains to supply output. #[test] @@ -136,6 +168,9 @@ fn grpc_encoding_rejects_untimed_runtime_output() { WaitOutcome::LiveCell(response.clone()), WaitOutcome::MissingCell(response), ] { - assert_eq!(wait_response(outcome), Err(MissingCodeModeHostDuration)); + assert_eq!( + wait_response(outcome).unwrap_err(), + MissingCodeModeHostDuration + ); } } diff --git a/codex-rs/code-mode-host/src/grpc/mod.rs b/codex-rs/code-mode-host/src/grpc/mod.rs index a76f825e89..088077e33c 100644 --- a/codex-rs/code-mode-host/src/grpc/mod.rs +++ b/codex-rs/code-mode-host/src/grpc/mod.rs @@ -248,9 +248,7 @@ impl GrpcCodeModeHost { } }; let outcome = outcome.with_code_mode_host_duration(received_at.elapsed()); - let response = conversions::wait_response(outcome) - .map_err(|error| Status::internal(error.to_string()))?; - Ok(Response::new(response)) + conversions::wait_response(outcome).map_err(|error| Status::internal(error.to_string())) } async fn cancel_wait_request( @@ -292,9 +290,7 @@ impl GrpcCodeModeHost { let _permit = self.state.request_permit()?; let outcome = session.terminate(CellId::new(request.cell_id)).await?; let outcome = outcome.with_code_mode_host_duration(received_at.elapsed()); - let response = conversions::wait_response(outcome) - .map_err(|error| Status::internal(error.to_string()))?; - Ok(Response::new(response)) + conversions::wait_response(outcome).map_err(|error| Status::internal(error.to_string())) } } diff --git a/codex-rs/code-mode-host/src/peer_tests.rs b/codex-rs/code-mode-host/src/peer_tests.rs index 8bf2e90553..6aa9ddfbbd 100644 --- a/codex-rs/code-mode-host/src/peer_tests.rs +++ b/codex-rs/code-mode-host/src/peer_tests.rs @@ -48,6 +48,7 @@ async fn start_cell_reports_when_initial_response_is_enqueued() { code_mode_host_duration: None, cell_id: cell_id.clone(), content_items: Vec::new(), + result_buffered_for: Duration::ZERO, error_text: None, }) .expect("initial response receiver"); diff --git a/codex-rs/code-mode-host/tests/grpc.rs b/codex-rs/code-mode-host/tests/grpc.rs index c81890998d..acbf3af2f2 100644 --- a/codex-rs/code-mode-host/tests/grpc.rs +++ b/codex-rs/code-mode-host/tests/grpc.rs @@ -146,6 +146,7 @@ fn text_response( content_items: vec![FunctionCallOutputContentItem::InputText { text: value.to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, } } @@ -235,6 +236,7 @@ async fn tcp_session_persists_values_and_forwards_tools_notifications_and_closur code_mode_host_duration: actual.code_mode_host_duration(), cell_id: cell_id("1"), content_items: Vec::new(), + result_buffered_for: Duration::ZERO, error_text: None, } ); @@ -359,6 +361,7 @@ async fn grpc_early_yield_preserves_cell_and_later_waits() -> Result<()> { content_items: vec![FunctionCallOutputContentItem::InputText { text: "done".to_string(), }], + result_buffered_for: response.result_buffered_for().unwrap(), error_text: None, }) ); @@ -988,6 +991,7 @@ async fn large_unary_tool_completion_does_not_block_an_independent_session() -> content_items: vec![FunctionCallOutputContentItem::InputText { text: "8388608".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, } ); @@ -1043,6 +1047,7 @@ async fn single_subscription_processes_slow_and_fast_tools_concurrently() -> Res content_items: vec![FunctionCallOutputContentItem::InputText { text: "8388608".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, } ); diff --git a/codex-rs/code-mode-host/tests/grpc_notifications.rs b/codex-rs/code-mode-host/tests/grpc_notifications.rs index 9b9500192f..42fac5a4bb 100644 --- a/codex-rs/code-mode-host/tests/grpc_notifications.rs +++ b/codex-rs/code-mode-host/tests/grpc_notifications.rs @@ -113,6 +113,7 @@ fn text_response( content_items: vec![FunctionCallOutputContentItem::InputText { text: value.to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, } } @@ -247,6 +248,7 @@ async fn completed_waits_drain_pending_notifications_before_returning() -> Resul content_items: vec![FunctionCallOutputContentItem::InputText { text: "done".to_string(), }], + result_buffered_for: actual.result_buffered_for().unwrap(), error_text: None, }) ); diff --git a/codex-rs/code-mode-host/tests/grpc_timing.rs b/codex-rs/code-mode-host/tests/grpc_timing.rs index ec067a842e..c59c17812f 100644 --- a/codex-rs/code-mode-host/tests/grpc_timing.rs +++ b/codex-rs/code-mode-host/tests/grpc_timing.rs @@ -288,6 +288,7 @@ async fn observation_timing_excludes_previous_requests_and_background_time() -> assert_eq!( missing, WaitOutcome::MissingCell(RuntimeResponse::Result { + result_buffered_for: Duration::ZERO, error_text: Some(format!("exec cell {cell_id} not found")), cell_id, content_items: Vec::new(), diff --git a/codex-rs/code-mode-host/tests/settled.rs b/codex-rs/code-mode-host/tests/settled.rs index fff96761ff..a724d9ec11 100644 --- a/codex-rs/code-mode-host/tests/settled.rs +++ b/codex-rs/code-mode-host/tests/settled.rs @@ -146,6 +146,9 @@ async fn assert_partial_results(provider: &dyn CodeModeSessionProvider) { content_items: vec![FunctionCallOutputContentItem::InputText { text: json!({"index": 0, "status": "fulfilled", "value": "slow"}).to_string(), }], + result_buffered_for: completed + .result_buffered_for() + .expect("completed result includes buffering duration"), error_text: None, }) ); diff --git a/codex-rs/code-mode-host/tests/stdio.rs b/codex-rs/code-mode-host/tests/stdio.rs index e07582e0d4..ce53e6b9ad 100644 --- a/codex-rs/code-mode-host/tests/stdio.rs +++ b/codex-rs/code-mode-host/tests/stdio.rs @@ -401,6 +401,7 @@ async fn remote_session_persists_values_forwards_delegates_and_controls_cells() code_mode_host_duration: actual.code_mode_host_duration(), cell_id: cell_id("1"), content_items: Vec::new(), + result_buffered_for: Duration::ZERO, error_text: None, } ); @@ -431,6 +432,7 @@ text(result.value); content_items: vec![FunctionCallOutputContentItem::InputText { text: "output".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, } ); @@ -698,6 +700,7 @@ return; content_items: vec![FunctionCallOutputContentItem::InputText { text: "hello".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, }) ); @@ -741,6 +744,7 @@ async fn oversized_execute_request_does_not_close_the_shared_host() { content_items: vec![FunctionCallOutputContentItem::InputText { text: "still alive".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, } ); @@ -791,6 +795,7 @@ try {{ content_items: vec![FunctionCallOutputContentItem::InputText { text: "argument rejected".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, } ); @@ -820,6 +825,7 @@ try { content_items: vec![FunctionCallOutputContentItem::InputText { text: "result rejected".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, } ); @@ -838,6 +844,7 @@ try { content_items: vec![FunctionCallOutputContentItem::InputText { text: "still alive".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, } ); @@ -884,6 +891,7 @@ async fn oversized_initial_response_does_not_close_the_shared_host() { content_items: vec![FunctionCallOutputContentItem::InputText { text: "still alive".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, } ); @@ -1075,6 +1083,7 @@ async fn child_process_loss_cleans_up_and_rebuilds_the_shared_host() { content_items: vec![FunctionCallOutputContentItem::InputText { text: "replacement".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, } ); diff --git a/codex-rs/code-mode-protocol/src/grpc/mod.rs b/codex-rs/code-mode-protocol/src/grpc/mod.rs index 08c88badea..f575cab9d6 100644 --- a/codex-rs/code-mode-protocol/src/grpc/mod.rs +++ b/codex-rs/code-mode-protocol/src/grpc/mod.rs @@ -5,3 +5,7 @@ pub use code_mode_proto::codex::code_mode::v1::*; tonic::include_proto!("codex.code_mode.v1"); pub const MAX_IDENTIFIER_BYTES: usize = 256; + +/// Decimal nanoseconds a completed result was buffered before a wait or terminate retrieved it. +/// Zero means immediate delivery. Hosts omit this metadata if no completed result was retrieved. +pub const RESULT_BUFFERED_FOR_NS_METADATA: &str = "codex-code-mode-result-buffered-for-ns"; 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 591b29e5dd..01eedf3b28 100644 --- a/codex-rs/code-mode-protocol/src/host/host_tests.rs +++ b/codex-rs/code-mode-protocol/src/host/host_tests.rs @@ -488,6 +488,7 @@ fn host_to_client_v1_variants_are_pinned() { code_mode_host_duration_ns: 0, cell_id: cell_id("missing-cell"), content_items: Vec::new(), + result_buffered_for_ns: 0, error_text: Some("cell not found".to_string()), }), }, @@ -499,6 +500,7 @@ fn host_to_client_v1_variants_are_pinned() { "cell_id": "missing-cell", "content_items": [], "error_text": "cell not found", + "result_buffered_for_ns": 0, "code_mode_host_duration_ns": 0, }, }, diff --git a/codex-rs/code-mode-protocol/src/host/payload.rs b/codex-rs/code-mode-protocol/src/host/payload.rs index 869aeb4c61..ea24adb15b 100644 --- a/codex-rs/code-mode-protocol/src/host/payload.rs +++ b/codex-rs/code-mode-protocol/src/host/payload.rs @@ -335,6 +335,8 @@ pub enum WireRuntimeResponse { cell_id: WireCellId, content_items: Vec, error_text: Option, + #[serde(default)] + result_buffered_for_ns: u64, code_mode_host_duration_ns: u64, }, } @@ -377,6 +379,7 @@ impl TryFrom for WireRuntimeResponse { cell_id, content_items, error_text, + result_buffered_for, code_mode_host_duration, } => { let code_mode_host_duration = @@ -385,6 +388,8 @@ impl TryFrom for WireRuntimeResponse { cell_id: cell_id.into(), content_items: content_items.into_iter().map(Into::into).collect(), error_text, + result_buffered_for_ns: u64::try_from(result_buffered_for.as_nanos()) + .unwrap_or(u64::MAX), code_mode_host_duration_ns: u64::try_from(code_mode_host_duration.as_nanos()) .unwrap_or(u64::MAX), } @@ -418,11 +423,13 @@ impl From for RuntimeResponse { cell_id, content_items, error_text, + result_buffered_for_ns, code_mode_host_duration_ns, } => Self::Result { cell_id: cell_id.into(), content_items: content_items.into_iter().map(Into::into).collect(), error_text, + result_buffered_for: Duration::from_nanos(result_buffered_for_ns), code_mode_host_duration: Some(Duration::from_nanos(code_mode_host_duration_ns)), }, } diff --git a/codex-rs/code-mode-protocol/src/runtime.rs b/codex-rs/code-mode-protocol/src/runtime.rs index afa96ebf29..66334cfeb9 100644 --- a/codex-rs/code-mode-protocol/src/runtime.rs +++ b/codex-rs/code-mode-protocol/src/runtime.rs @@ -43,6 +43,18 @@ pub enum WaitOutcome { } impl WaitOutcome { + /// How long a completed result waited for retrieval. Zero if it was delivered immediately, + /// or `None` if this wait did not retrieve a completed result. + pub fn result_buffered_for(&self) -> Option { + match self { + Self::LiveCell(RuntimeResponse::Result { + result_buffered_for, + .. + }) => Some(*result_buffered_for), + Self::LiveCell(_) | Self::MissingCell(_) => None, + } + } + /// Returns timing for this wait or termination request, when supplied by its host. pub fn code_mode_host_duration(&self) -> Option { match self { @@ -113,6 +125,9 @@ pub enum RuntimeResponse { cell_id: CellId, content_items: Vec, error_text: Option, + /// Time between completion and retrieval when the result was buffered. + #[serde(default, skip_serializing_if = "Duration::is_zero")] + result_buffered_for: Duration, #[serde(skip_serializing_if = "Option::is_none")] code_mode_host_duration: Option, }, diff --git a/codex-rs/code-mode-protocol/src/runtime_tests.rs b/codex-rs/code-mode-protocol/src/runtime_tests.rs index 9f20067b2c..5e72577c3c 100644 --- a/codex-rs/code-mode-protocol/src/runtime_tests.rs +++ b/codex-rs/code-mode-protocol/src/runtime_tests.rs @@ -33,6 +33,7 @@ fn code_mode_host_duration_survives_runtime_and_stdio_serialization() { RuntimeResponse::Result { cell_id: CellId::new("completed-cell".to_string()), content_items, + result_buffered_for: Duration::ZERO, error_text: Some("execution failed".to_string()), code_mode_host_duration: None, }, diff --git a/codex-rs/code-mode-runtime/src/cell_actor/mod.rs b/codex-rs/code-mode-runtime/src/cell_actor/mod.rs index 4480e133da..1ccdb01c14 100644 --- a/codex-rs/code-mode-runtime/src/cell_actor/mod.rs +++ b/codex-rs/code-mode-runtime/src/cell_actor/mod.rs @@ -315,6 +315,7 @@ async fn run_cell( .await; let event = CellEvent::Completed { content_items: std::mem::take(&mut content_items), + result_buffered_for: std::time::Duration::ZERO, error_text: Some("exec runtime ended unexpectedly".to_string()), }; let rejected_event = match host @@ -471,6 +472,7 @@ async fn run_cell( .await; let event = CellEvent::Completed { content_items: std::mem::take(&mut content_items), + result_buffered_for: std::time::Duration::ZERO, error_text, }; let rejected_event = match host diff --git a/codex-rs/code-mode-runtime/src/cell_actor/tests.rs b/codex-rs/code-mode-runtime/src/cell_actor/tests.rs index afb7494cc5..e85ea96eb3 100644 --- a/codex-rs/code-mode-runtime/src/cell_actor/tests.rs +++ b/codex-rs/code-mode-runtime/src/cell_actor/tests.rs @@ -208,6 +208,7 @@ async fn runtime_thread_panic_remains_a_cell_error_without_owner_supervision() { harness.initial_event_rx.await.expect("initial event"), Ok(CellEvent::Completed { content_items: Vec::new(), + result_buffered_for: Duration::ZERO, error_text: Some("exec runtime ended unexpectedly".to_string()), }) ); @@ -469,11 +470,12 @@ async fn dropped_pending_observer_preserves_the_frontier_for_the_next_observatio harness.task.await.unwrap(); } -#[tokio::test] +#[tokio::test(start_paused = true)] async fn only_the_first_termination_claims_a_buffered_completion() { let cell_state = CellState::new(CancellationToken::new()); let completion = CellEvent::Completed { content_items: Vec::new(), + result_buffered_for: Duration::ZERO, error_text: None, }; assert_eq!( @@ -510,6 +512,7 @@ async fn termination_claim_prevents_stored_value_commit() { let mut commit_ran = false; let completion = CellEvent::Completed { content_items: Vec::new(), + result_buffered_for: Duration::ZERO, error_text: None, }; @@ -533,11 +536,12 @@ async fn termination_claim_prevents_stored_value_commit() { assert_eq!(termination.await, Ok(terminated)); } -#[test] -fn failed_completion_delivery_rebuffers_the_event() { +#[tokio::test(start_paused = true)] +async fn failed_completion_delivery_rebuffers_the_event() { let cell_state = CellState::new(CancellationToken::new()); let event = CellEvent::Completed { content_items: Vec::new(), + result_buffered_for: Duration::ZERO, error_text: None, }; assert_eq!( @@ -564,13 +568,14 @@ fn failed_completion_delivery_rebuffers_the_event() { assert_eq!(response_rx.try_recv(), Ok(Ok(event))); } -#[test] -fn buffered_initial_yield_precedes_buffered_completion_for_yield_observer() { +#[tokio::test(start_paused = true)] +async fn buffered_initial_yield_precedes_buffered_completion_for_yield_observer() { let cell_state = CellState::new(CancellationToken::new()); let completion = CellEvent::Completed { content_items: vec![OutputItem::Text { text: "after".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, }; assert_eq!( @@ -610,8 +615,8 @@ fn buffered_initial_yield_precedes_buffered_completion_for_yield_observer() { assert_eq!(response_rx.try_recv(), Ok(Ok(completion))); } -#[test] -fn pending_observer_merges_initial_yield_and_completion_output() { +#[tokio::test(start_paused = true)] +async fn pending_observer_merges_initial_yield_and_completion_output() { let cell_state = CellState::new(CancellationToken::new()); assert_eq!( cell_state.commit_completion( @@ -619,6 +624,7 @@ fn pending_observer_merges_initial_yield_and_completion_output() { content_items: vec![OutputItem::Text { text: "after".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, }, Some(vec![OutputItem::Text { @@ -633,6 +639,8 @@ fn pending_observer_merges_initial_yield_and_completion_output() { CompletionDelivery::Buffered )); + let buffered = Duration::from_secs(30); + tokio::time::advance(buffered).await; let (response_tx, mut response_rx) = oneshot::channel(); assert!(matches!( cell_state.route_observation(ObserveMode::PendingFrontier, response_tx), @@ -649,18 +657,20 @@ fn pending_observer_merges_initial_yield_and_completion_output() { text: "after".to_string(), }, ], + result_buffered_for: buffered, error_text: None, })) ); } -#[test] -fn dropped_pending_observation_preserves_the_initial_yield_boundary() { +#[tokio::test(start_paused = true)] +async fn dropped_pending_observation_preserves_the_initial_yield_boundary() { let cell_state = CellState::new(CancellationToken::new()); let completion = CellEvent::Completed { content_items: vec![OutputItem::Text { text: "after".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, }; assert_eq!( diff --git a/codex-rs/code-mode-runtime/src/cell_actor/types.rs b/codex-rs/code-mode-runtime/src/cell_actor/types.rs index 2aa343022a..3f6be17aac 100644 --- a/codex-rs/code-mode-runtime/src/cell_actor/types.rs +++ b/codex-rs/code-mode-runtime/src/cell_actor/types.rs @@ -3,6 +3,7 @@ use std::future::Future; use std::pin::Pin; use std::sync::Arc; use std::sync::Mutex; +use tokio::time::Instant; use serde_json::Value as JsonValue; use tokio::sync::mpsc; @@ -123,6 +124,7 @@ enum CellPhase { // Set only when `yield_control()` races the create-to-first-observe handoff. pending_initial_yield_items: Option>, event: CellEvent, + completed_at: Instant, }, CompletionClaimed(CellEvent), Tombstone, @@ -186,8 +188,9 @@ impl CellState { CellPhase::Completed { pending_initial_yield_items, event, + completed_at, } => { - let event = prepend_initial_yield(event, pending_initial_yield_items); + let event = buffered_completion(event, pending_initial_yield_items, completed_at); *phase = CellPhase::CompletionClaimed(event.clone()); self.cancellation_token.cancel(); ready_event(event) @@ -217,6 +220,7 @@ impl CellState { *phase = CellPhase::Completed { pending_initial_yield_items, event, + completed_at: Instant::now(), }; CompletionCommit::Committed } @@ -229,12 +233,13 @@ impl CellState { .phase .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); - let (pending_initial_yield_items, event) = + let (pending_initial_yield_items, event, completed_at) = match std::mem::replace(&mut *phase, CellPhase::Tombstone) { CellPhase::Completed { pending_initial_yield_items, event, - } => (pending_initial_yield_items, event), + completed_at, + } => (pending_initial_yield_items, event, completed_at), previous => { *phase = previous; return CompletionDelivery::Rejected(response_tx); @@ -244,6 +249,7 @@ impl CellState { *phase = CellPhase::Completed { pending_initial_yield_items, event, + completed_at, }; return CompletionDelivery::Buffered; }; @@ -256,6 +262,7 @@ impl CellState { *phase = CellPhase::Completed { pending_initial_yield_items, event, + completed_at, }; CompletionDelivery::Buffered } @@ -282,12 +289,14 @@ impl CellState { CellPhase::Completed { pending_initial_yield_items: Some(content_items), event, + completed_at, } if matches!(mode, ObserveMode::YieldAfter(_)) => { match response_tx.send(Ok(CellEvent::Yielded { content_items })) { Ok(()) => { *phase = CellPhase::Completed { pending_initial_yield_items: None, event, + completed_at, }; ObservationDelivery::Buffered } @@ -295,6 +304,7 @@ impl CellState { *phase = CellPhase::Completed { pending_initial_yield_items: Some(content_items), event, + completed_at, }; ObservationDelivery::Buffered } @@ -309,9 +319,13 @@ impl CellState { CellPhase::Completed { pending_initial_yield_items, event, + completed_at, } => { - let delivered_event = - prepend_initial_yield(event.clone(), pending_initial_yield_items.clone()); + let delivered_event = buffered_completion( + event.clone(), + pending_initial_yield_items.clone(), + completed_at, + ); match response_tx.send(Ok(delivered_event)) { Ok(()) => { self.cancellation_token.cancel(); @@ -321,6 +335,7 @@ impl CellState { *phase = CellPhase::Completed { pending_initial_yield_items, event, + completed_at, }; ObservationDelivery::Buffered } @@ -364,7 +379,12 @@ impl CellState { CellPhase::Completed { pending_initial_yield_items, event, - } => Some(prepend_initial_yield(event, pending_initial_yield_items)), + completed_at, + } => Some(buffered_completion( + event, + pending_initial_yield_items, + completed_at, + )), CellPhase::CompletionClaimed(completed_event) => Some(completed_event), CellPhase::Tombstone => None, }; @@ -385,47 +405,29 @@ impl CellState { } } -fn prepend_initial_yield( - event: CellEvent, +fn buffered_completion( + mut event: CellEvent, pending_initial_yield_items: Option>, + completed_at: Instant, ) -> CellEvent { - let Some(mut pending_initial_yield_items) = pending_initial_yield_items else { - return event; - }; - match event { - CellEvent::Yielded { mut content_items } => { - pending_initial_yield_items.append(&mut content_items); - CellEvent::Yielded { - content_items: pending_initial_yield_items, - } - } - CellEvent::Pending { - mut content_items, - pending_tool_call_ids, - } => { - pending_initial_yield_items.append(&mut content_items); - CellEvent::Pending { - content_items: pending_initial_yield_items, - pending_tool_call_ids, - } - } + let content_items = match &mut event { CellEvent::Completed { - mut content_items, - error_text, + content_items, + result_buffered_for, + .. } => { - pending_initial_yield_items.append(&mut content_items); - CellEvent::Completed { - content_items: pending_initial_yield_items, - error_text, - } - } - CellEvent::Terminated { mut content_items } => { - pending_initial_yield_items.append(&mut content_items); - CellEvent::Terminated { - content_items: pending_initial_yield_items, - } + *result_buffered_for = completed_at.elapsed(); + content_items } + CellEvent::Yielded { content_items } + | CellEvent::Pending { content_items, .. } + | CellEvent::Terminated { content_items } => content_items, + }; + if let Some(mut initial) = pending_initial_yield_items { + initial.append(content_items); + *content_items = initial; } + event } pub(super) enum CellCommand { diff --git a/codex-rs/code-mode-runtime/src/service.rs b/codex-rs/code-mode-runtime/src/service.rs index 37fa8506c7..9e7265562f 100644 --- a/codex-rs/code-mode-runtime/src/service.rs +++ b/codex-rs/code-mode-runtime/src/service.rs @@ -374,10 +374,12 @@ fn runtime_response( runtime::CellEvent::Completed { content_items, error_text, + result_buffered_for, } => Ok(RuntimeResponse::Result { cell_id: cell_id.clone(), content_items: content_items.into_iter().map(output_item).collect(), error_text, + result_buffered_for, code_mode_host_duration: None, }), runtime::CellEvent::Terminated { content_items } => Ok(RuntimeResponse::Terminated { @@ -414,6 +416,7 @@ fn output_item(item: runtime::OutputItem) -> FunctionCallOutputContentItem { fn missing_cell_response(cell_id: CellId) -> RuntimeResponse { RuntimeResponse::Result { error_text: Some(format!("exec cell {cell_id} not found")), + result_buffered_for: Duration::ZERO, cell_id, content_items: Vec::new(), code_mode_host_duration: None, 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 ec4070e074..aa1579c3b1 100644 --- a/codex-rs/code-mode-runtime/src/service_audio_tests.rs +++ b/codex-rs/code-mode-runtime/src/service_audio_tests.rs @@ -80,6 +80,7 @@ audio({{type:"audio",mimeType:"audio/wav",data:{payload:?}}});"# code_mode_host_duration: None, cell_id: cell_id("1"), content_items: vec![expected; 3], + result_buffered_for: std::time::Duration::ZERO, error_text: None, }, "sample rate {sample_rate}, frames {frames}" @@ -145,21 +146,23 @@ async fn audio_helper_bounds_wav_data_and_preserves_outputs_after_yield() { expected.push(FunctionCallOutputContentItem::InputText { text: "after".to_string(), }); + let actual = service + .wait( + WaitRequest { + cell_id: cell_id("1"), + yield_time_ms: 60_000, + }, + /*preempt*/ None, + ) + .await + .unwrap(); assert_eq!( - service - .wait( - WaitRequest { - cell_id: cell_id("1"), - yield_time_ms: 60_000, - }, - /*preempt*/ None, - ) - .await - .unwrap(), + actual, WaitOutcome::LiveCell(RuntimeResponse::Result { code_mode_host_duration: None, cell_id: cell_id("1"), content_items: expected, + result_buffered_for: actual.result_buffered_for().unwrap(), error_text: None, }) ); 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 1fb094cd6b..224a8893e1 100644 --- a/codex-rs/code-mode-runtime/src/service_contract_tests.rs +++ b/codex-rs/code-mode-runtime/src/service_contract_tests.rs @@ -220,21 +220,23 @@ async fn yielded_cells_retain_their_own_delegate_until_closed() { delegate_a.release_tool(); assert_eq!(next_event(&mut events_a).await, DelegateEvent::ToolStarted); delegate_a.release_tool(); + let actual = service + .wait( + WaitRequest { + cell_id: cell_id("1"), + yield_time_ms: 60_000, + }, + /*preempt*/ None, + ) + .await + .unwrap(); assert_eq!( - service - .wait( - WaitRequest { - cell_id: cell_id("1"), - yield_time_ms: 60_000, - }, - /*preempt*/ None, - ) - .await - .unwrap(), + actual, WaitOutcome::LiveCell(RuntimeResponse::Result { code_mode_host_duration: None, cell_id: cell_id("1"), content_items: Vec::new(), + result_buffered_for: actual.result_buffered_for().unwrap(), error_text: None }) ); @@ -288,23 +290,25 @@ async fn yields_and_resumes() { }], } ); + let actual = service + .wait( + WaitRequest { + cell_id: cell_id("1"), + yield_time_ms: 60_000, + }, + /*preempt*/ None, + ) + .await + .unwrap(); assert_eq!( - service - .wait( - WaitRequest { - cell_id: cell_id("1"), - yield_time_ms: 60_000, - }, - /*preempt*/ None, - ) - .await - .unwrap(), + actual, WaitOutcome::LiveCell(RuntimeResponse::Result { code_mode_host_duration: None, cell_id: cell_id("1"), content_items: vec![FunctionCallOutputContentItem::InputText { text: "after".to_string(), }], + result_buffered_for: actual.result_buffered_for().unwrap(), error_text: None, }) ); @@ -356,6 +360,7 @@ text("after"); content_items: vec![FunctionCallOutputContentItem::InputText { text: "after".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, } )) @@ -413,14 +418,18 @@ async fn observed_natural_completion_wins_over_termination() { }) .await .unwrap(); + let actual = service.terminate(cell_id("1")).await.unwrap(); + let result_buffered_for = actual.result_buffered_for().unwrap(); + assert!(!result_buffered_for.is_zero()); assert_eq!( - service.terminate(cell_id("1")).await.unwrap(), + actual, WaitOutcome::LiveCell(RuntimeResponse::Result { code_mode_host_duration: None, cell_id: cell_id("1"), content_items: vec![FunctionCallOutputContentItem::InputText { text: "done".to_string(), }], + result_buffered_for, error_text: None, }) ); @@ -645,6 +654,7 @@ async fn natural_completion_cleans_up_callbacks_before_responding() { content_items: vec![FunctionCallOutputContentItem::InputText { text: "done".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, } ); diff --git a/codex-rs/code-mode-runtime/src/service_tests.rs b/codex-rs/code-mode-runtime/src/service_tests.rs index 9fb77e4fe6..fd78ad7866 100644 --- a/codex-rs/code-mode-runtime/src/service_tests.rs +++ b/codex-rs/code-mode-runtime/src/service_tests.rs @@ -139,6 +139,7 @@ async fn interrupt_yields_execute_and_wait_without_stopping_the_cell() { content_items: vec![FunctionCallOutputContentItem::InputText { text: "done".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, }) ); @@ -173,6 +174,7 @@ async fn execute_waits_for_nested_tool_during_yield_grace() { content_items: vec![FunctionCallOutputContentItem::InputText { text: "done".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, } ); @@ -254,6 +256,7 @@ async fn execute_and_wait_clamp_yield_grace_without_stopping_the_cell() { content_items: vec![FunctionCallOutputContentItem::InputText { text: "done".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, }) ); @@ -306,6 +309,7 @@ async fn wait_waits_for_nested_tool_during_yield_grace() { content_items: vec![FunctionCallOutputContentItem::InputText { text: "done".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, }) ); @@ -415,6 +419,7 @@ async fn tool_start_acknowledgement_reports_error_before_invocation() { code_mode_host_duration: None, cell_id: cell_id("1"), content_items: Vec::new(), + result_buffered_for: Duration::ZERO, error_text: Some("startup failed".to_string()), } ); @@ -572,6 +577,7 @@ async fn synchronous_exit_returns_successfully() { content_items: vec![FunctionCallOutputContentItem::InputText { text: "before".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, } ); @@ -667,6 +673,7 @@ throw { content_items: vec![FunctionCallOutputContentItem::InputText { text: "before".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, }, "{label}" @@ -688,6 +695,7 @@ throw { content_items: vec![FunctionCallOutputContentItem::InputText { text: "before".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, }, "{label}: next cell" @@ -735,23 +743,25 @@ text("after exit"); } ); delegate.release_tool(); + let actual = service + .wait( + WaitRequest { + cell_id: cell_id("1"), + yield_time_ms: 30_000, + }, + /*preempt*/ None, + ) + .await + .unwrap(); assert_eq!( - service - .wait( - WaitRequest { - cell_id: cell_id("1"), - yield_time_ms: 30_000, - }, - /*preempt*/ None, - ) - .await - .unwrap(), + actual, WaitOutcome::LiveCell(RuntimeResponse::Result { code_mode_host_duration: None, cell_id: cell_id("1"), content_items: vec![FunctionCallOutputContentItem::InputText { text: "after tool".to_string(), }], + result_buffered_for: actual.result_buffered_for().unwrap(), error_text: None, }) ); @@ -770,6 +780,7 @@ text("after exit"); content_items: vec![FunctionCallOutputContentItem::InputText { text: "after tool".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, } ); @@ -815,6 +826,7 @@ async fn stored_values_are_shared_between_cells_but_not_sessions() { code_mode_host_duration: None, cell_id: cell_id("1"), content_items: Vec::new(), + result_buffered_for: Duration::ZERO, error_text: None, } ); @@ -826,6 +838,7 @@ async fn stored_values_are_shared_between_cells_but_not_sessions() { content_items: vec![FunctionCallOutputContentItem::InputText { text: "visible".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, } ); @@ -837,6 +850,7 @@ async fn stored_values_are_shared_between_cells_but_not_sessions() { content_items: vec![FunctionCallOutputContentItem::InputText { text: "undefined".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, } ); @@ -879,6 +893,7 @@ text(load("key")); text: "null".to_string(), }, ], + result_buffered_for: Duration::ZERO, error_text: None, } ); @@ -1055,6 +1070,7 @@ async fn execute_to_pending_returns_completed_for_synchronous_results() { content_items: vec![FunctionCallOutputContentItem::InputText { text: "done".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, }) ); @@ -1333,6 +1349,7 @@ text("done"); content_items: vec![FunctionCallOutputContentItem::InputText { text: "done".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, } )) @@ -1481,6 +1498,7 @@ async fn v8_console_is_not_exposed_on_global_this() { content_items: vec![FunctionCallOutputContentItem::InputText { text: "false".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, } ); @@ -1522,6 +1540,7 @@ text(value); content_items: vec![FunctionCallOutputContentItem::InputText { text: "jeudi 2 janvier \u{e0} 03:04:05".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, } ); @@ -1562,6 +1581,7 @@ text(formatter.format(new Date("2025-01-02T03:04:05Z"))); content_items: vec![FunctionCallOutputContentItem::InputText { text: "jeudi 2 janvier \u{e0} 03:04:05".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, } ); @@ -1610,6 +1630,7 @@ text(JSON.stringify(returnsUndefined)); text: "[true,true,true,true]".to_string(), }, ], + result_buffered_for: Duration::ZERO, error_text: None, } ); @@ -1637,6 +1658,7 @@ async fn text_helper_serializes_objects() { content_items: vec![FunctionCallOutputContentItem::InputText { text: r#"{"json":true}"#.to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, } ); @@ -1686,6 +1708,7 @@ async fn output_helpers_surface_serialization_errors() { code_mode_host_duration: None, cell_id: cell_id("1"), content_items: Vec::new(), + result_buffered_for: Duration::ZERO, error_text: Some(error_text), } ); @@ -1729,6 +1752,7 @@ audio({ audio_url: "data:audio/wav;base64,YXVkaW8=".to_string(), }, ], + result_buffered_for: Duration::ZERO, error_text: None, } ); @@ -1761,6 +1785,7 @@ async fn audio_helper_rejects_non_data_urls() { code_mode_host_duration: None, cell_id: cell_id("1"), content_items: Vec::new(), + result_buffered_for: Duration::ZERO, error_text: Some( "Tool call failed: invalid audio output. Pass a base64 data URI instead" .to_string(), @@ -1798,6 +1823,7 @@ image({ image_url: "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg==".to_string(), detail: Some(crate::ImageDetail::Original), }], + result_buffered_for: Duration::ZERO, error_text: None, } ); @@ -1837,6 +1863,7 @@ generatedImage({ text: "generated image save hint".to_string(), }, ], + result_buffered_for: Duration::ZERO, error_text: None, } ); @@ -1874,6 +1901,7 @@ image( image_url: "data:image/png;base64,AAA".to_string(), detail: Some(crate::ImageDetail::Original), }], + result_buffered_for: Duration::ZERO, error_text: None, } ); @@ -1910,6 +1938,7 @@ image( image_url: "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg==".to_string(), detail: Some(crate::ImageDetail::High), }], + result_buffered_for: Duration::ZERO, error_text: None, } ); @@ -1944,6 +1973,7 @@ image({ image_url: "data:image/png;base64,AAA".to_string(), detail: Some(crate::ImageDetail::Low), }], + result_buffered_for: Duration::ZERO, error_text: None, } ); @@ -1977,6 +2007,7 @@ async fn image_helpers_reject_remote_urls() { code_mode_host_duration: None, cell_id: cell_id("1"), content_items: Vec::new(), + result_buffered_for: Duration::ZERO, error_text: Some( "Tool call failed: remote image URLs are not supported in tool outputs. Pass a base64 data URI instead".to_string(), ), @@ -2012,6 +2043,7 @@ async fn image_helpers_reject_invalid_image_outputs() { code_mode_host_duration: None, cell_id: cell_id("1"), content_items: Vec::new(), + result_buffered_for: Duration::ZERO, error_text: Some( "Tool call failed: invalid image output. Pass a base64 data URI instead" .to_string(), @@ -2047,6 +2079,7 @@ image({ code_mode_host_duration: None, cell_id: cell_id("1"), content_items: Vec::new(), + result_buffered_for: Duration::ZERO, error_text: Some("image detail must be one of: auto, low, high, original".to_string()), } ); @@ -2082,6 +2115,7 @@ image({ code_mode_host_duration: None, cell_id: cell_id("1"), content_items: Vec::new(), + result_buffered_for: Duration::ZERO, error_text: Some( "image expects a non-empty image URL string, an object with image_url and optional detail, or a raw MCP image block".to_string(), ), @@ -2110,6 +2144,7 @@ async fn wait_reports_missing_cell_separately_from_runtime_results() { code_mode_host_duration: None, cell_id: cell_id("missing"), content_items: Vec::new(), + result_buffered_for: Duration::ZERO, error_text: Some("exec cell missing not found".to_string()), }) ); diff --git a/codex-rs/code-mode-runtime/src/session_runtime/tests.rs b/codex-rs/code-mode-runtime/src/session_runtime/tests.rs index dc4395a975..9d8a3f28a1 100644 --- a/codex-rs/code-mode-runtime/src/session_runtime/tests.rs +++ b/codex-rs/code-mode-runtime/src/session_runtime/tests.rs @@ -83,6 +83,7 @@ async fn reports_cell_actor_panics_to_the_owner() { content_items: vec![OutputItem::Text { text: "done".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, }) ); @@ -107,6 +108,7 @@ async fn termination_rejects_a_waiting_store_commit_before_the_next_cell_can_loa content_items: vec![OutputItem::Text { text: "uncommitted output".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, }; @@ -163,6 +165,7 @@ async fn termination_rejects_a_waiting_store_commit_before_the_next_cell_can_loa content_items: vec![OutputItem::Text { text: "undefined".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, }) ); diff --git a/codex-rs/code-mode-runtime/src/session_runtime/types.rs b/codex-rs/code-mode-runtime/src/session_runtime/types.rs index d1cdd79cee..eb2541cb1e 100644 --- a/codex-rs/code-mode-runtime/src/session_runtime/types.rs +++ b/codex-rs/code-mode-runtime/src/session_runtime/types.rs @@ -66,6 +66,8 @@ pub(crate) enum CellEvent { Completed { content_items: Vec, error_text: Option, + /// Time the completed result spent waiting for an observer. + result_buffered_for: Duration, }, Terminated { content_items: Vec, diff --git a/codex-rs/code-mode-runtime/tests/array_sort.rs b/codex-rs/code-mode-runtime/tests/array_sort.rs index d1a25a9b81..4a9782cabe 100644 --- a/codex-rs/code-mode-runtime/tests/array_sort.rs +++ b/codex-rs/code-mode-runtime/tests/array_sort.rs @@ -79,6 +79,7 @@ text(JSON.stringify([3, 1, 2].sort((a, b) => a - b))); content_items: vec![FunctionCallOutputContentItem::InputText { text: "[1,2,3]".to_string(), }], + result_buffered_for: std::time::Duration::ZERO, error_text: None, } ); diff --git a/codex-rs/code-mode-runtime/tests/jit.rs b/codex-rs/code-mode-runtime/tests/jit.rs index 44359005c4..2c235e86e1 100644 --- a/codex-rs/code-mode-runtime/tests/jit.rs +++ b/codex-rs/code-mode-runtime/tests/jit.rs @@ -38,6 +38,7 @@ async fn code_mode_runs_with_jit_disabled() { code_mode_host_duration: None, cell_id, content_items: Vec::new(), + result_buffered_for: std::time::Duration::ZERO, error_text: None, } ); diff --git a/codex-rs/code-mode/src/grpc_session/conversion.rs b/codex-rs/code-mode/src/grpc_session/conversion.rs index 3208ab56a7..d78f3f571a 100644 --- a/codex-rs/code-mode/src/grpc_session/conversion.rs +++ b/codex-rs/code-mode/src/grpc_session/conversion.rs @@ -117,24 +117,42 @@ pub(super) fn runtime_response(outcome: grpc::ExecutionOutcome) -> Result Result { - match response +pub(super) fn wait_outcome( + response: tonic::Response, +) -> Result { + let (metadata, response, _) = response.into_parts(); + let outcome = match response .state .ok_or_else(|| "code-mode wait omitted its outcome".to_string())? { grpc::wait_response::State::LiveCell(response) => { - runtime_response(response).map(WaitOutcome::LiveCell) + let mut response = runtime_response(response)?; + if let RuntimeResponse::Result { + result_buffered_for, + .. + } = &mut response + { + // Older hosts may omit this metadata, including for completed results. + *result_buffered_for = metadata + .get(grpc::RESULT_BUFFERED_FOR_NS_METADATA) + .and_then(|value| value.to_str().ok()?.parse().ok()) + .map(Duration::from_nanos) + .unwrap_or_default(); + } + WaitOutcome::LiveCell(response) } grpc::wait_response::State::MissingCell(response) => { - runtime_response(response).map(WaitOutcome::MissingCell) + WaitOutcome::MissingCell(runtime_response(response)?) } - } + }; + Ok(outcome) } fn content_item(item: grpc::ContentItem) -> Result { diff --git a/codex-rs/code-mode/src/grpc_session/conversion_tests.rs b/codex-rs/code-mode/src/grpc_session/conversion_tests.rs index 21e815ff4a..d3250f46f3 100644 --- a/codex-rs/code-mode/src/grpc_session/conversion_tests.rs +++ b/codex-rs/code-mode/src/grpc_session/conversion_tests.rs @@ -137,6 +137,7 @@ fn runtime_response_decodes_mixed_content_items() { audio_url: "data:audio/wav;base64,AA==".to_string(), }, ], + result_buffered_for: Duration::ZERO, error_text: Some("warning".to_string()), }) ); @@ -158,7 +159,7 @@ fn wait_outcome_preserves_missing_cell_state() { }; assert_eq!( - wait_outcome(response), + wait_outcome(tonic::Response::new(response)), Ok(WaitOutcome::MissingCell(RuntimeResponse::Terminated { code_mode_host_duration: Some(Duration::ZERO), cell_id: CellId::new("missing".to_string()), @@ -204,7 +205,7 @@ fn invalid_output_enums_and_missing_oneofs_are_rejected() { }; assert!(runtime_response(invalid_image).is_err()); - assert!(wait_outcome(grpc::WaitResponse { state: None }).is_err()); + assert!(wait_outcome(tonic::Response::new(grpc::WaitResponse { state: None })).is_err()); } /// Zero is a valid measurement, and decoding must not round nanoseconds to @@ -227,10 +228,38 @@ fn host_timing_preserves_zero_and_nanosecond_precision() { }; assert_eq!(runtime_response(outcome.clone()).as_ref(), Ok(&expected)); assert_eq!( - wait_outcome(grpc::WaitResponse { + wait_outcome(tonic::Response::new(grpc::WaitResponse { state: Some(grpc::wait_response::State::LiveCell(outcome)), - }), + })), Ok(WaitOutcome::LiveCell(expected)) ); } } + +#[test] +fn completed_wait_reads_optional_buffered_duration_from_response_metadata() { + for (metadata, expected_ns) in [(None, 0), (Some("0"), 0), (Some("42"), 42)] { + let mut response = tonic::Response::new(grpc::WaitResponse { + state: Some(grpc::wait_response::State::LiveCell( + grpc::ExecutionOutcome { + cell_id: "cell".to_string(), + outcome: Some(grpc::execution_outcome::Outcome::Completed( + grpc::ExecutionCompleted::default(), + )), + ..Default::default() + }, + )), + }); + if let Some(value) = metadata { + response.metadata_mut().insert( + grpc::RESULT_BUFFERED_FOR_NS_METADATA, + value.parse().unwrap(), + ); + } + let wait = wait_outcome(response).unwrap(); + assert_eq!( + wait.result_buffered_for(), + Some(Duration::from_nanos(expected_ns)) + ); + } +} diff --git a/codex-rs/code-mode/src/grpc_session/generation.rs b/codex-rs/code-mode/src/grpc_session/generation.rs index 59d6108eee..04ccb6314a 100644 --- a/codex-rs/code-mode/src/grpc_session/generation.rs +++ b/codex-rs/code-mode/src/grpc_session/generation.rs @@ -103,11 +103,13 @@ fn public_runtime_response(generation: u64, response: RuntimeResponse) -> Runtim RuntimeResponse::Result { cell_id, content_items, + result_buffered_for, error_text, code_mode_host_duration, } => RuntimeResponse::Result { cell_id: public_cell_id(generation, &cell_id), content_items, + result_buffered_for, error_text, code_mode_host_duration, }, diff --git a/codex-rs/code-mode/src/grpc_session/generation_tests.rs b/codex-rs/code-mode/src/grpc_session/generation_tests.rs index 2c9bc3d20e..ffab286643 100644 --- a/codex-rs/code-mode/src/grpc_session/generation_tests.rs +++ b/codex-rs/code-mode/src/grpc_session/generation_tests.rs @@ -177,6 +177,7 @@ async fn reconnected_execution_maps_started_and_initial_response_ids() { content_items: vec![FunctionCallOutputContentItem::InputText { text: "result".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, }; let started = StartedCell::from_future(wire_id, async move { @@ -199,6 +200,7 @@ async fn reconnected_execution_maps_started_and_initial_response_ids() { content_items: vec![FunctionCallOutputContentItem::InputText { text: "result".to_string(), }], + result_buffered_for: Duration::ZERO, error_text: None, }) ); diff --git a/codex-rs/code-mode/src/grpc_session/operations.rs b/codex-rs/code-mode/src/grpc_session/operations.rs index f79703c3c7..56c4c48ef0 100644 --- a/codex-rs/code-mode/src/grpc_session/operations.rs +++ b/codex-rs/code-mode/src/grpc_session/operations.rs @@ -331,7 +331,7 @@ impl SessionInner { drop(watcher); cancellation.disarm(); self.prune_wait_slots(); - let outcome = conversion::wait_outcome(response?.into_inner())?; + let outcome = conversion::wait_outcome(response?)?; self.validate_wait_cell(&expected_cell_id, outcome).await } @@ -370,8 +370,7 @@ impl SessionInner { cell_id: cell_id.as_str().to_string(), })), ) - .await? - .into_inner(); + .await?; let outcome = conversion::wait_outcome(response)?; self.validate_wait_cell(&cell_id, outcome).await } diff --git a/codex-rs/code-mode/src/remote_session/connection/driver/cell_ids.rs b/codex-rs/code-mode/src/remote_session/connection/driver/cell_ids.rs index 69b408bd05..bb8234818d 100644 --- a/codex-rs/code-mode/src/remote_session/connection/driver/cell_ids.rs +++ b/codex-rs/code-mode/src/remote_session/connection/driver/cell_ids.rs @@ -78,11 +78,13 @@ pub(super) fn public_runtime_response( RuntimeResponse::Result { cell_id, content_items, + result_buffered_for, error_text, code_mode_host_duration, } => RuntimeResponse::Result { cell_id: public_cell_id_from_protocol(generation, &cell_id), content_items, + result_buffered_for, error_text, code_mode_host_duration, }, diff --git a/codex-rs/core/src/tools/code_mode/execute_handler.rs b/codex-rs/core/src/tools/code_mode/execute_handler.rs index 25661c6c35..13e351398a 100644 --- a/codex-rs/core/src/tools/code_mode/execute_handler.rs +++ b/codex-rs/core/src/tools/code_mode/execute_handler.rs @@ -144,7 +144,12 @@ impl CodeModeExecuteHandler { code_cell_trace.record_initial_response(&response); // Yielded cells keep running, so terminal lifecycle is only emitted // here when the first response also ended the runtime. - if !matches!(response, codex_code_mode::RuntimeResponse::Yielded { .. }) { + if matches!(response, codex_code_mode::RuntimeResponse::Yielded { .. }) { + telemetry.record_cell_timing( + &cell_id, + codex_extension_api::CodeModeCellTiming::ExecYielded, + ); + } else { code_cell_trace.record_ended(&response); exec.session .services diff --git a/codex-rs/core/src/tools/code_mode/telemetry.rs b/codex-rs/core/src/tools/code_mode/telemetry.rs index 6e851aa356..910ae0a361 100644 --- a/codex-rs/core/src/tools/code_mode/telemetry.rs +++ b/codex-rs/core/src/tools/code_mode/telemetry.rs @@ -87,6 +87,21 @@ impl CodeModeToolCallGuard { "code-mode host operation completed" ); } + + pub(super) fn record_cell_timing( + &self, + cell_id: &codex_code_mode::CellId, + timing: codex_extension_api::CodeModeCellTiming, + ) { + for observer in self.extensions.tool_lifecycle_contributors() { + observer.on_code_mode_cell_timing( + &self.thread_id, + &self.call_id, + cell_id.as_str(), + timing, + ); + } + } } impl Drop for CodeModeToolCallGuard { 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 b0e4479134..6ad9265ff9 100644 --- a/codex-rs/core/src/tools/code_mode/wait_handler.rs +++ b/codex-rs/core/src/tools/code_mode/wait_handler.rs @@ -9,6 +9,8 @@ use crate::tools::registry::CoreToolRuntime; use crate::tools::registry::PostToolUsePayload; use crate::tools::registry::PreToolUsePayload; use crate::tools::registry::ToolExecutor; +use codex_code_mode::RuntimeResponse; +use codex_extension_api::CodeModeCellTiming; use codex_tools::ToolName; use codex_tools::ToolSpec; @@ -128,7 +130,7 @@ impl CodeModeWaitHandler { exec.session .services .code_mode_service - .terminate(cell_id) + .terminate(cell_id.clone()) .await } else { exec.session @@ -136,7 +138,7 @@ impl CodeModeWaitHandler { .code_mode_service .wait( codex_code_mode::WaitRequest { - cell_id, + cell_id: cell_id.clone(), yield_time_ms: args.yield_time_ms, }, step_context.preempt.clone(), @@ -149,9 +151,9 @@ impl CodeModeWaitHandler { })?; if let codex_code_mode::WaitOutcome::LiveCell(response) = &wait_response { let runtime_cell_id = match response { - codex_code_mode::RuntimeResponse::Yielded { cell_id, .. } - | codex_code_mode::RuntimeResponse::Terminated { cell_id, .. } - | codex_code_mode::RuntimeResponse::Result { cell_id, .. } => cell_id, + RuntimeResponse::Yielded { cell_id, .. } + | RuntimeResponse::Terminated { cell_id, .. } + | RuntimeResponse::Result { cell_id, .. } => cell_id, }; tracing::Span::current().record("cell.id", trace_id(runtime_cell_id.as_str())); telemetry.cell_id = Some(runtime_cell_id.to_string()); @@ -159,7 +161,15 @@ impl CodeModeWaitHandler { .services .executed_tool_calls .register_cell(runtime_cell_id, &call_id); - if !matches!(response, codex_code_mode::RuntimeResponse::Yielded { .. }) { + if !matches!(response, RuntimeResponse::Yielded { .. }) { + telemetry.record_cell_timing( + runtime_cell_id, + CodeModeCellTiming::WaitFinished { + result_buffered_for: wait_response + .result_buffered_for() + .unwrap_or_default(), + }, + ); exec.session .services .rollout_thread_trace @@ -183,6 +193,8 @@ impl CodeModeWaitHandler { }, ); } + } else { + telemetry.record_cell_timing(&cell_id, CodeModeCellTiming::CellUnavailable); } if let Some(code_mode_host_duration) = wait_response.code_mode_host_duration() { telemetry.record_code_mode_host_duration(code_mode_host_duration); 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 086ed3902c..df8c510778 100644 --- a/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs +++ b/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs @@ -101,6 +101,7 @@ impl codex_code_mode::CodeModeSession for MissingCellCodeModeSession { Ok(codex_code_mode::WaitOutcome::MissingCell( codex_code_mode::RuntimeResponse::Result { code_mode_host_duration: None, + result_buffered_for: std::time::Duration::ZERO, error_text: Some(format!("exec cell {cell_id} not found")), cell_id, content_items: Vec::new(), diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index 9ba83b18ea..c2c3e60c34 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -48,6 +48,7 @@ pub use thread_lifecycle::ThreadReadyInput; pub use thread_lifecycle::ThreadResumeInput; pub use thread_lifecycle::ThreadStartInput; pub use thread_lifecycle::ThreadStopInput; +pub use tool_lifecycle::CodeModeCellTiming; pub use tool_lifecycle::CommandStartInput; pub use tool_lifecycle::McpToolContext; pub use tool_lifecycle::McpToolResultInput; @@ -439,6 +440,18 @@ pub trait ToolLifecycleContributor: Send + Sync { /// Observe handler or remote host duration as defined by `ToolTimingBoundary`. /// Handler timing includes cancellation. Observers must return promptly. fn on_tool_timing(&self, _input: ToolTimingInput<'_>) {} + + /// Observe the exec that yielded a cell and the wait that collected its terminal result + /// or found the cell unavailable. + /// `call_id` identifies that dispatch. `cell_id` stays the same across waits. + fn on_code_mode_cell_timing( + &self, + _thread_id: &str, + _call_id: &str, + _cell_id: &str, + _timing: CodeModeCellTiming, + ) { + } } /// Owns the complete approval decision, including whether to consult a reviewer. diff --git a/codex-rs/ext/extension-api/src/contributors/tool_lifecycle.rs b/codex-rs/ext/extension-api/src/contributors/tool_lifecycle.rs index 3b31eaaf56..21633a185a 100644 --- a/codex-rs/ext/extension-api/src/contributors/tool_lifecycle.rs +++ b/codex-rs/ext/extension-api/src/contributors/tool_lifecycle.rs @@ -227,6 +227,18 @@ pub enum ToolTimingBoundary { HostOperation, } +/// Lifecycle of work that continues after a code-mode `exec` yields. +#[derive(Clone, Copy, Debug)] +pub enum CodeModeCellTiming { + ExecYielded, + WaitFinished { + /// Time the finished result was buffered before this wait retrieved it. + result_buffered_for: std::time::Duration, + }, + /// The runtime reported that the cell is missing or closed; its completion time is unknown. + CellUnavailable, +} + /// A measured tool interval. Handler observations also cover cancellation; /// host-operation observations are available only when the host reports a duration. /// Observers must not block. diff --git a/codex-rs/ext/extension-api/src/lib.rs b/codex-rs/ext/extension-api/src/lib.rs index 5256baf5e0..73767fe43e 100644 --- a/codex-rs/ext/extension-api/src/lib.rs +++ b/codex-rs/ext/extension-api/src/lib.rs @@ -49,6 +49,7 @@ pub use codex_tools::parse_tool_input_schema_without_compaction; pub use contributors::ApprovalDecision; pub use contributors::ApprovalDecisionInput; pub use contributors::ApprovalReviewContributor; +pub use contributors::CodeModeCellTiming; pub use contributors::CommandStartInput; pub use contributors::ConfigContributor; pub use contributors::ContextContributor;