From 45288d994d8bbb4d91762b1a779aba65c065e65c Mon Sep 17 00:00:00 2001 From: Owen Lin Date: Tue, 25 Aug 2026 17:49:32 +0000 Subject: [PATCH] Start rollout migration after runtime feature enablement (#40668) ## Why App-server clients can enable `background_paginated_rollout_migration` after the server has started, so the startup-only migration path does not handle that transition. ## What changed - Allow runtime enablement of `background_paginated_rollout_migration`. - Start the local rollout migration when the feature transitions from disabled to enabled at runtime. ## Testing Update the app-server integration test to enable migration through the runtime feature API, wait for paginated history, and verify that cold resume preserves model context. GitOrigin-RevId: d695624a03b4f056be94fcaac2075f91896f19e5 --- .../request_processors/config_processor.rs | 22 ++++++- .../tests/suite/v2/rollout_migration.rs | 63 ++++++++++--------- codex-rs/core/src/thread_manager.rs | 21 +++++++ 3 files changed, 76 insertions(+), 30 deletions(-) diff --git a/codex-rs/app-server/src/request_processors/config_processor.rs b/codex-rs/app-server/src/request_processors/config_processor.rs index 81893abab4..278fc8e050 100644 --- a/codex-rs/app-server/src/request_processors/config_processor.rs +++ b/codex-rs/app-server/src/request_processors/config_processor.rs @@ -50,6 +50,7 @@ use codex_config::MatcherGroup as CoreMatcherGroup; use codex_config::ResidencyRequirement as CoreResidencyRequirement; use codex_config::SandboxModeRequirement as CoreSandboxModeRequirement; use codex_core::ThreadManager; +use codex_features::Feature; use codex_features::canonical_feature_for_key; use codex_features::feature_for_key; use codex_model_provider::create_model_provider; @@ -58,8 +59,12 @@ use codex_protocol::config_types::WebSearchMode; use serde_json::json; use std::path::PathBuf; +const BACKGROUND_PAGINATED_ROLLOUT_MIGRATION_FEATURE: &str = + "background_paginated_rollout_migration"; + const SUPPORTED_EXPERIMENTAL_FEATURE_ENABLEMENT: &[&str] = &[ "auth_elicitation", + BACKGROUND_PAGINATED_ROLLOUT_MIGRATION_FEATURE, "mcp_2026_07_28", "memories", "mentions_v2", @@ -287,6 +292,18 @@ impl ConfigRequestProcessor { return Ok(ExperimentalFeatureEnablementSetResponse { enablement }); } + // Most runtime features are read later from config. Background migration is a one-shot + // process-scoped task, so start it when runtime enablement first changes it to on. + let feature = Feature::BackgroundPaginatedRolloutMigration; + let should_start_background_rollout_migration = enablement + .get(BACKGROUND_PAGINATED_ROLLOUT_MIGRATION_FEATURE) + .is_some_and(|enabled| *enabled) + && !self + .load_latest_config(/*fallback_cwd*/ None) + .await? + .features + .enabled(feature); + self.config_manager .extend_runtime_feature_enablement( enablement @@ -295,7 +312,10 @@ impl ConfigRequestProcessor { ) .map_err(|_| internal_error("failed to update feature enablement"))?; - self.load_latest_config(/*fallback_cwd*/ None).await?; + let config = self.load_latest_config(/*fallback_cwd*/ None).await?; + if should_start_background_rollout_migration && config.features.enabled(feature) { + self.thread_manager.start_background_rollout_migration(); + } Ok(ExperimentalFeatureEnablementSetResponse { enablement }) } diff --git a/codex-rs/app-server/tests/suite/v2/rollout_migration.rs b/codex-rs/app-server/tests/suite/v2/rollout_migration.rs index 431534a175..da20bdbea8 100644 --- a/codex-rs/app-server/tests/suite/v2/rollout_migration.rs +++ b/codex-rs/app-server/tests/suite/v2/rollout_migration.rs @@ -1,19 +1,19 @@ +use std::collections::BTreeMap; + use anyhow::Result; use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; +use codex_app_server_protocol::ExperimentalFeatureEnablementSetParams; +use codex_app_server_protocol::ExperimentalFeatureEnablementSetResponse; use codex_app_server_protocol::ThreadHistoryMode; +use codex_app_server_protocol::ThreadReadParams; +use codex_app_server_protocol::ThreadReadResponse; use codex_app_server_protocol::ThreadResumeParams; use codex_app_server_protocol::ThreadResumeResponse; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; use codex_app_server_protocol::TurnStartParams; use codex_app_server_protocol::UserInput; -use codex_thread_store::LocalThreadStore; -use codex_thread_store::LocalThreadStoreConfig; -use codex_thread_store::RolloutMigrationMode; -use codex_thread_store::RolloutMigrationOptions; -use codex_thread_store::RolloutMigrationStatus; -use codex_utils_absolute_path::test_support::PathExt; use core_test_support::responses; use pretty_assertions::assert_eq; use tempfile::TempDir; @@ -22,7 +22,7 @@ use tokio::time::timeout; const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); #[tokio::test] -async fn migrated_legacy_thread_cold_resume_preserves_model_context() -> Result<()> { +async fn runtime_enabled_legacy_migration_preserves_cold_resume_model_context() -> Result<()> { let server = responses::start_mock_server().await; let response_mock = responses::mount_sse_sequence( &server, @@ -69,32 +69,37 @@ async fn migrated_legacy_thread_cold_resume_preserves_model_context() -> Result< .await??; timeout(DEFAULT_READ_TIMEOUT, primary.shutdown_gracefully()).await??; - let sqlite = codex_state::SqliteConfig::new_for_testing(codex_home.path().abs()); - let state_db = - codex_state::StateRuntime::init(sqlite.clone(), "mock_provider".to_string()).await?; - let store = LocalThreadStore::new( - LocalThreadStoreConfig { - codex_home: codex_home.path().to_path_buf(), - sqlite, - default_model_provider_id: "mock_provider".to_string(), - }, - Some(state_db), - ); - let report = store - .migrate_rollouts(RolloutMigrationOptions { - mode: RolloutMigrationMode::Apply, - max_mib_per_second: Some(1024), - ..RolloutMigrationOptions::default() - }) - .await?; - assert_eq!(report.outcomes.len(), 1); - assert_eq!(report.outcomes[0].status, RolloutMigrationStatus::Migrated); - drop(store); - let mut secondary = TestAppServer::builder() .with_codex_home(codex_home.path()) .build_initialized() .await?; + let enablement_id = secondary + .send_experimental_feature_enablement_set_request(ExperimentalFeatureEnablementSetParams { + enablement: BTreeMap::from([( + "background_paginated_rollout_migration".to_string(), + true, + )]), + }) + .await?; + let _: ExperimentalFeatureEnablementSetResponse = + timeout(DEFAULT_READ_TIMEOUT, secondary.read_response(enablement_id)).await??; + timeout(DEFAULT_READ_TIMEOUT, async { + loop { + let read_id = secondary + .send_thread_read_request(ThreadReadParams { + thread_id: thread.id.clone(), + include_turns: false, + }) + .await?; + let ThreadReadResponse { thread: read } = secondary.read_response(read_id).await?; + if read.history_mode == ThreadHistoryMode::Paginated { + return Ok::<(), anyhow::Error>(()); + } + tokio::task::yield_now().await; + } + }) + .await??; + let resume_id = secondary .send_thread_resume_request(ThreadResumeParams { thread_id: thread.id.clone(), diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 0377f2485e..20f134e31a 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -664,6 +664,27 @@ impl ThreadManager { self.state.environment_manager.clone() } + /// Starts the local rollout migration path after a runtime feature enablement. + /// + /// Startup config handles the initial launch in [`thread_store_from_config`]. This covers + /// clients that decide to enable background migration after constructing the app-server. + pub fn start_background_rollout_migration(&self) { + let Some(store) = self + .state + .thread_store + .as_any() + .downcast_ref::() + else { + return; + }; + let store = store.clone(); + tokio::spawn(async move { + if let Err(err) = store.migrate_rollouts_on_startup().await { + warn!("failed to migrate legacy rollouts on startup: {err}"); + } + }); + } + /// Refreshes every loaded thread and marks threads that are still being created. pub async fn invalidate_mcp_runtimes(&self) { self.invalidate_starting_mcp_runtimes();