codex/codex-rs/thread-store/src/local/pending_thread_metadata_tests.rs

use std::sync::Arc;
use std::time::Duration;

use codex_protocol::ThreadId;
use codex_protocol::models::BaseInstructions;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::ThreadHistoryMode;
use codex_protocol::protocol::ThreadMemoryMode;
use codex_protocol::protocol::UserMessageEvent;
use codex_rollout::RolloutItem;
use codex_rollout::RolloutRecorder;
use pretty_assertions::assert_eq;
use tempfile::TempDir;

use super::*;
use crate::CreateThreadParams;
use crate::LiveThread;
use crate::ThreadPersistenceMetadata;
use crate::UpdateThreadMetadataParams;
use crate::local::test_support::test_config;

#[tokio::test]
async fn naming_empty_paginated_thread_materializes_rollout() {
    let (_home, store, runtime) = store_with_runtime().await;
    let thread_id = ThreadId::new();
    let mut params = create_thread_params(thread_id);
    params.history_mode = ThreadHistoryMode::Paginated;
    let live_thread = LiveThread::create(store.clone(), params)
        .await
        .expect("create live thread");
    let rollout_path = store
        .live_rollout_path(thread_id)
        .await
        .expect("live rollout path");
    assert!(!rollout_path.exists());

    let named = live_thread
        .update_metadata(
            ThreadMetadataPatch {
                name: Some(Some("Scheduled task".to_string())),
                memory_mode: Some(ThreadMemoryMode::Disabled),
                ..Default::default()
            },
            /*include_archived*/ false,
        )
        .await
        .expect("name empty thread");

    assert!(rollout_path.exists());
    assert_eq!(named.name.as_deref(), Some("Scheduled task"));
    assert_eq!(named.history_mode, ThreadHistoryMode::Paginated);
    assert_eq!(
        runtime
            .get_thread_memory_mode(thread_id)
            .await
            .expect("read memory mode")
            .as_deref(),
        Some("disabled")
    );
    live_thread.shutdown().await.expect("shutdown named thread");
    let stored = store
        .read_thread(ReadThreadParams {
            thread_id,
            include_archived: false,
            include_history: false,
        })
        .await
        .expect("read durable empty thread");
    assert_eq!(stored.thread_id, thread_id);
    assert_eq!(stored.name, named.name);
    assert_eq!(stored.history_mode, ThreadHistoryMode::Paginated);
}

#[tokio::test]
async fn naming_empty_paginated_thread_propagates_rollout_persistence_failure() {
    let (_home, store, runtime) = store_with_runtime().await;
    let thread_id = ThreadId::new();
    let mut params = create_thread_params(thread_id);
    params.history_mode = ThreadHistoryMode::Paginated;
    let live_thread = LiveThread::create(store.clone(), params)
        .await
        .expect("create live thread");
    let rollout_path = store
        .live_rollout_path(thread_id)
        .await
        .expect("live rollout path");
    std::fs::create_dir_all(&rollout_path).expect("block rollout file creation");

    live_thread
        .update_metadata(
            ThreadMetadataPatch {
                name: Some(Some("Unpersistable task".to_string())),
                ..Default::default()
            },
            /*include_archived*/ false,
        )
        .await
        .expect_err("naming must fail when the thread cannot be persisted");
    assert_eq!(
        runtime
            .get_thread(thread_id)
            .await
            .expect("read thread metadata")
            .and_then(|metadata| metadata.name),
        None
    );
}

#[tokio::test]
async fn pending_thread_metadata_is_consumed_by_first_metadata_update() {
    let (_home, store, runtime) = store_with_runtime().await;
    let thread_id = ThreadId::new();
    let pending_patch = staged_model_patch();

    store
        .stage_pending_thread_metadata(thread_id, pending_patch.clone())
        .await
        .expect("stage pending metadata");
    let live_thread = LiveThread::create(store.clone(), create_thread_params(thread_id))
        .await
        .expect("create live thread");
    live_thread
        .append_items(&[user_message_item("materialize pending metadata")])
        .await
        .expect("append observed item");

    let metadata = runtime
        .get_thread(thread_id)
        .await
        .expect("sqlite metadata read")
        .expect("sqlite metadata");
    assert_eq!(metadata.model.as_deref(), Some("staged-model"));
    assert_eq!(metadata.model_provider, "test-provider");
    assert_eq!(metadata.source, "exec");
    assert_eq!(
        metadata.cwd.canonicalize().expect("canonical metadata cwd"),
        std::env::current_dir()
            .expect("current directory")
            .canonicalize()
            .expect("canonical current directory")
    );
    assert_pending_metadata_consumed(&store, thread_id).await;
}

