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 ef1b3f2b92..ecf189d7b7 100644 --- a/codex-rs/app-server-protocol/src/protocol/thread_history.rs +++ b/codex-rs/app-server-protocol/src/protocol/thread_history.rs @@ -311,11 +311,12 @@ impl ThreadHistoryBuilder { .map(|turn| turn.rollout_start_index) } - /// Shared reducer for persisted rollout replay and live current-turn tracking used by running - /// thread resume/rejoin. + /// Shared reducer for persisted rollout replay and in-memory current-turn + /// tracking used by running thread resume/rejoin. /// - /// In addition to persisted events, this must reduce typed item lifecycle and delta events that - /// reconstruct the latest in-progress turn while app-server remains connected to core. + /// This function should handle all EventMsg variants that can be persisted in a rollout file. + /// See `should_persist_event_msg` in `codex-rs/core/rollout/policy.rs`. + /// It also handles transient typed item lifecycle and delta events used by live resume. pub fn handle_event(&mut self, event: &EventMsg) { if self.consume_completed_item_legacy_event(event) { return; @@ -1388,6 +1389,7 @@ impl ThreadHistoryBuilder { } fn finish_current_turn(&mut self) { + self.pending_completed_item_legacy_events.clear(); if let Some(turn) = self.current_turn.take() { if turn.items.is_empty() && !turn.opened_explicitly && !turn.saw_compaction { return; @@ -1724,8 +1726,6 @@ mod tests { use codex_protocol::items::AgentMessageContent as CoreAgentMessageContent; use codex_protocol::items::AgentMessageItem as CoreAgentMessageItem; use codex_protocol::items::HookPromptFragment as CoreHookPromptFragment; - use codex_protocol::items::PlanItem as CorePlanItem; - use codex_protocol::items::ReasoningItem as CoreReasoningItem; use codex_protocol::items::SleepItem as CoreSleepItem; use codex_protocol::items::TurnItem as CoreTurnItem; use codex_protocol::items::UserMessageItem as CoreUserMessageItem; @@ -1743,20 +1743,13 @@ mod tests { use codex_protocol::protocol::CodexErrorInfo; use codex_protocol::protocol::CompactedItem; use codex_protocol::protocol::DynamicToolCallResponseEvent; - use codex_protocol::protocol::ExecCommandBeginEvent; use codex_protocol::protocol::ExecCommandEndEvent; - use codex_protocol::protocol::ExecCommandOutputDeltaEvent; use codex_protocol::protocol::ExecCommandSource; - use codex_protocol::protocol::ExecOutputStream; use codex_protocol::protocol::ItemStartedEvent; use codex_protocol::protocol::McpInvocation; use codex_protocol::protocol::McpToolCallEndEvent; use codex_protocol::protocol::PatchApplyBeginEvent; - use codex_protocol::protocol::PlanDeltaEvent; - use codex_protocol::protocol::ReasoningContentDeltaEvent; - use codex_protocol::protocol::ReasoningRawContentDeltaEvent; use codex_protocol::protocol::ThreadRolledBackEvent; - use codex_protocol::protocol::TokenCountEvent; use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::TurnAbortedEvent; use codex_protocol::protocol::TurnCompleteEvent; @@ -3402,174 +3395,6 @@ mod tests { }] ); } - - #[test] - fn active_turn_snapshot_accumulates_reasoning_deltas_by_index() { - let thread_id = ThreadId::new(); - let turn_id = "turn-1"; - let mut builder = ThreadHistoryBuilder::new(); - let events = [ - turn_started(turn_id), - EventMsg::ItemStarted(ItemStartedEvent { - thread_id, - turn_id: turn_id.to_string(), - item: CoreTurnItem::Reasoning(CoreReasoningItem { - id: "reason-1".to_string(), - summary_text: Vec::new(), - raw_content: Vec::new(), - }), - started_at_ms: 0, - }), - EventMsg::ReasoningContentDelta(ReasoningContentDeltaEvent { - thread_id: thread_id.to_string(), - turn_id: turn_id.to_string(), - item_id: "reason-1".to_string(), - delta: "think".to_string(), - summary_index: 0, - }), - EventMsg::ReasoningContentDelta(ReasoningContentDeltaEvent { - thread_id: thread_id.to_string(), - turn_id: turn_id.to_string(), - item_id: "reason-1".to_string(), - delta: " more".to_string(), - summary_index: 0, - }), - EventMsg::ReasoningContentDelta(ReasoningContentDeltaEvent { - thread_id: thread_id.to_string(), - turn_id: turn_id.to_string(), - item_id: "reason-1".to_string(), - delta: "next".to_string(), - summary_index: 1, - }), - EventMsg::ReasoningRawContentDelta(ReasoningRawContentDeltaEvent { - thread_id: thread_id.to_string(), - turn_id: turn_id.to_string(), - item_id: "reason-1".to_string(), - delta: "raw".to_string(), - content_index: 0, - }), - EventMsg::ReasoningRawContentDelta(ReasoningRawContentDeltaEvent { - thread_id: thread_id.to_string(), - turn_id: turn_id.to_string(), - item_id: "reason-1".to_string(), - delta: " detail".to_string(), - content_index: 0, - }), - EventMsg::ReasoningRawContentDelta(ReasoningRawContentDeltaEvent { - thread_id: thread_id.to_string(), - turn_id: turn_id.to_string(), - item_id: "reason-1".to_string(), - delta: "other".to_string(), - content_index: 1, - }), - ]; - for event in events { - builder.handle_event(&event); - } - - assert_eq!( - builder.active_turn_snapshot().expect("active turn").items, - vec![ThreadItem::Reasoning { - id: "reason-1".to_string(), - summary: vec!["think more".to_string(), "next".to_string()], - content: vec!["raw detail".to_string(), "other".to_string()], - }] - ); - } - - #[test] - fn active_turn_snapshot_accumulates_plan_deltas() { - let thread_id = ThreadId::new(); - let turn_id = "turn-1"; - let mut builder = ThreadHistoryBuilder::new(); - for event in [ - turn_started(turn_id), - EventMsg::ItemStarted(ItemStartedEvent { - thread_id, - turn_id: turn_id.to_string(), - item: CoreTurnItem::Plan(CorePlanItem { - id: "plan-1".to_string(), - text: String::new(), - }), - started_at_ms: 0, - }), - EventMsg::PlanDelta(PlanDeltaEvent { - thread_id: thread_id.to_string(), - turn_id: turn_id.to_string(), - item_id: "plan-1".to_string(), - delta: "first".to_string(), - }), - EventMsg::PlanDelta(PlanDeltaEvent { - thread_id: thread_id.to_string(), - turn_id: turn_id.to_string(), - item_id: "plan-1".to_string(), - delta: " step".to_string(), - }), - ] { - builder.handle_event(&event); - } - - assert_eq!( - builder.active_turn_snapshot().expect("active turn").items, - vec![ThreadItem::Plan { - id: "plan-1".to_string(), - text: "first step".to_string(), - }] - ); - } - - #[test] - fn active_turn_snapshot_accumulates_command_output_deltas() { - let turn_id = "turn-1"; - let mut builder = ThreadHistoryBuilder::new(); - for event in [ - turn_started(turn_id), - EventMsg::ExecCommandBegin(ExecCommandBeginEvent { - call_id: "command-1".to_string(), - process_id: Some("process-1".to_string()), - turn_id: turn_id.to_string(), - started_at_ms: 0, - command: vec!["echo".to_string(), "hello".to_string()], - cwd: test_path_buf("/tmp").abs().into(), - parsed_cmd: vec![ParsedCommand::Unknown { - cmd: "echo hello".to_string(), - }], - source: ExecCommandSource::Agent, - interaction_input: None, - }), - EventMsg::ExecCommandOutputDelta(ExecCommandOutputDeltaEvent { - call_id: "command-1".to_string(), - stream: ExecOutputStream::Stdout, - chunk: b"hel".to_vec(), - }), - EventMsg::ExecCommandOutputDelta(ExecCommandOutputDeltaEvent { - call_id: "command-1".to_string(), - stream: ExecOutputStream::Stdout, - chunk: b"lo".to_vec(), - }), - ] { - builder.handle_event(&event); - } - - assert_eq!( - builder.active_turn_snapshot().expect("active turn").items, - vec![ThreadItem::CommandExecution { - id: "command-1".to_string(), - command: "echo hello".to_string(), - cwd: test_path_buf("/tmp").abs().into(), - process_id: Some("process-1".to_string()), - source: CommandExecutionSource::Agent, - status: CommandExecutionStatus::InProgress, - command_actions: vec![CommandAction::Unknown { - command: "echo hello".to_string(), - }], - aggregated_output: Some("hello".to_string()), - exit_code: None, - duration_ms: None, - }] - ); - } - #[test] fn completed_agent_message_replaces_streamed_item_without_legacy_duplicate() { let thread_id = ThreadId::new(); @@ -3596,22 +3421,12 @@ mod tests { }), started_at_ms: 0, }), - EventMsg::AgentMessageContentDelta(AgentMessageContentDeltaEvent { - thread_id: thread_id.to_string(), - turn_id: turn_id.to_string(), - item_id: "msg-1".to_string(), - delta: "hello".to_string(), - }), EventMsg::ItemCompleted(ItemCompletedEvent { thread_id, turn_id: turn_id.to_string(), item: completed_item.clone(), completed_at_ms: 0, }), - EventMsg::TokenCount(TokenCountEvent { - info: None, - rate_limits: None, - }), completed_item .as_legacy_events(/*show_raw_agent_reasoning*/ true) .into_iter() @@ -3631,69 +3446,6 @@ mod tests { }] ); } - - #[test] - fn completed_reasoning_replaces_streamed_item_without_legacy_duplicate() { - let thread_id = ThreadId::new(); - let turn_id = "turn-1"; - let completed_item = CoreTurnItem::Reasoning(CoreReasoningItem { - id: "reason-1".to_string(), - summary_text: vec!["thinking".to_string()], - raw_content: vec!["details".to_string()], - }); - let mut builder = ThreadHistoryBuilder::new(); - let mut events = vec![ - turn_started(turn_id), - EventMsg::ItemStarted(ItemStartedEvent { - thread_id, - turn_id: turn_id.to_string(), - item: CoreTurnItem::Reasoning(CoreReasoningItem { - id: "reason-1".to_string(), - summary_text: Vec::new(), - raw_content: Vec::new(), - }), - started_at_ms: 0, - }), - EventMsg::ReasoningContentDelta(ReasoningContentDeltaEvent { - thread_id: thread_id.to_string(), - turn_id: turn_id.to_string(), - item_id: "reason-1".to_string(), - delta: "thinking".to_string(), - summary_index: 0, - }), - EventMsg::ReasoningRawContentDelta(ReasoningRawContentDeltaEvent { - thread_id: thread_id.to_string(), - turn_id: turn_id.to_string(), - item_id: "reason-1".to_string(), - delta: "details".to_string(), - content_index: 0, - }), - EventMsg::ItemCompleted(ItemCompletedEvent { - thread_id, - turn_id: turn_id.to_string(), - item: completed_item.clone(), - completed_at_ms: 0, - }), - EventMsg::TokenCount(TokenCountEvent { - info: None, - rate_limits: None, - }), - ]; - events.extend(completed_item.as_legacy_events(/*show_raw_agent_reasoning*/ true)); - for event in events { - builder.handle_event(&event); - } - - assert_eq!( - builder.active_turn_snapshot().expect("active turn").items, - vec![ThreadItem::Reasoning { - id: "reason-1".to_string(), - summary: vec!["thinking".to_string()], - content: vec!["details".to_string()], - }] - ); - } - fn turn_started(turn_id: &str) -> EventMsg { EventMsg::TurnStarted(TurnStartedEvent { turn_id: turn_id.to_string(), diff --git a/codex-rs/app-server/src/extensions.rs b/codex-rs/app-server/src/extensions.rs index bfd513d196..ed040aa8d9 100644 --- a/codex-rs/app-server/src/extensions.rs +++ b/codex-rs/app-server/src/extensions.rs @@ -21,12 +21,8 @@ use codex_protocol::ThreadId; use codex_protocol::error::CodexErr; use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; -use codex_protocol::protocol::RolloutItem; -use codex_protocol::protocol::ThreadGoalUpdatedEvent; use codex_rollout::state_db::StateDbHandle; -use codex_thread_store::AppendThreadItemsParams; use codex_thread_store::ThreadStore; -use tokio::sync::mpsc; use crate::outgoing_message::OutgoingMessageSender; use crate::thread_state::ThreadListenerCommand; @@ -103,23 +99,16 @@ where pub(crate) fn app_server_extension_event_sink( outgoing: Arc, thread_state_manager: ThreadStateManager, - thread_store: Arc, ) -> Arc { - let (fallback_goal_update_tx, fallback_goal_update_rx) = mpsc::unbounded_channel(); - tokio::spawn(run_fallback_goal_update_worker( - Arc::downgrade(&outgoing), - thread_store, - fallback_goal_update_rx, - )); Arc::new(AppServerExtensionEventSink { + outgoing, thread_state_manager, - fallback_goal_update_tx, }) } struct AppServerExtensionEventSink { + outgoing: Arc, thread_state_manager: ThreadStateManager, - fallback_goal_update_tx: mpsc::UnboundedSender, } impl ExtensionEventSink for AppServerExtensionEventSink { @@ -127,13 +116,16 @@ impl ExtensionEventSink for AppServerExtensionEventSink { match event.msg { EventMsg::ThreadGoalUpdated(thread_goal_event) => { let thread_id = thread_goal_event.thread_id; + let turn_id = thread_goal_event.turn_id; + let goal: ThreadGoal = thread_goal_event.goal.into(); if let Some(listener_command_tx) = self .thread_state_manager .current_listener_command_tx(thread_id) { - let command = ThreadListenerCommand::PersistAndEmitThreadGoalUpdated( - thread_goal_event.clone(), - ); + let command = ThreadListenerCommand::EmitThreadGoalUpdated { + turn_id: turn_id.clone(), + goal: goal.clone(), + }; if listener_command_tx.send(command).is_ok() { return; } @@ -141,15 +133,18 @@ impl ExtensionEventSink for AppServerExtensionEventSink { "failed to enqueue extension goal update for {thread_id}: listener command channel is closed" ); } - if self - .fallback_goal_update_tx - .send(thread_goal_event) - .is_err() - { - tracing::error!( - "failed to enqueue fallback extension goal update for {thread_id}: fallback worker is closed" - ); - } + let outgoing = Arc::clone(&self.outgoing); + tokio::spawn(async move { + outgoing + .send_server_notification(ServerNotification::ThreadGoalUpdated( + ThreadGoalUpdatedNotification { + thread_id: thread_id.to_string(), + turn_id, + goal, + }, + )) + .await; + }); } msg => { tracing::debug!(event_id = %event.id, ?msg, "dropping unsupported extension event"); @@ -158,47 +153,6 @@ impl ExtensionEventSink for AppServerExtensionEventSink { } } -async fn run_fallback_goal_update_worker( - outgoing: Weak, - thread_store: Arc, - mut goal_updates: mpsc::UnboundedReceiver, -) { - while let Some(event) = goal_updates.recv().await { - let thread_id = event.thread_id; - let notification = thread_goal_updated_notification(&event); - let item = RolloutItem::EventMsg(EventMsg::ThreadGoalUpdated(event)); - if let Err(err) = thread_store - .append_items(AppendThreadItemsParams { - thread_id, - items: vec![item], - }) - .await - { - tracing::error!( - "failed to persist fallback extension goal update for {thread_id}: {err}" - ); - continue; - } - let Some(outgoing) = outgoing.upgrade() else { - return; - }; - outgoing - .send_server_notification(ServerNotification::ThreadGoalUpdated(notification)) - .await; - } -} - -/// Converts a core goal event without losing its persisted turn attribution. -pub(crate) fn thread_goal_updated_notification( - event: &ThreadGoalUpdatedEvent, -) -> ThreadGoalUpdatedNotification { - ThreadGoalUpdatedNotification { - thread_id: event.thread_id.to_string(), - turn_id: event.turn_id.clone(), - goal: ThreadGoal::from(event.goal.clone()), - } -} - pub(crate) fn guardian_agent_spawner( thread_manager: Weak, ) -> impl AgentSpawner { @@ -224,15 +178,11 @@ mod tests { use codex_protocol::protocol::ThreadGoal as CoreThreadGoal; use codex_protocol::protocol::ThreadGoalStatus; use codex_protocol::protocol::ThreadGoalUpdatedEvent; - use codex_thread_store::InMemoryThreadStore; - use codex_thread_store::LoadThreadHistoryParams; use pretty_assertions::assert_eq; use tokio::sync::mpsc; use tokio::time::timeout; use super::*; - use crate::outgoing_message::OutgoingEnvelope; - use crate::outgoing_message::OutgoingMessage; #[tokio::test] async fn app_server_event_sink_uses_listener_fifo_for_goal_updates_and_clears() { @@ -245,11 +195,7 @@ mod tests { let thread_id = ThreadId::default(); let (listener_command_tx, mut listener_command_rx) = mpsc::unbounded_channel(); thread_state_manager.register_listener_command_tx(thread_id, listener_command_tx.clone()); - let sink = app_server_extension_event_sink( - outgoing, - thread_state_manager, - Arc::new(InMemoryThreadStore::default()), - ); + let sink = app_server_extension_event_sink(outgoing, thread_state_manager); for turn_id in ["turn-1", "turn-2"] { sink.emit(thread_goal_updated_event(thread_id, turn_id)); @@ -265,12 +211,8 @@ mod tests { .expect("timed out waiting for listener command") .expect("listener command channel closed unexpectedly"); match command { - ThreadListenerCommand::PersistAndEmitThreadGoalUpdated(event) => { - observed.push( - event - .turn_id - .expect("extension goal updates should include turn ids"), - ); + ThreadListenerCommand::EmitThreadGoalUpdated { turn_id, .. } => { + observed.push(turn_id.expect("extension goal updates should include turn ids")); } ThreadListenerCommand::EmitThreadGoalCleared => { observed.push("cleared".to_string()) @@ -289,86 +231,6 @@ mod tests { ); } - #[tokio::test] - async fn app_server_event_sink_fallback_persists_goal_updates_before_broadcast_in_order() { - let (outgoing_tx, mut outgoing_rx) = mpsc::channel(4); - let outgoing = Arc::new(OutgoingMessageSender::new( - outgoing_tx, - AnalyticsEventsClient::disabled(), - )); - let thread_store = Arc::new(InMemoryThreadStore::default()); - let thread_id = ThreadId::default(); - let sink = app_server_extension_event_sink( - outgoing.clone(), - ThreadStateManager::new(), - thread_store.clone(), - ); - - for turn_id in ["turn-1", "turn-2"] { - sink.emit(thread_goal_updated_event(thread_id, turn_id)); - } - - let mut notified_turn_ids = Vec::new(); - for _ in 0..2 { - let envelope = timeout(Duration::from_secs(1), outgoing_rx.recv()) - .await - .expect("timed out waiting for fallback goal update") - .expect("outgoing channel closed unexpectedly"); - let OutgoingEnvelope::Broadcast { message } = envelope else { - panic!("expected broadcast goal update"); - }; - let OutgoingMessage::AppServerNotification(ServerNotification::ThreadGoalUpdated( - notification, - )) = message - else { - panic!("expected thread goal update notification"); - }; - notified_turn_ids.push( - notification - .turn_id - .expect("extension goal updates should include turn ids"), - ); - } - - let history = thread_store - .load_history(LoadThreadHistoryParams { - thread_id, - include_archived: true, - }) - .await - .expect("fallback goal updates should be persisted"); - let persisted_turn_ids = history - .items - .into_iter() - .filter_map(|item| match item { - RolloutItem::EventMsg(EventMsg::ThreadGoalUpdated(event)) => event.turn_id, - _ => None, - }) - .collect::>(); - - assert_eq!(notified_turn_ids, vec!["turn-1", "turn-2"]); - assert_eq!(persisted_turn_ids, notified_turn_ids); - assert_eq!(thread_store.calls().await.append_items, 2); - } - - #[tokio::test] - async fn app_server_event_sink_fallback_worker_does_not_own_outgoing_lifetime() { - let (outgoing_tx, _outgoing_rx) = mpsc::channel(1); - let outgoing = Arc::new(OutgoingMessageSender::new( - outgoing_tx, - AnalyticsEventsClient::disabled(), - )); - let weak_outgoing = Arc::downgrade(&outgoing); - - let _sink = app_server_extension_event_sink( - outgoing, - ThreadStateManager::new(), - Arc::new(InMemoryThreadStore::default()), - ); - - assert!(weak_outgoing.upgrade().is_none()); - } - fn thread_goal_updated_event(thread_id: ThreadId, turn_id: &str) -> Event { Event { id: turn_id.to_string(), diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 48a4e6241b..41ecbfe561 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -357,7 +357,6 @@ impl MessageProcessor { event_sink: app_server_extension_event_sink( outgoing.clone(), thread_state_manager.clone(), - Arc::clone(&thread_store), ), auth_manager: auth_manager.clone(), state_db: state_db.clone(), diff --git a/codex-rs/app-server/src/outgoing_message.rs b/codex-rs/app-server/src/outgoing_message.rs index b1585a92f7..9ce815bafe 100644 --- a/codex-rs/app-server/src/outgoing_message.rs +++ b/codex-rs/app-server/src/outgoing_message.rs @@ -7,7 +7,6 @@ use std::time::UNIX_EPOCH; use codex_analytics::AnalyticsEventsClient; use codex_app_server_protocol::ClientResponsePayload; -use codex_app_server_protocol::ItemStartedNotification; use codex_app_server_protocol::JSONRPCErrorError; use codex_app_server_protocol::RequestId; use codex_app_server_protocol::Result; @@ -15,7 +14,6 @@ use codex_app_server_protocol::ServerNotification; use codex_app_server_protocol::ServerRequest; use codex_app_server_protocol::ServerRequestPayload; use codex_app_server_protocol::ServerResponse; -use codex_app_server_protocol::Turn; use codex_otel::span_w3c_trace_context; use codex_protocol::ThreadId; use codex_protocol::protocol::W3cTraceContext; @@ -27,7 +25,6 @@ use tracing::Instrument; use tracing::Span; use tracing::warn; -use self::active_item_lifecycle::ActiveItemLifecycleTracker; use crate::error_code::internal_error; use crate::server_request_error::TURN_TRANSITION_PENDING_REQUEST_ERROR_REASON; pub(crate) use codex_app_server_transport::ConnectionId; @@ -36,8 +33,6 @@ pub(crate) use codex_app_server_transport::OutgoingMessage; pub(crate) use codex_app_server_transport::OutgoingResponse; pub(crate) use codex_app_server_transport::QueuedOutgoingMessage; -mod active_item_lifecycle; - #[cfg(test)] use codex_protocol::account::PlanType; @@ -106,8 +101,6 @@ pub(crate) struct OutgoingMessageSender { /// We keep them here because this is where responses, errors, and /// disconnect cleanup all get handled. request_contexts: Mutex>, - /// Live typed item lifecycles used to restore client-only in-progress identity on reconnect. - active_item_lifecycles: Mutex, analytics_events_client: AnalyticsEventsClient, } @@ -165,9 +158,6 @@ impl ThreadScopedOutgoingMessageSender { } pub(crate) async fn send_server_notification(&self, notification: ServerNotification) { - self.outgoing - .note_thread_item_lifecycle(self.thread_id, ¬ification) - .await; self.outgoing .analytics_events_client .track_notification(notification.clone()); @@ -226,34 +216,10 @@ impl OutgoingMessageSender { sender, request_id_to_callback: Mutex::new(HashMap::new()), request_contexts: Mutex::new(HashMap::new()), - active_item_lifecycles: Mutex::new(ActiveItemLifecycleTracker::default()), analytics_events_client, } } - async fn note_thread_item_lifecycle( - &self, - thread_id: ThreadId, - notification: &ServerNotification, - ) { - self.active_item_lifecycles - .lock() - .await - .note_notification(thread_id, notification); - } - - /// Builds the active item starts that must follow a running-thread resume response. - pub(crate) async fn active_item_starts_for_turn( - &self, - thread_id: ThreadId, - active_turn: &Turn, - ) -> Vec { - self.active_item_lifecycles - .lock() - .await - .active_starts_for_turn(thread_id, active_turn) - } - pub(crate) async fn register_request_context(&self, request_context: RequestContext) { let mut request_contexts = self.request_contexts.lock().await; if request_contexts diff --git a/codex-rs/app-server/src/outgoing_message/active_item_lifecycle.rs b/codex-rs/app-server/src/outgoing_message/active_item_lifecycle.rs deleted file mode 100644 index 99e1d6903e..0000000000 --- a/codex-rs/app-server/src/outgoing_message/active_item_lifecycle.rs +++ /dev/null @@ -1,94 +0,0 @@ -use std::collections::HashMap; - -use codex_app_server_protocol::ItemStartedNotification; -use codex_app_server_protocol::ServerNotification; -use codex_app_server_protocol::Turn; -use codex_protocol::ThreadId; - -/// Tracks the typed item lifecycles emitted by app-server so a reconnecting client can rebuild -/// which materialized items are still active. -#[derive(Default)] -pub(super) struct ActiveItemLifecycleTracker { - started_by_thread: HashMap>, -} - -impl ActiveItemLifecycleTracker { - pub(super) fn note_notification( - &mut self, - thread_id: ThreadId, - notification: &ServerNotification, - ) { - match notification { - ServerNotification::ItemStarted(started) => { - let active_items = self.started_by_thread.entry(thread_id).or_default(); - if let Some(existing) = active_items.iter_mut().find(|existing| { - existing.turn_id == started.turn_id && existing.item.id() == started.item.id() - }) { - *existing = started.clone(); - } else { - active_items.push(started.clone()); - } - } - ServerNotification::ItemCompleted(completed) => { - if let Some(active_items) = self.started_by_thread.get_mut(&thread_id) { - active_items.retain(|started| { - started.turn_id != completed.turn_id - || started.item.id() != completed.item.id() - }); - } - } - ServerNotification::TurnCompleted(completed) => { - self.remove_turn(thread_id, &completed.turn.id); - } - ServerNotification::Error(error) if !error.will_retry => { - self.remove_turn(thread_id, &error.turn_id); - } - _ => return, - } - - if self - .started_by_thread - .get(&thread_id) - .is_some_and(Vec::is_empty) - { - self.started_by_thread.remove(&thread_id); - } - } - - /// Returns active start notifications with their item payload replaced by the latest - /// materialized content from the listener's turn snapshot. - pub(super) fn active_starts_for_turn( - &self, - thread_id: ThreadId, - active_turn: &Turn, - ) -> Vec { - let Some(active_items) = self.started_by_thread.get(&thread_id) else { - return Vec::new(); - }; - active_items - .iter() - .filter(|started| started.turn_id == active_turn.id) - .cloned() - .map(|mut started| { - if let Some(latest_item) = active_turn - .items - .iter() - .find(|item| item.id() == started.item.id()) - { - started.item = latest_item.clone(); - } - started - }) - .collect() - } - - fn remove_turn(&mut self, thread_id: ThreadId, turn_id: &str) { - if let Some(active_items) = self.started_by_thread.get_mut(&thread_id) { - active_items.retain(|started| started.turn_id != turn_id); - } - } -} - -#[cfg(test)] -#[path = "active_item_lifecycle_tests.rs"] -mod tests; diff --git a/codex-rs/app-server/src/outgoing_message/active_item_lifecycle_tests.rs b/codex-rs/app-server/src/outgoing_message/active_item_lifecycle_tests.rs deleted file mode 100644 index 1adf2d2e6e..0000000000 --- a/codex-rs/app-server/src/outgoing_message/active_item_lifecycle_tests.rs +++ /dev/null @@ -1,87 +0,0 @@ -use codex_app_server_protocol::ErrorNotification; -use codex_app_server_protocol::ItemCompletedNotification; -use codex_app_server_protocol::ItemStartedNotification; -use codex_app_server_protocol::ServerNotification; -use codex_app_server_protocol::ThreadItem; -use codex_app_server_protocol::Turn; -use codex_app_server_protocol::TurnCompletedNotification; -use codex_app_server_protocol::TurnError; -use codex_app_server_protocol::TurnItemsView; -use codex_app_server_protocol::TurnStatus; -use codex_protocol::ThreadId; -use pretty_assertions::assert_eq; - -use super::ActiveItemLifecycleTracker; - -#[test] -fn replays_latest_content_until_terminal_notification() { - let thread_id = ThreadId::new(); - let turn_id = "turn-1".to_string(); - let started = ItemStartedNotification { - thread_id: thread_id.to_string(), - turn_id: turn_id.clone(), - started_at_ms: 1, - item: ThreadItem::AgentMessage { - id: "message-1".to_string(), - text: String::new(), - phase: None, - memory_citation: None, - }, - }; - let latest_item = ThreadItem::AgentMessage { - id: "message-1".to_string(), - text: "latest".to_string(), - phase: None, - memory_citation: None, - }; - let active_turn = Turn { - id: turn_id.clone(), - items: vec![latest_item.clone()], - items_view: TurnItemsView::Full, - status: TurnStatus::InProgress, - error: None, - started_at: None, - completed_at: None, - duration_ms: None, - }; - let terminal_notifications = [ - ServerNotification::ItemCompleted(ItemCompletedNotification { - item: latest_item.clone(), - thread_id: thread_id.to_string(), - turn_id: turn_id.clone(), - completed_at_ms: 2, - }), - ServerNotification::TurnCompleted(TurnCompletedNotification { - thread_id: thread_id.to_string(), - turn: active_turn.clone(), - }), - ServerNotification::Error(ErrorNotification { - error: TurnError { - message: "terminal failure".to_string(), - codex_error_info: None, - additional_details: None, - }, - will_retry: false, - thread_id: thread_id.to_string(), - turn_id, - }), - ]; - - for terminal_notification in terminal_notifications { - let mut tracker = ActiveItemLifecycleTracker::default(); - tracker.note_notification(thread_id, &ServerNotification::ItemStarted(started.clone())); - assert_eq!( - tracker.active_starts_for_turn(thread_id, &active_turn), - vec![ItemStartedNotification { - item: latest_item.clone(), - ..started.clone() - }] - ); - - tracker.note_notification(thread_id, &terminal_notification); - assert_eq!( - tracker.active_starts_for_turn(thread_id, &active_turn), - Vec::new() - ); - } -} diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index fb07a09c67..d73a41706c 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -496,7 +496,6 @@ mod git_processor; mod initialize_processor; mod marketplace_processor; mod mcp_processor; -mod persisted_side_state_replay; mod plugins; mod process_exec_processor; mod remote_control_processor; @@ -537,8 +536,6 @@ use crate::thread_state::ConnectionCapabilities; use crate::thread_state::ThreadListenerCommand; use crate::thread_state::ThreadState; use crate::thread_state::ThreadStateManager; -use persisted_side_state_replay::send_persisted_side_state_to_connection; -use token_usage_replay::TokenUsageTurnOwner; use token_usage_replay::latest_token_usage_turn_id_from_rollout_items; use token_usage_replay::send_thread_token_usage_update_to_connection; diff --git a/codex-rs/app-server/src/request_processors/persisted_side_state_replay.rs b/codex-rs/app-server/src/request_processors/persisted_side_state_replay.rs deleted file mode 100644 index 6c5e095a41..0000000000 --- a/codex-rs/app-server/src/request_processors/persisted_side_state_replay.rs +++ /dev/null @@ -1,92 +0,0 @@ -//! Reconstructs typed lifecycle notifications that are durable outside turn items. -//! -//! Resume callers own response ordering. This module keeps rollout extraction pure and sends the -//! reconstructed notifications only to the connection whose state is being restored. - -use std::sync::Arc; - -use codex_app_server_protocol::HookCompletedNotification; -use codex_app_server_protocol::ServerNotification; -use codex_app_server_protocol::ThreadGoalUpdatedNotification; -use codex_app_server_protocol::guardian_auto_approval_review_notification; -use codex_protocol::ThreadId; -use codex_protocol::protocol::EventMsg; -use codex_protocol::protocol::RolloutItem; -use codex_protocol::protocol::ThreadGoalStatus; -use codex_rollout::is_persisted_rollout_item; - -use crate::outgoing_message::ConnectionId; -use crate::outgoing_message::OutgoingMessageSender; - -/// Extracts ordered typed notifications for persisted state that is not represented by turn items. -pub(super) fn persisted_side_state_notifications( - thread_id: ThreadId, - rollout_items: &[RolloutItem], - include_goal_updates: bool, -) -> Vec { - let mut latest_turn_id = ""; - let mut notifications = Vec::new(); - - for rollout_item in rollout_items { - if !is_persisted_rollout_item(rollout_item) { - continue; - } - let RolloutItem::EventMsg(event) = rollout_item else { - continue; - }; - if let EventMsg::TurnStarted(started) = event { - latest_turn_id = started.turn_id.as_str(); - continue; - } - if let EventMsg::GuardianAssessment(assessment) = event { - notifications.push(guardian_auto_approval_review_notification( - &thread_id, - latest_turn_id, - assessment, - )); - continue; - } - if let EventMsg::HookCompleted(completed) = event { - notifications.push(ServerNotification::HookCompleted( - HookCompletedNotification { - thread_id: thread_id.to_string(), - turn_id: completed.turn_id.clone(), - run: completed.run.clone().into(), - }, - )); - continue; - } - if include_goal_updates - && let EventMsg::ThreadGoalUpdated(updated) = event - && updated.goal.status == ThreadGoalStatus::Complete - && let Some(turn_id) = updated.turn_id.as_ref() - { - notifications.push(ServerNotification::ThreadGoalUpdated( - ThreadGoalUpdatedNotification { - thread_id: thread_id.to_string(), - turn_id: Some(turn_id.clone()), - goal: updated.goal.clone().into(), - }, - )); - } - } - - notifications -} - -/// Sends reconstructed persisted state only to the connection completing resume. -pub(super) async fn send_persisted_side_state_to_connection( - outgoing: &Arc, - connection_id: ConnectionId, - thread_id: ThreadId, - rollout_items: &[RolloutItem], - include_goal_updates: bool, -) { - for notification in - persisted_side_state_notifications(thread_id, rollout_items, include_goal_updates) - { - outgoing - .send_server_notification_to_connections(&[connection_id], notification) - .await; - } -} diff --git a/codex-rs/app-server/src/request_processors/thread_lifecycle.rs b/codex-rs/app-server/src/request_processors/thread_lifecycle.rs index 3e865d3d5c..4e288a7fb0 100644 --- a/codex-rs/app-server/src/request_processors/thread_lifecycle.rs +++ b/codex-rs/app-server/src/request_processors/thread_lifecycle.rs @@ -493,20 +493,6 @@ pub(super) async fn handle_thread_listener_command( )) .await; } - ThreadListenerCommand::PersistAndEmitThreadGoalUpdated(event) => { - let thread_id = event.thread_id; - let notification = crate::extensions::thread_goal_updated_notification(&event); - if let Err(err) = conversation - .append_rollout_items(&[RolloutItem::EventMsg(EventMsg::ThreadGoalUpdated(event))]) - .await - { - tracing::error!("failed to persist extension goal update for {thread_id}: {err}"); - return; - } - outgoing - .send_server_notification(ServerNotification::ThreadGoalUpdated(notification)) - .await; - } ThreadListenerCommand::EmitThreadGoalCleared => { outgoing .send_server_notification(ServerNotification::ThreadGoalCleared( @@ -552,23 +538,32 @@ pub(super) async fn handle_pending_thread_resume_request( pending: crate::thread_state::PendingThreadResumeRequest, ) { // Read persisted history in listener order so the resume snapshot cannot be older than a - // terminal event that already cleared the listener's active-turn projection. - let history_items = match conversation.load_history(/*include_archived*/ true).await { - Ok(history) => history.items, - Err(err) => { - outgoing - .send_error( - pending.request_id, - super::thread_processor::thread_store_resume_read_error(err), - ) - .await; - return; + // terminal event that already reached the rollout writer. + let history_items = if pending.include_turns || pending.initial_turns_page.is_some() { + match conversation.load_history(/*include_archived*/ true).await { + Ok(history) => history.items, + Err(err) => { + outgoing + .send_error( + pending.request_id, + super::thread_processor::thread_store_resume_read_error(err), + ) + .await; + return; + } } + } else { + Vec::new() }; let active_turn = { let state = thread_state.lock().await; state.active_turn_snapshot() }; + let mut reconstructed_turns = build_api_turns_from_rollout_items(&history_items); + let active_turn_is_current = active_turn.as_ref().is_none_or(|active_turn| { + merge_turn_history_with_active_turn(&mut reconstructed_turns, active_turn.clone()) + }); + let active_turn = active_turn.filter(|_| active_turn_is_current); tracing::debug!( thread_id = %conversation_id, request_id = ?pending.request_id, @@ -577,17 +572,17 @@ pub(super) async fn handle_pending_thread_resume_request( active_turn_status = ?active_turn.as_ref().map(|turn| &turn.status), "composing running thread resume response" ); - let has_live_in_progress_turn = - matches!(conversation.agent_status().await, AgentStatus::Running) + let has_live_in_progress_turn = active_turn_is_current + && (matches!(conversation.agent_status().await, AgentStatus::Running) || active_turn .as_ref() - .is_some_and(|turn| matches!(turn.status, TurnStatus::InProgress)); + .is_some_and(|turn| matches!(turn.status, TurnStatus::InProgress))); let request_id = pending.request_id; let connection_id = request_id.connection_id; let mut thread = pending.thread_summary; if pending.include_turns { - populate_thread_turns_from_history(&mut thread, &history_items, active_turn.as_ref()); + thread.turns = reconstructed_turns; } let thread_status = thread_watch_manager @@ -599,24 +594,16 @@ pub(super) async fn handle_pending_thread_resume_request( thread_status, has_live_in_progress_turn, ); - let initial_turns_sort_direction = pending - .initial_turns_page - .as_ref() - .and_then(|params| params.sort_direction) - .unwrap_or(SortDirection::Desc); - let mut initial_turns_page_token_usage_turn_id = None; + let token_usage_thread = pending.include_turns.then(|| thread.clone()); let mut initial_turns_page = if let Some(params) = pending.initial_turns_page.as_ref() { match super::thread_processor::build_thread_resume_initial_turns_page( &history_items, thread.status.clone(), has_live_in_progress_turn, - active_turn.clone(), + active_turn, params, ) { - Ok((page, token_usage_turn_id)) => { - initial_turns_page_token_usage_turn_id = token_usage_turn_id; - Some(page) - } + Ok(page) => Some(page), Err(error) => { outgoing.send_error(request_id, error).await; return; @@ -625,26 +612,6 @@ pub(super) async fn handle_pending_thread_resume_request( } else { None }; - let token_usage_thread = pending.include_turns.then(|| thread.clone()).or_else(|| { - initial_turns_page.as_ref().map(|page| { - let mut token_usage_thread = thread.clone(); - token_usage_thread.turns = page.data.clone(); - if initial_turns_sort_direction == SortDirection::Desc { - token_usage_thread.turns.reverse(); - } - token_usage_thread - }) - }); - let token_usage_turn_id = if initial_turns_page.is_some() { - initial_turns_page_token_usage_turn_id - } else { - token_usage_thread.as_ref().and_then(|token_usage_thread| { - latest_token_usage_turn_id_from_rollout_items( - &history_items, - token_usage_thread.turns.as_slice(), - ) - }) - }; if pending.redact_resume_payloads { redact_thread_resume_payloads(&mut thread.turns); if let Some(initial_turns_page) = initial_turns_page.as_mut() { @@ -718,24 +685,13 @@ pub(super) async fn handle_pending_thread_resume_request( initial_turns_page, }; outgoing.send_response(request_id, response).await; - if let Some(active_turn) = active_turn.as_ref() - && matches!(active_turn.status, TurnStatus::InProgress) - { - for started in outgoing - .active_item_starts_for_turn(conversation_id, active_turn) - .await - { - outgoing - .send_server_notification_to_connection_and_wait( - connection_id, - ServerNotification::ItemStarted(started), - ) - .await; - } - } - // An initial turns page already paid the reconstruction cost and provides the bounded turn - // identities needed for usage attribution. Metadata-only resume remains cheap. + // Match cold resume: metadata-only resume should attach the listener without + // paying the cost of turn reconstruction for historical usage replay. if let Some(token_usage_thread) = token_usage_thread { + let token_usage_turn_id = latest_token_usage_turn_id_from_rollout_items( + &history_items, + token_usage_thread.turns.as_slice(), + ); // Rejoining a loaded thread has the same UI contract as a cold resume, but // uses the live conversation state instead of reconstructing a new session. send_thread_token_usage_update_to_connection( @@ -748,14 +704,6 @@ pub(super) async fn handle_pending_thread_resume_request( ) .await; } - send_persisted_side_state_to_connection( - outgoing, - connection_id, - conversation_id, - &history_items, - pending.emit_thread_goal_update, - ) - .await; if pending.emit_thread_goal_update { if let Some(state_db) = pending.thread_goal_state_db { send_thread_goal_snapshot_notification(outgoing, conversation_id, &state_db).await; @@ -848,9 +796,22 @@ pub(super) async fn resolve_pending_server_request( .await; } -pub(super) fn merge_turn_history_with_active_turn(turns: &mut Vec, active_turn: Turn) { +/// Merges a current listener snapshot unless persisted terminal state proves it is stale. +/// Returns whether the active snapshot was accepted. +pub(super) fn merge_turn_history_with_active_turn( + turns: &mut Vec, + active_turn: Turn, +) -> bool { + if matches!(active_turn.status, TurnStatus::InProgress) + && turns + .iter() + .any(|turn| turn.id == active_turn.id && !matches!(turn.status, TurnStatus::InProgress)) + { + return false; + } turns.retain(|turn| turn.id != active_turn.id); turns.push(active_turn); + true } pub(super) fn set_thread_status_and_interrupt_stale_turns( diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index 43834fd9f5..8d5950db07 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -2387,7 +2387,6 @@ impl ThreadRequestProcessor { items_view: items_view.unwrap_or(TurnItemsView::Summary), }, ) - .map(|(response, _)| response) } async fn thread_items_list_response_inner( @@ -2819,11 +2818,7 @@ impl ThreadRequestProcessor { let active_permission_profile = thread_response_active_permission_profile( config_snapshot.active_permission_profile, ); - let initial_turns_sort_direction = initial_turns_page - .as_ref() - .and_then(|params| params.sort_direction) - .unwrap_or(SortDirection::Desc); - let mut initial_turns_page_token_usage_turn_id = None; + let token_usage_thread = include_turns.then(|| thread.clone()); let mut initial_turns_page = if let Some(params) = initial_turns_page.as_ref() { match build_thread_resume_initial_turns_page( response_history.get_rollout_items(), @@ -2832,10 +2827,7 @@ impl ThreadRequestProcessor { /*active_turn*/ None, params, ) { - Ok((page, token_usage_turn_id)) => { - initial_turns_page_token_usage_turn_id = token_usage_turn_id; - Some(page) - } + Ok(page) => Some(page), Err(error) => { self.outgoing.send_error(request_id, error).await; return Ok(()); @@ -2844,26 +2836,6 @@ impl ThreadRequestProcessor { } else { None }; - let token_usage_thread = include_turns.then(|| thread.clone()).or_else(|| { - initial_turns_page.as_ref().map(|page| { - let mut token_usage_thread = thread.clone(); - token_usage_thread.turns = page.data.clone(); - if initial_turns_sort_direction == SortDirection::Desc { - token_usage_thread.turns.reverse(); - } - token_usage_thread - }) - }); - let token_usage_turn_id = if initial_turns_page.is_some() { - initial_turns_page_token_usage_turn_id - } else { - token_usage_thread.as_ref().and_then(|token_usage_thread| { - latest_token_usage_turn_id_from_rollout_items( - response_history.get_rollout_items(), - token_usage_thread.turns.as_slice(), - ) - }) - }; if redact_resume_payloads { redact_thread_resume_payloads(&mut thread.turns); if let Some(initial_turns_page) = initial_turns_page.as_mut() { @@ -2890,10 +2862,13 @@ impl ThreadRequestProcessor { let connection_id = request_id.connection_id; self.outgoing.send_response(request_id, response).await; - // An initial turns page already paid the reconstruction cost and provides the - // bounded turn identities needed for usage attribution. Metadata-only resume - // remains cheap. + // `excludeTurns` is explicitly the cheap resume path, so avoid + // rebuilding history only to attribute a replayed usage update. if let Some(token_usage_thread) = token_usage_thread { + let token_usage_turn_id = latest_token_usage_turn_id_from_rollout_items( + response_history.get_rollout_items(), + token_usage_thread.turns.as_slice(), + ); // The client needs restored usage before it starts another turn. // Sending after the response preserves JSON-RPC request ordering while // still filling the status line before the next turn lifecycle begins. @@ -2907,14 +2882,6 @@ impl ThreadRequestProcessor { ) .await; } - send_persisted_side_state_to_connection( - &self.outgoing, - connection_id, - thread_id, - response_history.get_rollout_items(), - self.config.features.enabled(Feature::Goals), - ) - .await; self.thread_goal_processor .emit_resume_goal_snapshot_and_continue(thread_id, codex_thread.as_ref()) .await; @@ -3941,8 +3908,8 @@ fn build_thread_turns_page_response( has_live_running_thread: bool, active_turn: Option, options: ThreadTurnsPageOptions<'_>, -) -> Result<(ThreadTurnsListResponse, Option), JSONRPCErrorError> { - let (mut turns, token_usage_turn_id) = reconstruct_thread_turns_for_turns_list( +) -> Result { + let mut turns = reconstruct_thread_turns_for_turns_list( items, loaded_status, has_live_running_thread, @@ -3950,14 +3917,11 @@ fn build_thread_turns_page_response( ); apply_thread_turns_items_view(&mut turns, options.items_view); let page = paginate_thread_turns(turns, options.cursor, options.limit, options.sort_direction)?; - Ok(( - ThreadTurnsListResponse { - data: page.turns, - next_cursor: page.next_cursor, - backwards_cursor: page.backwards_cursor, - }, - token_usage_turn_id, - )) + Ok(ThreadTurnsListResponse { + data: page.turns, + next_cursor: page.next_cursor, + backwards_cursor: page.backwards_cursor, + }) } pub(super) fn build_thread_resume_initial_turns_page( @@ -3966,7 +3930,7 @@ pub(super) fn build_thread_resume_initial_turns_page( has_live_running_thread: bool, active_turn: Option, params: &ThreadResumeInitialTurnsPageParams, -) -> Result<(codex_app_server_protocol::TurnsPage, Option), JSONRPCErrorError> { +) -> Result { build_thread_turns_page_response( items, loaded_status, @@ -3979,7 +3943,7 @@ pub(super) fn build_thread_resume_initial_turns_page( items_view: params.items_view.unwrap_or(TurnItemsView::Summary), }, ) - .map(|(response, token_usage_turn_id)| (response.into(), token_usage_turn_id)) + .map(Into::into) } fn apply_thread_turns_items_view(turns: &mut [Turn], items_view: TurnItemsView) { @@ -4025,29 +3989,17 @@ fn reconstruct_thread_turns_for_turns_list( loaded_status: ThreadStatus, has_live_running_thread: bool, active_turn: Option, -) -> (Vec, Option) { +) -> Vec { let has_live_in_progress_turn = has_live_running_thread || active_turn .as_ref() .is_some_and(|turn| matches!(turn.status, TurnStatus::InProgress)); - let mut builder = ThreadHistoryBuilder::new(); - let mut token_usage_turn_owner = None; - for item in items { - if !is_persisted_rollout_item(item) { - continue; - } - if matches!(item, RolloutItem::EventMsg(EventMsg::TokenCount(_))) { - token_usage_turn_owner = TokenUsageTurnOwner::from_history_builder(&builder); - } - builder.handle_rollout_item(item); - } - let mut turns = builder.finish(); + let mut turns = build_api_turns_from_rollout_items(items); normalize_thread_turns_status(&mut turns, loaded_status, has_live_in_progress_turn); if let Some(active_turn) = active_turn { merge_turn_history_with_active_turn(&mut turns, active_turn); } - let token_usage_turn_id = token_usage_turn_owner.and_then(|owner| owner.resolve(&turns)); - (turns, token_usage_turn_id) + turns } fn normalize_thread_turns_status( diff --git a/codex-rs/app-server/src/request_processors/thread_processor_tests.rs b/codex-rs/app-server/src/request_processors/thread_processor_tests.rs index 4ebf061e3c..c1de4d795b 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor_tests.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor_tests.rs @@ -1,3 +1,30 @@ +#[test] +fn persisted_terminal_turn_wins_over_stale_active_snapshot() { + use codex_app_server_protocol::Turn; + use codex_app_server_protocol::TurnItemsView; + use codex_app_server_protocol::TurnStatus; + let persisted = Turn { + id: "turn-1".to_string(), + items: Vec::new(), + items_view: TurnItemsView::Full, + status: TurnStatus::Completed, + error: None, + started_at: Some(1), + completed_at: Some(2), + duration_ms: Some(1_000), + }; + let mut stale = persisted.clone(); + stale.status = TurnStatus::InProgress; + stale.completed_at = None; + stale.duration_ms = None; + let mut turns = vec![persisted]; + + assert!( + !super::super::thread_lifecycle::merge_turn_history_with_active_turn(&mut turns, stale,) + ); + assert_eq!(turns[0].status, TurnStatus::Completed); +} + mod thread_list_cwd_filter_tests { use super::super::normalize_thread_list_cwd_filters; use codex_app_server_protocol::ThreadListCwdFilter; @@ -314,7 +341,7 @@ mod thread_processor_behavior_tests { duration_ms: None, }; - let (turns, _) = reconstruct_thread_turns_for_turns_list( + let turns = reconstruct_thread_turns_for_turns_list( &persisted_items, ThreadStatus::Idle, /*has_live_running_thread*/ false, diff --git a/codex-rs/app-server/src/request_processors/token_usage_replay.rs b/codex-rs/app-server/src/request_processors/token_usage_replay.rs index 61c5706187..d4d2228d4b 100644 --- a/codex-rs/app-server/src/request_processors/token_usage_replay.rs +++ b/codex-rs/app-server/src/request_processors/token_usage_replay.rs @@ -61,31 +61,11 @@ pub(super) async fn send_thread_token_usage_update_to_connection( /// /// The id is preferred when it still appears in the rebuilt thread. The position is a /// fallback for histories whose implicit turn ids are regenerated during reconstruction. -pub(super) struct TokenUsageTurnOwner { +struct TokenUsageTurnOwner { id: String, position: Option, } -impl TokenUsageTurnOwner { - pub(super) fn from_history_builder(builder: &ThreadHistoryBuilder) -> Option { - let position = builder.active_turn_position(); - builder.active_turn_snapshot().map(|turn| Self { - id: turn.id, - position, - }) - } - - pub(super) fn resolve(self, turns: &[Turn]) -> Option { - if turns.iter().any(|turn| turn.id == self.id) { - Some(self.id) - } else { - self.position - .and_then(|position| turns.get(position)) - .map(|turn| turn.id.clone()) - } - } -} - pub(super) fn latest_token_usage_turn_id_from_rollout_items( rollout_items: &[RolloutItem], turns: &[Turn], @@ -95,12 +75,26 @@ pub(super) fn latest_token_usage_turn_id_from_rollout_items( for item in rollout_items { if matches!(item, RolloutItem::EventMsg(EventMsg::TokenCount(_))) { - token_usage_turn_owner = TokenUsageTurnOwner::from_history_builder(&builder); + token_usage_turn_owner = + builder + .active_turn_snapshot() + .map(|turn| TokenUsageTurnOwner { + id: turn.id, + position: builder.active_turn_position(), + }); } builder.handle_rollout_item(item); } - token_usage_turn_owner.and_then(|owner| owner.resolve(turns)) + let owner = token_usage_turn_owner?; + if turns.iter().any(|turn| turn.id == owner.id) { + Some(owner.id) + } else { + owner + .position + .and_then(|position| turns.get(position)) + .map(|turn| turn.id.clone()) + } } /// Chooses a fallback turn id that should own a replayed token usage update. diff --git a/codex-rs/app-server/src/thread_state.rs b/codex-rs/app-server/src/thread_state.rs index 5d27ed6fc2..0d204c9d50 100644 --- a/codex-rs/app-server/src/thread_state.rs +++ b/codex-rs/app-server/src/thread_state.rs @@ -11,7 +11,6 @@ use codex_core::ThreadConfigSnapshot; use codex_file_watcher::WatchRegistration; use codex_protocol::ThreadId; use codex_protocol::protocol::EventMsg; -use codex_protocol::protocol::ThreadGoalUpdatedEvent; use codex_rollout::state_db::StateDbHandle; use codex_utils_path_uri::LegacyAppPathString; use std::collections::HashMap; @@ -49,9 +48,6 @@ pub(crate) enum ThreadListenerCommand { turn_id: Option, goal: ThreadGoal, }, - // PersistAndEmitThreadGoalUpdated is used for extension-originated updates that have not yet - // passed through the live thread's ordered rollout writer. - PersistAndEmitThreadGoalUpdated(ThreadGoalUpdatedEvent), // EmitThreadGoalCleared is used to order app-server goal clears with running-thread resume responses. EmitThreadGoalCleared, // EmitThreadGoalSnapshot is used to read and emit the latest goal state in the listener order. diff --git a/codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs b/codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs index 14afd4419d..cc9d2d781b 100644 --- a/codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs +++ b/codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs @@ -9,7 +9,6 @@ use base64::engine::general_purpose::URL_SAFE_NO_PAD; use codex_app_server_protocol::ClientInfo; use codex_app_server_protocol::InitializeCapabilities; use codex_app_server_protocol::InitializeParams; -use codex_app_server_protocol::ItemStartedNotification; use codex_app_server_protocol::JSONRPCError; use codex_app_server_protocol::JSONRPCMessage; use codex_app_server_protocol::JSONRPCNotification; @@ -459,10 +458,21 @@ async fn websocket_reconnect_resumes_in_progress_agent_message_deltas() -> Resul drop(ws1); let mut ws2 = connect_websocket(bind_addr).await?; - send_initialize_request_with_experimental_api( + send_request( &mut ws2, + "initialize", /*id*/ 5, - "codex_chatgpt_ios_remote", + Some(serde_json::to_value(InitializeParams { + client_info: ClientInfo { + name: "codex_chatgpt_ios_remote".to_string(), + title: Some("WebSocket Test Client".to_string()), + version: "0.1.0".to_string(), + }, + capabilities: Some(InitializeCapabilities { + experimental_api: true, + ..Default::default() + }), + })?), ) .await?; read_response_for_id(&mut ws2, /*id*/ 5).await?; @@ -532,35 +542,6 @@ async fn websocket_reconnect_resumes_in_progress_agent_message_deltas() -> Resul matches!(item, ThreadItem::AgentMessage { id, text, .. } if id == "msg-1" && text == "hello") })); - let replayed_started = timeout( - Duration::from_secs(5), - read_notification_for_method(&mut ws2, "item/started"), - ) - .await - .context("resume did not replay the active item lifecycle")??; - let replayed_started: ItemStartedNotification = serde_json::from_value( - replayed_started - .params - .context("replayed item/started notification should include params")?, - )?; - assert_eq!( - ( - replayed_started.thread_id, - replayed_started.turn_id, - replayed_started.item, - ), - ( - thread_id.clone(), - active_turn.id, - ThreadItem::AgentMessage { - id: "msg-1".to_string(), - text: "hello".to_string(), - phase: None, - memory_citation: None, - }, - ) - ); - complete_turn_tx .send(()) .expect("streaming response should still be waiting for completion"); @@ -809,31 +790,6 @@ pub(super) async fn send_initialize_request( .await } -async fn send_initialize_request_with_experimental_api( - stream: &mut WsClient, - id: i64, - client_name: &str, -) -> Result<()> { - let params = InitializeParams { - client_info: ClientInfo { - name: client_name.to_string(), - title: Some("WebSocket Test Client".to_string()), - version: "0.1.0".to_string(), - }, - capabilities: Some(InitializeCapabilities { - experimental_api: true, - ..Default::default() - }), - }; - send_request( - stream, - "initialize", - id, - Some(serde_json::to_value(params)?), - ) - .await -} - async fn start_thread(stream: &mut WsClient, id: i64) -> Result { send_request( stream, diff --git a/codex-rs/app-server/tests/suite/v2/mod.rs b/codex-rs/app-server/tests/suite/v2/mod.rs index 7d67cea4b1..bbd9952754 100644 --- a/codex-rs/app-server/tests/suite/v2/mod.rs +++ b/codex-rs/app-server/tests/suite/v2/mod.rs @@ -69,7 +69,6 @@ mod thread_metadata_update; mod thread_name_websocket; mod thread_read; mod thread_resume; -mod thread_resume_side_state; mod thread_rollback; mod thread_settings_update; mod thread_shell_command; diff --git a/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs b/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs index 5b24b9c6b2..2c8914f90b 100644 --- a/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs +++ b/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs @@ -26,7 +26,6 @@ use codex_app_server::in_process::InProcessServerEvent; use codex_app_server::in_process::InProcessStartArgs; use codex_app_server_protocol::ClientInfo; use codex_app_server_protocol::ClientRequest; -use codex_app_server_protocol::InitializeCapabilities; use codex_app_server_protocol::InitializeParams; use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ServerNotification; @@ -35,13 +34,10 @@ use codex_app_server_protocol::ThreadDeleteResponse; use codex_app_server_protocol::ThreadItem; use codex_app_server_protocol::ThreadListParams; use codex_app_server_protocol::ThreadListResponse; -use codex_app_server_protocol::ThreadResumeInitialTurnsPageParams; use codex_app_server_protocol::ThreadResumeParams; use codex_app_server_protocol::ThreadResumeResponse; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; -use codex_app_server_protocol::TurnCompletedNotification; -use codex_app_server_protocol::TurnItemsView; use codex_app_server_protocol::TurnStartParams; use codex_app_server_protocol::TurnStatus; use codex_app_server_protocol::UserInput as V2UserInput; @@ -345,7 +341,8 @@ async fn running_thread_resume_includes_completion_after_initial_history_snapsho .await .context("waiting for turn/started before pausing resume history")?; - let mut snapshot_pause = thread_store.pause_next_read_thread_after_snapshot().await; + let (snapshot_taken, release_snapshot) = + thread_store.pause_next_read_thread_after_snapshot().await; let resume = tokio::spawn({ let sender = client.sender(); let thread_id = thread.id.clone(); @@ -355,12 +352,8 @@ async fn running_thread_resume_includes_completion_after_initial_history_snapsho request_id: RequestId::Integer(3), params: ThreadResumeParams { thread_id, - exclude_turns: true, - initial_turns_page: Some(ThreadResumeInitialTurnsPageParams { - limit: None, - sort_direction: None, - items_view: Some(TurnItemsView::Full), - }), + exclude_turns: false, + initial_turns_page: None, ..Default::default() }, }) @@ -368,21 +361,9 @@ async fn running_thread_resume_includes_completion_after_initial_history_snapsho } }); - if timeout( - DEFAULT_READ_TIMEOUT, - snapshot_pause.wait_until_snapshot_taken(), - ) - .await - .is_err() - { - if resume.is_finished() { - let response = resume.await?; - anyhow::bail!( - "thread/resume finished before capturing its initial history: {response:?}" - ); - } - anyhow::bail!("timed out waiting for thread/resume to capture its initial history"); - } + timeout(DEFAULT_READ_TIMEOUT, snapshot_taken) + .await + .context("waiting for thread/resume to capture its initial history")??; complete_turn_tx .send(()) .expect("streaming response should still be waiting for completion"); @@ -391,19 +372,18 @@ async fn running_thread_resume_includes_completion_after_initial_history_snapsho }) .await .context("waiting for turn/completed while the resume history read is paused")?; - snapshot_pause.release(); + release_snapshot + .send(()) + .expect("paused history read should still be waiting for release"); let response = timeout(DEFAULT_READ_TIMEOUT, resume) .await .context("waiting for thread/resume after releasing its initial history read")?? .expect("thread/resume transport should succeed") .expect("thread/resume should succeed"); - let ThreadResumeResponse { - initial_turns_page, .. - } = serde_json::from_value(response)?; - let resumed_turn = initial_turns_page - .expect("resume should include the requested initial turns page") - .data + let ThreadResumeResponse { thread, .. } = serde_json::from_value(response)?; + let resumed_turn = thread + .turns .into_iter() .find(|turn| { turn.items.iter().any(|item| { @@ -425,97 +405,6 @@ async fn running_thread_resume_includes_completion_after_initial_history_snapsho Ok(()) } -#[tokio::test] -async fn running_thread_resume_preserves_failed_turn_status_and_error() -> Result<()> { - let server = responses::start_mock_server().await; - let _response_mock = responses::mount_sse_once( - &server, - responses::sse_failed("resp-1", "server_error", "simulated failure"), - ) - .await; - let codex_home = TempDir::new()?; - let store_id = Uuid::new_v4().to_string(); - create_config_toml_with_thread_store(codex_home.path(), &server.uri(), &store_id)?; - - let _in_memory_store = InMemoryThreadStoreId { - store_id: store_id.clone(), - }; - let _thread_store = InMemoryThreadStore::for_id(store_id); - let mut client = start_in_process_server(codex_home.path()).await?; - - let response = client - .request(ClientRequest::ThreadStart { - request_id: RequestId::Integer(1), - params: ThreadStartParams::default(), - }) - .await? - .expect("thread/start should succeed"); - let ThreadStartResponse { thread, .. } = serde_json::from_value(response)?; - let response = client - .request(ClientRequest::TurnStart { - request_id: RequestId::Integer(2), - params: TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "Fail this turn".to_string(), - text_elements: Vec::new(), - }], - ..Default::default() - }, - }) - .await? - .expect("turn/start should succeed"); - let turn_id = serde_json::from_value::(response)? - .turn - .id; - let notification = wait_for_thread_notification(&mut client, &thread.id, |notification| { - matches!(notification, ServerNotification::TurnCompleted(_)) - }) - .await?; - let ServerNotification::TurnCompleted(TurnCompletedNotification { - turn: live_turn, .. - }) = notification - else { - unreachable!("notification predicate accepts only turn/completed"); - }; - assert_eq!(live_turn.id, turn_id); - assert_eq!(live_turn.status, TurnStatus::Failed); - let live_error = live_turn - .error - .expect("live failed turn should include its error"); - - let response = client - .request(ClientRequest::ThreadResume { - request_id: RequestId::Integer(3), - params: ThreadResumeParams { - thread_id: thread.id, - exclude_turns: true, - initial_turns_page: Some(ThreadResumeInitialTurnsPageParams { - limit: None, - sort_direction: None, - items_view: Some(TurnItemsView::Full), - }), - ..Default::default() - }, - }) - .await? - .expect("thread/resume should succeed"); - let response: ThreadResumeResponse = serde_json::from_value(response)?; - let resumed_turn = response - .initial_turns_page - .expect("resume should include the requested initial turns page") - .data - .into_iter() - .find(|turn| turn.id == turn_id) - .expect("resume should include the failed turn"); - assert_eq!(resumed_turn.status, TurnStatus::Failed); - assert_eq!(resumed_turn.error, Some(live_error)); - - client.shutdown().await?; - Ok(()) -} - async fn wait_for_thread_notification( client: &mut InProcessClientHandle, thread_id: &str, @@ -581,10 +470,7 @@ async fn start_in_process_client( title: None, version: "0.1.0".to_string(), }, - capabilities: Some(InitializeCapabilities { - experimental_api: true, - ..Default::default() - }), + capabilities: None, }, channel_capacity: in_process::DEFAULT_IN_PROCESS_CHANNEL_CAPACITY, }) diff --git a/codex-rs/app-server/tests/suite/v2/thread_resume.rs b/codex-rs/app-server/tests/suite/v2/thread_resume.rs index b4e0fe6b7c..ce95f4810c 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_resume.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_resume.rs @@ -30,7 +30,6 @@ use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ServerNotification; use codex_app_server_protocol::ServerRequest; use codex_app_server_protocol::SessionSource; -use codex_app_server_protocol::SortDirection; use codex_app_server_protocol::ThreadGoalClearResponse; use codex_app_server_protocol::ThreadGoalSetResponse; use codex_app_server_protocol::ThreadGoalStatus; @@ -48,7 +47,6 @@ use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; use codex_app_server_protocol::ThreadStatus; use codex_app_server_protocol::ThreadUnsubscribeParams; -use codex_app_server_protocol::TurnCompletedNotification; use codex_app_server_protocol::TurnItemsView; use codex_app_server_protocol::TurnStartParams; use codex_app_server_protocol::TurnStartResponse; @@ -1138,203 +1136,6 @@ async fn thread_resume_keeps_paused_goal_paused() -> Result<()> { Ok(()) } -#[tokio::test] -async fn thread_resume_after_server_restart_replays_completed_goal_with_original_turn_id() --> Result<()> { - let update_goal_arguments = serde_json::to_string(&json!({ "status": "complete" }))?; - let server = create_mock_responses_server_sequence_unchecked(vec![ - responses::sse(vec![ - responses::ev_response_created("complete-goal"), - responses::ev_function_call( - "call-complete-goal", - "update_goal", - &update_goal_arguments, - ), - responses::ev_completed("complete-goal"), - ]), - responses::sse(vec![ - responses::ev_response_created("finish-turn"), - responses::ev_assistant_message("finished", "Goal complete"), - responses::ev_completed("finish-turn"), - ]), - ]) - .await; - let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; - let config_path = codex_home.path().join("config.toml"); - let config = std::fs::read_to_string(&config_path)?; - std::fs::write( - &config_path, - config.replace("personality = true\n", "personality = true\ngoals = true\n"), - )?; - - let mut first_mcp = TestAppServer::new_without_managed_config(codex_home.path()).await?; - timeout(DEFAULT_READ_TIMEOUT, first_mcp.initialize()).await??; - - let start_id = first_mcp - .send_thread_start_request(ThreadStartParams { - model: Some("gpt-5.2-codex".to_string()), - ..Default::default() - }) - .await?; - let start_response = timeout( - DEFAULT_READ_TIMEOUT, - first_mcp.read_stream_until_response_message(RequestId::Integer(start_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response(start_response)?; - - let goal_id = first_mcp - .send_raw_request( - "thread/goal/set", - Some(json!({ - "threadId": thread.id, - "objective": "preserve completed goal attribution", - })), - ) - .await?; - let goal_response = timeout( - DEFAULT_READ_TIMEOUT, - first_mcp.read_stream_until_response_message(RequestId::Integer(goal_id)), - ) - .await??; - let _: ThreadGoalSetResponse = to_response(goal_response)?; - timeout( - DEFAULT_READ_TIMEOUT, - first_mcp.read_stream_until_notification_message("thread/goal/updated"), - ) - .await??; - - let turn_start_id = first_mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - input: vec![UserInput::Text { - text: "Complete the active goal".to_string(), - text_elements: Vec::new(), - }], - ..Default::default() - }) - .await?; - let turn_start_response = timeout( - DEFAULT_READ_TIMEOUT, - first_mcp.read_stream_until_response_message(RequestId::Integer(turn_start_id)), - ) - .await??; - let _: TurnStartResponse = to_response(turn_start_response)?; - let turn_started = timeout( - DEFAULT_READ_TIMEOUT, - first_mcp.read_stream_until_notification_message("turn/started"), - ) - .await??; - let turn_started: ServerNotification = turn_started.try_into()?; - let ServerNotification::TurnStarted(turn_started) = turn_started else { - anyhow::bail!("expected turn started notification"); - }; - let live_goal_update = timeout( - DEFAULT_READ_TIMEOUT, - first_mcp.read_stream_until_matching_notification( - "completed thread goal update", - |notification| { - notification.method == "thread/goal/updated" - && notification - .params - .as_ref() - .is_some_and(|params| params["goal"]["status"].as_str() == Some("complete")) - }, - ), - ) - .await??; - let live_goal_update: ServerNotification = live_goal_update.try_into()?; - let ServerNotification::ThreadGoalUpdated(live_goal_update) = live_goal_update else { - anyhow::bail!("expected completed thread goal update notification"); - }; - assert_eq!(live_goal_update.goal.status, ThreadGoalStatus::Complete); - assert_eq!( - live_goal_update.turn_id.as_deref(), - Some(turn_started.turn.id.as_str()), - ); - timeout( - DEFAULT_READ_TIMEOUT, - first_mcp.read_stream_until_notification_message("turn/completed"), - ) - .await??; - - drop(first_mcp); - - let mut second_mcp = TestAppServer::new_without_managed_config(codex_home.path()).await?; - timeout(DEFAULT_READ_TIMEOUT, second_mcp.initialize()).await??; - let resume_id = second_mcp - .send_thread_resume_request(ThreadResumeParams { - thread_id: thread.id.clone(), - exclude_turns: true, - initial_turns_page: Some(ThreadResumeInitialTurnsPageParams { - limit: Some(5), - sort_direction: Some(SortDirection::Desc), - items_view: Some(TurnItemsView::Full), - }), - ..Default::default() - }) - .await?; - let resume_response = timeout( - DEFAULT_READ_TIMEOUT, - second_mcp.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; - let _: ThreadResumeResponse = to_response(resume_response)?; - let replayed_goal_update = timeout( - DEFAULT_READ_TIMEOUT, - second_mcp.read_stream_until_notification_message("thread/goal/updated"), - ) - .await??; - let replayed_goal_update: ServerNotification = replayed_goal_update.try_into()?; - let ServerNotification::ThreadGoalUpdated(replayed_goal_update) = replayed_goal_update else { - anyhow::bail!("expected replayed thread goal update notification"); - }; - - assert_eq!(replayed_goal_update, live_goal_update); - - let loaded_resume_id = second_mcp - .send_thread_resume_request(ThreadResumeParams { - thread_id: thread.id, - exclude_turns: true, - initial_turns_page: Some(ThreadResumeInitialTurnsPageParams { - limit: Some(5), - sort_direction: Some(SortDirection::Desc), - items_view: Some(TurnItemsView::Full), - }), - ..Default::default() - }) - .await?; - let loaded_resume_response = timeout( - DEFAULT_READ_TIMEOUT, - second_mcp.read_stream_until_response_message(RequestId::Integer(loaded_resume_id)), - ) - .await??; - let _: ThreadResumeResponse = to_response(loaded_resume_response)?; - let replayed_loaded_goal_update = timeout( - DEFAULT_READ_TIMEOUT, - second_mcp.read_stream_until_matching_notification( - "loaded completed thread goal update", - |notification| { - notification.method == "thread/goal/updated" - && notification.params.as_ref().is_some_and(|params| { - params["turnId"].as_str() == live_goal_update.turn_id.as_deref() - }) - }, - ), - ) - .await??; - let replayed_loaded_goal_update: ServerNotification = replayed_loaded_goal_update.try_into()?; - let ServerNotification::ThreadGoalUpdated(replayed_loaded_goal_update) = - replayed_loaded_goal_update - else { - anyhow::bail!("expected replayed loaded thread goal update notification"); - }; - assert_eq!(replayed_loaded_goal_update, live_goal_update); - - Ok(()) -} - #[tokio::test] async fn thread_goal_set_preserves_budget_limited_same_objective() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; @@ -1878,7 +1679,7 @@ async fn thread_resume_emits_restored_token_usage_before_next_turn() -> Result<( } #[tokio::test] -async fn thread_resume_replays_restored_token_usage_with_initial_turns_page() -> Result<()> { +async fn thread_resume_skips_restored_token_usage_when_turns_are_excluded() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; create_config_toml(codex_home.path(), &server.uri())?; @@ -1897,12 +1698,6 @@ async fn thread_resume_replays_restored_token_usage_with_initial_turns_page() -> let first_resume_id = mcp .send_thread_resume_request(ThreadResumeParams { thread_id: conversation_id.clone(), - exclude_turns: true, - initial_turns_page: Some(ThreadResumeInitialTurnsPageParams { - limit: Some(5), - sort_direction: Some(SortDirection::Desc), - items_view: Some(TurnItemsView::Full), - }), ..Default::default() }) .await?; @@ -1911,19 +1706,9 @@ async fn thread_resume_replays_restored_token_usage_with_initial_turns_page() -> mcp.read_stream_until_response_message(RequestId::Integer(first_resume_id)), ) .await??; - let ThreadResumeResponse { - thread, - initial_turns_page, - .. - } = to_response::(first_resume_resp)?; - assert!(thread.turns.is_empty()); - let expected_turn_id = initial_turns_page - .expect("resume should include the requested initial turns page") - .data - .into_iter() - .next() - .expect("initial turns page should include the token usage owner") - .id; + let ThreadResumeResponse { thread, .. } = + to_response::(first_resume_resp)?; + let expected_turn_id = thread.turns[0].id.clone(); let first_note = timeout( DEFAULT_READ_TIMEOUT, @@ -1931,20 +1716,15 @@ async fn thread_resume_replays_restored_token_usage_with_initial_turns_page() -> ) .await??; let parsed: ServerNotification = first_note.try_into()?; - let ServerNotification::ThreadTokenUsageUpdated(first_notification) = parsed else { + let ServerNotification::ThreadTokenUsageUpdated(notification) = parsed else { panic!("expected thread/tokenUsage/updated notification"); }; - assert_eq!(first_notification.turn_id, expected_turn_id); + assert_eq!(notification.turn_id, expected_turn_id); let second_resume_id = mcp .send_thread_resume_request(ThreadResumeParams { thread_id: conversation_id, exclude_turns: true, - initial_turns_page: Some(ThreadResumeInitialTurnsPageParams { - limit: Some(5), - sort_direction: Some(SortDirection::Desc), - items_view: Some(TurnItemsView::Full), - }), ..Default::default() }) .await?; @@ -1955,30 +1735,19 @@ async fn thread_resume_replays_restored_token_usage_with_initial_turns_page() -> .await??; let ThreadResumeResponse { thread: resumed_again, - initial_turns_page, .. } = to_response::(second_resume_resp)?; assert!(resumed_again.turns.is_empty()); - assert_eq!( - initial_turns_page - .expect("resume should include the requested initial turns page") - .data - .into_iter() - .map(|turn| turn.id) - .collect::>(), - vec![expected_turn_id.clone()], - ); let second_note = timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("thread/tokenUsage/updated"), ) - .await??; - let parsed: ServerNotification = second_note.try_into()?; - let ServerNotification::ThreadTokenUsageUpdated(notification) = parsed else { - panic!("expected thread/tokenUsage/updated notification"); - }; - assert_eq!(notification, first_notification); + .await; + assert!( + second_note.is_err(), + "excludeTurns=true should not replay token usage" + ); Ok(()) } @@ -3861,100 +3630,6 @@ async fn thread_resume_supports_history_and_overrides() -> Result<()> { Ok(()) } -#[tokio::test] -async fn thread_resume_after_server_restart_preserves_failed_turn_status_and_error() -> Result<()> { - let server = responses::start_mock_server().await; - let _response_mock = responses::mount_sse_once( - &server, - responses::sse_failed("resp-1", "server_error", "simulated failure"), - ) - .await; - let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; - - let mut first_mcp = TestAppServer::new(codex_home.path()).await?; - timeout(DEFAULT_READ_TIMEOUT, first_mcp.initialize()).await??; - - let start_id = first_mcp - .send_thread_start_request(ThreadStartParams::default()) - .await?; - let start_response = timeout( - DEFAULT_READ_TIMEOUT, - first_mcp.read_stream_until_response_message(RequestId::Integer(start_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response(start_response)?; - - let turn_start_id = first_mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - input: vec![UserInput::Text { - text: "Fail before app-server restarts".to_string(), - text_elements: Vec::new(), - }], - ..Default::default() - }) - .await?; - let turn_start_response = timeout( - DEFAULT_READ_TIMEOUT, - first_mcp.read_stream_until_response_message(RequestId::Integer(turn_start_id)), - ) - .await??; - let TurnStartResponse { turn, .. } = to_response(turn_start_response)?; - let completed = timeout( - DEFAULT_READ_TIMEOUT, - first_mcp.read_stream_until_notification_message("turn/completed"), - ) - .await??; - let TurnCompletedNotification { - turn: live_turn, .. - } = serde_json::from_value( - completed - .params - .expect("turn/completed should include notification params"), - )?; - assert_eq!(live_turn.id, turn.id); - assert_eq!(live_turn.status, TurnStatus::Failed); - let live_error = live_turn - .error - .expect("live failed turn should include its error"); - - drop(first_mcp); - - let mut second_mcp = TestAppServer::new(codex_home.path()).await?; - timeout(DEFAULT_READ_TIMEOUT, second_mcp.initialize()).await??; - let resume_id = second_mcp - .send_thread_resume_request(ThreadResumeParams { - thread_id: thread.id, - exclude_turns: true, - initial_turns_page: Some(ThreadResumeInitialTurnsPageParams { - limit: Some(5), - sort_direction: Some(SortDirection::Desc), - items_view: Some(TurnItemsView::Full), - }), - ..Default::default() - }) - .await?; - let resume_response = timeout( - DEFAULT_READ_TIMEOUT, - second_mcp.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; - let ThreadResumeResponse { - initial_turns_page, .. - } = to_response(resume_response)?; - let resumed_turn = initial_turns_page - .expect("resume should include the requested initial turns page") - .data - .into_iter() - .find(|resumed| resumed.id == turn.id) - .expect("resume should reconstruct the failed turn after restart"); - assert_eq!(resumed_turn.status, TurnStatus::Failed); - assert_eq!(resumed_turn.error, Some(live_error)); - - Ok(()) -} - struct RestartedThreadFixture { mcp: TestAppServer, thread_id: String, diff --git a/codex-rs/app-server/tests/suite/v2/thread_resume_side_state.rs b/codex-rs/app-server/tests/suite/v2/thread_resume_side_state.rs deleted file mode 100644 index e232e5d954..0000000000 --- a/codex-rs/app-server/tests/suite/v2/thread_resume_side_state.rs +++ /dev/null @@ -1,296 +0,0 @@ -use std::path::Path; -use std::time::Duration; - -use anyhow::Result; -use app_test_support::TestAppServer; -use app_test_support::create_fake_rollout; -use app_test_support::create_mock_responses_server_repeating_assistant; -use app_test_support::rollout_path; -use app_test_support::test_absolute_path; -use app_test_support::to_response; -use codex_app_server_protocol::AutoReviewDecisionSource; -use codex_app_server_protocol::GuardianApprovalReview; -use codex_app_server_protocol::GuardianApprovalReviewAction; -use codex_app_server_protocol::GuardianApprovalReviewStatus; -use codex_app_server_protocol::GuardianCommandSource; -use codex_app_server_protocol::GuardianRiskLevel; -use codex_app_server_protocol::GuardianUserAuthorization; -use codex_app_server_protocol::HookCompletedNotification; -use codex_app_server_protocol::HookEventName; -use codex_app_server_protocol::HookExecutionMode; -use codex_app_server_protocol::HookHandlerType; -use codex_app_server_protocol::HookOutputEntry; -use codex_app_server_protocol::HookOutputEntryKind; -use codex_app_server_protocol::HookRunStatus; -use codex_app_server_protocol::HookRunSummary; -use codex_app_server_protocol::HookScope; -use codex_app_server_protocol::HookSource; -use codex_app_server_protocol::ItemGuardianApprovalReviewCompletedNotification; -use codex_app_server_protocol::JSONRPCResponse; -use codex_app_server_protocol::RequestId; -use codex_app_server_protocol::ServerNotification; -use codex_app_server_protocol::SortDirection; -use codex_app_server_protocol::ThreadResumeInitialTurnsPageParams; -use codex_app_server_protocol::ThreadResumeParams; -use codex_app_server_protocol::ThreadResumeResponse; -use codex_app_server_protocol::TurnItemsView; -use codex_protocol::protocol::EventMsg; -use codex_protocol::protocol::GuardianAssessmentAction; -use codex_protocol::protocol::GuardianAssessmentDecisionSource; -use codex_protocol::protocol::GuardianAssessmentEvent; -use codex_protocol::protocol::GuardianAssessmentStatus; -use codex_protocol::protocol::GuardianCommandSource as CoreGuardianCommandSource; -use codex_protocol::protocol::GuardianRiskLevel as CoreGuardianRiskLevel; -use codex_protocol::protocol::GuardianUserAuthorization as CoreGuardianUserAuthorization; -use codex_protocol::protocol::HookCompletedEvent; -use codex_protocol::protocol::HookEventName as CoreHookEventName; -use codex_protocol::protocol::HookExecutionMode as CoreHookExecutionMode; -use codex_protocol::protocol::HookHandlerType as CoreHookHandlerType; -use codex_protocol::protocol::HookOutputEntry as CoreHookOutputEntry; -use codex_protocol::protocol::HookOutputEntryKind as CoreHookOutputEntryKind; -use codex_protocol::protocol::HookRunStatus as CoreHookRunStatus; -use codex_protocol::protocol::HookRunSummary as CoreHookRunSummary; -use codex_protocol::protocol::HookScope as CoreHookScope; -use codex_protocol::protocol::HookSource as CoreHookSource; -use codex_protocol::protocol::RolloutItem; -use codex_protocol::protocol::TurnStartedEvent; -use codex_rollout::append_rollout_item_to_path; -use pretty_assertions::assert_eq; -use tempfile::TempDir; -use tokio::time::timeout; - -use super::connection_handling_websocket::create_config_toml; - -const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(10); -const REPLAY_TIMEOUT: Duration = Duration::from_secs(1); -const FILENAME_TIMESTAMP: &str = "2025-01-05T12-00-00"; -const META_TIMESTAMP: &str = "2025-01-05T12:00:00Z"; - -#[tokio::test] -async fn thread_resume_replays_persisted_terminal_auto_approval_review() -> Result<()> { - let server = create_mock_responses_server_repeating_assistant("Done").await; - let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri(), "never")?; - - let turn_id = "turn-offline-auto-review"; - let review_id = "review-offline-command"; - let target_item_id = "command-offline"; - let command = "rm -f /tmp/offline-review.sqlite"; - let cwd = test_absolute_path("/tmp"); - let thread_id = create_rollout_with_side_state( - codex_home.path(), - turn_id, - EventMsg::GuardianAssessment(GuardianAssessmentEvent { - id: review_id.to_string(), - target_item_id: Some(target_item_id.to_string()), - turn_id: turn_id.to_string(), - started_at_ms: 1_000, - completed_at_ms: Some(1_042), - status: GuardianAssessmentStatus::Approved, - risk_level: Some(CoreGuardianRiskLevel::Low), - user_authorization: Some(CoreGuardianUserAuthorization::High), - rationale: Some("The saved command is safe.".to_string()), - decision_source: Some(GuardianAssessmentDecisionSource::Agent), - action: GuardianAssessmentAction::Command { - source: CoreGuardianCommandSource::Shell, - command: command.to_string(), - cwd: cwd.clone(), - }, - }), - ) - .await?; - - let mut app_server = TestAppServer::new(codex_home.path()).await?; - timeout(DEFAULT_READ_TIMEOUT, app_server.initialize()).await??; - let expected = ItemGuardianApprovalReviewCompletedNotification { - thread_id, - turn_id: turn_id.to_string(), - started_at_ms: 1_000, - completed_at_ms: 1_042, - review_id: review_id.to_string(), - target_item_id: Some(target_item_id.to_string()), - decision_source: AutoReviewDecisionSource::Agent, - review: GuardianApprovalReview { - status: GuardianApprovalReviewStatus::Approved, - risk_level: Some(GuardianRiskLevel::Low), - user_authorization: Some(GuardianUserAuthorization::High), - rationale: Some("The saved command is safe.".to_string()), - }, - action: GuardianApprovalReviewAction::Command { - source: GuardianCommandSource::Shell, - command: command.to_string(), - cwd, - }, - }; - - // The first resume loads the thread; the second rejoins that same running thread. - for _ in 0..2 { - let response = - resume_with_initial_turns_page(&mut app_server, expected.thread_id.as_str()).await?; - assert!(response.thread.turns.is_empty()); - assert!(response.initial_turns_page.is_some()); - - let notification = timeout( - REPLAY_TIMEOUT, - app_server.read_stream_until_notification_message("item/autoApprovalReview/completed"), - ) - .await??; - let parsed: ServerNotification = notification.try_into()?; - let ServerNotification::ItemGuardianApprovalReviewCompleted(actual) = parsed else { - unreachable!("filtered notification must be a completed auto-approval review"); - }; - assert_eq!(actual, expected); - } - - Ok(()) -} - -#[tokio::test] -async fn thread_resume_replays_persisted_completed_hook() -> Result<()> { - let server = create_mock_responses_server_repeating_assistant("Done").await; - let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri(), "never")?; - - let turn_id = "turn-offline-hook"; - let run_id = "hook-run-offline"; - let source_path = test_absolute_path("/tmp/offline-hook.sh"); - let thread_id = create_rollout_with_side_state( - codex_home.path(), - turn_id, - EventMsg::HookCompleted(HookCompletedEvent { - turn_id: Some(turn_id.to_string()), - run: CoreHookRunSummary { - id: run_id.to_string(), - event_name: CoreHookEventName::PostToolUse, - handler_type: CoreHookHandlerType::Command, - execution_mode: CoreHookExecutionMode::Sync, - scope: CoreHookScope::Turn, - source_path: source_path.clone(), - source: CoreHookSource::Project, - display_order: 7, - status: CoreHookRunStatus::Completed, - status_message: Some("Offline hook completed".to_string()), - started_at: 1_700_000_000, - completed_at: Some(1_700_000_042), - duration_ms: Some(42_000), - entries: vec![ - CoreHookOutputEntry { - kind: CoreHookOutputEntryKind::Feedback, - text: "Apply the saved formatting feedback.".to_string(), - }, - CoreHookOutputEntry { - kind: CoreHookOutputEntryKind::Context, - text: "Persist this hook context across reconnect.".to_string(), - }, - ], - }, - }), - ) - .await?; - - let mut app_server = TestAppServer::new(codex_home.path()).await?; - timeout(DEFAULT_READ_TIMEOUT, app_server.initialize()).await??; - let expected = HookCompletedNotification { - thread_id, - turn_id: Some(turn_id.to_string()), - run: HookRunSummary { - id: run_id.to_string(), - event_name: HookEventName::PostToolUse, - handler_type: HookHandlerType::Command, - execution_mode: HookExecutionMode::Sync, - scope: HookScope::Turn, - source_path, - source: HookSource::Project, - display_order: 7, - status: HookRunStatus::Completed, - status_message: Some("Offline hook completed".to_string()), - started_at: 1_700_000_000, - completed_at: Some(1_700_000_042), - duration_ms: Some(42_000), - entries: vec![ - HookOutputEntry { - kind: HookOutputEntryKind::Feedback, - text: "Apply the saved formatting feedback.".to_string(), - }, - HookOutputEntry { - kind: HookOutputEntryKind::Context, - text: "Persist this hook context across reconnect.".to_string(), - }, - ], - }, - }; - - // The first resume loads the thread; the second rejoins that same running thread. - for _ in 0..2 { - let response = - resume_with_initial_turns_page(&mut app_server, expected.thread_id.as_str()).await?; - assert!(response.thread.turns.is_empty()); - assert!(response.initial_turns_page.is_some()); - - let notification = timeout( - REPLAY_TIMEOUT, - app_server.read_stream_until_notification_message("hook/completed"), - ) - .await??; - let parsed: ServerNotification = notification.try_into()?; - let ServerNotification::HookCompleted(actual) = parsed else { - unreachable!("filtered notification must be a completed hook run"); - }; - assert_eq!(actual, expected); - } - - Ok(()) -} - -async fn create_rollout_with_side_state( - codex_home: &Path, - turn_id: &str, - event: EventMsg, -) -> Result { - let thread_id = create_fake_rollout( - codex_home, - FILENAME_TIMESTAMP, - META_TIMESTAMP, - "Saved user message", - Some("mock_provider"), - /*git_info*/ None, - )?; - let path = rollout_path(codex_home, FILENAME_TIMESTAMP, thread_id.as_str()); - append_rollout_item_to_path( - path.as_path(), - &RolloutItem::EventMsg(EventMsg::TurnStarted(TurnStartedEvent { - turn_id: turn_id.to_string(), - trace_id: None, - started_at: Some(1), - model_context_window: None, - collaboration_mode_kind: Default::default(), - })), - ) - .await?; - append_rollout_item_to_path(path.as_path(), &RolloutItem::EventMsg(event)).await?; - Ok(thread_id) -} - -async fn resume_with_initial_turns_page( - app_server: &mut TestAppServer, - thread_id: &str, -) -> Result { - let request_id = app_server - .send_thread_resume_request(ThreadResumeParams { - thread_id: thread_id.to_string(), - exclude_turns: true, - initial_turns_page: Some(ThreadResumeInitialTurnsPageParams { - limit: Some(5), - sort_direction: Some(SortDirection::Desc), - items_view: Some(TurnItemsView::Full), - }), - ..Default::default() - }) - .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - app_server.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - to_response(response) -} diff --git a/codex-rs/rollout/src/policy.rs b/codex-rs/rollout/src/policy.rs index d06702e321..ec8db3f6ff 100644 --- a/codex-rs/rollout/src/policy.rs +++ b/codex-rs/rollout/src/policy.rs @@ -107,16 +107,9 @@ pub fn should_persist_event_msg(ev: &EventMsg) -> bool { | codex_protocol::items::TurnItem::Sleep(_) ) } - EventMsg::Error(error) => error.affects_turn_status(), - EventMsg::GuardianAssessment(event) => match event.status { - codex_protocol::protocol::GuardianAssessmentStatus::InProgress => false, - codex_protocol::protocol::GuardianAssessmentStatus::Approved - | codex_protocol::protocol::GuardianAssessmentStatus::Denied - | codex_protocol::protocol::GuardianAssessmentStatus::TimedOut - | codex_protocol::protocol::GuardianAssessmentStatus::Aborted => true, - }, - EventMsg::HookCompleted(_) => true, - EventMsg::ExecCommandEnd(_) + EventMsg::Error(_) + | EventMsg::GuardianAssessment(_) + | EventMsg::ExecCommandEnd(_) | EventMsg::ViewImageToolCall(_) | EventMsg::CollabAgentSpawnEnd(_) | EventMsg::CollabAgentInteractionEnd(_) @@ -161,6 +154,7 @@ pub fn should_persist_event_msg(ev: &EventMsg) -> bool { | EventMsg::DeprecationNotice(_) | EventMsg::ItemStarted(_) | EventMsg::HookStarted(_) + | EventMsg::HookCompleted(_) | EventMsg::AgentMessageContentDelta(_) | EventMsg::PlanDelta(_) | EventMsg::ReasoningContentDelta(_) @@ -173,111 +167,3 @@ pub fn should_persist_event_msg(ev: &EventMsg) -> bool { | EventMsg::CollabResumeBegin(_) => false, } } - -#[cfg(test)] -mod tests { - use super::*; - use codex_protocol::protocol::CodexErrorInfo; - use codex_protocol::protocol::ErrorEvent; - use codex_protocol::protocol::GuardianAssessmentAction; - use codex_protocol::protocol::GuardianAssessmentEvent; - use codex_protocol::protocol::GuardianAssessmentStatus; - use codex_protocol::protocol::HookCompletedEvent; - use codex_protocol::protocol::HookEventName; - use codex_protocol::protocol::HookExecutionMode; - use codex_protocol::protocol::HookHandlerType; - use codex_protocol::protocol::HookRunStatus; - use codex_protocol::protocol::HookRunSummary; - use codex_protocol::protocol::HookScope; - use codex_protocol::protocol::HookSource; - use codex_protocol::protocol::HookStartedEvent; - use codex_protocol::protocol::NetworkApprovalProtocol; - use pretty_assertions::assert_eq; - - #[test] - fn persists_only_errors_that_define_terminal_turn_state() { - let terminal_error = EventMsg::Error(ErrorEvent { - message: "simulated failure".to_string(), - codex_error_info: Some(CodexErrorInfo::InternalServerError), - }); - let request_error = EventMsg::Error(ErrorEvent { - message: "rollback failed".to_string(), - codex_error_info: Some(CodexErrorInfo::ThreadRollbackFailed), - }); - - assert!(should_persist_event_msg(&terminal_error)); - assert!(!should_persist_event_msg(&request_error)); - } - - #[test] - fn persists_only_terminal_guardian_assessments() { - let mut assessment = GuardianAssessmentEvent { - id: "review-1".to_string(), - target_item_id: None, - turn_id: "turn-1".to_string(), - started_at_ms: 1, - completed_at_ms: None, - status: GuardianAssessmentStatus::InProgress, - risk_level: None, - user_authorization: None, - rationale: None, - decision_source: None, - action: GuardianAssessmentAction::NetworkAccess { - target: "https://example.com".to_string(), - host: "example.com".to_string(), - protocol: NetworkApprovalProtocol::Https, - port: 443, - }, - }; - - for (status, expected) in [ - (GuardianAssessmentStatus::InProgress, false), - (GuardianAssessmentStatus::Approved, true), - (GuardianAssessmentStatus::Denied, true), - (GuardianAssessmentStatus::TimedOut, true), - (GuardianAssessmentStatus::Aborted, true), - ] { - assessment.status = status; - assert_eq!( - should_persist_event_msg(&EventMsg::GuardianAssessment(assessment.clone())), - expected, - ); - } - } - - #[test] - fn persists_completed_hook_events_but_not_started_events() { - let run = HookRunSummary { - id: "hook-1".to_string(), - event_name: HookEventName::PostToolUse, - handler_type: HookHandlerType::Command, - execution_mode: HookExecutionMode::Sync, - scope: HookScope::Turn, - source_path: std::env::current_dir() - .expect("current directory should be available") - .try_into() - .expect("current directory should be absolute"), - source: HookSource::Project, - display_order: 0, - status: HookRunStatus::Completed, - status_message: None, - started_at: 1, - completed_at: Some(2), - duration_ms: Some(1_000), - entries: Vec::new(), - }; - - assert!(!should_persist_event_msg(&EventMsg::HookStarted( - HookStartedEvent { - turn_id: Some("turn-1".to_string()), - run: run.clone(), - }, - ))); - assert!(should_persist_event_msg(&EventMsg::HookCompleted( - HookCompletedEvent { - turn_id: Some("turn-1".to_string()), - run, - }, - ))); - } -} diff --git a/codex-rs/thread-store/src/in_memory.rs b/codex-rs/thread-store/src/in_memory.rs index 07dd066c50..828e324694 100644 --- a/codex-rs/thread-store/src/in_memory.rs +++ b/codex-rs/thread-store/src/in_memory.rs @@ -233,30 +233,6 @@ pub struct InMemoryThreadStore { state: tokio::sync::Mutex, } -/// Controls a one-shot pause after [`InMemoryThreadStore::read_thread`] has captured its result. -/// -/// This lets integration tests advance concurrent appends while a reader still owns an older -/// snapshot, without holding the store lock and preventing those appends. -pub struct InMemoryReadThreadSnapshotPause { - snapshot_taken: oneshot::Receiver<()>, - release: oneshot::Sender<()>, -} - -impl InMemoryReadThreadSnapshotPause { - /// Waits until the paused read has captured its result and released the store lock. - pub async fn wait_until_snapshot_taken(&mut self) { - assert!( - (&mut self.snapshot_taken).await.is_ok(), - "paused in-memory thread read should signal its snapshot" - ); - } - - /// Allows the paused read to return its captured result. - pub fn release(self) { - let _ = self.release.send(()); - } -} - struct PendingReadThreadSnapshotPause { snapshot_taken: oneshot::Sender<()>, release: oneshot::Receiver<()>, @@ -295,7 +271,9 @@ impl InMemoryThreadStore { } /// Pauses the next `read_thread` call after it captures a snapshot. - pub async fn pause_next_read_thread_after_snapshot(&self) -> InMemoryReadThreadSnapshotPause { + pub async fn pause_next_read_thread_after_snapshot( + &self, + ) -> (oneshot::Receiver<()>, oneshot::Sender<()>) { let (snapshot_taken_tx, snapshot_taken) = oneshot::channel(); let (release, release_rx) = oneshot::channel(); let mut state = self.state.lock().await; @@ -307,10 +285,7 @@ impl InMemoryThreadStore { snapshot_taken: snapshot_taken_tx, release: release_rx, }); - InMemoryReadThreadSnapshotPause { - snapshot_taken, - release, - } + (snapshot_taken, release) } async fn create_thread(&self, params: CreateThreadParams) -> ThreadStoreResult<()> { diff --git a/codex-rs/thread-store/src/lib.rs b/codex-rs/thread-store/src/lib.rs index 11a7a270db..34a0ab2ba1 100644 --- a/codex-rs/thread-store/src/lib.rs +++ b/codex-rs/thread-store/src/lib.rs @@ -14,7 +14,6 @@ mod types; pub use error::ThreadStoreError; pub use error::ThreadStoreResult; -pub use in_memory::InMemoryReadThreadSnapshotPause; pub use in_memory::InMemoryThreadStore; pub use in_memory::InMemoryThreadStoreCalls; pub use live_thread::LiveThread;