diff --git a/codex-rs/rollout/src/lib.rs b/codex-rs/rollout/src/lib.rs index 7cf140dc20..450113218b 100644 --- a/codex-rs/rollout/src/lib.rs +++ b/codex-rs/rollout/src/lib.rs @@ -8,6 +8,7 @@ pub(crate) mod compression; pub(crate) mod config; pub(crate) mod list; pub(crate) mod metadata; +mod model_context; mod ordinal; mod persistence_metrics; pub(crate) mod policy; @@ -59,6 +60,8 @@ pub use list::read_session_meta_line; pub use list::read_thread_item_from_rollout; pub use list::rollout_date_parts; pub use metadata::builder_from_items; +pub use model_context::ModelContextScan; +pub use model_context::ModelContextScanProgress; pub use persistence_metrics::RolloutPersistenceBatchMeasurement; pub use persistence_metrics::RolloutPersistenceTelemetry; pub use persistence_metrics::measure_and_filter_rollout_items; @@ -68,6 +71,8 @@ pub use policy::should_persist_response_item_for_memories; pub use recorder::RolloutRecorder; pub use recorder::RolloutRecorderParams; pub use recorder::append_rollout_item_to_path; +pub use reverse_jsonl_scanner::ReverseJsonlScanner; +pub use reverse_jsonl_scanner::ScanOutcome; pub use search::first_rollout_content_match_snippet; pub use search::search_rollout_matches; pub use search::search_rollout_paths; diff --git a/codex-rs/rollout/src/model_context.rs b/codex-rs/rollout/src/model_context.rs new file mode 100644 index 0000000000..45fc5de6bf --- /dev/null +++ b/codex-rs/rollout/src/model_context.rs @@ -0,0 +1,197 @@ +use codex_protocol::items::TurnItem; +use codex_protocol::models::ResponseItem; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::InterAgentCommunication; +use codex_protocol::protocol::RolloutItem; +use codex_protocol::protocol::SessionMetaLine; + +/// Whether a reverse model-context scan needs more rollout items. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ModelContextScanProgress { + /// The reader should provide the next older rollout item. + Continue, + /// The scan has collected a safe bounded suffix. + Complete, +} + +/// Accumulates newest-to-oldest rollout items until they are sufficient to reconstruct the latest +/// model context. +/// +/// Storage implementations own how they fetch older items. Local JSONL readers and future +/// reverse-paged cloud readers can both feed their items through this scan to share the cutoff +/// rules and chronological replay assembly. +/// +/// The scan stops once it has both: +/// +/// - `saw_compaction`: a `CompactedItem` with `replacement_history` and `window_number`; +/// - `saw_completed_turn_context`: a completed user turn with a compatible `TurnContextItem`. +/// +/// If the scan reaches the beginning before finding a bounded cutoff, it has already collected +/// the complete replay and so we can return that directly. +/// +/// `TurnContextItem` does not identify whether it came from a user turn, so one only counts after +/// the same turn also proves a user-turn boundary: a paginated +/// `ItemCompleted(UserMessage)` marker, agent message, or inter-agent message. Paginated writers +/// persist that marker for real user turns; older rollouts without it conservatively scan to the +/// beginning. A raw `role=user` response item is not sufficient because contextual user fragments +/// use that role but do not count as turn boundaries during reconstruction. The compaction restores +/// model-visible items; the turn context restores previous settings (`model`, `comp_hash`, and +/// `realtime_active`) and the reference baseline. +/// +/// These paginated shapes disable the bounded cutoff: +/// +/// - compaction without `replacement_history` or `window_number`; +/// - rollback markers; +/// +/// When one appears, the scanner continues to the beginning and returns the complete replay. +#[derive(Debug, Default)] +pub struct ModelContextScan { + items_newest_first: Vec, + saw_compaction: bool, + saw_completed_turn_context: bool, + must_scan_to_start: bool, + active_segment: ActiveTurnSegment, +} + +impl ModelContextScan { + /// Adds the next newest-to-oldest rollout item and reports whether the reader can stop. + pub fn push(&mut self, item: RolloutItem) -> ModelContextScanProgress { + let progress = self.observe(&item); + self.items_newest_first.push(item); + progress + } + + /// Returns the collected items in chronological order with canonical head metadata. + /// + /// Call this after the reader reaches the beginning of its source or after [`Self::push`] + /// returns [`ModelContextScanProgress::Complete`]. + pub fn finish(mut self, session_meta: SessionMetaLine) -> Vec { + self.items_newest_first.reverse(); + if self.has_bounded_cutoff() { + // A bounded scan stops before reaching the head. Prepend the separately loaded head + // SessionMeta, which remains canonical when copied fork history contains later + // metadata. + self.items_newest_first + .insert(0, RolloutItem::SessionMeta(session_meta)); + } + self.items_newest_first + } + + fn observe(&mut self, item: &RolloutItem) -> ModelContextScanProgress { + if self.must_scan_to_start { + return ModelContextScanProgress::Continue; + } + + match item { + RolloutItem::Compacted(compacted) + if compacted.replacement_history.is_none() || compacted.window_number.is_none() => + { + self.must_scan_to_start = true; + } + RolloutItem::Compacted(_) => { + self.saw_compaction = true; + } + RolloutItem::EventMsg(EventMsg::ThreadRolledBack(_)) => { + // Paginated threads reject rollback. Keep old rollouts correct rather than + // duplicating rollback survival semantics in this bounded selector. + self.must_scan_to_start = true; + } + RolloutItem::EventMsg(EventMsg::ItemCompleted(event)) => { + if self.active_segment.turn_id.is_none() { + self.active_segment.turn_id = Some(event.turn_id.clone()); + } + if turn_ids_are_compatible( + self.active_segment.turn_id.as_deref(), + Some(event.turn_id.as_str()), + ) { + self.active_segment.has_user_turn |= + matches!(&event.item, TurnItem::UserMessage(_)); + } + } + RolloutItem::EventMsg(EventMsg::TurnComplete(event)) => { + self.active_segment + .turn_id + .get_or_insert_with(|| event.turn_id.clone()); + } + RolloutItem::EventMsg(EventMsg::TurnAborted(event)) => { + if let Some(turn_id) = &event.turn_id { + self.active_segment + .turn_id + .get_or_insert_with(|| turn_id.clone()); + } + } + RolloutItem::EventMsg(EventMsg::TurnStarted(event)) => { + if turn_ids_are_compatible( + self.active_segment.turn_id.as_deref(), + Some(event.turn_id.as_str()), + ) { + self.finalize_active_segment(); + } + } + RolloutItem::TurnContext(context) => { + if self.active_segment.turn_id.is_none() { + self.active_segment.turn_id = context.turn_id.clone(); + } + if turn_ids_are_compatible( + self.active_segment.turn_id.as_deref(), + context.turn_id.as_deref(), + ) { + self.active_segment.has_turn_context = true; + } + } + RolloutItem::ResponseItem(response_item) => { + self.active_segment.has_user_turn |= + response_item_counts_as_user_turn(response_item); + } + RolloutItem::InterAgentCommunication(_) => { + self.active_segment.has_user_turn = true; + } + RolloutItem::EventMsg(EventMsg::UserMessage(_)) => { + self.active_segment.has_user_turn = true; + } + RolloutItem::EventMsg(_) + | RolloutItem::SessionMeta(_) + | RolloutItem::InterAgentCommunicationMetadata { .. } + | RolloutItem::WorldState(_) => {} + } + + if self.has_bounded_cutoff() { + ModelContextScanProgress::Complete + } else { + ModelContextScanProgress::Continue + } + } + + fn finalize_active_segment(&mut self) { + if self.active_segment.has_user_turn && self.active_segment.has_turn_context { + self.saw_completed_turn_context = true; + } + self.active_segment = ActiveTurnSegment::default(); + } + + fn has_bounded_cutoff(&self) -> bool { + !self.must_scan_to_start && self.saw_compaction && self.saw_completed_turn_context + } +} + +#[derive(Debug, Default)] +struct ActiveTurnSegment { + turn_id: Option, + has_user_turn: bool, + has_turn_context: bool, +} + +fn turn_ids_are_compatible(active_turn_id: Option<&str>, item_turn_id: Option<&str>) -> bool { + active_turn_id + .is_none_or(|turn_id| item_turn_id.is_none_or(|item_turn_id| item_turn_id == turn_id)) +} + +fn response_item_counts_as_user_turn(response_item: &ResponseItem) -> bool { + match response_item { + ResponseItem::AgentMessage { .. } => true, + ResponseItem::Message { role, content, .. } => { + role == "assistant" && InterAgentCommunication::is_message_content(content) + } + _ => false, + } +} diff --git a/codex-rs/rollout/src/reverse_jsonl_scanner.rs b/codex-rs/rollout/src/reverse_jsonl_scanner.rs index c84809ca1a..530e1fdd06 100644 --- a/codex-rs/rollout/src/reverse_jsonl_scanner.rs +++ b/codex-rs/rollout/src/reverse_jsonl_scanner.rs @@ -8,14 +8,16 @@ use serde::de::DeserializeOwned; const READ_CHUNK_SIZE: usize = 8 * 1024; #[derive(Debug)] -pub(crate) enum ScanOutcome { +pub enum ScanOutcome { + /// The record was valid JSON and deserialized as the requested type. Parsed(T), + /// The record was not valid JSON for the requested type. #[allow(dead_code)] Rejected(serde_json::Error), } /// Read-only scanner for newline-delimited JSON records, starting from the end. -pub(crate) struct ReverseJsonlScanner { +pub struct ReverseJsonlScanner { reader: R, next_chunk_end: u64, chunk_position: usize, @@ -27,7 +29,7 @@ impl ReverseJsonlScanner where R: Read + Seek, { - pub(crate) fn new(mut reader: R) -> io::Result { + pub fn new(mut reader: R) -> io::Result { let next_chunk_end = reader.seek(SeekFrom::End(0))?; Ok(Self { reader, @@ -42,7 +44,7 @@ where /// /// I/O failures are returned as [`Err`]. Invalid JSON records are returned as /// [`ScanOutcome::Rejected`], and the scanner remains usable. - pub(crate) fn scan_next(&mut self) -> io::Result>> + pub fn scan_next(&mut self) -> io::Result>> where T: DeserializeOwned, { diff --git a/codex-rs/thread-store/src/in_memory.rs b/codex-rs/thread-store/src/in_memory.rs index 76f8727c8b..cbaccceb36 100644 --- a/codex-rs/thread-store/src/in_memory.rs +++ b/codex-rs/thread-store/src/in_memory.rs @@ -27,6 +27,7 @@ use crate::LoadThreadHistoryParams; use crate::ReadThreadByRolloutPathParams; use crate::ReadThreadParams; use crate::ResumeThreadParams; +use crate::StoredModelContext; use crate::StoredThread; use crate::StoredThreadHistory; use crate::ThreadMetadataPatch; @@ -368,6 +369,7 @@ pub struct InMemoryThreadStoreCalls { pub shutdown_thread: usize, pub discard_thread: usize, pub load_history: usize, + pub load_latest_model_context: usize, pub read_thread: usize, pub read_thread_with_history: usize, pub read_thread_by_rollout_path: usize, @@ -521,6 +523,25 @@ impl InMemoryThreadStore { }) } + async fn load_latest_model_context( + &self, + params: LoadThreadHistoryParams, + ) -> ThreadStoreResult { + let mut state = self.state.lock().await; + state.calls.load_latest_model_context += 1; + let items = + state + .histories + .get(¶ms.thread_id) + .ok_or(ThreadStoreError::ThreadNotFound { + thread_id: params.thread_id, + })?; + Ok(StoredModelContext { + thread_id: params.thread_id, + items: items.clone(), + }) + } + async fn read_thread(&self, params: ReadThreadParams) -> ThreadStoreResult { let mut state = self.state.lock().await; state.calls.read_thread += 1; @@ -659,6 +680,13 @@ impl ThreadStore for InMemoryThreadStore { Box::pin(InMemoryThreadStore::load_history(self, params)) } + fn load_latest_model_context( + &self, + params: LoadThreadHistoryParams, + ) -> ThreadStoreFuture<'_, StoredModelContext> { + Box::pin(InMemoryThreadStore::load_latest_model_context(self, params)) + } + fn read_thread(&self, params: ReadThreadParams) -> ThreadStoreFuture<'_, StoredThread> { Box::pin(InMemoryThreadStore::read_thread(self, params)) } diff --git a/codex-rs/thread-store/src/lib.rs b/codex-rs/thread-store/src/lib.rs index 34a0ab2ba1..7dcf07c7a2 100644 --- a/codex-rs/thread-store/src/lib.rs +++ b/codex-rs/thread-store/src/lib.rs @@ -39,6 +39,7 @@ pub use types::ReadThreadParams; pub use types::ResumeThreadParams; pub use types::SearchThreadsParams; pub use types::SortDirection; +pub use types::StoredModelContext; pub use types::StoredThread; pub use types::StoredThreadHistory; pub use types::StoredThreadItem; diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs index d4275ff79e..e7e3f1e5d4 100644 --- a/codex-rs/thread-store/src/local/mod.rs +++ b/codex-rs/thread-store/src/local/mod.rs @@ -4,6 +4,7 @@ mod delete_thread; mod helpers; mod list_threads; mod live_writer; +mod model_context; mod read_thread; mod search_threads; mod unarchive_thread; @@ -32,6 +33,7 @@ use crate::ReadThreadByRolloutPathParams; use crate::ReadThreadParams; use crate::ResumeThreadParams; use crate::SearchThreadsParams; +use crate::StoredModelContext; use crate::StoredThread; use crate::StoredThreadHistory; use crate::ThreadPage; @@ -277,6 +279,13 @@ impl ThreadStore for LocalThreadStore { Box::pin(LocalThreadStore::load_history(self, params)) } + fn load_latest_model_context( + &self, + params: LoadThreadHistoryParams, + ) -> ThreadStoreFuture<'_, StoredModelContext> { + Box::pin(async move { model_context::load_latest_model_context(self, params).await }) + } + fn read_thread(&self, params: ReadThreadParams) -> ThreadStoreFuture<'_, StoredThread> { Box::pin(async move { read_thread::read_thread(self, params).await }) } diff --git a/codex-rs/thread-store/src/local/model_context.rs b/codex-rs/thread-store/src/local/model_context.rs new file mode 100644 index 0000000000..7230432d67 --- /dev/null +++ b/codex-rs/thread-store/src/local/model_context.rs @@ -0,0 +1,121 @@ +use std::fs::File; +use std::io; +use std::path::Path; +use std::path::PathBuf; + +use codex_protocol::protocol::RolloutItem; +use codex_protocol::protocol::RolloutLine; +use codex_protocol::protocol::SessionMetaLine; +use codex_protocol::protocol::ThreadHistoryMode; +use codex_rollout::ModelContextScan; +use codex_rollout::ModelContextScanProgress; +use codex_rollout::ReverseJsonlScanner; +use codex_rollout::ScanOutcome; + +use super::LocalThreadStore; +use super::read_thread; +use crate::LoadThreadHistoryParams; +use crate::StoredModelContext; +use crate::ThreadStoreError; +use crate::ThreadStoreResult; + +#[cfg(test)] +#[path = "model_context_tests.rs"] +mod tests; + +/// Loads rollout items needed to reconstruct the latest model-visible context. +/// +/// Plain paginated JSONL rollouts use a reverse scan. When it finds both a usable replacement- +/// history checkpoint and the completed user-turn context needed for resume metadata, the returned +/// replay starts with the canonical head `SessionMeta` followed by that newest suffix. When no +/// bounded cutoff is available, the scan continues to the beginning and returns the complete +/// replay it already accumulated. +/// +/// Legacy and compressed rollout shapes keep the existing full-history path. +pub(super) async fn load_latest_model_context( + store: &LocalThreadStore, + params: LoadThreadHistoryParams, +) -> ThreadStoreResult { + let path = read_thread::resolve_rollout_path(store, params.thread_id, params.include_archived) + .await? + .ok_or_else(|| ThreadStoreError::InvalidRequest { + message: format!("no rollout found for thread id {}", params.thread_id), + })?; + + let session_meta = codex_rollout::read_session_meta_line(path.as_path()) + .await + .map_err(|err| ThreadStoreError::Internal { + message: format!("failed to read session metadata {}: {err}", path.display()), + })?; + if session_meta.meta.id != params.thread_id { + return Err(ThreadStoreError::InvalidRequest { + message: format!( + "rollout at {} belongs to thread {}, not {}", + path.display(), + session_meta.meta.id, + params.thread_id + ), + }); + } + + let items = if matches!(session_meta.meta.history_mode, ThreadHistoryMode::Paginated) + && !path + .file_name() + .and_then(|file_name| file_name.to_str()) + .is_some_and(|file_name| file_name.ends_with(".jsonl.zst")) + { + scan_model_context_from_end(path, session_meta).await? + } else { + read_thread::load_history_items(path.as_path()).await? + }; + + Ok(StoredModelContext { + thread_id: params.thread_id, + items, + }) +} + +async fn scan_model_context_from_end( + path: PathBuf, + session_meta: SessionMetaLine, +) -> ThreadStoreResult> { + let path_for_scan = path.clone(); + let scan = tokio::task::spawn_blocking(move || { + scan_model_context_from_end_blocking(&path_for_scan, session_meta) + }) + .await + .map_err(|err| ThreadStoreError::Internal { + message: format!("failed to join model context scan: {err}"), + })?; + match scan { + Ok(items) => Ok(items), + Err(err) if err.kind() == io::ErrorKind::NotFound => { + // Compression can replace the resolved plain rollout with its compressed sibling + // before the blocking reverse scanner opens it. The forward loader re-resolves that + // representation transition and already supports compressed rollouts. + read_thread::load_history_items(path.as_path()).await + } + Err(err) => Err(ThreadStoreError::Internal { + message: format!("failed to scan model context {}: {err}", path.display()), + }), + } +} + +fn scan_model_context_from_end_blocking( + path: &Path, + session_meta: SessionMetaLine, +) -> io::Result> { + let mut scan = ModelContextScan::default(); + let mut scanner = ReverseJsonlScanner::new(File::open(path)?)?; + while let Some(outcome) = scanner.scan_next::()? { + let ScanOutcome::Parsed(line) = outcome else { + continue; + }; + match scan.push(line.item) { + ModelContextScanProgress::Continue => {} + ModelContextScanProgress::Complete => break, + } + } + + Ok(scan.finish(session_meta)) +} diff --git a/codex-rs/thread-store/src/local/model_context_tests.rs b/codex-rs/thread-store/src/local/model_context_tests.rs new file mode 100644 index 0000000000..89a7c0bdaf --- /dev/null +++ b/codex-rs/thread-store/src/local/model_context_tests.rs @@ -0,0 +1,398 @@ +use std::fs::OpenOptions; +use std::io::Write; +use std::path::Path; +use std::path::PathBuf; + +use codex_protocol::config_types::ReasoningSummary; +use codex_protocol::items::TurnItem; +use codex_protocol::items::UserMessageItem; +use codex_protocol::models::AgentMessageInputContent; +use codex_protocol::models::ContentItem; +use codex_protocol::models::ResponseItem; +use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::CompactedItem; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::ItemCompletedEvent; +use codex_protocol::protocol::RolloutItem; +use codex_protocol::protocol::RolloutLine; +use codex_protocol::protocol::SandboxPolicy; +use codex_protocol::protocol::ThreadHistoryMode; +use codex_protocol::protocol::TurnCompleteEvent; +use codex_protocol::protocol::TurnContextItem; +use codex_protocol::protocol::TurnStartedEvent; +use codex_protocol::user_input::UserInput; +use pretty_assertions::assert_eq; +use tempfile::TempDir; +use uuid::Uuid; + +use super::*; +use crate::ThreadStore; +use crate::local::test_support::test_config; +use crate::local::test_support::write_session_file_with_history_mode; + +#[tokio::test] +async fn loads_latest_checkpoint_with_required_turn_metadata() { + let home = TempDir::new().expect("temp dir"); + let uuid = Uuid::from_u128(/*v*/ 1001); + let thread_id = codex_protocol::ThreadId::from_string(&uuid.to_string()).expect("thread id"); + write_paginated_rollout( + home.path(), + "2025-01-03T13-00-00", + uuid, + [ + turn_started("turn-1"), + user_message("older turn"), + completed_user_message("turn-1", "older turn"), + turn_context(home.path(), "turn-1"), + compacted("older checkpoint", Some(Vec::new())), + turn_complete("turn-1"), + turn_started("turn-2"), + user_message("latest turn"), + completed_user_message("turn-2", "latest turn"), + turn_context(home.path(), "turn-2"), + compacted("latest checkpoint", Some(Vec::new())), + turn_complete("turn-2"), + ], + ); + let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + + let context = store + .load_latest_model_context(LoadThreadHistoryParams { + thread_id, + include_archived: false, + }) + .await + .expect("load model context"); + + assert!(matches!( + context.items.first(), + Some(RolloutItem::SessionMeta(_)) + )); + assert!(context.items.iter().any(|item| { + matches!(item, RolloutItem::Compacted(compacted) if compacted.message == "latest checkpoint") + })); + assert!(!context.items.iter().any(|item| { + matches!(item, RolloutItem::Compacted(compacted) if compacted.message == "older checkpoint") + })); + assert!(context.items.iter().any(|item| { + matches!(item, RolloutItem::TurnContext(context) if context.turn_id.as_deref() == Some("turn-2")) + })); +} + +#[tokio::test] +async fn loads_turn_metadata_across_an_older_checkpoint() { + let home = TempDir::new().expect("temp dir"); + let uuid = Uuid::from_u128(/*v*/ 1006); + let thread_id = codex_protocol::ThreadId::from_string(&uuid.to_string()).expect("thread id"); + write_paginated_rollout( + home.path(), + "2025-01-03T13-00-05", + uuid, + [ + turn_started("turn-0"), + user_message("oldest turn"), + completed_user_message("turn-0", "oldest turn"), + turn_context(home.path(), "turn-0"), + turn_complete("turn-0"), + turn_started("turn-1"), + user_message("metadata turn"), + completed_user_message("turn-1", "metadata turn"), + turn_context(home.path(), "turn-1"), + compacted("older checkpoint", Some(Vec::new())), + turn_complete("turn-1"), + turn_started("turn-2"), + compacted("latest checkpoint", Some(Vec::new())), + turn_complete("turn-2"), + ], + ); + let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + + let context = store + .load_latest_model_context(LoadThreadHistoryParams { + thread_id, + include_archived: false, + }) + .await + .expect("load model context"); + + assert!(context.items.iter().any(|item| { + matches!(item, RolloutItem::Compacted(compacted) if compacted.message == "latest checkpoint") + })); + assert!(context.items.iter().any(|item| { + matches!(item, RolloutItem::TurnContext(context) if context.turn_id.as_deref() == Some("turn-1")) + })); + assert!(!context.items.iter().any(|item| { + matches!(item, RolloutItem::TurnContext(context) if context.turn_id.as_deref() == Some("turn-0")) + })); +} + +#[tokio::test] +async fn returns_scanned_full_history_for_unsupported_compaction() { + let home = TempDir::new().expect("temp dir"); + let uuid = Uuid::from_u128(/*v*/ 1002); + let path = write_paginated_rollout( + home.path(), + "2025-01-03T13-00-01", + uuid, + [ + turn_started("turn-1"), + user_message("turn"), + completed_user_message("turn-1", "turn"), + turn_context(home.path(), "turn-1"), + compacted("usable checkpoint", Some(Vec::new())), + compacted("legacy checkpoint", /*replacement_history*/ None), + turn_complete("turn-1"), + ], + ); + + assert_reverse_scan_matches_full_history(path.as_path()).await; +} + +#[tokio::test] +async fn returns_scanned_full_history_at_bof_without_checkpoint() { + let home = TempDir::new().expect("temp dir"); + let uuid = Uuid::from_u128(/*v*/ 1003); + let path = write_paginated_rollout( + home.path(), + "2025-01-03T13-00-02", + uuid, + [ + turn_started("turn-1"), + user_message("turn"), + completed_user_message("turn-1", "turn"), + turn_context(home.path(), "turn-1"), + turn_complete("turn-1"), + ], + ); + + assert_reverse_scan_matches_full_history(path.as_path()).await; +} + +#[tokio::test] +async fn uses_agent_message_turn_context_without_scanning_older_turn() { + let home = TempDir::new().expect("temp dir"); + let uuid = Uuid::from_u128(/*v*/ 1004); + let thread_id = codex_protocol::ThreadId::from_string(&uuid.to_string()).expect("thread id"); + write_paginated_rollout( + home.path(), + "2025-01-03T13-00-03", + uuid, + [ + turn_started("turn-1"), + user_message("older turn"), + completed_user_message("turn-1", "older turn"), + turn_context(home.path(), "turn-1"), + compacted("checkpoint", Some(Vec::new())), + turn_complete("turn-1"), + turn_started("turn-2"), + turn_context(home.path(), "turn-2"), + agent_message("child done"), + turn_complete("turn-2"), + ], + ); + let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + + let context = store + .load_latest_model_context(LoadThreadHistoryParams { + thread_id, + include_archived: false, + }) + .await + .expect("load model context"); + + assert!(context.items.iter().any(|item| { + matches!(item, RolloutItem::TurnContext(context) if context.turn_id.as_deref() == Some("turn-2")) + })); + assert!(!context.items.iter().any(|item| { + matches!(item, RolloutItem::TurnContext(context) if context.turn_id.as_deref() == Some("turn-1")) + })); +} + +#[tokio::test] +async fn ignores_contextual_user_messages_when_selecting_turn_context() { + let home = TempDir::new().expect("temp dir"); + let uuid = Uuid::from_u128(/*v*/ 1005); + let thread_id = codex_protocol::ThreadId::from_string(&uuid.to_string()).expect("thread id"); + write_paginated_rollout( + home.path(), + "2025-01-03T13-00-04", + uuid, + [ + turn_started("turn-1"), + user_message("real user turn"), + completed_user_message("turn-1", "real user turn"), + turn_context(home.path(), "turn-1"), + compacted("checkpoint", Some(Vec::new())), + turn_complete("turn-1"), + turn_started("turn-2"), + contextual_user_message(), + turn_context(home.path(), "turn-2"), + turn_complete("turn-2"), + ], + ); + let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + + let context = store + .load_latest_model_context(LoadThreadHistoryParams { + thread_id, + include_archived: false, + }) + .await + .expect("load model context"); + + assert!(context.items.iter().any(|item| { + matches!(item, RolloutItem::TurnContext(context) if context.turn_id.as_deref() == Some("turn-1")) + })); +} + +fn write_paginated_rollout( + home: &Path, + timestamp: &str, + uuid: Uuid, + items: [RolloutItem; N], +) -> PathBuf { + let path = + write_session_file_with_history_mode(home, timestamp, uuid, ThreadHistoryMode::Paginated) + .expect("write session file"); + append_items(path.as_path(), items); + path +} + +async fn assert_reverse_scan_matches_full_history(path: &Path) { + let session_meta = codex_rollout::read_session_meta_line(path) + .await + .expect("read session metadata"); + let items = + scan_model_context_from_end_blocking(path, session_meta).expect("scan model context"); + let full_items = read_thread::load_history_items(path) + .await + .expect("load full history"); + + assert_eq!( + serde_json::to_value(items).expect("serialize scanned items"), + serde_json::to_value(full_items).expect("serialize full items") + ); +} + +fn append_items(path: &Path, items: [RolloutItem; N]) { + let mut file = OpenOptions::new() + .append(true) + .open(path) + .expect("open session file"); + for item in items { + let line = RolloutLine { + timestamp: "2025-01-03T13:00:01Z".to_string(), + ordinal: None, + item, + }; + writeln!( + file, + "{}", + serde_json::to_string(&line).expect("serialize line") + ) + .expect("append rollout line"); + } +} + +fn turn_started(turn_id: &str) -> RolloutItem { + RolloutItem::EventMsg(EventMsg::TurnStarted(TurnStartedEvent { + turn_id: turn_id.to_string(), + trace_id: None, + started_at: None, + model_context_window: Some(128_000), + collaboration_mode_kind: Default::default(), + })) +} + +fn turn_complete(turn_id: &str) -> RolloutItem { + RolloutItem::EventMsg(EventMsg::TurnComplete(TurnCompleteEvent { + turn_id: turn_id.to_string(), + last_agent_message: None, + error: None, + started_at: None, + completed_at: None, + duration_ms: None, + time_to_first_token_ms: None, + })) +} + +fn user_message(message: &str) -> RolloutItem { + RolloutItem::ResponseItem(ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::InputText { + text: message.to_string(), + }], + phase: None, + internal_chat_message_metadata_passthrough: None, + }) +} + +fn contextual_user_message() -> RolloutItem { + user_message("context only") +} + +fn completed_user_message(turn_id: &str, message: &str) -> RolloutItem { + RolloutItem::EventMsg(EventMsg::ItemCompleted(ItemCompletedEvent { + thread_id: codex_protocol::ThreadId::from_string("00000000-0000-0000-0000-000000000000") + .expect("fixture thread id"), + turn_id: turn_id.to_string(), + item: TurnItem::UserMessage(UserMessageItem { + id: format!("user-{turn_id}"), + client_id: None, + content: vec![UserInput::Text { + text: message.to_string(), + text_elements: Vec::new(), + }], + }), + completed_at_ms: 0, + })) +} + +fn agent_message(message: &str) -> RolloutItem { + RolloutItem::ResponseItem(ResponseItem::AgentMessage { + id: None, + author: "worker".to_string(), + recipient: "root".to_string(), + content: vec![AgentMessageInputContent::InputText { + text: message.to_string(), + }], + internal_chat_message_metadata_passthrough: None, + }) +} + +fn turn_context(root: &Path, turn_id: &str) -> RolloutItem { + RolloutItem::TurnContext(TurnContextItem { + turn_id: Some(turn_id.to_string()), + cwd: serde_json::from_value(serde_json::json!(root)).expect("absolute cwd"), + workspace_roots: None, + current_date: None, + timezone: None, + approval_policy: AskForApproval::Never, + approvals_reviewer: None, + sandbox_policy: SandboxPolicy::new_read_only_policy(), + permission_profile: None, + network: None, + file_system_sandbox_policy: None, + model: "test-model".to_string(), + comp_hash: None, + personality: None, + collaboration_mode: None, + multi_agent_version: None, + multi_agent_mode: None, + realtime_active: None, + effort: None, + summary: ReasoningSummary::Auto, + }) +} + +fn compacted(message: &str, replacement_history: Option>) -> RolloutItem { + RolloutItem::Compacted(CompactedItem { + message: message.to_string(), + replacement_history, + window_number: Some(1), + first_window_id: None, + previous_window_id: None, + window_id: None, + }) +} diff --git a/codex-rs/thread-store/src/local/read_thread.rs b/codex-rs/thread-store/src/local/read_thread.rs index 27ce3076e9..035774b04e 100644 --- a/codex-rs/thread-store/src/local/read_thread.rs +++ b/codex-rs/thread-store/src/local/read_thread.rs @@ -206,7 +206,7 @@ async fn attach_history_if_requested( Ok(()) } -async fn resolve_rollout_path( +pub(super) async fn resolve_rollout_path( store: &LocalThreadStore, thread_id: codex_protocol::ThreadId, include_archived: bool, @@ -291,7 +291,7 @@ async fn read_thread_from_rollout_path( Ok(thread) } -async fn load_history_items( +pub(super) async fn load_history_items( path: &std::path::Path, ) -> ThreadStoreResult> { let (items, _, _) = RolloutRecorder::load_rollout_items(path) diff --git a/codex-rs/thread-store/src/local/test_support.rs b/codex-rs/thread-store/src/local/test_support.rs index a930067fc4..3537555ee4 100644 --- a/codex-rs/thread-store/src/local/test_support.rs +++ b/codex-rs/thread-store/src/local/test_support.rs @@ -111,15 +111,17 @@ pub(super) fn write_session_file_with_fork( }, }); writeln!(file, "{meta}")?; - let user_event = serde_json::json!({ - "timestamp": ts, - "type": "event_msg", - "payload": { - "type": "user_message", - "message": first_user_message, - "kind": "plain", - }, - }); - writeln!(file, "{user_event}")?; + if matches!(history_mode, ThreadHistoryMode::Legacy) { + let user_event = serde_json::json!({ + "timestamp": ts, + "type": "event_msg", + "payload": { + "type": "user_message", + "message": first_user_message, + "kind": "plain", + }, + }); + writeln!(file, "{user_event}")?; + } Ok(path) } diff --git a/codex-rs/thread-store/src/local/update_thread_metadata.rs b/codex-rs/thread-store/src/local/update_thread_metadata.rs index a6ea76e7e0..800e1799c0 100644 --- a/codex-rs/thread-store/src/local/update_thread_metadata.rs +++ b/codex-rs/thread-store/src/local/update_thread_metadata.rs @@ -820,6 +820,7 @@ mod tests { ThreadHistoryMode::Paginated, ) .expect("session file"); + let original_rollout = std::fs::read_to_string(&path).expect("read rollout"); let runtime = codex_state::StateRuntime::init( home.path().to_path_buf(), config.default_model_provider_id.clone(), @@ -845,7 +846,10 @@ mod tests { } )); - assert_eq!(last_rollout_item(path.as_path())["type"], "event_msg"); + assert_eq!( + std::fs::read_to_string(&path).expect("read rollout"), + original_rollout + ); assert_eq!( runtime .get_thread_memory_mode(thread_id) diff --git a/codex-rs/thread-store/src/store.rs b/codex-rs/thread-store/src/store.rs index 67ee603f49..67834ef93d 100644 --- a/codex-rs/thread-store/src/store.rs +++ b/codex-rs/thread-store/src/store.rs @@ -17,6 +17,7 @@ use crate::ReadThreadByRolloutPathParams; use crate::ReadThreadParams; use crate::ResumeThreadParams; use crate::SearchThreadsParams; +use crate::StoredModelContext; use crate::StoredThread; use crate::StoredThreadHistory; use crate::ThreadPage; @@ -76,6 +77,20 @@ pub trait ThreadStore: Any + Send + Sync { params: LoadThreadHistoryParams, ) -> ThreadStoreFuture<'_, StoredThreadHistory>; + /// Loads the persisted rollout items needed to reconstruct the latest model-visible context. + /// + /// Implementations that cannot perform a targeted read may return the full persisted history. + fn load_latest_model_context( + &self, + _params: LoadThreadHistoryParams, + ) -> ThreadStoreFuture<'_, StoredModelContext> { + Box::pin(async { + Err(ThreadStoreError::Unsupported { + operation: "load_latest_model_context", + }) + }) + } + /// Reads a thread summary and optionally its persisted history. fn read_thread(&self, params: ReadThreadParams) -> ThreadStoreFuture<'_, StoredThread>; diff --git a/codex-rs/thread-store/src/types.rs b/codex-rs/thread-store/src/types.rs index 317170442c..74f21872ef 100644 --- a/codex-rs/thread-store/src/types.rs +++ b/codex-rs/thread-store/src/types.rs @@ -159,6 +159,19 @@ pub struct StoredThreadHistory { pub items: Vec, } +/// Persisted rollout items needed to reconstruct the latest model-visible context. +/// +/// Local stores may return only a resumable suffix while stores without targeted reads may return +/// the full persisted history. In either case, `items` remain in replay order and are suitable for +/// the existing rollout reconstruction path. +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct StoredModelContext { + /// Thread id represented by the model context. + pub thread_id: ThreadId, + /// Persisted rollout items in replay order. + pub items: Vec, +} + /// Parameters for reading a thread summary and optionally its replay history. #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] pub struct ReadThreadParams {