Codex Handbook
tui/src/session_resume.rs 311 lines
//! Resolve saved-session state needed before resuming or forking a thread.//!//! The app-server API owns normal thread lifecycle data. This module coordinates//! the TUI-specific cwd prompt and falls back to local rollout metadata only//! before the app server has resumed the selected thread.use std::io;use std::path::Path;use std::path::PathBuf;use crate::cwd_prompt;use crate::cwd_prompt::CwdPromptAction;use crate::cwd_prompt::CwdPromptOutcome;use crate::cwd_prompt::CwdSelection;use crate::tui::Tui;use codex_protocol::ThreadId;use codex_rollout::open_rollout_line_reader;use codex_state::StateRuntime;use codex_utils_path as path_utils;use serde::Deserialize;use serde_json::Value;#[derive(Default)]struct RolloutResumeState {    thread_id: Option<ThreadId>,    cwd: Option<PathBuf>,    model: Option<String>,}#[derive(Deserialize)]struct SessionMetadata {    id: ThreadId,    cwd: PathBuf,}#[derive(Deserialize)]struct TurnContextResumeState {    cwd: PathBuf,    model: String,}#[derive(Deserialize)]struct RawRecord {    #[serde(rename = "type")]    item_type: String,    payload: Option<Value>,}pub(crate) enum ResolveCwdOutcome {    Continue(Option<PathBuf>),    Exit,}pub(crate) async fn resolve_session_thread_id(    path: &Path,    id_str_if_uuid: Option<&str>,) -> Option<ThreadId> {    match id_str_if_uuid {        Some(id_str) => ThreadId::from_string(id_str).ok(),        None => read_rollout_resume_state(path)            .await            .ok()            .and_then(|state| state.thread_id),    }}pub(crate) async fn read_session_model(    state_db_ctx: Option<&StateRuntime>,    thread_id: ThreadId,    path: Option<&Path>,) -> Option<String> {    if let Some(state_db_ctx) = state_db_ctx        && let Ok(Some(metadata)) = state_db_ctx.get_thread(thread_id).await        && let Some(model) = metadata.model    {        return Some(model);    }    let path = path?;    read_rollout_resume_state(path)        .await        .ok()        .and_then(|state| state.model)}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>,    action: CwdPromptAction,    allow_prompt: bool,) -> color_eyre::Result<ResolveCwdOutcome> {    let Some(history_cwd) = read_session_cwd(state_db_ctx, thread_id, path).await else {        return Ok(ResolveCwdOutcome::Continue(None));    };    if allow_prompt && cwds_differ(current_cwd, &history_cwd) {        let selection_outcome =            cwd_prompt::run_cwd_selection_prompt(tui, action, current_cwd, &history_cwd).await?;        return Ok(match selection_outcome {            CwdPromptOutcome::Selection(CwdSelection::Current) => {                ResolveCwdOutcome::Continue(Some(current_cwd.to_path_buf()))            }            CwdPromptOutcome::Selection(CwdSelection::Session) => {                ResolveCwdOutcome::Continue(Some(history_cwd))            }            CwdPromptOutcome::Exit => ResolveCwdOutcome::Exit,        });    }    Ok(ResolveCwdOutcome::Continue(Some(history_cwd)))}async fn read_session_cwd(    state_db_ctx: Option<&StateRuntime>,    thread_id: ThreadId,    path: Option<&Path>,) -> Option<PathBuf> {    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        }    }}pub(crate) fn cwds_differ(current_cwd: &Path, session_cwd: &Path) -> bool {    !path_utils::paths_match_after_normalization(current_cwd, session_cwd)}async fn read_rollout_resume_state(path: &Path) -> io::Result<RolloutResumeState> {    let mut reader = open_rollout_line_reader(path).await?;    let mut state = RolloutResumeState::default();    let mut saw_record = false;    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;        };        saw_record = true;        let Some(payload) = record.payload else {            continue;        };        match record.item_type.as_str() {            "session_meta" if state.thread_id.is_none() => {                if let Ok(metadata) = serde_json::from_value::<SessionMetadata>(payload) {                    state.thread_id = Some(metadata.id);                    state.cwd.get_or_insert(metadata.cwd);                }            }            "turn_context" => {                if let Ok(turn_context) = serde_json::from_value::<TurnContextResumeState>(payload)                {                    state.cwd = Some(turn_context.cwd);                    state.model = Some(turn_context.model);                }            }            _ => {}        }    }    if saw_record {        Ok(state)    } else {        Err(io::Error::other(format!(            "rollout at {} is empty",            path.display()        )))    }}#[cfg(test)]mod tests {    use super::*;    use pretty_assertions::assert_eq;    use tempfile::TempDir;    fn rollout_line(        timestamp: &str,        item_type: &str,        payload: serde_json::Value,    ) -> serde_json::Value {        serde_json::json!({            "timestamp": timestamp,            "type": item_type,            "payload": payload,        })    }    fn write_rollout_lines(path: &Path, lines: &[serde_json::Value]) -> std::io::Result<()> {        let mut text = String::new();        for line in lines {            text.push_str(&serde_json::to_string(line).expect("serialize rollout"));            text.push('\n');        }        std::fs::write(path, text)    }    #[tokio::test]    async fn rollout_resume_state_prefers_latest_turn_context() -> std::io::Result<()> {        let temp_dir = TempDir::new()?;        let thread_id = ThreadId::new();        let original = temp_dir.path().join("original");        let latest = temp_dir.path().join("latest");        let rollout_path = temp_dir.path().join("rollout.jsonl");        write_rollout_lines(            &rollout_path,            &[                rollout_line(                    "t0",                    "session_meta",                    serde_json::json!({                        "id": thread_id,                        "cwd": original,                        "originator": "test",                        "cli_version": "test",                    }),                ),                rollout_line(                    "t1",                    "turn_context",                    serde_json::json!({ "cwd": temp_dir.path().join("middle"), "model": "middle" }),                ),                rollout_line(                    "t2",                    "turn_context",                    serde_json::json!({ "cwd": latest.clone(), "model": "latest" }),                ),            ],        )?;        let state = read_rollout_resume_state(&rollout_path).await?;        assert_eq!(state.thread_id, Some(thread_id));        assert_eq!(state.cwd, Some(latest));        assert_eq!(state.model, Some("latest".to_string()));        Ok(())    }    #[tokio::test]    async fn rollout_resume_state_falls_back_to_session_meta() -> std::io::Result<()> {        let temp_dir = TempDir::new()?;        let thread_id = ThreadId::new();        let cwd = temp_dir.path().join("session");        let rollout_path = temp_dir.path().join("rollout.jsonl");        write_rollout_lines(            &rollout_path,            &[rollout_line(                "t0",                "session_meta",                serde_json::json!({                    "id": thread_id,                    "cwd": cwd.clone(),                    "originator": "test",                    "cli_version": "test",                }),            )],        )?;        let state = read_rollout_resume_state(&rollout_path).await?;        assert_eq!(state.thread_id, Some(thread_id));        assert_eq!(state.cwd, Some(cwd));        assert_eq!(state.model, None);        Ok(())    }    #[tokio::test]    async fn rollout_resume_state_skips_malformed_lines() -> std::io::Result<()> {        let temp_dir = TempDir::new()?;        let thread_id = ThreadId::new();        let cwd = temp_dir.path().join("session");        let rollout_path = temp_dir.path().join("rollout.jsonl");        let valid_line = serde_json::to_string(&rollout_line(            "t0",            "session_meta",            serde_json::json!({                "id": thread_id,                "cwd": cwd.clone(),                "originator": "test",                "cli_version": "test",            }),        ))        .expect("serialize rollout line");        std::fs::write(&rollout_path, format!("{valid_line}\n{{"))?;        let state = read_rollout_resume_state(&rollout_path).await?;        assert_eq!(state.thread_id, Some(thread_id));        assert_eq!(state.cwd, Some(cwd));        Ok(())    }}