reduce resume and fork startup work

This commit is contained in:
anais
2026-06-15 22:23:26 -07:00
parent 4b7746130b
commit 93f761da72
10 changed files with 501 additions and 90 deletions
@@ -2600,10 +2600,11 @@ impl ThreadRequestProcessor {
self.resume_thread_from_history(history.as_slice())
.await
.map(|thread_history| (thread_history, None))
} 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)))
} else if let Some(mut stored_thread) = stored_thread_from_running_probe {
match self.take_stored_thread_initial_history(&mut stored_thread) {
Ok(thread_history) => Ok((thread_history, Some(*stored_thread))),
Err(error) => Err(error),
}
} else {
self.resume_thread_from_rollout(&thread_id, path.as_ref())
.await
@@ -2751,9 +2752,10 @@ impl ThreadRequestProcessor {
config_snapshot.active_permission_profile,
);
let token_usage_thread = include_turns.then(|| thread.clone());
let response_history_items = initial_history_rollout_items(&response_history);
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(),
response_history_items,
thread.status.clone(),
/*has_live_running_thread*/ false,
/*active_turn*/ None,
@@ -2791,15 +2793,18 @@ impl ThreadRequestProcessor {
initial_turns_page,
};
let connection_id = request_id.connection_id;
self.outgoing.send_response(request_id, response).await;
// `excludeTurns` is explicitly the cheap resume path, so avoid
// 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(),
let token_usage_replay = token_usage_thread.map(|token_usage_thread| {
let turn_id = latest_token_usage_turn_id_from_rollout_items(
response_history_items,
token_usage_thread.turns.as_slice(),
);
(token_usage_thread, turn_id)
});
let connection_id = request_id.connection_id;
self.outgoing.send_response(request_id, response).await;
if let Some((token_usage_thread, token_usage_turn_id)) = token_usage_replay {
// The client needs restored usage before it starts another turn.
// Sending after the response preserves JSON-RPC request ordering while
// still filling the status line before the next turn lifecycle begins.
@@ -2955,10 +2960,11 @@ impl ThreadRequestProcessor {
}
let redact_resume_payloads =
should_redact_thread_resume_payloads(app_server_client_name.as_deref());
let mut source_thread = source_thread;
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 {existing_thread_id} did not include persisted history"
@@ -3054,12 +3060,10 @@ impl ThreadRequestProcessor {
thread_id: &str,
path: Option<&PathBuf>,
) -> Result<(InitialHistory, StoredThread), JSONRPCErrorError> {
let stored_thread = self
let mut stored_thread = self
.read_stored_thread_for_resume(thread_id, path, /*include_history*/ true)
.await?;
let history = self
.stored_thread_to_initial_history(&stored_thread)
.await?;
let history = self.take_stored_thread_initial_history(&mut stored_thread)?;
Ok((history, stored_thread))
}
@@ -3104,15 +3108,15 @@ impl ThreadRequestProcessor {
}
#[tracing::instrument(level = "trace", skip_all)]
async fn stored_thread_to_initial_history(
fn take_stored_thread_initial_history(
&self,
stored_thread: &StoredThread,
stored_thread: &mut StoredThread,
) -> Result<InitialHistory, JSONRPCErrorError> {
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"
@@ -3172,33 +3176,10 @@ impl ThreadRequestProcessor {
let thread = match thread_history {
InitialHistory::Resumed(resumed) => {
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
})
};
if let Some(mut stored_thread) = resume_source_thread {
if stored_thread.preview.is_empty() {
stored_thread.preview = preview_from_rollout_items(&resumed.history);
}
Ok(thread_from_stored_thread(
stored_thread,
fallback_provider,
@@ -3246,10 +3227,10 @@ impl ThreadRequestProcessor {
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 = initial_history_rollout_items(thread_history);
populate_thread_turns_from_history(
&mut thread,
&history_items,
history_items,
/*active_turn*/ None,
);
}
@@ -3258,6 +3239,9 @@ impl ThreadRequestProcessor {
}
async fn attach_thread_name(&self, thread_id: ThreadId, thread: &mut Thread) {
if thread.name.is_some() {
return;
}
if let Ok(stored_thread) = self
.thread_store
.read_thread(StoreReadThreadParams {
@@ -3306,7 +3290,7 @@ impl ThreadRequestProcessor {
"`permissions` cannot be combined with `sandbox`",
));
}
let source_thread = self
let mut source_thread = self
.read_stored_thread_for_resume(&thread_id, path.as_ref(), /*include_history*/ true)
.await?;
let source_thread_id = source_thread.thread_id;
@@ -3316,13 +3300,14 @@ impl ThreadRequestProcessor {
.and_then(codex_core::util::normalize_thread_name);
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"
))
})?;
let response_history_items = (ephemeral || include_turns).then(|| history_items.clone());
let history_cwd = Some(source_thread.cwd.clone());
// Persist Windows sandbox mode.
@@ -3384,7 +3369,7 @@ impl ThreadRequestProcessor {
config,
InitialHistory::Resumed(ResumedHistory {
conversation_id: source_thread_id,
history: history_items.clone(),
history: history_items,
rollout_path: source_thread.rollout_path.clone(),
}),
thread_source.map(Into::into),
@@ -3449,18 +3434,19 @@ impl ThreadRequestProcessor {
)
} else {
let config_snapshot = forked_thread.config_snapshot().await;
let history_items = response_history_items.as_deref().unwrap_or(&[]);
let mut thread = build_thread_from_snapshot(
thread_id,
session_configured.session_id.to_string(),
&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,
);
}
@@ -3517,7 +3503,7 @@ impl ThreadRequestProcessor {
// instead of rebuilding history only to attribute a historical update.
if let Some(token_usage_thread) = token_usage_thread {
let token_usage_turn_id = latest_token_usage_turn_id_from_rollout_items(
&history_items,
response_history_items.as_deref().unwrap_or(&[]),
token_usage_thread.turns.as_slice(),
);
// Mirror the resume contract for forks: the new thread is usable as soon
@@ -3726,6 +3712,14 @@ fn thread_backwards_cursor_for_sort_key(
Some(timestamp.to_rfc3339_opts(SecondsFormat::Millis, true))
}
fn initial_history_rollout_items(history: &InitialHistory) -> &[RolloutItem] {
match history {
InitialHistory::New | InitialHistory::Cleared => &[],
InitialHistory::Resumed(resumed) => resumed.history.as_slice(),
InitialHistory::Forked(items) => items.as_slice(),
}
}
struct ThreadTurnsPage {
pub(super) turns: Vec<Turn>,
pub(super) next_cursor: Option<String>,
@@ -4300,6 +4294,25 @@ fn preview_from_rollout_items(items: &[RolloutItem]) -> String {
Some(codex_protocol::items::TurnItem::UserMessage(user)) => Some(user.message()),
_ => None,
},
RolloutItem::EventMsg(codex_protocol::protocol::EventMsg::UserMessage(event)) => {
let message = event.message.trim();
if !message.is_empty() {
Some(event.message.clone())
} else if event
.images
.as_ref()
.is_some_and(|images| !images.is_empty())
|| !event.local_images.is_empty()
{
Some("[Image]".to_string())
} else {
None
}
}
RolloutItem::EventMsg(codex_protocol::protocol::EventMsg::ThreadGoalUpdated(event)) => {
let objective = event.goal.objective.trim();
(!objective.is_empty()).then(|| objective.to_string())
}
_ => None,
})
.map(|preview| match preview.find(USER_MESSAGE_BEGIN) {
+22 -2
View File
@@ -915,8 +915,18 @@ impl App {
(chat_widget, None)
}
SessionSelection::Resume(target_session) => {
let rollout_path = match target_session.path.as_deref() {
Some(path) => {
crate::session_resume::verified_rollout_path_for_thread(
path,
target_session.thread_id,
)
.await
}
None => None,
};
let resumed = app_server
.resume_thread(config.clone(), target_session.thread_id)
.resume_thread_with_path(config.clone(), target_session.thread_id, rollout_path)
.await
.map_err(|err| session_start_error("resume", &target_session, err))?;
let init = crate::chatwidget::ChatWidgetInit {
@@ -954,8 +964,18 @@ impl App {
/*inc*/ 1,
&[("source", "cli_subcommand")],
);
let rollout_path = match target_session.path.as_deref() {
Some(path) => {
crate::session_resume::verified_rollout_path_for_thread(
path,
target_session.thread_id,
)
.await
}
None => None,
};
let forked = app_server
.fork_thread(config.clone(), target_session.thread_id)
.fork_thread_with_path(config.clone(), target_session.thread_id, rollout_path)
.await
.map_err(|err| session_start_error("fork", &target_session, err))?;
let init = crate::chatwidget::ChatWidgetInit {
+8 -1
View File
@@ -137,7 +137,14 @@ impl App {
tui.frame_requester().schedule_frame();
}
AppEvent::ResumeSessionByIdOrName(id_or_name) => {
match crate::lookup_session_target_with_app_server(app_server, &id_or_name).await? {
match crate::lookup_session_target_with_app_server(
app_server,
&id_or_name,
self.config.codex_home.as_path(),
self.state_db.as_deref(),
)
.await?
{
Some(target_session) => {
return self
.resume_target_session(tui, app_server, target_session)
+16 -3
View File
@@ -747,8 +747,7 @@ impl App {
tui,
self.state_db.as_deref(),
&current_cwd,
target_session.thread_id,
target_session.path.as_deref(),
&target_session,
CwdPromptAction::Resume,
/*allow_prompt*/ true,
)
@@ -782,8 +781,22 @@ impl App {
self.chat_widget.thread_name(),
self.chat_widget.rollout_path().as_deref(),
);
let rollout_path = match target_session.path.as_deref() {
Some(path) => {
crate::session_resume::verified_rollout_path_for_thread(
path,
target_session.thread_id,
)
.await
}
None => None,
};
match app_server
.resume_thread(resume_config.clone(), target_session.thread_id)
.resume_thread_with_path(
resume_config.clone(),
target_session.thread_id,
rollout_path,
)
.await
{
Ok(resumed) => {
+1
View File
@@ -4879,6 +4879,7 @@ fn session_start_error_surfaces_archived_guidance_without_rollout_path() {
"/Users/me/.codex/archived_sessions/rollout.jsonl",
)),
thread_id,
cwd: None,
};
let expected = format!(
"session {thread_id} is archived. Run `codex unarchive {thread_id}` to unarchive it first."
+8
View File
@@ -19,6 +19,7 @@ fn startup_waiting_gate_is_only_for_fresh_or_exit_session_selection() {
crate::resume_picker::SessionTarget {
path: Some(PathBuf::from("/tmp/restore")),
thread_id: ThreadId::new(),
cwd: None,
}
)),
false
@@ -28,6 +29,7 @@ fn startup_waiting_gate_is_only_for_fresh_or_exit_session_selection() {
crate::resume_picker::SessionTarget {
path: Some(PathBuf::from("/tmp/fork")),
thread_id: ThreadId::new(),
cwd: None,
}
)),
false
@@ -39,10 +41,12 @@ fn startup_paused_goal_prompt_gate_is_only_for_quiet_resume() {
let resume = SessionSelection::Resume(crate::resume_picker::SessionTarget {
path: Some(PathBuf::from("/tmp/restore")),
thread_id: ThreadId::new(),
cwd: None,
});
let fork = SessionSelection::Fork(crate::resume_picker::SessionTarget {
path: Some(PathBuf::from("/tmp/fork")),
thread_id: ThreadId::new(),
cwd: None,
});
let no_images: Vec<PathBuf> = Vec::new();
let initial_images = vec![PathBuf::from("/tmp/image.png")];
@@ -111,6 +115,7 @@ fn startup_waiting_gate_not_applied_for_resume_or_fork_session_selection() {
crate::resume_picker::SessionTarget {
path: Some(PathBuf::from("/tmp/restore")),
thread_id: ThreadId::new(),
cwd: None,
},
));
assert_eq!(
@@ -124,6 +129,7 @@ fn startup_waiting_gate_not_applied_for_resume_or_fork_session_selection() {
crate::resume_picker::SessionTarget {
path: Some(PathBuf::from("/tmp/fork")),
thread_id: ThreadId::new(),
cwd: None,
},
));
assert_eq!(
@@ -266,6 +272,7 @@ async fn ignore_same_thread_resume_reports_noop_for_current_thread() {
let ignored = app.ignore_same_thread_resume(&crate::resume_picker::SessionTarget {
path: Some(test_path_buf("/tmp/project")),
thread_id,
cwd: None,
});
assert!(ignored);
@@ -290,6 +297,7 @@ async fn ignore_same_thread_resume_allows_reattaching_displayed_inactive_thread(
let ignored = app.ignore_same_thread_resume(&crate::resume_picker::SessionTarget {
path: Some(test_path_buf("/tmp/project")),
thread_id,
cwd: None,
});
assert!(!ignored);
+37
View File
@@ -454,6 +454,16 @@ impl AppServerSession {
&mut self,
config: Config,
thread_id: ThreadId,
) -> Result<AppServerStartedThread> {
self.resume_thread_with_path(config, thread_id, /*path*/ None)
.await
}
pub(crate) async fn resume_thread_with_path(
&mut self,
config: Config,
thread_id: ThreadId,
path: Option<std::path::PathBuf>,
) -> Result<AppServerStartedThread> {
let request_id = self.next_request_id();
let session_config = self.session_config_with_effective_service_tier(&config);
@@ -464,6 +474,7 @@ impl AppServerSession {
params: thread_resume_params_from_config(
session_config,
thread_id,
path,
self.thread_params_mode(),
self.remote_cwd_override.as_deref(),
),
@@ -486,6 +497,16 @@ impl AppServerSession {
&mut self,
config: Config,
thread_id: ThreadId,
) -> Result<AppServerStartedThread> {
self.fork_thread_with_path(config, thread_id, /*path*/ None)
.await
}
pub(crate) async fn fork_thread_with_path(
&mut self,
config: Config,
thread_id: ThreadId,
path: Option<std::path::PathBuf>,
) -> Result<AppServerStartedThread> {
let request_id = self.next_request_id();
let session_config = self.session_config_with_effective_service_tier(&config);
@@ -496,6 +517,7 @@ impl AppServerSession {
params: thread_fork_params_from_config(
session_config,
thread_id,
path,
self.thread_params_mode(),
self.remote_cwd_override.as_deref(),
),
@@ -1404,6 +1426,7 @@ fn thread_start_params_from_config(
fn thread_resume_params_from_config(
config: Config,
thread_id: ThreadId,
path: Option<std::path::PathBuf>,
thread_params_mode: ThreadParamsMode,
remote_cwd_override: Option<&std::path::Path>,
) -> ThreadResumeParams {
@@ -1419,6 +1442,7 @@ fn thread_resume_params_from_config(
.flatten();
ThreadResumeParams {
thread_id: thread_id.to_string(),
path,
model: config.model.clone(),
model_provider: thread_params_mode.model_provider_from_config(&config),
service_tier: service_tier_override_from_config(&config),
@@ -1439,6 +1463,7 @@ fn thread_resume_params_from_config(
fn thread_fork_params_from_config(
config: Config,
thread_id: ThreadId,
path: Option<std::path::PathBuf>,
thread_params_mode: ThreadParamsMode,
remote_cwd_override: Option<&std::path::Path>,
) -> ThreadForkParams {
@@ -1454,6 +1479,7 @@ fn thread_fork_params_from_config(
.flatten();
ThreadForkParams {
thread_id: thread_id.to_string(),
path,
model: config.model.clone(),
model_provider: thread_params_mode.model_provider_from_config(&config),
service_tier: service_tier_override_from_config(&config),
@@ -1989,12 +2015,14 @@ mod tests {
let resume = thread_resume_params_from_config(
config.clone(),
thread_id,
/*path*/ None,
ThreadParamsMode::Remote,
/*remote_cwd_override*/ None,
);
let fork = thread_fork_params_from_config(
config,
thread_id,
/*path*/ None,
ThreadParamsMode::Remote,
/*remote_cwd_override*/ None,
);
@@ -2106,12 +2134,14 @@ mod tests {
let resume = thread_resume_params_from_config(
config.clone(),
thread_id,
/*path*/ None,
ThreadParamsMode::Remote,
Some(remote_cwd.as_path()),
);
let fork = thread_fork_params_from_config(
config,
thread_id,
/*path*/ None,
ThreadParamsMode::Remote,
Some(remote_cwd.as_path()),
);
@@ -2157,12 +2187,14 @@ mod tests {
let resume = thread_resume_params_from_config(
config.clone(),
thread_id,
/*path*/ None,
ThreadParamsMode::Embedded,
/*remote_cwd_override*/ None,
);
let fork = thread_fork_params_from_config(
config,
thread_id,
/*path*/ None,
ThreadParamsMode::Embedded,
/*remote_cwd_override*/ None,
);
@@ -2217,6 +2249,7 @@ mod tests {
let params = thread_fork_params_from_config(
config,
thread_id,
/*path*/ None,
ThreadParamsMode::Embedded,
/*remote_cwd_override*/ None,
);
@@ -2244,12 +2277,14 @@ mod tests {
let control_resume = thread_resume_params_from_config(
config.clone(),
thread_id,
/*path*/ None,
ThreadParamsMode::Embedded,
/*remote_cwd_override*/ None,
);
let control_fork = thread_fork_params_from_config(
config.clone(),
thread_id,
/*path*/ None,
ThreadParamsMode::Embedded,
/*remote_cwd_override*/ None,
);
@@ -2273,12 +2308,14 @@ mod tests {
let treatment_resume = thread_resume_params_from_config(
config.clone(),
thread_id,
/*path*/ None,
ThreadParamsMode::Embedded,
/*remote_cwd_override*/ None,
);
let treatment_fork = thread_fork_params_from_config(
config,
thread_id,
/*path*/ None,
ThreadParamsMode::Embedded,
/*remote_cwd_override*/ None,
);
+112 -8
View File
@@ -585,6 +585,7 @@ fn session_target_from_app_server_thread(
Ok(thread_id) => Some(resume_picker::SessionTarget {
path: thread.path,
thread_id,
cwd: Some(thread.cwd.to_path_buf()),
}),
Err(err) => {
warn!(
@@ -597,6 +598,41 @@ fn session_target_from_app_server_thread(
}
}
async fn lookup_local_session_target_by_id(
thread_id: ThreadId,
codex_home: &Path,
state_db_ctx: Option<&codex_state::StateRuntime>,
) -> Option<resume_picker::SessionTarget> {
let state_db_ctx = state_db_ctx?;
let id = thread_id.to_string();
let metadata = state_db_ctx.get_thread(thread_id).await.ok().flatten();
if let Some(metadata) = metadata.as_ref()
&& let Some(verified_path) = crate::session_resume::verified_rollout_path_for_thread(
metadata.rollout_path.as_path(),
thread_id,
)
.await
{
return Some(resume_picker::SessionTarget {
path: Some(verified_path),
thread_id,
cwd: Some(metadata.cwd.clone()),
});
}
let path =
codex_rollout::find_any_thread_path_by_id_str(codex_home, id.as_str(), Some(state_db_ctx))
.await
.ok()
.flatten()?;
let verified_path =
crate::session_resume::verified_rollout_path_for_thread(path.as_path(), thread_id).await;
Some(resume_picker::SessionTarget {
path: verified_path,
thread_id,
cwd: metadata.map(|metadata| metadata.cwd),
})
}
pub(crate) fn resume_source_kinds(include_non_interactive: bool) -> Vec<ThreadSourceKind> {
let mut source_kinds = vec![ThreadSourceKind::Cli, ThreadSourceKind::VsCode];
if include_non_interactive {
@@ -646,6 +682,8 @@ async fn lookup_session_target_by_name_with_app_server(
async fn lookup_session_target_with_app_server(
app_server: &mut AppServerSession,
id_or_name: &str,
codex_home: &Path,
state_db_ctx: Option<&codex_state::StateRuntime>,
) -> color_eyre::Result<Option<resume_picker::SessionTarget>> {
if Uuid::parse_str(id_or_name).is_ok() {
let thread_id = match ThreadId::from_string(id_or_name) {
@@ -659,6 +697,12 @@ async fn lookup_session_target_with_app_server(
return Ok(None);
}
};
if !app_server.uses_remote_workspace()
&& let Some(target) =
lookup_local_session_target_by_id(thread_id, codex_home, state_db_ctx).await
{
return Ok(Some(target));
}
return match app_server
.thread_read(thread_id, /*include_turns*/ false)
.await
@@ -702,10 +746,10 @@ async fn lookup_latest_session_target_with_app_server(
.data
.into_iter()
.find_map(session_target_from_app_server_thread);
if target.as_ref().is_some_and(|target| {
uses_remote_workspace || target.path.as_deref().is_some_and(std::path::Path::exists)
}) {
return Ok(target);
if let Some(target) = target
&& (uses_remote_workspace || target.path.is_some())
{
return Ok(Some(target));
}
}
Ok(None)
@@ -1474,7 +1518,14 @@ async fn run_ratatui_app(
let Some(startup_app_server) = app_server.as_mut() else {
unreachable!("app server should be initialized for --fork <id>");
};
match lookup_session_target_with_app_server(startup_app_server, id_str).await? {
match lookup_session_target_with_app_server(
startup_app_server,
id_str,
config.codex_home.as_path(),
state_db.as_deref(),
)
.await?
{
Some(target_session) => resume_picker::SessionSelection::Fork(target_session),
None => {
shutdown_app_server_if_present(app_server.take()).await;
@@ -1531,7 +1582,14 @@ async fn run_ratatui_app(
let Some(startup_app_server) = app_server.as_mut() else {
unreachable!("app server should be initialized for --resume <id>");
};
match lookup_session_target_with_app_server(startup_app_server, id_str).await? {
match lookup_session_target_with_app_server(
startup_app_server,
id_str,
config.codex_home.as_path(),
state_db.as_deref(),
)
.await?
{
Some(target_session) => resume_picker::SessionSelection::Resume(target_session),
None => {
shutdown_app_server_if_present(app_server.take()).await;
@@ -1609,8 +1667,7 @@ async fn run_ratatui_app(
&mut tui,
state_db.as_deref(),
&current_cwd,
target_session.thread_id,
target_session.path.as_deref(),
target_session,
action,
allow_prompt,
)
@@ -2152,6 +2209,7 @@ mod tests {
let target = crate::resume_picker::SessionTarget {
path: None,
thread_id,
cwd: None,
};
assert_eq!(target.display_label(), format!("thread {thread_id}"));
@@ -2849,6 +2907,52 @@ mod tests {
.await
}
#[tokio::test]
async fn lookup_local_session_target_by_id_uses_state_db_path() -> color_eyre::Result<()> {
let temp_dir = TempDir::new()?;
let config = build_config(&temp_dir).await?;
let uuid = Uuid::new_v4();
let thread_id = ThreadId::from_string(&uuid.to_string())?;
let rollout_path = temp_dir
.path()
.join("sessions/2025/02/02")
.join(format!("rollout-2025-02-02T10-00-00-{uuid}.jsonl"));
std::fs::create_dir_all(rollout_path.parent().expect("rollout parent"))?;
std::fs::write(&rollout_path, "")?;
let state_runtime = codex_state::StateRuntime::init(
config.codex_home.to_path_buf(),
config.model_provider_id.clone(),
)
.await
.map_err(std::io::Error::other)?;
let mut builder = codex_state::ThreadMetadataBuilder::new(
thread_id,
rollout_path.clone(),
chrono::Utc::now(),
serde_json::from_value(serde_json::json!("cli"))
.expect("cli session source should deserialize"),
);
builder.model_provider = Some(config.model_provider_id.clone());
let metadata = builder.build(config.model_provider_id.as_str());
state_runtime
.upsert_thread(&metadata)
.await
.map_err(std::io::Error::other)?;
let target = lookup_local_session_target_by_id(
thread_id,
config.codex_home.as_path(),
Some(state_runtime.as_ref()),
)
.await
.expect("state db path should resolve");
assert_eq!(target.thread_id, thread_id);
assert_eq!(target.path, Some(rollout_path));
Ok(())
}
#[tokio::test]
async fn embedded_app_server_start_failure_is_returned() -> color_eyre::Result<()> {
let temp_dir = TempDir::new()?;
+15 -3
View File
@@ -82,6 +82,7 @@ const PICKER_LIST_HORIZONTAL_INSET: u16 = 4;
pub struct SessionTarget {
pub path: Option<PathBuf>,
pub thread_id: ThreadId,
pub cwd: Option<PathBuf>,
}
impl SessionTarget {
@@ -128,8 +129,17 @@ impl SessionPickerAction {
}
}
fn selection(self, path: Option<PathBuf>, thread_id: ThreadId) -> SessionSelection {
let target_session = SessionTarget { path, thread_id };
fn selection(
self,
path: Option<PathBuf>,
thread_id: ThreadId,
cwd: Option<PathBuf>,
) -> SessionSelection {
let target_session = SessionTarget {
path,
thread_id,
cwd,
};
match self {
SessionPickerAction::Resume => SessionSelection::Resume(target_session),
SessionPickerAction::Fork => SessionSelection::Fork(target_session),
@@ -1107,6 +1117,7 @@ impl PickerState {
_ if self.list_keymap.accept.is_pressed(key) => {
if let Some(row) = self.filtered_rows.get(self.selected) {
let path = row.path.clone();
let cwd = row.cwd.clone();
let thread_id = match row.thread_id {
Some(thread_id) => Some(thread_id),
None => match path.as_ref() {
@@ -1118,7 +1129,7 @@ impl PickerState {
},
};
if let Some(thread_id) = thread_id {
return Ok(Some(self.action.selection(path, thread_id)));
return Ok(Some(self.action.selection(path, thread_id, cwd)));
}
self.inline_error = Some(match path {
Some(path) => {
@@ -5710,6 +5721,7 @@ session_picker_view = "dense"
Some(SessionSelection::Resume(SessionTarget {
path: None,
thread_id: selected_thread_id,
..
})) => assert_eq!(selected_thread_id, thread_id),
other => panic!("unexpected selection: {other:?}"),
}
+213 -17
View File
@@ -5,6 +5,9 @@
//! before the app server has resumed the selected thread.
use std::io;
use std::io::Read;
use std::io::Seek;
use std::io::SeekFrom;
use std::path::Path;
use std::path::PathBuf;
@@ -12,6 +15,7 @@ use crate::cwd_prompt;
use crate::cwd_prompt::CwdPromptAction;
use crate::cwd_prompt::CwdPromptOutcome;
use crate::cwd_prompt::CwdSelection;
use crate::resume_picker::SessionTarget;
use crate::tui::Tui;
use codex_protocol::ThreadId;
use codex_rollout::open_rollout_line_reader;
@@ -20,6 +24,35 @@ use codex_utils_path as path_utils;
use serde::Deserialize;
use serde_json::Value;
const ROLLOUT_RESUME_TAIL_READ_CHUNK_SIZE: usize = 64 * 1024;
const TURN_CONTEXT_NEEDLE: &[u8] = b"turn_context";
pub(crate) async fn verified_rollout_path_for_thread(
path: &Path,
thread_id: ThreadId,
) -> Option<PathBuf> {
let existing_path = codex_rollout::existing_rollout_path(path).await?;
if rollout_file_name_matches_thread_id(existing_path.as_path(), thread_id)
|| codex_rollout::read_session_meta_line(existing_path.as_path())
.await
.is_ok_and(|session_meta| session_meta.meta.id == thread_id)
{
return Some(codex_rollout::plain_rollout_path(existing_path.as_path()));
}
None
}
fn rollout_file_name_matches_thread_id(path: &Path, thread_id: ThreadId) -> bool {
let Some(file_name) = path.file_name().and_then(|name| name.to_str()) else {
return false;
};
let plain_suffix = format!("{thread_id}.jsonl");
let compressed_suffix = format!("{plain_suffix}.zst");
file_name.starts_with("rollout-")
&& (file_name.ends_with(plain_suffix.as_str())
|| file_name.ends_with(compressed_suffix.as_str()))
}
#[derive(Default)]
struct RolloutResumeState {
thread_id: Option<ThreadId>,
@@ -87,12 +120,24 @@ pub(crate) async fn resolve_cwd_for_resume_or_fork(
tui: &mut Tui,
state_db_ctx: Option<&StateRuntime>,
current_cwd: &Path,
thread_id: ThreadId,
path: Option<&Path>,
target_session: &SessionTarget,
action: CwdPromptAction,
allow_prompt: bool,
) -> color_eyre::Result<ResolveCwdOutcome> {
let Some(history_cwd) = read_session_cwd(state_db_ctx, thread_id, path).await else {
let history_cwd = if target_session.path.is_some() {
read_session_cwd(
state_db_ctx,
target_session.thread_id,
target_session.path.as_deref(),
)
.await
.or_else(|| target_session.cwd.clone())
} else if let Some(cwd) = target_session.cwd.as_ref() {
Some(cwd.clone())
} else {
read_session_cwd(state_db_ctx, target_session.thread_id, /*path*/ None).await
};
let Some(history_cwd) = history_cwd else {
return Ok(ResolveCwdOutcome::Continue(None));
};
if allow_prompt && cwds_differ(current_cwd, &history_cwd) {
@@ -116,25 +161,30 @@ async fn read_session_cwd(
thread_id: ThreadId,
path: Option<&Path>,
) -> Option<PathBuf> {
if let Some(path) = path {
match read_rollout_resume_state(path).await {
Ok(state) => {
if state.thread_id == Some(thread_id) && state.cwd.is_some() {
return state.cwd;
}
}
Err(err) => {
let rollout_path = path.display().to_string();
tracing::warn!(
%rollout_path,
%err,
"Failed to read session metadata from rollout"
);
}
}
}
if let Some(state_db_ctx) = state_db_ctx
&& let Ok(Some(metadata)) = state_db_ctx.get_thread(thread_id).await
{
return Some(metadata.cwd);
}
let path = path?;
match read_rollout_resume_state(path).await {
Ok(state) => state.cwd,
Err(err) => {
let rollout_path = path.display().to_string();
tracing::warn!(
%rollout_path,
%err,
"Failed to read session metadata from rollout"
);
None
}
}
None
}
pub(crate) fn cwds_differ(current_cwd: &Path, session_cwd: &Path) -> bool {
@@ -142,6 +192,152 @@ pub(crate) fn cwds_differ(current_cwd: &Path, session_cwd: &Path) -> bool {
}
async fn read_rollout_resume_state(path: &Path) -> io::Result<RolloutResumeState> {
if let Some(state) = try_read_rollout_resume_state_fast(path).await? {
return Ok(state);
}
read_rollout_resume_state_full(path).await
}
async fn try_read_rollout_resume_state_fast(path: &Path) -> io::Result<Option<RolloutResumeState>> {
let Some(existing_path) = codex_rollout::existing_rollout_path(path).await else {
return Ok(None);
};
if is_compressed_rollout_path(existing_path.as_path()) {
return Ok(None);
}
let Some(session_meta) = read_initial_session_metadata(existing_path.as_path()).await? else {
return Ok(None);
};
let latest_turn_context =
match latest_turn_context_from_plain_rollout(existing_path.as_path()).await {
Ok(turn_context) => turn_context,
Err(err) if err.kind() == io::ErrorKind::NotFound => return Ok(None),
Err(err) => return Err(err),
};
let mut state = RolloutResumeState {
thread_id: Some(session_meta.id),
cwd: Some(session_meta.cwd),
..Default::default()
};
if let Some(turn_context) = latest_turn_context {
state.cwd = Some(turn_context.cwd);
state.model = Some(turn_context.model);
}
Ok(Some(state))
}
async fn read_initial_session_metadata(path: &Path) -> io::Result<Option<SessionMetadata>> {
let mut reader = open_rollout_line_reader(path).await?;
while let Some(line) = reader.next_line().await? {
let trimmed = line.trim();
if trimmed.is_empty() {
continue;
}
let Ok(record) = serde_json::from_str::<RawRecord>(trimmed) else {
continue;
};
let Some(payload) = record.payload else {
return Ok(None);
};
if record.item_type != "session_meta" {
return Ok(None);
}
return Ok(serde_json::from_value::<SessionMetadata>(payload).ok());
}
Ok(None)
}
async fn latest_turn_context_from_plain_rollout(
path: &Path,
) -> io::Result<Option<TurnContextResumeState>> {
let path = path.to_path_buf();
tokio::task::spawn_blocking(move || {
latest_turn_context_from_plain_rollout_blocking(path.as_path())
})
.await
.map_err(io::Error::other)?
}
fn latest_turn_context_from_plain_rollout_blocking(
path: &Path,
) -> io::Result<Option<TurnContextResumeState>> {
let mut file = std::fs::File::open(path)?;
let mut remaining = file.metadata()?.len();
let mut line_rev = Vec::new();
let mut buf = vec![0u8; ROLLOUT_RESUME_TAIL_READ_CHUNK_SIZE];
while remaining > 0 {
let read_size = usize::try_from(remaining.min(ROLLOUT_RESUME_TAIL_READ_CHUNK_SIZE as u64))
.map_err(io::Error::other)?;
remaining -= read_size as u64;
file.seek(SeekFrom::Start(remaining))?;
file.read_exact(&mut buf[..read_size])?;
for &byte in buf[..read_size].iter().rev() {
if byte == b'\n' {
if let Some(turn_context) = parse_turn_context_from_rev_line(&mut line_rev)? {
return Ok(Some(turn_context));
}
} else {
line_rev.push(byte);
}
}
}
parse_turn_context_from_rev_line(&mut line_rev)
}
fn parse_turn_context_from_rev_line(
line_rev: &mut Vec<u8>,
) -> io::Result<Option<TurnContextResumeState>> {
if line_rev.is_empty() {
return Ok(None);
}
line_rev.reverse();
let line = std::mem::take(line_rev);
let trimmed = trim_ascii_whitespace(line.as_slice());
if trimmed.is_empty() || !contains_bytes(trimmed, TURN_CONTEXT_NEEDLE) {
return Ok(None);
}
let Ok(record) = serde_json::from_slice::<RawRecord>(trimmed) else {
return Ok(None);
};
if record.item_type != "turn_context" {
return Ok(None);
}
let Some(payload) = record.payload else {
return Ok(None);
};
Ok(serde_json::from_value::<TurnContextResumeState>(payload).ok())
}
fn is_compressed_rollout_path(path: &Path) -> bool {
path.file_name()
.and_then(|name| name.to_str())
.is_some_and(|name| name.ends_with(".jsonl.zst"))
}
fn contains_bytes(haystack: &[u8], needle: &[u8]) -> bool {
haystack
.windows(needle.len())
.any(|window| window == needle)
}
fn trim_ascii_whitespace(bytes: &[u8]) -> &[u8] {
let mut start = 0usize;
let mut end = bytes.len();
while start < end && bytes[start].is_ascii_whitespace() {
start += 1;
}
while end > start && bytes[end - 1].is_ascii_whitespace() {
end -= 1;
}
&bytes[start..end]
}
async fn read_rollout_resume_state_full(path: &Path) -> io::Result<RolloutResumeState> {
let mut reader = open_rollout_line_reader(path).await?;
let mut state = RolloutResumeState::default();
let mut saw_record = false;