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 30221d9442..c1062c2cdc 100644 --- a/codex-rs/app-server-protocol/src/protocol/thread_history.rs +++ b/codex-rs/app-server-protocol/src/protocol/thread_history.rs @@ -252,6 +252,14 @@ impl ThreadHistoryBuilder { } } + pub fn with_initial_rollout_index(initial_rollout_index: usize) -> Self { + Self { + current_rollout_index: initial_rollout_index, + next_rollout_index: initial_rollout_index, + ..Self::new() + } + } + pub fn reset(&mut self) { *self = Self::new(); } diff --git a/codex-rs/app-server/src/request_processors/external_agent_session_import.rs b/codex-rs/app-server/src/request_processors/external_agent_session_import.rs index d4abb155ec..3c89f44310 100644 --- a/codex-rs/app-server/src/request_processors/external_agent_session_import.rs +++ b/codex-rs/app-server/src/request_processors/external_agent_session_import.rs @@ -214,6 +214,7 @@ impl ExternalAgentSessionImporter { }, dynamic_tools: Vec::new(), multi_agent_version: Some(MultiAgentVersion::V1), + initial_rollout_copy: None, metadata: ThreadPersistenceMetadata { cwd: Some(cwd.clone()), model_provider: model_provider.clone(), 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 a3e4b49ce8..b52ed59d04 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -373,6 +373,33 @@ enum RunningThreadResumeResult { NotRunning(Option>), } +enum ResumeHistoryExtent { + Full, + CompactedCheckpoint, +} + +/// Small resume-history shape retained after the full rollout history is moved +/// into core session startup. +enum ResumeResponseHistory { + Resumed { conversation_id: ThreadId }, + Forked { preview: String }, + Missing, +} + +impl ResumeResponseHistory { + fn from_initial_history(history: &InitialHistory) -> Self { + match history { + InitialHistory::Resumed(resumed) => Self::Resumed { + conversation_id: resumed.conversation_id, + }, + InitialHistory::Forked(items) => Self::Forked { + preview: preview_from_rollout_items(items), + }, + InitialHistory::New | InitialHistory::Cleared => Self::Missing, + } + } +} + impl ThreadRequestProcessor { #[allow(clippy::too_many_arguments)] pub(crate) fn new( @@ -2348,15 +2375,14 @@ impl ThreadRequestProcessor { let thread_uuid = ThreadId::from_string(&thread_id) .map_err(|err| invalid_request(format!("invalid thread id: {err}")))?; - let items = self - .load_thread_turns_list_history(thread_uuid) - .await - .map_err(thread_read_view_error)?; + let cursor = cursor.as_deref(); + let sort_direction = sort_direction.unwrap_or(SortDirection::Desc); + let items_view = items_view.unwrap_or(TurnItemsView::Summary); + // This API optimizes network transfer by letting clients page through a - // thread's turns incrementally, but it still replays the entire rollout on - // every request. Rollback and compaction events can change earlier turns, so - // the server has to rebuild the full turn list until turn metadata is indexed - // separately. + // thread's turns incrementally. The newest notLoaded page can be served + // from a rollout suffix because it does not need item payloads or earlier + // turn item indexes. Other shapes continue to replay the whole rollout. let loaded_thread = self.thread_manager.get_thread(thread_uuid).await.ok(); let has_live_running_thread = match loaded_thread.as_ref() { Some(thread) => matches!(thread.agent_status().await, AgentStatus::Running), @@ -2372,18 +2398,44 @@ impl ThreadRequestProcessor { } else { None }; + let loaded_status = self + .thread_watch_manager + .loaded_status_for_thread(&thread_uuid.to_string()) + .await; + + if loaded_thread.is_none() + && !has_live_running_thread + && active_turn.is_none() + && cursor.is_none() + && sort_direction == SortDirection::Desc + && items_view == TurnItemsView::NotLoaded + && let Some(response) = self + .try_build_recent_not_loaded_turns_page_from_store( + thread_uuid, + loaded_status.clone(), + limit, + ) + .await? + { + return Ok(response); + } + + // Rollback and compaction events can change earlier turns, so all other + // requests rebuild the full turn list until turn metadata is indexed separately. + let items = self + .load_thread_turns_list_history(thread_uuid) + .await + .map_err(thread_read_view_error)?; build_thread_turns_page_response( &items, - self.thread_watch_manager - .loaded_status_for_thread(&thread_uuid.to_string()) - .await, + loaded_status, has_live_running_thread, active_turn, ThreadTurnsPageOptions { - cursor: cursor.as_deref(), + cursor, limit, - sort_direction: sort_direction.unwrap_or(SortDirection::Desc), - items_view: items_view.unwrap_or(TurnItemsView::Summary), + sort_direction, + items_view, }, ) } @@ -2445,6 +2497,35 @@ impl ThreadRequestProcessor { .map_err(|err| thread_turns_list_history_load_error(thread_id, err)) } + async fn try_build_recent_not_loaded_turns_page_from_store( + &self, + thread_id: ThreadId, + loaded_status: ThreadStatus, + limit: Option, + ) -> Result, JSONRPCErrorError> { + let stored_thread = match self + .thread_store + .read_thread(StoreReadThreadParams { + thread_id, + include_archived: true, + include_history: false, + }) + .await + { + Ok(stored_thread) => stored_thread, + Err(_) => return Ok(None), + }; + let Some(rollout_path) = stored_thread.rollout_path else { + return Ok(None); + }; + try_build_recent_not_loaded_turns_page_from_rollout_path( + rollout_path.as_path(), + loaded_status, + limit, + ) + .await + } + pub(crate) fn thread_created_receiver(&self) -> broadcast::Receiver { self.thread_manager.subscribe_thread_created() } @@ -2599,21 +2680,34 @@ impl ThreadRequestProcessor { initial_turns_page, } = params; let include_turns = !exclude_turns; + let prefer_compacted_history = !include_turns + && initial_turns_page + .as_ref() + .is_none_or(thread_resume_initial_page_is_recent_not_loaded); let resume_result = if let Some(history) = history { self.resume_thread_from_history(history.as_slice()) .await - .map(|thread_history| (thread_history, None)) + .map(|thread_history| (thread_history, None, ResumeHistoryExtent::Full)) } else if let Some(stored_thread) = stored_thread_from_running_probe { - self.stored_thread_to_initial_history(&stored_thread) - .await - .map(|thread_history| (thread_history, Some(*stored_thread))) + self.resume_thread_from_stored_thread( + &thread_id, + path.as_ref(), + *stored_thread, + prefer_compacted_history, + ) + .await + .map(|(thread_history, stored_thread, history_extent)| { + (thread_history, Some(stored_thread), history_extent) + }) } else { - self.resume_thread_from_rollout(&thread_id, path.as_ref()) + self.resume_thread_from_rollout(&thread_id, path.as_ref(), prefer_compacted_history) .await - .map(|(thread_history, stored_thread)| (thread_history, Some(stored_thread))) + .map(|(thread_history, stored_thread, history_extent)| { + (thread_history, Some(stored_thread), history_extent) + }) }; - let (thread_history, resume_source_thread) = match resume_result { + let (thread_history, resume_source_thread, history_extent) = match resume_result { Ok(value) => value, Err(error) => { self.outgoing.send_error(request_id, error).await; @@ -2658,7 +2752,54 @@ impl ThreadRequestProcessor { } }; - let response_history = thread_history.clone(); + let response_history = ResumeResponseHistory::from_initial_history(&thread_history); + let mut initial_turns_recent_items = None; + if !include_turns + && let Some(params) = initial_turns_page.as_ref() + && thread_resume_initial_page_is_recent_not_loaded(params) + && let InitialHistory::Resumed(resumed) = &thread_history + && let Some(rollout_path) = resumed.rollout_path.as_ref() + { + match try_read_recent_not_loaded_rollout_items_from_rollout_path( + rollout_path.as_path(), + params.limit, + ) + .await + { + Ok(recent) => initial_turns_recent_items = recent, + Err(error) => { + self.outgoing.send_error(request_id, error).await; + return Ok(()); + } + } + } + let response_history_items = (include_turns + || (initial_turns_page.is_some() && initial_turns_recent_items.is_none())) + .then(|| { + if matches!(history_extent, ResumeHistoryExtent::CompactedCheckpoint) { + None + } else { + Some(thread_history.get_rollout_items()) + } + }) + .flatten(); + let response_history_items = if response_history_items.is_none() + && initial_turns_page.is_some() + && initial_turns_recent_items.is_none() + { + match self + .load_full_resume_history_items(&thread_id, path.as_ref()) + .await + { + Ok(items) => Some(items), + Err(error) => { + self.outgoing.send_error(request_id, error).await; + return Ok(()); + } + } + } else { + response_history_items + }; match self .thread_manager @@ -2695,24 +2836,13 @@ impl ThreadRequestProcessor { self.outgoing.send_error(request_id, error).await; return Ok(()); }; - // Auto-attach a thread listener when resuming a thread. - log_listener_attach_result( - self.ensure_conversation_listener( - thread_id, - request_id.connection_id, - /*raw_events_enabled*/ false, - ) - .await, - thread_id, - request_id.connection_id, - "thread", - ); let mut thread = match self .load_thread_from_resume_source_or_send_internal( thread_id, codex_thread.as_ref(), &response_history, + response_history_items.as_deref(), rollout_path.as_path(), resume_source_thread, include_turns, @@ -2757,17 +2887,33 @@ impl ThreadRequestProcessor { ); let token_usage_thread = include_turns.then(|| thread.clone()); let mut initial_turns_page = if let Some(params) = initial_turns_page.as_ref() { - match build_thread_resume_initial_turns_page( - &response_history.get_rollout_items(), - thread.status.clone(), - /*has_live_running_thread*/ false, - /*active_turn*/ None, - params, - ) { - Ok(page) => Some(page), - Err(error) => { - self.outgoing.send_error(request_id, error).await; - return Ok(()); + if let Some(recent) = initial_turns_recent_items.take() { + match build_recent_not_loaded_turns_page_response( + recent, + thread.status.clone(), + params.limit, + ) { + Ok(page) => Some(page.into()), + Err(error) => { + self.outgoing.send_error(request_id, error).await; + return Ok(()); + } + } + } else { + match build_thread_resume_initial_turns_page( + response_history_items.as_deref().expect( + "response history items loaded for resume initial turns page", + ), + thread.status.clone(), + /*has_live_running_thread*/ false, + /*active_turn*/ None, + params, + ) { + Ok(page) => Some(page), + Err(error) => { + self.outgoing.send_error(request_id, error).await; + return Ok(()); + } } } } else { @@ -2803,7 +2949,9 @@ impl ThreadRequestProcessor { // rebuilding history only to attribute a replayed usage update. if let Some(token_usage_thread) = token_usage_thread { let token_usage_turn_id = latest_token_usage_turn_id_from_rollout_items( - &response_history.get_rollout_items(), + response_history_items + .as_deref() + .expect("response history items loaded for resume token usage update"), token_usage_thread.turns.as_slice(), ); // The client needs restored usage before it starts another turn. @@ -2819,6 +2967,21 @@ impl ThreadRequestProcessor { ) .await; } + tokio::task::yield_now().await; + // Auto-attach after the response is queued. The attach is + // best-effort logging, and idle resumes do not need it to + // construct the response. + log_listener_attach_result( + self.ensure_conversation_listener( + thread_id, + connection_id, + /*raw_events_enabled*/ false, + ) + .await, + thread_id, + connection_id, + "thread", + ); self.thread_goal_processor .emit_resume_goal_snapshot_and_continue(thread_id, codex_thread.as_ref()) .await; @@ -2883,16 +3046,33 @@ impl ThreadRequestProcessor { .await?; Some((existing_thread_id, existing_thread, source_thread)) } else { + let cold_resume_needs_full_history = !params.exclude_turns + || params + .initial_turns_page + .as_ref() + .is_some_and(|params| !thread_resume_initial_page_is_recent_not_loaded(params)); let source_thread = self .read_stored_thread_for_resume( ¶ms.thread_id, params.path.as_ref(), - /*include_history*/ true, + cold_resume_needs_full_history, ) .await?; let existing_thread_id = source_thread.thread_id; match self.thread_manager.get_thread(existing_thread_id).await { - Ok(existing_thread) => Some((existing_thread_id, existing_thread, source_thread)), + Ok(existing_thread) => { + let source_thread = if source_thread.history.is_some() { + source_thread + } else { + self.read_stored_thread_for_resume( + ¶ms.thread_id, + params.path.as_ref(), + /*include_history*/ true, + ) + .await? + }; + Some((existing_thread_id, existing_thread, source_thread)) + } Err(_) => { return Ok(RunningThreadResumeResult::NotRunning(Some(Box::new( source_thread, @@ -3059,14 +3239,80 @@ impl ThreadRequestProcessor { &self, thread_id: &str, path: Option<&PathBuf>, - ) -> Result<(InitialHistory, StoredThread), JSONRPCErrorError> { + prefer_compacted_history: bool, + ) -> Result<(InitialHistory, StoredThread, ResumeHistoryExtent), JSONRPCErrorError> { let stored_thread = self - .read_stored_thread_for_resume(thread_id, path, /*include_history*/ true) + .read_stored_thread_for_resume( + thread_id, + path, + /*include_history*/ !prefer_compacted_history, + ) .await?; - let history = self - .stored_thread_to_initial_history(&stored_thread) + self.resume_thread_from_stored_thread( + thread_id, + path, + stored_thread, + prefer_compacted_history, + ) + .await + } + + async fn resume_thread_from_stored_thread( + &self, + thread_id: &str, + path: Option<&PathBuf>, + mut stored_thread: StoredThread, + prefer_compacted_history: bool, + ) -> Result<(InitialHistory, StoredThread, ResumeHistoryExtent), JSONRPCErrorError> { + let mut response_thread = self + .read_stored_thread_for_resume(thread_id, path, /*include_history*/ false) .await?; - Ok((history, stored_thread)) + response_thread.history = None; + let history_extent = if stored_thread.history.is_some() { + ResumeHistoryExtent::Full + } else if prefer_compacted_history { + match stored_thread.rollout_path.clone() { + Some(rollout_path) => { + match try_read_compacted_resume_rollout_items_from_rollout_path( + rollout_path.as_path(), + stored_thread.thread_id, + ) + .await? + { + Some(items) => { + stored_thread.history = Some(codex_thread_store::StoredThreadHistory { + thread_id: stored_thread.thread_id, + items, + }); + ResumeHistoryExtent::CompactedCheckpoint + } + None => { + stored_thread = self + .read_stored_thread_for_resume( + thread_id, path, /*include_history*/ true, + ) + .await?; + ResumeHistoryExtent::Full + } + } + } + None => { + stored_thread = self + .read_stored_thread_for_resume( + thread_id, path, /*include_history*/ true, + ) + .await?; + ResumeHistoryExtent::Full + } + } + } else { + stored_thread = self + .read_stored_thread_for_resume(thread_id, path, /*include_history*/ true) + .await?; + ResumeHistoryExtent::Full + }; + let history = self.take_stored_thread_initial_history(&mut stored_thread)?; + Ok((history, response_thread, history_extent)) } async fn read_stored_thread_for_resume( @@ -3109,16 +3355,62 @@ impl ThreadRequestProcessor { Ok(stored_thread) } - #[tracing::instrument(level = "trace", skip_all)] - async fn stored_thread_to_initial_history( + async fn read_stored_thread_for_resume_with_compacted_history( &self, - stored_thread: &StoredThread, + thread_id: &str, + path: Option<&PathBuf>, + ) -> Result, JSONRPCErrorError> { + let mut stored_thread = self + .read_stored_thread_for_resume(thread_id, path, /*include_history*/ false) + .await?; + let Some(rollout_path) = stored_thread.rollout_path.clone() else { + return Ok(None); + }; + let Some(items) = try_read_compacted_resume_rollout_items_from_rollout_path( + rollout_path.as_path(), + stored_thread.thread_id, + ) + .await? + else { + return Ok(None); + }; + stored_thread.history = Some(codex_thread_store::StoredThreadHistory { + thread_id: stored_thread.thread_id, + items, + }); + Ok(Some(stored_thread)) + } + + async fn load_full_resume_history_items( + &self, + thread_id: &str, + path: Option<&PathBuf>, + ) -> Result, JSONRPCErrorError> { + let mut stored_thread = self + .read_stored_thread_for_resume(thread_id, path, /*include_history*/ true) + .await?; + stored_thread + .history + .take() + .map(|history| history.items) + .ok_or_else(|| { + internal_error(format!( + "thread {} did not include persisted history", + stored_thread.thread_id + )) + }) + } + + #[tracing::instrument(level = "trace", skip_all)] + fn take_stored_thread_initial_history( + &self, + stored_thread: &mut StoredThread, ) -> Result { let thread_id = stored_thread.thread_id; let history = stored_thread .history - .as_ref() - .map(|history| history.items.clone()) + .take() + .map(|history| history.items) .ok_or_else(|| { internal_error(format!( "thread {thread_id} did not include persisted history" @@ -3168,98 +3460,84 @@ impl ThreadRequestProcessor { &self, thread_id: ThreadId, thread: &CodexThread, - thread_history: &InitialHistory, + response_history: &ResumeResponseHistory, + response_history_items: Option<&[RolloutItem]>, rollout_path: &Path, resume_source_thread: Option, include_turns: bool, ) -> std::result::Result { let config_snapshot = thread.config_snapshot().await; let session_id = thread.session_configured().session_id.to_string(); - let thread = match thread_history { - InitialHistory::Resumed(resumed) => { + let (mut thread, name_loaded_from_store) = match response_history { + ResumeResponseHistory::Resumed { conversation_id } => { let fallback_provider = config_snapshot.model_provider_id.as_str(); - if let Some(stored_thread) = resume_source_thread { - let stored_thread = - if let Some(rollout_path) = stored_thread.rollout_path.clone() { - self.thread_store - .read_thread_by_rollout_path(StoreReadThreadByRolloutPathParams { - rollout_path, - include_archived: true, - include_history: false, - }) - .await - .unwrap_or(StoredThread { - history: None, - ..stored_thread - }) - } else { - self.thread_store - .read_thread(StoreReadThreadParams { - thread_id: stored_thread.thread_id, - include_archived: true, - include_history: false, - }) - .await - .unwrap_or(StoredThread { - history: None, - ..stored_thread - }) - }; - Ok(thread_from_stored_thread( - stored_thread, - fallback_provider, - &self.config.cwd, - ) - .0) + if let Some(mut stored_thread) = resume_source_thread { + stored_thread.history = None; + Ok(( + thread_from_stored_thread( + stored_thread, + fallback_provider, + &self.config.cwd, + ) + .0, + true, + )) } else { match self .thread_store .read_thread(StoreReadThreadParams { - thread_id: resumed.conversation_id, + thread_id: *conversation_id, include_archived: true, include_history: false, }) .await { - Ok(stored_thread) => Ok(thread_from_stored_thread( - stored_thread, - fallback_provider, - &self.config.cwd, - ) - .0), + Ok(stored_thread) => Ok(( + thread_from_stored_thread( + stored_thread, + fallback_provider, + &self.config.cwd, + ) + .0, + true, + )), Err(read_err) => { Err(format!("failed to read thread from store: {read_err}")) } } } } - InitialHistory::Forked(items) => { + ResumeResponseHistory::Forked { preview } => { let mut thread = build_thread_from_snapshot( thread_id, session_id.clone(), &config_snapshot, Some(rollout_path.into()), ); - thread.preview = preview_from_rollout_items(items); - Ok(thread) + thread.preview = preview.clone(); + Ok((thread, false)) } - InitialHistory::New | InitialHistory::Cleared => Err(format!( + ResumeResponseHistory::Missing => Err(format!( "failed to build resume response for thread {thread_id}: initial history missing" )), - }; - let mut thread = thread?; + }?; thread.id = thread_id.to_string(); thread.session_id = session_id; thread.path = Some(rollout_path.to_path_buf()); if include_turns { - let history_items = thread_history.get_rollout_items(); + let history_items = response_history_items.ok_or_else(|| { + "failed to build resume response: history items missing for inline turns" + .to_string() + })?; populate_thread_turns_from_history( &mut thread, - &history_items, + history_items, /*active_turn*/ None, ); } - self.attach_thread_name(thread_id, &mut thread).await; + if !name_loaded_from_store { + self.attach_thread_name(thread_id, &mut thread).await; + } Ok(thread) } @@ -3313,18 +3591,57 @@ impl ThreadRequestProcessor { "`permissions` cannot be combined with `sandbox`", )); } - let source_thread = self - .read_stored_thread_for_resume(&thread_id, path.as_ref(), /*include_history*/ true) - .await?; + let optimized_source = if !include_turns && !ephemeral { + match self + .read_stored_thread_for_resume_with_compacted_history(&thread_id, path.as_ref()) + .await? + { + Some(stored_thread) => stored_thread + .rollout_path + .clone() + .map(|rollout_path| (stored_thread, rollout_path)), + None => None, + } + } else { + None + }; + let (mut source_thread, initial_rollout_copy) = match optimized_source { + Some((source_thread, rollout_path)) => (source_thread, Some(rollout_path)), + None => ( + self.read_stored_thread_for_resume( + &thread_id, + path.as_ref(), + /*include_history*/ true, + ) + .await?, + None, + ), + }; let source_thread_id = source_thread.thread_id; let source_thread_name = source_thread .name .as_deref() .and_then(codex_core::util::normalize_thread_name); + let source_thread_preview = source_thread.preview.clone(); + let response_history_items = if include_turns || ephemeral { + Some( + source_thread + .history + .as_ref() + .map(|history| history.items.clone()) + .ok_or_else(|| { + internal_error(format!( + "thread {source_thread_id} did not include persisted history" + )) + })?, + ) + } else { + None + }; let history_items = source_thread .history - .as_ref() - .map(|history| history.items.clone()) + .take() + .map(|history| history.items) .ok_or_else(|| { internal_error(format!( "thread {source_thread_id} did not include persisted history" @@ -3379,33 +3696,50 @@ impl ThreadRequestProcessor { let fallback_model_provider = config.model_provider_id.clone(); + let initial_history = InitialHistory::Resumed(ResumedHistory { + conversation_id: source_thread_id, + history: history_items, + rollout_path: source_thread.rollout_path.clone(), + }); + let parent_trace = self.request_trace_context(&request_id).await; + let thread_source = thread_source.map(Into::into); + let fork_result = if let Some(initial_rollout_copy) = initial_rollout_copy { + self.thread_manager + .fork_thread_from_history_with_initial_rollout_copy( + ForkSnapshot::Interrupted, + config, + initial_history, + thread_source, + parent_trace, + supports_openai_form_elicitation, + initial_rollout_copy, + ) + .await + } else { + self.thread_manager + .fork_thread_from_history( + ForkSnapshot::Interrupted, + config, + initial_history, + thread_source, + parent_trace, + supports_openai_form_elicitation, + ) + .await + }; + let NewThread { thread_id, thread: forked_thread, session_configured, .. - } = self - .thread_manager - .fork_thread_from_history( - ForkSnapshot::Interrupted, - config, - InitialHistory::Resumed(ResumedHistory { - conversation_id: source_thread_id, - history: history_items.clone(), - rollout_path: source_thread.rollout_path.clone(), - }), - thread_source.map(Into::into), - self.request_trace_context(&request_id).await, - supports_openai_form_elicitation, - ) - .await - .map_err(|err| match err { - CodexErr::Io(_) | CodexErr::Json(_) => { - invalid_request(format!("failed to load thread {source_thread_id}: {err}")) - } - CodexErr::InvalidRequest(message) => invalid_request(message), - err => internal_error(format!("error forking thread: {err}")), - })?; + } = fork_result.map_err(|err| match err { + CodexErr::Io(_) | CodexErr::Json(_) => { + invalid_request(format!("failed to load thread {source_thread_id}: {err}")) + } + CodexErr::InvalidRequest(message) => invalid_request(message), + err => internal_error(format!("error forking thread: {err}")), + })?; Self::set_app_server_client_info( forked_thread.as_ref(), @@ -3431,31 +3765,38 @@ impl ThreadRequestProcessor { let instruction_sources = forked_thread.legacy_instruction_sources().await; - // Auto-attach a conversation listener when forking a thread. - log_listener_attach_result( - self.ensure_conversation_listener( - thread_id, - request_id.connection_id, - /*raw_events_enabled*/ false, - ) - .await, - thread_id, - request_id.connection_id, - "thread", - ); - - // Persistent forks materialize their own rollout immediately. Ephemeral forks stay - // pathless, so they rebuild their visible history from the copied source history instead. - let mut thread = if session_configured.rollout_path.is_some() { - let stored_thread = self - .read_stored_thread_for_new_fork(thread_id, include_turns) - .await?; - self.stored_thread_to_api_thread( - stored_thread, - fallback_model_provider.as_str(), - include_turns, - ) + // Persistent forks materialize their own rollout immediately. When inline turns are + // excluded, response metadata can be built from the live snapshot without rereading the + // child rollout that startup just wrote. Ephemeral forks stay pathless, so they rebuild + // their visible history from the copied source history instead. + let mut thread = if let Some(rollout_path) = session_configured.rollout_path.clone() { + if include_turns { + let stored_thread = self + .read_stored_thread_for_new_fork(thread_id, include_turns) + .await?; + self.stored_thread_to_api_thread( + stored_thread, + fallback_model_provider.as_str(), + include_turns, + ) + } else { + let config_snapshot = forked_thread.config_snapshot().await; + let mut thread = build_thread_from_snapshot( + thread_id, + session_configured.session_id.to_string(), + &config_snapshot, + Some(rollout_path), + ); + thread.preview = source_thread_preview.clone(); + thread.forked_from_id = Some(source_thread_id.to_string()); + thread + } } else { + let history_items = response_history_items.as_ref().ok_or_else(|| { + internal_error(format!( + "thread {source_thread_id} did not retain response history for ephemeral fork" + )) + })?; let config_snapshot = forked_thread.config_snapshot().await; let mut thread = build_thread_from_snapshot( thread_id, @@ -3463,12 +3804,12 @@ impl ThreadRequestProcessor { &config_snapshot, /*path*/ None, ); - thread.preview = preview_from_rollout_items(&history_items); + thread.preview = preview_from_rollout_items(history_items); thread.forked_from_id = Some(source_thread_id.to_string()); if include_turns { populate_thread_turns_from_history( &mut thread, - &history_items, + history_items, /*active_turn*/ None, ); } @@ -3525,8 +3866,13 @@ impl ThreadRequestProcessor { // `excludeTurns` is the cheap fork path, so skip restored usage replay // instead of rebuilding history only to attribute a historical update. if let Some(token_usage_thread) = token_usage_thread { + let history_items = response_history_items.as_deref().ok_or_else(|| { + internal_error(format!( + "thread {source_thread_id} did not retain response history for token usage replay" + )) + })?; let token_usage_turn_id = latest_token_usage_turn_id_from_rollout_items( - &history_items, + history_items, token_usage_thread.turns.as_slice(), ); // Mirror the resume contract for forks: the new thread is usable as soon @@ -3542,6 +3888,22 @@ impl ThreadRequestProcessor { .await; } + tokio::task::yield_now().await; + // Auto-attach after the response is queued. The attach is still + // performed before the follow-up notification, but it no longer sits on + // the fork response latency path. + log_listener_attach_result( + self.ensure_conversation_listener( + thread_id, + connection_id, + /*raw_events_enabled*/ false, + ) + .await, + thread_id, + connection_id, + "thread", + ); + self.outgoing .send_server_notification(ServerNotification::ThreadStarted(notif)) .await; @@ -3742,6 +4104,379 @@ struct ThreadTurnsPage { pub(super) backwards_cursor: Option, } +struct RecentRolloutItems { + initial_rollout_index: usize, + items: Vec, +} + +fn thread_resume_initial_page_is_recent_not_loaded( + params: &ThreadResumeInitialTurnsPageParams, +) -> bool { + params.sort_direction.unwrap_or(SortDirection::Desc) == SortDirection::Desc + && params.items_view.unwrap_or(TurnItemsView::Summary) == TurnItemsView::NotLoaded +} + +async fn try_read_recent_not_loaded_rollout_items_from_rollout_path( + rollout_path: &Path, + limit: Option, +) -> Result, JSONRPCErrorError> { + let Some(existing_path) = codex_rollout::existing_rollout_path(rollout_path).await else { + return Ok(None); + }; + if !is_plain_jsonl_rollout_path(existing_path.as_path()) { + return Ok(None); + } + let page_size = thread_turns_page_size(limit); + let recent = tokio::task::spawn_blocking(move || { + read_recent_not_loaded_rollout_items(existing_path.as_path(), page_size) + }) + .await + .map_err(|err| internal_error(format!("failed to join rollout suffix read: {err}")))?; + match recent { + Ok(recent) => Ok(recent), + Err(_) => Ok(None), + } +} + +async fn try_build_recent_not_loaded_turns_page_from_rollout_path( + rollout_path: &Path, + loaded_status: ThreadStatus, + limit: Option, +) -> Result, JSONRPCErrorError> { + let Some(recent) = + try_read_recent_not_loaded_rollout_items_from_rollout_path(rollout_path, limit).await? + else { + return Ok(None); + }; + build_recent_not_loaded_turns_page_response(recent, loaded_status, limit).map(Some) +} + +fn is_plain_jsonl_rollout_path(path: &Path) -> bool { + path.file_name() + .and_then(|name| name.to_str()) + .is_some_and(|name| name.ends_with(".jsonl")) +} + +fn rollout_line_has_fork_parent_rollout_ref(line: &str) -> bool { + let Ok(value) = serde_json::from_str::(line) else { + return false; + }; + let payload = value.get("payload").unwrap_or(&value); + payload + .get("fork_parent_rollout_path") + .and_then(serde_json::Value::as_str) + .is_some_and(|path| !path.is_empty()) + && payload + .get("fork_parent_rollout_byte_len") + .and_then(serde_json::Value::as_u64) + .is_some_and(|byte_len| byte_len > 0) +} + +async fn try_read_compacted_resume_rollout_items_from_rollout_path( + rollout_path: &Path, + expected_thread_id: ThreadId, +) -> Result>, JSONRPCErrorError> { + let Some(existing_path) = codex_rollout::existing_rollout_path(rollout_path).await else { + return Ok(None); + }; + if !is_plain_jsonl_rollout_path(existing_path.as_path()) { + return Ok(None); + } + let items = tokio::task::spawn_blocking(move || { + read_compacted_resume_rollout_items(existing_path.as_path(), expected_thread_id) + }) + .await + .map_err(|err| internal_error(format!("failed to join compacted rollout read: {err}")))?; + match items { + Ok(items) => Ok(items), + Err(_) => Ok(None), + } +} + +fn read_compacted_resume_rollout_items( + path: &Path, + expected_thread_id: ThreadId, +) -> std::io::Result>> { + let Some(mut prefix_items) = read_compacted_resume_prefix_items(path, expected_thread_id)? + else { + return Ok(None); + }; + let Some(mut suffix_items) = read_compacted_resume_suffix_items(path)? else { + return Ok(None); + }; + + suffix_items.reverse(); + prefix_items.reserve(suffix_items.len()); + prefix_items.extend(suffix_items); + Ok(Some(prefix_items)) +} + +fn read_compacted_resume_prefix_items( + path: &Path, + expected_thread_id: ThreadId, +) -> std::io::Result>> { + use std::io::BufRead; + + let file = std::fs::File::open(path)?; + let reader = std::io::BufReader::new(file); + let mut prefix_items = Vec::with_capacity(3); + let mut saw_session_meta = false; + let mut saw_user_event = false; + let mut saw_user_response_item = false; + let mut checked_first_record_for_fork_parent = false; + for line in reader.lines() { + let line = line?; + let trimmed = line.trim(); + if trimmed.is_empty() { + continue; + } + if !checked_first_record_for_fork_parent { + checked_first_record_for_fork_parent = true; + if rollout_line_has_fork_parent_rollout_ref(trimmed) { + return Ok(None); + } + } + let rollout_line = + match serde_json::from_str::(trimmed) { + Ok(rollout_line) => rollout_line, + Err(_) => return Ok(None), + }; + match &rollout_line.item { + RolloutItem::SessionMeta(meta_line) if meta_line.meta.id == expected_thread_id + && !saw_session_meta => { + saw_session_meta = true; + prefix_items.push(rollout_line.item); + } + RolloutItem::SessionMeta(_) => return Ok(None), + RolloutItem::EventMsg(EventMsg::UserMessage(_)) if !saw_user_event => { + saw_user_event = true; + prefix_items.push(rollout_line.item); + } + RolloutItem::ResponseItem(codex_protocol::models::ResponseItem::Message { + role, + .. + }) if role == "user" && !saw_user_response_item => { + saw_user_response_item = true; + prefix_items.push(rollout_line.item); + } + item if rollout_item_needs_full_resume_load(item) => return Ok(None), + _ => {} + } + if saw_session_meta && saw_user_event && saw_user_response_item { + return Ok(Some(prefix_items)); + } + } + Ok(saw_session_meta.then_some(prefix_items)) +} + +fn read_compacted_resume_suffix_items(path: &Path) -> std::io::Result>> { + let mut suffix_items = Vec::new(); + let mut found_checkpoint = false; + let mut should_fallback = false; + + read_rollout_lines_from_end(path, |line| { + let trimmed = line.trim(); + if trimmed.is_empty() { + return Ok(true); + } + let rollout_line = + match serde_json::from_str::(trimmed) { + Ok(rollout_line) => rollout_line, + Err(_) => { + should_fallback = true; + return Ok(false); + } + }; + if rollout_item_needs_full_resume_load(&rollout_line.item) { + should_fallback = true; + return Ok(false); + } + let is_checkpoint = rollout_item_is_reconstruction_checkpoint(&rollout_line.item); + suffix_items.push(rollout_line.item); + if is_checkpoint { + found_checkpoint = true; + return Ok(false); + } + Ok(true) + })?; + + if should_fallback || !found_checkpoint { + return Ok(None); + } + + Ok(Some(suffix_items)) +} + +fn read_rollout_lines_from_end(path: &Path, mut visit: F) -> std::io::Result<()> +where + F: FnMut(&str) -> std::io::Result, +{ + use std::io::Read; + use std::io::Seek; + use std::io::SeekFrom; + + const CHUNK_SIZE: u64 = 64 * 1024; + + let mut file = std::fs::File::open(path)?; + let mut position = file.metadata()?.len(); + let mut pending = Vec::new(); + + while position > 0 { + let read_len = CHUNK_SIZE.min(position); + position -= read_len; + + let mut chunk = vec![0; read_len as usize]; + file.seek(SeekFrom::Start(position))?; + file.read_exact(&mut chunk)?; + if !pending.is_empty() { + chunk.extend_from_slice(&pending); + pending.clear(); + } + + let mut end = chunk.len(); + while let Some(newline_index) = chunk[..end].iter().rposition(|byte| *byte == b'\n') { + let line = std::str::from_utf8(&chunk[newline_index + 1..end]) + .map_err(|err| IoError::new(std::io::ErrorKind::InvalidData, err))?; + if !visit(line)? { + return Ok(()); + } + end = newline_index; + } + pending.extend_from_slice(&chunk[..end]); + } + + if !pending.is_empty() { + let line = std::str::from_utf8(&pending) + .map_err(|err| IoError::new(std::io::ErrorKind::InvalidData, err))?; + let _ = visit(line)?; + } + + Ok(()) +} + +fn rollout_item_is_reconstruction_checkpoint(item: &RolloutItem) -> bool { + matches!(item, RolloutItem::Compacted(compacted) if compacted.replacement_history.is_some()) +} + +fn rollout_item_needs_full_resume_load(item: &RolloutItem) -> bool { + match item { + RolloutItem::ResponseItem(codex_protocol::models::ResponseItem::Other) => true, + RolloutItem::Compacted(compacted) => { + compacted + .replacement_history + .as_ref() + .is_some_and(|history| { + history + .iter() + .any(|item| matches!(item, codex_protocol::models::ResponseItem::Other)) + }) + } + _ => false, + } +} + +fn read_recent_not_loaded_rollout_items( + path: &Path, + page_size: usize, +) -> std::io::Result> { + let contents = std::fs::read_to_string(path)?; + if contents + .lines() + .find(|line| !line.trim().is_empty()) + .is_some_and(rollout_line_has_fork_parent_rollout_ref) + { + return Ok(None); + } + let desired_turns = page_size.saturating_add(1); + let mut record_index_after = contents + .lines() + .filter(|line| !line.trim().is_empty()) + .count(); + let mut suffix_start_index = record_index_after; + let mut suffix_items = Vec::new(); + let mut turn_openings = 0usize; + + for line in contents.lines().rev() { + let trimmed = line.trim(); + if trimmed.is_empty() { + continue; + } + record_index_after = record_index_after.saturating_sub(1); + let rollout_line = + match serde_json::from_str::(trimmed) { + Ok(rollout_line) => rollout_line, + Err(_) => return Ok(None), + }; + if rollout_item_disables_recent_not_loaded_fast_path(&rollout_line.item) { + return Ok(None); + } + if rollout_item_opens_implicit_turn(&rollout_line.item) { + turn_openings = turn_openings.saturating_add(1); + } + suffix_start_index = record_index_after; + suffix_items.push(rollout_line.item); + if turn_openings >= desired_turns { + break; + } + } + + suffix_items.reverse(); + Ok(Some(RecentRolloutItems { + initial_rollout_index: suffix_start_index, + items: suffix_items, + })) +} + +fn rollout_item_opens_implicit_turn(item: &RolloutItem) -> bool { + matches!(item, RolloutItem::EventMsg(EventMsg::UserMessage(_))) +} + +fn rollout_item_disables_recent_not_loaded_fast_path(item: &RolloutItem) -> bool { + matches!( + item, + RolloutItem::EventMsg( + EventMsg::ThreadRolledBack(_) + | EventMsg::TurnStarted(_) + | EventMsg::TurnComplete(_) + | EventMsg::TurnAborted(_) + ) + ) +} + +fn build_recent_not_loaded_turns_page_response( + recent: RecentRolloutItems, + loaded_status: ThreadStatus, + limit: Option, +) -> Result { + let mut builder = + ThreadHistoryBuilder::with_initial_rollout_index(recent.initial_rollout_index); + for item in &recent.items { + if is_persisted_rollout_item(item) { + builder.handle_rollout_item(item); + } + } + let mut turns = builder.finish(); + normalize_thread_turns_status( + &mut turns, + loaded_status, + /*has_live_in_progress_turn*/ false, + ); + apply_thread_turns_items_view(&mut turns, TurnItemsView::NotLoaded); + let page = paginate_thread_turns(turns, None, limit, SortDirection::Desc)?; + Ok(ThreadTurnsListResponse { + data: page.turns, + next_cursor: page.next_cursor, + backwards_cursor: page.backwards_cursor, + }) +} + +fn thread_turns_page_size(limit: Option) -> usize { + limit + .map(|value| value as usize) + .unwrap_or(THREAD_TURNS_DEFAULT_LIMIT) + .clamp(1, THREAD_TURNS_MAX_LIMIT) +} + #[derive(serde::Serialize, serde::Deserialize)] #[serde(rename_all = "camelCase")] struct ThreadTurnsCursor { @@ -3764,10 +4499,7 @@ fn paginate_thread_turns( } let anchor = cursor.map(parse_thread_turns_cursor).transpose()?; - let page_size = limit - .map(|value| value as usize) - .unwrap_or(THREAD_TURNS_DEFAULT_LIMIT) - .clamp(1, THREAD_TURNS_MAX_LIMIT); + let page_size = thread_turns_page_size(limit); let anchor_index = anchor .as_ref() diff --git a/codex-rs/app-server/src/request_processors/thread_processor_tests.rs b/codex-rs/app-server/src/request_processors/thread_processor_tests.rs index 0ac9b7fa24..e9879de43e 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor_tests.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor_tests.rs @@ -324,6 +324,160 @@ mod thread_processor_behavior_tests { assert_eq!(turns.last(), Some(&active_turn)); } + fn compacted_resume_test_line(item: RolloutItem) -> codex_protocol::protocol::RolloutLine { + codex_protocol::protocol::RolloutLine { + timestamp: "2025-09-05T16:53:11.850Z".to_string(), + item, + } + } + + fn compacted_resume_test_session(thread_id: ThreadId) -> RolloutItem { + RolloutItem::SessionMeta(codex_protocol::protocol::SessionMetaLine { + meta: codex_protocol::protocol::SessionMeta { + id: thread_id, + timestamp: "2025-09-05T16:53:11.850Z".to_string(), + ..codex_protocol::protocol::SessionMeta::default() + }, + git: None, + }) + } + + fn compacted_resume_test_user_event(text: &str) -> RolloutItem { + RolloutItem::EventMsg(EventMsg::UserMessage( + codex_protocol::protocol::UserMessageEvent { + message: text.to_string(), + ..Default::default() + }, + )) + } + + fn compacted_resume_test_agent_event(text: &str) -> RolloutItem { + RolloutItem::EventMsg(EventMsg::AgentMessage( + codex_protocol::protocol::AgentMessageEvent { + message: text.to_string(), + phase: None, + memory_citation: None, + }, + )) + } + + fn compacted_resume_test_message( + role: &str, + text: &str, + ) -> codex_protocol::models::ResponseItem { + let content = if role == "user" { + vec![codex_protocol::models::ContentItem::InputText { + text: text.to_string(), + }] + } else { + vec![codex_protocol::models::ContentItem::OutputText { + text: text.to_string(), + }] + }; + codex_protocol::models::ResponseItem::Message { + id: None, + role: role.to_string(), + content, + phase: None, + } + } + + fn compacted_resume_test_compaction(message: &str, window_id: u64) -> RolloutItem { + RolloutItem::Compacted(codex_protocol::protocol::CompactedItem { + message: message.to_string(), + replacement_history: Some(vec![compacted_resume_test_message("assistant", message)]), + window_id: Some(window_id), + }) + } + + fn compacted_resume_test_rollout_contents(items: Vec) -> Result { + let mut contents = String::new(); + for item in items { + contents.push_str(&serde_json::to_string(&compacted_resume_test_line(item))?); + contents.push('\n'); + } + Ok(contents) + } + + fn compacted_resume_test_replacement_text(item: &RolloutItem) -> Option<&str> { + let RolloutItem::Compacted(compacted) = item else { + return None; + }; + let [codex_protocol::models::ResponseItem::Message { content, .. }] = + compacted.replacement_history.as_deref()? + else { + return None; + }; + let [codex_protocol::models::ContentItem::OutputText { text }] = content.as_slice() else { + return None; + }; + Some(text.as_str()) + } + + #[test] + fn compacted_resume_rollout_items_reads_latest_checkpoint_suffix_across_chunks() -> Result<()> { + let temp_dir = TempDir::new()?; + let path = temp_dir.path().join("rollout.jsonl"); + let thread_id = ThreadId::from_string("bfd12a78-5900-467b-9bc5-d3d35df08191")?; + let large_tail_message = "tail".repeat(20_000); + let contents = compacted_resume_test_rollout_contents(vec![ + compacted_resume_test_session(thread_id), + compacted_resume_test_user_event("first user event"), + RolloutItem::ResponseItem(compacted_resume_test_message("user", "first user response")), + compacted_resume_test_compaction("older checkpoint", 1), + compacted_resume_test_user_event("middle user event"), + compacted_resume_test_compaction("latest checkpoint", 2), + compacted_resume_test_agent_event(&large_tail_message), + ])?; + std::fs::write(&path, contents)?; + + let items = read_compacted_resume_rollout_items(path.as_path(), thread_id)? + .expect("compacted resume items"); + + assert_eq!(items.len(), 5); + assert!(matches!(items[0], RolloutItem::SessionMeta(_))); + assert!(matches!( + items[1], + RolloutItem::EventMsg(EventMsg::UserMessage(_)) + )); + assert!(matches!(items[2], RolloutItem::ResponseItem(_))); + assert_eq!( + compacted_resume_test_replacement_text(&items[3]), + Some("latest checkpoint") + ); + assert!(matches!( + &items[4], + RolloutItem::EventMsg(EventMsg::AgentMessage(event)) + if event.message == large_tail_message + )); + assert_eq!( + items + .iter() + .filter(|item| matches!(item, RolloutItem::Compacted(_))) + .count(), + 1 + ); + Ok(()) + } + + #[test] + fn compacted_resume_rollout_items_falls_back_on_malformed_tail_line() -> Result<()> { + let temp_dir = TempDir::new()?; + let path = temp_dir.path().join("rollout.jsonl"); + let thread_id = ThreadId::from_string("bfd12a78-5900-467b-9bc5-d3d35df08191")?; + let mut contents = compacted_resume_test_rollout_contents(vec![ + compacted_resume_test_session(thread_id), + compacted_resume_test_user_event("first user event"), + RolloutItem::ResponseItem(compacted_resume_test_message("user", "first user response")), + compacted_resume_test_compaction("latest checkpoint", 2), + ])?; + contents.push_str("{not json}\n"); + std::fs::write(&path, contents)?; + + assert!(read_compacted_resume_rollout_items(path.as_path(), thread_id)?.is_none()); + Ok(()) + } + #[test] fn validate_dynamic_tools_rejects_empty_namespace() { let tools = vec![dynamic_tool( diff --git a/codex-rs/app-server/tests/suite/conversation_summary.rs b/codex-rs/app-server/tests/suite/conversation_summary.rs index 6ad9f1c633..80fea41e58 100644 --- a/codex-rs/app-server/tests/suite/conversation_summary.rs +++ b/codex-rs/app-server/tests/suite/conversation_summary.rs @@ -130,6 +130,7 @@ async fn get_conversation_summary_by_thread_id_reads_pathless_store_thread() -> base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), multi_agent_version: None, + initial_rollout_copy: None, metadata: ThreadPersistenceMetadata { cwd: None, model_provider: "test-provider".to_string(), 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 baf11aff5e..4fba9d178b 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 @@ -156,6 +156,7 @@ async fn thread_delete_with_non_local_thread_store_does_not_create_local_persist base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), multi_agent_version: None, + initial_rollout_copy: None, metadata: ThreadPersistenceMetadata { cwd: Some(codex_home.path().to_path_buf()), model_provider: "mock_provider".to_string(), diff --git a/codex-rs/app-server/tests/suite/v2/thread_read.rs b/codex-rs/app-server/tests/suite/v2/thread_read.rs index a6ed5e4a63..df6246bd53 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_read.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_read.rs @@ -1366,6 +1366,7 @@ async fn seed_pathless_store_thread( base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), multi_agent_version: None, + initial_rollout_copy: None, metadata: ThreadPersistenceMetadata { cwd: None, model_provider: "test-provider".to_string(), diff --git a/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs b/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs index 7267ab073a..4b2c210173 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs @@ -217,6 +217,7 @@ async fn thread_unarchive_preserves_pathless_store_metadata() -> Result<()> { base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), multi_agent_version: None, + initial_rollout_copy: None, metadata: ThreadPersistenceMetadata { cwd: None, model_provider: "test-provider".to_string(), diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 9f5b867f75..677330b925 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -101,6 +101,7 @@ pub(crate) async fn run_codex_thread_interactive( conversation_history, session_source: SessionSource::SubAgent(subagent_source.clone()), forked_from_thread_id, + initial_rollout_copy: None, parent_thread_id: Some(parent_session.thread_id), thread_source: Some(ThreadSource::Subagent), agent_control: parent_session.services.agent_control.clone(), diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 0792334b30..bffd64d0cf 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -423,6 +423,7 @@ pub(crate) struct CodexSpawnArgs { pub(crate) conversation_history: InitialHistory, pub(crate) session_source: SessionSource, pub(crate) forked_from_thread_id: Option, + pub(crate) initial_rollout_copy: Option, pub(crate) parent_thread_id: Option, pub(crate) thread_source: Option, pub(crate) agent_control: AgentControl, @@ -512,6 +513,7 @@ impl Codex { conversation_history, session_source, forked_from_thread_id, + initial_rollout_copy, parent_thread_id, thread_source, agent_control, @@ -644,6 +646,7 @@ impl Codex { app_server_client_version: None, session_source, forked_from_thread_id, + initial_rollout_copy, parent_thread_id, thread_source, dynamic_tools, @@ -1259,12 +1262,15 @@ impl Session { } async fn record_initial_history(&self, conversation_history: InitialHistory) { - let is_subagent = { + let (is_subagent, skip_fork_history_persistence) = { let state = self.state.lock().await; - state - .session_configuration - .session_source - .is_non_root_agent() + ( + state + .session_configuration + .session_source + .is_non_root_agent(), + state.session_configuration.initial_rollout_copy.is_some(), + ) }; let has_prior_user_turns = initial_history_has_prior_user_turns(&conversation_history); { @@ -1331,7 +1337,12 @@ impl Session { } // If persisting, persist all rollout items as-is (the store filters). - if !rollout_items.is_empty() { + if skip_fork_history_persistence { + let suffix = Self::fork_rollout_copy_suffix(&rollout_items); + if !suffix.is_empty() { + self.persist_rollout_items(suffix).await; + } + } else if !rollout_items.is_empty() { self.persist_rollout_items(&rollout_items).await; } @@ -1346,6 +1357,17 @@ impl Session { } } + fn fork_rollout_copy_suffix(rollout_items: &[RolloutItem]) -> &[RolloutItem] { + let Some(RolloutItem::EventMsg(EventMsg::TurnAborted(event))) = rollout_items.last() else { + return &[]; + }; + if event.reason == TurnAbortReason::Interrupted { + &rollout_items[rollout_items.len() - 1..] + } else { + &[] + } + } + #[instrument( level = "trace", skip_all, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 0c07f5e951..922af9c6d0 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -102,6 +102,8 @@ pub(crate) struct SessionConfiguration { pub(super) session_source: SessionSource, /// Immediate history source copied into this thread, when this thread was forked. pub(super) forked_from_thread_id: Option, + /// Optional rollout JSONL copied into this thread after its own session metadata. + pub(super) initial_rollout_copy: Option, /// Immediate control/spawn parent for this thread, when it has one. pub(super) parent_thread_id: Option, /// Optional analytics source classification for this thread. @@ -557,6 +559,9 @@ impl Session { }, dynamic_tools: session_configuration.dynamic_tools.clone(), multi_agent_version: initial_multi_agent_version, + initial_rollout_copy: session_configuration + .initial_rollout_copy + .clone(), metadata: ThreadPersistenceMetadata { cwd: Some(config.cwd.to_path_buf()), model_provider: config.model_provider_id.clone(), @@ -570,10 +575,13 @@ impl Session { LiveThread::create(Arc::clone(&thread_store), params).await? } InitialHistory::Resumed(resumed_history) => { + let local_resume_with_path = resumed_history.rollout_path.is_some() + && thread_store.as_any().is::(); let params = ResumeThreadParams { thread_id: resumed_history.conversation_id, rollout_path: resumed_history.rollout_path.clone(), - history: Some(resumed_history.history.clone()), + history: (!local_resume_with_path) + .then(|| resumed_history.history.clone()), include_archived: true, metadata: ThreadPersistenceMetadata { cwd: Some(config.cwd.to_path_buf()), @@ -585,7 +593,12 @@ impl Session { }, }, }; - LiveThread::resume(Arc::clone(&thread_store), params).await? + LiveThread::resume_with_observed_history( + Arc::clone(&thread_store), + params, + Some(resumed_history.history.as_slice()), + ) + .await? } }; Ok(Some(live_thread)) diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index bca6be6a74..a1fe0b18c7 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -3378,6 +3378,7 @@ async fn set_rate_limits_retains_previous_credits() { app_server_client_version: None, session_source: SessionSource::Exec, forked_from_thread_id: None, + initial_rollout_copy: None, parent_thread_id: None, thread_source: None, dynamic_tools: Vec::new(), @@ -3485,6 +3486,7 @@ async fn set_rate_limits_updates_plan_type_when_present() { app_server_client_version: None, session_source: SessionSource::Exec, forked_from_thread_id: None, + initial_rollout_copy: None, parent_thread_id: None, thread_source: None, dynamic_tools: Vec::new(), @@ -3738,6 +3740,7 @@ async fn attach_thread_persistence(session: &mut Session) -> PathBuf { base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), multi_agent_version: None, + initial_rollout_copy: None, metadata: ThreadPersistenceMetadata { cwd: Some(config.cwd.to_path_buf()), model_provider: config.model_provider_id.clone(), @@ -4013,6 +4016,7 @@ pub(crate) async fn make_session_configuration_for_tests() -> SessionConfigurati app_server_client_version: None, session_source: SessionSource::Exec, forked_from_thread_id: None, + initial_rollout_copy: None, parent_thread_id: None, thread_source: None, dynamic_tools: Vec::new(), @@ -4880,6 +4884,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() { app_server_client_version: None, session_source: SessionSource::Exec, forked_from_thread_id: None, + initial_rollout_copy: None, parent_thread_id: None, thread_source: None, dynamic_tools: Vec::new(), @@ -4993,6 +4998,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { app_server_client_version: None, session_source: SessionSource::Exec, forked_from_thread_id: None, + initial_rollout_copy: None, parent_thread_id: None, thread_source: None, dynamic_tools: Vec::new(), @@ -5239,6 +5245,7 @@ async fn make_session_with_config_and_rx( app_server_client_version: None, session_source: SessionSource::Exec, forked_from_thread_id: None, + initial_rollout_copy: None, parent_thread_id: None, thread_source: None, dynamic_tools: Vec::new(), @@ -5346,6 +5353,7 @@ async fn make_session_with_history_source_and_agent_control_and_rx( app_server_client_version: None, session_source: session_source.clone(), forked_from_thread_id: None, + initial_rollout_copy: None, parent_thread_id: None, thread_source: None, dynamic_tools: Vec::new(), @@ -6578,6 +6586,7 @@ async fn shutdown_complete_does_not_append_to_thread_store_after_shutdown() { base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), multi_agent_version: None, + initial_rollout_copy: None, metadata: ThreadPersistenceMetadata { cwd: Some(config.cwd.to_path_buf()), model_provider: config.model_provider_id.clone(), @@ -7053,6 +7062,7 @@ where app_server_client_version: None, session_source: SessionSource::Exec, forked_from_thread_id: None, + initial_rollout_copy: None, parent_thread_id: None, thread_source: None, dynamic_tools, diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index 7011e6d0f5..b050cbb308 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -720,6 +720,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { GUARDIAN_REVIEWER_NAME.to_string(), )), forked_from_thread_id: None, + initial_rollout_copy: None, parent_thread_id: None, thread_source: None, agent_control: AgentControl::default(), diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 3a0e374750..131cb48cd9 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -954,6 +954,32 @@ impl ThreadManager { thread_source, parent_trace, supports_openai_form_elicitation, + /*initial_rollout_copy*/ None, + ) + .await + } + + pub async fn fork_thread_from_history_with_initial_rollout_copy( + &self, + snapshot: S, + config: Config, + history: InitialHistory, + thread_source: Option, + parent_trace: Option, + supports_openai_form_elicitation: bool, + initial_rollout_copy: PathBuf, + ) -> CodexResult + where + S: Into, + { + self.fork_thread_with_initial_history( + snapshot.into(), + config, + history, + thread_source, + parent_trace, + supports_openai_form_elicitation, + Some(initial_rollout_copy), ) .await } @@ -966,6 +992,7 @@ impl ThreadManager { thread_source: Option, parent_trace: Option, supports_openai_form_elicitation: bool, + initial_rollout_copy: Option, ) -> CodexResult { // `forked_from_id()` describes this history's existing lineage. When // forking a resumed thread, the child copies the resumed thread itself. @@ -999,23 +1026,30 @@ impl ThreadManager { &config.cwd, ); let agent_control = self.agent_control_for_config(&config); - Box::pin(self.state.spawn_thread( - config, - history, - Arc::clone(&self.state.auth_manager), - agent_control, - /*parent_thread_id*/ None, - source_thread_id, - thread_source, - Vec::new(), - /*metrics_service_name*/ None, - initial_multi_agent_mode, - parent_trace, - environments, - /*thread_extension_init*/ ExtensionDataInit::default(), - supports_openai_form_elicitation, - /*user_shell_override*/ None, - )) + Box::pin( + self.state + .spawn_thread_with_source_and_initial_rollout_copy( + config, + history, + Arc::clone(&self.state.auth_manager), + agent_control, + self.state.session_source.clone(), + /*parent_thread_id*/ None, + source_thread_id, + thread_source, + Vec::new(), + /*metrics_service_name*/ None, + initial_multi_agent_mode, + /*inherited_environments*/ None, + /*inherited_exec_policy*/ None, + parent_trace, + environments, + /*thread_extension_init*/ ExtensionDataInit::default(), + supports_openai_form_elicitation, + /*user_shell_override*/ None, + initial_rollout_copy, + ), + ) .await } @@ -1428,6 +1462,53 @@ impl ThreadManagerState { thread_extension_init: ExtensionDataInit, supports_openai_form_elicitation: bool, user_shell_override: Option, + ) -> CodexResult { + Box::pin(self.spawn_thread_with_source_and_initial_rollout_copy( + config, + initial_history, + auth_manager, + agent_control, + session_source, + parent_thread_id, + forked_from_thread_id, + thread_source, + dynamic_tools, + metrics_service_name, + initial_multi_agent_mode, + inherited_environments, + inherited_exec_policy, + parent_trace, + environments, + thread_extension_init, + supports_openai_form_elicitation, + user_shell_override, + /*initial_rollout_copy*/ None, + )) + .await + } + + #[allow(clippy::too_many_arguments)] + pub(crate) async fn spawn_thread_with_source_and_initial_rollout_copy( + &self, + config: Config, + initial_history: InitialHistory, + auth_manager: Arc, + agent_control: AgentControl, + session_source: SessionSource, + parent_thread_id: Option, + forked_from_thread_id: Option, + thread_source: Option, + dynamic_tools: Vec, + metrics_service_name: Option, + initial_multi_agent_mode: Option, + inherited_environments: Option, + inherited_exec_policy: Option>, + parent_trace: Option, + environments: Vec, + thread_extension_init: ExtensionDataInit, + supports_openai_form_elicitation: bool, + user_shell_override: Option, + initial_rollout_copy: Option, ) -> CodexResult { let is_resumed_thread = matches!(&initial_history, InitialHistory::Resumed(_)); if let InitialHistory::Resumed(resumed) = &initial_history { @@ -1482,6 +1563,7 @@ impl ThreadManagerState { conversation_history: initial_history, session_source, forked_from_thread_id, + initial_rollout_copy, parent_thread_id, thread_source, agent_control, diff --git a/codex-rs/rollout/src/lib.rs b/codex-rs/rollout/src/lib.rs index 23be07996c..d1b067233b 100644 --- a/codex-rs/rollout/src/lib.rs +++ b/codex-rs/rollout/src/lib.rs @@ -63,9 +63,11 @@ pub use metadata::builder_from_items; pub use policy::is_persisted_rollout_item; pub use policy::persisted_rollout_items; pub use policy::should_persist_response_item_for_memories; +pub use recorder::ForkParentRolloutRef; pub use recorder::RolloutRecorder; pub use recorder::RolloutRecorderParams; pub use recorder::append_rollout_item_to_path; +pub use recorder::read_fork_parent_rollout_ref; 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/list.rs b/codex-rs/rollout/src/list.rs index f5013e7aae..49c27926ed 100644 --- a/codex-rs/rollout/src/list.rs +++ b/codex-rs/rollout/src/list.rs @@ -1258,6 +1258,25 @@ fn event_msg_preview(event: &EventMsg) -> Option { /// Read the SessionMetaLine from the head of a rollout file for reuse by /// callers that need the session metadata (e.g. to derive a cwd for config). pub async fn read_session_meta_line(path: &Path) -> io::Result { + let mut lines = compression::open_rollout_line_reader(path).await?; + while let Some(line) = lines.next_line().await? { + let trimmed = line.trim(); + if trimmed.is_empty() { + continue; + } + + if let Ok(RolloutLine { + item: RolloutItem::SessionMeta(session_meta_line), + .. + }) = serde_json::from_str::(trimmed) + { + return Ok(session_meta_line); + } + break; + } + + // Preserve the legacy tolerant behavior for unusual files whose first + // parseable summary record is still session metadata. let head = read_head_for_summary(path).await?; let Some(first) = head.first() else { return Err(io::Error::other(format!( diff --git a/codex-rs/rollout/src/recorder.rs b/codex-rs/rollout/src/recorder.rs index 7ba111e944..c0f2cd4fab 100644 --- a/codex-rs/rollout/src/recorder.rs +++ b/codex-rs/rollout/src/recorder.rs @@ -50,6 +50,7 @@ use crate::state_db; use crate::state_db::StateDbHandle; use codex_git_utils::collect_git_info; use codex_git_utils::get_git_repo_root; +use codex_protocol::models::ResponseItem; use codex_protocol::protocol::GitInfo as ProtocolGitInfo; use codex_protocol::protocol::InitialHistory; use codex_protocol::protocol::MultiAgentVersion; @@ -63,6 +64,15 @@ use codex_protocol::protocol::ThreadSource; use codex_state::StateRuntime; use codex_utils_path as path_utils; +const FORK_PARENT_ROLLOUT_PATH_FIELD: &str = "fork_parent_rollout_path"; +const FORK_PARENT_ROLLOUT_BYTE_LEN_FIELD: &str = "fork_parent_rollout_byte_len"; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ForkParentRolloutRef { + pub path: PathBuf, + pub byte_len: u64, +} + /// Writes canonical session rollout items to JSONL. /// /// Rollouts are recorded as JSONL and can be inspected with tools such as: @@ -89,6 +99,7 @@ pub enum RolloutRecorderParams { base_instructions: BaseInstructions, dynamic_tools: Vec, multi_agent_version: Option, + initial_rollout_copy: Option, }, Resume { path: PathBuf, @@ -175,6 +186,7 @@ impl RolloutRecorderParams { base_instructions, dynamic_tools, multi_agent_version: None, + initial_rollout_copy: None, } } @@ -192,6 +204,17 @@ impl RolloutRecorderParams { self } + pub fn with_initial_rollout_copy(mut self, initial_rollout_copy: Option) -> Self { + if let Self::Create { + initial_rollout_copy: copy, + .. + } = &mut self + { + *copy = initial_rollout_copy; + } + self + } + pub fn resume(path: PathBuf) -> Self { Self::Resume { path } } @@ -694,7 +717,8 @@ impl RolloutRecorder { config: &impl RolloutConfigView, params: RolloutRecorderParams, ) -> std::io::Result { - let (file, deferred_log_file_info, rollout_path, meta) = match params { + let (file, deferred_log_file_info, rollout_path, meta, initial_rollout_copy) = match params + { RolloutRecorderParams::Create { conversation_id, forked_from_id, @@ -704,6 +728,7 @@ impl RolloutRecorder { base_instructions, dynamic_tools, multi_agent_version, + initial_rollout_copy, } => { let log_file_info = precompute_log_file_info(config, conversation_id)?; let path = log_file_info.path.clone(); @@ -742,7 +767,13 @@ impl RolloutRecorder { multi_agent_version, }; - (None, Some(log_file_info), path, Some(session_meta)) + ( + None, + Some(log_file_info), + path, + Some(session_meta), + initial_rollout_copy, + ) } RolloutRecorderParams::Resume { path } => { let path = compression::materialize_rollout_for_append(path.as_path()).await?; @@ -756,6 +787,7 @@ impl RolloutRecorder { None, path, None, + None, ) } }; @@ -779,6 +811,7 @@ impl RolloutRecorder { deferred_log_file_info, rx, meta, + initial_rollout_copy, cwd, rollout_path_for_spawn.clone(), ) @@ -878,35 +911,25 @@ impl RolloutRecorder { continue; } saw_non_empty_line = true; - let mut v: Value = match serde_json::from_str(&line) { - Ok(v) => v, - Err(e) => { - warn!("failed to parse line as JSON: {line:?}, error: {e}"); - parse_errors = parse_errors.saturating_add(1); - continue; - } - }; - if strip_legacy_ghost_snapshot_rollout_line(&mut v) { - trace!("skipping legacy ghost_snapshot rollout line"); - continue; - } - - // Parse the rollout line structure - match serde_json::from_value::(v.clone()) { - Ok(rollout_line) => { - let item = rollout_line.item; - // Use the FIRST SessionMeta encountered in the file as the canonical - // thread id and main session information. Keep all items intact. - if thread_id.is_none() - && let RolloutItem::SessionMeta(session_meta_line) = &item + match serde_json::from_str::(&line) { + Ok(rollout_line) if rollout_item_may_need_legacy_cleanup(&rollout_line.item) => { + if let Some(rollout_line) = + parse_rollout_line_after_legacy_cleanup(&line, &mut parse_errors, None) { - thread_id = Some(session_meta_line.meta.id); + push_loaded_rollout_item(&mut items, &mut thread_id, rollout_line.item); } - items.push(item); } - Err(e) => { - trace!("failed to parse rollout line: {e}"); - parse_errors = parse_errors.saturating_add(1); + Ok(rollout_line) => { + push_loaded_rollout_item(&mut items, &mut thread_id, rollout_line.item); + } + Err(rollout_error) => { + if let Some(rollout_line) = parse_rollout_line_after_legacy_cleanup( + &line, + &mut parse_errors, + Some(&rollout_error), + ) { + push_loaded_rollout_item(&mut items, &mut thread_id, rollout_line.item); + } } } } @@ -923,6 +946,65 @@ impl RolloutRecorder { Ok((items, thread_id, parse_errors)) } + pub async fn load_rollout_items_prefix( + path: &Path, + byte_len: u64, + ) -> std::io::Result<(Vec, Option, usize)> { + trace!("Resuming rollout prefix from {path:?}, byte_len={byte_len}"); + let mut items: Vec = Vec::new(); + let mut thread_id: Option = None; + let mut parse_errors = 0usize; + let mut bytes_read = 0u64; + let mut reader = compression::open_rollout_line_reader(path).await?; + let mut saw_non_empty_line = false; + while bytes_read < byte_len { + let Some(line) = reader.next_line().await? else { + return Err(IoError::other(format!( + "rollout prefix ended before {byte_len} bytes" + ))); + }; + let line_len = u64::try_from(line.len()) + .unwrap_or(u64::MAX) + .saturating_add(1); + if bytes_read.saturating_add(line_len) > byte_len { + return Err(IoError::other(format!( + "rollout prefix boundary {byte_len} splits a JSONL line" + ))); + } + bytes_read = bytes_read.saturating_add(line_len); + if line.trim().is_empty() { + continue; + } + saw_non_empty_line = true; + match serde_json::from_str::(&line) { + Ok(rollout_line) if rollout_item_may_need_legacy_cleanup(&rollout_line.item) => { + if let Some(rollout_line) = + parse_rollout_line_after_legacy_cleanup(&line, &mut parse_errors, None) + { + push_loaded_rollout_item(&mut items, &mut thread_id, rollout_line.item); + } + } + Ok(rollout_line) => { + push_loaded_rollout_item(&mut items, &mut thread_id, rollout_line.item); + } + Err(rollout_error) => { + if let Some(rollout_line) = parse_rollout_line_after_legacy_cleanup( + &line, + &mut parse_errors, + Some(&rollout_error), + ) { + push_loaded_rollout_item(&mut items, &mut thread_id, rollout_line.item); + } + } + } + } + if !saw_non_empty_line { + return Err(IoError::other("empty session file prefix")); + } + + Ok((items, thread_id, parse_errors)) + } + pub async fn get_rollout_history(path: &Path) -> std::io::Result { let (items, thread_id, _parse_errors) = Self::load_rollout_items(path).await?; let conversation_id = thread_id @@ -968,6 +1050,101 @@ impl RolloutRecorder { } } +pub async fn read_fork_parent_rollout_ref( + path: &Path, +) -> std::io::Result> { + let mut reader = compression::open_rollout_line_reader(path).await?; + while let Some(line) = reader.next_line().await? { + let trimmed = line.trim(); + if trimmed.is_empty() { + continue; + } + return Ok(parse_fork_parent_rollout_ref_from_line(trimmed)); + } + Ok(None) +} + +pub fn parse_fork_parent_rollout_ref_from_line(line: &str) -> Option { + let value: Value = serde_json::from_str(line).ok()?; + let payload = value.get("payload").unwrap_or(&value); + let path = payload + .get(FORK_PARENT_ROLLOUT_PATH_FIELD) + .and_then(Value::as_str) + .filter(|path| !path.is_empty())?; + let byte_len = payload + .get(FORK_PARENT_ROLLOUT_BYTE_LEN_FIELD) + .and_then(Value::as_u64) + .filter(|byte_len| *byte_len > 0)?; + Some(ForkParentRolloutRef { + path: PathBuf::from(path), + byte_len, + }) +} + +fn push_loaded_rollout_item( + items: &mut Vec, + thread_id: &mut Option, + item: RolloutItem, +) { + // Use the FIRST SessionMeta encountered in the file as the canonical + // thread id and main session information. Keep all items intact. + if thread_id.is_none() + && let RolloutItem::SessionMeta(session_meta_line) = &item + { + *thread_id = Some(session_meta_line.meta.id); + } + items.push(item); +} + +fn rollout_item_may_need_legacy_cleanup(item: &RolloutItem) -> bool { + match item { + RolloutItem::ResponseItem(ResponseItem::Other) => true, + RolloutItem::Compacted(compacted) => { + compacted + .replacement_history + .as_ref() + .is_some_and(|history| { + history + .iter() + .any(|item| matches!(item, ResponseItem::Other)) + }) + } + _ => false, + } +} + +fn parse_rollout_line_after_legacy_cleanup( + line: &str, + parse_errors: &mut usize, + rollout_error: Option<&serde_json::Error>, +) -> Option { + let mut value: Value = match serde_json::from_str(line) { + Ok(value) => value, + Err(err) => { + warn!("failed to parse line as JSON: {line:?}, error: {err}"); + *parse_errors = parse_errors.saturating_add(1); + return None; + } + }; + if strip_legacy_ghost_snapshot_rollout_line(&mut value) { + trace!("skipping legacy ghost_snapshot rollout line"); + return None; + } + + match serde_json::from_value::(value) { + Ok(rollout_line) => Some(rollout_line), + Err(err) => { + if let Some(rollout_error) = rollout_error { + trace!("failed to parse rollout line: {rollout_error}; fallback error: {err}"); + } else { + trace!("failed to parse rollout line after legacy cleanup: {err}"); + } + *parse_errors = parse_errors.saturating_add(1); + None + } + } +} + fn strip_legacy_ghost_snapshot_rollout_line(value: &mut Value) -> bool { match value.get("type").and_then(Value::as_str) { Some("response_item") => value @@ -1461,6 +1638,7 @@ struct RolloutWriterState { deferred_log_file_info: Option, pending_items: Vec, meta: Option, + initial_rollout_copy: Option, cwd: PathBuf, rollout_path: PathBuf, last_logged_error: Option, @@ -1471,6 +1649,7 @@ impl RolloutWriterState { file: Option, deferred_log_file_info: Option, meta: Option, + initial_rollout_copy: Option, cwd: PathBuf, rollout_path: PathBuf, ) -> Self { @@ -1479,6 +1658,7 @@ impl RolloutWriterState { deferred_log_file_info, pending_items: Vec::new(), meta, + initial_rollout_copy, cwd, rollout_path, last_logged_error: None, @@ -1584,14 +1764,41 @@ impl RolloutWriterState { let Some(session_meta) = self.meta.as_ref().cloned() else { return Ok(()); }; - write_session_meta(self.writer.as_mut(), session_meta, &self.cwd).await?; + let fork_parent_rollout_ref = match self.initial_rollout_copy.as_deref() { + Some(source_path) => fork_parent_rollout_ref_for_initial_copy(source_path).await?, + None => None, + }; + write_session_meta( + self.writer.as_mut(), + session_meta, + &self.cwd, + fork_parent_rollout_ref.as_ref(), + ) + .await?; + if fork_parent_rollout_ref.is_some() { + self.initial_rollout_copy = None; + } self.meta = None; Ok(()) } + async fn copy_initial_rollout_if_needed(&mut self) -> std::io::Result<()> { + let Some(source_path) = self.initial_rollout_copy.as_ref().cloned() else { + return Ok(()); + }; + let Some(writer) = self.writer.as_mut() else { + return Err(IoError::other("rollout writer is not open")); + }; + let mut source = tokio::fs::File::open(source_path.as_path()).await?; + tokio::io::copy(&mut source, &mut writer.file).await?; + self.initial_rollout_copy = None; + Ok(()) + } + async fn write_pending_once(&mut self) -> std::io::Result<()> { self.ensure_writer_open().await?; self.write_session_meta_if_needed().await?; + self.copy_initial_rollout_if_needed().await?; self.write_pending_items_once().await?; @@ -1629,10 +1836,18 @@ async fn rollout_writer( deferred_log_file_info: Option, mut rx: mpsc::Receiver, meta: Option, + initial_rollout_copy: Option, cwd: PathBuf, rollout_path: PathBuf, ) -> std::io::Result<()> { - let mut state = RolloutWriterState::new(file, deferred_log_file_info, meta, cwd, rollout_path); + let mut state = RolloutWriterState::new( + file, + deferred_log_file_info, + meta, + initial_rollout_copy, + cwd, + rollout_path, + ); // Process rollout commands while let Some(cmd) = rx.recv().await { @@ -1662,10 +1877,36 @@ async fn rollout_writer( Ok(()) } +async fn fork_parent_rollout_ref_for_initial_copy( + source_path: &Path, +) -> std::io::Result> { + let Some(existing_path) = compression::existing_rollout_path(source_path).await else { + return Ok(None); + }; + if !is_plain_jsonl_rollout_path(existing_path.as_path()) { + return Ok(None); + } + let metadata = tokio::fs::metadata(existing_path.as_path()).await?; + if !metadata.is_file() || metadata.len() == 0 { + return Ok(None); + } + Ok(Some(ForkParentRolloutRef { + path: existing_path, + byte_len: metadata.len(), + })) +} + +fn is_plain_jsonl_rollout_path(path: &Path) -> bool { + path.file_name() + .and_then(|name| name.to_str()) + .is_some_and(|name| name.ends_with(".jsonl")) +} + async fn write_session_meta( mut writer: Option<&mut JsonlWriter>, session_meta: SessionMeta, cwd: &Path, + fork_parent_rollout_ref: Option<&ForkParentRolloutRef>, ) -> std::io::Result<()> { let git_info = if get_git_repo_root(cwd).is_some() { collect_git_info(cwd).await.map(|info| ProtocolGitInfo { @@ -1681,9 +1922,10 @@ async fn write_session_meta( git: git_info, }; - let rollout_item = RolloutItem::SessionMeta(session_meta_line); if let Some(writer) = writer.as_mut() { - writer.write_rollout_item(&rollout_item).await?; + writer + .write_session_meta_line(&session_meta_line, fork_parent_rollout_ref) + .await?; } Ok(()) } @@ -1717,7 +1959,52 @@ struct RolloutLineRef<'a> { item: &'a RolloutItem, } +#[derive(serde::Serialize)] +struct SessionMetaLineRef<'a> { + timestamp: String, + #[serde(rename = "type")] + item_type: &'static str, + payload: SessionMetaPayloadRef<'a>, +} + +#[derive(serde::Serialize)] +struct SessionMetaPayloadRef<'a> { + #[serde(flatten)] + meta: &'a SessionMeta, + #[serde(skip_serializing_if = "Option::is_none")] + git: Option<&'a ProtocolGitInfo>, + #[serde(skip_serializing_if = "Option::is_none")] + fork_parent_rollout_path: Option<&'a PathBuf>, + #[serde(skip_serializing_if = "Option::is_none")] + fork_parent_rollout_byte_len: Option, +} + impl JsonlWriter { + async fn write_session_meta_line( + &mut self, + session_meta_line: &SessionMetaLine, + fork_parent_rollout_ref: Option<&ForkParentRolloutRef>, + ) -> std::io::Result<()> { + let timestamp_format: &[FormatItem] = format_description!( + "[year]-[month]-[day]T[hour]:[minute]:[second].[subsecond digits:3]Z" + ); + let timestamp = OffsetDateTime::now_utc() + .format(timestamp_format) + .map_err(|e| IoError::other(format!("failed to format timestamp: {e}")))?; + let line = SessionMetaLineRef { + timestamp, + item_type: "session_meta", + payload: SessionMetaPayloadRef { + meta: &session_meta_line.meta, + git: session_meta_line.git.as_ref(), + fork_parent_rollout_path: fork_parent_rollout_ref.map(|reference| &reference.path), + fork_parent_rollout_byte_len: fork_parent_rollout_ref + .map(|reference| reference.byte_len), + }, + }; + self.write_line(&line).await + } + async fn write_rollout_item(&mut self, rollout_item: &RolloutItem) -> std::io::Result<()> { let timestamp_format: &[FormatItem] = format_description!( "[year]-[month]-[day]T[hour]:[minute]:[second].[subsecond digits:3]Z" diff --git a/codex-rs/rollout/src/recorder_tests.rs b/codex-rs/rollout/src/recorder_tests.rs index 8701578ebf..49f334ddfd 100644 --- a/codex-rs/rollout/src/recorder_tests.rs +++ b/codex-rs/rollout/src/recorder_tests.rs @@ -506,6 +506,7 @@ async fn writer_state_retries_write_error_before_reporting_flush_success() -> st Some(tokio::fs::File::from_std(read_only_file)), /*deferred_log_file_info*/ None, /*meta*/ None, + /*initial_rollout_copy*/ None, home.path().to_path_buf(), rollout_path.clone(), ); diff --git a/codex-rs/thread-store/src/in_memory.rs b/codex-rs/thread-store/src/in_memory.rs index a60409b740..b420e0f7b6 100644 --- a/codex-rs/thread-store/src/in_memory.rs +++ b/codex-rs/thread-store/src/in_memory.rs @@ -119,6 +119,7 @@ mod tests { base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), multi_agent_version: None, + initial_rollout_copy: None, metadata: ThreadPersistenceMetadata { cwd: None, model_provider: "test-provider".to_string(), diff --git a/codex-rs/thread-store/src/live_thread.rs b/codex-rs/thread-store/src/live_thread.rs index 6b94807511..c2eb875817 100644 --- a/codex-rs/thread-store/src/live_thread.rs +++ b/codex-rs/thread-store/src/live_thread.rs @@ -100,13 +100,28 @@ impl LiveThread { pub async fn resume( thread_store: Arc, - mut params: ResumeThreadParams, + params: ResumeThreadParams, + ) -> ThreadStoreResult { + Self::resume_with_observed_history(thread_store, params, None).await + } + + pub async fn resume_with_observed_history( + thread_store: Arc, + params: ResumeThreadParams, + observed_history: Option<&[RolloutItem]>, ) -> ThreadStoreResult { let thread_id = params.thread_id; - let should_load_history = params.history.is_none(); let include_archived = params.include_archived; - thread_store.resume_thread(params.clone()).await?; - if should_load_history { + let metadata = params.metadata.clone(); + let metadata_sync = observed_history + .or(params.history.as_deref()) + .map(|history| { + ThreadMetadataSync::for_resume_parts(thread_id, &metadata, Some(history)) + }); + thread_store.resume_thread(params).await?; + let metadata_sync = if let Some(metadata_sync) = metadata_sync { + metadata_sync + } else { match thread_store .load_history(LoadThreadHistoryParams { thread_id, @@ -114,7 +129,11 @@ impl LiveThread { }) .await { - Ok(history) => params.history = Some(history.items), + Ok(history) => ThreadMetadataSync::for_resume_parts( + thread_id, + &metadata, + Some(history.items.as_slice()), + ), Err(err) => { if let Err(discard_err) = thread_store.discard_thread(thread_id).await { warn!( @@ -124,8 +143,7 @@ impl LiveThread { return Err(err); } } - } - let metadata_sync = ThreadMetadataSync::for_resume(¶ms); + }; Ok(Self { thread_id, thread_store, diff --git a/codex-rs/thread-store/src/local/create_thread.rs b/codex-rs/thread-store/src/local/create_thread.rs index 05924881d5..a849da4fea 100644 --- a/codex-rs/thread-store/src/local/create_thread.rs +++ b/codex-rs/thread-store/src/local/create_thread.rs @@ -36,7 +36,8 @@ pub(super) async fn create_thread( params.base_instructions, params.dynamic_tools, ) - .with_multi_agent_version(params.multi_agent_version), + .with_multi_agent_version(params.multi_agent_version) + .with_initial_rollout_copy(params.initial_rollout_copy), ) .await .map_err(|err| ThreadStoreError::Internal { diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs index 19a0d2955b..056d9444a3 100644 --- a/codex-rs/thread-store/src/local/mod.rs +++ b/codex-rs/thread-store/src/local/mod.rs @@ -1132,6 +1132,7 @@ mod tests { base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), multi_agent_version: None, + initial_rollout_copy: None, metadata: thread_metadata(), } } diff --git a/codex-rs/thread-store/src/local/read_thread.rs b/codex-rs/thread-store/src/local/read_thread.rs index a1e1c59190..f27ca0c688 100644 --- a/codex-rs/thread-store/src/local/read_thread.rs +++ b/codex-rs/thread-store/src/local/read_thread.rs @@ -39,12 +39,8 @@ pub(super) async fn read_thread( metadata.rollout_path.as_path(), ))) && (!params.include_history - || sqlite_rollout_path_can_load_history_for_thread( - store, - &metadata.rollout_path, - thread_id, - ) - .await) + || sqlite_rollout_path_can_load_history_for_thread(&metadata.rollout_path, thread_id) + .await) { let metadata_sandbox_policy = metadata.sandbox_policy.clone(); let mut thread = stored_thread_from_sqlite_metadata(store, metadata).await; @@ -87,7 +83,6 @@ pub(super) async fn read_thread( } async fn sqlite_rollout_path_can_load_history_for_thread( - store: &LocalThreadStore, path: &std::path::Path, thread_id: codex_protocol::ThreadId, ) -> bool { @@ -96,10 +91,11 @@ async fn sqlite_rollout_path_can_load_history_for_thread( } // SQLite metadata can outlive a moved/recreated rollout path. When history is // requested, verify the path still resolves to the requested thread before - // trusting it as the source replay. - read_thread_from_rollout_path(store, path.to_path_buf()) + // trusting it as the source replay. The session metadata line is enough for + // this identity check and avoids doing a summary scan before full replay. + read_session_meta_line(path) .await - .is_ok_and(|thread| thread.thread_id == thread_id) + .is_ok_and(|meta_line| meta_line.meta.id == thread_id) } pub(super) async fn read_thread_by_rollout_path( @@ -288,6 +284,75 @@ async fn load_history_items( .map_err(|err| ThreadStoreError::Internal { message: format!("failed to load thread history {}: {err}", path.display()), })?; + expand_fork_parent_rollout_ref(path, items).await +} + +async fn expand_fork_parent_rollout_ref( + path: &std::path::Path, + child_items: Vec, +) -> ThreadStoreResult> { + let Some(parent_ref) = codex_rollout::read_fork_parent_rollout_ref(path) + .await + .map_err(|err| ThreadStoreError::Internal { + message: format!( + "failed to read fork parent reference {}: {err}", + path.display() + ), + })? + else { + return Ok(child_items); + }; + + let Some(child_forked_from_id) = child_items.iter().find_map(|item| match item { + codex_protocol::protocol::RolloutItem::SessionMeta(meta_line) => { + meta_line.meta.forked_from_id + } + _ => None, + }) else { + return Err(ThreadStoreError::Internal { + message: format!( + "failed to load fork parent history for {}: child fork parent missing", + path.display() + ), + }); + }; + let Some(parent_path) = codex_rollout::existing_rollout_path(parent_ref.path.as_path()).await + else { + return Err(ThreadStoreError::Internal { + message: format!( + "failed to load fork parent history for {}: parent rollout {} does not exist", + path.display(), + parent_ref.path.display() + ), + }); + }; + let (parent_items, parent_thread_id, _) = + RolloutRecorder::load_rollout_items_prefix(parent_path.as_path(), parent_ref.byte_len) + .await + .map_err(|err| ThreadStoreError::Internal { + message: format!( + "failed to load fork parent history {} for {}: {err}", + parent_path.display(), + path.display() + ), + })?; + if parent_thread_id != Some(child_forked_from_id) { + return Err(ThreadStoreError::Internal { + message: format!( + "failed to load fork parent history for {}: parent identity mismatch", + path.display() + ), + }); + } + + let mut child_items = child_items.into_iter(); + let Some(child_meta) = child_items.next() else { + return Ok(parent_items); + }; + let mut items = Vec::with_capacity(1 + parent_items.len() + child_items.size_hint().0); + items.push(child_meta); + items.extend(parent_items); + items.extend(child_items); Ok(items) } diff --git a/codex-rs/thread-store/src/thread_metadata_sync.rs b/codex-rs/thread-store/src/thread_metadata_sync.rs index b45bb2d635..94a1d2934d 100644 --- a/codex-rs/thread-store/src/thread_metadata_sync.rs +++ b/codex-rs/thread-store/src/thread_metadata_sync.rs @@ -16,8 +16,10 @@ use codex_protocol::protocol::UserMessageEvent; use crate::CreateThreadParams; use crate::GitInfoPatch; +#[cfg(test)] use crate::ResumeThreadParams; use crate::ThreadMetadataPatch; +use crate::ThreadPersistenceMetadata; const IMAGE_ONLY_USER_MESSAGE_PLACEHOLDER: &str = "[Image]"; #[cfg(not(test))] @@ -90,11 +92,23 @@ impl ThreadMetadataSync { } } + #[cfg(test)] pub(crate) fn for_resume(params: &ResumeThreadParams) -> Self { + Self::for_resume_parts( + params.thread_id, + ¶ms.metadata, + params.history.as_deref(), + ) + } + + pub(crate) fn for_resume_parts( + thread_id: ThreadId, + metadata: &ThreadPersistenceMetadata, + history: Option<&[RolloutItem]>, + ) -> Self { let mut sync = Self { - thread_id: params.thread_id, - cwd_seen: params - .metadata + thread_id, + cwd_seen: metadata .cwd .as_ref() .is_some_and(|cwd| !cwd.as_os_str().is_empty()), @@ -107,7 +121,7 @@ impl ThreadMetadataSync { defer_create_update_until_history_exists: false, defer_resume_update_until_append: false, }; - if let Some(history) = params.history.as_deref() { + if let Some(history) = history { let update = sync.observe_resume_history(history); sync.merge_pending_update(update); sync.defer_resume_update_until_append = sync.pending_update.is_some(); diff --git a/codex-rs/thread-store/src/types.rs b/codex-rs/thread-store/src/types.rs index 2e3c60723f..9b171176d2 100644 --- a/codex-rs/thread-store/src/types.rs +++ b/codex-rs/thread-store/src/types.rs @@ -81,6 +81,10 @@ pub struct CreateThreadParams { pub dynamic_tools: Vec, /// Multi-agent runtime selected when the thread was created. pub multi_agent_version: Option, + /// Optional existing rollout JSONL to seed into the new rollout. Local + /// stores may persist this as a copy-on-write parent reference when safe. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub initial_rollout_copy: Option, /// Metadata captured for the newly created thread. pub metadata: ThreadPersistenceMetadata, }