#[tokio::test]
async fn pending_rollout_compatible_metadata_does_not_deadlock() {
    let (_home, store, runtime) = store_with_runtime().await;
    let thread_id = ThreadId::new();

    store
        .stage_pending_thread_metadata(
            thread_id,
            ThreadMetadataPatch {
                name: Some(Some("staged-name".to_string())),
                model_provider: Some("staged-provider".to_string()),
                ..Default::default()
            },
        )
        .await
        .expect("stage pending metadata");
    let live_thread = LiveThread::create(store.clone(), create_thread_params(thread_id))
        .await
        .expect("create live thread");
    tokio::time::timeout(
        Duration::from_secs(5),
        live_thread.update_metadata(
            ThreadMetadataPatch {
                model: Some("updated-model".to_string()),
                ..Default::default()
            },
            /*include_archived*/ false,
        ),
    )
    .await
    .expect("pending metadata update should not deadlock")
    .expect("update metadata before rollout materialization");

    let metadata = runtime
        .get_thread(thread_id)
        .await
        .expect("sqlite metadata read")
        .expect("sqlite metadata");
    assert_eq!(metadata.title, "staged-name");
    assert_eq!(metadata.model_provider, "test-provider");
    assert_eq!(metadata.model.as_deref(), Some("updated-model"));
    assert_pending_metadata_consumed(&store, thread_id).await;
}

#[tokio::test]
async fn pending_rollout_compatible_metadata_is_written_with_observed_update() {
    let (_home, store, _runtime) = store_with_runtime().await;
    let thread_id = ThreadId::new();

    store
        .stage_pending_thread_metadata(
            thread_id,
            ThreadMetadataPatch {
                memory_mode: Some(ThreadMemoryMode::Disabled),
                ..Default::default()
            },
        )
        .await
        .expect("stage pending metadata");
    let _live_thread = LiveThread::create(store.clone(), create_thread_params(thread_id))
        .await
        .expect("create live thread");
    store
        .update_thread_metadata(UpdateThreadMetadataParams {
            thread_id,
            patch: ThreadMetadataPatch {
                model: Some("observed-model".to_string()),
                ..Default::default()
            },
            include_archived: false,
        })
        .await
        .expect("apply observed metadata");

    let rollout_path = store
        .live_rollout_path(thread_id)
        .await
        .expect("rollout path");
    let (items, _, _) = RolloutRecorder::load_rollout_items(rollout_path.as_path())
        .await
        .expect("load durable rollout");
    let memory_mode = items.iter().rev().find_map(|item| match item {
        RolloutItem::SessionMeta(session_meta) => session_meta.meta.memory_mode.as_deref(),
        _ => None,
    });

    assert_eq!(memory_mode, Some("disabled"));
    assert_pending_metadata_consumed(&store, thread_id).await;
}

#[tokio::test]
async fn deferred_reads_preserve_creation_metadata_and_updates() {
    for history_mode in [ThreadHistoryMode::Legacy, ThreadHistoryMode::Paginated] {
        let (_home, store, runtime) = store_with_runtime().await;
        let thread_id = ThreadId::new();
        let mut params = create_thread_params(thread_id);
        params.history_mode = history_mode;
        params.forked_from_id = Some(ThreadId::new());
        params.parent_thread_id = Some(ThreadId::new());
        let live_thread = LiveThread::create(store.clone(), params.clone())
            .await
            .expect("create live thread");
        let rollout_path = store
            .live_rollout_path(thread_id)
            .await
            .expect("live rollout path");
        let initial = live_thread
            .read_thread(
                /*include_archived*/ false, /*include_history*/ false,
            )
            .await
            .expect_err("untouched threads still have no stored record");
        assert!(matches!(
            initial,
            ThreadStoreError::InvalidRequest { message }
                if message == format!("no rollout found for thread id {thread_id}")
        ));
        assert!(!rollout_path.exists());
        assert_eq!(
            runtime.get_thread(thread_id).await.expect("read index"),
            None
        );

        let updated = live_thread
            .update_metadata(
                ThreadMetadataPatch {
                    daybreak_enabled: Some(true),
                    ..Default::default()
                },
                /*include_archived*/ false,
            )
            .await
            .expect("update metadata without materializing history");
        let read = live_thread
            .read_thread(
                /*include_archived*/ false, /*include_history*/ false,
            )
            .await
            .expect("read updated metadata");
        assert_eq!(
            serde_json::to_value(read).expect("serialize metadata read"),
            serde_json::to_value(&updated).expect("serialize update result")
        );
        assert_eq!(
            (
                updated.forked_from_id,
                updated.parent_thread_id,
                updated.originator.clone(),
                updated.history_mode,
                updated.daybreak_enabled,
            ),
            (
                params.forked_from_id,
                params.parent_thread_id,
                Some(params.originator),
                history_mode,
                Some(true),
            )
        );
        assert!(!rollout_path.exists());

        live_thread
            .persist(PersistContext::Standard)
            .await
            .expect("persist history");
        live_thread.shutdown().await.expect("shutdown writer");
        let durable = store
            .read_thread(ReadThreadParams {
                thread_id,
                include_archived: false,
                include_history: false,
            })
            .await
            .expect("read durable metadata");
        assert_eq!(
            serde_json::to_value(durable).expect("serialize durable metadata"),
            serde_json::to_value(updated).expect("serialize initial metadata")
        );
    }
}

