From 3bb7ee893560335c562607264530b39e284f7fe2 Mon Sep 17 00:00:00 2001 From: Albin Cassirer Date: Mon, 20 Apr 2026 17:50:01 -0700 Subject: [PATCH] Accurately capture the time that custom.exec calls blocks the main thread. --- .../src/tools/code_mode/execute_handler.rs | 25 ++++--- codex-rs/trace/src/reduce.rs | 70 ++++++++++++++++--- 2 files changed, 73 insertions(+), 22 deletions(-) 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 aafe41a128..55641d5628 100644 --- a/codex-rs/core/src/tools/code_mode/execute_handler.rs +++ b/codex-rs/core/src/tools/code_mode/execute_handler.rs @@ -78,19 +78,18 @@ impl CodeModeExecuteHandler { return Err(FunctionCallError::RespondToModel(err)); } }; - if !matches!(&response, codex_code_mode::RuntimeResponse::Yielded { .. }) { - let result_payload = codex_trace::write_payload( - "code_cell_result", - &code_cell_response_payload(&response), - ); - emit_code_cell_ended( - &exec, - &runtime_cell_id, - &response, - /*model_visible_wait_call_id*/ None, - result_payload.as_ref(), - ); - } + let result_payload = + codex_trace::write_payload("code_cell_result", &code_cell_response_payload(&response)); + // A yielded response is the end of the model-visible custom `exec` + // call, not the end of the runtime cell. Emit it anyway so the trace + // can distinguish "the model got a cell id" from "the JS finished". + emit_code_cell_ended( + &exec, + &runtime_cell_id, + &response, + /*model_visible_wait_call_id*/ None, + result_payload.as_ref(), + ); handle_runtime_response(&exec, response, args.max_output_tokens, started_at) .await .map_err(FunctionCallError::RespondToModel) diff --git a/codex-rs/trace/src/reduce.rs b/codex-rs/trace/src/reduce.rs index 61b4d1b434..50cd33da74 100644 --- a/codex-rs/trace/src/reduce.rs +++ b/codex-rs/trace/src/reduce.rs @@ -501,13 +501,23 @@ impl Reducer<'_> { }; if let Some(cell) = self.trace.code_cells.get_mut(&code_cell_id) { let runtime_status = field_str(event, "status").unwrap_or("completed"); - set_execution_end( - &mut cell.execution, - event.wall_time_unix_ms, - event.seq, - code_cell_execution_status(runtime_status), - ); cell.runtime_status = code_cell_runtime_status(runtime_status); + if matches!(cell.runtime_status, CodeCellRuntimeStatus::Yielded) { + // Yielding is a partial result: the model-visible custom + // `exec` call has returned a cell id, but the runtime cell is + // still alive and may later complete through `wait`. + cell.initial_response_at_unix_ms = Some(event.wall_time_unix_ms); + cell.initial_response_seq = Some(event.seq); + cell.yielded_at_unix_ms = Some(event.wall_time_unix_ms); + cell.yielded_seq = Some(event.seq); + } else { + set_execution_end( + &mut cell.execution, + event.wall_time_unix_ms, + event.seq, + code_cell_execution_status(runtime_status), + ); + } if let Some(wait_call_id) = trace_field_str(event, "model_visible_wait_call", "id") .filter(|call_id| !call_id.is_empty()) { @@ -2892,6 +2902,14 @@ mod tests { temp.path().join("payloads/code-cell-invocation.json"), serde_json::to_vec(&json!({ "source_js": code_call["input"] }))?, )?; + std::fs::write( + temp.path().join("payloads/code-cell-yielded.json"), + serde_json::to_vec(&json!({ + "cell_id": "runtime-cell-1", + "status": "yielded", + "content_items": [{ "type": "input_text", "text": "Execution yielded. Call wait." }] + }))?, + )?; std::fs::write( temp.path().join("payloads/wait-invocation.json"), serde_json::to_vec(&json!({ @@ -2948,9 +2966,25 @@ mod tests { "raw_payload.invocation.kind": "code_cell_invocation" }), )?; + // The initial custom `exec` call can return before the cell finishes. + // This event records that yield point without closing the full runtime + // execution window. write_event( &mut events, 4, + "codex.code_cell.ended", + json!({ + "thread.id": "thread-1", + "code_cell.runtime_id": "runtime-cell-1", + "status": "yielded", + "raw_payload.result.id": "raw_payload:code-cell-yielded", + "raw_payload.result.path": "payloads/code-cell-yielded.json", + "raw_payload.result.kind": "code_cell_result" + }), + )?; + write_event( + &mut events, + 5, "codex.tool.started", json!({ "thread.id": "thread-1", @@ -2966,7 +3000,7 @@ mod tests { )?; write_event( &mut events, - 5, + 6, "codex.tool.ended", json!({ "tool.call_id": "tool:call-wait", @@ -2979,7 +3013,7 @@ mod tests { )?; write_event( &mut events, - 6, + 7, "codex.code_cell.ended", json!({ "thread.id": "thread-1", @@ -2990,7 +3024,7 @@ mod tests { )?; write_event( &mut events, - 7, + 8, "codex.inference.started", json!({ "thread.id": "thread-1", @@ -3006,11 +3040,29 @@ mod tests { assert_eq!( json!({ + "execution": trace["code_cells"]["code_cell:call-code"]["execution"].clone(), + "runtime_status": trace["code_cells"]["code_cell:call-code"]["runtime_status"].clone(), + "initial_response_at_unix_ms": trace["code_cells"]["code_cell:call-code"]["initial_response_at_unix_ms"].clone(), + "initial_response_seq": trace["code_cells"]["code_cell:call-code"]["initial_response_seq"].clone(), + "yielded_at_unix_ms": trace["code_cells"]["code_cell:call-code"]["yielded_at_unix_ms"].clone(), + "yielded_seq": trace["code_cells"]["code_cell:call-code"]["yielded_seq"].clone(), "output_item_ids": trace["code_cells"]["code_cell:call-code"]["output_item_ids"].clone(), "wait_tool_call_ids": trace["code_cells"]["code_cell:call-code"]["wait_tool_call_ids"].clone(), "wait_output_producers": trace["conversation_items"]["item:4"]["produced_by"].clone(), }), json!({ + "execution": { + "started_at_unix_ms": 1776420000000i64, + "ended_at_unix_ms": 1776420000000i64, + "status": "completed", + "started_seq": 3, + "ended_seq": 7, + }, + "runtime_status": "completed", + "initial_response_at_unix_ms": 1776420000000i64, + "initial_response_seq": 4, + "yielded_at_unix_ms": 1776420000000i64, + "yielded_seq": 4, "output_item_ids": ["item:2", "item:4"], "wait_tool_call_ids": ["tool:call-wait"], "wait_output_producers": [