agent/codex
2
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

[codex] Route live thread writes through ThreadStore (#18882)

Browse Source 
Begin migrating the thread write codepaths to ThreadStore.

This starts using ThreadStore inside of core session code, not only in
the app server code.

Rework the interfaces around thread recording/persistence. We're left
with the following:

* `ThreadManager`: owns the process-level registry of loaded threads and
handles cross-thread orchestration: start, resume, fork, lookup, remove,
and route ops to running CodexThreads.
* `CodexThread`: represents one loaded/running thread from the outside.
It is the handle app-server and callers use to submit ops, inspect
session metadata, and shut the thread down.
* `LiveThread`: session-owned persistence lifecycle handle for one
active thread. Core session code uses it to append rollout items,
materialize lazy persistence, flush, shutdown, discard init-failed
writers, and load that thread’s persisted history.
* `ThreadStore`: storage backend abstraction. It answers “how are
threads persisted, read, listed, updated, archived?” Local and remote
implementations live behind this trait.
* `LocalThreadStore`: local ThreadStore implementation. It owns the
file/sqlite-specific details and keeps RolloutRecorder as a local
implementation detail.

This is a few too many Thread abstractions for my liking, but they do
all represent different concepts / needs / layers.

Migration note: in places where the core code explicitly requires a
path, rather than a thread ID, throw an error if we're running with a
remote store.

Cover the new local live-writer lifecycle with focused tests and
preserve app-server thread-start behavior, including ephemeral pathless
sessions.
This commit is contained in:
Tom
2026-04-23 10:17:09 -07:00
committed by GitHub
Unverified
parent 3d3028a5a9
commit f1923a38b1
23 changed files with 1789 additions and 889 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/Cargo.lock
Generated
+1
View File
@@ -3380,6 +3380,7 @@ dependencies = [
"tonic",
"tonic-prost",
"tonic-prost-build",
"tracing",
"uuid",
]
codex-rs/core/src/codex_delegate.rs
+1
View File
@@ -95,6 +95,7 @@ pub(crate) async fn run_codex_thread_interactive(
inherited_rollout_trace: codex_rollout_trace::RolloutTraceRecorder::disabled(),
parent_trace: None,
analytics_events_client: Some(parent_session.services.analytics_events_client.clone()),
thread_store: Arc::clone(&parent_session.services.thread_store),
}))
.await?;
if parent_session.enabled(codex_features::Feature::GeneralAnalytics) {
codex-rs/core/src/guardian/review_session.rs
+1 -1
View File
@@ -751,7 +751,7 @@ async fn load_rollout_items_for_fork(
session: &Session,
) -> anyhow::Result<Option<Vec<RolloutItem>>> {
session.flush_rollout().await?;
let Some(rollout_path) = session.current_rollout_path().await else {
let Some(rollout_path) = session.current_rollout_path().await? else {
return Ok(None);
};
let history = RolloutRecorder::get_rollout_history(rollout_path.as_path()).await?;
codex-rs/core/src/realtime_context.rs
-1
View File
@@ -9,7 +9,6 @@ use codex_thread_store::ListThreadsParams;
use codex_thread_store::SortDirection;
use codex_thread_store::StoredThread;
use codex_thread_store::ThreadSortKey;
use codex_thread_store::ThreadStore;
use codex_utils_absolute_path::AbsolutePathBuf;
use codex_utils_output_truncation::TruncationPolicy;
use codex_utils_output_truncation::truncate_text;
codex-rs/core/src/rollout.rs
-5
View File
@@ -50,12 +50,7 @@ pub(crate) mod list {
pub use codex_rollout::find_thread_path_by_id_str;
}
pub(crate) mod metadata {
pub(crate) use codex_rollout::builder_from_items;
}
pub(crate) mod policy {
pub use codex_rollout::EventPersistenceMode;
pub use codex_rollout::should_persist_response_item_for_memories;
}
codex-rs/core/src/session/handlers.rs
+30 -75
View File
@@ -27,8 +27,6 @@ use codex_features::Feature;
use codex_utils_absolute_path::AbsolutePathBuf;
use crate::review_prompts::resolve_review_request;
use crate::rollout::RolloutRecorder;
use crate::rollout::read_session_meta_line;
use crate::tasks::CompactTask;
use crate::tasks::UndoTask;
use crate::tasks::UserShellCommandMode;
@@ -757,36 +755,25 @@ pub async fn thread_rollback(sess: &Arc<Session>, sub_id: String, num_turns: u32
}
let turn_context = sess.new_default_turn_with_sub_id(sub_id).await;
let rollout_path = {
let recorder = {
let guard = sess.services.rollout.lock().await;
guard.clone()
};
let Some(recorder) = recorder else {
let live_thread = match sess.live_thread_for_persistence("rollback thread") {
Ok(live_thread) => live_thread,
Err(_) => {
sess.send_event_raw(Event {
id: turn_context.sub_id.clone(),
msg: EventMsg::Error(ErrorEvent {
message: "thread rollback requires a persisted rollout path".to_string(),
message: "thread rollback requires persisted thread history".to_string(),
codex_error_info: Some(CodexErrorInfo::ThreadRollbackFailed),
}),
})
.await;
return;
};
recorder.rollout_path().to_path_buf()
}
};
if let Some(recorder) = {
let guard = sess.services.rollout.lock().await;
guard.clone()
} && let Err(err) = recorder.flush().await
{
if let Err(err) = live_thread.flush().await {
sess.send_event_raw(Event {
id: turn_context.sub_id.clone(),
msg: EventMsg::Error(ErrorEvent {
message: format!(
"failed to flush rollout `{}` for rollback replay: {err}",
rollout_path.display()
),
message: format!("failed to flush thread persistence for rollback replay: {err}"),
codex_error_info: Some(CodexErrorInfo::ThreadRollbackFailed),
}),
})
@@ -794,16 +781,13 @@ pub async fn thread_rollback(sess: &Arc<Session>, sub_id: String, num_turns: u32
return;
}
let initial_history = match RolloutRecorder::get_rollout_history(rollout_path.as_path()).await {
let stored_history = match live_thread.load_history(/*include_archived*/ false).await {
Ok(history) => history,
Err(err) => {
sess.send_event_raw(Event {
id: turn_context.sub_id.clone(),
msg: EventMsg::Error(ErrorEvent {
message: format!(
"failed to load rollout `{}` for rollback replay: {err}",
rollout_path.display()
),
message: format!("failed to load thread history for rollback replay: {err}"),
codex_error_info: Some(CodexErrorInfo::ThreadRollbackFailed),
}),
})
@@ -814,8 +798,8 @@ pub async fn thread_rollback(sess: &Arc<Session>, sub_id: String, num_turns: u32
let rollback_event = ThreadRolledBackEvent { num_turns };
let rollback_msg = EventMsg::ThreadRolledBack(rollback_event.clone());
let replay_items = initial_history
.get_rollout_items()
let replay_items = stored_history
.items
.into_iter()
.chain(std::iter::once(RolloutItem::EventMsg(rollback_msg.clone())))
.collect::<Vec<_>>();
@@ -850,14 +834,12 @@ async fn persist_thread_name_update(
) -> anyhow::Result<EventMsg> {
let msg = EventMsg::ThreadNameUpdated(event);
let item = RolloutItem::EventMsg(msg.clone());
let recorder = {
let guard = sess.services.rollout.lock().await;
guard.clone()
}
.ok_or_else(|| anyhow::anyhow!("Session persistence is disabled; cannot rename thread."))?;
recorder.persist().await?;
recorder.record_items(std::slice::from_ref(&item)).await?;
recorder.flush().await?;
let live_thread = sess.live_thread_for_persistence("rename thread")?;
live_thread.persist().await?;
live_thread
.append_items(std::slice::from_ref(&item))
.await?;
live_thread.flush().await?;
Ok(msg)
}
@@ -865,36 +847,13 @@ pub(super) async fn persist_thread_memory_mode_update(
sess: &Arc<Session>,
mode: ThreadMemoryMode,
) -> anyhow::Result<()> {
let recorder = {
let guard = sess.services.rollout.lock().await;
guard.clone()
}
.ok_or_else(|| {
anyhow::anyhow!("Session persistence is disabled; cannot update thread memory mode.")
})?;
recorder.persist().await?;
recorder.flush().await?;
let rollout_path = recorder.rollout_path().to_path_buf();
let mut session_meta = read_session_meta_line(rollout_path.as_path()).await?;
if session_meta.meta.id != sess.conversation_id {
anyhow::bail!(
"rollout session metadata id mismatch: expected {}, found {}",
sess.conversation_id,
session_meta.meta.id
);
}
session_meta.meta.memory_mode = Some(
match mode {
ThreadMemoryMode::Enabled => "enabled",
ThreadMemoryMode::Disabled => "disabled",
}
.to_string(),
);
let item = RolloutItem::SessionMeta(session_meta);
recorder.record_items(std::slice::from_ref(&item)).await?;
recorder.flush().await?;
let live_thread = sess.live_thread_for_persistence("update thread memory mode")?;
live_thread.persist().await?;
live_thread.flush().await?;
live_thread
.update_memory_mode(mode, /*include_archived*/ false)
.await?;
live_thread.flush().await?;
Ok(())
}
@@ -996,20 +955,16 @@ pub async fn shutdown(sess: &Arc<Session>, sub_id: String) -> bool {
&[],
);
// Gracefully flush and shutdown rollout recorder on session end so tests
// that inspect the rollout file do not race with the background writer.
let recorder_opt = {
let mut guard = sess.services.rollout.lock().await;
guard.take()
};
if let Some(rec) = recorder_opt
&& let Err(e) = rec.shutdown().await
// Gracefully flush and shutdown thread persistence on session end so tests
// that inspect durable state do not race with the background writer.
if let Some(live_thread) = sess.live_thread()
&& let Err(e) = live_thread.shutdown().await
{
warn!("failed to shutdown rollout recorder: {e}");
warn!("failed to shutdown thread persistence: {e}");
let event = Event {
id: sub_id.clone(),
msg: EventMsg::Error(ErrorEvent {
message: "Failed to shutdown rollout recorder".to_string(),
message: "Failed to shutdown thread persistence".to_string(),
codex_error_info: Some(CodexErrorInfo::Other),
}),
};
codex-rs/core/src/session/mod.rs
+41 -32
View File
@@ -120,14 +120,19 @@ use codex_protocol::request_permissions::RequestPermissionsResponse;
use codex_protocol::request_user_input::RequestUserInputArgs;
use codex_protocol::request_user_input::RequestUserInputResponse;
use codex_rmcp_client::ElicitationResponse;
use codex_rollout::RolloutConfig;
use codex_rollout::state_db;
use codex_rollout_trace::RolloutTraceRecorder;
use codex_rollout_trace::ThreadStartedTraceMetadata;
use codex_sandboxing::policy_transforms::intersect_permission_profiles;
use codex_shell_command::parse_command::parse_command;
use codex_terminal_detection::user_agent;
use codex_thread_store::CreateThreadParams;
use codex_thread_store::LiveThread;
use codex_thread_store::LiveThreadInitGuard;
use codex_thread_store::LocalThreadStore;
use codex_thread_store::ResumeThreadParams;
use codex_thread_store::ThreadEventPersistenceMode;
use codex_thread_store::ThreadStore;
use codex_utils_output_truncation::TruncationPolicy;
use futures::future::BoxFuture;
use futures::future::Shared;
@@ -267,11 +272,7 @@ use crate::mcp::McpManager;
use crate::memories;
use crate::network_policy_decision::execpolicy_network_rule_amendment;
use crate::plugins::PluginsManager;
use crate::rollout::RolloutRecorder;
use crate::rollout::RolloutRecorderParams;
use crate::rollout::map_session_init_error;
use crate::rollout::metadata;
use crate::rollout::policy::EventPersistenceMode;
use crate::session_startup_prewarm::SessionStartupPrewarmHandle;
use crate::shell;
use crate::shell_snapshot::ShellSnapshot;
@@ -406,6 +407,7 @@ pub(crate) struct CodexSpawnArgs {
pub(crate) user_shell_override: Option<shell::Shell>,
pub(crate) parent_trace: Option<W3cTraceContext>,
pub(crate) analytics_events_client: Option<AnalyticsEventsClient>,
pub(crate) thread_store: Arc<dyn ThreadStore>,
}
pub(crate) const INITIAL_SUBMIT_ID: &str = "";
@@ -461,6 +463,7 @@ impl Codex {
inherited_rollout_trace,
parent_trace: _,
analytics_events_client,
thread_store,
} = args;
let (tx_sub, rx_sub) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY);
let (tx_event, rx_event) = async_channel::unbounded();
@@ -663,6 +666,7 @@ impl Codex {
agent_control,
environment_manager,
analytics_events_client,
thread_store,
inherited_rollout_trace,
)
.await
@@ -1023,33 +1027,37 @@ impl Session {
self.services.state_db.clone()
}
pub(crate) fn live_thread_for_persistence(
&self,
operation: &str,
) -> anyhow::Result<&LiveThread> {
self.live_thread()
.ok_or_else(|| anyhow::anyhow!("Session persistence is disabled; cannot {operation}."))
}
pub(crate) fn live_thread(&self) -> Option<&LiveThread> {
self.services.live_thread.as_ref()
}
/// Flush rollout writes and return the final durability-barrier result.
pub(crate) async fn flush_rollout(&self) -> std::io::Result<()> {
let recorder = {
let guard = self.services.rollout.lock().await;
guard.clone()
};
if let Some(recorder) = recorder {
recorder.flush().await
if let Some(live_thread) = self.live_thread() {
live_thread.flush().await.map_err(std::io::Error::other)
} else {
Ok(())
}
}
pub(crate) async fn try_ensure_rollout_materialized(&self) -> std::io::Result<()> {
let recorder = {
let guard = self.services.rollout.lock().await;
guard.clone()
};
if let Some(rec) = recorder {
rec.persist().await?;
if let Some(live_thread) = self.live_thread() {
live_thread.persist().await.map_err(std::io::Error::other)?;
}
Ok(())
}
pub(crate) async fn ensure_rollout_materialized(&self) {
if let Err(e) = self.try_ensure_rollout_materialized().await {
warn!("failed to materialize rollout recorder: {e}");
warn!("failed to materialize thread persistence: {e}");
}
}
@@ -1211,7 +1219,7 @@ impl Session {
state.set_token_info(Some(info));
}
// If persisting, persist all rollout items as-is (recorder filters)
// If persisting, persist all rollout items as-is (the store filters).
if !rollout_items.is_empty() {
self.persist_rollout_items(&rollout_items).await;
}
@@ -1554,7 +1562,7 @@ impl Session {
}
pub(crate) async fn send_event_raw(&self, event: Event) {
// Persist the event into rollout (recorder filters as needed)
// Persist the event into rollout storage (the store filters as needed).
let rollout_items = vec![RolloutItem::EventMsg(event.msg.clone())];
self.persist_rollout_items(&rollout_items).await;
self.deliver_event_raw(event).await;
@@ -2666,12 +2674,8 @@ impl Session {
}
pub(crate) async fn persist_rollout_items(&self, items: &[RolloutItem]) {
let recorder = {
let guard = self.services.rollout.lock().await;
guard.clone()
};
if let Some(rec) = recorder
&& let Err(e) = rec.record_items(items).await
if let Some(live_thread) = self.live_thread()
&& let Err(e) = live_thread.append_items(items).await
{
error!("failed to record rollout items: {e:#}");
}
@@ -3190,17 +3194,22 @@ impl Session {
Arc::clone(&self.services.user_shell)
}
pub(crate) async fn current_rollout_path(&self) -> Option<PathBuf> {
let recorder = {
let guard = self.services.rollout.lock().await;
guard.clone()
pub(crate) async fn current_rollout_path(&self) -> anyhow::Result<Option<PathBuf>> {
let Some(live_thread) = self.live_thread() else {
return Ok(None);
};
recorder.map(|recorder| recorder.rollout_path().to_path_buf())
live_thread.local_rollout_path().await.map_err(Into::into)
}
pub(crate) async fn hook_transcript_path(&self) -> Option<PathBuf> {
self.ensure_rollout_materialized().await;
self.current_rollout_path().await
match self.current_rollout_path().await {
Ok(path) => path,
Err(err) => {
warn!("{err}");
None
}
}
}
pub(crate) async fn take_pending_session_start_source(
codex-rs/core/src/session/session.rs
+595 -565
View File
File diff suppressed because it is too large Load Diff
codex-rs/core/src/session/tests.rs
+63 -107
View File
@@ -45,8 +45,6 @@ use codex_protocol::request_permissions::PermissionGrantScope;
use codex_protocol::request_permissions::RequestPermissionProfile;
use tracing::Span;
use crate::RolloutRecorderParams;
use crate::rollout::policy::EventPersistenceMode;
use crate::rollout::recorder::RolloutRecorder;
use crate::state::TaskKind;
use crate::tasks::SessionTask;
@@ -1705,8 +1703,11 @@ async fn record_initial_history_forked_hydrates_previous_turn_settings() {
#[tokio::test]
async fn thread_rollback_drops_last_turn_from_history() {
let (sess, tc, rx) = make_session_and_context_with_rx().await;
let rollout_path = attach_rollout_recorder(&sess).await;
let (mut sess, tc, rx) = make_session_and_context_with_rx().await;
let rollout_path = attach_thread_persistence(
Arc::get_mut(&mut sess).expect("session should not have additional references"),
)
.await;
let initial_context = sess.build_initial_context(tc.as_ref()).await;
let turn_1 = vec![
@@ -1769,8 +1770,11 @@ async fn thread_rollback_drops_last_turn_from_history() {
#[tokio::test]
async fn thread_rollback_clears_history_when_num_turns_exceeds_existing_turns() {
let (sess, tc, rx) = make_session_and_context_with_rx().await;
attach_rollout_recorder(&sess).await;
let (mut sess, tc, rx) = make_session_and_context_with_rx().await;
attach_thread_persistence(
Arc::get_mut(&mut sess).expect("session should not have additional references"),
)
.await;
let initial_context = sess.build_initial_context(tc.as_ref()).await;
let turn_1 = vec![user_message("turn 1 user")];
@@ -1795,7 +1799,7 @@ async fn thread_rollback_clears_history_when_num_turns_exceeds_existing_turns()
}
#[tokio::test]
async fn thread_rollback_fails_without_persisted_rollout_path() {
async fn thread_rollback_fails_without_persisted_thread_history() {
let (sess, tc, rx) = make_session_and_context_with_rx().await;
let initial_context = sess.build_initial_context(tc.as_ref()).await;
@@ -1807,7 +1811,7 @@ async fn thread_rollback_fails_without_persisted_rollout_path() {
let error_event = wait_for_thread_rollback_failed(&rx).await;
assert_eq!(
error_event.message,
"thread rollback requires a persisted rollout path"
"thread rollback requires persisted thread history"
);
assert_eq!(
error_event.codex_error_info,
@@ -1818,8 +1822,11 @@ async fn thread_rollback_fails_without_persisted_rollout_path() {
#[tokio::test]
async fn thread_rollback_recomputes_previous_turn_settings_and_reference_context_from_replay() {
let (sess, tc, rx) = make_session_and_context_with_rx().await;
attach_rollout_recorder(&sess).await;
let (mut sess, tc, rx) = make_session_and_context_with_rx().await;
attach_thread_persistence(
Arc::get_mut(&mut sess).expect("session should not have additional references"),
)
.await;
let first_context_item = tc.to_turn_context_item();
let first_turn_id = first_context_item
@@ -1929,8 +1936,11 @@ async fn thread_rollback_recomputes_previous_turn_settings_and_reference_context
#[tokio::test]
async fn thread_rollback_restores_cleared_reference_context_item_after_compaction() {
let (sess, tc, rx) = make_session_and_context_with_rx().await;
attach_rollout_recorder(&sess).await;
let (mut sess, tc, rx) = make_session_and_context_with_rx().await;
attach_thread_persistence(
Arc::get_mut(&mut sess).expect("session should not have additional references"),
)
.await;
let first_context_item = tc.to_turn_context_item();
let first_turn_id = first_context_item
@@ -2034,8 +2044,11 @@ async fn thread_rollback_restores_cleared_reference_context_item_after_compactio
#[tokio::test]
async fn thread_rollback_persists_marker_and_replays_cumulatively() {
let (sess, tc, rx) = make_session_and_context_with_rx().await;
let rollout_path = attach_rollout_recorder(&sess).await;
let (mut sess, tc, rx) = make_session_and_context_with_rx().await;
let rollout_path = attach_thread_persistence(
Arc::get_mut(&mut sess).expect("session should not have additional references"),
)
.await;
let turn_context_item = tc.to_turn_context_item();
sess.persist_rollout_items(&[
@@ -2580,34 +2593,31 @@ async fn wait_for_thread_rollback_failed(rx: &async_channel::Receiver<Event>) ->
}
}
async fn attach_rollout_recorder(session: &Arc<Session>) -> PathBuf {
let config = session.get_config().await;
let recorder = RolloutRecorder::new(
config.as_ref(),
RolloutRecorderParams::new(
session.conversation_id,
/*forked_from_id*/ None,
SessionSource::Exec,
BaseInstructions::default(),
Vec::new(),
EventPersistenceMode::Limited,
),
/*state_db_ctx*/ None,
/*state_builder*/ None,
async fn attach_thread_persistence(session: &mut Session) -> PathBuf {
let live_thread = LiveThread::create(
Arc::clone(&session.services.thread_store),
CreateThreadParams {
thread_id: session.conversation_id,
forked_from_id: None,
source: SessionSource::Exec,
base_instructions: BaseInstructions::default(),
dynamic_tools: Vec::new(),
event_persistence_mode: ThreadEventPersistenceMode::Limited,
},
)
.await
.expect("create rollout recorder");
let rollout_path = recorder.rollout_path().to_path_buf();
{
let mut rollout = session.services.rollout.lock().await;
*rollout = Some(recorder);
}
.expect("create thread persistence");
session.services.live_thread = Some(live_thread);
session.ensure_rollout_materialized().await;
session
.flush_rollout()
.await
.expect("attached rollout should flush");
rollout_path
session
.current_rollout_path()
.await
.expect("load rollout path")
.expect("thread should have rollout path")
}
fn text_block(s: &str) -> serde_json::Value {
@@ -3137,6 +3147,9 @@ async fn session_new_fails_when_zsh_fork_enabled_without_zsh_path() {
AgentControl::default(),
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
/*analytics_events_client*/ None,
Arc::new(codex_thread_store::LocalThreadStore::new(
codex_rollout::RolloutConfig::from_view(config.as_ref()),
)),
RolloutTraceRecorder::disabled(),
)
.await;
@@ -3259,7 +3272,6 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
legacy_notify_argv: config.notify.clone(),
..HooksConfig::default()
}),
rollout: Mutex::new(None),
rollout_trace: RolloutTraceRecorder::disabled(),
user_shell: Arc::new(default_user_shell()),
shell_snapshot_tx: watch::channel(None).0,
@@ -3280,9 +3292,10 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
network_proxy: None,
network_approval: Arc::clone(&network_approval),
state_db: None,
thread_store: codex_thread_store::LocalThreadStore::new(
live_thread: None,
thread_store: Arc::new(codex_thread_store::LocalThreadStore::new(
codex_rollout::RolloutConfig::from_view(config.as_ref()),
),
)),
model_client: ModelClient::new(
Some(auth_manager.clone()),
conversation_id,
@@ -3456,6 +3469,9 @@ async fn make_session_with_config_and_rx(
AgentControl::default(),
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
/*analytics_events_client*/ None,
Arc::new(codex_thread_store::LocalThreadStore::new(
codex_rollout::RolloutConfig::from_view(config.as_ref()),
)),
RolloutTraceRecorder::disabled(),
)
.await?;
@@ -4571,7 +4587,6 @@ pub(crate) async fn make_session_and_context_with_dynamic_tools_and_rx(
legacy_notify_argv: config.notify.clone(),
..HooksConfig::default()
}),
rollout: Mutex::new(None),
rollout_trace: RolloutTraceRecorder::disabled(),
user_shell: Arc::new(default_user_shell()),
shell_snapshot_tx: watch::channel(None).0,
@@ -4592,9 +4607,10 @@ pub(crate) async fn make_session_and_context_with_dynamic_tools_and_rx(
network_proxy: None,
network_approval: Arc::clone(&network_approval),
state_db: None,
thread_store: codex_thread_store::LocalThreadStore::new(
live_thread: None,
thread_store: Arc::new(codex_thread_store::LocalThreadStore::new(
codex_rollout::RolloutConfig::from_view(config.as_ref()),
),
)),
model_client: ModelClient::new(
Some(Arc::clone(&auth_manager)),
conversation_id,
@@ -5488,7 +5504,7 @@ async fn record_context_updates_and_set_reference_context_item_reinjects_full_co
#[tokio::test]
async fn record_context_updates_and_set_reference_context_item_persists_baseline_without_emitting_diffs()
{
let (session, previous_context) = make_session_and_context().await;
let (mut session, previous_context) = make_session_and_context().await;
let next_model = if previous_context.model_info.slug == "gpt-5.4" {
"gpt-5.2"
} else {
@@ -5502,27 +5518,7 @@ async fn record_context_updates_and_set_reference_context_item_persists_baseline
let mut state = session.state.lock().await;
state.set_reference_context_item(Some(previous_context_item.clone()));
}
let config = session.get_config().await;
let recorder = RolloutRecorder::new(
config.as_ref(),
RolloutRecorderParams::new(
ThreadId::default(),
/*forked_from_id*/ None,
SessionSource::Exec,
BaseInstructions::default(),
Vec::new(),
EventPersistenceMode::Limited,
),
/*state_db_ctx*/ None,
/*state_builder*/ None,
)
.await
.expect("create rollout recorder");
let rollout_path = recorder.rollout_path().to_path_buf();
{
let mut rollout = session.services.rollout.lock().await;
*rollout = Some(recorder);
}
let rollout_path = attach_thread_persistence(&mut session).await;
let update_items = session
.build_settings_update_items(Some(&previous_context_item), &turn_context)
@@ -5567,30 +5563,10 @@ async fn record_context_updates_and_set_reference_context_item_persists_baseline
#[tokio::test]
async fn record_context_updates_and_set_reference_context_item_persists_split_file_system_policy_to_rollout()
{
let (session, mut turn_context) = make_session_and_context().await;
let (mut session, mut turn_context) = make_session_and_context().await;
let file_system_sandbox_policy = file_system_policy_with_unreadable_glob(&turn_context);
turn_context.file_system_sandbox_policy = file_system_sandbox_policy.clone();
let config = session.get_config().await;
let recorder = RolloutRecorder::new(
config.as_ref(),
RolloutRecorderParams::new(
ThreadId::default(),
/*forked_from_id*/ None,
SessionSource::Exec,
BaseInstructions::default(),
Vec::new(),
EventPersistenceMode::Limited,
),
/*state_db_ctx*/ None,
/*state_builder*/ None,
)
.await
.expect("create rollout recorder");
let rollout_path = recorder.rollout_path().to_path_buf();
{
let mut rollout = session.services.rollout.lock().await;
*rollout = Some(recorder);
}
let rollout_path = attach_thread_persistence(&mut session).await;
session
.record_context_updates_and_set_reference_context_item(&turn_context)
@@ -5640,7 +5616,7 @@ async fn build_initial_context_prepends_model_switch_message() {
#[tokio::test]
async fn record_context_updates_and_set_reference_context_item_persists_full_reinjection_to_rollout()
{
let (session, previous_context) = make_session_and_context().await;
let (mut session, previous_context) = make_session_and_context().await;
let next_model = if previous_context.model_info.slug == "gpt-5.4" {
"gpt-5.2"
} else {
@@ -5649,27 +5625,7 @@ async fn record_context_updates_and_set_reference_context_item_persists_full_rei
let turn_context = previous_context
.with_model(next_model.to_string(), &session.services.models_manager)
.await;
let config = session.get_config().await;
let recorder = RolloutRecorder::new(
config.as_ref(),
RolloutRecorderParams::new(
ThreadId::default(),
/*forked_from_id*/ None,
SessionSource::Exec,
BaseInstructions::default(),
Vec::new(),
EventPersistenceMode::Limited,
),
/*state_db_ctx*/ None,
/*state_builder*/ None,
)
.await
.expect("create rollout recorder");
let rollout_path = recorder.rollout_path().to_path_buf();
{
let mut rollout = session.services.rollout.lock().await;
*rollout = Some(recorder);
}
let rollout_path = attach_thread_persistence(&mut session).await;
session
.persist_rollout_items(&[RolloutItem::EventMsg(EventMsg::UserMessage(
codex-rs/core/src/session/tests/guardian_tests.rs
+4
View File
@@ -744,6 +744,9 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() {
));
let mcp_manager = Arc::new(McpManager::new(Arc::clone(&plugins_manager)));
let skills_watcher = Arc::new(SkillsWatcher::noop());
let thread_store = Arc::new(codex_thread_store::LocalThreadStore::new(
codex_rollout::RolloutConfig::from_view(&config),
));
let CodexSpawnOk { codex, .. } = Codex::spawn(CodexSpawnArgs {
config,
@@ -768,6 +771,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() {
user_shell_override: None,
parent_trace: None,
analytics_events_client: None,
thread_store,
})
.await
.expect("spawn guardian subagent");
codex-rs/core/src/state/service.rs
+4 -4
View File
@@ -1,7 +1,6 @@
use std::collections::HashMap;
use std::sync::Arc;
use crate::RolloutRecorder;
use crate::SkillsManager;
use crate::agent::AgentControl;
use crate::client::ModelClient;
@@ -25,7 +24,8 @@ use codex_models_manager::manager::ModelsManager;
use codex_otel::SessionTelemetry;
use codex_rollout::state_db::StateDbHandle;
use codex_rollout_trace::RolloutTraceRecorder;
use codex_thread_store::LocalThreadStore;
use codex_thread_store::LiveThread;
use codex_thread_store::ThreadStore;
use std::path::PathBuf;
use tokio::runtime::Handle;
use tokio::sync::Mutex;
@@ -43,7 +43,6 @@ pub(crate) struct SessionServices {
pub(crate) main_execve_wrapper_exe: Option<PathBuf>,
pub(crate) analytics_events_client: AnalyticsEventsClient,
pub(crate) hooks: Hooks,
pub(crate) rollout: Mutex<Option<RolloutRecorder>>,
pub(crate) rollout_trace: RolloutTraceRecorder,
pub(crate) user_shell: Arc<crate::shell::Shell>,
pub(crate) shell_snapshot_tx: watch::Sender<Option<Arc<crate::shell_snapshot::ShellSnapshot>>>,
@@ -64,7 +63,8 @@ pub(crate) struct SessionServices {
pub(crate) network_proxy: Option<StartedNetworkProxy>,
pub(crate) network_approval: Arc<NetworkApprovalService>,
pub(crate) state_db: Option<StateDbHandle>,
pub(crate) thread_store: LocalThreadStore,
pub(crate) live_thread: Option<LiveThread>,
pub(crate) thread_store: Arc<dyn ThreadStore>,
/// Session-scoped model client shared across turns.
pub(crate) model_client: ModelClient,
pub(crate) code_mode_service: CodeModeService,
codex-rs/core/src/thread_manager.rs
+31
View File
@@ -44,7 +44,11 @@ use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::TurnAbortReason;
use codex_protocol::protocol::TurnAbortedEvent;
use codex_protocol::protocol::W3cTraceContext;
use codex_rollout::RolloutConfig;
use codex_state::DirectionalThreadSpawnEdgeStatus;
use codex_thread_store::LocalThreadStore;
use codex_thread_store::RemoteThreadStore;
use codex_thread_store::ThreadStore;
use codex_utils_absolute_path::AbsolutePathBuf;
use futures::StreamExt;
use futures::stream::FuturesUnordered;
@@ -237,6 +241,13 @@ pub fn build_models_manager(
))
}
fn configured_thread_store(config: &Config) -> Arc<dyn ThreadStore> {
if let Some(endpoint) = config.experimental_thread_store_endpoint.clone() {
return Arc::new(RemoteThreadStore::new(endpoint));
}
Arc::new(LocalThreadStore::new(RolloutConfig::from_view(config)))
}
impl ThreadManager {
pub fn new(
config: &Config,
@@ -514,8 +525,10 @@ impl ThreadManager {
metrics_service_name: Option<String>,
parent_trace: Option<W3cTraceContext>,
) -> CodexResult<NewThread> {
let thread_store = configured_thread_store(&config);
Box::pin(self.state.spawn_thread(
config,
thread_store,
initial_history,
Arc::clone(&self.state.auth_manager),
self.agent_control(),
@@ -554,8 +567,10 @@ impl ThreadManager {
persist_extended_history: bool,
parent_trace: Option<W3cTraceContext>,
) -> CodexResult<NewThread> {
let thread_store = configured_thread_store(&config);
Box::pin(self.state.spawn_thread(
config,
thread_store,
initial_history,
auth_manager,
self.agent_control(),
@@ -573,8 +588,10 @@ impl ThreadManager {
config: Config,
user_shell_override: crate::shell::Shell,
) -> CodexResult<NewThread> {
let thread_store = configured_thread_store(&config);
Box::pin(self.state.spawn_thread(
config,
thread_store,
InitialHistory::New,
Arc::clone(&self.state.auth_manager),
self.agent_control(),
@@ -595,8 +612,10 @@ impl ThreadManager {
user_shell_override: crate::shell::Shell,
) -> CodexResult<NewThread> {
let initial_history = RolloutRecorder::get_rollout_history(&rollout_path).await?;
let thread_store = configured_thread_store(&config);
Box::pin(self.state.spawn_thread(
config,
thread_store,
initial_history,
auth_manager,
self.agent_control(),
@@ -703,8 +722,10 @@ impl ThreadManager {
}
}
};
let thread_store = configured_thread_store(&config);
Box::pin(self.state.spawn_thread(
config,
thread_store,
history,
Arc::clone(&self.state.auth_manager),
self.agent_control(),
@@ -801,8 +822,10 @@ impl ThreadManagerState {
inherited_shell_snapshot: Option<Arc<ShellSnapshot>>,
inherited_exec_policy: Option<Arc<crate::exec_policy::ExecPolicyManager>>,
) -> CodexResult<NewThread> {
let thread_store = configured_thread_store(&config);
Box::pin(self.spawn_thread_with_source(
config,
thread_store,
InitialHistory::New,
Arc::clone(&self.auth_manager),
agent_control,
@@ -828,8 +851,10 @@ impl ThreadManagerState {
inherited_exec_policy: Option<Arc<crate::exec_policy::ExecPolicyManager>>,
) -> CodexResult<NewThread> {
let initial_history = RolloutRecorder::get_rollout_history(&rollout_path).await?;
let thread_store = configured_thread_store(&config);
Box::pin(self.spawn_thread_with_source(
config,
thread_store,
initial_history,
Arc::clone(&self.auth_manager),
agent_control,
@@ -856,8 +881,10 @@ impl ThreadManagerState {
inherited_shell_snapshot: Option<Arc<ShellSnapshot>>,
inherited_exec_policy: Option<Arc<crate::exec_policy::ExecPolicyManager>>,
) -> CodexResult<NewThread> {
let thread_store = configured_thread_store(&config);
Box::pin(self.spawn_thread_with_source(
config,
thread_store,
initial_history,
Arc::clone(&self.auth_manager),
agent_control,
@@ -878,6 +905,7 @@ impl ThreadManagerState {
pub(crate) async fn spawn_thread(
&self,
config: Config,
thread_store: Arc<dyn ThreadStore>,
initial_history: InitialHistory,
auth_manager: Arc<AuthManager>,
agent_control: AgentControl,
@@ -889,6 +917,7 @@ impl ThreadManagerState {
) -> CodexResult<NewThread> {
Box::pin(self.spawn_thread_with_source(
config,
thread_store,
initial_history,
auth_manager,
agent_control,
@@ -908,6 +937,7 @@ impl ThreadManagerState {
pub(crate) async fn spawn_thread_with_source(
&self,
config: Config,
thread_store: Arc<dyn ThreadStore>,
initial_history: InitialHistory,
auth_manager: Arc<AuthManager>,
agent_control: AgentControl,
@@ -957,6 +987,7 @@ impl ThreadManagerState {
user_shell_override,
parent_trace,
analytics_events_client: self.analytics_events_client.clone(),
thread_store,
})
.await?;
self.finalize_thread_spawn(codex, thread_id, watch_registration)
codex-rs/thread-store/Cargo.toml
+2
View File
@@ -26,8 +26,10 @@ prost = "0.14.3"
serde = { workspace = true, features = ["derive"] }
serde_json = { workspace = true }
thiserror = { workspace = true }
tokio = { workspace = true }
tonic = { workspace = true }
tonic-prost = { workspace = true }
tracing = { workspace = true }
[dev-dependencies]
pretty_assertions = { workspace = true }
codex-rs/thread-store/src/lib.rs
+4 -3
View File
@@ -5,16 +5,17 @@
//! any other backing store.
mod error;
mod live_thread;
mod local;
mod recorder;
mod remote;
mod store;
mod types;
pub use error::ThreadStoreError;
pub use error::ThreadStoreResult;
pub use live_thread::LiveThread;
pub use live_thread::LiveThreadInitGuard;
pub use local::LocalThreadStore;
pub use recorder::ThreadRecorder;
pub use remote::RemoteThreadStore;
pub use store::ThreadStore;
pub use types::AppendThreadItemsParams;
@@ -25,7 +26,7 @@ pub use types::ListThreadsParams;
pub use types::LoadThreadHistoryParams;
pub use types::OptionalStringPatch;
pub use types::ReadThreadParams;
pub use types::ResumeThreadRecorderParams;
pub use types::ResumeThreadParams;
pub use types::SortDirection;
pub use types::StoredThread;
pub use types::StoredThreadHistory;
codex-rs/thread-store/src/live_thread.rs
+176
View File
@@ -0,0 +1,176 @@
use std::path::PathBuf;
use std::sync::Arc;
use codex_protocol::ThreadId;
use codex_protocol::protocol::RolloutItem;
use codex_protocol::protocol::ThreadMemoryMode;
use tracing::warn;
use crate::AppendThreadItemsParams;
use crate::CreateThreadParams;
use crate::LoadThreadHistoryParams;
use crate::LocalThreadStore;
use crate::ResumeThreadParams;
use crate::StoredThreadHistory;
use crate::ThreadMetadataPatch;
use crate::ThreadStore;
use crate::ThreadStoreResult;
use crate::UpdateThreadMetadataParams;
/// Handle for an active thread's persistence lifecycle.
///
/// `LiveThread` keeps lifecycle decisions with the caller while delegating storage details to
/// [`ThreadStore`]. Local stores may use a rollout file internally and remote stores may use a
/// service, but session code should only need this handle for the active thread.
#[derive(Clone)]
pub struct LiveThread {
thread_id: ThreadId,
thread_store: Arc<dyn ThreadStore>,
}
/// Owns a live thread while session initialization is still fallible.
///
/// If initialization returns early after persistence has been opened, dropping this guard discards
/// the live writer without forcing lazy in-memory state to become durable. Call [`commit`] once the
/// session owns the live thread for normal operation.
pub struct LiveThreadInitGuard {
live_thread: Option<LiveThread>,
}
impl LiveThreadInitGuard {
pub fn new(live_thread: Option<LiveThread>) -> Self {
Self { live_thread }
}
pub fn as_ref(&self) -> Option<&LiveThread> {
self.live_thread.as_ref()
}
pub fn commit(&mut self) {
self.live_thread = None;
}
pub async fn discard(&mut self) {
let Some(live_thread) = self.live_thread.take() else {
return;
};
if let Err(err) = live_thread.discard().await {
warn!("failed to discard thread persistence for failed session init: {err}");
}
}
}
impl Drop for LiveThreadInitGuard {
fn drop(&mut self) {
let Some(live_thread) = self.live_thread.take() else {
return;
};
let Ok(handle) = tokio::runtime::Handle::try_current() else {
warn!("failed to discard thread persistence for failed session init: no Tokio runtime");
return;
};
handle.spawn(async move {
if let Err(err) = live_thread.discard().await {
warn!("failed to discard thread persistence for failed session init: {err}");
}
});
}
}
impl LiveThread {
pub async fn create(
thread_store: Arc<dyn ThreadStore>,
params: CreateThreadParams,
) -> ThreadStoreResult<Self> {
let thread_id = params.thread_id;
thread_store.create_thread(params).await?;
Ok(Self {
thread_id,
thread_store,
})
}
pub async fn resume(
thread_store: Arc<dyn ThreadStore>,
params: ResumeThreadParams,
) -> ThreadStoreResult<Self> {
let thread_id = params.thread_id;
thread_store.resume_thread(params).await?;
Ok(Self {
thread_id,
thread_store,
})
}
pub async fn append_items(&self, items: &[RolloutItem]) -> ThreadStoreResult<()> {
self.thread_store
.append_items(AppendThreadItemsParams {
thread_id: self.thread_id,
items: items.to_vec(),
})
.await
}
pub async fn persist(&self) -> ThreadStoreResult<()> {
self.thread_store.persist_thread(self.thread_id).await
}
pub async fn flush(&self) -> ThreadStoreResult<()> {
self.thread_store.flush_thread(self.thread_id).await
}
pub async fn shutdown(&self) -> ThreadStoreResult<()> {
self.thread_store.shutdown_thread(self.thread_id).await
}
pub async fn discard(&self) -> ThreadStoreResult<()> {
self.thread_store.discard_thread(self.thread_id).await
}
pub async fn load_history(
&self,
include_archived: bool,
) -> ThreadStoreResult<StoredThreadHistory> {
self.thread_store
.load_history(LoadThreadHistoryParams {
thread_id: self.thread_id,
include_archived,
})
.await
}
pub async fn update_memory_mode(
&self,
mode: ThreadMemoryMode,
include_archived: bool,
) -> ThreadStoreResult<()> {
self.thread_store
.update_thread_metadata(UpdateThreadMetadataParams {
thread_id: self.thread_id,
patch: ThreadMetadataPatch {
memory_mode: Some(mode),
..Default::default()
},
include_archived,
})
.await?;
Ok(())
}
/// Returns the live local rollout path for legacy local-only callers.
///
/// Remote stores do not expose rollout files, so they return `Ok(None)`.
pub async fn local_rollout_path(&self) -> ThreadStoreResult<Option<PathBuf>> {
let Some(local_store) = self
.thread_store
.as_any()
.downcast_ref::<LocalThreadStore>()
else {
return Ok(None);
};
local_store
.live_rollout_path(self.thread_id)
.await
.map(Some)
}
}
codex-rs/thread-store/src/local/create_thread.rs
+41
View File
@@ -0,0 +1,41 @@
use super::LocalThreadStore;
use crate::CreateThreadParams;
use crate::ThreadEventPersistenceMode;
use crate::ThreadStoreError;
use crate::ThreadStoreResult;
use codex_rollout::EventPersistenceMode;
use codex_rollout::RolloutRecorder;
use codex_rollout::RolloutRecorderParams;
pub(super) async fn create_thread(
store: &LocalThreadStore,
params: CreateThreadParams,
) -> ThreadStoreResult<RolloutRecorder> {
let state_db_ctx = store.state_db().await;
let recorder = RolloutRecorder::new(
&store.config,
RolloutRecorderParams::new(
params.thread_id,
params.forked_from_id,
params.source,
params.base_instructions,
params.dynamic_tools,
event_persistence_mode(params.event_persistence_mode),
),
state_db_ctx,
/*state_builder*/ None,
)
.await
.map_err(|err| ThreadStoreError::Internal {
message: format!("failed to initialize local thread recorder: {err}"),
})?;
Ok(recorder)
}
pub(super) fn event_persistence_mode(mode: ThreadEventPersistenceMode) -> EventPersistenceMode {
match mode {
ThreadEventPersistenceMode::Limited => EventPersistenceMode::Limited,
ThreadEventPersistenceMode::Extended => EventPersistenceMode::Extended,
}
}
codex-rs/thread-store/src/local/live_writer.rs
+152
View File
@@ -0,0 +1,152 @@
use std::path::PathBuf;
use codex_protocol::ThreadId;
use codex_rollout::RolloutRecorder;
use codex_rollout::RolloutRecorderParams;
use codex_rollout::builder_from_items;
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, history) = match (params.rollout_path, params.history) {
(Some(rollout_path), history) => (rollout_path, history),
(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?;
let rollout_path = thread
.rollout_path
.ok_or_else(|| ThreadStoreError::Internal {
message: format!("thread {} does not have a rollout path", params.thread_id),
})?;
(
rollout_path,
history.or_else(|| thread.history.map(|history| history.items)),
)
}
};
let state_builder = history
.as_deref()
.and_then(|items| builder_from_items(items, rollout_path.as_path()));
let state_db_ctx = store.state_db().await;
let recorder = RolloutRecorder::new(
&store.config,
RolloutRecorderParams::resume(
rollout_path,
create_thread::event_persistence_mode(params.event_persistence_mode),
),
state_db_ctx,
state_builder,
)
.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<()> {
store
.live_recorder(params.thread_id)
.await?
.record_items(params.items.as_slice())
.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)
}
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)
}
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)?;
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())
}
fn thread_store_io_error(err: std::io::Error) -> ThreadStoreError {
ThreadStoreError::Internal {
message: err.to_string(),
}
}
codex-rs/thread-store/src/local/mod.rs
+444 -23
View File
@@ -1,6 +1,8 @@
mod archive_thread;
mod create_thread;
mod helpers;
mod list_threads;
mod live_writer;
mod read_thread;
mod unarchive_thread;
mod update_thread_metadata;
@@ -9,7 +11,16 @@ mod update_thread_metadata;
mod test_support;
use async_trait::async_trait;
use codex_protocol::ThreadId;
use codex_rollout::RolloutConfig;
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 tokio::sync::OnceCell;
use crate::AppendThreadItemsParams;
use crate::ArchiveThreadParams;
@@ -17,32 +28,56 @@ use crate::CreateThreadParams;
use crate::ListThreadsParams;
use crate::LoadThreadHistoryParams;
use crate::ReadThreadParams;
use crate::ResumeThreadRecorderParams;
use crate::ResumeThreadParams;
use crate::StoredThread;
use crate::StoredThreadHistory;
use crate::ThreadPage;
use crate::ThreadRecorder;
use crate::ThreadStore;
use crate::ThreadStoreError;
use crate::ThreadStoreResult;
use crate::UpdateThreadMetadataParams;
/// Local filesystem/SQLite-backed implementation of [`ThreadStore`].
#[derive(Clone, Debug)]
#[derive(Clone)]
pub struct LocalThreadStore {
pub(super) config: RolloutConfig,
live_recorders: Arc<Mutex<HashMap<ThreadId, RolloutRecorder>>>,
state_db: Arc<OnceCell<StateDbHandle>>,
}
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 from the rollout configuration used by existing local persistence.
pub fn new(config: RolloutConfig) -> Self {
Self { config }
Self {
config,
live_recorders: Arc::new(Mutex::new(HashMap::new())),
state_db: Arc::new(OnceCell::new()),
}
}
/// Return the state DB handle used by local rollout writers.
pub async fn state_db(&self) -> Option<StateDbHandle> {
self.state_db
.get_or_try_init(|| async {
codex_rollout::state_db::init(&self.config).await.ok_or(())
})
.await
.ok()
.cloned()
}
/// Read a local rollout-backed thread by path.
pub async fn read_thread_by_rollout_path(
&self,
rollout_path: std::path::PathBuf,
rollout_path: PathBuf,
include_archived: bool,
include_history: bool,
) -> ThreadStoreResult<StoredThread> {
@@ -54,6 +89,51 @@ impl LocalThreadStore {
)
.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_trait]
@@ -62,29 +142,65 @@ impl ThreadStore for LocalThreadStore {
self
}
async fn create_thread(
&self,
_params: CreateThreadParams,
) -> ThreadStoreResult<Box<dyn ThreadRecorder>> {
unsupported("create_thread")
async fn create_thread(&self, params: CreateThreadParams) -> ThreadStoreResult<()> {
live_writer::create_thread(self, params).await
}
async fn resume_thread_recorder(
&self,
_params: ResumeThreadRecorderParams,
) -> ThreadStoreResult<Box<dyn ThreadRecorder>> {
unsupported("resume_thread_recorder")
async fn resume_thread(&self, params: ResumeThreadParams) -> ThreadStoreResult<()> {
live_writer::resume_thread(self, params).await
}
async fn append_items(&self, _params: AppendThreadItemsParams) -> ThreadStoreResult<()> {
unsupported("append_items")
async fn append_items(&self, params: AppendThreadItemsParams) -> ThreadStoreResult<()> {
live_writer::append_items(self, params).await
}
async fn persist_thread(&self, thread_id: ThreadId) -> ThreadStoreResult<()> {
live_writer::persist_thread(self, thread_id).await
}
async fn flush_thread(&self, thread_id: ThreadId) -> ThreadStoreResult<()> {
live_writer::flush_thread(self, thread_id).await
}
async fn shutdown_thread(&self, thread_id: ThreadId) -> ThreadStoreResult<()> {
live_writer::shutdown_thread(self, thread_id).await
}
async fn discard_thread(&self, thread_id: ThreadId) -> ThreadStoreResult<()> {
live_writer::discard_thread(self, thread_id).await
}
async fn load_history(
&self,
_params: LoadThreadHistoryParams,
params: LoadThreadHistoryParams,
) -> ThreadStoreResult<StoredThreadHistory> {
unsupported("load_history")
if let Ok(rollout_path) = live_writer::rollout_path(self, params.thread_id).await {
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(&self, params: ReadThreadParams) -> ThreadStoreResult<StoredThread> {
@@ -114,8 +230,313 @@ impl ThreadStore for LocalThreadStore {
}
}
fn unsupported<T>(operation: &str) -> ThreadStoreResult<T> {
Err(ThreadStoreError::Internal {
message: format!("local thread store does not implement {operation} in this slice"),
})
#[cfg(test)]
mod tests {
use codex_protocol::ThreadId;
use codex_protocol::models::BaseInstructions;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::RolloutItem;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::UserMessageEvent;
use tempfile::TempDir;
use super::*;
use crate::ThreadEventPersistenceMode;
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()));
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 discard_thread_drops_unmaterialized_live_writer() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()));
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());
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);
resumed_store
.resume_thread(ResumeThreadParams {
thread_id,
rollout_path: None,
history: None,
include_archived: true,
event_persistence_mode: ThreadEventPersistenceMode::Limited,
})
.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()));
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 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()));
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,
event_persistence_mode: ThreadEventPersistenceMode::Limited,
})
.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 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()));
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,
event_persistence_mode: ThreadEventPersistenceMode::Limited,
})
.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 history = store
.load_history(LoadThreadHistoryParams {
thread_id,
include_archived: false,
})
.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"
)
}));
}
fn create_thread_params(thread_id: ThreadId) -> CreateThreadParams {
CreateThreadParams {
thread_id,
forked_from_id: None,
source: SessionSource::Exec,
base_instructions: BaseInstructions::default(),
dynamic_tools: Vec::new(),
event_persistence_mode: ThreadEventPersistenceMode::Limited,
}
}
fn user_message_item(message: &str) -> RolloutItem {
RolloutItem::EventMsg(EventMsg::UserMessage(UserMessageEvent {
message: message.to_string(),
images: None,
local_images: Vec::new(),
text_elements: Vec::new(),
}))
}
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
)
}));
}
}
codex-rs/thread-store/src/local/update_thread_metadata.rs
+144 -13
View File
@@ -5,7 +5,7 @@ use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::RolloutItem;
use codex_protocol::protocol::ThreadMemoryMode;
use codex_protocol::protocol::ThreadNameUpdatedEvent;
use codex_rollout::StateDbHandle;
use codex_rollout::ARCHIVED_SESSIONS_SUBDIR;
use codex_rollout::append_rollout_item_to_path;
use codex_rollout::append_thread_name;
use codex_rollout::find_archived_thread_path_by_id_str;
@@ -13,6 +13,7 @@ use codex_rollout::find_thread_path_by_id_str;
use codex_rollout::read_session_meta_line;
use super::LocalThreadStore;
use super::live_writer;
use crate::ReadThreadParams;
use crate::StoredThread;
use crate::ThreadStoreError;
@@ -53,7 +54,7 @@ pub(super) async fn update_thread_metadata(
.await?;
}
let state_db_ctx = open_state_db_for_direct_thread_lookup(store).await;
let state_db_ctx = store.state_db().await;
codex_rollout::state_db::reconcile_rollout(
state_db_ctx.as_deref(),
resolved_rollout_path.path.as_path(),
@@ -65,7 +66,7 @@ pub(super) async fn update_thread_metadata(
)
.await;
read_thread::read_thread(
match read_thread::read_thread(
store,
ReadThreadParams {
thread_id,
@@ -74,6 +75,18 @@ pub(super) async fn update_thread_metadata(
},
)
.await
{
Ok(thread) => Ok(thread),
Err(_) => {
read_thread::read_thread_by_rollout_path(
store,
resolved_rollout_path.path,
params.include_archived,
/*include_history*/ false,
)
.await
}
}
}
async fn apply_thread_name(
@@ -127,15 +140,6 @@ async fn apply_thread_memory_mode(
})
}
async fn open_state_db_for_direct_thread_lookup(store: &LocalThreadStore) -> Option<StateDbHandle> {
codex_state::StateRuntime::init(
store.config.sqlite_home.clone(),
store.config.model_provider_id.clone(),
)
.await
.ok()
}
fn memory_mode_as_str(mode: ThreadMemoryMode) -> &'static str {
match mode {
ThreadMemoryMode::Enabled => "enabled",
@@ -148,6 +152,11 @@ async fn resolve_rollout_path(
thread_id: ThreadId,
include_archived: bool,
) -> ThreadStoreResult<ResolvedRolloutPath> {
if let Ok(path) = live_writer::rollout_path(store, thread_id).await {
let archived = rollout_path_is_archived(store, path.as_path());
return Ok(ResolvedRolloutPath { path, archived });
}
let active_path =
find_thread_path_by_id_str(store.config.codex_home.as_path(), &thread_id.to_string())
.await
@@ -179,6 +188,10 @@ async fn resolve_rollout_path(
})
}
fn rollout_path_is_archived(store: &LocalThreadStore, path: &std::path::Path) -> bool {
path.starts_with(store.config.codex_home.join(ARCHIVED_SESSIONS_SUBDIR))
}
#[cfg(test)]
mod tests {
use pretty_assertions::assert_eq;
@@ -187,6 +200,8 @@ mod tests {
use uuid::Uuid;
use super::*;
use crate::ResumeThreadParams;
use crate::ThreadEventPersistenceMode;
use crate::ThreadMetadataPatch;
use crate::ThreadStore;
use crate::local::LocalThreadStore;
@@ -231,11 +246,18 @@ mod tests {
#[tokio::test]
async fn update_thread_metadata_sets_memory_mode_on_active_rollout() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()));
let config = test_config(home.path());
let store = LocalThreadStore::new(config.clone());
let uuid = Uuid::from_u128(302);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let path =
write_session_file(home.path(), "2025-01-03T14-30-00", uuid).expect("session file");
let runtime = codex_state::StateRuntime::init(
home.path().to_path_buf(),
config.model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let thread = store
.update_thread_metadata(UpdateThreadMetadataParams {
@@ -254,6 +276,51 @@ mod tests {
assert_eq!(appended["type"], "session_meta");
assert_eq!(appended["payload"]["id"], thread_id.to_string());
assert_eq!(appended["payload"]["memory_mode"], "disabled");
let memory_mode = runtime
.get_thread_memory_mode(thread_id)
.await
.expect("thread memory mode should be readable");
assert_eq!(memory_mode.as_deref(), Some("disabled"));
}
#[tokio::test]
async fn update_thread_metadata_uses_live_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()));
let uuid = Uuid::from_u128(307);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let path = write_session_file(external_home.path(), "2025-01-03T14-45-00", uuid)
.expect("external session file");
store
.resume_thread(ResumeThreadParams {
thread_id,
rollout_path: Some(path.clone()),
history: None,
include_archived: true,
event_persistence_mode: ThreadEventPersistenceMode::Limited,
})
.await
.expect("resume external live thread");
let thread = store
.update_thread_metadata(UpdateThreadMetadataParams {
thread_id,
patch: ThreadMetadataPatch {
memory_mode: Some(ThreadMemoryMode::Disabled),
..Default::default()
},
include_archived: false,
})
.await
.expect("set memory mode on external live thread");
assert_eq!(thread.thread_id, thread_id);
assert!(thread.rollout_path.is_some());
let appended = last_rollout_item(path.as_path());
assert_eq!(appended["type"], "session_meta");
assert_eq!(appended["payload"]["memory_mode"], "disabled");
}
#[tokio::test]
@@ -385,6 +452,70 @@ mod tests {
);
}
#[tokio::test]
async fn update_thread_metadata_keeps_live_archived_thread_archived_in_sqlite() {
let home = TempDir::new().expect("temp dir");
let config = test_config(home.path());
let store = LocalThreadStore::new(config.clone());
let uuid = Uuid::from_u128(308);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let archived_path = write_archived_session_file(home.path(), "2025-01-03T16-30-00", uuid)
.expect("archived session file");
let runtime = codex_state::StateRuntime::init(
home.path().to_path_buf(),
config.model_provider_id.clone(),
)
.await
.expect("state db should initialize");
runtime
.mark_backfill_complete(/*last_watermark*/ None)
.await
.expect("backfill should be complete");
codex_rollout::state_db::reconcile_rollout(
Some(runtime.as_ref()),
archived_path.as_path(),
config.model_provider_id.as_str(),
/*builder*/ None,
&[],
/*archived_only*/ Some(true),
/*new_thread_memory_mode*/ None,
)
.await;
store
.resume_thread(ResumeThreadParams {
thread_id,
rollout_path: Some(archived_path.clone()),
history: None,
include_archived: true,
event_persistence_mode: ThreadEventPersistenceMode::Limited,
})
.await
.expect("resume archived live thread");
let thread = store
.update_thread_metadata(UpdateThreadMetadataParams {
thread_id,
patch: ThreadMetadataPatch {
name: Some("Live archived title".to_string()),
..Default::default()
},
include_archived: true,
})
.await
.expect("set archived thread name");
assert!(thread.archived_at.is_some());
assert!(
runtime
.get_thread(thread_id)
.await
.expect("get metadata")
.expect("metadata")
.archived_at
.is_some()
);
}
fn last_rollout_item(path: &std::path::Path) -> Value {
let last_line = std::fs::read_to_string(path)
.expect("read rollout")
codex-rs/thread-store/src/recorder.rs
-28
View File
@@ -1,28 +0,0 @@
use async_trait::async_trait;
use codex_protocol::ThreadId;
use codex_protocol::protocol::RolloutItem;
use crate::ThreadStoreResult;
/// Live append handle for a thread.
///
/// This is the storage-neutral version of the existing rollout recorder API. The local
/// implementation is expected to wrap `codex_rollout::RolloutRecorder` and preserve its lazy
/// materialization, filtering, flush, and shutdown behavior.
#[async_trait]
pub trait ThreadRecorder: Send + Sync {
/// Returns the thread id this recorder appends to.
fn thread_id(&self) -> ThreadId;
/// Queues items for persistence according to this recorder's filtering policy.
async fn record_items(&self, items: &[RolloutItem]) -> ThreadStoreResult<()>;
/// Materializes the thread if persistence is lazy, then persists all queued items.
async fn persist(&self) -> ThreadStoreResult<()>;
/// Flushes all queued items and returns once they are durable/readable.
async fn flush(&self) -> ThreadStoreResult<()>;
/// Flushes pending items and closes the recorder.
async fn shutdown(&self) -> ThreadStoreResult<()>;
}
codex-rs/thread-store/src/remote/mod.rs
+21 -11
View File
@@ -2,6 +2,7 @@ mod helpers;
mod list_threads;
use async_trait::async_trait;
use codex_protocol::ThreadId;
use crate::AppendThreadItemsParams;
use crate::ArchiveThreadParams;
@@ -9,11 +10,10 @@ use crate::CreateThreadParams;
use crate::ListThreadsParams;
use crate::LoadThreadHistoryParams;
use crate::ReadThreadParams;
use crate::ResumeThreadRecorderParams;
use crate::ResumeThreadParams;
use crate::StoredThread;
use crate::StoredThreadHistory;
use crate::ThreadPage;
use crate::ThreadRecorder;
use crate::ThreadStore;
use crate::ThreadStoreError;
use crate::ThreadStoreResult;
@@ -52,24 +52,34 @@ impl ThreadStore for RemoteThreadStore {
self
}
async fn create_thread(
&self,
_params: CreateThreadParams,
) -> ThreadStoreResult<Box<dyn ThreadRecorder>> {
async fn create_thread(&self, _params: CreateThreadParams) -> ThreadStoreResult<()> {
Err(not_implemented("create_thread"))
}
async fn resume_thread_recorder(
&self,
_params: ResumeThreadRecorderParams,
) -> ThreadStoreResult<Box<dyn ThreadRecorder>> {
Err(not_implemented("resume_thread_recorder"))
async fn resume_thread(&self, _params: ResumeThreadParams) -> ThreadStoreResult<()> {
Err(not_implemented("resume_thread"))
}
async fn append_items(&self, _params: AppendThreadItemsParams) -> ThreadStoreResult<()> {
Err(not_implemented("append_items"))
}
async fn persist_thread(&self, _thread_id: ThreadId) -> ThreadStoreResult<()> {
Err(not_implemented("persist_thread"))
}
async fn flush_thread(&self, _thread_id: ThreadId) -> ThreadStoreResult<()> {
Err(not_implemented("flush_thread"))
}
async fn shutdown_thread(&self, _thread_id: ThreadId) -> ThreadStoreResult<()> {
Err(not_implemented("shutdown_thread"))
}
async fn discard_thread(&self, _thread_id: ThreadId) -> ThreadStoreResult<()> {
Err(not_implemented("discard_thread"))
}
async fn load_history(
&self,
_params: LoadThreadHistoryParams,
codex-rs/thread-store/src/store.rs
+24 -15
View File
@@ -1,6 +1,6 @@
use std::any::Any;
use async_trait::async_trait;
use codex_protocol::ThreadId;
use std::any::Any;
use crate::AppendThreadItemsParams;
use crate::ArchiveThreadParams;
@@ -8,11 +8,10 @@ use crate::CreateThreadParams;
use crate::ListThreadsParams;
use crate::LoadThreadHistoryParams;
use crate::ReadThreadParams;
use crate::ResumeThreadRecorderParams;
use crate::ResumeThreadParams;
use crate::StoredThread;
use crate::StoredThreadHistory;
use crate::ThreadPage;
use crate::ThreadRecorder;
use crate::ThreadStoreResult;
use crate::UpdateThreadMetadataParams;
@@ -23,21 +22,31 @@ pub trait ThreadStore: Any + Send + Sync {
/// make sense for a concrete store implementation.
fn as_any(&self) -> &dyn Any;
/// Creates a new thread and returns a live recorder for future appends.
async fn create_thread(
&self,
params: CreateThreadParams,
) -> ThreadStoreResult<Box<dyn ThreadRecorder>>;
/// Creates a new live thread.
async fn create_thread(&self, params: CreateThreadParams) -> ThreadStoreResult<()>;
/// Reopens a live recorder for an existing thread.
async fn resume_thread_recorder(
&self,
params: ResumeThreadRecorderParams,
) -> ThreadStoreResult<Box<dyn ThreadRecorder>>;
/// Reopens an existing thread for live appends.
async fn resume_thread(&self, params: ResumeThreadParams) -> ThreadStoreResult<()>;
/// Appends items to a stored thread outside the live-recorder path.
/// Appends items to a live thread.
async fn append_items(&self, params: AppendThreadItemsParams) -> ThreadStoreResult<()>;
/// Materializes the thread if persistence is lazy, then persists all queued items.
async fn persist_thread(&self, thread_id: ThreadId) -> ThreadStoreResult<()>;
/// Flushes all queued items and returns once they are durable/readable.
async fn flush_thread(&self, thread_id: ThreadId) -> ThreadStoreResult<()>;
/// Flushes pending items and closes the live thread writer.
async fn shutdown_thread(&self, thread_id: ThreadId) -> ThreadStoreResult<()>;
/// Discards the live thread writer without forcing pending in-memory items to become durable.
///
/// Core calls this when session initialization fails after a live writer has been created.
/// Implementations should release any live writer resources for the thread while preserving
/// already-durable thread data.
async fn discard_thread(&self, thread_id: ThreadId) -> ThreadStoreResult<()>;
/// Loads persisted history for resume, fork, rollback, and memory jobs.
async fn load_history(
&self,
codex-rs/thread-store/src/types.rs
+10 -6
View File
@@ -26,7 +26,7 @@ pub enum ThreadEventPersistenceMode {
Extended,
}
/// Parameters required to create a persisted thread and its recorder.
/// Parameters required to create a persisted thread.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct CreateThreadParams {
/// Thread id generated by Codex before opening persistence.
@@ -39,22 +39,26 @@ pub struct CreateThreadParams {
pub base_instructions: BaseInstructions,
/// Dynamic tools available to the thread at startup.
pub dynamic_tools: Vec<DynamicToolSpec>,
/// Whether the recorder should persist the extended event surface.
/// Whether persistence should include the extended event surface.
pub event_persistence_mode: ThreadEventPersistenceMode,
}
/// Parameters required to reopen persistence for an existing thread.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ResumeThreadRecorderParams {
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ResumeThreadParams {
/// Existing thread id whose future items should be appended.
pub thread_id: ThreadId,
/// Known local rollout path when the caller resumed from a specific file.
pub rollout_path: Option<PathBuf>,
/// Known replay history for the resumed thread, if already loaded by the caller.
pub history: Option<Vec<RolloutItem>>,
/// Whether archived threads may be reopened.
pub include_archived: bool,
/// Whether the recorder should persist the extended event surface.
/// Whether persistence should include the extended event surface.
pub event_persistence_mode: ThreadEventPersistenceMode,
}
/// Parameters for appending rollout items outside a live recorder.
/// Parameters for appending rollout items to a live thread.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct AppendThreadItemsParams {
/// Thread id to append to.