#[tokio::test]
async fn idle_shutdown_clears_pending_thread_metadata_without_materializing() {
    let (_home, store, runtime) = store_with_runtime().await;
    let thread_id = ThreadId::new();

    store
        .stage_pending_thread_metadata(thread_id, staged_model_patch())
        .await
        .expect("stage pending metadata");
    let live_thread = LiveThread::create(store.clone(), create_thread_params(thread_id))
        .await
        .expect("create live thread");
    let read_error = live_thread
        .read_thread(
            /*include_archived*/ false, /*include_history*/ false,
        )
        .await
        .expect_err("untouched threads still have no stored record");
    assert!(matches!(
        read_error,
        ThreadStoreError::InvalidRequest { message }
            if message == format!("no rollout found for thread id {thread_id}")
    ));
    assert!(
        !store
            .live_rollout_path(thread_id)
            .await
            .expect("live rollout path")
            .exists()
    );
    live_thread.shutdown().await.expect("shutdown idle thread");

    assert_eq!(
        runtime
            .get_thread(thread_id)
            .await
            .expect("sqlite metadata read"),
        None
    );
    assert_pending_metadata_consumed(&store, thread_id).await;
}

#[tokio::test]
async fn discarding_live_thread_clears_pending_metadata() {
    let (_home, store, _runtime) = store_with_runtime().await;
    let thread_id = ThreadId::new();

    store
        .stage_pending_thread_metadata(thread_id, staged_model_patch())
        .await
        .expect("stage pending metadata");
    let live_thread = LiveThread::create(store.clone(), create_thread_params(thread_id))
        .await
        .expect("create live thread");
    live_thread.discard().await.expect("discard live thread");

    assert_pending_metadata_consumed(&store, thread_id).await;
}

#[tokio::test]
async fn pending_thread_metadata_requires_state_db() {
    let home = TempDir::new().expect("temp dir");
    let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);

    let err = store
        .stage_pending_thread_metadata(ThreadId::new(), staged_model_patch())
        .await
        .expect_err("reject pending metadata without state db");
    assert!(matches!(
        err,
        ThreadStoreError::InvalidRequest { message }
            if message == "pending thread metadata requires a state db"
    ));
}

#[tokio::test]
async fn pending_thread_metadata_rejects_rollout_path() {
    let (home, store, _runtime) = store_with_runtime().await;

    let err = store
        .stage_pending_thread_metadata(
            ThreadId::new(),
            ThreadMetadataPatch {
                rollout_path: Some(home.path().join("wrong-rollout.jsonl")),
                ..Default::default()
            },
        )
        .await
        .expect_err("reject host-owned rollout path");
    assert!(matches!(
        err,
        ThreadStoreError::InvalidRequest { message }
            if message == "pending thread metadata cannot set rollout_path"
    ));
}

async fn store_with_runtime() -> (
    TempDir,
    Arc<LocalThreadStore>,
    Arc<codex_state::StateRuntime>,
) {
    let home = TempDir::new().expect("temp dir");
    let config = test_config(home.path());
    let runtime = codex_state::StateRuntime::init(
        config.sqlite.clone(),
        config.default_model_provider_id.clone(),
    )
    .await
    .expect("state db should initialize");
    let store = Arc::new(LocalThreadStore::new(config, Some(runtime.clone())));
    (home, store, runtime)
}

async fn assert_pending_metadata_consumed(store: &LocalThreadStore, thread_id: ThreadId) {
    store
        .stage_pending_thread_metadata(thread_id, staged_model_patch())
        .await
        .expect("restage metadata after consumption");
    store
        .remove_pending_thread_metadata(thread_id)
        .await
        .expect("remove restaged metadata");
}

fn staged_model_patch() -> ThreadMetadataPatch {
    ThreadMetadataPatch {
        model: Some("staged-model".to_string()),
        ..Default::default()
    }
}

fn create_thread_params(thread_id: ThreadId) -> CreateThreadParams {
    CreateThreadParams {
        attached_thread: false,
        creator_user_id: None,
        creator_account_id: None,
        session_id: thread_id.into(),
        thread_id,
        extra_config: None,
        forked_from_id: None,
        parent_thread_id: None,
        source: SessionSource::Exec,
        thread_source: None,
        originator: "test_originator".to_string(),
        base_instructions: BaseInstructions::default(),
        dynamic_tools: Vec::new(),
        selected_capability_roots: Vec::new(),
        multi_agent_version: None,
        history_mode: ThreadHistoryMode::Legacy,
        history_base: None,
        subagent_history_start_ordinal: None,
        initial_window_id: uuid::Uuid::now_v7().to_string(),
        runtime_workspace_roots: None,
        metadata: ThreadPersistenceMetadata {
            cwd: Some(std::env::current_dir().expect("cwd")),
            model_provider: "test-provider".to_string(),
            memory_mode: ThreadMemoryMode::Enabled,
        },
    }
}

fn user_message_item(message: &str) -> RolloutItem {
    RolloutItem::EventMsg(EventMsg::UserMessage(UserMessageEvent {
        client_id: None,
        message: message.to_string(),
        images: None,
        local_images: Vec::new(),
        text_elements: Vec::new(),
        ..Default::default()
    }))
}