From 0e597b272c0437353e1fabeba87d65d3a5709f91 Mon Sep 17 00:00:00 2001 From: Javier Soto Date: Wed, 24 Jun 2026 17:06:18 -0700 Subject: [PATCH] Harden live resume reconstruction --- .../src/protocol/thread_history.rs | 261 ++--- .../src/protocol/thread_history_live_tests.rs | 958 ++++++++++++++++++ .../request_processors/thread_lifecycle.rs | 32 +- .../request_processors/thread_processor.rs | 28 +- .../thread_resume_reconciliation_tests.rs | 582 +++++++++++ .../suite/v2/connection_handling_websocket.rs | 157 ++- .../tests/suite/v2/remote_thread_store.rs | 120 ++- 7 files changed, 1955 insertions(+), 183 deletions(-) create mode 100644 codex-rs/app-server-protocol/src/protocol/thread_history_live_tests.rs create mode 100644 codex-rs/app-server/src/request_processors/thread_resume_reconciliation_tests.rs 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 1abacf9dc7..79cb13713c 100644 --- a/codex-rs/app-server-protocol/src/protocol/thread_history.rs +++ b/codex-rs/app-server-protocol/src/protocol/thread_history.rs @@ -56,6 +56,7 @@ use codex_protocol::protocol::ViewImageToolCallEvent; use codex_protocol::protocol::WebSearchBeginEvent; use codex_protocol::protocol::WebSearchEndEvent; use std::collections::HashMap; +use std::collections::HashSet; use std::collections::VecDeque; use tracing::warn; use uuid::Uuid; @@ -235,6 +236,9 @@ pub struct ThreadHistoryBuilder { active_change_set: Option, /// Legacy events generated after typed item completion, retained across unrelated live events. pending_completed_item_legacy_events: VecDeque, + completed_item_ids_by_turn: HashMap>, + command_turn_ids: HashMap, + pending_command_output: HashMap>, } impl Default for ThreadHistoryBuilder { @@ -253,6 +257,9 @@ impl ThreadHistoryBuilder { next_rollout_index: 0, active_change_set: None, pending_completed_item_legacy_events: VecDeque::new(), + completed_item_ids_by_turn: HashMap::new(), + command_turn_ids: HashMap::new(), + pending_command_output: HashMap::new(), } } @@ -594,6 +601,10 @@ impl ThreadHistoryBuilder { } fn handle_item_started(&mut self, payload: &ItemStartedEvent) { + let item_id = payload.item.id(); + if self.is_item_completed(&payload.turn_id, &item_id) { + return; + } match &payload.item { codex_protocol::items::TurnItem::AgentMessage(_) | codex_protocol::items::TurnItem::Plan(_) @@ -619,7 +630,10 @@ impl ThreadHistoryBuilder { &mut self, payload: &codex_protocol::protocol::AgentMessageContentDeltaEvent, ) { - self.update_item_in_active_turn(&payload.turn_id, &payload.item_id, |item| { + if self.is_item_completed(&payload.turn_id, &payload.item_id) { + return; + } + self.update_item_in_turn_id(&payload.turn_id, &payload.item_id, |item| { let ThreadItem::AgentMessage { text, .. } = item else { return false; }; @@ -629,7 +643,10 @@ impl ThreadHistoryBuilder { } fn handle_plan_delta(&mut self, payload: &codex_protocol::protocol::PlanDeltaEvent) { - self.update_item_in_active_turn(&payload.turn_id, &payload.item_id, |item| { + if self.is_item_completed(&payload.turn_id, &payload.item_id) { + return; + } + self.update_item_in_turn_id(&payload.turn_id, &payload.item_id, |item| { let ThreadItem::Plan { text, .. } = item else { return false; }; @@ -642,7 +659,10 @@ impl ThreadHistoryBuilder { &mut self, payload: &codex_protocol::protocol::ReasoningContentDeltaEvent, ) { - self.update_item_in_active_turn(&payload.turn_id, &payload.item_id, |item| { + if self.is_item_completed(&payload.turn_id, &payload.item_id) { + return; + } + self.update_item_in_turn_id(&payload.turn_id, &payload.item_id, |item| { let ThreadItem::Reasoning { summary, .. } = item else { return false; }; @@ -654,7 +674,10 @@ impl ThreadHistoryBuilder { &mut self, payload: &codex_protocol::protocol::ReasoningRawContentDeltaEvent, ) { - self.update_item_in_active_turn(&payload.turn_id, &payload.item_id, |item| { + if self.is_item_completed(&payload.turn_id, &payload.item_id) { + return; + } + self.update_item_in_turn_id(&payload.turn_id, &payload.item_id, |item| { let ThreadItem::Reasoning { content, .. } = item else { return false; }; @@ -666,24 +689,40 @@ impl ThreadHistoryBuilder { &mut self, payload: &codex_protocol::protocol::ExecCommandOutputDeltaEvent, ) { - let Some(turn_id) = self.current_turn.as_ref().map(|turn| turn.id.clone()) else { + let Some(turn_id) = self.command_turn_ids.get(&payload.call_id).cloned() else { return; }; - self.update_item_in_active_turn(&turn_id, &payload.call_id, |item| { + if self.is_item_completed(&turn_id, &payload.call_id) { + return; + } + let output = { + let pending = self + .pending_command_output + .entry(payload.call_id.clone()) + .or_default(); + pending.extend_from_slice(&payload.chunk); + take_decodable_utf8(pending) + }; + let Some(output) = output else { + return; + }; + self.update_item_in_turn_id(&turn_id, &payload.call_id, |item| { let ThreadItem::CommandExecution { aggregated_output, .. } = item else { return false; }; - aggregated_output - .get_or_insert_default() - .push_str(&String::from_utf8_lossy(&payload.chunk)); + aggregated_output.get_or_insert_default().push_str(&output); true }); } fn handle_item_completed(&mut self, payload: &ItemCompletedEvent) { + let item_id = payload.item.id(); + if self.is_item_completed(&payload.turn_id, &item_id) { + return; + } match &payload.item { codex_protocol::items::TurnItem::AgentMessage(_) | codex_protocol::items::TurnItem::Reasoning(_) => { @@ -696,6 +735,7 @@ impl ThreadHistoryBuilder { .item .as_legacy_events(/*show_raw_agent_reasoning*/ true), ); + self.mark_item_completed(&payload.turn_id, &item_id); } codex_protocol::items::TurnItem::Plan(plan) => { if plan.text.is_empty() { @@ -705,12 +745,14 @@ impl ThreadHistoryBuilder { &payload.turn_id, ThreadItem::from(payload.item.clone()), ); + self.mark_item_completed(&payload.turn_id, &item_id); } codex_protocol::items::TurnItem::Sleep(_) => { self.upsert_item_in_turn_id( &payload.turn_id, ThreadItem::from(payload.item.clone()), ); + self.mark_item_completed(&payload.turn_id, &item_id); } codex_protocol::items::TurnItem::UserMessage(_) | codex_protocol::items::TurnItem::HookPrompt(_) @@ -754,11 +796,20 @@ impl ThreadHistoryBuilder { } fn handle_exec_command_begin(&mut self, payload: &ExecCommandBeginEvent) { + if self.is_item_completed(&payload.turn_id, &payload.call_id) { + return; + } + self.command_turn_ids + .insert(payload.call_id.clone(), payload.turn_id.clone()); + self.pending_command_output.remove(&payload.call_id); let item = build_command_execution_begin_item(payload); self.upsert_item_in_turn_id(&payload.turn_id, item); } fn handle_exec_command_end(&mut self, payload: &ExecCommandEndEvent) { + if self.is_item_completed(&payload.turn_id, &payload.call_id) { + return; + } let item = build_command_execution_end_item(payload); // Command completions can arrive out of order. Unified exec may return // while a PTY is still running, then emit ExecCommandEnd later from a @@ -766,6 +817,10 @@ impl ThreadHistoryBuilder { // newer user turn may already have started. Route by event turn_id so // replay preserves the original turn association. self.upsert_item_in_turn_id(&payload.turn_id, item); + self.command_turn_ids + .insert(payload.call_id.clone(), payload.turn_id.clone()); + self.pending_command_output.remove(&payload.call_id); + self.mark_item_completed(&payload.turn_id, &payload.call_id); } fn handle_guardian_assessment(&mut self, payload: &GuardianAssessmentEvent) { @@ -1491,30 +1546,51 @@ impl ThreadHistoryBuilder { } } - fn update_item_in_active_turn( + fn update_item_in_turn_id( &mut self, turn_id: &str, item_id: &str, update: impl FnOnce(&mut ThreadItem) -> bool, ) { let tracking_changes = self.is_tracking_changes(); - let changed_item = { - let Some(turn) = self.current_turn.as_mut().filter(|turn| turn.id == turn_id) else { + let changed_item = + if let Some(turn) = self.current_turn.as_mut().filter(|turn| turn.id == turn_id) { + let Some(item) = turn.items.iter_mut().find(|item| item.id() == item_id) else { + return; + }; + if !update(item) { + return; + } + tracking_changes.then(|| (turn.id.clone(), item.clone())) + } else if let Some(turn) = self.turns.iter_mut().find(|turn| turn.id == turn_id) { + let Some(item) = turn.items.iter_mut().find(|item| item.id() == item_id) else { + return; + }; + if !update(item) { + return; + } + tracking_changes.then(|| (turn.id.clone(), item.clone())) + } else { return; }; - let Some(item) = turn.items.iter_mut().find(|item| item.id() == item_id) else { - return; - }; - if !update(item) { - return; - } - tracking_changes.then(|| (turn.id.clone(), item.clone())) - }; if let Some((turn_id, item)) = changed_item { self.record_changed_item(turn_id, item); } } + fn is_item_completed(&self, turn_id: &str, item_id: &str) -> bool { + self.completed_item_ids_by_turn + .get(turn_id) + .is_some_and(|item_ids| item_ids.contains(item_id)) + } + + fn mark_item_completed(&mut self, turn_id: &str, item_id: &str) { + self.completed_item_ids_by_turn + .entry(turn_id.to_string()) + .or_default() + .insert(item_id.to_string()); + } + fn is_tracking_changes(&self) -> bool { self.active_change_set.is_some() } @@ -1586,13 +1662,42 @@ fn append_indexed_delta(parts: &mut Vec, index: i64, delta: &str) -> boo let Ok(index) = usize::try_from(index) else { return false; }; - if parts.len() <= index { - parts.resize(index + 1, String::new()); + if index > parts.len() { + return false; + } + if index == parts.len() { + parts.push(String::new()); } parts[index].push_str(delta); true } +fn take_decodable_utf8(bytes: &mut Vec) -> Option { + let mut output = String::new(); + loop { + match std::str::from_utf8(bytes) { + Ok(text) => { + output.push_str(text); + bytes.clear(); + break; + } + Err(error) => { + let valid_up_to = error.valid_up_to(); + if valid_up_to > 0 { + output.push_str(&String::from_utf8_lossy(&bytes[..valid_up_to])); + bytes.drain(..valid_up_to); + } + let Some(error_len) = error.error_len() else { + break; + }; + output.push('\u{FFFD}'); + bytes.drain(..error_len); + } + } + } + (!output.is_empty()).then_some(output) +} + fn completed_item_legacy_events_match(expected: &EventMsg, actual: &EventMsg) -> bool { match (expected, actual) { (EventMsg::AgentMessage(expected), EventMsg::AgentMessage(actual)) => { @@ -1717,14 +1822,16 @@ impl From<&PendingTurn> for Turn { } } +#[cfg(test)] +#[path = "thread_history_live_tests.rs"] +mod live_tests; + #[cfg(test)] mod tests { use super::*; use crate::protocol::v2::CommandExecutionSource; use codex_protocol::ThreadId; use codex_protocol::dynamic_tools::DynamicToolCallOutputContentItem as CoreDynamicToolCallOutputContentItem; - 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::SleepItem as CoreSleepItem; use codex_protocol::items::TurnItem as CoreTurnItem; @@ -1735,7 +1842,6 @@ mod tests { use codex_protocol::models::MessagePhase as CoreMessagePhase; use codex_protocol::models::WebSearchAction as CoreWebSearchAction; use codex_protocol::parse_command::ParsedCommand; - use codex_protocol::protocol::AgentMessageContentDeltaEvent; use codex_protocol::protocol::AgentMessageEvent; use codex_protocol::protocol::AgentReasoningEvent; use codex_protocol::protocol::AgentReasoningRawContentEvent; @@ -3351,111 +3457,6 @@ mod tests { ); } - #[test] - fn active_turn_snapshot_accumulates_agent_message_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::AgentMessage(CoreAgentMessageItem { - id: "msg-1".to_string(), - content: Vec::new(), - phase: None, - memory_citation: None, - }), - 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::AgentMessageContentDelta(AgentMessageContentDeltaEvent { - thread_id: thread_id.to_string(), - turn_id: turn_id.to_string(), - item_id: "msg-1".to_string(), - delta: "world".to_string(), - }), - ] { - builder.handle_event(&event); - } - - assert_eq!( - builder.active_turn_snapshot().expect("active turn").items, - vec![ThreadItem::AgentMessage { - id: "msg-1".to_string(), - text: "hello world".to_string(), - phase: None, - memory_citation: None, - }] - ); - } - #[test] - fn completed_agent_message_replaces_streamed_item_without_legacy_duplicate() { - let thread_id = ThreadId::new(); - let turn_id = "turn-1"; - let completed_item = CoreTurnItem::AgentMessage(CoreAgentMessageItem { - id: "msg-1".to_string(), - content: vec![CoreAgentMessageContent::Text { - text: "hello".to_string(), - }], - phase: None, - memory_citation: None, - }); - 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::AgentMessage(CoreAgentMessageItem { - id: "msg-1".to_string(), - content: Vec::new(), - phase: None, - memory_citation: None, - }), - started_at_ms: 0, - }), - EventMsg::ItemCompleted(ItemCompletedEvent { - thread_id, - turn_id: turn_id.to_string(), - item: completed_item.clone(), - completed_at_ms: 0, - }), - completed_item - .as_legacy_events(/*show_raw_agent_reasoning*/ true) - .into_iter() - .next() - .expect("agent message should produce one legacy event"), - ] { - builder.handle_event(&event); - } - - assert_eq!( - builder.active_turn_snapshot().expect("active turn").items, - vec![ThreadItem::AgentMessage { - id: "msg-1".to_string(), - text: "hello".to_string(), - phase: None, - memory_citation: None, - }] - ); - } - fn turn_started(turn_id: &str) -> EventMsg { - EventMsg::TurnStarted(TurnStartedEvent { - turn_id: turn_id.to_string(), - trace_id: None, - started_at: None, - model_context_window: None, - collaboration_mode_kind: Default::default(), - }) - } - #[test] fn late_turn_complete_does_not_close_active_turn() { let events = vec![ diff --git a/codex-rs/app-server-protocol/src/protocol/thread_history_live_tests.rs b/codex-rs/app-server-protocol/src/protocol/thread_history_live_tests.rs new file mode 100644 index 0000000000..8b5b2a8010 --- /dev/null +++ b/codex-rs/app-server-protocol/src/protocol/thread_history_live_tests.rs @@ -0,0 +1,958 @@ +use super::*; +use crate::protocol::v2::CommandAction; +use crate::protocol::v2::CommandExecutionSource; +use crate::protocol::v2::CommandExecutionStatus; +use codex_protocol::ThreadId; +use codex_protocol::items::AgentMessageContent as CoreAgentMessageContent; +use codex_protocol::items::AgentMessageItem as CoreAgentMessageItem; +use codex_protocol::items::PlanItem as CorePlanItem; +use codex_protocol::items::ReasoningItem as CoreReasoningItem; +use codex_protocol::items::TurnItem as CoreTurnItem; +use codex_protocol::models::MessagePhase; +use codex_protocol::parse_command::ParsedCommand; +use codex_protocol::protocol::AgentMessageContentDeltaEvent; +use codex_protocol::protocol::AgentMessageEvent; +use codex_protocol::protocol::ExecCommandBeginEvent; +use codex_protocol::protocol::ExecCommandEndEvent; +use codex_protocol::protocol::ExecCommandOutputDeltaEvent; +use codex_protocol::protocol::ExecCommandSource; +use codex_protocol::protocol::ExecCommandStatus as CoreExecCommandStatus; +use codex_protocol::protocol::ExecOutputStream; +use codex_protocol::protocol::ItemCompletedEvent; +use codex_protocol::protocol::ItemStartedEvent; +use codex_protocol::protocol::PlanDeltaEvent; +use codex_protocol::protocol::ReasoningContentDeltaEvent; +use codex_protocol::protocol::ReasoningRawContentDeltaEvent; +use codex_protocol::protocol::TokenCountEvent; +use codex_protocol::protocol::TurnCompleteEvent; +use codex_protocol::protocol::TurnStartedEvent; +use codex_utils_absolute_path::test_support::PathBufExt; +use codex_utils_absolute_path::test_support::test_path_buf; +use pretty_assertions::assert_eq; +use std::time::Duration; + +const TURN_A: &str = "turn-a"; +const TURN_B: &str = "turn-b"; +const AGENT_ID: &str = "agent-1"; +const PLAN_ID: &str = "plan-1"; +const REASONING_ID: &str = "reasoning-1"; +const COMMAND_ID: &str = "command-1"; + +#[test] +fn live_agent_message_state_is_cumulative() { + let mut builder = started_builder(TURN_A); + let started_item = core_agent(AGENT_ID, &[], Some(MessagePhase::Commentary)); + let mut expected = vec![agent_item(AGENT_ID, "", Some(MessagePhase::Commentary))]; + + apply_single_item(&mut builder, item_started(TURN_A, started_item), &expected); + expected[0] = agent_item(AGENT_ID, "hello ", Some(MessagePhase::Commentary)); + apply_single_item( + &mut builder, + agent_delta(TURN_A, AGENT_ID, "hello "), + &expected, + ); + apply_no_change(&mut builder, EventMsg::TokenCount(empty_token_count())); + expected[0] = agent_item(AGENT_ID, "hello world", Some(MessagePhase::Commentary)); + apply_single_item( + &mut builder, + agent_delta(TURN_A, AGENT_ID, "world"), + &expected, + ); + + let completed_item = core_agent( + AGENT_ID, + &["canonical ", "answer"], + Some(MessagePhase::FinalAnswer), + ); + expected[0] = agent_item( + AGENT_ID, + "canonical answer", + Some(MessagePhase::FinalAnswer), + ); + apply_single_item( + &mut builder, + item_completed(TURN_A, completed_item.clone()), + &expected, + ); + + for legacy_event in completed_item.as_legacy_events(/*show_raw_agent_reasoning*/ true) { + apply_no_change(&mut builder, legacy_event); + } +} + +#[test] +fn live_plan_state_is_cumulative_and_completion_is_authoritative() { + let mut builder = started_builder(TURN_A); + let mut expected = vec![plan_item(PLAN_ID, "")]; + + apply_single_item( + &mut builder, + item_started(TURN_A, core_plan(PLAN_ID, "")), + &expected, + ); + expected[0] = plan_item(PLAN_ID, "draft "); + apply_single_item( + &mut builder, + plan_delta(TURN_A, PLAN_ID, "draft "), + &expected, + ); + expected[0] = plan_item(PLAN_ID, "draft plan"); + apply_single_item(&mut builder, plan_delta(TURN_A, PLAN_ID, "plan"), &expected); + + expected[0] = plan_item(PLAN_ID, "canonical plan"); + apply_single_item( + &mut builder, + item_completed(TURN_A, core_plan(PLAN_ID, "canonical plan")), + &expected, + ); + apply_no_change(&mut builder, item_completed(TURN_A, core_plan(PLAN_ID, ""))); +} + +#[test] +fn live_reasoning_state_is_cumulative_by_index() { + let mut builder = started_builder(TURN_A); + let mut expected = vec![reasoning_item(REASONING_ID, &[], &[])]; + + apply_single_item( + &mut builder, + item_started(TURN_A, core_reasoning(REASONING_ID, &[], &[])), + &expected, + ); + expected[0] = reasoning_item(REASONING_ID, &["think"], &[]); + apply_single_item( + &mut builder, + reasoning_summary_delta(TURN_A, REASONING_ID, 0, "think"), + &expected, + ); + expected[0] = reasoning_item(REASONING_ID, &["think more"], &[]); + apply_single_item( + &mut builder, + reasoning_summary_delta(TURN_A, REASONING_ID, 0, " more"), + &expected, + ); + expected[0] = reasoning_item(REASONING_ID, &["think more", "next"], &[]); + apply_single_item( + &mut builder, + reasoning_summary_delta(TURN_A, REASONING_ID, 1, "next"), + &expected, + ); + expected[0] = reasoning_item(REASONING_ID, &["think more", "next"], &["raw"]); + apply_single_item( + &mut builder, + reasoning_raw_delta(TURN_A, REASONING_ID, 0, "raw"), + &expected, + ); + expected[0] = reasoning_item(REASONING_ID, &["think more", "next"], &["raw detail"]); + apply_single_item( + &mut builder, + reasoning_raw_delta(TURN_A, REASONING_ID, 0, " detail"), + &expected, + ); + + let completed_item = core_reasoning( + REASONING_ID, + &["canonical summary", "second summary"], + &["canonical raw", "second raw"], + ); + expected[0] = reasoning_item( + REASONING_ID, + &["canonical summary", "second summary"], + &["canonical raw", "second raw"], + ); + apply_single_item( + &mut builder, + item_completed(TURN_A, completed_item.clone()), + &expected, + ); + + let mut legacy_events = completed_item.as_legacy_events(/*show_raw_agent_reasoning*/ true); + legacy_events.reverse(); + for legacy_event in legacy_events { + apply_no_change(&mut builder, legacy_event); + } +} + +#[test] +fn live_command_state_is_cumulative_by_chunk() { + let mut builder = started_builder(TURN_A); + let mut expected = vec![running_command(/*aggregated_output*/ None)]; + + apply_single_item(&mut builder, command_begin(TURN_A), &expected); + expected[0] = running_command(Some("out ")); + apply_single_item( + &mut builder, + command_delta(ExecOutputStream::Stdout, b"out "), + &expected, + ); + expected[0] = running_command(Some("out warning ")); + apply_single_item( + &mut builder, + command_delta(ExecOutputStream::Stderr, b"warning "), + &expected, + ); + expected[0] = running_command(Some("out warning done")); + apply_single_item( + &mut builder, + command_delta(ExecOutputStream::Stdout, b"done"), + &expected, + ); + + expected[0] = completed_command("canonical output"); + apply_single_item( + &mut builder, + command_end(TURN_A, "canonical output"), + &expected, + ); +} + +#[test] +fn mismatched_live_events_are_noops() { + let mut builder = started_builder(TURN_A); + let expected = vec![ + agent_item(AGENT_ID, "", None), + plan_item(PLAN_ID, ""), + reasoning_item(REASONING_ID, &[], &[]), + running_command(/*aggregated_output*/ None), + ]; + for event in [ + item_started(TURN_A, core_agent(AGENT_ID, &[], None)), + item_started(TURN_A, core_plan(PLAN_ID, "")), + item_started(TURN_A, core_reasoning(REASONING_ID, &[], &[])), + command_begin(TURN_A), + ] { + builder.handle_event(&event); + } + assert_active_turn(&builder, TURN_A, expected.clone(), TurnStatus::InProgress); + + for event in [ + agent_delta(TURN_B, AGENT_ID, "wrong turn"), + plan_delta(TURN_A, "unknown-plan", "wrong item"), + reasoning_summary_delta(TURN_A, AGENT_ID, 0, "wrong type"), + reasoning_raw_delta(TURN_A, REASONING_ID, -1, "negative index"), + EventMsg::ExecCommandOutputDelta(ExecCommandOutputDeltaEvent { + call_id: "unknown-command".to_string(), + stream: ExecOutputStream::Stdout, + chunk: b"wrong item".to_vec(), + }), + ] { + assert_eq!( + apply_no_change_with_expected(&mut builder, event, expected.clone()), + ThreadHistoryChangeSet::default() + ); + } +} + +#[test] +fn nonmatching_agent_legacy_event_is_preserved() { + let mut builder = started_builder(TURN_A); + let completed_item = core_agent(AGENT_ID, &["same text"], Some(MessagePhase::FinalAnswer)); + let completed = agent_item(AGENT_ID, "same text", Some(MessagePhase::FinalAnswer)); + apply_single_item( + &mut builder, + item_completed(TURN_A, completed_item.clone()), + std::slice::from_ref(&completed), + ); + + let commentary = agent_item("item-1", "same text", Some(MessagePhase::Commentary)); + assert_eq!( + apply_with_expected( + &mut builder, + EventMsg::AgentMessage(AgentMessageEvent { + message: "same text".to_string(), + phase: Some(MessagePhase::Commentary), + memory_citation: None, + }), + vec![completed.clone(), commentary.clone()], + ), + single_item_change(TURN_A, commentary.clone()) + ); + assert_eq!( + apply_no_change_with_expected( + &mut builder, + completed_item + .as_legacy_events(/*show_raw_agent_reasoning*/ true) + .into_iter() + .next() + .expect("agent completion should generate one legacy event"), + vec![completed, commentary], + ), + ThreadHistoryChangeSet::default() + ); +} + +#[test] +fn legacy_expectations_clear_at_turn_boundary() { + let mut builder = started_builder(TURN_A); + let completed_item = core_agent(AGENT_ID, &["repeated text"], None); + builder.handle_event(&item_completed(TURN_A, completed_item)); + builder.handle_event(&turn_completed(TURN_A)); + builder.handle_event(&turn_started(TURN_B)); + + let expected = vec![agent_item("item-1", "repeated text", None)]; + assert_eq!( + apply_with_expected( + &mut builder, + EventMsg::AgentMessage(AgentMessageEvent { + message: "repeated text".to_string(), + phase: None, + memory_citation: None, + }), + expected.clone(), + ), + single_item_change(TURN_B, expected[0].clone()) + ); +} + +#[test] +fn interleaved_completions_consume_only_their_legacy_events() { + let mut builder = started_builder(TURN_A); + let first = core_agent("agent-1", &["first"], None); + let second = core_agent("agent-2", &["second"], None); + let expected = vec![ + agent_item("agent-1", "first", None), + agent_item("agent-2", "second", None), + ]; + builder.handle_event(&item_completed(TURN_A, first.clone())); + builder.handle_event(&item_completed(TURN_A, second.clone())); + assert_active_turn(&builder, TURN_A, expected.clone(), TurnStatus::InProgress); + + let second_legacy = second + .as_legacy_events(/*show_raw_agent_reasoning*/ true) + .into_iter() + .next() + .expect("agent completion should generate one legacy event"); + assert_eq!( + apply_no_change_with_expected(&mut builder, second_legacy, expected.clone()), + ThreadHistoryChangeSet::default() + ); + assert_eq!( + apply_no_change_with_expected( + &mut builder, + EventMsg::TokenCount(empty_token_count()), + expected.clone(), + ), + ThreadHistoryChangeSet::default() + ); + let first_legacy = first + .as_legacy_events(/*show_raw_agent_reasoning*/ true) + .into_iter() + .next() + .expect("agent completion should generate one legacy event"); + assert_eq!( + apply_no_change_with_expected(&mut builder, first_legacy, expected), + ThreadHistoryChangeSet::default() + ); +} + +#[test] +fn late_completion_updates_only_its_original_turn() { + let mut builder = started_builder(TURN_A); + builder.handle_event(&item_started( + TURN_A, + core_agent(AGENT_ID, &[], Some(MessagePhase::Commentary)), + )); + builder.handle_event(&turn_completed(TURN_A)); + builder.handle_event(&turn_started(TURN_B)); + + let completed_item = core_agent(AGENT_ID, &["final"], Some(MessagePhase::FinalAnswer)); + let completed = agent_item(AGENT_ID, "final", Some(MessagePhase::FinalAnswer)); + assert_eq!( + builder.handle_event_with_changes(&item_completed(TURN_A, completed_item.clone())), + single_item_change(TURN_A, completed.clone()) + ); + assert_turn(&builder, TURN_A, vec![completed], TurnStatus::Completed); + assert_active_turn(&builder, TURN_B, Vec::new(), TurnStatus::InProgress); + + for legacy_event in completed_item.as_legacy_events(/*show_raw_agent_reasoning*/ true) { + assert_eq!( + builder.handle_event_with_changes(&legacy_event), + ThreadHistoryChangeSet::default() + ); + assert_active_turn(&builder, TURN_B, Vec::new(), TurnStatus::InProgress); + } +} + +#[test] +fn batched_live_changes_coalesce_to_authoritative_completion() { + let mut builder = started_builder(TURN_A); + let completed_item = core_agent(AGENT_ID, &["canonical"], None); + let mut rollout_items = vec![ + RolloutItem::EventMsg(item_started(TURN_A, core_agent(AGENT_ID, &[], None))), + RolloutItem::EventMsg(agent_delta(TURN_A, AGENT_ID, "partial")), + RolloutItem::EventMsg(item_completed(TURN_A, completed_item.clone())), + ]; + rollout_items.extend( + completed_item + .as_legacy_events(/*show_raw_agent_reasoning*/ true) + .into_iter() + .map(RolloutItem::EventMsg), + ); + let completed = agent_item(AGENT_ID, "canonical", None); + + assert_eq!( + builder.handle_rollout_items_with_changes(&rollout_items), + single_item_change(TURN_A, completed.clone()) + ); + assert_active_turn(&builder, TURN_A, vec![completed], TurnStatus::InProgress); +} + +#[test] +fn completed_agent_message_ignores_late_delta() { + let mut builder = started_builder(TURN_A); + let completed = agent_item(AGENT_ID, "final", None); + builder.handle_event(&item_completed( + TURN_A, + core_agent(AGENT_ID, &["final"], None), + )); + + assert_eq!( + apply_no_change(&mut builder, agent_delta(TURN_A, AGENT_ID, " late")), + ThreadHistoryChangeSet::default() + ); + assert_active_turn(&builder, TURN_A, vec![completed], TurnStatus::InProgress); +} + +#[test] +fn completed_plan_ignores_late_delta() { + let mut builder = started_builder(TURN_A); + let completed = plan_item(PLAN_ID, "final"); + builder.handle_event(&item_completed(TURN_A, core_plan(PLAN_ID, "final"))); + + assert_eq!( + apply_no_change(&mut builder, plan_delta(TURN_A, PLAN_ID, " late")), + ThreadHistoryChangeSet::default() + ); + assert_active_turn(&builder, TURN_A, vec![completed], TurnStatus::InProgress); +} + +#[test] +fn completed_reasoning_ignores_late_deltas() { + let mut builder = started_builder(TURN_A); + let completed = reasoning_item(REASONING_ID, &["final summary"], &["final raw"]); + builder.handle_event(&item_completed( + TURN_A, + core_reasoning(REASONING_ID, &["final summary"], &["final raw"]), + )); + + assert_eq!( + apply_no_change( + &mut builder, + reasoning_summary_delta(TURN_A, REASONING_ID, 0, " late") + ), + ThreadHistoryChangeSet::default() + ); + assert_eq!( + apply_no_change( + &mut builder, + reasoning_raw_delta(TURN_A, REASONING_ID, 0, " late") + ), + ThreadHistoryChangeSet::default() + ); + assert_active_turn(&builder, TURN_A, vec![completed], TurnStatus::InProgress); +} + +#[test] +fn completed_command_ignores_late_begin_and_output() { + let mut builder = started_builder(TURN_A); + builder.handle_event(&command_end(TURN_A, "final")); + + apply_no_change(&mut builder, command_begin(TURN_A)); + assert_eq!( + apply_no_change( + &mut builder, + command_delta(ExecOutputStream::Stdout, b" late") + ), + ThreadHistoryChangeSet::default() + ); + assert_active_turn( + &builder, + TURN_A, + vec![completed_command("final")], + TurnStatus::InProgress, + ); +} + +#[test] +fn completed_command_ignores_duplicate_end() { + let mut builder = started_builder(TURN_A); + builder.handle_event(&command_end(TURN_A, "final")); + + assert_eq!( + apply_no_change(&mut builder, command_end(TURN_A, "stale duplicate")), + ThreadHistoryChangeSet::default() + ); + assert_active_turn( + &builder, + TURN_A, + vec![completed_command("final")], + TurnStatus::InProgress, + ); +} + +#[test] +fn completed_item_ignores_late_start() { + let mut builder = started_builder(TURN_A); + let completed = agent_item(AGENT_ID, "final", None); + builder.handle_event(&item_completed( + TURN_A, + core_agent(AGENT_ID, &["final"], None), + )); + + assert_eq!( + apply_no_change( + &mut builder, + item_started(TURN_A, core_agent(AGENT_ID, &[], None)) + ), + ThreadHistoryChangeSet::default() + ); + assert_active_turn(&builder, TURN_A, vec![completed], TurnStatus::InProgress); +} + +#[test] +fn command_output_preserves_utf8_across_every_chunk_boundary() { + let text = "a😀éz"; + let bytes = text.as_bytes(); + + for first_split in 0..=bytes.len() { + for second_split in first_split..=bytes.len() { + let mut builder = started_builder(TURN_A); + builder.handle_event(&command_begin(TURN_A)); + for chunk in [ + &bytes[..first_split], + &bytes[first_split..second_split], + &bytes[second_split..], + ] { + builder.handle_event(&command_delta(ExecOutputStream::Stdout, chunk)); + } + assert_active_turn( + &builder, + TURN_A, + vec![running_command(Some(text))], + TurnStatus::InProgress, + ); + } + } +} + +#[test] +fn command_output_replaces_invalid_utf8_without_dropping_surrounding_bytes() { + for (chunks, expected) in [ + (vec![b"a".to_vec(), vec![0xFF], b"b".to_vec()], "a�b"), + (vec![vec![0xF0], vec![0x28, 0x8C], vec![0x28]], "�(�("), + ] { + let mut builder = started_builder(TURN_A); + builder.handle_event(&command_begin(TURN_A)); + for chunk in chunks { + builder.handle_event(&command_delta(ExecOutputStream::Stdout, &chunk)); + } + assert_active_turn( + &builder, + TURN_A, + vec![running_command(Some(expected))], + TurnStatus::InProgress, + ); + } +} + +#[test] +fn completed_item_tracking_is_scoped_to_its_turn() { + let mut builder = started_builder(TURN_A); + builder.handle_event(&item_completed( + TURN_A, + core_agent(AGENT_ID, &["first"], None), + )); + builder.handle_event(&turn_completed(TURN_A)); + builder.handle_event(&turn_started(TURN_B)); + + let second = agent_item(AGENT_ID, "second", None); + builder.handle_event(&item_started(TURN_B, core_agent(AGENT_ID, &[], None))); + assert_eq!( + builder.handle_event_with_changes(&agent_delta(TURN_B, AGENT_ID, "second")), + single_item_change(TURN_B, second.clone()) + ); + assert_turn( + &builder, + TURN_A, + vec![agent_item(AGENT_ID, "first", None)], + TurnStatus::Completed, + ); + assert_active_turn(&builder, TURN_B, vec![second], TurnStatus::InProgress); +} + +#[test] +fn late_command_output_updates_original_turn() { + let mut builder = started_builder(TURN_A); + builder.handle_event(&command_begin(TURN_A)); + builder.handle_event(&turn_completed(TURN_A)); + builder.handle_event(&turn_started(TURN_B)); + + assert_eq!( + builder.handle_event_with_changes(&command_delta(ExecOutputStream::Stdout, b"late")), + single_item_change(TURN_A, running_command(Some("late"))) + ); + assert_turn( + &builder, + TURN_A, + vec![running_command(Some("late"))], + TurnStatus::Completed, + ); + assert_active_turn(&builder, TURN_B, Vec::new(), TurnStatus::InProgress); +} + +#[test] +fn reasoning_delta_rejects_sparse_index() { + let mut builder = started_builder(TURN_A); + let expected = vec![reasoning_item(REASONING_ID, &[], &[])]; + builder.handle_event(&item_started( + TURN_A, + core_reasoning(REASONING_ID, &[], &[]), + )); + + assert_eq!( + apply_no_change( + &mut builder, + reasoning_summary_delta(TURN_A, REASONING_ID, 2, "sparse") + ), + ThreadHistoryChangeSet::default() + ); + assert_active_turn(&builder, TURN_A, expected, TurnStatus::InProgress); +} + +fn started_builder(turn_id: &str) -> ThreadHistoryBuilder { + let mut builder = ThreadHistoryBuilder::new(); + assert_eq!( + builder.handle_event_with_changes(&turn_started(turn_id)), + ThreadHistoryChangeSet { + changed_items: Vec::new(), + changed_turns: vec![ThreadHistoryTurnChange { + turn_id: turn_id.to_string(), + status: TurnStatus::InProgress, + error: None, + started_at: None, + completed_at: None, + duration_ms: None, + }], + removed_turn_ids: Vec::new(), + } + ); + assert_active_turn(&builder, turn_id, Vec::new(), TurnStatus::InProgress); + builder +} + +fn apply_single_item( + builder: &mut ThreadHistoryBuilder, + event: EventMsg, + expected_items: &[ThreadItem], +) { + let expected_item = expected_items + .first() + .expect("event should leave one materialized item"); + assert_eq!(expected_items.len(), 1); + let turn_id = builder.active_turn_snapshot().expect("active turn").id; + assert_eq!( + builder.handle_event_with_changes(&event), + single_item_change(&turn_id, expected_item.clone()) + ); + assert_active_turn( + builder, + &turn_id, + expected_items.to_vec(), + TurnStatus::InProgress, + ); +} + +fn apply_with_expected( + builder: &mut ThreadHistoryBuilder, + event: EventMsg, + expected_items: Vec, +) -> ThreadHistoryChangeSet { + let changes = builder.handle_event_with_changes(&event); + let turn_id = builder.active_turn_snapshot().expect("active turn").id; + assert_active_turn(builder, &turn_id, expected_items, TurnStatus::InProgress); + changes +} + +fn apply_no_change(builder: &mut ThreadHistoryBuilder, event: EventMsg) -> ThreadHistoryChangeSet { + let expected_items = builder.active_turn_snapshot().expect("active turn").items; + apply_no_change_with_expected(builder, event, expected_items) +} + +fn apply_no_change_with_expected( + builder: &mut ThreadHistoryBuilder, + event: EventMsg, + expected_items: Vec, +) -> ThreadHistoryChangeSet { + let changes = apply_with_expected(builder, event, expected_items); + assert_eq!(changes, ThreadHistoryChangeSet::default()); + changes +} + +fn assert_active_turn( + builder: &ThreadHistoryBuilder, + turn_id: &str, + items: Vec, + status: TurnStatus, +) { + assert_eq!( + builder.active_turn_snapshot(), + Some(turn(turn_id, items, status)) + ); +} + +fn assert_turn( + builder: &ThreadHistoryBuilder, + turn_id: &str, + items: Vec, + status: TurnStatus, +) { + assert_eq!( + builder.turn_snapshot(turn_id), + Some(turn(turn_id, items, status)) + ); +} + +fn turn(turn_id: &str, items: Vec, status: TurnStatus) -> Turn { + Turn { + id: turn_id.to_string(), + items, + items_view: TurnItemsView::Full, + status, + error: None, + started_at: None, + completed_at: None, + duration_ms: None, + } +} + +fn single_item_change(turn_id: &str, item: ThreadItem) -> ThreadHistoryChangeSet { + ThreadHistoryChangeSet { + changed_items: vec![ThreadHistoryItemChange { + turn_id: turn_id.to_string(), + item, + }], + changed_turns: Vec::new(), + removed_turn_ids: Vec::new(), + } +} + +fn turn_started(turn_id: &str) -> EventMsg { + EventMsg::TurnStarted(TurnStartedEvent { + turn_id: turn_id.to_string(), + trace_id: None, + started_at: None, + model_context_window: None, + collaboration_mode_kind: Default::default(), + }) +} + +fn turn_completed(turn_id: &str) -> EventMsg { + EventMsg::TurnComplete(TurnCompleteEvent { + turn_id: turn_id.to_string(), + last_agent_message: None, + completed_at: None, + duration_ms: None, + time_to_first_token_ms: None, + }) +} + +fn item_started(turn_id: &str, item: CoreTurnItem) -> EventMsg { + EventMsg::ItemStarted(ItemStartedEvent { + thread_id: ThreadId::new(), + turn_id: turn_id.to_string(), + item, + started_at_ms: 0, + }) +} + +fn item_completed(turn_id: &str, item: CoreTurnItem) -> EventMsg { + EventMsg::ItemCompleted(ItemCompletedEvent { + thread_id: ThreadId::new(), + turn_id: turn_id.to_string(), + item, + completed_at_ms: 0, + }) +} + +fn core_agent(id: &str, fragments: &[&str], phase: Option) -> CoreTurnItem { + CoreTurnItem::AgentMessage(CoreAgentMessageItem { + id: id.to_string(), + content: fragments + .iter() + .map(|text| CoreAgentMessageContent::Text { + text: (*text).to_string(), + }) + .collect(), + phase, + memory_citation: None, + }) +} + +fn core_plan(id: &str, text: &str) -> CoreTurnItem { + CoreTurnItem::Plan(CorePlanItem { + id: id.to_string(), + text: text.to_string(), + }) +} + +fn core_reasoning(id: &str, summary: &[&str], content: &[&str]) -> CoreTurnItem { + CoreTurnItem::Reasoning(CoreReasoningItem { + id: id.to_string(), + summary_text: summary.iter().map(ToString::to_string).collect(), + raw_content: content.iter().map(ToString::to_string).collect(), + }) +} + +fn agent_delta(turn_id: &str, item_id: &str, delta: &str) -> EventMsg { + EventMsg::AgentMessageContentDelta(AgentMessageContentDeltaEvent { + thread_id: "thread-1".to_string(), + turn_id: turn_id.to_string(), + item_id: item_id.to_string(), + delta: delta.to_string(), + }) +} + +fn plan_delta(turn_id: &str, item_id: &str, delta: &str) -> EventMsg { + EventMsg::PlanDelta(PlanDeltaEvent { + thread_id: "thread-1".to_string(), + turn_id: turn_id.to_string(), + item_id: item_id.to_string(), + delta: delta.to_string(), + }) +} + +fn reasoning_summary_delta( + turn_id: &str, + item_id: &str, + summary_index: i64, + delta: &str, +) -> EventMsg { + EventMsg::ReasoningContentDelta(ReasoningContentDeltaEvent { + thread_id: "thread-1".to_string(), + turn_id: turn_id.to_string(), + item_id: item_id.to_string(), + delta: delta.to_string(), + summary_index, + }) +} + +fn reasoning_raw_delta(turn_id: &str, item_id: &str, content_index: i64, delta: &str) -> EventMsg { + EventMsg::ReasoningRawContentDelta(ReasoningRawContentDeltaEvent { + thread_id: "thread-1".to_string(), + turn_id: turn_id.to_string(), + item_id: item_id.to_string(), + delta: delta.to_string(), + content_index, + }) +} + +fn command_begin(turn_id: &str) -> EventMsg { + EventMsg::ExecCommandBegin(ExecCommandBeginEvent { + call_id: COMMAND_ID.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, + }) +} + +fn command_delta(stream: ExecOutputStream, chunk: &[u8]) -> EventMsg { + EventMsg::ExecCommandOutputDelta(ExecCommandOutputDeltaEvent { + call_id: COMMAND_ID.to_string(), + stream, + chunk: chunk.to_vec(), + }) +} + +fn command_end(turn_id: &str, aggregated_output: &str) -> EventMsg { + EventMsg::ExecCommandEnd(ExecCommandEndEvent { + call_id: COMMAND_ID.to_string(), + process_id: Some("process-1".to_string()), + turn_id: turn_id.to_string(), + completed_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, + stdout: aggregated_output.to_string(), + stderr: String::new(), + aggregated_output: aggregated_output.to_string(), + exit_code: 0, + duration: Duration::from_millis(5), + formatted_output: aggregated_output.to_string(), + status: CoreExecCommandStatus::Completed, + }) +} + +fn agent_item(id: &str, text: &str, phase: Option) -> ThreadItem { + ThreadItem::AgentMessage { + id: id.to_string(), + text: text.to_string(), + phase, + memory_citation: None, + } +} + +fn plan_item(id: &str, text: &str) -> ThreadItem { + ThreadItem::Plan { + id: id.to_string(), + text: text.to_string(), + } +} + +fn reasoning_item(id: &str, summary: &[&str], content: &[&str]) -> ThreadItem { + ThreadItem::Reasoning { + id: id.to_string(), + summary: summary.iter().map(ToString::to_string).collect(), + content: content.iter().map(ToString::to_string).collect(), + } +} + +fn running_command(aggregated_output: Option<&str>) -> ThreadItem { + ThreadItem::CommandExecution { + id: COMMAND_ID.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: aggregated_output.map(ToString::to_string), + exit_code: None, + duration_ms: None, + } +} + +fn completed_command(aggregated_output: &str) -> ThreadItem { + ThreadItem::CommandExecution { + id: COMMAND_ID.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::Completed, + command_actions: vec![CommandAction::Unknown { + command: "echo hello".to_string(), + }], + aggregated_output: Some(aggregated_output.to_string()), + exit_code: Some(0), + duration_ms: Some(5), + } +} + +fn empty_token_count() -> TokenCountEvent { + TokenCountEvent { + info: None, + rate_limits: None, + } +} 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 4e288a7fb0..9788d934e4 100644 --- a/codex-rs/app-server/src/request_processors/thread_lifecycle.rs +++ b/codex-rs/app-server/src/request_processors/thread_lifecycle.rs @@ -819,13 +819,31 @@ pub(super) fn set_thread_status_and_interrupt_stale_turns( loaded_status: ThreadStatus, has_live_in_progress_turn: bool, ) { - let status = resolve_thread_status(loaded_status, has_live_in_progress_turn); - if !matches!(status, ThreadStatus::Active { .. }) { - for turn in &mut thread.turns { - if matches!(turn.status, TurnStatus::InProgress) { - turn.status = TurnStatus::Interrupted; - } - } + let mut status = resolve_thread_status(loaded_status, has_live_in_progress_turn); + if !has_live_in_progress_turn + && matches!(&status, ThreadStatus::Active { active_flags } if active_flags.is_empty()) + { + status = ThreadStatus::Idle; } + + interrupt_stale_in_progress_turns( + &mut thread.turns, + matches!(status, ThreadStatus::Active { .. }), + ); thread.status = status; } + +pub(super) fn interrupt_stale_in_progress_turns(turns: &mut [Turn], preserve_latest: bool) { + let current_turn_index = preserve_latest + .then(|| { + turns + .iter() + .rposition(|turn| matches!(turn.status, TurnStatus::InProgress)) + }) + .flatten(); + for (index, turn) in turns.iter_mut().enumerate() { + if matches!(turn.status, TurnStatus::InProgress) && Some(index) != current_turn_index { + turn.status = TurnStatus::Interrupted; + } + } +} 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 8d5950db07..f886927137 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -3990,15 +3990,16 @@ fn reconstruct_thread_turns_for_turns_list( has_live_running_thread: bool, active_turn: 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 turns = build_api_turns_from_rollout_items(items); + let active_turn_is_current = active_turn.as_ref().is_none_or(|active_turn| { + merge_turn_history_with_active_turn(&mut turns, active_turn.clone()) + }); + let has_live_in_progress_turn = active_turn_is_current + && (has_live_running_thread + || active_turn + .as_ref() + .is_some_and(|turn| matches!(turn.status, TurnStatus::InProgress))); 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); - } turns } @@ -4008,14 +4009,7 @@ fn normalize_thread_turns_status( has_live_in_progress_turn: bool, ) { let status = resolve_thread_status(loaded_status, has_live_in_progress_turn); - if matches!(status, ThreadStatus::Active { .. }) { - return; - } - for turn in turns { - if matches!(turn.status, TurnStatus::InProgress) { - turn.status = TurnStatus::Interrupted; - } - } + interrupt_stale_in_progress_turns(turns, matches!(status, ThreadStatus::Active { .. })); } enum ThreadReadViewError { @@ -4489,3 +4483,7 @@ fn build_thread_from_loaded_snapshot( #[cfg(test)] #[path = "thread_processor_tests.rs"] mod thread_processor_tests; + +#[cfg(test)] +#[path = "thread_resume_reconciliation_tests.rs"] +mod thread_resume_reconciliation_tests; diff --git a/codex-rs/app-server/src/request_processors/thread_resume_reconciliation_tests.rs b/codex-rs/app-server/src/request_processors/thread_resume_reconciliation_tests.rs new file mode 100644 index 0000000000..aa9c45f7a5 --- /dev/null +++ b/codex-rs/app-server/src/request_processors/thread_resume_reconciliation_tests.rs @@ -0,0 +1,582 @@ +use super::super::thread_lifecycle::merge_turn_history_with_active_turn; +use super::super::thread_lifecycle::set_thread_status_and_interrupt_stale_turns; +use super::build_thread_resume_initial_turns_page; +use codex_app_server_protocol::SessionSource; +use codex_app_server_protocol::SortDirection; +use codex_app_server_protocol::Thread; +use codex_app_server_protocol::ThreadActiveFlag; +use codex_app_server_protocol::ThreadItem; +use codex_app_server_protocol::ThreadResumeInitialTurnsPageParams; +use codex_app_server_protocol::ThreadStatus; +use codex_app_server_protocol::Turn; +use codex_app_server_protocol::TurnItemsView; +use codex_app_server_protocol::TurnStatus; +use codex_app_server_protocol::UserInput; +use codex_protocol::protocol::AgentMessageEvent; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::RolloutItem; +use codex_protocol::protocol::TurnCompleteEvent; +use codex_protocol::protocol::TurnStartedEvent; +use codex_protocol::protocol::UserMessageEvent; +use codex_utils_absolute_path::AbsolutePathBuf; +use pretty_assertions::assert_eq; + +#[test] +fn persisted_terminal_statuses_reject_same_id_in_progress_snapshot() { + for persisted_status in [ + TurnStatus::Completed, + TurnStatus::Interrupted, + TurnStatus::Failed, + ] { + let older_turn = turn("older", TurnStatus::Completed, "older persisted item"); + let persisted_turn = turn("current", persisted_status, "authoritative persisted item"); + let stale_active_turn = turn("current", TurnStatus::InProgress, "stale active item"); + let mut turns = vec![older_turn.clone(), persisted_turn.clone()]; + + let accepted = merge_turn_history_with_active_turn(&mut turns, stale_active_turn); + + assert_eq!(accepted, false); + assert_eq!(turns, vec![older_turn, persisted_turn]); + } +} + +#[test] +fn current_snapshots_replace_or_append_without_duplicates() { + struct Case { + persisted_turns: Vec, + active_turn: Turn, + expected_turns: Vec, + } + + let older_turn = turn("older", TurnStatus::Completed, "older"); + let persisted_partial = turn("current", TurnStatus::InProgress, "persisted partial"); + let active_partial = turn("current", TurnStatus::InProgress, "active partial"); + let active_completed = turn("current", TurnStatus::Completed, "active final"); + let active_interrupted = turn("current", TurnStatus::Interrupted, "active interrupted"); + let active_failed = turn("current", TurnStatus::Failed, "active failed"); + let next_active = turn("next", TurnStatus::InProgress, "next active"); + let cases = [ + Case { + persisted_turns: Vec::new(), + active_turn: active_partial.clone(), + expected_turns: vec![active_partial.clone()], + }, + Case { + persisted_turns: vec![older_turn.clone(), persisted_partial.clone()], + active_turn: active_partial.clone(), + expected_turns: vec![older_turn.clone(), active_partial], + }, + Case { + persisted_turns: vec![older_turn.clone(), persisted_partial.clone()], + active_turn: active_completed.clone(), + expected_turns: vec![older_turn.clone(), active_completed], + }, + Case { + persisted_turns: vec![older_turn.clone(), persisted_partial.clone()], + active_turn: active_interrupted.clone(), + expected_turns: vec![older_turn.clone(), active_interrupted], + }, + Case { + persisted_turns: vec![older_turn.clone(), persisted_partial], + active_turn: active_failed.clone(), + expected_turns: vec![older_turn.clone(), active_failed], + }, + Case { + persisted_turns: vec![older_turn.clone()], + active_turn: next_active.clone(), + expected_turns: vec![older_turn, next_active], + }, + ]; + + for case in cases { + let mut turns = case.persisted_turns; + + let accepted = merge_turn_history_with_active_turn(&mut turns, case.active_turn); + + assert_eq!(accepted, true); + assert_eq!(turns, case.expected_turns); + assert_eq!( + turns + .iter() + .map(|turn| turn.id.as_str()) + .collect::>() + .len(), + turns.len(), + "reconciliation must leave at most one snapshot for each turn id" + ); + } +} + +#[test] +fn thread_status_normalization_matrix() { + struct Case { + loaded_status: ThreadStatus, + has_live_turn: bool, + expected_status: ThreadStatus, + expected_turn_status: TurnStatus, + } + + let active = ThreadStatus::Active { + active_flags: Vec::new(), + }; + let waiting = ThreadStatus::Active { + active_flags: vec![ThreadActiveFlag::WaitingOnApproval], + }; + let cases = [ + Case { + loaded_status: ThreadStatus::Idle, + has_live_turn: false, + expected_status: ThreadStatus::Idle, + expected_turn_status: TurnStatus::Interrupted, + }, + Case { + loaded_status: ThreadStatus::Idle, + has_live_turn: true, + expected_status: active.clone(), + expected_turn_status: TurnStatus::InProgress, + }, + Case { + loaded_status: ThreadStatus::NotLoaded, + has_live_turn: false, + expected_status: ThreadStatus::NotLoaded, + expected_turn_status: TurnStatus::Interrupted, + }, + Case { + loaded_status: ThreadStatus::NotLoaded, + has_live_turn: true, + expected_status: active, + expected_turn_status: TurnStatus::InProgress, + }, + Case { + loaded_status: waiting.clone(), + has_live_turn: false, + expected_status: waiting.clone(), + expected_turn_status: TurnStatus::InProgress, + }, + Case { + loaded_status: waiting.clone(), + has_live_turn: true, + expected_status: waiting, + expected_turn_status: TurnStatus::InProgress, + }, + Case { + loaded_status: ThreadStatus::SystemError, + has_live_turn: false, + expected_status: ThreadStatus::SystemError, + expected_turn_status: TurnStatus::Interrupted, + }, + Case { + loaded_status: ThreadStatus::SystemError, + has_live_turn: true, + expected_status: ThreadStatus::SystemError, + expected_turn_status: TurnStatus::Interrupted, + }, + ]; + + for case in cases { + let original_thread = thread(vec![turn("current", TurnStatus::InProgress, "active item")]); + let mut thread = original_thread.clone(); + + set_thread_status_and_interrupt_stale_turns( + &mut thread, + case.loaded_status, + case.has_live_turn, + ); + + let mut expected_thread = original_thread; + expected_thread.status = case.expected_status; + expected_thread.turns[0].status = case.expected_turn_status; + assert_eq!(thread, expected_thread); + } +} + +#[test] +fn persisted_terminal_turn_invalidates_stale_active_thread_status() { + let history_items = completed_rollout("current", "persisted user", "persisted answer"); + let stale_active_turn = turn_with_user_and_agent_items( + "current", + TurnStatus::InProgress, + "stale user", + "stale partial", + ); + + let (resumed_thread, accepted_active_turn, has_live_turn) = project_full_resume( + &history_items, + Some(stale_active_turn), + ThreadStatus::Active { + active_flags: Vec::new(), + }, + /*agent_running*/ true, + ); + + assert_eq!(accepted_active_turn, None); + assert_eq!(has_live_turn, false); + assert_eq!(resumed_thread.status, ThreadStatus::Idle); + assert_eq!(resumed_thread.turns[0].status, TurnStatus::Completed); +} + +#[test] +fn distinct_persisted_in_progress_turn_is_interrupted_when_new_turn_is_live() { + let persisted_turn = turn( + "incomplete-old-turn", + TurnStatus::InProgress, + "persisted partial", + ); + let active_turn = turn("new-live-turn", TurnStatus::InProgress, "live partial"); + let mut turns = vec![persisted_turn.clone()]; + assert!(merge_turn_history_with_active_turn( + &mut turns, + active_turn.clone() + )); + let mut resumed_thread = thread(turns); + + set_thread_status_and_interrupt_stale_turns( + &mut resumed_thread, + ThreadStatus::Idle, + /*has_live_in_progress_turn*/ true, + ); + + let mut expected_persisted_turn = persisted_turn; + expected_persisted_turn.status = TurnStatus::Interrupted; + let mut expected_thread = thread(vec![expected_persisted_turn, active_turn]); + expected_thread.status = ThreadStatus::Active { + active_flags: Vec::new(), + }; + assert_eq!(resumed_thread, expected_thread); + assert_eq!( + resumed_thread + .turns + .iter() + .filter(|turn| matches!(turn.status, TurnStatus::InProgress)) + .count(), + 1 + ); +} + +#[test] +fn full_turns_and_initial_page_apply_the_same_system_error_status() { + let active_turn = turn("current", TurnStatus::InProgress, "live partial"); + let mut resumed_thread = thread(vec![active_turn.clone()]); + set_thread_status_and_interrupt_stale_turns( + &mut resumed_thread, + ThreadStatus::SystemError, + /*has_live_in_progress_turn*/ true, + ); + + let initial_page = build_thread_resume_initial_turns_page( + &[], + resumed_thread.status.clone(), + /*has_live_running_thread*/ true, + Some(active_turn), + &full_ascending_page_params(), + ) + .expect("initial turns page should build"); + + assert_eq!(initial_page.data, resumed_thread.turns); +} + +#[test] +fn full_turns_and_initial_page_interrupt_the_same_stale_in_progress_turn() { + let history_items = incomplete_rollout("persisted", "persisted user", "persisted partial"); + let active_turn = turn_with_user_and_agent_items( + "active", + TurnStatus::InProgress, + "active user", + "active partial", + ); + let (resumed_thread, accepted_active_turn, has_live_turn) = project_full_resume( + &history_items, + Some(active_turn), + ThreadStatus::Idle, + /*agent_running*/ true, + ); + + let initial_page = build_thread_resume_initial_turns_page( + &history_items, + resumed_thread.status.clone(), + has_live_turn, + accepted_active_turn, + &full_ascending_page_params(), + ) + .expect("initial turns page should build"); + + assert_eq!(initial_page.data, resumed_thread.turns); + assert_eq!( + initial_page + .data + .iter() + .filter(|turn| matches!(turn.status, TurnStatus::InProgress)) + .count(), + 1 + ); +} + +#[test] +fn full_turns_and_initial_page_share_ids_order_status_and_full_items() { + let history_items = completed_rollout("persisted", "persisted user", "persisted answer"); + let active_turn = turn_with_user_and_agent_items( + "active", + TurnStatus::InProgress, + "active user", + "active partial", + ); + let (resumed_thread, accepted_active_turn, has_live_turn) = project_full_resume( + &history_items, + Some(active_turn), + ThreadStatus::Idle, + /*agent_running*/ true, + ); + + let initial_page = build_thread_resume_initial_turns_page( + &history_items, + resumed_thread.status.clone(), + has_live_turn, + accepted_active_turn, + &full_ascending_page_params(), + ) + .expect("initial turns page should build"); + + assert_eq!(initial_page.data, resumed_thread.turns); +} + +#[test] +fn initial_page_item_views_do_not_change_turn_identity_or_status() { + let active_turn = turn_with_user_and_agent_items( + "active", + TurnStatus::InProgress, + "active user", + "active partial", + ); + + for items_view in [ + TurnItemsView::NotLoaded, + TurnItemsView::Summary, + TurnItemsView::Full, + ] { + let page = build_thread_resume_initial_turns_page( + &[], + ThreadStatus::Active { + active_flags: Vec::new(), + }, + /*has_live_running_thread*/ true, + Some(active_turn.clone()), + &ThreadResumeInitialTurnsPageParams { + limit: None, + sort_direction: Some(SortDirection::Asc), + items_view: Some(items_view), + }, + ) + .expect("initial turns page should build"); + + let mut expected_turn = active_turn.clone(); + match items_view { + TurnItemsView::NotLoaded => expected_turn.items.clear(), + TurnItemsView::Summary | TurnItemsView::Full => {} + } + expected_turn.items_view = items_view; + assert_eq!(page.data, vec![expected_turn]); + } +} + +#[test] +fn stale_active_snapshot_is_excluded_from_both_full_turns_and_initial_page() { + let history_items = completed_rollout("current", "persisted user", "persisted answer"); + let stale_active_turn = turn_with_user_and_agent_items( + "current", + TurnStatus::InProgress, + "stale user", + "stale partial", + ); + let (resumed_thread, accepted_active_turn, has_live_turn) = project_full_resume( + &history_items, + Some(stale_active_turn), + ThreadStatus::Idle, + /*agent_running*/ true, + ); + assert_eq!(accepted_active_turn, None); + assert_eq!(has_live_turn, false); + + let initial_page = build_thread_resume_initial_turns_page( + &history_items, + resumed_thread.status.clone(), + has_live_turn, + accepted_active_turn, + &full_ascending_page_params(), + ) + .expect("initial turns page should build"); + + assert_eq!(initial_page.data, resumed_thread.turns); + assert_eq!(resumed_thread.status, ThreadStatus::Idle); + assert_eq!(resumed_thread.turns[0].status, TurnStatus::Completed); +} + +fn project_full_resume( + history_items: &[RolloutItem], + active_turn: Option, + loaded_status: ThreadStatus, + agent_running: bool, +) -> (Thread, Option, bool) { + let mut turns = codex_app_server_protocol::build_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 turns, active_turn.clone()) + }); + let active_turn = active_turn.filter(|_| active_turn_is_current); + let has_live_turn = active_turn_is_current + && (agent_running + || active_turn + .as_ref() + .is_some_and(|turn| matches!(turn.status, TurnStatus::InProgress))); + let mut thread = thread(turns); + set_thread_status_and_interrupt_stale_turns(&mut thread, loaded_status, has_live_turn); + (thread, active_turn, has_live_turn) +} + +fn full_ascending_page_params() -> ThreadResumeInitialTurnsPageParams { + ThreadResumeInitialTurnsPageParams { + limit: None, + sort_direction: Some(SortDirection::Asc), + items_view: Some(TurnItemsView::Full), + } +} + +fn turn(id: &str, status: TurnStatus, item_text: &str) -> Turn { + Turn { + id: id.to_string(), + items: vec![ThreadItem::AgentMessage { + id: format!("{id}-item"), + text: item_text.to_string(), + phase: None, + memory_citation: None, + }], + items_view: TurnItemsView::Full, + error: None, + started_at: Some(1), + completed_at: (!matches!(status, TurnStatus::InProgress)).then_some(2), + duration_ms: (!matches!(status, TurnStatus::InProgress)).then_some(1_000), + status, + } +} + +fn turn_with_user_and_agent_items( + id: &str, + status: TurnStatus, + user_text: &str, + agent_text: &str, +) -> Turn { + Turn { + id: id.to_string(), + items: vec![ + ThreadItem::UserMessage { + id: format!("{id}-user"), + client_id: None, + content: vec![UserInput::Text { + text: user_text.to_string(), + text_elements: Vec::new(), + }], + }, + ThreadItem::AgentMessage { + id: format!("{id}-agent"), + text: agent_text.to_string(), + phase: None, + memory_citation: None, + }, + ], + items_view: TurnItemsView::Full, + error: None, + started_at: Some(1), + completed_at: (!matches!(status, TurnStatus::InProgress)).then_some(2), + duration_ms: (!matches!(status, TurnStatus::InProgress)).then_some(1_000), + status, + } +} + +fn completed_rollout(turn_id: &str, user_text: &str, agent_text: &str) -> Vec { + [ + 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(), + }), + EventMsg::UserMessage(UserMessageEvent { + client_id: None, + message: user_text.to_string(), + images: None, + local_images: Vec::new(), + text_elements: Vec::new(), + ..Default::default() + }), + EventMsg::AgentMessage(AgentMessageEvent { + message: agent_text.to_string(), + phase: None, + memory_citation: None, + }), + EventMsg::TurnComplete(TurnCompleteEvent { + turn_id: turn_id.to_string(), + last_agent_message: Some(agent_text.to_string()), + completed_at: Some(2), + duration_ms: Some(1_000), + time_to_first_token_ms: None, + }), + ] + .into_iter() + .map(RolloutItem::EventMsg) + .collect() +} + +fn incomplete_rollout(turn_id: &str, user_text: &str, agent_text: &str) -> Vec { + [ + 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(), + }), + EventMsg::UserMessage(UserMessageEvent { + client_id: None, + message: user_text.to_string(), + images: None, + local_images: Vec::new(), + text_elements: Vec::new(), + ..Default::default() + }), + EventMsg::AgentMessage(AgentMessageEvent { + message: agent_text.to_string(), + phase: None, + memory_citation: None, + }), + ] + .into_iter() + .map(RolloutItem::EventMsg) + .collect() +} + +fn thread(turns: Vec) -> Thread { + Thread { + id: "thread".to_string(), + extra: None, + session_id: "session".to_string(), + forked_from_id: None, + parent_thread_id: None, + preview: String::new(), + ephemeral: false, + model_provider: "test".to_string(), + created_at: 0, + updated_at: 0, + recency_at: None, + status: ThreadStatus::Idle, + path: None, + cwd: AbsolutePathBuf::from_absolute_path(std::env::temp_dir()) + .expect("temporary directory should be absolute"), + cli_version: "test".to_string(), + source: SessionSource::Cli, + thread_source: None, + agent_nickname: None, + agent_role: None, + git_info: None, + name: None, + turns, + } +} 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 0867532e55..9289d4bef1 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 @@ -28,8 +28,13 @@ use codex_app_server_protocol::ThreadTurnsListParams; use codex_app_server_protocol::ThreadTurnsListResponse; use codex_app_server_protocol::TurnItemsView; use codex_app_server_protocol::TurnStartParams; +use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::TurnStatus; use codex_app_server_protocol::UserInput; +use codex_core::RolloutRecorder; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::InitialHistory; +use codex_protocol::protocol::RolloutItem; use core_test_support::responses; use core_test_support::streaming_sse::StreamingSseChunk; use core_test_support::streaming_sse::start_streaming_sse_server; @@ -392,7 +397,7 @@ async fn websocket_disconnect_keeps_last_subscribed_thread_loaded_until_idle_tim } #[tokio::test] -async fn websocket_reconnect_resumes_in_progress_agent_message_deltas() -> Result<()> { +async fn websocket_reconnect_recovers_in_progress_and_terminal_turn_state() -> Result<()> { let (emit_deltas_tx, emit_deltas_rx) = oneshot::channel(); let (complete_turn_tx, complete_turn_rx) = oneshot::channel(); let (server, _) = start_streaming_sse_server(vec![vec![ @@ -426,7 +431,22 @@ async fn websocket_reconnect_resumes_in_progress_agent_message_deltas() -> Resul let mut ws1 = connect_websocket(bind_addr).await?; send_initialize_request(&mut ws1, /*id*/ 1, "ws_thread_owner").await?; read_response_for_id(&mut ws1, /*id*/ 1).await?; - let thread_id = start_thread(&mut ws1, /*id*/ 2).await?; + send_request( + &mut ws1, + "thread/start", + /*id*/ 2, + Some(serde_json::to_value(ThreadStartParams { + model: Some("mock-model".to_string()), + ..Default::default() + })?), + ) + .await?; + let response = read_response_for_id(&mut ws1, /*id*/ 2).await?; + let ThreadStartResponse { thread, .. } = to_response(response)?; + let thread_id = thread.id; + let rollout_path = thread + .path + .context("started thread should have a rollout path")?; send_request( &mut ws1, @@ -442,26 +462,21 @@ async fn websocket_reconnect_resumes_in_progress_agent_message_deltas() -> Resul })?), ) .await?; - read_response_for_id(&mut ws1, /*id*/ 3).await?; + let response = read_response_for_id(&mut ws1, /*id*/ 3).await?; + let TurnStartResponse { turn } = to_response(response)?; + let turn_id = turn.id; read_notification_for_method(&mut ws1, "item/started").await?; - send_request( - &mut ws1, - "thread/unsubscribe", - /*id*/ 4, - Some(json!({ "threadId": thread_id })), - ) - .await?; - read_response_for_id(&mut ws1, /*id*/ 4).await?; - - ws1.close(None).await.context("failed to close websocket")?; drop(ws1); + emit_deltas_tx + .send(()) + .expect("streaming response should still be waiting to emit deltas"); let mut ws2 = connect_websocket(bind_addr).await?; send_request( &mut ws2, "initialize", - /*id*/ 5, + /*id*/ 4, Some(serde_json::to_value(InitializeParams { client_info: ClientInfo { name: "codex_chatgpt_ios_remote".to_string(), @@ -475,13 +490,9 @@ async fn websocket_reconnect_resumes_in_progress_agent_message_deltas() -> Resul })?), ) .await?; - read_response_for_id(&mut ws2, /*id*/ 5).await?; + read_response_for_id(&mut ws2, /*id*/ 4).await?; - emit_deltas_tx - .send(()) - .expect("streaming response should still be waiting to emit deltas"); - - let mut request_id = 6; + let mut request_id = 5; loop { send_request( &mut ws2, @@ -500,14 +511,15 @@ async fn websocket_reconnect_resumes_in_progress_agent_message_deltas() -> Resul let ThreadTurnsListResponse { data, .. } = to_response(response)?; request_id += 1; if data.iter().any(|turn| { - turn.status == TurnStatus::InProgress + turn.id == turn_id + && turn.status == TurnStatus::InProgress && turn.items.iter().any(|item| { - matches!(item, ThreadItem::AgentMessage { id, text, .. } if id == "msg-1" && text == "hello world") + matches!(item, ThreadItem::AgentMessage { text, .. } if text == "hello world") }) }) { break; } - if request_id == 26 { + if request_id == 25 { bail!("thread never reflected deltas emitted with no subscribed client"); } } @@ -536,17 +548,102 @@ async fn websocket_reconnect_resumes_in_progress_agent_message_deltas() -> Resul .expect("resume should include the requested initial turns page") .data .into_iter() - .find(|turn| turn.status == TurnStatus::InProgress) - .expect("resume should include the in-progress turn"); - assert!(active_turn.items.iter().any(|item| { - matches!(item, ThreadItem::AgentMessage { id, text, .. } if id == "msg-1" && text == "hello world") - })); + .find(|turn| turn.id == turn_id && turn.status == TurnStatus::InProgress) + .expect("resume should include the in-progress turn updated while disconnected"); + assert_eq!( + active_turn + .items + .iter() + .filter_map(|item| match item { + ThreadItem::AgentMessage { text, .. } => Some(text.as_str()), + _ => None, + }) + .collect::>(), + vec!["hello world"] + ); + drop(ws2); complete_turn_tx .send(()) .expect("streaming response should still be waiting for completion"); - read_notification_for_method(&mut ws2, "item/completed").await?; - read_notification_for_method(&mut ws2, "turn/completed").await?; + timeout(DEFAULT_READ_TIMEOUT, async { + loop { + let InitialHistory::Resumed(history) = + RolloutRecorder::get_rollout_history(&rollout_path).await? + else { + bail!("started thread should have resumable rollout history"); + }; + if history.history.iter().any(|item| { + matches!( + item, + RolloutItem::EventMsg(EventMsg::TurnComplete(event)) + if event.turn_id == turn_id + ) + }) { + return Ok::<(), anyhow::Error>(()); + } + sleep(Duration::from_millis(10)).await; + } + }) + .await + .context("turn did not finish persisting while the websocket was disconnected")??; + + let mut ws3 = connect_websocket(bind_addr).await?; + send_request( + &mut ws3, + "initialize", + /*id*/ 6, + 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 ws3, /*id*/ 6).await?; + send_request( + &mut ws3, + "thread/resume", + /*id*/ 7, + Some(serde_json::to_value(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 response = read_response_for_id(&mut ws3, /*id*/ 7).await?; + let ThreadResumeResponse { + initial_turns_page, .. + } = to_response(response)?; + let completed_turn = initial_turns_page + .expect("resume should include the requested initial turns page") + .data + .into_iter() + .find(|turn| turn.id == turn_id && turn.status == TurnStatus::Completed) + .expect("resume should include the turn completed while disconnected"); + assert_eq!( + completed_turn + .items + .iter() + .filter_map(|item| match item { + ThreadItem::AgentMessage { text, .. } => Some(text.as_str()), + _ => None, + }) + .collect::>(), + vec!["hello world"] + ); process .kill() 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 2c8914f90b..3fadfb2b38 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,18 +26,22 @@ 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; +use codex_app_server_protocol::SortDirection; use codex_app_server_protocol::ThreadDeleteParams; 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::TurnItemsView; use codex_app_server_protocol::TurnStartParams; use codex_app_server_protocol::TurnStatus; use codex_app_server_protocol::UserInput as V2UserInput; @@ -405,6 +409,117 @@ async fn running_thread_resume_includes_completion_after_initial_history_snapsho Ok(()) } +#[tokio::test] +async fn running_thread_resume_reads_history_for_every_response_shape_that_needs_it() -> Result<()> +{ + struct Case { + exclude_turns: bool, + include_initial_page: bool, + expected_history_reads: usize, + } + + let server = create_mock_responses_server_repeating_assistant("Done").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 thread_store = InMemoryThreadStore::for_id(store_id.clone()); + let _in_memory_store = InMemoryThreadStoreId { 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)?; + 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: "Materialize one completed turn".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }, + }) + .await? + .expect("turn/start should succeed"); + wait_for_thread_notification(&mut client, &thread.id, |notification| { + matches!(notification, ServerNotification::TurnCompleted(_)) + }) + .await?; + + let cases = [ + Case { + exclude_turns: true, + include_initial_page: false, + expected_history_reads: 0, + }, + Case { + exclude_turns: false, + include_initial_page: false, + expected_history_reads: 1, + }, + Case { + exclude_turns: true, + include_initial_page: true, + expected_history_reads: 1, + }, + Case { + exclude_turns: false, + include_initial_page: true, + expected_history_reads: 1, + }, + ]; + + for (index, case) in cases.into_iter().enumerate() { + let reads_before = thread_store.calls().await.load_history; + let response = client + .request(ClientRequest::ThreadResume { + request_id: RequestId::Integer(3 + index as i64), + params: ThreadResumeParams { + thread_id: thread.id.clone(), + exclude_turns: case.exclude_turns, + initial_turns_page: case.include_initial_page.then_some( + ThreadResumeInitialTurnsPageParams { + limit: None, + sort_direction: Some(SortDirection::Asc), + items_view: Some(TurnItemsView::Full), + }, + ), + ..Default::default() + }, + }) + .await? + .expect("thread/resume should succeed"); + let response: ThreadResumeResponse = serde_json::from_value(response)?; + + assert_eq!( + thread_store.calls().await.load_history, + reads_before + case.expected_history_reads + ); + assert_eq!(response.thread.turns.is_empty(), case.exclude_turns); + assert_eq!( + response.initial_turns_page.is_some(), + case.include_initial_page + ); + if let Some(page) = response.initial_turns_page { + assert!(!page.data.is_empty()); + if !case.exclude_turns { + assert_eq!(page.data, response.thread.turns); + } + } + } + + client.shutdown().await?; + Ok(()) +} + async fn wait_for_thread_notification( client: &mut InProcessClientHandle, thread_id: &str, @@ -470,7 +585,10 @@ async fn start_in_process_client( title: None, version: "0.1.0".to_string(), }, - capabilities: None, + capabilities: Some(InitializeCapabilities { + experimental_api: true, + ..Default::default() + }), }, channel_capacity: in_process::DEFAULT_IN_PROCESS_CHANNEL_CAPACITY, })