openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-21 17:48:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

Track code-mode cell lifecycle and result buffering time (#55653)

Browse Source 
## 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
This commit is contained in:
Ruslan Nigmatullin
2026-10-20 18:53:06 +00:00
committed by copyberry
parent 801bd4d773
commit 87f3c3284e
38 changed files with 397 additions and 123 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/code-mode-host/src/grpc/conversions.rs
+12 -2
View File
@@ -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<proto::WaitResponse, MissingCodeModeHostDuration> {
) -> Result<Response<proto::WaitResponse>, 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 {
codex-rs/code-mode-host/src/grpc/conversions_tests.rs
+40 -5
View File
@@ -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
);
}
}
codex-rs/code-mode-host/src/grpc/mod.rs
+2 -6
View File
@@ -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()))
}
}
codex-rs/code-mode-host/src/peer_tests.rs
+1
View File
@@ -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");
codex-rs/code-mode-host/tests/grpc.rs
+5
View File
@@ -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,
}
);
codex-rs/code-mode-host/tests/grpc_notifications.rs
+2
View File
@@ -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,
})
);
codex-rs/code-mode-host/tests/grpc_timing.rs
+1
View File
@@ -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(),
codex-rs/code-mode-host/tests/settled.rs
+3
View File
@@ -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,
})
);
codex-rs/code-mode-host/tests/stdio.rs
+9
View File
@@ -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,
}
);
codex-rs/code-mode-protocol/src/grpc/mod.rs
+4
View File
@@ -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";
codex-rs/code-mode-protocol/src/host/host_tests.rs
+2
View File
@@ -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,
},
},
codex-rs/code-mode-protocol/src/host/payload.rs
+7
View File
@@ -335,6 +335,8 @@ pub enum WireRuntimeResponse {
cell_id: WireCellId,
content_items: Vec<WireContentItem>,
error_text: Option<String>,
#[serde(default)]
result_buffered_for_ns: u64,
code_mode_host_duration_ns: u64,
},
}
@@ -377,6 +379,7 @@ impl TryFrom<RuntimeResponse> 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<RuntimeResponse> 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<WireRuntimeResponse> 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)),
},
}
codex-rs/code-mode-protocol/src/runtime.rs
+15
View File
@@ -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<Duration> {
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<Duration> {
match self {
@@ -113,6 +125,9 @@ pub enum RuntimeResponse {
cell_id: CellId,
content_items: Vec<FunctionCallOutputContentItem>,
error_text: Option<String>,
/// 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<Duration>,
},
codex-rs/code-mode-protocol/src/runtime_tests.rs
+1
View File
@@ -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,
},
codex-rs/code-mode-runtime/src/cell_actor/mod.rs
+2
View File
@@ -315,6 +315,7 @@ async fn run_cell<H: CellHost>(
.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<H: CellHost>(
.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
codex-rs/code-mode-runtime/src/cell_actor/tests.rs
+19 -9
View File
@@ -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!(
codex-rs/code-mode-runtime/src/cell_actor/types.rs
+43 -41
View File
@@ -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<Vec<OutputItem>>,
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<Vec<OutputItem>>,
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 {
codex-rs/code-mode-runtime/src/service.rs
+3
View File
@@ -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,
codex-rs/code-mode-runtime/src/service_audio_tests.rs
+13 -10
View File
@@ -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,
})
);
codex-rs/code-mode-runtime/src/service_contract_tests.rs
+31 -21
View File
@@ -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,
}
);
codex-rs/code-mode-runtime/src/service_tests.rs
+45 -10
View File
@@ -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()),
})
);
codex-rs/code-mode-runtime/src/session_runtime/tests.rs
+3
View File
@@ -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,
})
);
codex-rs/code-mode-runtime/src/session_runtime/types.rs
+2
View File
@@ -66,6 +66,8 @@ pub(crate) enum CellEvent {
Completed {
content_items: Vec<OutputItem>,
error_text: Option<String>,
/// Time the completed result spent waiting for an observer.
result_buffered_for: Duration,
},
Terminated {
content_items: Vec<OutputItem>,
codex-rs/code-mode-runtime/tests/array_sort.rs
+1
View File
@@ -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,
}
);
codex-rs/code-mode-runtime/tests/jit.rs
+1
View File
@@ -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,
}
);
codex-rs/code-mode/src/grpc_session/conversion.rs
+23 -5
View File
@@ -117,24 +117,42 @@ pub(super) fn runtime_response(outcome: grpc::ExecutionOutcome) -> Result<Runtim
cell_id,
content_items,
error_text: completed.error_text,
result_buffered_for: Duration::ZERO,
code_mode_host_duration,
},
};
Ok(response)
}
pub(super) fn wait_outcome(response: grpc::WaitResponse) -> Result<WaitOutcome, String> {
match response
pub(super) fn wait_outcome(
response: tonic::Response<grpc::WaitResponse>,
) -> Result<WaitOutcome, String> {
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<FunctionCallOutputContentItem, String> {
codex-rs/code-mode/src/grpc_session/conversion_tests.rs
+33 -4
View File
@@ -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))
);
}
}
codex-rs/code-mode/src/grpc_session/generation.rs
+2
View File
@@ -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,
},
codex-rs/code-mode/src/grpc_session/generation_tests.rs
+2
View File
@@ -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,
})
);
codex-rs/code-mode/src/grpc_session/operations.rs
+2 -3
View File
@@ -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
}
codex-rs/code-mode/src/remote_session/connection/driver/cell_ids.rs
+2
View File
@@ -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,
},
codex-rs/core/src/tools/code_mode/execute_handler.rs
+6 -1
View File
@@ -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
codex-rs/core/src/tools/code_mode/telemetry.rs
+15
View File
@@ -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 {
codex-rs/core/src/tools/code_mode/wait_handler.rs
+18 -6
View File
@@ -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);
codex-rs/core/src/tools/tool_dispatch_trace_tests.rs
+1
View File
@@ -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(),
codex-rs/ext/extension-api/src/contributors.rs
+13
View File
@@ -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.
codex-rs/ext/extension-api/src/contributors/tool_lifecycle.rs
+12
View File
@@ -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.
codex-rs/ext/extension-api/src/lib.rs
+1
View File
@@ -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;