diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index ef04a0e41..1ac5afb56 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3380,6 +3380,7 @@ dependencies = [ "tonic", "tonic-prost", "tonic-prost-build", + "tracing", "uuid", ] diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 631c83974..1a30d3263 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -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) { diff --git a/codex-rs/core/src/guardian/review_session.rs b/codex-rs/core/src/guardian/review_session.rs index 50123ef50..96778d0de 100644 --- a/codex-rs/core/src/guardian/review_session.rs +++ b/codex-rs/core/src/guardian/review_session.rs @@ -751,7 +751,7 @@ async fn load_rollout_items_for_fork( session: &Session, ) -> anyhow::Result>> { 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?; diff --git a/codex-rs/core/src/realtime_context.rs b/codex-rs/core/src/realtime_context.rs index bc769d29f..f9bb49a42 100644 --- a/codex-rs/core/src/realtime_context.rs +++ b/codex-rs/core/src/realtime_context.rs @@ -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; diff --git a/codex-rs/core/src/rollout.rs b/codex-rs/core/src/rollout.rs index 2405b83fc..4d282cdb5 100644 --- a/codex-rs/core/src/rollout.rs +++ b/codex-rs/core/src/rollout.rs @@ -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; } diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 9c91f7818..c0f9620f3 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -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, 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, 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, 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::>(); @@ -850,14 +834,12 @@ async fn persist_thread_name_update( ) -> anyhow::Result { 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, 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, 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), }), }; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 829e8d128..22a322b2a 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -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, pub(crate) parent_trace: Option, pub(crate) analytics_events_client: Option, + pub(crate) thread_store: Arc, } 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 { - let recorder = { - let guard = self.services.rollout.lock().await; - guard.clone() + pub(crate) async fn current_rollout_path(&self) -> anyhow::Result> { + 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 { 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( diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 95740cc4d..42e98ea58 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -272,6 +272,7 @@ impl Session { agent_control: AgentControl, environment_manager: Arc, analytics_events_client: Option, + thread_store: Arc, inherited_rollout_trace: RolloutTraceRecorder, ) -> anyhow::Result> { debug!( @@ -281,38 +282,16 @@ impl Session { ); let forked_from_id = initial_history.forked_from_id(); - let (conversation_id, rollout_params) = match &initial_history { + let event_persistence_mode = if session_configuration.persist_extended_history { + ThreadEventPersistenceMode::Extended + } else { + ThreadEventPersistenceMode::Limited + }; + let conversation_id = match &initial_history { InitialHistory::New | InitialHistory::Cleared | InitialHistory::Forked(_) => { - let conversation_id = ThreadId::default(); - ( - conversation_id, - RolloutRecorderParams::new( - conversation_id, - forked_from_id, - session_source, - BaseInstructions { - text: session_configuration.base_instructions.clone(), - }, - session_configuration.dynamic_tools.clone(), - if session_configuration.persist_extended_history { - EventPersistenceMode::Extended - } else { - EventPersistenceMode::Limited - }, - ), - ) + ThreadId::default() } - InitialHistory::Resumed(resumed_history) => ( - resumed_history.conversation_id, - RolloutRecorderParams::resume( - resumed_history.rollout_path.clone(), - if session_configuration.persist_extended_history { - EventPersistenceMode::Extended - } else { - EventPersistenceMode::Limited - }, - ), - ), + InitialHistory::Resumed(resumed_history) => resumed_history.conversation_id, }; let window_generation = match &initial_history { InitialHistory::Resumed(resumed_history) => u64::try_from( @@ -325,37 +304,68 @@ impl Session { .unwrap_or(u64::MAX), InitialHistory::New | InitialHistory::Cleared | InitialHistory::Forked(_) => 0, }; - let state_builder = match &initial_history { - InitialHistory::Resumed(resumed) => metadata::builder_from_items( - resumed.history.as_slice(), - resumed.rollout_path.as_path(), - ), - InitialHistory::New | InitialHistory::Cleared | InitialHistory::Forked(_) => None, - }; - // Kick off independent async setup tasks in parallel to reduce startup latency. // - // - initialize RolloutRecorder with new or resumed session info + // - initialize thread persistence with new or resumed session info // - perform default shell discovery // - load history metadata (skipped for subagents) - let rollout_fut = async { + let thread_persistence_fut = async { if config.ephemeral { - Ok::<_, anyhow::Error>((None, None)) + Ok::<_, anyhow::Error>(None) } else { - let state_db_ctx = state_db::init(&config).await; - let rollout_recorder = RolloutRecorder::new( - &config, - rollout_params, - state_db_ctx.clone(), - state_builder.clone(), - ) - .await?; - Ok((Some(rollout_recorder), state_db_ctx)) + let live_thread = match &initial_history { + InitialHistory::New | InitialHistory::Cleared | InitialHistory::Forked(_) => { + LiveThread::create( + Arc::clone(&thread_store), + CreateThreadParams { + thread_id: conversation_id, + forked_from_id, + source: session_source, + base_instructions: BaseInstructions { + text: session_configuration.base_instructions.clone(), + }, + dynamic_tools: session_configuration.dynamic_tools.clone(), + event_persistence_mode, + }, + ) + .await? + } + InitialHistory::Resumed(resumed_history) => { + LiveThread::resume( + Arc::clone(&thread_store), + ResumeThreadParams { + thread_id: resumed_history.conversation_id, + rollout_path: Some(resumed_history.rollout_path.clone()), + history: Some(resumed_history.history.clone()), + include_archived: true, + event_persistence_mode, + }, + ) + .await? + } + }; + Ok(Some(live_thread)) } } .instrument(info_span!( - "session_init.rollout", - otel.name = "session_init.rollout", + "session_init.thread_persistence", + otel.name = "session_init.thread_persistence", + session_init.ephemeral = config.ephemeral, + )); + let state_db_fut = async { + if config.ephemeral { + None + } else if let Some(local_store) = + thread_store.as_any().downcast_ref::() + { + local_store.state_db().await + } else { + None + } + } + .instrument(info_span!( + "session_init.state_db", + otel.name = "session_init.state_db", session_init.ephemeral = config.ephemeral, )); @@ -397,542 +407,562 @@ impl Session { // Join all independent futures. let ( - rollout_recorder_and_state_db, + thread_persistence_result, + state_db_ctx, (history_log_id, history_entry_count), (auth, mcp_servers, auth_statuses), - ) = tokio::join!(rollout_fut, history_meta_fut, auth_and_mcp_fut); + ) = tokio::join!( + thread_persistence_fut, + state_db_fut, + history_meta_fut, + auth_and_mcp_fut + ); - let (rollout_recorder, state_db_ctx) = rollout_recorder_and_state_db.map_err(|e| { - error!("failed to initialize rollout recorder: {e:#}"); - e - })?; - let rollout_path = rollout_recorder - .as_ref() - .map(|rec| rec.rollout_path().to_path_buf()); - let trace_agent_path = session_configuration - .session_source - .get_agent_path() - .unwrap_or_else(codex_protocol::AgentPath::root); - let trace_task_name = - (!trace_agent_path.is_root()).then(|| trace_agent_path.name().to_string()); - let trace_metadata = ThreadStartedTraceMetadata { - thread_id: conversation_id.to_string(), - agent_path: trace_agent_path.to_string(), - task_name: trace_task_name, - nickname: session_configuration.session_source.get_nickname(), - agent_role: session_configuration.session_source.get_agent_role(), - session_source: session_configuration.session_source.clone(), - cwd: session_configuration.cwd.to_path_buf(), - rollout_path: rollout_path.clone(), - model: session_configuration.collaboration_mode.model().to_string(), - provider_name: config.model_provider_id.clone(), - approval_policy: session_configuration.approval_policy.value().to_string(), - sandbox_policy: format!("{:?}", session_configuration.sandbox_policy.get()), - }; - let rollout_trace = if matches!( - session_configuration.session_source, - SessionSource::SubAgent(SubAgentSource::ThreadSpawn { .. }) - ) { - // Spawned child threads are part of their root rollout tree. If - // the parent had no trace recorder, do not create an orphan child - // bundle that looks like an independent rollout. - inherited_rollout_trace - } else { - RolloutTraceRecorder::create_root_or_disabled(conversation_id) - }; - rollout_trace.record_thread_started(trace_metadata); + let mut live_thread_init = + LiveThreadInitGuard::new(thread_persistence_result.map_err(|e| { + error!("failed to initialize thread persistence: {e:#}"); + e + })?); + let session_result: anyhow::Result> = async { + let rollout_path = if let Some(live_thread) = live_thread_init.as_ref() { + live_thread.local_rollout_path().await? + } else { + None + }; + let trace_agent_path = session_configuration + .session_source + .get_agent_path() + .unwrap_or_else(codex_protocol::AgentPath::root); + let trace_task_name = + (!trace_agent_path.is_root()).then(|| trace_agent_path.name().to_string()); + let trace_metadata = ThreadStartedTraceMetadata { + thread_id: conversation_id.to_string(), + agent_path: trace_agent_path.to_string(), + task_name: trace_task_name, + nickname: session_configuration.session_source.get_nickname(), + agent_role: session_configuration.session_source.get_agent_role(), + session_source: session_configuration.session_source.clone(), + cwd: session_configuration.cwd.to_path_buf(), + rollout_path: rollout_path.clone(), + model: session_configuration.collaboration_mode.model().to_string(), + provider_name: config.model_provider_id.clone(), + approval_policy: session_configuration.approval_policy.value().to_string(), + sandbox_policy: format!("{:?}", session_configuration.sandbox_policy.get()), + }; + let rollout_trace = if matches!( + session_configuration.session_source, + SessionSource::SubAgent(SubAgentSource::ThreadSpawn { .. }) + ) { + // Spawned child threads are part of their root rollout tree. If + // the parent had no trace recorder, do not create an orphan child + // bundle that looks like an independent rollout. + inherited_rollout_trace + } else { + RolloutTraceRecorder::create_root_or_disabled(conversation_id) + }; + rollout_trace.record_thread_started(trace_metadata); - let mut post_session_configured_events = Vec::::new(); + let mut post_session_configured_events = Vec::::new(); - for usage in config.features.legacy_feature_usages() { - post_session_configured_events.push(Event { - id: INITIAL_SUBMIT_ID.to_owned(), - msg: EventMsg::DeprecationNotice(DeprecationNoticeEvent { - summary: usage.summary.clone(), - details: usage.details.clone(), - }), - }); - } - if crate::config::uses_deprecated_instructions_file(&config.config_layer_stack) { - post_session_configured_events.push(Event { - id: INITIAL_SUBMIT_ID.to_owned(), - msg: EventMsg::DeprecationNotice(DeprecationNoticeEvent { - summary: "`experimental_instructions_file` is deprecated and ignored. Use `model_instructions_file` instead." - .to_string(), - details: Some( - "Move the setting to `model_instructions_file` in config.toml (or under a profile) to load instructions from a file." + for usage in config.features.legacy_feature_usages() { + post_session_configured_events.push(Event { + id: INITIAL_SUBMIT_ID.to_owned(), + msg: EventMsg::DeprecationNotice(DeprecationNoticeEvent { + summary: usage.summary.clone(), + details: usage.details.clone(), + }), + }); + } + if crate::config::uses_deprecated_instructions_file(&config.config_layer_stack) { + post_session_configured_events.push(Event { + id: INITIAL_SUBMIT_ID.to_owned(), + msg: EventMsg::DeprecationNotice(DeprecationNoticeEvent { + summary: "`experimental_instructions_file` is deprecated and ignored. Use `model_instructions_file` instead." .to_string(), - ), - }), - }); - } - for message in &config.startup_warnings { - post_session_configured_events.push(Event { - id: "".to_owned(), - msg: EventMsg::Warning(WarningEvent { - message: message.clone(), - }), - }); - } - let config_path = config.codex_home.join(CONFIG_TOML_FILE); - if let Some(event) = unstable_features_warning_event( - config - .config_layer_stack - .effective_config() - .get("features") - .and_then(TomlValue::as_table), - config.suppress_unstable_features_warning, - &config.features, - &config_path.display().to_string(), - ) { - post_session_configured_events.push(event); - } - if config.permissions.approval_policy.value() == AskForApproval::OnFailure { - post_session_configured_events.push(Event { - id: "".to_owned(), - msg: EventMsg::Warning(WarningEvent { - message: "`on-failure` approval policy is deprecated and will be removed in a future release. Use `on-request` for interactive approvals or `never` for non-interactive runs.".to_string(), - }), - }); - } + details: Some( + "Move the setting to `model_instructions_file` in config.toml (or under a profile) to load instructions from a file." + .to_string(), + ), + }), + }); + } + for message in &config.startup_warnings { + post_session_configured_events.push(Event { + id: "".to_owned(), + msg: EventMsg::Warning(WarningEvent { + message: message.clone(), + }), + }); + } + let config_path = config.codex_home.join(CONFIG_TOML_FILE); + if let Some(event) = unstable_features_warning_event( + config + .config_layer_stack + .effective_config() + .get("features") + .and_then(TomlValue::as_table), + config.suppress_unstable_features_warning, + &config.features, + &config_path.display().to_string(), + ) { + post_session_configured_events.push(event); + } + if config.permissions.approval_policy.value() == AskForApproval::OnFailure { + post_session_configured_events.push(Event { + id: "".to_owned(), + msg: EventMsg::Warning(WarningEvent { + message: "`on-failure` approval policy is deprecated and will be removed in a future release. Use `on-request` for interactive approvals or `never` for non-interactive runs.".to_string(), + }), + }); + } - let auth = auth.as_ref(); - let auth_mode = auth.map(CodexAuth::auth_mode).map(TelemetryAuthMode::from); - let account_id = auth.and_then(CodexAuth::get_account_id); - let account_email = auth.and_then(CodexAuth::get_account_email); - let originator = originator().value; - let terminal_type = user_agent(); - let session_model = session_configuration.collaboration_mode.model().to_string(); - let auth_env_telemetry = collect_auth_env_telemetry( - &session_configuration.provider, - auth_manager.codex_api_key_env_enabled(), - ); - let mut session_telemetry = SessionTelemetry::new( - conversation_id, - session_model.as_str(), - session_model.as_str(), - account_id.clone(), - account_email.clone(), - auth_mode, - originator.clone(), - config.otel.log_user_prompt, - terminal_type.clone(), - session_configuration.session_source.clone(), - ) - .with_auth_env(auth_env_telemetry.to_otel_metadata()); - if let Some(service_name) = session_configuration.metrics_service_name.as_deref() { - session_telemetry = session_telemetry.with_metrics_service_name(service_name); - } - let network_proxy_audit_metadata = NetworkProxyAuditMetadata { - conversation_id: Some(conversation_id.to_string()), - app_version: Some(env!("CARGO_PKG_VERSION").to_string()), - user_account_id: account_id, - auth_mode: auth_mode.map(|mode| mode.to_string()), - originator: Some(originator), - user_email: account_email, - terminal_type: Some(terminal_type), - model: Some(session_model.clone()), - slug: Some(session_model), - }; - config.features.emit_metrics(&session_telemetry); - session_telemetry.counter( - THREAD_STARTED_METRIC, - /*inc*/ 1, - &[( - "is_git", - if get_git_repo_root(&session_configuration.cwd).is_some() { - "true" + let auth = auth.as_ref(); + let auth_mode = auth.map(CodexAuth::auth_mode).map(TelemetryAuthMode::from); + let account_id = auth.and_then(CodexAuth::get_account_id); + let account_email = auth.and_then(CodexAuth::get_account_email); + let originator = originator().value; + let terminal_type = user_agent(); + let session_model = session_configuration.collaboration_mode.model().to_string(); + let auth_env_telemetry = collect_auth_env_telemetry( + &session_configuration.provider, + auth_manager.codex_api_key_env_enabled(), + ); + let mut session_telemetry = SessionTelemetry::new( + conversation_id, + session_model.as_str(), + session_model.as_str(), + account_id.clone(), + account_email.clone(), + auth_mode, + originator.clone(), + config.otel.log_user_prompt, + terminal_type.clone(), + session_configuration.session_source.clone(), + ) + .with_auth_env(auth_env_telemetry.to_otel_metadata()); + if let Some(service_name) = session_configuration.metrics_service_name.as_deref() { + session_telemetry = session_telemetry.with_metrics_service_name(service_name); + } + let network_proxy_audit_metadata = NetworkProxyAuditMetadata { + conversation_id: Some(conversation_id.to_string()), + app_version: Some(env!("CARGO_PKG_VERSION").to_string()), + user_account_id: account_id, + auth_mode: auth_mode.map(|mode| mode.to_string()), + originator: Some(originator), + user_email: account_email, + terminal_type: Some(terminal_type), + model: Some(session_model.clone()), + slug: Some(session_model), + }; + config.features.emit_metrics(&session_telemetry); + session_telemetry.counter( + THREAD_STARTED_METRIC, + /*inc*/ 1, + &[( + "is_git", + if get_git_repo_root(&session_configuration.cwd).is_some() { + "true" + } else { + "false" + }, + )], + ); + + session_telemetry.conversation_starts( + config.model_provider.name.as_str(), + session_configuration.collaboration_mode.reasoning_effort(), + config + .model_reasoning_summary + .unwrap_or(ReasoningSummaryConfig::Auto), + config.model_context_window, + config.model_auto_compact_token_limit, + config.permissions.approval_policy.value(), + config.permissions.sandbox_policy.get().clone(), + mcp_servers.keys().map(String::as_str).collect(), + config.active_profile.clone(), + ); + + let use_zsh_fork_shell = config.features.enabled(Feature::ShellZshFork); + let mut default_shell = if let Some(user_shell_override) = + session_configuration.user_shell_override.clone() + { + user_shell_override + } else if use_zsh_fork_shell { + let zsh_path = config.zsh_path.as_ref().ok_or_else(|| { + anyhow::anyhow!( + "zsh fork feature enabled, but `zsh_path` is not configured; set `zsh_path` in config.toml" + ) + })?; + let zsh_path = zsh_path.to_path_buf(); + shell::get_shell(shell::ShellType::Zsh, Some(&zsh_path)).ok_or_else(|| { + anyhow::anyhow!( + "zsh fork feature enabled, but zsh_path `{}` is not usable; set `zsh_path` to a valid zsh executable", + zsh_path.display() + ) + })? + } else { + shell::default_user_shell() + }; + // Create the mutable state for the Session. + let shell_snapshot_tx = if config.features.enabled(Feature::ShellSnapshot) { + if let Some(snapshot) = session_configuration.inherited_shell_snapshot.clone() { + let (tx, rx) = watch::channel(Some(snapshot)); + default_shell.shell_snapshot = rx; + tx } else { - "false" - }, - )], - ); - - session_telemetry.conversation_starts( - config.model_provider.name.as_str(), - session_configuration.collaboration_mode.reasoning_effort(), - config - .model_reasoning_summary - .unwrap_or(ReasoningSummaryConfig::Auto), - config.model_context_window, - config.model_auto_compact_token_limit, - config.permissions.approval_policy.value(), - config.permissions.sandbox_policy.get().clone(), - mcp_servers.keys().map(String::as_str).collect(), - config.active_profile.clone(), - ); - - let use_zsh_fork_shell = config.features.enabled(Feature::ShellZshFork); - let mut default_shell = if let Some(user_shell_override) = - session_configuration.user_shell_override.clone() - { - user_shell_override - } else if use_zsh_fork_shell { - let zsh_path = config.zsh_path.as_ref().ok_or_else(|| { - anyhow::anyhow!( - "zsh fork feature enabled, but `zsh_path` is not configured; set `zsh_path` in config.toml" - ) - })?; - let zsh_path = zsh_path.to_path_buf(); - shell::get_shell(shell::ShellType::Zsh, Some(&zsh_path)).ok_or_else(|| { - anyhow::anyhow!( - "zsh fork feature enabled, but zsh_path `{}` is not usable; set `zsh_path` to a valid zsh executable", - zsh_path.display() - ) - })? - } else { - shell::default_user_shell() - }; - // Create the mutable state for the Session. - let shell_snapshot_tx = if config.features.enabled(Feature::ShellSnapshot) { - if let Some(snapshot) = session_configuration.inherited_shell_snapshot.clone() { - let (tx, rx) = watch::channel(Some(snapshot)); + ShellSnapshot::start_snapshotting( + config.codex_home.clone(), + conversation_id, + session_configuration.cwd.clone(), + &mut default_shell, + session_telemetry.clone(), + ) + } + } else { + let (tx, rx) = watch::channel(None); default_shell.shell_snapshot = rx; tx + }; + let thread_name = + thread_title_from_state_db(state_db_ctx.as_ref(), &config.codex_home, conversation_id) + .instrument(info_span!( + "session_init.thread_name_lookup", + otel.name = "session_init.thread_name_lookup", + )) + .await; + session_configuration.thread_name = thread_name.clone(); + let state = SessionState::new(session_configuration.clone()); + let managed_network_requirements_configured = config + .config_layer_stack + .requirements_toml() + .network + .is_some(); + let managed_network_requirements_enabled = config.managed_network_requirements_enabled(); + let network_approval = Arc::new(NetworkApprovalService::default()); + // The managed proxy can call back into core for allowlist-miss decisions. + let network_policy_decider_session = if managed_network_requirements_configured { + config + .permissions + .network + .as_ref() + .map(|_| Arc::new(RwLock::new(std::sync::Weak::::new()))) } else { - ShellSnapshot::start_snapshotting( - config.codex_home.clone(), - conversation_id, - session_configuration.cwd.clone(), - &mut default_shell, - session_telemetry.clone(), - ) - } - } else { - let (tx, rx) = watch::channel(None); - default_shell.shell_snapshot = rx; - tx - }; - let thread_name = - thread_title_from_state_db(state_db_ctx.as_ref(), &config.codex_home, conversation_id) - .instrument(info_span!( - "session_init.thread_name_lookup", - otel.name = "session_init.thread_name_lookup", - )) - .await; - session_configuration.thread_name = thread_name.clone(); - let state = SessionState::new(session_configuration.clone()); - let managed_network_requirements_configured = config - .config_layer_stack - .requirements_toml() - .network - .is_some(); - let managed_network_requirements_enabled = config.managed_network_requirements_enabled(); - let network_approval = Arc::new(NetworkApprovalService::default()); - // The managed proxy can call back into core for allowlist-miss decisions. - let network_policy_decider_session = if managed_network_requirements_configured { - config - .permissions - .network - .as_ref() - .map(|_| Arc::new(RwLock::new(std::sync::Weak::::new()))) - } else { - None - }; - let blocked_request_observer = if managed_network_requirements_configured { - config - .permissions - .network - .as_ref() - .map(|_| build_blocked_request_observer(Arc::clone(&network_approval))) - } else { - None - }; - let network_policy_decider = - network_policy_decider_session - .as_ref() - .map(|network_policy_decider_session| { - build_network_policy_decider( - Arc::clone(&network_approval), - Arc::clone(network_policy_decider_session), + None + }; + let blocked_request_observer = if managed_network_requirements_configured { + config + .permissions + .network + .as_ref() + .map(|_| build_blocked_request_observer(Arc::clone(&network_approval))) + } else { + None + }; + let network_policy_decider = + network_policy_decider_session + .as_ref() + .map(|network_policy_decider_session| { + build_network_policy_decider( + Arc::clone(&network_approval), + Arc::clone(network_policy_decider_session), + ) + }); + let (network_proxy, session_network_proxy) = + if let Some(spec) = config.permissions.network.as_ref() { + let current_exec_policy = exec_policy.current(); + let (network_proxy, session_network_proxy) = Self::start_managed_network_proxy( + spec, + current_exec_policy.as_ref(), + config.permissions.sandbox_policy.get(), + network_policy_decider.as_ref().map(Arc::clone), + blocked_request_observer.as_ref().map(Arc::clone), + managed_network_requirements_configured, + network_proxy_audit_metadata, ) + .instrument(info_span!( + "session_init.network_proxy", + otel.name = "session_init.network_proxy", + session_init.managed_network_requirements_enabled = + managed_network_requirements_enabled, + )) + .await?; + (Some(network_proxy), Some(session_network_proxy)) + } else { + (None, None) + }; + + let mut hook_shell_argv = + default_shell.derive_exec_args("", /*use_login_shell*/ false); + let hook_shell_program = hook_shell_argv.remove(0); + let _ = hook_shell_argv.pop(); + let hooks = Hooks::new(HooksConfig { + legacy_notify_argv: config.notify.clone(), + feature_enabled: config.features.enabled(Feature::CodexHooks), + config_layer_stack: Some(config.config_layer_stack.clone()), + shell_program: Some(hook_shell_program), + shell_args: hook_shell_argv, + }); + for warning in hooks.startup_warnings() { + post_session_configured_events.push(Event { + id: INITIAL_SUBMIT_ID.to_owned(), + msg: EventMsg::Warning(WarningEvent { + message: warning.clone(), + }), }); - let (network_proxy, session_network_proxy) = - if let Some(spec) = config.permissions.network.as_ref() { - let current_exec_policy = exec_policy.current(); - let (network_proxy, session_network_proxy) = Self::start_managed_network_proxy( - spec, - current_exec_policy.as_ref(), - config.permissions.sandbox_policy.get(), - network_policy_decider.as_ref().map(Arc::clone), - blocked_request_observer.as_ref().map(Arc::clone), - managed_network_requirements_configured, - network_proxy_audit_metadata, + } + + let installation_id = resolve_installation_id(&config.codex_home).await?; + let analytics_events_client = analytics_events_client.unwrap_or_else(|| { + AnalyticsEventsClient::new( + Arc::clone(&auth_manager), + config.chatgpt_base_url.trim_end_matches('/').to_string(), + config.analytics_enabled, ) - .instrument(info_span!( - "session_init.network_proxy", - otel.name = "session_init.network_proxy", - session_init.managed_network_requirements_enabled = - managed_network_requirements_enabled, - )) - .await?; - (Some(network_proxy), Some(session_network_proxy)) + }); + let services = SessionServices { + // Initialize the MCP connection manager with an uninitialized + // instance. It will be replaced with one created via + // McpConnectionManager::new() once all its constructor args are + // available. This also ensures `SessionConfigured` is emitted + // before any MCP-related events. It is reasonable to consider + // changing this to use Option or OnceCell, though the current + // setup is straightforward enough and performs well. + mcp_connection_manager: Arc::new(RwLock::new(McpConnectionManager::new_uninitialized( + &config.permissions.approval_policy, + &config.permissions.sandbox_policy, + ))), + mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()), + unified_exec_manager: UnifiedExecProcessManager::new( + config.background_terminal_max_timeout, + ), + shell_zsh_path: config.zsh_path.clone(), + main_execve_wrapper_exe: config.main_execve_wrapper_exe.clone(), + analytics_events_client, + hooks, + rollout_trace, + user_shell: Arc::new(default_shell), + shell_snapshot_tx, + show_raw_agent_reasoning: config.show_raw_agent_reasoning, + exec_policy, + auth_manager: Arc::clone(&auth_manager), + session_telemetry, + models_manager: Arc::clone(&models_manager), + tool_approvals: Mutex::new(ApprovalStore::default()), + guardian_rejections: Mutex::new(HashMap::new()), + guardian_rejection_circuit_breaker: Mutex::new(Default::default()), + runtime_handle: tokio::runtime::Handle::current(), + skills_manager, + plugins_manager: Arc::clone(&plugins_manager), + mcp_manager: Arc::clone(&mcp_manager), + skills_watcher, + agent_control, + network_proxy, + network_approval: Arc::clone(&network_approval), + state_db: state_db_ctx.clone(), + live_thread: live_thread_init.as_ref().cloned(), + thread_store: Arc::clone(&thread_store), + model_client: ModelClient::new( + Some(Arc::clone(&auth_manager)), + conversation_id, + installation_id, + session_configuration.provider.clone(), + session_configuration.session_source.clone(), + config.model_verbosity, + config.features.enabled(Feature::EnableRequestCompression), + config.features.enabled(Feature::RuntimeMetrics), + Self::build_model_client_beta_features_header(config.as_ref()), + ), + code_mode_service: crate::tools::code_mode::CodeModeService::new( + config.js_repl_node_path.clone(), + ), + environment_manager, + }; + services + .model_client + .set_window_generation(window_generation); + let js_repl = Arc::new(JsReplHandle::with_node_path( + config.js_repl_node_path.clone(), + config.js_repl_node_module_dirs.clone(), + )); + let (out_of_band_elicitation_paused, _out_of_band_elicitation_paused_rx) = + watch::channel(false); + + let (mailbox, mailbox_rx) = Mailbox::new(); + let sess = Arc::new(Session { + conversation_id, + tx_event: tx_event.clone(), + agent_status, + out_of_band_elicitation_paused, + state: Mutex::new(state), + managed_network_proxy_refresh_lock: Semaphore::new(/*permits*/ 1), + features: config.features.clone(), + pending_mcp_server_refresh_config: Mutex::new(None), + conversation: Arc::new(RealtimeConversationManager::new()), + active_turn: Mutex::new(None), + mailbox, + mailbox_rx: Mutex::new(mailbox_rx), + idle_pending_input: Mutex::new(Vec::new()), + guardian_review_session: GuardianReviewSessionManager::default(), + services, + js_repl, + next_internal_sub_id: AtomicU64::new(0), + }); + if let Some(network_policy_decider_session) = network_policy_decider_session { + let mut guard = network_policy_decider_session.write().await; + *guard = Arc::downgrade(&sess); + } + // Dispatch the SessionConfiguredEvent first and then report any errors. + // If resuming, include converted initial messages in the payload so UIs can render them immediately. + let initial_messages = initial_history.get_event_msgs(); + let permission_profile = if matches!( + session_configuration.file_system_sandbox_policy.kind, + FileSystemSandboxKind::ExternalSandbox + ) { + None } else { - (None, None) + Some(session_configuration.permission_profile()) + }; + let events = std::iter::once(Event { + id: INITIAL_SUBMIT_ID.to_owned(), + msg: EventMsg::SessionConfigured(SessionConfiguredEvent { + session_id: conversation_id, + forked_from_id, + thread_name: session_configuration.thread_name.clone(), + model: session_configuration.collaboration_mode.model().to_string(), + model_provider_id: config.model_provider_id.clone(), + service_tier: session_configuration.service_tier, + approval_policy: session_configuration.approval_policy.value(), + approvals_reviewer: session_configuration.approvals_reviewer, + sandbox_policy: session_configuration.sandbox_policy.get().clone(), + permission_profile, + cwd: session_configuration.cwd.clone(), + reasoning_effort: session_configuration.collaboration_mode.reasoning_effort(), + history_log_id, + history_entry_count, + initial_messages, + network_proxy: session_network_proxy.filter(|_| { + Self::managed_network_proxy_active_for_sandbox_policy( + session_configuration.sandbox_policy.get(), + ) + }), + rollout_path, + }), + }) + .chain(post_session_configured_events.into_iter()); + for event in events { + sess.send_event_raw(event).await; + } + + // Start the watcher after SessionConfigured so it cannot emit earlier events. + sess.start_skills_watcher_listener(); + let mut required_mcp_servers: Vec = mcp_servers + .iter() + .filter(|(_, server)| server.enabled && server.required) + .map(|(name, _)| name.clone()) + .collect(); + required_mcp_servers.sort(); + let enabled_mcp_server_count = mcp_servers.values().filter(|server| server.enabled).count(); + let required_mcp_server_count = required_mcp_servers.len(); + let tool_plugin_provenance = mcp_manager.tool_plugin_provenance(config.as_ref()).await; + { + let mut cancel_guard = sess.services.mcp_startup_cancellation_token.lock().await; + cancel_guard.cancel(); + *cancel_guard = CancellationToken::new(); + } + let (mcp_connection_manager, cancel_token) = McpConnectionManager::new( + &mcp_servers, + config.mcp_oauth_credentials_store_mode, + auth_statuses.clone(), + &session_configuration.approval_policy, + INITIAL_SUBMIT_ID.to_owned(), + tx_event.clone(), + session_configuration.sandbox_policy.get().clone(), + McpRuntimeEnvironment::new( + sess.services + .environment_manager + .default_environment() + .unwrap_or_else(|| sess.services.environment_manager.local_environment()), + session_configuration.cwd.to_path_buf(), + ), + config.codex_home.to_path_buf(), + codex_apps_tools_cache_key(auth), + tool_plugin_provenance, + ) + .instrument(info_span!( + "session_init.mcp_manager_init", + otel.name = "session_init.mcp_manager_init", + session_init.enabled_mcp_server_count = enabled_mcp_server_count, + session_init.required_mcp_server_count = required_mcp_server_count, + )) + .await; + { + let mut manager_guard = sess.services.mcp_connection_manager.write().await; + *manager_guard = mcp_connection_manager; + } + { + let mut cancel_guard = sess.services.mcp_startup_cancellation_token.lock().await; + if cancel_guard.is_cancelled() { + cancel_token.cancel(); + } + *cancel_guard = cancel_token; + } + if !required_mcp_servers.is_empty() { + let failures = sess + .services + .mcp_connection_manager + .read() + .await + .required_startup_failures(&required_mcp_servers) + .instrument(info_span!( + "session_init.required_mcp_wait", + otel.name = "session_init.required_mcp_wait", + session_init.required_mcp_server_count = required_mcp_server_count, + )) + .await; + if !failures.is_empty() { + let details = failures + .iter() + .map(|failure| format!("{}: {}", failure.server, failure.error)) + .collect::>() + .join("; "); + anyhow::bail!("required MCP servers failed to initialize: {details}"); + } + } + sess.schedule_startup_prewarm(session_configuration.base_instructions.clone()) + .await; + let session_start_source = match &initial_history { + InitialHistory::Resumed(_) => codex_hooks::SessionStartSource::Resume, + InitialHistory::New | InitialHistory::Forked(_) => { + codex_hooks::SessionStartSource::Startup + } + InitialHistory::Cleared => codex_hooks::SessionStartSource::Clear, }; - let mut hook_shell_argv = - default_shell.derive_exec_args("", /*use_login_shell*/ false); - let hook_shell_program = hook_shell_argv.remove(0); - let _ = hook_shell_argv.pop(); - let hooks = Hooks::new(HooksConfig { - legacy_notify_argv: config.notify.clone(), - feature_enabled: config.features.enabled(Feature::CodexHooks), - config_layer_stack: Some(config.config_layer_stack.clone()), - shell_program: Some(hook_shell_program), - shell_args: hook_shell_argv, - }); - for warning in hooks.startup_warnings() { - post_session_configured_events.push(Event { - id: INITIAL_SUBMIT_ID.to_owned(), - msg: EventMsg::Warning(WarningEvent { - message: warning.clone(), - }), - }); - } + // record_initial_history can emit events. We record only after the SessionConfiguredEvent is emitted. + sess.record_initial_history(initial_history).await; + { + let mut state = sess.state.lock().await; + state.set_pending_session_start_source(Some(session_start_source)); + } - let installation_id = resolve_installation_id(&config.codex_home).await?; - let analytics_events_client = analytics_events_client.unwrap_or_else(|| { - AnalyticsEventsClient::new( - Arc::clone(&auth_manager), - config.chatgpt_base_url.trim_end_matches('/').to_string(), - config.analytics_enabled, - ) - }); - let services = SessionServices { - // Initialize the MCP connection manager with an uninitialized - // instance. It will be replaced with one created via - // McpConnectionManager::new() once all its constructor args are - // available. This also ensures `SessionConfigured` is emitted - // before any MCP-related events. It is reasonable to consider - // changing this to use Option or OnceCell, though the current - // setup is straightforward enough and performs well. - mcp_connection_manager: Arc::new(RwLock::new(McpConnectionManager::new_uninitialized( - &config.permissions.approval_policy, - &config.permissions.sandbox_policy, - ))), - mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()), - unified_exec_manager: UnifiedExecProcessManager::new( - config.background_terminal_max_timeout, - ), - shell_zsh_path: config.zsh_path.clone(), - main_execve_wrapper_exe: config.main_execve_wrapper_exe.clone(), - analytics_events_client, - hooks, - rollout: Mutex::new(rollout_recorder), - rollout_trace, - user_shell: Arc::new(default_shell), - shell_snapshot_tx, - show_raw_agent_reasoning: config.show_raw_agent_reasoning, - exec_policy, - auth_manager: Arc::clone(&auth_manager), - session_telemetry, - models_manager: Arc::clone(&models_manager), - tool_approvals: Mutex::new(ApprovalStore::default()), - guardian_rejections: Mutex::new(HashMap::new()), - guardian_rejection_circuit_breaker: Mutex::new(Default::default()), - runtime_handle: tokio::runtime::Handle::current(), - skills_manager, - plugins_manager: Arc::clone(&plugins_manager), - mcp_manager: Arc::clone(&mcp_manager), - skills_watcher, - agent_control, - network_proxy, - network_approval: Arc::clone(&network_approval), - state_db: state_db_ctx.clone(), - thread_store: LocalThreadStore::new(RolloutConfig::from_view(config.as_ref())), - model_client: ModelClient::new( - Some(Arc::clone(&auth_manager)), - conversation_id, - installation_id, - session_configuration.provider.clone(), - session_configuration.session_source.clone(), - config.model_verbosity, - config.features.enabled(Feature::EnableRequestCompression), - config.features.enabled(Feature::RuntimeMetrics), - Self::build_model_client_beta_features_header(config.as_ref()), - ), - code_mode_service: crate::tools::code_mode::CodeModeService::new( - config.js_repl_node_path.clone(), - ), - environment_manager, - }; - services - .model_client - .set_window_generation(window_generation); - let js_repl = Arc::new(JsReplHandle::with_node_path( - config.js_repl_node_path.clone(), - config.js_repl_node_module_dirs.clone(), - )); - let (out_of_band_elicitation_paused, _out_of_band_elicitation_paused_rx) = - watch::channel(false); + memories::start_memories_startup_task( + &sess, + Arc::clone(&config), + &session_configuration.session_source, + ); - let (mailbox, mailbox_rx) = Mailbox::new(); - let sess = Arc::new(Session { - conversation_id, - tx_event: tx_event.clone(), - agent_status, - out_of_band_elicitation_paused, - state: Mutex::new(state), - managed_network_proxy_refresh_lock: Semaphore::new(/*permits*/ 1), - features: config.features.clone(), - pending_mcp_server_refresh_config: Mutex::new(None), - conversation: Arc::new(RealtimeConversationManager::new()), - active_turn: Mutex::new(None), - mailbox, - mailbox_rx: Mutex::new(mailbox_rx), - idle_pending_input: Mutex::new(Vec::new()), - guardian_review_session: GuardianReviewSessionManager::default(), - services, - js_repl, - next_internal_sub_id: AtomicU64::new(0), - }); - if let Some(network_policy_decider_session) = network_policy_decider_session { - let mut guard = network_policy_decider_session.write().await; - *guard = Arc::downgrade(&sess); + Ok(sess) } - // Dispatch the SessionConfiguredEvent first and then report any errors. - // If resuming, include converted initial messages in the payload so UIs can render them immediately. - let initial_messages = initial_history.get_event_msgs(); - let permission_profile = if matches!( - session_configuration.file_system_sandbox_policy.kind, - FileSystemSandboxKind::ExternalSandbox - ) { - None - } else { - Some(session_configuration.permission_profile()) - }; - let events = std::iter::once(Event { - id: INITIAL_SUBMIT_ID.to_owned(), - msg: EventMsg::SessionConfigured(SessionConfiguredEvent { - session_id: conversation_id, - forked_from_id, - thread_name: session_configuration.thread_name.clone(), - model: session_configuration.collaboration_mode.model().to_string(), - model_provider_id: config.model_provider_id.clone(), - service_tier: session_configuration.service_tier, - approval_policy: session_configuration.approval_policy.value(), - approvals_reviewer: session_configuration.approvals_reviewer, - sandbox_policy: session_configuration.sandbox_policy.get().clone(), - permission_profile, - cwd: session_configuration.cwd.clone(), - reasoning_effort: session_configuration.collaboration_mode.reasoning_effort(), - history_log_id, - history_entry_count, - initial_messages, - network_proxy: session_network_proxy.filter(|_| { - Self::managed_network_proxy_active_for_sandbox_policy( - session_configuration.sandbox_policy.get(), - ) - }), - rollout_path, - }), - }) - .chain(post_session_configured_events.into_iter()); - for event in events { - sess.send_event_raw(event).await; - } - - // Start the watcher after SessionConfigured so it cannot emit earlier events. - sess.start_skills_watcher_listener(); - let mut required_mcp_servers: Vec = mcp_servers - .iter() - .filter(|(_, server)| server.enabled && server.required) - .map(|(name, _)| name.clone()) - .collect(); - required_mcp_servers.sort(); - let enabled_mcp_server_count = mcp_servers.values().filter(|server| server.enabled).count(); - let required_mcp_server_count = required_mcp_servers.len(); - let tool_plugin_provenance = mcp_manager.tool_plugin_provenance(config.as_ref()).await; - { - let mut cancel_guard = sess.services.mcp_startup_cancellation_token.lock().await; - cancel_guard.cancel(); - *cancel_guard = CancellationToken::new(); - } - let (mcp_connection_manager, cancel_token) = McpConnectionManager::new( - &mcp_servers, - config.mcp_oauth_credentials_store_mode, - auth_statuses.clone(), - &session_configuration.approval_policy, - INITIAL_SUBMIT_ID.to_owned(), - tx_event.clone(), - session_configuration.sandbox_policy.get().clone(), - McpRuntimeEnvironment::new( - sess.services - .environment_manager - .default_environment() - .unwrap_or_else(|| sess.services.environment_manager.local_environment()), - session_configuration.cwd.to_path_buf(), - ), - config.codex_home.to_path_buf(), - codex_apps_tools_cache_key(auth), - tool_plugin_provenance, - ) - .instrument(info_span!( - "session_init.mcp_manager_init", - otel.name = "session_init.mcp_manager_init", - session_init.enabled_mcp_server_count = enabled_mcp_server_count, - session_init.required_mcp_server_count = required_mcp_server_count, - )) .await; - { - let mut manager_guard = sess.services.mcp_connection_manager.write().await; - *manager_guard = mcp_connection_manager; - } - { - let mut cancel_guard = sess.services.mcp_startup_cancellation_token.lock().await; - if cancel_guard.is_cancelled() { - cancel_token.cancel(); + match session_result { + Ok(sess) => { + live_thread_init.commit(); + Ok(sess) } - *cancel_guard = cancel_token; - } - if !required_mcp_servers.is_empty() { - let failures = sess - .services - .mcp_connection_manager - .read() - .await - .required_startup_failures(&required_mcp_servers) - .instrument(info_span!( - "session_init.required_mcp_wait", - otel.name = "session_init.required_mcp_wait", - session_init.required_mcp_server_count = required_mcp_server_count, - )) - .await; - if !failures.is_empty() { - let details = failures - .iter() - .map(|failure| format!("{}: {}", failure.server, failure.error)) - .collect::>() - .join("; "); - return Err(anyhow::anyhow!( - "required MCP servers failed to initialize: {details}" - )); + Err(err) => { + live_thread_init.discard().await; + Err(err) } } - sess.schedule_startup_prewarm(session_configuration.base_instructions.clone()) - .await; - let session_start_source = match &initial_history { - InitialHistory::Resumed(_) => codex_hooks::SessionStartSource::Resume, - InitialHistory::New | InitialHistory::Forked(_) => { - codex_hooks::SessionStartSource::Startup - } - InitialHistory::Cleared => codex_hooks::SessionStartSource::Clear, - }; - - // record_initial_history can emit events. We record only after the SessionConfiguredEvent is emitted. - sess.record_initial_history(initial_history).await; - { - let mut state = sess.state.lock().await; - state.set_pending_session_start_source(Some(session_start_source)); - } - - memories::start_memories_startup_task( - &sess, - Arc::clone(&config), - &session_configuration.session_source, - ); - - Ok(sess) } } diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 4e47596b8..b8e046a95 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -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) -> } } -async fn attach_rollout_recorder(session: &Arc) -> 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( diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index 22db9f393..45ebc2636 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -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"); diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index f93a44d7f..2c62e04c8 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -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, pub(crate) analytics_events_client: AnalyticsEventsClient, pub(crate) hooks: Hooks, - pub(crate) rollout: Mutex>, pub(crate) rollout_trace: RolloutTraceRecorder, pub(crate) user_shell: Arc, pub(crate) shell_snapshot_tx: watch::Sender>>, @@ -64,7 +63,8 @@ pub(crate) struct SessionServices { pub(crate) network_proxy: Option, pub(crate) network_approval: Arc, pub(crate) state_db: Option, - pub(crate) thread_store: LocalThreadStore, + pub(crate) live_thread: Option, + pub(crate) thread_store: Arc, /// Session-scoped model client shared across turns. pub(crate) model_client: ModelClient, pub(crate) code_mode_service: CodeModeService, diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 4b8e9cc53..af7c3023e 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -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 { + 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, parent_trace: Option, ) -> CodexResult { + 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, ) -> CodexResult { + 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 { + 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 { 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>, inherited_exec_policy: Option>, ) -> CodexResult { + 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>, ) -> CodexResult { 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>, inherited_exec_policy: Option>, ) -> CodexResult { + 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, initial_history: InitialHistory, auth_manager: Arc, agent_control: AgentControl, @@ -889,6 +917,7 @@ impl ThreadManagerState { ) -> CodexResult { 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, initial_history: InitialHistory, auth_manager: Arc, 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) diff --git a/codex-rs/thread-store/Cargo.toml b/codex-rs/thread-store/Cargo.toml index 61e4e4145..3a0428f20 100644 --- a/codex-rs/thread-store/Cargo.toml +++ b/codex-rs/thread-store/Cargo.toml @@ -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 } diff --git a/codex-rs/thread-store/src/lib.rs b/codex-rs/thread-store/src/lib.rs index cb87bdd04..15ec238fa 100644 --- a/codex-rs/thread-store/src/lib.rs +++ b/codex-rs/thread-store/src/lib.rs @@ -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; diff --git a/codex-rs/thread-store/src/live_thread.rs b/codex-rs/thread-store/src/live_thread.rs new file mode 100644 index 000000000..bcce1c764 --- /dev/null +++ b/codex-rs/thread-store/src/live_thread.rs @@ -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, +} + +/// 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, +} + +impl LiveThreadInitGuard { + pub fn new(live_thread: Option) -> 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, + params: CreateThreadParams, + ) -> ThreadStoreResult { + 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, + params: ResumeThreadParams, + ) -> ThreadStoreResult { + 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 { + 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> { + let Some(local_store) = self + .thread_store + .as_any() + .downcast_ref::() + else { + return Ok(None); + }; + local_store + .live_rollout_path(self.thread_id) + .await + .map(Some) + } +} diff --git a/codex-rs/thread-store/src/local/create_thread.rs b/codex-rs/thread-store/src/local/create_thread.rs new file mode 100644 index 000000000..69f19adc6 --- /dev/null +++ b/codex-rs/thread-store/src/local/create_thread.rs @@ -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 { + 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, + } +} diff --git a/codex-rs/thread-store/src/local/live_writer.rs b/codex-rs/thread-store/src/local/live_writer.rs new file mode 100644 index 000000000..fd8cd93c7 --- /dev/null +++ b/codex-rs/thread-store/src/local/live_writer.rs @@ -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 { + 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(), + } +} diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs index c7c62f4cf..e5f73ff83 100644 --- a/codex-rs/thread-store/src/local/mod.rs +++ b/codex-rs/thread-store/src/local/mod.rs @@ -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>>, + state_db: Arc>, +} + +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 { + 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 { @@ -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 { + live_writer::rollout_path(self, thread_id).await + } + + pub(super) async fn live_recorder( + &self, + thread_id: ThreadId, + ) -> ThreadStoreResult { + 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> { - 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> { - 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 { - 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 { @@ -114,8 +230,313 @@ impl ThreadStore for LocalThreadStore { } } -fn unsupported(operation: &str) -> ThreadStoreResult { - 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 + ) + })); + } } diff --git a/codex-rs/thread-store/src/local/update_thread_metadata.rs b/codex-rs/thread-store/src/local/update_thread_metadata.rs index fae90e017..52c937c6b 100644 --- a/codex-rs/thread-store/src/local/update_thread_metadata.rs +++ b/codex-rs/thread-store/src/local/update_thread_metadata.rs @@ -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 { - 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 { + 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") diff --git a/codex-rs/thread-store/src/recorder.rs b/codex-rs/thread-store/src/recorder.rs deleted file mode 100644 index 03b02e80c..000000000 --- a/codex-rs/thread-store/src/recorder.rs +++ /dev/null @@ -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<()>; -} diff --git a/codex-rs/thread-store/src/remote/mod.rs b/codex-rs/thread-store/src/remote/mod.rs index 5be760870..b3ddeecf7 100644 --- a/codex-rs/thread-store/src/remote/mod.rs +++ b/codex-rs/thread-store/src/remote/mod.rs @@ -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> { + async fn create_thread(&self, _params: CreateThreadParams) -> ThreadStoreResult<()> { Err(not_implemented("create_thread")) } - async fn resume_thread_recorder( - &self, - _params: ResumeThreadRecorderParams, - ) -> ThreadStoreResult> { - 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, diff --git a/codex-rs/thread-store/src/store.rs b/codex-rs/thread-store/src/store.rs index 56cbd05f5..def8bf082 100644 --- a/codex-rs/thread-store/src/store.rs +++ b/codex-rs/thread-store/src/store.rs @@ -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>; + /// 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>; + /// 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, diff --git a/codex-rs/thread-store/src/types.rs b/codex-rs/thread-store/src/types.rs index 1f12bea64..537b09320 100644 --- a/codex-rs/thread-store/src/types.rs +++ b/codex-rs/thread-store/src/types.rs @@ -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, - /// 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, + /// Known replay history for the resumed thread, if already loaded by the caller. + pub history: Option>, /// 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.