From 3f76e88a480f15258eab3512e7eafca77bbe80ee Mon Sep 17 00:00:00 2001 From: jif Date: Tue, 8 Sep 2026 10:28:09 +0000 Subject: [PATCH] Add configurable memory versions with isolated storage (#43797) ## What changed - Add `memories.version` with `v1` as the default and `v2` as an opt-in selection. - Route memory generation, summaries, and retrieval tools through the selected version. Store `v2` artifacts in `memories_v2` and lazily create a separate SQLite database, keeping jobs and outputs isolated while sharing the thread catalog. - Preserve the initial memory version across extension config updates so summaries and retrieval tools use the same namespace. - Clear both versions on memory reset and remove thread memory from both stores on thread deletion. Recognize both artifact roots in shell usage telemetry. ## Testing Add regression coverage for version defaults and validation, independent job claims and outputs, deletion and reset across versions, and shell usage classification for both memory roots. GitOrigin-RevId: aa799bf87ab6ec10c8f5668213931e57547d080b --- codex-rs/Cargo.lock | 1 + .../request_processors/thread_processor.rs | 10 +- codex-rs/config/src/types.rs | 7 ++ codex-rs/config/src/types_tests.rs | 18 +++ codex-rs/core/config.schema.json | 15 +++ codex-rs/core/src/config/config_tests.rs | 2 + codex-rs/core/src/config/mod.rs | 3 +- codex-rs/ext/memories/Cargo.toml | 1 + codex-rs/ext/memories/src/extension.rs | 19 ++- codex-rs/ext/memories/src/local.rs | 6 - codex-rs/ext/memories/src/prompts.rs | 4 +- codex-rs/ext/memories/src/prompts_tests.rs | 2 +- codex-rs/ext/memories/src/tests.rs | 5 + codex-rs/memories/read/src/usage.rs | 5 + codex-rs/memories/read/src/usage_tests.rs | 22 ++++ codex-rs/memories/write/src/control.rs | 1 + codex-rs/memories/write/src/phase1.rs | 15 +-- codex-rs/memories/write/src/phase2.rs | 97 +++++---------- codex-rs/memories/write/src/runtime.rs | 20 +++- codex-rs/memories/write/src/start.rs | 14 +-- codex-rs/protocol/src/lib.rs | 2 + codex-rs/protocol/src/memory_version.rs | 23 ++++ codex-rs/state/src/runtime.rs | 56 +++++++-- codex-rs/state/src/runtime/memory_versions.rs | 56 +++++++++ .../src/runtime/memory_versions_tests.rs | 113 ++++++++++++++++++ codex-rs/state/src/runtime/threads.rs | 2 +- codex-rs/state/src/sqlite.rs | 22 +++- 27 files changed, 422 insertions(+), 119 deletions(-) create mode 100644 codex-rs/memories/read/src/usage_tests.rs create mode 100644 codex-rs/protocol/src/memory_version.rs create mode 100644 codex-rs/state/src/runtime/memory_versions.rs create mode 100644 codex-rs/state/src/runtime/memory_versions_tests.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index fcdb21be0d..f642b4591b 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3893,6 +3893,7 @@ dependencies = [ "codex-extension-api", "codex-features", "codex-otel", + "codex-protocol", "codex-tools", "codex-utils-absolute-path", "codex-utils-output-truncation", diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index 4442a95440..842ee1bb2a 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -1866,13 +1866,9 @@ impl ThreadRequestProcessor { .clone() .ok_or_else(|| internal_error("sqlite state db unavailable for memory reset"))?; - state_db - .memories() - .clear_memory_data() - .await - .map_err(|err| { - internal_error(format!("failed to clear memory rows in memories db: {err}")) - })?; + state_db.clear_all_memory_data().await.map_err(|err| { + internal_error(format!("failed to clear memory rows in memories db: {err}")) + })?; clear_memory_roots_contents(&self.config.codex_home) .await diff --git a/codex-rs/config/src/types.rs b/codex-rs/config/src/types.rs index dbb76178b2..15f02473e5 100644 --- a/codex-rs/config/src/types.rs +++ b/codex-rs/config/src/types.rs @@ -286,10 +286,14 @@ pub struct ToolSuggestConfig { pub disabled_tools: Vec, } +pub use codex_protocol::MemoryVersion; + /// Memories settings loaded from config.toml. #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Default, JsonSchema)] #[schemars(deny_unknown_fields)] pub struct MemoriesToml { + /// Selects the memory pipeline; v1 remains the default. + pub version: Option, /// When `true`, external context sources mark the thread `memory_mode` as `"polluted"`. #[serde(alias = "no_memories_if_mcp_or_web_search")] pub disable_on_external_context: Option, @@ -323,6 +327,7 @@ pub struct MemoriesToml { /// Effective memories settings after defaults are applied. #[derive(Debug, Clone, PartialEq, Eq, Serialize)] pub struct MemoriesConfig { + pub version: MemoryVersion, pub disable_on_external_context: bool, pub generate_memories: bool, pub use_memories: bool, @@ -340,6 +345,7 @@ pub struct MemoriesConfig { impl Default for MemoriesConfig { fn default() -> Self { Self { + version: MemoryVersion::V1, disable_on_external_context: false, generate_memories: true, use_memories: true, @@ -360,6 +366,7 @@ impl From for MemoriesConfig { fn from(toml: MemoriesToml) -> Self { let defaults = Self::default(); Self { + version: toml.version.unwrap_or(defaults.version), disable_on_external_context: toml .disable_on_external_context .unwrap_or(defaults.disable_on_external_context), diff --git a/codex-rs/config/src/types_tests.rs b/codex-rs/config/src/types_tests.rs index 2c3f69d986..3e1082a07b 100644 --- a/codex-rs/config/src/types_tests.rs +++ b/codex-rs/config/src/types_tests.rs @@ -86,3 +86,21 @@ fn memories_config_clamps_rate_limit_remaining_threshold() { } ); } + +#[test] +fn memories_version_selects_pipeline_without_changing_other_defaults() { + for (source, version) in [ + ("", MemoryVersion::V1), + ("version = \"v2\"", MemoryVersion::V2), + ] { + let parsed: MemoriesToml = toml::from_str(source).expect("parse memories config"); + assert_eq!( + MemoriesConfig::from(parsed), + MemoriesConfig { + version, + ..Default::default() + } + ); + } + assert!(toml::from_str::("version = \"v3\"").is_err()); +} diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 0e16b958be..02688af380 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -2128,10 +2128,25 @@ "use_memories": { "description": "When `false`, skip injecting memory usage instructions into developer prompts.", "type": "boolean" + }, + "version": { + "allOf": [ + { + "$ref": "#/definitions/MemoryVersion" + } + ], + "description": "Selects the memory pipeline; v1 remains the default." } }, "type": "object" }, + "MemoryVersion": { + "enum": [ + "v1", + "v2" + ], + "type": "string" + }, "ModelAvailabilityNuxConfig": { "additionalProperties": { "format": "uint32", diff --git a/codex-rs/core/src/config/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index 5cf603df74..50bea9b56b 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -334,6 +334,7 @@ consolidation_model = "gpt-5.2" toml::from_str::(memories).expect("TOML deserialization should succeed"); assert_eq!( Some(MemoriesToml { + version: None, disable_on_external_context: Some(true), generate_memories: Some(false), use_memories: Some(false), @@ -360,6 +361,7 @@ consolidation_model = "gpt-5.2" assert_eq!( config.memories, MemoriesConfig { + version: codex_protocol::MemoryVersion::V1, disable_on_external_context: true, generate_memories: false, use_memories: false, diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index d223baa950..f555f74584 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -87,7 +87,6 @@ use codex_mcp::McpPluginAttribution; use codex_mcp::McpProtocolMode; use codex_mcp::McpServerRegistration; use codex_mcp::ResolvedMcpCatalog; -use codex_memories_read::memory_root; use codex_model_provider::ProviderCapabilities; use codex_model_provider_info::LEGACY_OLLAMA_CHAT_PROVIDER_ID; use codex_model_provider_info::ModelProviderInfo; @@ -3428,7 +3427,7 @@ impl Config { } let memories_config: MemoriesConfig = cfg.memories.clone().unwrap_or_default().into(); - let memories_root = memory_root(&codex_home); + let memories_root = codex_home.join(memories_config.version.directory_name()); let profiles_are_active = effective_permission_selection.profiles_are_active( default_permissions_override.as_deref(), diff --git a/codex-rs/ext/memories/Cargo.toml b/codex-rs/ext/memories/Cargo.toml index e283140b15..9111571fa4 100644 --- a/codex-rs/ext/memories/Cargo.toml +++ b/codex-rs/ext/memories/Cargo.toml @@ -17,6 +17,7 @@ codex-core = { workspace = true } codex-extension-api = { workspace = true } codex-features = { workspace = true } codex-otel = { workspace = true } +codex-protocol = { workspace = true } codex-tools = { workspace = true } codex-utils-absolute-path = { workspace = true } codex-utils-output-truncation = { workspace = true } diff --git a/codex-rs/ext/memories/src/extension.rs b/codex-rs/ext/memories/src/extension.rs index c383a7e655..c58533112e 100644 --- a/codex-rs/ext/memories/src/extension.rs +++ b/codex-rs/ext/memories/src/extension.rs @@ -13,6 +13,7 @@ use codex_extension_api::ThreadStartInput; use codex_extension_api::ToolContributor; use codex_features::Feature; use codex_otel::MetricsClient; +use codex_protocol::MemoryVersion; use codex_utils_absolute_path::AbsolutePathBuf; use crate::local::LocalMemoriesBackend; @@ -36,6 +37,7 @@ pub(crate) struct MemoriesExtensionConfig { pub(crate) enabled: bool, pub(crate) dedicated_tools: bool, pub(crate) codex_home: AbsolutePathBuf, + pub(crate) version: MemoryVersion, } impl MemoriesExtensionConfig { @@ -44,6 +46,7 @@ impl MemoriesExtensionConfig { enabled: config.features.enabled(Feature::MemoryTool) && config.memories.use_memories, dedicated_tools: config.memories.dedicated_tools, codex_home: config.codex_home.clone(), + version: config.memories.version, } } } @@ -62,7 +65,7 @@ impl ContextContributor for MemoriesExtension { return Vec::new(); } - build_memory_tool_developer_instructions(&config.codex_home) + build_memory_tool_developer_instructions(&config.codex_home, config.version) .await .map(|instructions| { PromptFragment::developer_policy( @@ -97,7 +100,12 @@ impl ConfigContributor for MemoriesExtension { _previous_config: &Config, new_config: &Config, ) { - thread_store.insert(MemoriesExtensionConfig::from_config(new_config)); + let mut config = MemoriesExtensionConfig::from_config(new_config); + if let Some(previous) = thread_store.get::() { + // The initial summary and retrieval tools must use the same namespace. + config.version = previous.version; + } + thread_store.insert(config); } } @@ -117,7 +125,12 @@ impl ToolContributor for MemoriesExtension { } tools::memory_tools( - LocalMemoriesBackend::from_codex_home(&config.codex_home), + LocalMemoriesBackend::from_memory_root( + config + .codex_home + .join(config.version.directory_name()) + .to_path_buf(), + ), self.metrics_client.clone(), ) } diff --git a/codex-rs/ext/memories/src/local.rs b/codex-rs/ext/memories/src/local.rs index f4b6cbf1c4..3fb76d47a4 100644 --- a/codex-rs/ext/memories/src/local.rs +++ b/codex-rs/ext/memories/src/local.rs @@ -2,8 +2,6 @@ use std::path::Component; use std::path::Path; use std::path::PathBuf; -use codex_utils_absolute_path::AbsolutePathBuf; - use crate::backend::AddAdHocMemoryNoteRequest; use crate::backend::AddAdHocMemoryNoteResponse; use crate::backend::ListMemoriesRequest; @@ -27,10 +25,6 @@ pub(crate) struct LocalMemoriesBackend { } impl LocalMemoriesBackend { - pub(crate) fn from_codex_home(codex_home: &AbsolutePathBuf) -> Self { - Self::from_memory_root(codex_home.join("memories").to_path_buf()) - } - pub(crate) fn from_memory_root(root: impl Into) -> Self { Self { root: root.into() } } diff --git a/codex-rs/ext/memories/src/prompts.rs b/codex-rs/ext/memories/src/prompts.rs index 6546206b1c..e307934d6a 100644 --- a/codex-rs/ext/memories/src/prompts.rs +++ b/codex-rs/ext/memories/src/prompts.rs @@ -1,4 +1,5 @@ use crate::MEMORY_TOOL_DEVELOPER_INSTRUCTIONS_SUMMARY_TOKEN_LIMIT; +use codex_protocol::MemoryVersion; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_output_truncation::TruncationPolicy; use codex_utils_output_truncation::truncate_text; @@ -26,8 +27,9 @@ fn parse_embedded_template(source: &'static str, template_name: &str) -> Templat /// [MEMORY_TOOL_DEVELOPER_INSTRUCTIONS_SUMMARY_TOKEN_LIMIT]. pub(crate) async fn build_memory_tool_developer_instructions( codex_home: &AbsolutePathBuf, + version: MemoryVersion, ) -> Option { - let base_path = codex_home.join("memories"); + let base_path = codex_home.join(version.directory_name()); let memory_summary_path = base_path.join("memory_summary.md"); let memory_summary = fs::read_to_string(&memory_summary_path) .await diff --git a/codex-rs/ext/memories/src/prompts_tests.rs b/codex-rs/ext/memories/src/prompts_tests.rs index b4a8cef134..7a9c49f50c 100644 --- a/codex-rs/ext/memories/src/prompts_tests.rs +++ b/codex-rs/ext/memories/src/prompts_tests.rs @@ -17,7 +17,7 @@ async fn build_memory_tool_developer_instructions_renders_embedded_template() { .await .unwrap(); - let instructions = build_memory_tool_developer_instructions(&codex_home) + let instructions = build_memory_tool_developer_instructions(&codex_home, MemoryVersion::V1) .await .unwrap(); diff --git a/codex-rs/ext/memories/src/tests.rs b/codex-rs/ext/memories/src/tests.rs index 0dc714b696..a4f780d3d4 100644 --- a/codex-rs/ext/memories/src/tests.rs +++ b/codex-rs/ext/memories/src/tests.rs @@ -60,6 +60,7 @@ fn tools_are_not_contributed_when_disabled() { let extension = MemoriesExtension::default(); let thread_store = ExtensionData::new("thread"); thread_store.insert(MemoriesExtensionConfig { + version: codex_protocol::MemoryVersion::V1, enabled: false, dedicated_tools: true, codex_home: test_path_buf("/tmp/codex-home").abs(), @@ -77,6 +78,7 @@ fn tools_are_not_contributed_when_dedicated_tools_disabled() { let extension = MemoriesExtension::default(); let thread_store = ExtensionData::new("thread"); thread_store.insert(MemoriesExtensionConfig { + version: codex_protocol::MemoryVersion::V1, enabled: true, dedicated_tools: false, codex_home: test_path_buf("/tmp/codex-home").abs(), @@ -94,6 +96,7 @@ fn tools_are_contributed_when_enabled_with_dedicated_tools() { let extension = MemoriesExtension::default(); let thread_store = ExtensionData::new("thread"); thread_store.insert(MemoriesExtensionConfig { + version: codex_protocol::MemoryVersion::V1, enabled: true, dedicated_tools: true, codex_home: test_path_buf("/tmp/codex-home").abs(), @@ -123,6 +126,7 @@ fn install_registers_dedicated_tool_contributor() { let registry = builder.build(); let thread_store = ExtensionData::new("thread"); thread_store.insert(MemoriesExtensionConfig { + version: codex_protocol::MemoryVersion::V1, enabled: true, dedicated_tools: true, codex_home: test_path_buf("/tmp/codex-home").abs(), @@ -183,6 +187,7 @@ async fn prompt_contribution_uses_memory_summary_when_enabled() { let extension = MemoriesExtension::default(); let thread_store = ExtensionData::new("thread"); thread_store.insert(MemoriesExtensionConfig { + version: codex_protocol::MemoryVersion::V1, enabled: true, dedicated_tools: false, codex_home: tempdir.path().abs(), diff --git a/codex-rs/memories/read/src/usage.rs b/codex-rs/memories/read/src/usage.rs index cb9503efca..3cac16fb26 100644 --- a/codex-rs/memories/read/src/usage.rs +++ b/codex-rs/memories/read/src/usage.rs @@ -44,6 +44,7 @@ pub fn memories_usage_kinds_from_command(command: &str) -> Vec Option { + let path = path.replace("memories_v2/", "memories/"); if path.contains("memories/MEMORY.md") { Some(MemoriesUsageKind::MemoryMd) } else if path.contains("memories/memory_summary.md") { @@ -58,3 +59,7 @@ fn get_memory_kind(path: String) -> Option { None } } + +#[cfg(test)] +#[path = "usage_tests.rs"] +mod tests; diff --git a/codex-rs/memories/read/src/usage_tests.rs b/codex-rs/memories/read/src/usage_tests.rs new file mode 100644 index 0000000000..32e487ea8e --- /dev/null +++ b/codex-rs/memories/read/src/usage_tests.rs @@ -0,0 +1,22 @@ +//! Regression coverage for versioned memory paths in shell usage telemetry. + +use super::MemoriesUsageKind; +use super::memories_usage_kinds_from_command; +use pretty_assertions::assert_eq; + +#[test] +fn classifies_reads_in_both_memory_versions() { + for root in ["memories", "memories_v2"] { + let command = format!( + "cat /tmp/.codex/{root}/memory_summary.md && cat /tmp/.codex/{root}/rollout_summaries/thread.md" + ); + assert_eq!( + memories_usage_kinds_from_command(&command), + vec![ + MemoriesUsageKind::MemorySummary, + MemoriesUsageKind::RolloutSummaries, + ], + "memory root: {root}" + ); + } +} diff --git a/codex-rs/memories/write/src/control.rs b/codex-rs/memories/write/src/control.rs index bc5d2dc90e..9d0da1a0b7 100644 --- a/codex-rs/memories/write/src/control.rs +++ b/codex-rs/memories/write/src/control.rs @@ -3,6 +3,7 @@ use std::path::Path; pub async fn clear_memory_roots_contents(codex_home: &Path) -> std::io::Result<()> { for memory_root in [ codex_home.join("memories"), + codex_home.join("memories_v2"), codex_home.join("memories_extensions"), ] { clear_memory_root_contents(memory_root.as_path()).await?; diff --git a/codex-rs/memories/write/src/phase1.rs b/codex-rs/memories/write/src/phase1.rs index 92bd4ad782..9f7d7ff1ea 100644 --- a/codex-rs/memories/write/src/phase1.rs +++ b/codex-rs/memories/write/src/phase1.rs @@ -110,10 +110,9 @@ pub async fn run(context: Arc, config: Arc) { /// Prune old un-used "dead" raw memories. pub async fn prune(context: &MemoryStartupContext, config: &Config) { - if let Some(db) = context.state_db() { + if let Some(db) = context.memory_store().await { let max_unused_days = config.memories.max_unused_days; match db - .memories() .prune_stage1_outputs_for_retention(max_unused_days, crate::stage_one::PRUNE_BATCH_SIZE) .await { @@ -151,7 +150,7 @@ async fn claim_startup_jobs( context: &MemoryStartupContext, memories_config: &MemoriesConfig, ) -> Option> { - let Some(state_db) = context.state_db() else { + let Some(state_db) = context.memory_store().await else { // This should not happen. warn!("state db unavailable while claiming phase-1 startup jobs; skipping"); return None; @@ -163,7 +162,6 @@ async fn claim_startup_jobs( .collect::>(); match state_db - .memories() .claim_stage1_jobs_for_startup( context.thread_id(), codex_state::Stage1StartupClaimParams { @@ -335,9 +333,8 @@ mod job { reason: &str, ) { tracing::warn!("Phase 1 job failed for thread {thread_id}: {reason}"); - if let Some(state_db) = context.state_db() { + if let Some(state_db) = context.memory_store().await { let _ = state_db - .memories() .mark_stage1_job_failed( thread_id, ownership_token, @@ -353,12 +350,11 @@ mod job { thread_id: codex_protocol::ThreadId, ownership_token: &str, ) -> JobOutcome { - let Some(state_db) = context.state_db() else { + let Some(state_db) = context.memory_store().await else { return JobOutcome::Failed; }; if state_db - .memories() .mark_stage1_job_succeeded_no_output(thread_id, ownership_token) .await .unwrap_or(false) @@ -378,12 +374,11 @@ mod job { rollout_summary: &str, rollout_slug: Option<&str>, ) -> JobOutcome { - let Some(state_db) = context.state_db() else { + let Some(state_db) = context.memory_store().await else { return JobOutcome::Failed; }; if state_db - .memories() .mark_stage1_job_succeeded( thread_id, ownership_token, diff --git a/codex-rs/memories/write/src/phase2.rs b/codex-rs/memories/write/src/phase2.rs index a69190b83b..8b8eb5dd5e 100644 --- a/codex-rs/memories/write/src/phase2.rs +++ b/codex-rs/memories/write/src/phase2.rs @@ -1,5 +1,4 @@ use crate::build_consolidation_prompt; -use crate::memory_root; use crate::metrics::MEMORY_PHASE_TWO_E2E_MS; use crate::metrics::MEMORY_PHASE_TWO_INPUT; use crate::metrics::MEMORY_PHASE_TWO_JOBS; @@ -26,8 +25,8 @@ use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::TokenUsage; use codex_protocol::user_input::UserInput; +use codex_state::MemoryStore; use codex_state::Stage1Output; -use codex_state::StateRuntime; use std::collections::HashMap; use std::path::Path; use std::sync::Arc; @@ -53,16 +52,18 @@ pub async fn run( ) { let phase_two_e2e_timer = context.start_timer(MEMORY_PHASE_TWO_E2E_MS); - let Some(db) = context.state_db() else { + let Some(db) = context.memory_store().await else { // This should not happen. return; }; - let root = memory_root(&config.codex_home); + let root = config + .codex_home + .join(config.memories.version.directory_name()); let max_raw_memories = config.memories.max_raw_memories_for_consolidation; let max_unused_days = config.memories.max_unused_days; // 1. Claim the global Phase 2 lock before touching the memory workspace. - let claim = match job::claim(context.as_ref(), db.as_ref()).await { + let claim = match job::claim(context.as_ref(), &db).await { Ok(claim) => claim, Err(e) => { context.counter(MEMORY_PHASE_TWO_JOBS, /*inc*/ 1, &[("status", e)]); @@ -73,13 +74,7 @@ pub async fn run( // 2. Ensure the memories root has a git baseline repository. if let Err(err) = prepare_memory_workspace(&root).await { tracing::error!("failed preparing memory workspace: {err}"); - job::failed( - context.as_ref(), - db.as_ref(), - &claim, - "failed_prepare_workspace", - ) - .await; + job::failed(context.as_ref(), &db, &claim, "failed_prepare_workspace").await; return; } @@ -91,32 +86,19 @@ pub async fn run( ) else { // If we can't get the config, we can't consolidate. tracing::error!("failed to get agent config"); - job::failed( - context.as_ref(), - db.as_ref(), - &claim, - "failed_sandbox_policy", - ) - .await; + job::failed(context.as_ref(), &db, &claim, "failed_sandbox_policy").await; return; }; // 4. Load current DB-backed Phase 2 inputs. let raw_memories = match db - .memories() .get_phase2_input_selection(max_raw_memories, max_unused_days) .await { Ok(raw_memories) => raw_memories, Err(err) => { tracing::error!("failed to list stage1 outputs from global: {err}"); - job::failed( - context.as_ref(), - db.as_ref(), - &claim, - "failed_load_stage1_outputs", - ) - .await; + job::failed(context.as_ref(), &db, &claim, "failed_load_stage1_outputs").await; return; } }; @@ -128,7 +110,7 @@ pub async fn run( tracing::error!("failed syncing phase2 workspace inputs: {err}"); job::failed( context.as_ref(), - db.as_ref(), + &db, &claim, "failed_sync_workspace_inputs", ) @@ -141,13 +123,7 @@ pub async fn run( Ok(diff) => diff, Err(err) => { tracing::error!("failed checking memory workspace changes: {err}"); - job::failed( - context.as_ref(), - db.as_ref(), - &claim, - "failed_workspace_status", - ) - .await; + job::failed(context.as_ref(), &db, &claim, "failed_workspace_status").await; return; } }; @@ -156,7 +132,7 @@ pub async fn run( // We check only after sync of the file system. job::succeed( context.as_ref(), - db.as_ref(), + &db, &claim, new_watermark, &raw_memories, @@ -169,13 +145,7 @@ pub async fn run( // 7. Persist the diff for the consolidation agent to inspect. if let Err(err) = write_workspace_diff(&root, &workspace_diff).await { tracing::error!("failed writing memory workspace diff file: {err}"); - job::failed( - context.as_ref(), - db.as_ref(), - &claim, - "failed_workspace_diff_file", - ) - .await; + job::failed(context.as_ref(), &db, &claim, "failed_workspace_diff_file").await; return; } @@ -188,7 +158,7 @@ pub async fn run( Ok(agent) => agent, Err(err) => { tracing::error!("failed to spawn global memory consolidation agent: {err}"); - job::failed(context.as_ref(), db.as_ref(), &claim, "failed_spawn_agent").await; + job::failed(context.as_ref(), &db, &claim, "failed_spawn_agent").await; return; } }; @@ -227,10 +197,9 @@ mod job { pub(super) async fn claim( context: &MemoryStartupContext, - db: &StateRuntime, + db: &MemoryStore, ) -> Result { let claim = db - .memories() .try_claim_global_phase2_job(context.thread_id(), crate::stage_two::JOB_LEASE_SECONDS) .await .map_err(|e| { @@ -263,23 +232,21 @@ mod job { pub(super) async fn failed( context: &MemoryStartupContext, - db: &StateRuntime, + db: &MemoryStore, claim: &Claim, reason: &'static str, ) { context.counter(MEMORY_PHASE_TWO_JOBS, /*inc*/ 1, &[("status", reason)]); if matches!( - db.memories() - .mark_global_phase2_job_failed( - &claim.token, - reason, - crate::stage_two::JOB_RETRY_DELAY_SECONDS, - ) - .await, + db.mark_global_phase2_job_failed( + &claim.token, + reason, + crate::stage_two::JOB_RETRY_DELAY_SECONDS, + ) + .await, Ok(false) ) { let _ = db - .memories() .mark_global_phase2_job_failed_if_unowned( &claim.token, reason, @@ -291,15 +258,14 @@ mod job { pub(super) async fn succeed( context: &MemoryStartupContext, - db: &StateRuntime, + db: &MemoryStore, claim: &Claim, completion_watermark: i64, selected_outputs: &[codex_state::Stage1Output], reason: &'static str, ) -> bool { context.counter(MEMORY_PHASE_TWO_JOBS, /*inc*/ 1, &[("status", reason)]); - db.memories() - .mark_global_phase2_job_succeeded(&claim.token, completion_watermark, selected_outputs) + db.mark_global_phase2_job_succeeded(&claim.token, completion_watermark, selected_outputs) .await .unwrap_or(false) } @@ -314,7 +280,9 @@ mod agent { parent_permission_profile: PermissionProfile, provider: &dyn ModelProvider, ) -> Option { - let root = memory_root(&config.codex_home); + let root = config + .codex_home + .join(config.memories.version.directory_name()); let mut agent_config = config.clone(); agent_config.cwd = root.clone(); @@ -388,11 +356,10 @@ mod agent { agent: SpawnedConsolidationAgent, phase_two_e2e_timer: Option, ) { - let Some(db) = context.state_db() else { - return; - }; - tokio::spawn(async move { + let Some(db) = context.memory_store().await else { + return; + }; let _phase_two_e2e_timer = phase_two_e2e_timer; let SpawnedConsolidationAgent { thread_id, thread } = agent; @@ -437,7 +404,6 @@ mod agent { if agent_completed && artifacts_valid { // Do not reset the workspace baseline if we lost the lock. let still_owns_lock = match db - .memories() .heartbeat_global_phase2_job( &claim.token, crate::stage_two::JOB_LEASE_SECONDS, @@ -490,7 +456,7 @@ mod agent { } async fn loop_agent( - db: Arc, + db: MemoryStore, token: String, thread_id: ThreadId, thread: &codex_core::CodexThread, @@ -526,7 +492,6 @@ mod agent { } _ = heartbeat_interval.tick() => { match db - .memories() .heartbeat_global_phase2_job( &token, crate::stage_two::JOB_LEASE_SECONDS, diff --git a/codex-rs/memories/write/src/runtime.rs b/codex-rs/memories/write/src/runtime.rs index c4cf8c7d75..d35c5fbc78 100644 --- a/codex-rs/memories/write/src/runtime.rs +++ b/codex-rs/memories/write/src/runtime.rs @@ -22,6 +22,7 @@ use codex_model_provider::SharedModelProvider; use codex_model_provider::create_model_provider; use codex_otel::SessionTelemetry; use codex_otel::TelemetryAuthMode; +use codex_protocol::MemoryVersion; use codex_protocol::SessionId; use codex_protocol::ThreadId; use codex_protocol::config_types::ReasoningSummary; @@ -33,7 +34,7 @@ use codex_protocol::protocol::ThreadSource; use codex_protocol::protocol::TokenUsage; use codex_protocol::user_input::UserInput; use codex_rollout_trace::InferenceTraceContext; -use codex_state::StateRuntime; +use codex_state::MemoryStore; use codex_terminal_detection::user_agent; use futures::StreamExt; use std::sync::Arc; @@ -68,6 +69,7 @@ impl StageOneRequestContext { } pub(crate) struct MemoryStartupContext { + version: MemoryVersion, thread_id: ThreadId, thread: Arc, thread_manager: Arc, @@ -173,6 +175,7 @@ impl MemoryStartupContext { ); Self { + version: config.memories.version, thread_id, thread, thread_manager, @@ -186,8 +189,19 @@ impl MemoryStartupContext { self.thread_id } - pub(crate) fn state_db(&self) -> Option> { - self.thread.state_db() + pub(crate) async fn memory_store(&self) -> Option { + match self + .thread + .state_db()? + .memories_for_version(self.version) + .await + { + Ok(store) => Some(store), + Err(err) => { + tracing::warn!("failed opening memory store: {err}"); + None + } + } } pub(crate) fn provider(&self) -> &dyn ModelProvider { diff --git a/codex-rs/memories/write/src/start.rs b/codex-rs/memories/write/src/start.rs index d7c964bd23..8a766d86c0 100644 --- a/codex-rs/memories/write/src/start.rs +++ b/codex-rs/memories/write/src/start.rs @@ -1,7 +1,6 @@ use crate::ensure_layout; use crate::extensions::seed_extension_instructions; use crate::guard; -use crate::memory_root; use crate::metrics::MEMORY_STARTUP; use crate::phase1; use crate::phase2; @@ -46,13 +45,14 @@ pub fn start_memories_startup_task( source.clone(), )); - if context.state_db().is_none() { - warn!("state db unavailable for memories startup pipeline; skipping"); - return; - } - tokio::spawn(async move { - let root = memory_root(&config.codex_home); + if context.memory_store().await.is_none() { + warn!("state db unavailable for memories startup pipeline; skipping"); + return; + } + let root = config + .codex_home + .join(config.memories.version.directory_name()); if let Err(err) = ensure_layout(&root).await { warn!("failed preparing memories root: {err}"); return; diff --git a/codex-rs/protocol/src/lib.rs b/codex-rs/protocol/src/lib.rs index 10af9c7317..3c72be065f 100644 --- a/codex-rs/protocol/src/lib.rs +++ b/codex-rs/protocol/src/lib.rs @@ -32,6 +32,8 @@ pub mod mcp; pub mod mcp_approval_meta; pub mod mcp_policy; pub mod memory_citation; +mod memory_version; +pub use memory_version::MemoryVersion; pub mod models; pub mod network_policy; pub mod num_format; diff --git a/codex-rs/protocol/src/memory_version.rs b/codex-rs/protocol/src/memory_version.rs new file mode 100644 index 0000000000..3968bd199a --- /dev/null +++ b/codex-rs/protocol/src/memory_version.rs @@ -0,0 +1,23 @@ +//! Selects a coherent memory pipeline and its generated-artifact namespace. + +use schemars::JsonSchema; +use serde::Deserialize; +use serde::Serialize; + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] +#[serde(rename_all = "lowercase")] +pub enum MemoryVersion { + #[default] + V1, + V2, +} + +impl MemoryVersion { + /// Sibling roots keep v1 cleanup and rollback independent of v2 artifacts. + pub fn directory_name(self) -> &'static str { + match self { + Self::V1 => "memories", + Self::V2 => "memories_v2", + } + } +} diff --git a/codex-rs/state/src/runtime.rs b/codex-rs/state/src/runtime.rs index e7fa3cffef..1c571c0104 100644 --- a/codex-rs/state/src/runtime.rs +++ b/codex-rs/state/src/runtime.rs @@ -44,6 +44,7 @@ mod external_agent_config_imports; mod goals; mod logs; mod memories; +mod memory_versions; mod projects; mod queued_items; mod recovery; @@ -92,6 +93,7 @@ pub struct StateRuntime { logs_pool: Arc, thread_goals: GoalStore, memories: MemoryStore, + memories_v2: Arc>, thread_queue: SqliteQueueStore, thread_updated_at_millis: Arc, thread_recency_at_millis: Arc, @@ -133,6 +135,7 @@ impl StateRuntime { let goals_path = sqlite.goals_db_path(); let memories_path = sqlite.memories_db_path(); let queue_path = sqlite.queue_db_path(); + let has_memories_v2 = tokio::fs::try_exists(sqlite.memories_v2_db_path()).await?; let pool = match sqlite .open_state_db(&state_migrator, telemetry_override) .await @@ -251,6 +254,7 @@ impl StateRuntime { let runtime = Arc::new(Self { thread_goals: GoalStore::new(Arc::clone(&goals_pool)), memories: MemoryStore::new(Arc::clone(&memories_pool), Arc::clone(&pool)), + memories_v2: Arc::new(tokio::sync::OnceCell::new()), thread_queue: SqliteQueueStore::new(queue_pool), pool, logs_pool, @@ -259,6 +263,16 @@ impl StateRuntime { thread_updated_at_millis: Arc::new(AtomicI64::new(thread_updated_at_millis)), thread_recency_at_millis: Arc::new(AtomicI64::new(thread_recency_at_millis)), }); + // Existing v2 state must participate in startup corruption recovery. + // Keep creation lazy for users who have never used v2. + if has_memories_v2 + && let Err(err) = runtime + .memories_for_version(codex_protocol::MemoryVersion::V2) + .await + { + runtime.close().await; + return Err(err); + } if let Err(err) = runtime.run_logs_startup_maintenance().await { warn!( "failed to run startup maintenance for logs db at {}: {err}", @@ -290,24 +304,44 @@ impl StateRuntime { pub async fn close(&self) { self.thread_queue.close().await; self.memories.close().await; + if let Some(memories) = self.memories_v2.get() { + memories.close().await; + } self.thread_goals.close().await; self.logs_pool.close().await; self.pool.close().await; } pub async fn clear_memory_data_in_sqlite_home(sqlite: &SqliteConfig) -> anyhow::Result { - let memories_path = sqlite.memories_db_path(); - if !tokio::fs::try_exists(&memories_path).await? { - return Ok(false); + let mut cleared = false; + for version in [ + codex_protocol::MemoryVersion::V1, + codex_protocol::MemoryVersion::V2, + ] { + let path = match version { + codex_protocol::MemoryVersion::V1 => sqlite.memories_db_path(), + codex_protocol::MemoryVersion::V2 => sqlite.memories_v2_db_path(), + }; + if !tokio::fs::try_exists(path).await? { + continue; + } + let pool = match version { + codex_protocol::MemoryVersion::V1 => { + sqlite + .open_memories_db( + &runtime_memories_migrator(), + /*telemetry_override*/ None, + ) + .await? + } + codex_protocol::MemoryVersion::V2 => sqlite.open_memories_v2_db().await?, + }; + let result = memories::clear_memory_data_in_pool(&pool).await; + pool.close().await; + result?; + cleared = true; } - - let memories_migrator = runtime_memories_migrator(); - let pool = sqlite - .open_memories_db(&memories_migrator, /*telemetry_override*/ None) - .await?; - memories::clear_memory_data_in_pool(&pool).await?; - pool.close().await; - Ok(true) + Ok(cleared) } } diff --git a/codex-rs/state/src/runtime/memory_versions.rs b/codex-rs/state/src/runtime/memory_versions.rs new file mode 100644 index 0000000000..ac7530397e --- /dev/null +++ b/codex-rs/state/src/runtime/memory_versions.rs @@ -0,0 +1,56 @@ +//! Creates isolated v2 memory state on demand while sharing the source thread catalog. +//! Reset and thread deletion always cover every existing version. + +use super::MemoryStore; +use super::StateRuntime; +use codex_protocol::MemoryVersion; +use codex_protocol::ThreadId; +use std::sync::Arc; + +impl StateRuntime { + pub async fn memories_for_version( + &self, + version: MemoryVersion, + ) -> anyhow::Result { + match version { + MemoryVersion::V1 => Ok(self.memories.clone()), + MemoryVersion::V2 => self + .memories_v2 + .get_or_try_init(|| async { + let pool = self.sqlite.open_memories_v2_db().await?; + Ok(MemoryStore::new(Arc::new(pool), Arc::clone(&self.pool))) + }) + .await + .cloned(), + } + } + + pub async fn clear_all_memory_data(&self) -> anyhow::Result<()> { + self.memories.clear_memory_data().await?; + if tokio::fs::try_exists(self.sqlite.memories_v2_db_path()).await? { + self.memories_for_version(MemoryVersion::V2) + .await? + .clear_memory_data() + .await?; + } + Ok(()) + } + + pub(super) async fn delete_versioned_thread_memory( + &self, + thread_id: ThreadId, + ) -> anyhow::Result<()> { + self.memories.delete_thread_memory(thread_id).await?; + if tokio::fs::try_exists(self.sqlite.memories_v2_db_path()).await? { + self.memories_for_version(MemoryVersion::V2) + .await? + .delete_thread_memory(thread_id) + .await?; + } + Ok(()) + } +} + +#[cfg(test)] +#[path = "memory_versions_tests.rs"] +mod tests; diff --git a/codex-rs/state/src/runtime/memory_versions_tests.rs b/codex-rs/state/src/runtime/memory_versions_tests.rs new file mode 100644 index 0000000000..767661f61b --- /dev/null +++ b/codex-rs/state/src/runtime/memory_versions_tests.rs @@ -0,0 +1,113 @@ +use super::*; +use crate::SqliteConfig; +use crate::Stage1JobClaimOutcome; +use crate::runtime::test_support::test_thread_metadata; +use codex_utils_absolute_path::test_support::PathExt; +use pretty_assertions::assert_eq; + +#[tokio::test] +async fn versions_isolate_jobs_outputs_and_reset_without_losing_threads() -> anyhow::Result<()> { + let home = crate::runtime::test_support::unique_temp_dir(); + let sqlite = SqliteConfig::new_for_testing(home.as_path().abs()); + let db = StateRuntime::init(sqlite.clone(), "test-provider".to_string()).await?; + assert!(!sqlite.memories_v2_db_path().exists()); + let thread_id = ThreadId::new(); + let metadata = test_thread_metadata(home.as_path(), thread_id, home.as_path().join("project")); + db.upsert_thread(&metadata).await?; + + for version in [MemoryVersion::V1, MemoryVersion::V2] { + let store = db.memories_for_version(version).await?; + let Stage1JobClaimOutcome::Claimed { ownership_token } = store + .try_claim_stage1_job( + thread_id, + thread_id, + metadata.updated_at.timestamp(), + /*lease_seconds*/ 60, + /*max_running_jobs*/ 1, + ) + .await? + else { + panic!("each version must independently claim the same source") + }; + assert!( + store + .mark_stage1_job_succeeded( + thread_id, + &ownership_token, + metadata.updated_at.timestamp(), + "raw", + version.directory_name(), + /*rollout_slug*/ None + ) + .await? + ); + } + for version in [MemoryVersion::V1, MemoryVersion::V2] { + let outputs = db + .memories_for_version(version) + .await? + .list_stage1_outputs_for_global(/*n*/ 10) + .await?; + assert_eq!( + outputs + .iter() + .map(|output| output.rollout_summary.as_str()) + .collect::>(), + vec![version.directory_name()] + ); + } + // Reopening under v1 must still clear existing v2 state. + db.close().await; + let db = StateRuntime::init(sqlite.clone(), "test-provider".to_string()).await?; + db.delete_thread(thread_id).await?; + for version in [MemoryVersion::V1, MemoryVersion::V2] { + assert!( + db.memories_for_version(version) + .await? + .list_stage1_outputs_for_global(/*n*/ 10) + .await? + .is_empty() + ); + } + db.upsert_thread(&metadata).await?; + for version in [MemoryVersion::V1, MemoryVersion::V2] { + let store = db.memories_for_version(version).await?; + let Stage1JobClaimOutcome::Claimed { ownership_token } = store + .try_claim_stage1_job( + thread_id, + thread_id, + metadata.updated_at.timestamp(), + /*lease_seconds*/ 60, + /*max_running_jobs*/ 1, + ) + .await? + else { + panic!("deletion must clear the job watermark") + }; + store + .mark_stage1_job_succeeded( + thread_id, + &ownership_token, + metadata.updated_at.timestamp(), + "raw", + "summary", + /*rollout_slug*/ None, + ) + .await?; + } + db.clear_all_memory_data().await?; + for version in [MemoryVersion::V1, MemoryVersion::V2] { + assert!( + db.memories_for_version(version) + .await? + .list_stage1_outputs_for_global(/*n*/ 10) + .await? + .is_empty() + ); + } + assert!(db.get_thread(thread_id).await?.is_some()); + db.close().await; + assert!(StateRuntime::clear_memory_data_in_sqlite_home(&sqlite).await?); + tokio::fs::remove_dir_all(home).await?; + Ok(()) +} diff --git a/codex-rs/state/src/runtime/threads.rs b/codex-rs/state/src/runtime/threads.rs index 99838ed644..d1886bfee2 100644 --- a/codex-rs/state/src/runtime/threads.rs +++ b/codex-rs/state/src/runtime/threads.rs @@ -1155,7 +1155,7 @@ ON CONFLICT(id) DO UPDATE SET .execute(self.logs_pool.as_ref()) .await?; self.thread_queue.delete_thread_queue(*thread_id).await?; - self.memories.delete_thread_memory(*thread_id).await?; + self.delete_versioned_thread_memory(*thread_id).await?; self.thread_goals.delete_thread_goal(*thread_id).await?; } diff --git a/codex-rs/state/src/sqlite.rs b/codex-rs/state/src/sqlite.rs index 536d56befa..f05caf1f7e 100644 --- a/codex-rs/state/src/sqlite.rs +++ b/codex-rs/state/src/sqlite.rs @@ -80,6 +80,12 @@ const MEMORIES_DB: RuntimeDbSpec = RuntimeDbSpec { migrate_phase: "migrate_memories", }; +const MEMORIES_V2_DB: RuntimeDbSpec = RuntimeDbSpec { + label: "memories v2 DB", + filename: "memories_v2_1.sqlite", + ..MEMORIES_DB +}; + const QUEUE_DB: RuntimeDbSpec = RuntimeDbSpec { label: "queue DB", filename: QUEUE_DB_FILENAME, @@ -96,11 +102,12 @@ const THREAD_HISTORY_DB: RuntimeDbSpec = RuntimeDbSpec { migrate_phase: "migrate_thread_history", }; -const RUNTIME_DBS: [RuntimeDbSpec; 6] = [ +const RUNTIME_DBS: [RuntimeDbSpec; 7] = [ STATE_DB, LOGS_DB, GOALS_DB, MEMORIES_DB, + MEMORIES_V2_DB, QUEUE_DB, THREAD_HISTORY_DB, ]; @@ -150,6 +157,19 @@ impl SqliteConfig { MEMORIES_DB.path(self.home()) } + pub(crate) fn memories_v2_db_path(&self) -> PathBuf { + MEMORIES_V2_DB.path(self.home()) + } + + pub(crate) async fn open_memories_v2_db(&self) -> anyhow::Result { + self.open_runtime_db( + MEMORIES_V2_DB, + &crate::migrations::runtime_memories_migrator(), + /*telemetry_override*/ None, + ) + .await + } + /// Return the path to the durable user-message queue database. pub fn queue_db_path(&self) -> PathBuf { QUEUE_DB.path(self.home())