From 03573b051e96e6d028fbed4e4ec76889f6ba3ae4 Mon Sep 17 00:00:00 2001 From: Jeremy Rose Date: Fri, 1 May 2026 14:54:29 -0700 Subject: [PATCH] derive running command durations in history builder --- .../src/protocol/thread_history.rs | 55 ++++++++++++++++++- .../app-server/src/codex_message_processor.rs | 11 ---- codex-rs/app-server/src/thread_state.rs | 52 +----------------- 3 files changed, 56 insertions(+), 62 deletions(-) diff --git a/codex-rs/app-server-protocol/src/protocol/thread_history.rs b/codex-rs/app-server-protocol/src/protocol/thread_history.rs index 64307c24bf..44d80bcdc0 100644 --- a/codex-rs/app-server-protocol/src/protocol/thread_history.rs +++ b/codex-rs/app-server-protocol/src/protocol/thread_history.rs @@ -54,6 +54,8 @@ use codex_protocol::protocol::ViewImageToolCallEvent; use codex_protocol::protocol::WebSearchBeginEvent; use codex_protocol::protocol::WebSearchEndEvent; use std::collections::HashMap; +use std::time::SystemTime; +use std::time::UNIX_EPOCH; use tracing::warn; use uuid::Uuid; @@ -82,9 +84,45 @@ pub fn build_turns_from_rollout_items(items: &[RolloutItem]) -> Vec { builder.finish() } +fn now_unix_timestamp_ms() -> i64 { + let millis = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis(); + i64::try_from(millis).unwrap_or(i64::MAX) +} + +fn elapsed_duration_ms(started_at_ms: i64) -> i64 { + now_unix_timestamp_ms().saturating_sub(started_at_ms) +} + +fn materialize_running_command_durations( + running_command_started_at_ms: &HashMap, + turns: &mut [Turn], +) { + for turn in turns { + for item in &mut turn.items { + let ThreadItem::CommandExecution { + id, + status: CommandExecutionStatus::InProgress, + duration_ms, + .. + } = item + else { + continue; + }; + let Some(started_at_ms) = running_command_started_at_ms.get(id) else { + continue; + }; + *duration_ms = Some(elapsed_duration_ms(*started_at_ms)); + } + } +} + pub struct ThreadHistoryBuilder { turns: Vec, current_turn: Option, + running_command_started_at_ms: HashMap, next_item_index: i64, current_rollout_index: usize, next_rollout_index: usize, @@ -101,6 +139,7 @@ impl ThreadHistoryBuilder { Self { turns: Vec::new(), current_turn: None, + running_command_started_at_ms: HashMap::new(), next_item_index: 1, current_rollout_index: 0, next_rollout_index: 0, @@ -113,14 +152,21 @@ impl ThreadHistoryBuilder { pub fn finish(mut self) -> Vec { self.finish_current_turn(); + materialize_running_command_durations(&self.running_command_started_at_ms, &mut self.turns); self.turns } pub fn active_turn_snapshot(&self) -> Option { - self.current_turn + let mut turn = self + .current_turn .as_ref() .map(Turn::from) - .or_else(|| self.turns.last().cloned()) + .or_else(|| self.turns.last().cloned())?; + materialize_running_command_durations( + &self.running_command_started_at_ms, + std::slice::from_mut(&mut turn), + ); + Some(turn) } /// Returns the index of the active turn snapshot within the finished turn list. @@ -405,11 +451,16 @@ impl ThreadHistoryBuilder { } fn handle_exec_command_begin(&mut self, payload: &ExecCommandBeginEvent) { + self.running_command_started_at_ms.insert( + payload.call_id.clone(), + payload.started_at_ms.unwrap_or_else(now_unix_timestamp_ms), + ); let item = build_command_execution_begin_item(payload); self.upsert_item_in_turn_id(&payload.turn_id, item); } fn handle_exec_command_end(&mut self, payload: &ExecCommandEndEvent) { + self.running_command_started_at_ms.remove(&payload.call_id); let item = build_command_execution_end_item(payload); // Command completions can arrive out of order. Unified exec may return // while a PTY is still running, then emit ExecCommandEnd later from a diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index 42e19b6d41..441a80e254 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -4006,12 +4006,6 @@ impl CodexMessageProcessor { } else { false }; - if include_turns && has_live_in_progress_turn { - let thread_state = self.thread_state_manager.thread_state(thread_id).await; - let state = thread_state.lock().await; - state.refresh_running_command_durations(&mut thread.turns); - } - let thread_status = self .thread_watch_manager .loaded_status_for_thread(&thread.id) @@ -8432,11 +8426,6 @@ async fn handle_pending_thread_resume_request( .await; return; } - { - let state = thread_state.lock().await; - state.refresh_running_command_durations(&mut thread.turns); - } - let thread_status = thread_watch_manager .loaded_status_for_thread(&thread.id) .await; diff --git a/codex-rs/app-server/src/thread_state.rs b/codex-rs/app-server/src/thread_state.rs index f99d64f466..dddbcf483b 100644 --- a/codex-rs/app-server/src/thread_state.rs +++ b/codex-rs/app-server/src/thread_state.rs @@ -1,11 +1,8 @@ use crate::outgoing_message::ConnectionId; use crate::outgoing_message::ConnectionRequestId; -use chrono::Utc; -use codex_app_server_protocol::CommandExecutionStatus; use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ThreadGoal; use codex_app_server_protocol::ThreadHistoryBuilder; -use codex_app_server_protocol::ThreadItem; use codex_app_server_protocol::Turn; use codex_app_server_protocol::TurnError; use codex_core::CodexThread; @@ -27,14 +24,6 @@ use tracing::error; type PendingInterruptQueue = Vec; -fn now_unix_timestamp_ms() -> i64 { - Utc::now().timestamp_millis() -} - -fn elapsed_duration_ms(started_at_ms: i64) -> i64 { - now_unix_timestamp_ms().saturating_sub(started_at_ms) -} - pub(crate) struct PendingThreadResumeRequest { pub(crate) request_id: ConnectionRequestId, pub(crate) history_items: Vec, @@ -85,7 +74,6 @@ pub(crate) struct ThreadState { pub(crate) cancel_tx: Option>, pub(crate) experimental_raw_events: bool, pub(crate) listener_generation: u64, - running_command_started_at_ms: HashMap, listener_command_tx: Option>, current_turn_history: ThreadHistoryBuilder, listener_thread: Option>, @@ -134,46 +122,12 @@ impl ThreadState { } pub(crate) fn active_turn_snapshot(&self) -> Option { - let mut turn = self.current_turn_history.active_turn_snapshot()?; - self.refresh_running_command_durations(std::slice::from_mut(&mut turn)); - Some(turn) - } - - pub(crate) fn refresh_running_command_durations(&self, turns: &mut [Turn]) { - for turn in turns { - for item in &mut turn.items { - let ThreadItem::CommandExecution { - id, - status: CommandExecutionStatus::InProgress, - duration_ms, - .. - } = item - else { - continue; - }; - let Some(started_at) = self.running_command_started_at_ms.get(id) else { - continue; - }; - *duration_ms = Some(elapsed_duration_ms(*started_at)); - } - } + self.current_turn_history.active_turn_snapshot() } pub(crate) fn track_current_turn_event(&mut self, event_turn_id: &str, event: &EventMsg) { - match event { - EventMsg::TurnStarted(payload) => { - self.turn_summary.started_at = payload.started_at; - } - EventMsg::ExecCommandBegin(payload) => { - self.running_command_started_at_ms.insert( - payload.call_id.clone(), - payload.started_at_ms.unwrap_or_else(now_unix_timestamp_ms), - ); - } - EventMsg::ExecCommandEnd(payload) => { - self.running_command_started_at_ms.remove(&payload.call_id); - } - _ => {} + if let EventMsg::TurnStarted(payload) = event { + self.turn_summary.started_at = payload.started_at; } self.current_turn_history.handle_event(event); if matches!(event, EventMsg::TurnAborted(_) | EventMsg::TurnComplete(_))