mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
jif
8f02973d25
## Why
`selectedCapabilityRoots` is durable thread intent: “use this capability
root from environment `worker`.”
The important product assumption is:
> One environment ID always names the same logical executor and stable
contents.
`worker` does not silently change from executor A to an unrelated
executor B. The process-local connection handle for `worker` can still
be replaced while Codex is running, though, for example when
`environment/add` registers a fresh handle for the same logical
environment.
The thread should persist only the stable selection. Each model step
should pair that selection with the exact ready handle captured for that
step.
## The boundary
```text
persisted thread intent
plugin@1 -> environment "worker"
|
| capture the current step
v
model-step view
unavailable, or
plugin@1 + worker's exact captured ready handle
```
The environment ID is the stable identity and cache key. The
`Arc<Environment>` is only a process-local handle retained so consumers
of one model step use the same captured environment. It is never
persisted and it does not imply different environment contents.
## What changes
### Persist the stable selection
Selected roots are written into `SessionMeta` and restored with the
thread. Forked subagents inherit the same selections, including
bounded-history forks.
Only stable data is persisted: root ID, environment ID, and root path.
### Capture readiness together with the exact handle
The environment snapshot records:
```rust
environment_id -> Some(Arc<Environment>) // ready in this step
environment_id -> None // still starting in this step
```
This prevents readiness and execution from coming from different
registry snapshots.
For example:
```text
step snapshot: worker -> handle A, ready
environment/add: worker -> fresh handle B for the same logical environment
current step: plugin@1 still uses captured handle A
```
Without carrying handle A in the snapshot, the resolver could combine “A
was ready” with handle B and treat B as ready before it had finished
starting.
This does not change cache invalidation. Stable capability metadata
remains identified by environment ID and capability root. Replacing a
process-local handle under the same stable environment ID does not
invalidate or rediscover that metadata.
### Resolve availability per model step
- A ready captured environment produces resolved roots using its
captured handle.
- A starting, missing, or failed environment is omitted from that step.
- A selected lazy environment that is outside the turn's captured
environment set is asked to start, and a later step can observe it as
ready.
- No capability files are scanned here.
Transient transport disconnects remain the remote client's reconnect
concern. This PR models initial attachment/readiness; it does not add
live socket-connectivity state.
## Example
```text
thread selection: plugin@1 -> environment "worker"
step 1: worker is starting -> plugin@1 unavailable
step 2: worker is ready -> plugin@1 resolves through worker's captured handle
step 3: fresh local handle -> current step remains pinned; a later step captures its own view
```
Temporary unavailability does not discard the durable selection. Later
PRs can retain stable metadata caches while projecting only currently
available capabilities into model-visible World State.
## Compatibility
The app-server request shape does not change. Older rollouts without
`selected_capability_roots` deserialize to an empty list.
## Stack
1. **This PR:** persist stable selected roots and resolve them through
an exact model-step handle.
2. #29960: cache stable skill metadata and project available skills into
World State.
3. #29946: cache stable plugin declarations and manage the separate live
MCP runtime.
1174 lines
41 KiB
Rust
1174 lines
41 KiB
Rust
mod archive_thread;
|
|
mod create_thread;
|
|
mod delete_thread;
|
|
mod helpers;
|
|
mod list_threads;
|
|
mod live_writer;
|
|
mod read_thread;
|
|
mod search_threads;
|
|
mod unarchive_thread;
|
|
mod update_thread_metadata;
|
|
|
|
#[cfg(test)]
|
|
mod test_support;
|
|
|
|
use codex_protocol::ThreadId;
|
|
use codex_rollout::RolloutRecorder;
|
|
use codex_rollout::StateDbHandle;
|
|
use std::collections::HashMap;
|
|
use std::collections::hash_map::Entry;
|
|
use std::path::PathBuf;
|
|
use std::sync::Arc;
|
|
use tokio::sync::Mutex;
|
|
|
|
use crate::AppendThreadItemsParams;
|
|
use crate::ArchiveThreadParams;
|
|
use crate::CreateThreadParams;
|
|
use crate::DeleteThreadParams;
|
|
use crate::ListThreadsParams;
|
|
use crate::LoadThreadHistoryParams;
|
|
use crate::ReadThreadByRolloutPathParams;
|
|
use crate::ReadThreadParams;
|
|
use crate::ResumeThreadParams;
|
|
use crate::SearchThreadsParams;
|
|
use crate::StoredThread;
|
|
use crate::StoredThreadHistory;
|
|
use crate::ThreadPage;
|
|
use crate::ThreadSearchPage;
|
|
use crate::ThreadStore;
|
|
use crate::ThreadStoreError;
|
|
use crate::ThreadStoreFuture;
|
|
use crate::ThreadStoreResult;
|
|
use crate::UpdateThreadMetadataParams;
|
|
|
|
/// Local filesystem/SQLite-backed implementation of [`ThreadStore`].
|
|
///
|
|
/// Local storage has two compatibility surfaces. Rollout JSONL files are the
|
|
/// durable replay format and remain readable without SQLite, including older
|
|
/// files that encode metadata in `SessionMeta` items and name-index entries.
|
|
/// The SQLite state DB, when available, is the queryable metadata index used by
|
|
/// list/read paths for fast lookup.
|
|
///
|
|
/// Live appends still write canonical JSONL history, but append-derived
|
|
/// metadata is observed above the store and applied through
|
|
/// [`ThreadStore::update_thread_metadata`]. This implementation applies that
|
|
/// patch literally to SQLite while keeping the JSONL/name-index compatibility
|
|
/// behavior needed for SQLite-less reads, repair, and old local rollout files.
|
|
#[derive(Clone)]
|
|
pub struct LocalThreadStore {
|
|
pub(super) config: LocalThreadStoreConfig,
|
|
live_recorders: Arc<Mutex<HashMap<ThreadId, RolloutRecorder>>>,
|
|
state_db: Option<StateDbHandle>,
|
|
}
|
|
|
|
/// Process-scoped configuration for local thread storage.
|
|
///
|
|
/// This describes where local storage lives. New-thread rollout metadata such
|
|
/// as cwd, provider, and memory mode is supplied when live persistence is opened.
|
|
#[derive(Clone, Debug, PartialEq, Eq)]
|
|
pub struct LocalThreadStoreConfig {
|
|
pub codex_home: PathBuf,
|
|
pub sqlite_home: PathBuf,
|
|
/// Provider used only when older local metadata does not contain one.
|
|
pub default_model_provider_id: String,
|
|
}
|
|
|
|
impl LocalThreadStoreConfig {
|
|
pub fn from_config(config: &impl codex_rollout::RolloutConfigView) -> Self {
|
|
Self {
|
|
codex_home: config.codex_home().to_path_buf(),
|
|
sqlite_home: config.sqlite_home().to_path_buf(),
|
|
default_model_provider_id: config.model_provider_id().to_string(),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl std::fmt::Debug for LocalThreadStore {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
f.debug_struct("LocalThreadStore")
|
|
.field("config", &self.config)
|
|
.finish_non_exhaustive()
|
|
}
|
|
}
|
|
|
|
impl LocalThreadStore {
|
|
/// Create a local store using an already initialized state DB handle.
|
|
pub fn new(config: LocalThreadStoreConfig, state_db: Option<StateDbHandle>) -> Self {
|
|
Self {
|
|
config,
|
|
live_recorders: Arc::new(Mutex::new(HashMap::new())),
|
|
state_db,
|
|
}
|
|
}
|
|
|
|
/// Return the state DB handle used by local rollout writers.
|
|
pub async fn state_db(&self) -> Option<StateDbHandle> {
|
|
self.state_db.clone()
|
|
}
|
|
|
|
/// Read a local rollout-backed thread by path.
|
|
pub async fn read_thread_by_rollout_path(
|
|
&self,
|
|
rollout_path: PathBuf,
|
|
include_archived: bool,
|
|
include_history: bool,
|
|
) -> ThreadStoreResult<StoredThread> {
|
|
read_thread::read_thread_by_rollout_path(
|
|
self,
|
|
rollout_path,
|
|
include_archived,
|
|
include_history,
|
|
)
|
|
.await
|
|
}
|
|
|
|
/// Return the live local rollout path for legacy local-only code paths.
|
|
pub async fn live_rollout_path(&self, thread_id: ThreadId) -> ThreadStoreResult<PathBuf> {
|
|
live_writer::rollout_path(self, thread_id).await
|
|
}
|
|
|
|
pub(super) async fn live_recorder(
|
|
&self,
|
|
thread_id: ThreadId,
|
|
) -> ThreadStoreResult<RolloutRecorder> {
|
|
self.live_recorders
|
|
.lock()
|
|
.await
|
|
.get(&thread_id)
|
|
.cloned()
|
|
.ok_or(ThreadStoreError::ThreadNotFound { thread_id })
|
|
}
|
|
|
|
pub(super) async fn ensure_live_recorder_absent(
|
|
&self,
|
|
thread_id: ThreadId,
|
|
) -> ThreadStoreResult<()> {
|
|
if self.live_recorders.lock().await.contains_key(&thread_id) {
|
|
return Err(ThreadStoreError::InvalidRequest {
|
|
message: format!("thread {thread_id} already has a live local writer"),
|
|
});
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
pub(super) async fn insert_live_recorder(
|
|
&self,
|
|
thread_id: ThreadId,
|
|
recorder: RolloutRecorder,
|
|
) -> ThreadStoreResult<()> {
|
|
match self.live_recorders.lock().await.entry(thread_id) {
|
|
Entry::Occupied(entry) => Err(ThreadStoreError::InvalidRequest {
|
|
message: format!("thread {} already has a live local writer", entry.key()),
|
|
}),
|
|
Entry::Vacant(entry) => {
|
|
entry.insert(recorder);
|
|
Ok(())
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn load_history(
|
|
&self,
|
|
params: LoadThreadHistoryParams,
|
|
) -> ThreadStoreResult<StoredThreadHistory> {
|
|
if let Ok(rollout_path) = live_writer::rollout_path(self, params.thread_id).await {
|
|
if !params.include_archived
|
|
&& helpers::rollout_path_is_archived(
|
|
self.config.codex_home.as_path(),
|
|
rollout_path.as_path(),
|
|
)
|
|
{
|
|
return Err(ThreadStoreError::InvalidRequest {
|
|
message: format!("thread {} is archived", params.thread_id),
|
|
});
|
|
}
|
|
return read_thread::read_thread_by_rollout_path(
|
|
self,
|
|
rollout_path,
|
|
/*include_archived*/ true,
|
|
/*include_history*/ true,
|
|
)
|
|
.await?
|
|
.history
|
|
.ok_or_else(|| ThreadStoreError::Internal {
|
|
message: format!("failed to load history for thread {}", params.thread_id),
|
|
});
|
|
}
|
|
|
|
read_thread::read_thread(
|
|
self,
|
|
ReadThreadParams {
|
|
thread_id: params.thread_id,
|
|
include_archived: params.include_archived,
|
|
include_history: true,
|
|
},
|
|
)
|
|
.await?
|
|
.history
|
|
.ok_or_else(|| ThreadStoreError::Internal {
|
|
message: format!("failed to load history for thread {}", params.thread_id),
|
|
})
|
|
}
|
|
|
|
async fn read_thread_by_rollout_path_params(
|
|
&self,
|
|
params: ReadThreadByRolloutPathParams,
|
|
) -> ThreadStoreResult<StoredThread> {
|
|
read_thread::read_thread_by_rollout_path(
|
|
self,
|
|
params.rollout_path,
|
|
params.include_archived,
|
|
params.include_history,
|
|
)
|
|
.await
|
|
}
|
|
}
|
|
|
|
impl ThreadStore for LocalThreadStore {
|
|
fn as_any(&self) -> &dyn std::any::Any {
|
|
self
|
|
}
|
|
|
|
fn create_thread(&self, params: CreateThreadParams) -> ThreadStoreFuture<'_, ()> {
|
|
Box::pin(async move { live_writer::create_thread(self, params).await })
|
|
}
|
|
|
|
fn resume_thread(&self, params: ResumeThreadParams) -> ThreadStoreFuture<'_, ()> {
|
|
Box::pin(async move { live_writer::resume_thread(self, params).await })
|
|
}
|
|
|
|
fn append_items(&self, params: AppendThreadItemsParams) -> ThreadStoreFuture<'_, ()> {
|
|
Box::pin(async move { live_writer::append_items(self, params).await })
|
|
}
|
|
|
|
fn persist_thread(&self, thread_id: ThreadId) -> ThreadStoreFuture<'_, ()> {
|
|
Box::pin(async move { live_writer::persist_thread(self, thread_id).await })
|
|
}
|
|
|
|
fn flush_thread(&self, thread_id: ThreadId) -> ThreadStoreFuture<'_, ()> {
|
|
Box::pin(async move { live_writer::flush_thread(self, thread_id).await })
|
|
}
|
|
|
|
fn shutdown_thread(&self, thread_id: ThreadId) -> ThreadStoreFuture<'_, ()> {
|
|
Box::pin(async move { live_writer::shutdown_thread(self, thread_id).await })
|
|
}
|
|
|
|
fn discard_thread(&self, thread_id: ThreadId) -> ThreadStoreFuture<'_, ()> {
|
|
Box::pin(async move { live_writer::discard_thread(self, thread_id).await })
|
|
}
|
|
|
|
fn load_history(
|
|
&self,
|
|
params: LoadThreadHistoryParams,
|
|
) -> ThreadStoreFuture<'_, StoredThreadHistory> {
|
|
Box::pin(LocalThreadStore::load_history(self, params))
|
|
}
|
|
|
|
fn read_thread(&self, params: ReadThreadParams) -> ThreadStoreFuture<'_, StoredThread> {
|
|
Box::pin(async move { read_thread::read_thread(self, params).await })
|
|
}
|
|
|
|
fn read_thread_by_rollout_path(
|
|
&self,
|
|
params: ReadThreadByRolloutPathParams,
|
|
) -> ThreadStoreFuture<'_, StoredThread> {
|
|
Box::pin(LocalThreadStore::read_thread_by_rollout_path_params(
|
|
self, params,
|
|
))
|
|
}
|
|
|
|
fn list_threads(&self, params: ListThreadsParams) -> ThreadStoreFuture<'_, ThreadPage> {
|
|
Box::pin(async move { list_threads::list_threads(self, params).await })
|
|
}
|
|
|
|
fn search_threads(
|
|
&self,
|
|
params: SearchThreadsParams,
|
|
) -> ThreadStoreFuture<'_, ThreadSearchPage> {
|
|
Box::pin(async move { search_threads::search_threads(self, params).await })
|
|
}
|
|
|
|
fn update_thread_metadata(
|
|
&self,
|
|
params: UpdateThreadMetadataParams,
|
|
) -> ThreadStoreFuture<'_, StoredThread> {
|
|
Box::pin(async move { update_thread_metadata::update_thread_metadata(self, params).await })
|
|
}
|
|
|
|
fn archive_thread(&self, params: ArchiveThreadParams) -> ThreadStoreFuture<'_, ()> {
|
|
Box::pin(async move { archive_thread::archive_thread(self, params).await })
|
|
}
|
|
|
|
fn unarchive_thread(&self, params: ArchiveThreadParams) -> ThreadStoreFuture<'_, StoredThread> {
|
|
Box::pin(async move { unarchive_thread::unarchive_thread(self, params).await })
|
|
}
|
|
|
|
fn delete_thread(&self, params: DeleteThreadParams) -> ThreadStoreFuture<'_, ()> {
|
|
Box::pin(async move { delete_thread::delete_thread(self, params).await })
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use std::sync::Arc;
|
|
|
|
use codex_protocol::ThreadId;
|
|
use codex_protocol::models::BaseInstructions;
|
|
use codex_protocol::models::FunctionCallOutputPayload;
|
|
use codex_protocol::models::MessagePhase;
|
|
use codex_protocol::models::ResponseItem;
|
|
use codex_protocol::protocol::AgentMessageEvent;
|
|
use codex_protocol::protocol::EventMsg;
|
|
use codex_protocol::protocol::RolloutItem;
|
|
use codex_protocol::protocol::SessionSource;
|
|
use codex_protocol::protocol::ThreadMemoryMode;
|
|
use codex_protocol::protocol::TurnCompleteEvent;
|
|
use codex_protocol::protocol::TurnStartedEvent;
|
|
use codex_protocol::protocol::UserMessageEvent;
|
|
use tempfile::TempDir;
|
|
|
|
use super::*;
|
|
use crate::LiveThread;
|
|
use crate::ThreadPersistenceMetadata;
|
|
use crate::local::test_support::test_config;
|
|
use crate::local::test_support::write_archived_session_file;
|
|
use crate::local::test_support::write_session_file;
|
|
|
|
#[tokio::test]
|
|
async fn live_writer_lifecycle_writes_and_closes() {
|
|
let home = TempDir::new().expect("temp dir");
|
|
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
|
|
let thread_id = ThreadId::default();
|
|
|
|
store
|
|
.create_thread(create_thread_params(thread_id))
|
|
.await
|
|
.expect("create live thread");
|
|
let rollout_path = store
|
|
.live_rollout_path(thread_id)
|
|
.await
|
|
.expect("load rollout path");
|
|
|
|
store
|
|
.append_items(AppendThreadItemsParams {
|
|
thread_id,
|
|
items: vec![user_message_item("first live write")],
|
|
})
|
|
.await
|
|
.expect("append live item");
|
|
store
|
|
.persist_thread(thread_id)
|
|
.await
|
|
.expect("persist live thread");
|
|
store
|
|
.flush_thread(thread_id)
|
|
.await
|
|
.expect("flush live thread");
|
|
|
|
assert_rollout_contains_message(rollout_path.as_path(), "first live write").await;
|
|
|
|
store
|
|
.shutdown_thread(thread_id)
|
|
.await
|
|
.expect("shutdown live thread");
|
|
let err = store
|
|
.append_items(AppendThreadItemsParams {
|
|
thread_id,
|
|
items: vec![user_message_item("write after shutdown")],
|
|
})
|
|
.await
|
|
.expect_err("shutdown should remove the live thread writer");
|
|
assert!(
|
|
matches!(err, ThreadStoreError::ThreadNotFound { thread_id: missing } if missing == thread_id)
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn raw_append_items_does_not_update_sqlite_metadata() {
|
|
// This pins the ThreadStore contract: raw appends are history-only. Callers that need
|
|
// metadata updates must use LiveThread or call update_thread_metadata explicitly.
|
|
let home = TempDir::new().expect("temp dir");
|
|
let config = test_config(home.path());
|
|
let runtime = codex_state::StateRuntime::init(
|
|
config.sqlite_home.clone(),
|
|
config.default_model_provider_id.clone(),
|
|
)
|
|
.await
|
|
.expect("state db should initialize");
|
|
let store = LocalThreadStore::new(config, Some(runtime.clone()));
|
|
let thread_id = ThreadId::default();
|
|
|
|
store
|
|
.create_thread(create_thread_params(thread_id))
|
|
.await
|
|
.expect("create live thread");
|
|
store
|
|
.append_items(AppendThreadItemsParams {
|
|
thread_id,
|
|
items: vec![user_message_item("raw append")],
|
|
})
|
|
.await
|
|
.expect("append raw item");
|
|
store.flush_thread(thread_id).await.expect("flush thread");
|
|
|
|
assert_eq!(
|
|
runtime
|
|
.get_thread(thread_id)
|
|
.await
|
|
.expect("sqlite metadata read"),
|
|
None
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn live_thread_observes_appended_items_into_sqlite_metadata() {
|
|
let home = TempDir::new().expect("temp dir");
|
|
let config = test_config(home.path());
|
|
let runtime = codex_state::StateRuntime::init(
|
|
config.sqlite_home.clone(),
|
|
config.default_model_provider_id.clone(),
|
|
)
|
|
.await
|
|
.expect("state db should initialize");
|
|
let store = Arc::new(LocalThreadStore::new(config, Some(runtime.clone())));
|
|
let thread_id = ThreadId::default();
|
|
let live_thread = LiveThread::create(store.clone(), create_thread_params(thread_id))
|
|
.await
|
|
.expect("create live thread");
|
|
|
|
live_thread
|
|
.append_items(&[user_message_item("observed append")])
|
|
.await
|
|
.expect("append observed item");
|
|
live_thread.flush().await.expect("flush thread");
|
|
|
|
let metadata = runtime
|
|
.get_thread(thread_id)
|
|
.await
|
|
.expect("sqlite metadata read")
|
|
.expect("sqlite metadata");
|
|
assert_eq!(
|
|
metadata.first_user_message.as_deref(),
|
|
Some("observed append")
|
|
);
|
|
assert_eq!(metadata.preview.as_deref(), Some("observed append"));
|
|
assert_eq!(metadata.title, "observed append");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn live_thread_output_advances_updated_at_but_not_recency_at() {
|
|
let home = TempDir::new().expect("temp dir");
|
|
let config = test_config(home.path());
|
|
let runtime = codex_state::StateRuntime::init(
|
|
config.sqlite_home.clone(),
|
|
config.default_model_provider_id.clone(),
|
|
)
|
|
.await
|
|
.expect("state db should initialize");
|
|
let store = Arc::new(LocalThreadStore::new(config, Some(runtime.clone())));
|
|
let thread_id = ThreadId::default();
|
|
let live_thread = LiveThread::create(store, create_thread_params(thread_id))
|
|
.await
|
|
.expect("create live thread");
|
|
|
|
live_thread
|
|
.append_items(&[user_message_item("start thread")])
|
|
.await
|
|
.expect("append initial user message");
|
|
live_thread.flush().await.expect("flush thread");
|
|
let before_turn_start = runtime
|
|
.get_thread(thread_id)
|
|
.await
|
|
.expect("sqlite metadata read")
|
|
.expect("sqlite metadata");
|
|
|
|
live_thread
|
|
.append_items(&[RolloutItem::EventMsg(EventMsg::TurnStarted(
|
|
TurnStartedEvent {
|
|
turn_id: "turn-1".to_string(),
|
|
trace_id: None,
|
|
started_at: None,
|
|
model_context_window: None,
|
|
collaboration_mode_kind: Default::default(),
|
|
},
|
|
))])
|
|
.await
|
|
.expect("append turn start");
|
|
live_thread.flush().await.expect("flush thread");
|
|
let after_turn_start = runtime
|
|
.get_thread(thread_id)
|
|
.await
|
|
.expect("sqlite metadata read")
|
|
.expect("sqlite metadata");
|
|
assert!(after_turn_start.recency_at > before_turn_start.recency_at);
|
|
|
|
live_thread
|
|
.append_items(&[
|
|
RolloutItem::EventMsg(EventMsg::AgentMessage(AgentMessageEvent {
|
|
message: "commentary".to_string(),
|
|
phase: Some(MessagePhase::Commentary),
|
|
memory_citation: None,
|
|
})),
|
|
RolloutItem::ResponseItem(ResponseItem::FunctionCallOutput {
|
|
id: None,
|
|
call_id: "call-1".to_string(),
|
|
output: FunctionCallOutputPayload::from_text("tool output".to_string()),
|
|
internal_chat_message_metadata_passthrough: None,
|
|
}),
|
|
RolloutItem::EventMsg(EventMsg::TokenCount(
|
|
codex_protocol::protocol::TokenCountEvent {
|
|
info: None,
|
|
rate_limits: None,
|
|
},
|
|
)),
|
|
RolloutItem::EventMsg(EventMsg::TurnComplete(TurnCompleteEvent {
|
|
turn_id: "turn-1".to_string(),
|
|
last_agent_message: None,
|
|
completed_at: None,
|
|
duration_ms: None,
|
|
time_to_first_token_ms: None,
|
|
})),
|
|
])
|
|
.await
|
|
.expect("append post-start items");
|
|
live_thread.flush().await.expect("flush thread");
|
|
let completed = runtime
|
|
.get_thread(thread_id)
|
|
.await
|
|
.expect("sqlite metadata read")
|
|
.expect("sqlite metadata");
|
|
|
|
assert!(completed.updated_at > after_turn_start.updated_at);
|
|
assert_eq!(completed.recency_at, after_turn_start.recency_at);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn live_thread_shutdown_does_not_materialize_empty_thread_metadata() {
|
|
let home = TempDir::new().expect("temp dir");
|
|
let config = test_config(home.path());
|
|
let runtime = codex_state::StateRuntime::init(
|
|
config.sqlite_home.clone(),
|
|
config.default_model_provider_id.clone(),
|
|
)
|
|
.await
|
|
.expect("state db should initialize");
|
|
let store = Arc::new(LocalThreadStore::new(config, Some(runtime.clone())));
|
|
let thread_id = ThreadId::default();
|
|
let live_thread = LiveThread::create(store.clone(), create_thread_params(thread_id))
|
|
.await
|
|
.expect("create live thread");
|
|
let rollout_path = store
|
|
.live_rollout_path(thread_id)
|
|
.await
|
|
.expect("live rollout path");
|
|
|
|
live_thread.shutdown().await.expect("shutdown thread");
|
|
|
|
assert!(
|
|
!tokio::fs::try_exists(rollout_path.as_path())
|
|
.await
|
|
.expect("rollout path should be checkable")
|
|
);
|
|
assert_eq!(
|
|
runtime
|
|
.get_thread(thread_id)
|
|
.await
|
|
.expect("sqlite metadata read"),
|
|
None
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn live_thread_shutdown_with_buffered_items_materializes_before_metadata_read() {
|
|
let home = TempDir::new().expect("temp dir");
|
|
let config = test_config(home.path());
|
|
let runtime = codex_state::StateRuntime::init(
|
|
config.sqlite_home.clone(),
|
|
config.default_model_provider_id.clone(),
|
|
)
|
|
.await
|
|
.expect("state db should initialize");
|
|
let store = Arc::new(LocalThreadStore::new(config, Some(runtime.clone())));
|
|
let thread_id = ThreadId::default();
|
|
let live_thread = LiveThread::create(store.clone(), create_thread_params(thread_id))
|
|
.await
|
|
.expect("create live thread");
|
|
let rollout_path = store
|
|
.live_rollout_path(thread_id)
|
|
.await
|
|
.expect("live rollout path");
|
|
|
|
live_thread
|
|
.append_items(&[RolloutItem::EventMsg(EventMsg::TokenCount(
|
|
codex_protocol::protocol::TokenCountEvent {
|
|
info: None,
|
|
rate_limits: None,
|
|
},
|
|
))])
|
|
.await
|
|
.expect("append metadata-only item");
|
|
live_thread.shutdown().await.expect("shutdown thread");
|
|
|
|
assert!(
|
|
tokio::fs::try_exists(rollout_path.as_path())
|
|
.await
|
|
.expect("rollout path should be checkable")
|
|
);
|
|
let metadata = runtime
|
|
.get_thread(thread_id)
|
|
.await
|
|
.expect("sqlite metadata read")
|
|
.expect("sqlite metadata");
|
|
assert_eq!(metadata.rollout_path, rollout_path);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn live_thread_resume_loads_history_before_observing_metadata() {
|
|
let home = TempDir::new().expect("temp dir");
|
|
let config = test_config(home.path());
|
|
let runtime = codex_state::StateRuntime::init(
|
|
config.sqlite_home.clone(),
|
|
config.default_model_provider_id.clone(),
|
|
)
|
|
.await
|
|
.expect("state db should initialize");
|
|
let store = Arc::new(LocalThreadStore::new(config, Some(runtime.clone())));
|
|
let uuid = uuid::Uuid::from_u128(401);
|
|
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
|
|
let rollout_path =
|
|
write_session_file(home.path(), "2025-01-03T17-00-00", uuid).expect("session file");
|
|
let live_thread = LiveThread::resume(
|
|
store,
|
|
ResumeThreadParams {
|
|
thread_id,
|
|
rollout_path: Some(rollout_path),
|
|
history: None,
|
|
include_archived: false,
|
|
metadata: ThreadPersistenceMetadata {
|
|
cwd: Some(home.path().to_path_buf()),
|
|
model_provider: "different-provider".to_string(),
|
|
memory_mode: ThreadMemoryMode::Enabled,
|
|
},
|
|
},
|
|
)
|
|
.await
|
|
.expect("resume live thread");
|
|
|
|
live_thread
|
|
.append_items(&[user_message_item("new live append")])
|
|
.await
|
|
.expect("append after resume");
|
|
|
|
let metadata = runtime
|
|
.get_thread(thread_id)
|
|
.await
|
|
.expect("sqlite metadata read")
|
|
.expect("sqlite metadata");
|
|
assert_eq!(
|
|
metadata.created_at.to_rfc3339(),
|
|
"2025-01-03T17:00:00+00:00"
|
|
);
|
|
assert_eq!(metadata.model_provider, "test-provider");
|
|
assert_eq!(
|
|
metadata.first_user_message.as_deref(),
|
|
Some("Hello from user")
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn live_thread_resume_loads_history_from_explicit_external_rollout_path() {
|
|
let home = TempDir::new().expect("temp dir");
|
|
let external_home = TempDir::new().expect("external temp dir");
|
|
let config = test_config(home.path());
|
|
let runtime = codex_state::StateRuntime::init(
|
|
config.sqlite_home.clone(),
|
|
config.default_model_provider_id.clone(),
|
|
)
|
|
.await
|
|
.expect("state db should initialize");
|
|
let store = Arc::new(LocalThreadStore::new(config, Some(runtime.clone())));
|
|
let uuid = uuid::Uuid::from_u128(402);
|
|
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
|
|
let rollout_path = write_session_file(external_home.path(), "2025-01-03T17-30-00", uuid)
|
|
.expect("external session file");
|
|
let live_thread = LiveThread::resume(
|
|
store,
|
|
ResumeThreadParams {
|
|
thread_id,
|
|
rollout_path: Some(rollout_path),
|
|
history: None,
|
|
include_archived: false,
|
|
metadata: ThreadPersistenceMetadata {
|
|
cwd: Some(home.path().to_path_buf()),
|
|
model_provider: "different-provider".to_string(),
|
|
memory_mode: ThreadMemoryMode::Enabled,
|
|
},
|
|
},
|
|
)
|
|
.await
|
|
.expect("resume external live thread");
|
|
|
|
live_thread
|
|
.append_items(&[user_message_item("new external append")])
|
|
.await
|
|
.expect("append after external resume");
|
|
|
|
let metadata = runtime
|
|
.get_thread(thread_id)
|
|
.await
|
|
.expect("sqlite metadata read")
|
|
.expect("sqlite metadata");
|
|
assert_eq!(
|
|
metadata.created_at.to_rfc3339(),
|
|
"2025-01-03T17:30:00+00:00"
|
|
);
|
|
assert_eq!(metadata.model_provider, "test-provider");
|
|
assert_eq!(
|
|
metadata.first_user_message.as_deref(),
|
|
Some("Hello from user")
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn create_thread_rejects_missing_cwd() {
|
|
let home = TempDir::new().expect("temp dir");
|
|
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
|
|
let thread_id = ThreadId::default();
|
|
let mut params = create_thread_params(thread_id);
|
|
params.metadata.cwd = None;
|
|
|
|
let err = store
|
|
.create_thread(params)
|
|
.await
|
|
.expect_err("local thread store should require cwd");
|
|
|
|
assert!(matches!(
|
|
err,
|
|
ThreadStoreError::InvalidRequest { message }
|
|
if message == "local thread store requires a cwd"
|
|
));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn discard_thread_drops_unmaterialized_live_writer() {
|
|
let home = TempDir::new().expect("temp dir");
|
|
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
|
|
let thread_id = ThreadId::default();
|
|
|
|
store
|
|
.create_thread(create_thread_params(thread_id))
|
|
.await
|
|
.expect("create live thread");
|
|
let rollout_path = store
|
|
.live_rollout_path(thread_id)
|
|
.await
|
|
.expect("load rollout path");
|
|
store
|
|
.discard_thread(thread_id)
|
|
.await
|
|
.expect("discard live thread");
|
|
|
|
assert!(
|
|
!tokio::fs::try_exists(rollout_path.as_path())
|
|
.await
|
|
.expect("check rollout path")
|
|
);
|
|
let err = store
|
|
.append_items(AppendThreadItemsParams {
|
|
thread_id,
|
|
items: vec![user_message_item("write after discard")],
|
|
})
|
|
.await
|
|
.expect_err("discard should remove the live thread writer");
|
|
assert!(
|
|
matches!(err, ThreadStoreError::ThreadNotFound { thread_id: missing } if missing == thread_id)
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn resume_thread_reopens_live_writer_and_appends() {
|
|
let home = TempDir::new().expect("temp dir");
|
|
let config = test_config(home.path());
|
|
let thread_id = ThreadId::default();
|
|
|
|
let first_store = LocalThreadStore::new(config.clone(), /*state_db*/ None);
|
|
first_store
|
|
.create_thread(create_thread_params(thread_id))
|
|
.await
|
|
.expect("create initial thread");
|
|
first_store
|
|
.append_items(AppendThreadItemsParams {
|
|
thread_id,
|
|
items: vec![user_message_item("before resume")],
|
|
})
|
|
.await
|
|
.expect("append initial item");
|
|
first_store
|
|
.persist_thread(thread_id)
|
|
.await
|
|
.expect("persist initial thread");
|
|
first_store
|
|
.flush_thread(thread_id)
|
|
.await
|
|
.expect("flush initial thread");
|
|
let rollout_path = first_store
|
|
.live_rollout_path(thread_id)
|
|
.await
|
|
.expect("load rollout path");
|
|
first_store
|
|
.shutdown_thread(thread_id)
|
|
.await
|
|
.expect("shutdown initial writer");
|
|
|
|
let resumed_store = LocalThreadStore::new(config, /*state_db*/ None);
|
|
resumed_store
|
|
.resume_thread(ResumeThreadParams {
|
|
thread_id,
|
|
rollout_path: None,
|
|
history: None,
|
|
include_archived: true,
|
|
metadata: thread_metadata(),
|
|
})
|
|
.await
|
|
.expect("resume live thread");
|
|
resumed_store
|
|
.append_items(AppendThreadItemsParams {
|
|
thread_id,
|
|
items: vec![user_message_item("after resume")],
|
|
})
|
|
.await
|
|
.expect("append resumed item");
|
|
resumed_store
|
|
.flush_thread(thread_id)
|
|
.await
|
|
.expect("flush resumed thread");
|
|
|
|
assert_rollout_contains_message(rollout_path.as_path(), "before resume").await;
|
|
assert_rollout_contains_message(rollout_path.as_path(), "after resume").await;
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn create_thread_rejects_duplicate_live_writer() {
|
|
let home = TempDir::new().expect("temp dir");
|
|
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
|
|
let thread_id = ThreadId::default();
|
|
|
|
store
|
|
.create_thread(create_thread_params(thread_id))
|
|
.await
|
|
.expect("create live thread");
|
|
|
|
let err = store
|
|
.create_thread(create_thread_params(thread_id))
|
|
.await
|
|
.expect_err("duplicate live writer should fail");
|
|
|
|
assert!(matches!(err, ThreadStoreError::InvalidRequest { .. }));
|
|
assert!(err.to_string().contains("already has a live local writer"));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn resume_thread_rejects_duplicate_live_writer() {
|
|
let home = TempDir::new().expect("temp dir");
|
|
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
|
|
let thread_id = ThreadId::default();
|
|
|
|
store
|
|
.create_thread(create_thread_params(thread_id))
|
|
.await
|
|
.expect("create live thread");
|
|
let rollout_path = store
|
|
.live_rollout_path(thread_id)
|
|
.await
|
|
.expect("live rollout path");
|
|
let err = store
|
|
.resume_thread(ResumeThreadParams {
|
|
thread_id,
|
|
rollout_path: Some(rollout_path),
|
|
history: None,
|
|
include_archived: true,
|
|
metadata: thread_metadata(),
|
|
})
|
|
.await
|
|
.expect_err("duplicate live resume should fail");
|
|
assert!(matches!(err, ThreadStoreError::InvalidRequest { .. }));
|
|
assert!(err.to_string().contains("already has a live local writer"));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn resume_thread_rejects_missing_cwd() {
|
|
let home = TempDir::new().expect("temp dir");
|
|
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
|
|
let uuid = uuid::Uuid::from_u128(407);
|
|
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
|
|
let rollout_path =
|
|
write_session_file(home.path(), "2025-01-04T11-30-00", uuid).expect("session file");
|
|
let err = store
|
|
.resume_thread(ResumeThreadParams {
|
|
thread_id,
|
|
rollout_path: Some(rollout_path),
|
|
history: None,
|
|
include_archived: true,
|
|
metadata: ThreadPersistenceMetadata {
|
|
cwd: None,
|
|
model_provider: "test-provider".to_string(),
|
|
memory_mode: ThreadMemoryMode::Enabled,
|
|
},
|
|
})
|
|
.await
|
|
.expect_err("missing cwd should fail");
|
|
|
|
assert!(matches!(err, ThreadStoreError::InvalidRequest { .. }));
|
|
assert!(err.to_string().contains("requires a cwd"));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn load_history_uses_live_writer_rollout_path() {
|
|
let home = TempDir::new().expect("temp dir");
|
|
let external_home = TempDir::new().expect("external temp dir");
|
|
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
|
|
let uuid = uuid::Uuid::from_u128(404);
|
|
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
|
|
let rollout_path = write_session_file(external_home.path(), "2025-01-04T10-00-00", uuid)
|
|
.expect("external session file");
|
|
|
|
store
|
|
.resume_thread(ResumeThreadParams {
|
|
thread_id,
|
|
rollout_path: Some(rollout_path),
|
|
history: None,
|
|
include_archived: true,
|
|
metadata: thread_metadata(),
|
|
})
|
|
.await
|
|
.expect("resume live thread");
|
|
store
|
|
.append_items(AppendThreadItemsParams {
|
|
thread_id,
|
|
items: vec![user_message_item("external history item")],
|
|
})
|
|
.await
|
|
.expect("append live item");
|
|
store
|
|
.flush_thread(thread_id)
|
|
.await
|
|
.expect("flush live thread");
|
|
|
|
let history = store
|
|
.load_history(LoadThreadHistoryParams {
|
|
thread_id,
|
|
include_archived: false,
|
|
})
|
|
.await
|
|
.expect("load external live history");
|
|
|
|
assert!(history.items.iter().any(|item| {
|
|
matches!(
|
|
item,
|
|
RolloutItem::EventMsg(EventMsg::UserMessage(event)) if event.message == "external history item"
|
|
)
|
|
}));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn read_thread_uses_live_writer_rollout_path_for_external_resume() {
|
|
let home = TempDir::new().expect("temp dir");
|
|
let external_home = TempDir::new().expect("external temp dir");
|
|
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
|
|
let uuid = uuid::Uuid::from_u128(406);
|
|
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
|
|
let rollout_path = write_session_file(external_home.path(), "2025-01-04T11-00-00", uuid)
|
|
.expect("external session file");
|
|
|
|
store
|
|
.resume_thread(ResumeThreadParams {
|
|
thread_id,
|
|
rollout_path: Some(rollout_path.clone()),
|
|
history: None,
|
|
include_archived: true,
|
|
metadata: thread_metadata(),
|
|
})
|
|
.await
|
|
.expect("resume live thread");
|
|
|
|
let thread = store
|
|
.read_thread(ReadThreadParams {
|
|
thread_id,
|
|
include_archived: false,
|
|
include_history: true,
|
|
})
|
|
.await
|
|
.expect("read external live thread");
|
|
|
|
assert_eq!(thread.rollout_path, Some(rollout_path));
|
|
assert!(thread.history.expect("history").items.iter().any(|item| {
|
|
matches!(
|
|
item,
|
|
RolloutItem::EventMsg(EventMsg::UserMessage(event)) if event.message == "Hello from user"
|
|
)
|
|
}));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn load_history_uses_live_writer_rollout_path_for_archived_source() {
|
|
let home = TempDir::new().expect("temp dir");
|
|
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
|
|
let uuid = uuid::Uuid::from_u128(405);
|
|
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
|
|
let rollout_path = write_archived_session_file(home.path(), "2025-01-04T10-30-00", uuid)
|
|
.expect("archived session file");
|
|
|
|
store
|
|
.resume_thread(ResumeThreadParams {
|
|
thread_id,
|
|
rollout_path: Some(rollout_path),
|
|
history: None,
|
|
include_archived: true,
|
|
metadata: thread_metadata(),
|
|
})
|
|
.await
|
|
.expect("resume live archived thread");
|
|
store
|
|
.append_items(AppendThreadItemsParams {
|
|
thread_id,
|
|
items: vec![user_message_item("archived live history item")],
|
|
})
|
|
.await
|
|
.expect("append live item");
|
|
store
|
|
.flush_thread(thread_id)
|
|
.await
|
|
.expect("flush live thread");
|
|
|
|
let err = store
|
|
.read_thread(ReadThreadParams {
|
|
thread_id,
|
|
include_archived: false,
|
|
include_history: false,
|
|
})
|
|
.await
|
|
.expect_err("active-only read should reject archived live thread");
|
|
assert!(matches!(err, ThreadStoreError::InvalidRequest { .. }));
|
|
|
|
let err = store
|
|
.load_history(LoadThreadHistoryParams {
|
|
thread_id,
|
|
include_archived: false,
|
|
})
|
|
.await
|
|
.expect_err("active-only history should reject archived live thread");
|
|
assert!(matches!(err, ThreadStoreError::InvalidRequest { .. }));
|
|
assert!(err.to_string().contains("archived"));
|
|
|
|
let history = store
|
|
.load_history(LoadThreadHistoryParams {
|
|
thread_id,
|
|
include_archived: true,
|
|
})
|
|
.await
|
|
.expect("load archived live history");
|
|
|
|
assert!(history.items.iter().any(|item| {
|
|
matches!(
|
|
item,
|
|
RolloutItem::EventMsg(EventMsg::UserMessage(event)) if event.message == "archived live history item"
|
|
)
|
|
}));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn read_thread_by_rollout_path_includes_history() {
|
|
let home = TempDir::new().expect("temp dir");
|
|
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
|
|
let thread_id = ThreadId::default();
|
|
|
|
store
|
|
.create_thread(create_thread_params(thread_id))
|
|
.await
|
|
.expect("create thread");
|
|
store
|
|
.append_items(AppendThreadItemsParams {
|
|
thread_id,
|
|
items: vec![user_message_item("path read")],
|
|
})
|
|
.await
|
|
.expect("append item");
|
|
store.flush_thread(thread_id).await.expect("flush thread");
|
|
let rollout_path = store
|
|
.live_rollout_path(thread_id)
|
|
.await
|
|
.expect("load rollout path");
|
|
|
|
let thread = store
|
|
.read_thread_by_rollout_path(
|
|
rollout_path,
|
|
/*include_archived*/ true,
|
|
/*include_history*/ true,
|
|
)
|
|
.await
|
|
.expect("read thread by rollout path");
|
|
|
|
assert_eq!(thread.thread_id, thread_id);
|
|
assert_eq!(
|
|
thread
|
|
.history
|
|
.expect("history")
|
|
.items
|
|
.into_iter()
|
|
.filter(|item| matches!(item, RolloutItem::EventMsg(EventMsg::UserMessage(_))))
|
|
.count(),
|
|
1
|
|
);
|
|
}
|
|
|
|
fn create_thread_params(thread_id: ThreadId) -> CreateThreadParams {
|
|
CreateThreadParams {
|
|
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,
|
|
initial_window_id: uuid::Uuid::now_v7().to_string(),
|
|
metadata: thread_metadata(),
|
|
}
|
|
}
|
|
|
|
fn thread_metadata() -> ThreadPersistenceMetadata {
|
|
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()
|
|
}))
|
|
}
|
|
|
|
async fn assert_rollout_contains_message(path: &std::path::Path, expected: &str) {
|
|
let (items, _, _) = RolloutRecorder::load_rollout_items(path)
|
|
.await
|
|
.expect("load rollout items");
|
|
assert!(items.iter().any(|item| {
|
|
matches!(
|
|
item,
|
|
RolloutItem::EventMsg(EventMsg::UserMessage(event)) if event.message == expected
|
|
)
|
|
}));
|
|
}
|
|
}
|