From 793dc92ae00f8d894be172e531cfbf26d62e1f07 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Tue, 31 Mar 2026 18:38:42 +0200 Subject: [PATCH] simplify --- .../app-server/src/codex_message_processor.rs | 13 ++- codex-rs/core/src/lib.rs | 1 + codex-rs/exec/src/lib.rs | 23 +---- codex-rs/rollout/src/file_io.rs | 97 +++++++++++++------ codex-rs/rollout/src/file_io_tests.rs | 38 ++++++-- codex-rs/rollout/src/lib.rs | 1 + codex-rs/rollout/src/list.rs | 25 +++++ codex-rs/rollout/src/recorder.rs | 91 ++++++++++------- codex-rs/rollout/src/recorder_tests.rs | 10 +- codex-rs/rollout/src/tests.rs | 43 ++++++++ codex-rs/tui/src/lib.rs | 79 ++++++++++----- 11 files changed, 302 insertions(+), 119 deletions(-) diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index e4fb6eb250..be6dd31667 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -232,6 +232,7 @@ use codex_core::plugins::PluginUninstallError as CorePluginUninstallError; use codex_core::plugins::load_plugin_apps; use codex_core::plugins::load_plugin_mcp_servers; use codex_core::read_head_for_summary; +use codex_core::read_latest_turn_context; use codex_core::read_session_meta_line; use codex_core::rollout_date_parts; use codex_core::sandboxing::SandboxPermissions; @@ -8177,6 +8178,10 @@ async fn read_history_cwd_from_state_db( thread_id: Option, rollout_path: &Path, ) -> Option { + if let Ok(Some(turn_context)) = read_latest_turn_context(rollout_path).await { + return Some(turn_context.cwd); + } + if let Some(state_db_ctx) = get_state_db(config).await && let Some(thread_id) = thread_id && let Ok(Some(metadata)) = state_db_ctx.get_thread(thread_id).await @@ -8412,6 +8417,12 @@ pub(crate) async fn read_summary_from_rollout( .unwrap_or_else(|| fallback_provider.to_string()); let git_info = git.as_ref().map(map_git_info); let updated_at = updated_at.or_else(|| timestamp.clone()); + let resume_cwd = read_latest_turn_context(path) + .await + .ok() + .flatten() + .map(|turn_context| turn_context.cwd) + .unwrap_or_else(|| session_meta.cwd.clone()); Ok(ConversationSummary { conversation_id: session_meta.id, @@ -8420,7 +8431,7 @@ pub(crate) async fn read_summary_from_rollout( path: path.to_path_buf(), preview: String::new(), model_provider, - cwd: session_meta.cwd, + cwd: resume_cwd, cli_version: session_meta.cli_version, source: session_meta.source, git_info, diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index c35d80e211..7026a1d19e 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -179,6 +179,7 @@ pub use rollout::list::ThreadSortKey; pub use rollout::list::ThreadsPage; pub use rollout::list::parse_cursor; pub use rollout::list::read_head_for_summary; +pub use rollout::list::read_latest_turn_context; pub use rollout::list::read_session_meta_line; pub use rollout::policy::EventPersistenceMode; pub use rollout::read_nonempty_rollout_text; diff --git a/codex-rs/exec/src/lib.rs b/codex-rs/exec/src/lib.rs index 71a37cc6aa..b17fe1aa69 100644 --- a/codex-rs/exec/src/lib.rs +++ b/codex-rs/exec/src/lib.rs @@ -65,6 +65,7 @@ use codex_core::config_loader::LoaderOverrides; use codex_core::config_loader::format_config_error_with_source; use codex_core::format_exec_policy_error_with_source; use codex_core::path_utils; +use codex_core::read_latest_turn_context; use codex_feedback::CodexFeedback; use codex_git_utils::get_git_repo_root; use codex_otel::set_parent_from_context; @@ -74,7 +75,6 @@ use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::ReviewRequest; use codex_protocol::protocol::ReviewTarget; use codex_protocol::protocol::RolloutItem; -use codex_protocol::protocol::RolloutLine; use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::SessionConfiguredEvent; use codex_protocol::protocol::SessionSource; @@ -1116,30 +1116,13 @@ fn all_thread_source_kinds() -> Vec { async fn latest_thread_cwd(thread: &AppServerThread) -> PathBuf { if let Some(path) = thread.path.as_deref() - && let Some(cwd) = parse_latest_turn_context_cwd(path).await + && let Ok(Some(turn_context)) = read_latest_turn_context(path).await { - return cwd; + return turn_context.cwd; } thread.cwd.clone() } -async fn parse_latest_turn_context_cwd(path: &Path) -> Option { - let text = tokio::fs::read_to_string(path).await.ok()?; - for line in text.lines().rev() { - let trimmed = line.trim(); - if trimmed.is_empty() { - continue; - } - let Ok(rollout_line) = serde_json::from_str::(trimmed) else { - continue; - }; - if let RolloutItem::TurnContext(item) = rollout_line.item { - return Some(item.cwd); - } - } - None -} - fn cwds_match(current_cwd: &Path, session_cwd: &Path) -> bool { match ( path_utils::normalize_for_path_comparison(current_cwd), diff --git a/codex-rs/rollout/src/file_io.rs b/codex-rs/rollout/src/file_io.rs index 88857745dd..4e61b1aa19 100644 --- a/codex-rs/rollout/src/file_io.rs +++ b/codex-rs/rollout/src/file_io.rs @@ -3,6 +3,7 @@ use std::fs::OpenOptions; use std::io; use std::io::BufRead; use std::io::BufReader; +use std::io::ErrorKind; use std::io::Read; use std::io::Write; use std::path::Path; @@ -26,20 +27,13 @@ impl RolloutFileEncoding { path.file_name() .and_then(|file_name| file_name.to_str()) .and_then(file_encoding_from_name) - .unwrap_or_else(|| Self::PlainJsonl) - } - - pub(crate) fn preferred_suffix(self) -> &'static str { - match self { - Self::PlainJsonl => ROLLOUT_FILE_SUFFIX, - Self::ZstdJsonl => COMPRESSED_ROLLOUT_FILE_SUFFIX, - } + .unwrap_or(Self::PlainJsonl) } } /// Returns the suffix used for newly created rollout files. pub fn preferred_rollout_file_suffix() -> &'static str { - RolloutFileEncoding::ZstdJsonl.preferred_suffix() + COMPRESSED_ROLLOUT_FILE_SUFFIX } /// Returns true when `name` matches a rollout filename in either supported encoding. @@ -60,23 +54,6 @@ pub fn strip_rollout_file_suffix(name: &str) -> Option<&str> { .or_else(|| name.strip_suffix(ROLLOUT_FILE_SUFFIX)) } -pub(crate) fn append_text(path: &Path, text: &str) -> io::Result<()> { - match RolloutFileEncoding::for_path(path) { - RolloutFileEncoding::PlainJsonl => { - let mut file = OpenOptions::new().append(true).create(true).open(path)?; - file.write_all(text.as_bytes())?; - file.flush() - } - RolloutFileEncoding::ZstdJsonl => { - let file = OpenOptions::new().append(true).create(true).open(path)?; - let mut encoder = zstd::stream::write::Encoder::new(file, DEFAULT_ZSTD_LEVEL)?; - encoder.write_all(text.as_bytes())?; - let mut file = encoder.finish()?; - file.flush() - } - } -} - /// Reads the full rollout file contents, transparently handling plain and zstd-compressed files. pub fn read_rollout_text(path: &Path) -> io::Result { let mut text = String::new(); @@ -87,7 +64,11 @@ pub fn read_rollout_text(path: &Path) -> io::Result { RolloutFileEncoding::ZstdJsonl => { let file = File::open(path)?; let mut decoder = zstd::stream::read::Decoder::new(file)?; - decoder.read_to_string(&mut text)?; + match decoder.read_to_string(&mut text) { + Ok(_) => {} + Err(err) if err.kind() == ErrorKind::UnexpectedEof && !text.is_empty() => {} + Err(err) => return Err(err), + } } } Ok(text) @@ -116,11 +97,64 @@ pub(crate) struct RolloutLineReader { inner: RolloutLineReaderInner, } +pub(crate) struct RolloutAppendWriter { + inner: RolloutAppendWriterInner, +} + enum RolloutLineReaderInner { Plain(BufReader), Zstd(BufReader>>), } +enum RolloutAppendWriterInner { + Plain(File), + Zstd(zstd::stream::write::Encoder<'static, File>), +} + +impl RolloutAppendWriter { + pub(crate) fn open(path: &Path) -> io::Result { + let Some(parent) = path.parent() else { + return Err(io::Error::other(format!( + "rollout path has no parent: {}", + path.display() + ))); + }; + std::fs::create_dir_all(parent)?; + let file = OpenOptions::new().append(true).create(true).open(path)?; + let inner = match RolloutFileEncoding::for_path(path) { + RolloutFileEncoding::PlainJsonl => RolloutAppendWriterInner::Plain(file), + RolloutFileEncoding::ZstdJsonl => RolloutAppendWriterInner::Zstd( + zstd::stream::write::Encoder::new(file, DEFAULT_ZSTD_LEVEL)?, + ), + }; + Ok(Self { inner }) + } + + pub(crate) fn append_text(&mut self, text: &str) -> io::Result<()> { + match &mut self.inner { + RolloutAppendWriterInner::Plain(file) => { + file.write_all(text.as_bytes())?; + file.flush() + } + RolloutAppendWriterInner::Zstd(encoder) => { + encoder.write_all(text.as_bytes())?; + encoder.flush()?; + encoder.get_mut().flush() + } + } + } + + pub(crate) fn finish(self) -> io::Result<()> { + match self.inner { + RolloutAppendWriterInner::Plain(mut file) => file.flush(), + RolloutAppendWriterInner::Zstd(encoder) => { + let mut file = encoder.finish()?; + file.flush() + } + } + } +} + impl RolloutLineReader { pub(crate) fn open(path: &Path) -> io::Result { let inner = match RolloutFileEncoding::for_path(path) { @@ -140,7 +174,14 @@ impl RolloutLineReader { let mut line = String::new(); let bytes_read = match &mut self.inner { RolloutLineReaderInner::Plain(reader) => reader.read_line(&mut line)?, - RolloutLineReaderInner::Zstd(reader) => reader.read_line(&mut line)?, + RolloutLineReaderInner::Zstd(reader) => match reader.read_line(&mut line) { + Ok(bytes_read) => bytes_read, + Err(err) if err.kind() == ErrorKind::UnexpectedEof && !line.is_empty() => { + line.len() + } + Err(err) if err.kind() == ErrorKind::UnexpectedEof => 0, + Err(err) => return Err(err), + }, }; if bytes_read == 0 { return Ok(None); diff --git a/codex-rs/rollout/src/file_io_tests.rs b/codex-rs/rollout/src/file_io_tests.rs index 2fdee0d97a..1444cc466b 100644 --- a/codex-rs/rollout/src/file_io_tests.rs +++ b/codex-rs/rollout/src/file_io_tests.rs @@ -1,7 +1,9 @@ use super::*; +use crate::file_io::RolloutAppendWriter; use crate::file_io::RolloutLineReader; -use crate::file_io::append_text; +use std::fs::File; use std::io; +use std::io::Write; use tempfile::TempDir; #[test] @@ -18,14 +20,14 @@ fn strip_rollout_file_suffix_supports_both_formats() { } #[test] -fn compressed_appends_are_read_back_as_one_stream() -> io::Result<()> { +fn plain_rollouts_are_still_readable() -> io::Result<()> { let temp_dir = TempDir::new()?; let path = temp_dir .path() - .join("rollout-2026-01-01T00-00-00-thread.jsonl.zst"); - - append_text(path.as_path(), "{\"a\":1}\n")?; - append_text(path.as_path(), "{\"b\":2}\n")?; + .join("rollout-2026-01-01T00-00-00-thread.jsonl"); + let mut file = File::create(&path)?; + writeln!(file, "{{\"a\":1}}")?; + writeln!(file, "{{\"b\":2}}")?; assert_eq!(read_rollout_text(path.as_path())?, "{\"a\":1}\n{\"b\":2}\n"); @@ -35,3 +37,27 @@ fn compressed_appends_are_read_back_as_one_stream() -> io::Result<()> { assert_eq!(reader.next_line()?, None); Ok(()) } + +#[test] +fn compressed_appends_are_read_back_from_open_stream() -> io::Result<()> { + let temp_dir = TempDir::new()?; + let path = temp_dir + .path() + .join("rollout-2026-01-01T00-00-00-thread.jsonl.zst"); + + let mut writer = RolloutAppendWriter::open(path.as_path())?; + writer.append_text("{\"a\":1}\n")?; + assert_eq!(read_rollout_text(path.as_path())?, "{\"a\":1}\n"); + + writer.append_text("{\"b\":2}\n")?; + assert_eq!(read_rollout_text(path.as_path())?, "{\"a\":1}\n{\"b\":2}\n"); + + writer.finish()?; + assert_eq!(read_rollout_text(path.as_path())?, "{\"a\":1}\n{\"b\":2}\n"); + + let mut reader = RolloutLineReader::open(path.as_path())?; + assert_eq!(reader.next_line()?, Some("{\"a\":1}".to_string())); + assert_eq!(reader.next_line()?, Some("{\"b\":2}".to_string())); + assert_eq!(reader.next_line()?, None); + Ok(()) +} diff --git a/codex-rs/rollout/src/lib.rs b/codex-rs/rollout/src/lib.rs index c0cc2b4a56..ee11be4df0 100644 --- a/codex-rs/rollout/src/lib.rs +++ b/codex-rs/rollout/src/lib.rs @@ -41,6 +41,7 @@ pub use list::find_archived_thread_path_by_id_str; pub use list::find_thread_path_by_id_str; #[deprecated(note = "use find_thread_path_by_id_str")] pub use list::find_thread_path_by_id_str as find_conversation_path_by_id_str; +pub use list::read_latest_turn_context; pub use list::rollout_date_parts; pub use policy::EventPersistenceMode; pub use recorder::RolloutRecorder; diff --git a/codex-rs/rollout/src/list.rs b/codex-rs/rollout/src/list.rs index 75dc5d8af1..0bce47871f 100644 --- a/codex-rs/rollout/src/list.rs +++ b/codex-rs/rollout/src/list.rs @@ -19,6 +19,7 @@ use super::ARCHIVED_SESSIONS_SUBDIR; use super::SESSIONS_SUBDIR; use super::file_io::RolloutLineReader; use super::file_io::is_rollout_file_name; +use super::file_io::read_rollout_text; use super::file_io::strip_rollout_file_suffix; use crate::protocol::EventMsg; use crate::state_db; @@ -28,6 +29,7 @@ use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::RolloutLine; use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::TurnContextItem; use codex_protocol::protocol::USER_MESSAGE_BEGIN; /// Returned page of thread (thread) summaries. @@ -1161,6 +1163,29 @@ pub async fn read_session_meta_line(path: &Path) -> io::Result }) } +/// Read the latest TurnContext item from a rollout file, if present. +pub async fn read_latest_turn_context(path: &Path) -> io::Result> { + let path = path.to_path_buf(); + tokio::task::spawn_blocking(move || { + let text = read_rollout_text(path.as_path())?; + for line in text.lines().rev() { + let trimmed = line.trim(); + if trimmed.is_empty() { + continue; + } + let Ok(rollout_line) = serde_json::from_str::(trimmed) else { + continue; + }; + if let RolloutItem::TurnContext(item) = rollout_line.item { + return Ok(Some(item)); + } + } + Ok(None) + }) + .await + .map_err(|err| io::Error::other(format!("failed to read latest turn context: {err}")))? +} + async fn file_modified_time(path: &Path) -> io::Result> { let meta = tokio::fs::metadata(path).await?; let modified = meta.modified().ok(); diff --git a/codex-rs/rollout/src/recorder.rs b/codex-rs/rollout/src/recorder.rs index cac83ec872..5d75f610c7 100644 --- a/codex-rs/rollout/src/recorder.rs +++ b/codex-rs/rollout/src/recorder.rs @@ -25,8 +25,8 @@ use tracing::warn; use super::ARCHIVED_SESSIONS_SUBDIR; use super::SESSIONS_SUBDIR; +use super::file_io::RolloutAppendWriter; use super::file_io::RolloutLineReader; -use super::file_io::append_text as append_rollout_text; use super::file_io::preferred_rollout_file_suffix; use super::list::Cursor; use super::list::ThreadItem; @@ -38,6 +38,7 @@ use super::list::get_threads; use super::list::get_threads_in_root; use super::list::parse_cursor; use super::list::parse_timestamp_uuid_from_filename; +use super::list::read_latest_turn_context; use super::metadata; use super::policy::EventPersistenceMode; use super::policy::is_persisted_response_item; @@ -434,7 +435,7 @@ impl RolloutRecorder { } => { tokio::fs::metadata(&path).await?; ( - Some(RolloutFileWriter::new(path.clone())), + Some(RolloutFileWriter::open(path.clone()).await?), None, path, None, @@ -782,7 +783,7 @@ async fn rollout_writer( "deferred rollout recorder missing log file metadata", )); }; - writer = Some(RolloutFileWriter::new(log_file_info.path)); + writer = Some(RolloutFileWriter::open(log_file_info.path).await?); if let Some(session_meta) = meta.take() { write_session_meta( @@ -826,11 +827,22 @@ async fn rollout_writer( let _ = ack.send(()); } RolloutCmd::Shutdown { ack } => { + if let Some(writer) = writer.as_mut() + && let Err(err) = writer.finish().await + { + let _ = ack.send(()); + return Err(err); + } let _ = ack.send(()); + return Ok(()); } } } + if let Some(writer) = writer.as_mut() { + writer.finish().await?; + } + Ok(()) } @@ -947,7 +959,11 @@ async fn sync_thread_state_after_write( } struct RolloutFileWriter { - path: PathBuf, + inner: Option, +} + +struct BlockingRolloutFileWriter { + writer: RolloutAppendWriter, } #[derive(serde::Serialize)] @@ -958,8 +974,15 @@ struct RolloutLineRef<'a> { } impl RolloutFileWriter { - fn new(path: PathBuf) -> Self { - Self { path } + async fn open(path: PathBuf) -> std::io::Result { + tokio::task::spawn_blocking(move || { + let writer = RolloutAppendWriter::open(path.as_path())?; + Ok(Self { + inner: Some(BlockingRolloutFileWriter { writer }), + }) + }) + .await + .map_err(|err| IoError::other(format!("failed to open rollout writer: {err}")))? } async fn write_rollout_item(&mut self, rollout_item: &RolloutItem) -> std::io::Result<()> { @@ -975,7 +998,7 @@ impl RolloutFileWriter { let timestamp_format: &[FormatItem] = format_description!( "[year]-[month]-[day]T[hour]:[minute]:[second].[subsecond digits:3]Z" ); - let mut batch = String::new(); + let mut lines = Vec::with_capacity(rollout_items.len()); for rollout_item in rollout_items { let timestamp = OffsetDateTime::now_utc() .format(timestamp_format) @@ -986,25 +1009,31 @@ impl RolloutFileWriter { }; let mut json = serde_json::to_string(&line)?; json.push('\n'); - batch.push_str(json.as_str()); + lines.push(json); } - self.write_text(batch).await - } - - async fn write_text(&mut self, text: String) -> std::io::Result<()> { - let path = self.path.clone(); - tokio::task::spawn_blocking(move || { - let Some(parent) = path.parent() else { - return Err(IoError::other(format!( - "rollout path has no parent: {}", - path.display() - ))); - }; - std::fs::create_dir_all(parent)?; - append_rollout_text(path.as_path(), text.as_str()) + let Some(writer) = self.inner.take() else { + return Err(IoError::other("rollout writer missing inner state")); + }; + let writer = tokio::task::spawn_blocking(move || { + let mut writer = writer; + for line in lines { + writer.writer.append_text(line.as_str())?; + } + Ok::(writer) }) .await - .map_err(|err| IoError::other(format!("failed to write rollout line: {err}")))? + .map_err(|err| IoError::other(format!("failed to write rollout line: {err}")))??; + self.inner = Some(writer); + Ok(()) + } + + async fn finish(&mut self) -> std::io::Result<()> { + let Some(writer) = self.inner.take() else { + return Ok(()); + }; + tokio::task::spawn_blocking(move || writer.writer.finish()) + .await + .map_err(|err| IoError::other(format!("failed to finish rollout writer: {err}")))? } } @@ -1074,20 +1103,12 @@ async fn resume_candidate_matches_cwd( cwd: &Path, default_provider: &str, ) -> bool { - if cached_cwd.is_some_and(|session_cwd| cwd_matches(session_cwd, cwd)) { - return true; + if let Ok(Some(turn_context)) = read_latest_turn_context(rollout_path).await { + return cwd_matches(turn_context.cwd.as_path(), cwd); } - if let Ok((items, _, _)) = RolloutRecorder::load_rollout_items(rollout_path).await - && let Some(latest_turn_context_cwd) = items.iter().rev().find_map(|item| match item { - RolloutItem::TurnContext(turn_context) => Some(turn_context.cwd.as_path()), - RolloutItem::SessionMeta(_) - | RolloutItem::ResponseItem(_) - | RolloutItem::Compacted(_) - | RolloutItem::EventMsg(_) => None, - }) - { - return cwd_matches(latest_turn_context_cwd, cwd); + if cached_cwd.is_some_and(|session_cwd| cwd_matches(session_cwd, cwd)) { + return true; } metadata::extract_metadata_from_rollout(rollout_path, default_provider) diff --git a/codex-rs/rollout/src/recorder_tests.rs b/codex-rs/rollout/src/recorder_tests.rs index 1214122ad6..0ef39edf5c 100644 --- a/codex-rs/rollout/src/recorder_tests.rs +++ b/codex-rs/rollout/src/recorder_tests.rs @@ -2,7 +2,7 @@ use super::*; use crate::config::RolloutConfig; -use crate::file_io::append_text; +use crate::file_io::RolloutAppendWriter; use crate::file_io::read_rollout_text; use chrono::TimeZone; use codex_protocol::config_types::ReasoningSummary as ReasoningSummaryConfig; @@ -518,8 +518,10 @@ async fn load_rollout_items_preserves_parsed_items_when_zstd_tail_is_truncated() }, }); - append_text(&rollout_path, &format!("{session_meta}\n"))?; - append_text(&rollout_path, &format!("{user_event}\n"))?; + let mut writer = RolloutAppendWriter::open(&rollout_path)?; + writer.append_text(&format!("{session_meta}\n"))?; + writer.append_text(&format!("{user_event}\n"))?; + drop(writer); let file = File::options().write(true).open(&rollout_path)?; let truncated_len = file.metadata()?.len().saturating_sub(1); @@ -544,7 +546,7 @@ async fn load_rollout_items_preserves_parsed_items_when_zstd_tail_is_truncated() Some("test-provider") ); assert!(session_meta_line.git.is_none()); - assert_eq!(parse_errors, 1); + assert_eq!(parse_errors, 0); Ok(()) } diff --git a/codex-rs/rollout/src/tests.rs b/codex-rs/rollout/src/tests.rs index 8bb943b2a5..28296fcbc8 100644 --- a/codex-rs/rollout/src/tests.rs +++ b/codex-rs/rollout/src/tests.rs @@ -27,6 +27,7 @@ use crate::list::ThreadSortKey; use crate::list::ThreadsPage; use crate::list::get_threads; use crate::list::read_head_for_summary; +use crate::list::read_latest_turn_context; use crate::rollout_date_parts; use anyhow::Result; use codex_protocol::ThreadId; @@ -38,6 +39,7 @@ use codex_protocol::protocol::RolloutLine; use codex_protocol::protocol::SessionMeta; use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::TurnContextItem; use codex_protocol::protocol::UserMessageEvent; const NO_SOURCE_FILTER: &[SessionSource] = &[]; @@ -90,6 +92,47 @@ async fn insert_state_db_thread( .expect("state db upsert should succeed"); } +#[tokio::test] +async fn read_latest_turn_context_reads_compressed_rollout() -> Result<()> { + let temp = TempDir::new()?; + let path = temp.path().join("rollout.jsonl.zst"); + let mut encoder = zstd::stream::write::Encoder::new(File::create(&path)?, 0)?; + for (idx, cwd) in ["/tmp/first", "/tmp/second"].into_iter().enumerate() { + let line = RolloutLine { + timestamp: format!("t{idx}"), + item: RolloutItem::TurnContext(TurnContextItem { + turn_id: Some(format!("turn-{idx}")), + trace_id: None, + cwd: cwd.into(), + current_date: None, + timezone: None, + approval_policy: codex_protocol::protocol::AskForApproval::OnRequest, + sandbox_policy: codex_protocol::protocol::SandboxPolicy::new_read_only_policy(), + network: None, + model: format!("model-{idx}"), + personality: None, + collaboration_mode: None, + realtime_active: None, + effort: None, + summary: codex_protocol::config_types::ReasoningSummary::Auto, + user_instructions: None, + developer_instructions: None, + final_output_json_schema: None, + truncation_policy: None, + }), + }; + writeln!(encoder, "{}", serde_json::to_string(&line)?)?; + } + let _ = encoder.finish()?; + + let turn_context = read_latest_turn_context(path.as_path()) + .await? + .expect("latest turn context"); + assert_eq!(turn_context.cwd, Path::new("/tmp/second")); + assert_eq!(turn_context.model, "model-1"); + Ok(()) +} + // TODO(jif) fix // #[tokio::test] // async fn list_threads_prefers_state_db_when_available() { diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index 02b743fa8a..d8a69576da 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -37,6 +37,7 @@ use codex_core::config_loader::format_config_error_with_source; use codex_core::default_client::set_default_client_residency_requirement; use codex_core::format_exec_policy_error_with_source; use codex_core::path_utils; +use codex_core::read_latest_turn_context; use codex_core::read_session_meta_line; use codex_core::state_db::get_state_db; use codex_core::windows_sandbox::WindowsSandboxLevelExt; @@ -45,9 +46,7 @@ use codex_protocol::config_types::AltScreenMode; use codex_protocol::config_types::SandboxMode; use codex_protocol::config_types::WindowsSandboxLevel; use codex_protocol::protocol::AskForApproval; -use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::RolloutLine; -use codex_protocol::protocol::TurnContextItem; use codex_state::log_db; use codex_terminal_detection::Multiplexer; use codex_terminal_detection::terminal_info; @@ -1373,6 +1372,12 @@ pub(crate) async fn read_session_cwd( thread_id: ThreadId, path: Option<&Path>, ) -> Option { + if let Some(path) = path + && let Ok(Some(turn_context)) = read_latest_turn_context(path).await + { + return Some(turn_context.cwd); + } + if let Some(state_db_ctx) = get_state_db(config).await && let Ok(Some(metadata)) = state_db_ctx.get_thread(thread_id).await { @@ -1380,14 +1385,12 @@ pub(crate) async fn read_session_cwd( } // Prefer the latest TurnContext cwd so resume/fork reflects the most recent - // session directory (for the changed-cwd prompt) when DB data is unavailable. + // session directory (for the changed-cwd prompt). When the rollout cannot be + // read, fall back to persisted thread metadata and then session metadata. // The alternative would be mutating the SessionMeta line when the session cwd // changes, but the rollout is an append-only JSONL log and rewriting the head // would be error-prone. let path = path?; - if let Some(cwd) = read_latest_turn_context(path).await.map(|item| item.cwd) { - return Some(cwd); - } match read_session_meta_line(path).await { Ok(meta_line) => Some(meta_line.meta.cwd), Err(err) => { @@ -1407,6 +1410,12 @@ pub(crate) async fn read_session_model( thread_id: ThreadId, path: Option<&Path>, ) -> Option { + if let Some(path) = path + && let Ok(Some(turn_context)) = read_latest_turn_context(path).await + { + return Some(turn_context.model); + } + if let Some(state_db_ctx) = get_state_db(config).await && let Ok(Some(metadata)) = state_db_ctx.get_thread(thread_id).await && let Some(model) = metadata.model @@ -1414,27 +1423,8 @@ pub(crate) async fn read_session_model( return Some(model); } - let path = path?; - read_latest_turn_context(path).await.map(|item| item.model) -} - -async fn read_latest_turn_context(path: &Path) -> Option { - let text = tokio::fs::read_to_string(path).await.ok()?; - for line in text.lines().rev() { - let trimmed = line.trim(); - if trimmed.is_empty() { - continue; - } - let Ok(rollout_line) = serde_json::from_str::(trimmed) else { - continue; - }; - if let RolloutItem::TurnContext(item) = rollout_line.item { - return Some(item); - } - } None } - pub(crate) fn cwds_differ(current_cwd: &Path, session_cwd: &Path) -> bool { match ( path_utils::normalize_for_path_comparison(current_cwd), @@ -2044,6 +2034,45 @@ mod tests { Ok(()) } + #[tokio::test] + async fn read_session_cwd_prefers_latest_turn_context_from_compressed_rollout() + -> std::io::Result<()> { + let temp_dir = TempDir::new()?; + let config = build_config(&temp_dir).await?; + let first = temp_dir.path().join("first"); + let second = temp_dir.path().join("second"); + std::fs::create_dir_all(&first)?; + std::fs::create_dir_all(&second)?; + + let rollout_path = temp_dir.path().join("rollout.jsonl.zst"); + let mut encoder = + zstd::stream::write::Encoder::new(std::fs::File::create(&rollout_path)?, 0)?; + for line in [ + RolloutLine { + timestamp: "t0".to_string(), + item: RolloutItem::TurnContext(build_turn_context(&config, first)), + }, + RolloutLine { + timestamp: "t1".to_string(), + item: RolloutItem::TurnContext(build_turn_context(&config, second.clone())), + }, + ] { + use std::io::Write; + writeln!( + encoder, + "{}", + serde_json::to_string(&line).expect("serialize rollout") + )?; + } + let _ = encoder.finish()?; + + let cwd = read_session_cwd(&config, ThreadId::new(), Some(&rollout_path)) + .await + .expect("expected cwd"); + assert_eq!(cwd, second); + Ok(()) + } + #[tokio::test] async fn should_prompt_when_meta_matches_current_but_latest_turn_differs() -> std::io::Result<()> {