Codex Handbook
thread-store/src/local/live_writer.rs 206 lines
use std::path::PathBuf;use codex_protocol::ThreadId;use codex_protocol::protocol::ThreadMemoryMode;use codex_rollout::RolloutConfig;use codex_rollout::RolloutRecorder;use codex_rollout::RolloutRecorderParams;use codex_rollout::persisted_rollout_items;use tracing::warn;use super::LocalThreadStore;use super::create_thread;use crate::AppendThreadItemsParams;use crate::CreateThreadParams;use crate::ReadThreadParams;use crate::ResumeThreadParams;use crate::ThreadStoreError;use crate::ThreadStoreResult;pub(super) async fn create_thread(    store: &LocalThreadStore,    params: CreateThreadParams,) -> ThreadStoreResult<()> {    let thread_id = params.thread_id;    store.ensure_live_recorder_absent(thread_id).await?;    let recorder = create_thread::create_thread(store, params).await?;    store.insert_live_recorder(thread_id, recorder).await}pub(super) async fn resume_thread(    store: &LocalThreadStore,    params: ResumeThreadParams,) -> ThreadStoreResult<()> {    store.ensure_live_recorder_absent(params.thread_id).await?;    let rollout_path = match (params.rollout_path, params.history) {        (Some(rollout_path), _history) => rollout_path,        (None, history) => {            let thread = super::read_thread::read_thread(                store,                ReadThreadParams {                    thread_id: params.thread_id,                    include_archived: params.include_archived,                    include_history: history.is_none(),                },            )            .await?;            thread                .rollout_path                .ok_or_else(|| ThreadStoreError::Internal {                    message: format!("thread {} does not have a rollout path", params.thread_id),                })?        }    };    let cwd = params        .metadata        .cwd        .clone()        .ok_or_else(|| ThreadStoreError::InvalidRequest {            message: "local thread store requires a cwd".to_string(),        })?;    let config = RolloutConfig {        codex_home: store.config.codex_home.clone(),        sqlite_home: store.config.sqlite_home.clone(),        cwd,        model_provider_id: params.metadata.model_provider.clone(),        generate_memories: matches!(params.metadata.memory_mode, ThreadMemoryMode::Enabled),    };    let recorder = RolloutRecorder::new(&config, RolloutRecorderParams::resume(rollout_path))        .await        .map_err(|err| ThreadStoreError::Internal {            message: format!("failed to resume local thread recorder: {err}"),        })?;    store.insert_live_recorder(params.thread_id, recorder).await}pub(super) async fn append_items(    store: &LocalThreadStore,    params: AppendThreadItemsParams,) -> ThreadStoreResult<()> {    let canonical_items = persisted_rollout_items(params.items.as_slice());    if canonical_items.is_empty() {        return Ok(());    }    let recorder = store.live_recorder(params.thread_id).await?;    recorder        .record_canonical_items(canonical_items.as_slice())        .await        .map_err(thread_store_io_error)?;    // LiveThread applies metadata immediately after append_items returns. Wait for the local    // writer so SQLite never gets ahead of JSONL for accepted live appends.    recorder.flush().await.map_err(thread_store_io_error)}pub(super) async fn persist_thread(    store: &LocalThreadStore,    thread_id: ThreadId,) -> ThreadStoreResult<()> {    store        .live_recorder(thread_id)        .await?        .persist()        .await        .map_err(thread_store_io_error)?;    sync_materialized_rollout_path(store, thread_id).await}pub(super) async fn flush_thread(    store: &LocalThreadStore,    thread_id: ThreadId,) -> ThreadStoreResult<()> {    store        .live_recorder(thread_id)        .await?        .flush()        .await        .map_err(thread_store_io_error)?;    sync_materialized_rollout_path(store, thread_id).await}pub(super) async fn shutdown_thread(    store: &LocalThreadStore,    thread_id: ThreadId,) -> ThreadStoreResult<()> {    let recorder = store.live_recorder(thread_id).await?;    recorder.shutdown().await.map_err(thread_store_io_error)?;    sync_materialized_rollout_path(store, thread_id).await?;    store.live_recorders.lock().await.remove(&thread_id);    Ok(())}pub(super) async fn discard_thread(    store: &LocalThreadStore,    thread_id: ThreadId,) -> ThreadStoreResult<()> {    store        .live_recorders        .lock()        .await        .remove(&thread_id)        .map(|_| ())        .ok_or(ThreadStoreError::ThreadNotFound { thread_id })}pub(super) async fn rollout_path(    store: &LocalThreadStore,    thread_id: ThreadId,) -> ThreadStoreResult<PathBuf> {    Ok(store        .live_recorders        .lock()        .await        .get(&thread_id)        .ok_or(ThreadStoreError::ThreadNotFound { thread_id })?        .rollout_path()        .to_path_buf())}async fn sync_materialized_rollout_path(    store: &LocalThreadStore,    thread_id: ThreadId,) -> ThreadStoreResult<()> {    let rollout_path = rollout_path(store, thread_id).await?;    if codex_rollout::existing_rollout_path(rollout_path.as_path())        .await        .is_none()    {        return Ok(());    }    let Some(state_db) = store.state_db().await else {        return Ok(());    };    let result: ThreadStoreResult<()> = async {        let Some(mut metadata) =            state_db                .get_thread(thread_id)                .await                .map_err(|err| ThreadStoreError::Internal {                    message: format!("failed to read thread metadata for {thread_id}: {err}"),                })?        else {            return Ok(());        };        if metadata.rollout_path != rollout_path {            metadata.rollout_path = rollout_path;            state_db                .upsert_thread(&metadata)                .await                .map_err(|err| ThreadStoreError::Internal {                    message: format!("failed to update thread metadata for {thread_id}: {err}"),                })?;        }        Ok(())    }    .await;    if let Err(err) = result {        warn!("failed to sync materialized rollout path for thread {thread_id}: {err}");    }    Ok(())}fn thread_store_io_error(err: std::io::Error) -> ThreadStoreError {    ThreadStoreError::Internal {        message: err.to_string(),    }}