Make thread store process-scoped (#19474)

- Build one app-server process ThreadStore from startup config and share it with ThreadManager and CodexMessageProcessor. - Remove per-thread/fork store reconstruction so effective thread config cannot switch the persistence backend. - Add params to ThreadStore create/resume for specifying thread metadata, since otherwise the metadata from store creation would be used (incorrectly).
2026-07-01 00:31:56 +08:00 · 2026-04-30 21:24:59 -07:00
parent f50c02d7bc
commit fe05acad23
55 changed files with 1076 additions and 514 deletions
@@ -28,6 +28,7 @@ use codex_login::AuthManager;
 use codex_login::CodexAuth;
 use codex_model_provider::create_model_provider;
 use codex_model_provider_info::ModelProviderInfo;
+use codex_model_provider_info::OPENAI_PROVIDER_ID;
 use codex_models_manager::manager::RefreshStrategy;
 use codex_models_manager::manager::SharedModelsManager;
 use codex_protocol::ThreadId;
@@ -50,12 +51,15 @@ use codex_protocol::protocol::TurnAbortReason;
 use codex_protocol::protocol::TurnAbortedEvent;
 use codex_protocol::protocol::TurnEnvironmentSelection;
 use codex_protocol::protocol::W3cTraceContext;
-use codex_rollout::RolloutConfig;
 use codex_state::DirectionalThreadSpawnEdgeStatus;
 use codex_thread_store::InMemoryThreadStore;
 use codex_thread_store::LocalThreadStore;
+use codex_thread_store::LocalThreadStoreConfig;
+use codex_thread_store::ReadThreadParams;
 use codex_thread_store::RemoteThreadStore;
+use codex_thread_store::StoredThread;
 use codex_thread_store::ThreadStore;
+use codex_thread_store::ThreadStoreError;
 use codex_utils_absolute_path::AbsolutePathBuf;
 use futures::StreamExt;
 use futures::stream::FuturesUnordered;
@@ -211,7 +215,6 @@ pub struct ThreadManager {

 pub struct StartThreadOptions {
    pub config: Config,
-    pub thread_store: Arc<dyn ThreadStore>,
    pub initial_history: InitialHistory,
    pub session_source: Option<SessionSource>,
    pub dynamic_tools: Vec<codex_protocol::dynamic_tools::DynamicToolSpec>,
@@ -221,10 +224,9 @@ pub struct StartThreadOptions {
    pub environments: Vec<TurnEnvironmentSelection>,
 }

-pub(crate) struct ResumeThreadFromRolloutOptions {
+pub(crate) struct ResumeThreadWithHistoryOptions {
    pub(crate) config: Config,
-    pub(crate) thread_store: Arc<dyn ThreadStore>,
-    pub(crate) rollout_path: PathBuf,
+    pub(crate) initial_history: InitialHistory,
    pub(crate) agent_control: AgentControl,
    pub(crate) session_source: SessionSource,
    pub(crate) inherited_shell_snapshot: Option<Arc<ShellSnapshot>>,
@@ -244,6 +246,7 @@ pub(crate) struct ThreadManagerState {
    plugins_manager: Arc<PluginsManager>,
    mcp_manager: Arc<McpManager>,
    skills_watcher: Arc<SkillsWatcher>,
+    thread_store: Arc<dyn ThreadStore>,
    session_source: SessionSource,
    analytics_events_client: Option<AnalyticsEventsClient>,
    // Captures submitted ops for testing purpose when test mode is enabled.
@@ -263,9 +266,9 @@ pub fn build_models_manager(

 pub fn thread_store_from_config(config: &Config) -> Arc<dyn ThreadStore> {
    match &config.experimental_thread_store {
-        ThreadStoreConfig::Local => {
-            Arc::new(LocalThreadStore::new(RolloutConfig::from_view(config)))
-        }
+        ThreadStoreConfig::Local => Arc::new(LocalThreadStore::new(
+            LocalThreadStoreConfig::from_config(config),
+        )),
        ThreadStoreConfig::Remote { endpoint } => Arc::new(RemoteThreadStore::new(endpoint)),
        ThreadStoreConfig::InMemory { id } => InMemoryThreadStore::for_id(id),
    }
@@ -278,6 +281,7 @@ impl ThreadManager {
        session_source: SessionSource,
        environment_manager: Arc<EnvironmentManager>,
        analytics_events_client: Option<AnalyticsEventsClient>,
+        thread_store: Arc<dyn ThreadStore>,
    ) -> Self {
        let codex_home = config.codex_home.clone();
        let restriction_product = session_source.restriction_product();
@@ -303,6 +307,7 @@ impl ThreadManager {
                plugins_manager,
                mcp_manager,
                skills_watcher,
+                thread_store,
                auth_manager,
                session_source,
                analytics_events_client,
@@ -363,6 +368,14 @@ impl ThreadManager {
            restriction_product,
        ));
        let skills_watcher = build_skills_watcher(Arc::clone(&skills_manager));
+        // This test constructor has no Config input. Tests that need a non-local
+        // process store should construct ThreadManager::new with an explicit store.
+        let thread_store: Arc<dyn ThreadStore> =
+            Arc::new(LocalThreadStore::new(LocalThreadStoreConfig {
+                codex_home: codex_home.clone(),
+                sqlite_home: codex_home.clone(),
+                default_model_provider_id: OPENAI_PROVIDER_ID.to_string(),
+            }));
        Self {
            state: Arc::new(ThreadManagerState {
                threads: Arc::new(RwLock::new(HashMap::new())),
@@ -374,6 +387,7 @@ impl ThreadManager {
                plugins_manager,
                mcp_manager,
                skills_watcher,
+                thread_store,
                auth_manager,
                session_source: SessionSource::Exec,
                analytics_events_client: None,
@@ -517,16 +531,11 @@ impl ThreadManager {
        Ok(subtree_thread_ids)
    }

-    pub async fn start_thread(
-        &self,
-        config: Config,
-        thread_store: Arc<dyn ThreadStore>,
-    ) -> CodexResult<NewThread> {
+    pub async fn start_thread(&self, config: Config) -> CodexResult<NewThread> {
        // Box delegated thread-spawn futures so these convenience wrappers do
        // not inline the full spawn path into every caller's async state.
        Box::pin(self.start_thread_with_tools(
            config,
-            thread_store,
            Vec::new(),
            /*persist_extended_history*/ false,
        ))
@@ -536,7 +545,6 @@ impl ThreadManager {
    pub async fn start_thread_with_tools(
        &self,
        config: Config,
-        thread_store: Arc<dyn ThreadStore>,
        dynamic_tools: Vec<codex_protocol::dynamic_tools::DynamicToolSpec>,
        persist_extended_history: bool,
    ) -> CodexResult<NewThread> {
@@ -546,7 +554,6 @@ impl ThreadManager {
        );
        Box::pin(self.start_thread_with_options(StartThreadOptions {
            config,
-            thread_store,
            initial_history: InitialHistory::New,
            session_source: None,
            dynamic_tools,
@@ -567,7 +574,6 @@ impl ThreadManager {
            .unwrap_or_else(|| self.state.session_source.clone());
        Box::pin(self.state.spawn_thread_with_source(
            options.config,
-            options.thread_store,
            options.initial_history,
            Arc::clone(&self.state.auth_manager),
            self.agent_control(),
@@ -587,7 +593,6 @@ impl ThreadManager {
    pub async fn resume_thread_from_rollout(
        &self,
        config: Config,
-        thread_store: Arc<dyn ThreadStore>,
        rollout_path: PathBuf,
        auth_manager: Arc<AuthManager>,
        parent_trace: Option<W3cTraceContext>,
@@ -595,7 +600,6 @@ impl ThreadManager {
        let initial_history = RolloutRecorder::get_rollout_history(&rollout_path).await?;
        Box::pin(self.resume_thread_with_history(
            config,
-            thread_store,
            initial_history,
            auth_manager,
            /*persist_extended_history*/ false,
@@ -607,7 +611,6 @@ impl ThreadManager {
    pub async fn resume_thread_with_history(
        &self,
        config: Config,
-        thread_store: Arc<dyn ThreadStore>,
        initial_history: InitialHistory,
        auth_manager: Arc<AuthManager>,
        persist_extended_history: bool,
@@ -619,7 +622,6 @@ impl ThreadManager {
        );
        Box::pin(self.state.spawn_thread(
            config,
-            thread_store,
            initial_history,
            auth_manager,
            self.agent_control(),
@@ -636,7 +638,6 @@ impl ThreadManager {
    pub(crate) async fn start_thread_with_user_shell_override_for_tests(
        &self,
        config: Config,
-        thread_store: Arc<dyn ThreadStore>,
        user_shell_override: crate::shell::Shell,
    ) -> CodexResult<NewThread> {
        let environments = default_thread_environment_selections(
@@ -645,7 +646,6 @@ impl ThreadManager {
        );
        Box::pin(self.state.spawn_thread(
            config,
-            thread_store,
            InitialHistory::New,
            Arc::clone(&self.state.auth_manager),
            self.agent_control(),
@@ -662,7 +662,6 @@ impl ThreadManager {
    pub(crate) async fn resume_thread_from_rollout_with_user_shell_override_for_tests(
        &self,
        config: Config,
-        thread_store: Arc<dyn ThreadStore>,
        rollout_path: PathBuf,
        auth_manager: Arc<AuthManager>,
        user_shell_override: crate::shell::Shell,
@@ -674,7 +673,6 @@ impl ThreadManager {
        );
        Box::pin(self.state.spawn_thread(
            config,
-            thread_store,
            initial_history,
            auth_manager,
            self.agent_control(),
@@ -754,7 +752,6 @@ impl ThreadManager {
        &self,
        snapshot: S,
        config: Config,
-        thread_store: Arc<dyn ThreadStore>,
        path: PathBuf,
        persist_extended_history: bool,
        parent_trace: Option<W3cTraceContext>,
@@ -767,7 +764,6 @@ impl ThreadManager {
        self.fork_thread_from_history(
            snapshot,
            config,
-            thread_store,
            history,
            persist_extended_history,
            parent_trace,
@@ -780,7 +776,6 @@ impl ThreadManager {
        &self,
        snapshot: S,
        config: Config,
-        thread_store: Arc<dyn ThreadStore>,
        history: InitialHistory,
        persist_extended_history: bool,
        parent_trace: Option<W3cTraceContext>,
@@ -791,7 +786,6 @@ impl ThreadManager {
        self.fork_thread_with_initial_history(
            snapshot.into(),
            config,
-            thread_store,
            history,
            persist_extended_history,
            parent_trace,
@@ -803,7 +797,6 @@ impl ThreadManager {
        &self,
        snapshot: ForkSnapshot,
        config: Config,
-        thread_store: Arc<dyn ThreadStore>,
        history: InitialHistory,
        persist_extended_history: bool,
        parent_trace: Option<W3cTraceContext>,
@@ -816,7 +809,6 @@ impl ThreadManager {
        );
        Box::pin(self.state.spawn_thread(
            config,
-            thread_store,
            history,
            Arc::clone(&self.state.auth_manager),
            self.agent_control(),
@@ -865,6 +857,31 @@ impl ThreadManagerState {
        }
    }

+    pub(crate) async fn read_stored_thread(
+        &self,
+        params: ReadThreadParams,
+    ) -> CodexResult<StoredThread> {
+        let thread_id = params.thread_id;
+        self.thread_store
+            .read_thread(params)
+            .await
+            .map_err(|err| match err {
+                ThreadStoreError::ThreadNotFound { thread_id } => {
+                    CodexErr::ThreadNotFound(thread_id)
+                }
+                ThreadStoreError::InvalidRequest { message } => {
+                    if message.starts_with("no rollout found for thread id ") {
+                        CodexErr::ThreadNotFound(thread_id)
+                    } else {
+                        CodexErr::Fatal(format!(
+                            "failed to read stored thread {thread_id}: invalid thread-store request: {message}"
+                        ))
+                    }
+                }
+                err => CodexErr::Fatal(format!("failed to read stored thread {thread_id}: {err}")),
+            })
+    }
+
    /// Send an operation to a thread by ID.
    pub(crate) async fn send_op(&self, thread_id: ThreadId, op: Op) -> CodexResult<String> {
        let thread = self.get_thread(thread_id).await?;
@@ -896,12 +913,10 @@ impl ThreadManagerState {
    pub(crate) async fn spawn_new_thread(
        &self,
        config: Config,
-        thread_store: Arc<dyn ThreadStore>,
        agent_control: AgentControl,
    ) -> CodexResult<NewThread> {
        Box::pin(self.spawn_new_thread_with_source(
            config,
-            thread_store,
            agent_control,
            self.session_source.clone(),
            /*persist_extended_history*/ false,
@@ -917,7 +932,6 @@ impl ThreadManagerState {
    pub(crate) async fn spawn_new_thread_with_source(
        &self,
        config: Config,
-        thread_store: Arc<dyn ThreadStore>,
        agent_control: AgentControl,
        session_source: SessionSource,
        persist_extended_history: bool,
@@ -931,7 +945,6 @@ impl ThreadManagerState {
        });
        Box::pin(self.spawn_thread_with_source(
            config,
-            thread_store,
            InitialHistory::New,
            Arc::clone(&self.auth_manager),
            agent_control,
@@ -948,25 +961,22 @@ impl ThreadManagerState {
        .await
    }

-    pub(crate) async fn resume_thread_from_rollout_with_source(
+    pub(crate) async fn resume_thread_with_history_with_source(
        &self,
-        options: ResumeThreadFromRolloutOptions,
+        options: ResumeThreadWithHistoryOptions,
    ) -> CodexResult<NewThread> {
-        let ResumeThreadFromRolloutOptions {
+        let ResumeThreadWithHistoryOptions {
            config,
-            thread_store,
-            rollout_path,
+            initial_history,
            agent_control,
            session_source,
            inherited_shell_snapshot,
            inherited_exec_policy,
        } = options;
-        let initial_history = RolloutRecorder::get_rollout_history(&rollout_path).await?;
        let environments =
            default_thread_environment_selections(self.environment_manager.as_ref(), &config.cwd);
        Box::pin(self.spawn_thread_with_source(
            config,
-            thread_store,
            initial_history,
            Arc::clone(&self.auth_manager),
            agent_control,
@@ -987,7 +997,6 @@ impl ThreadManagerState {
    pub(crate) async fn fork_thread_with_source(
        &self,
        config: Config,
-        thread_store: Arc<dyn ThreadStore>,
        initial_history: InitialHistory,
        agent_control: AgentControl,
        session_source: SessionSource,
@@ -1001,7 +1010,6 @@ impl ThreadManagerState {
        });
        Box::pin(self.spawn_thread_with_source(
            config,
-            thread_store,
            initial_history,
            Arc::clone(&self.auth_manager),
            agent_control,
@@ -1023,7 +1031,6 @@ impl ThreadManagerState {
    pub(crate) async fn spawn_thread(
        &self,
        config: Config,
-        thread_store: Arc<dyn ThreadStore>,
        initial_history: InitialHistory,
        auth_manager: Arc<AuthManager>,
        agent_control: AgentControl,
@@ -1036,7 +1043,6 @@ impl ThreadManagerState {
    ) -> CodexResult<NewThread> {
        Box::pin(self.spawn_thread_with_source(
            config,
-            thread_store,
            initial_history,
            auth_manager,
            agent_control,
@@ -1057,7 +1063,6 @@ impl ThreadManagerState {
    pub(crate) async fn spawn_thread_with_source(
        &self,
        config: Config,
-        thread_store: Arc<dyn ThreadStore>,
        initial_history: InitialHistory,
        auth_manager: Arc<AuthManager>,
        agent_control: AgentControl,
@@ -1072,6 +1077,27 @@ impl ThreadManagerState {
        user_shell_override: Option<crate::shell::Shell>,
    ) -> CodexResult<NewThread> {
        let is_resumed_thread = matches!(&initial_history, InitialHistory::Resumed(_));
+        if let InitialHistory::Resumed(resumed) = &initial_history {
+            let mut threads = self.threads.write().await;
+            if let Some(thread) = threads.get(&resumed.conversation_id).cloned() {
+                if thread.is_running() {
+                    if let Some(requested_rollout_path) = resumed.rollout_path.as_deref()
+                        && thread.rollout_path().as_deref() != Some(requested_rollout_path)
+                    {
+                        return Err(CodexErr::InvalidRequest(format!(
+                            "thread {} is already running with a different rollout path",
+                            resumed.conversation_id
+                        )));
+                    }
+                    return Ok(NewThread {
+                        thread_id: resumed.conversation_id,
+                        session_configured: thread.session_configured(),
+                        thread,
+                    });
+                }
+                threads.remove(&resumed.conversation_id);
+            }
+        }
        let environment =
            selected_primary_environment(self.environment_manager.as_ref(), &environments)?;
        let watch_registration = match environment.as_ref() {
@@ -1115,7 +1141,7 @@ impl ThreadManagerState {
            parent_trace,
            environments,
            analytics_events_client: self.analytics_events_client.clone(),
-            thread_store,
+            thread_store: Arc::clone(&self.thread_store),
        })
        .await?;
        let new_thread = self
@@ -1147,20 +1173,31 @@ impl ThreadManagerState {
            }
        };

-        let thread = Arc::new(CodexThread::new(
-            codex,
-            session_configured.rollout_path.clone(),
-            session_source,
-            watch_registration,
-        ));
-        let mut threads = self.threads.write().await;
-        threads.insert(thread_id, thread.clone());
+        {
+            let mut threads = self.threads.write().await;
+            if let std::collections::hash_map::Entry::Vacant(e) = threads.entry(thread_id) {
+                let thread = Arc::new(CodexThread::new(
+                    codex,
+                    session_configured.clone(),
+                    session_configured.rollout_path.clone(),
+                    session_source,
+                    watch_registration,
+                ));
+                e.insert(thread.clone());
+                return Ok(NewThread {
+                    thread_id,
+                    thread,
+                    session_configured,
+                });
+            }
+        }

-        Ok(NewThread {
-            thread_id,
-            thread,
-            session_configured,
-        })
+        if let Err(err) = codex.shutdown_and_wait().await {
+            warn!("failed to shut down duplicate thread {thread_id}: {err}");
+        }
+        Err(CodexErr::InvalidRequest(format!(
+            "thread {thread_id} is already running"
+        )))
    }

    pub(crate) fn notify_thread_created(&self, thread_id: ThreadId) {