diff --git a/codex-rs/thread-store/src/in_memory.rs b/codex-rs/thread-store/src/in_memory.rs index cb3f198ff3..49ce5a7a26 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, @@ -517,6 +519,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; @@ -655,6 +676,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 5939ea606e..cd17d9bbf6 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 read_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 { read_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/read_model_context.rs b/codex-rs/thread-store/src/local/read_model_context.rs new file mode 100644 index 0000000000..49bbce5489 --- /dev/null +++ b/codex-rs/thread-store/src/local/read_model_context.rs @@ -0,0 +1,298 @@ +use std::fs::File; +use std::io; +use std::io::Read; +use std::io::Seek; +use std::io::SeekFrom; +use std::path::Path; +use std::path::PathBuf; + +use codex_protocol::models::ResponseItem; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::RolloutItem; +use codex_protocol::protocol::RolloutLine; +use codex_protocol::protocol::SessionMetaLine; +use codex_protocol::protocol::ThreadHistoryMode; +use tracing::debug; + +use super::LocalThreadStore; +use super::helpers::rollout_path_is_archived; +use super::read_thread; +use crate::LoadThreadHistoryParams; +use crate::StoredModelContext; +use crate::ThreadStoreError; +use crate::ThreadStoreResult; + +#[cfg(test)] +#[path = "read_model_context_tests.rs"] +mod tests; + +const READ_CHUNK_SIZE: usize = 64 * 1024; + +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), + })?; + if !params.include_archived + && rollout_path_is_archived(store.config.codex_home.as_path(), path.as_path()) + { + return Err(ThreadStoreError::InvalidRequest { + message: format!("thread {} is archived", 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")) + { + match scan_bounded_model_context(path.clone(), session_meta.clone()).await? { + Some(items) => items, + None => { + debug!( + thread_id = %params.thread_id, + rollout_path = %path.display(), + "falling back to full rollout load for model context" + ); + read_thread::load_history_items(path.as_path()).await? + } + } + } else { + read_thread::load_history_items(path.as_path()).await? + }; + + Ok(StoredModelContext { + thread_id: params.thread_id, + items, + }) +} + +async fn scan_bounded_model_context( + path: PathBuf, + session_meta: SessionMetaLine, +) -> ThreadStoreResult>> { + let path_for_error = path.clone(); + tokio::task::spawn_blocking(move || scan_bounded_model_context_blocking(&path, session_meta)) + .await + .map_err(|err| ThreadStoreError::Internal { + message: format!("failed to join model context scan: {err}"), + })? + .map_err(|err| ThreadStoreError::Internal { + message: format!( + "failed to scan model context {}: {err}", + path_for_error.display() + ), + }) +} + +fn scan_bounded_model_context_blocking( + path: &Path, + session_meta: SessionMetaLine, +) -> io::Result>> { + let mut selector = ModelContextSelector::default(); + let mut items_newest_first = Vec::new(); + scan_rollout_from_end(path, |item| { + let selection = selector.observe(&item); + items_newest_first.push(item); + Ok(selection) + })?; + + if !selector.is_complete() { + return Ok(None); + } + + items_newest_first.reverse(); + // The head SessionMeta is canonical even when copied fork history contains later metadata. + // A successful bounded scan stops at a turn boundary after the head, so this does not + // duplicate the rollout's own first line. + items_newest_first.insert(0, RolloutItem::SessionMeta(session_meta)); + Ok(Some(items_newest_first)) +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum ScanControl { + Continue, + Stop, + Fallback, +} + +#[derive(Debug, Default)] +struct ModelContextSelector { + saw_checkpoint: bool, + saw_resume_metadata: bool, + active_segment: ActiveSegment, + fallback: bool, +} + +impl ModelContextSelector { + fn observe(&mut self, item: &RolloutItem) -> ScanControl { + match item { + RolloutItem::Compacted(compacted) => { + if compacted.replacement_history.is_none() || compacted.window_number.is_none() { + self.fallback = true; + return ScanControl::Fallback; + } + if self.saw_checkpoint { + // A second checkpoint before a usable turn boundary means the selector cannot + // prove that the first checkpoint belongs to a surviving replay segment. + self.fallback = true; + return ScanControl::Fallback; + } + self.saw_checkpoint = true; + } + RolloutItem::EventMsg(EventMsg::ThreadRolledBack(_)) => { + // Paginated threads reject rollback. Keep old rollouts correct rather than + // duplicating rollback survival semantics in this storage-only selector. + self.fallback = true; + return ScanControl::Fallback; + } + 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 |= + matches!(response_item, ResponseItem::Message { role, .. } if role == "user"); + } + 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.is_complete() { + ScanControl::Stop + } else { + ScanControl::Continue + } + } + + fn finalize_active_segment(&mut self) { + if self.active_segment.has_user_turn && self.active_segment.has_turn_context { + self.saw_resume_metadata = true; + } + self.active_segment = ActiveSegment::default(); + } + + fn is_complete(&self) -> bool { + !self.fallback && self.saw_checkpoint && self.saw_resume_metadata + } +} + +#[derive(Debug, Default)] +struct ActiveSegment { + 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 scan_rollout_from_end( + path: &Path, + mut visit_item: impl FnMut(RolloutItem) -> io::Result, +) -> io::Result<()> { + let mut file = File::open(path)?; + let mut remaining = file.metadata()?.len(); + let mut line_reversed = Vec::new(); + let mut buffer = vec![0u8; READ_CHUNK_SIZE]; + + while remaining > 0 { + let read_size = + usize::try_from(remaining.min(READ_CHUNK_SIZE as u64)).map_err(io::Error::other)?; + remaining -= read_size as u64; + file.seek(SeekFrom::Start(remaining))?; + file.read_exact(&mut buffer[..read_size])?; + + for &byte in buffer[..read_size].iter().rev() { + if byte == b'\n' { + if process_reversed_line(&mut line_reversed, &mut visit_item)? { + return Ok(()); + } + } else { + line_reversed.push(byte); + } + } + } + + let _ = process_reversed_line(&mut line_reversed, &mut visit_item)?; + Ok(()) +} + +fn process_reversed_line( + line_reversed: &mut Vec, + visit_item: &mut impl FnMut(RolloutItem) -> io::Result, +) -> io::Result { + if line_reversed.is_empty() { + return Ok(false); + } + line_reversed.reverse(); + let line = std::str::from_utf8(line_reversed) + .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?; + let parsed = serde_json::from_str::(line.trim()); + line_reversed.clear(); + let Ok(line) = parsed else { + return Ok(false); + }; + Ok(matches!( + visit_item(line.item)?, + ScanControl::Stop | ScanControl::Fallback + )) +} diff --git a/codex-rs/thread-store/src/local/read_model_context_tests.rs b/codex-rs/thread-store/src/local/read_model_context_tests.rs new file mode 100644 index 0000000000..7ae79ea23c --- /dev/null +++ b/codex-rs/thread-store/src/local/read_model_context_tests.rs @@ -0,0 +1,242 @@ +use std::fs::OpenOptions; +use std::io::Write; +use std::path::Path; + +use codex_protocol::config_types::ReasoningSummary; +use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::CompactedItem; +use codex_protocol::protocol::EventMsg; +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::protocol::UserMessageEvent; +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(1001); + let thread_id = codex_protocol::ThreadId::from_string(&uuid.to_string()).expect("thread id"); + let path = write_session_file_with_history_mode( + home.path(), + "2025-01-03T13-00-00", + uuid, + ThreadHistoryMode::Paginated, + ) + .expect("write session file"); + append_items( + path.as_path(), + [ + turn_started("turn-1"), + user_message("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"), + turn_context(home.path(), "turn-2"), + compacted("latest checkpoint", Some(Vec::new())), + turn_complete("turn-2"), + ], + ); + let store = LocalThreadStore::new(test_config(home.path()), None); + + let context = store + .load_latest_model_context(LoadThreadHistoryParams { + thread_id, + include_archived: false, + }) + .await + .expect("load model context"); + + assert_eq!(context.thread_id, thread_id); + 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 falls_back_to_full_history_for_compaction_without_replacement_history() { + let home = TempDir::new().expect("temp dir"); + let uuid = Uuid::from_u128(1002); + let thread_id = codex_protocol::ThreadId::from_string(&uuid.to_string()).expect("thread id"); + let path = write_session_file_with_history_mode( + home.path(), + "2025-01-03T13-00-01", + uuid, + ThreadHistoryMode::Paginated, + ) + .expect("write session file"); + append_items( + path.as_path(), + [ + turn_started("turn-1"), + user_message("turn"), + turn_context(home.path(), "turn-1"), + compacted("usable checkpoint", Some(Vec::new())), + compacted("legacy checkpoint", None), + turn_complete("turn-1"), + ], + ); + let full_items = read_thread::load_history_items(path.as_path()) + .await + .expect("load full history"); + let store = LocalThreadStore::new(test_config(home.path()), None); + + let context = store + .load_latest_model_context(LoadThreadHistoryParams { + thread_id, + include_archived: false, + }) + .await + .expect("load model context"); + + assert_eq!(context.items.len(), full_items.len()); + assert!(context.items.iter().any(|item| { + matches!(item, RolloutItem::Compacted(compacted) if compacted.message == "usable checkpoint") + })); +} + +#[tokio::test] +async fn ignores_malformed_tail_lines_before_selecting_checkpoint() { + let home = TempDir::new().expect("temp dir"); + let uuid = Uuid::from_u128(1003); + let thread_id = codex_protocol::ThreadId::from_string(&uuid.to_string()).expect("thread id"); + let path = write_session_file_with_history_mode( + home.path(), + "2025-01-03T13-00-02", + uuid, + ThreadHistoryMode::Paginated, + ) + .expect("write session file"); + append_items( + path.as_path(), + [ + turn_started("turn-1"), + user_message("turn"), + turn_context(home.path(), "turn-1"), + compacted("checkpoint", Some(Vec::new())), + turn_complete("turn-1"), + ], + ); + let mut file = OpenOptions::new() + .append(true) + .open(path.as_path()) + .expect("open session file"); + writeln!(file, "not-json").expect("append malformed line"); + let store = LocalThreadStore::new(test_config(home.path()), 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 == "checkpoint") + })); +} + +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(), + 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, + completed_at: None, + duration_ms: None, + time_to_first_token_ms: None, + })) +} + +fn user_message(message: &str) -> RolloutItem { + RolloutItem::EventMsg(EventMsg::UserMessage(UserMessageEvent { + message: message.to_string(), + ..Default::default() + })) +} + +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, + 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/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 4e688eca6f..e89db12200 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 {