github/codex
4
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-20 08:19:50 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add configurable memory versions with isolated storage (#43797)

Browse Source 
## 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
This commit is contained in:
jif
2026-09-08 10:55:48 +00:00
committed by copyberry
parent 35d9e4bc4d
commit 3f76e88a48
27 changed files with 422 additions and 119 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/Cargo.lock
Generated
+1
View File
@@ -3893,6 +3893,7 @@ dependencies = [
"codex-extension-api",
"codex-features",
"codex-otel",
"codex-protocol",
"codex-tools",
"codex-utils-absolute-path",
"codex-utils-output-truncation",
codex-rs/app-server/src/request_processors/thread_processor.rs
+3 -7
View File
@@ -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
codex-rs/config/src/types.rs
+7
View File
@@ -286,10 +286,14 @@ pub struct ToolSuggestConfig {
pub disabled_tools: Vec<ToolSuggestDisabledTool>,
}
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<MemoryVersion>,
/// 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<bool>,
@@ -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<MemoriesToml> 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),
codex-rs/config/src/types_tests.rs
+18
View File
@@ -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::<MemoriesToml>("version = \"v3\"").is_err());
}
codex-rs/core/config.schema.json
+15
View File
@@ -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",
codex-rs/core/src/config/config_tests.rs
+2
View File
@@ -334,6 +334,7 @@ consolidation_model = "gpt-5.2"
toml::from_str::<ConfigToml>(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,
codex-rs/core/src/config/mod.rs
+1 -2
View File
@@ -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(),
codex-rs/ext/memories/Cargo.toml
+1
View File
@@ -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 }
codex-rs/ext/memories/src/extension.rs
+16 -3
View File
@@ -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<Config> 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::<MemoriesExtensionConfig>() {
// 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(),
)
}
codex-rs/ext/memories/src/local.rs
-6
View File
@@ -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<PathBuf>) -> Self {
Self { root: root.into() }
}
codex-rs/ext/memories/src/prompts.rs
+3 -1
View File
@@ -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<String> {
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
codex-rs/ext/memories/src/prompts_tests.rs
+1 -1
View File
@@ -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();
codex-rs/ext/memories/src/tests.rs
+5
View File
@@ -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(),
codex-rs/memories/read/src/usage.rs
+5
View File
@@ -44,6 +44,7 @@ pub fn memories_usage_kinds_from_command(command: &str) -> Vec<MemoriesUsageKind
}
fn get_memory_kind(path: String) -> Option<MemoriesUsageKind> {
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<MemoriesUsageKind> {
None
}
}
#[cfg(test)]
#[path = "usage_tests.rs"]
mod tests;
codex-rs/memories/read/src/usage_tests.rs
+22
View File
@@ -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}"
);
}
}
codex-rs/memories/write/src/control.rs
+1
View File
@@ -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?;
codex-rs/memories/write/src/phase1.rs
+5 -10
View File
@@ -110,10 +110,9 @@ pub async fn run(context: Arc<MemoryStartupContext>, config: Arc<Config>) {
/// 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<Vec<codex_state::Stage1JobClaim>> {
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::<Vec<_>>();
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,
codex-rs/memories/write/src/phase2.rs
+31 -66
View File
@@ -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<Claim, &'static str> {
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<Config> {
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<codex_otel::Timer>,
) {
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<StateRuntime>,
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,
codex-rs/memories/write/src/runtime.rs
+17 -3
View File
@@ -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<CodexThread>,
thread_manager: Arc<ThreadManager>,
@@ -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<Arc<StateRuntime>> {
self.thread.state_db()
pub(crate) async fn memory_store(&self) -> Option<MemoryStore> {
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 {
codex-rs/memories/write/src/start.rs
+7 -7
View File
@@ -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;
codex-rs/protocol/src/lib.rs
+2
View File
@@ -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;
codex-rs/protocol/src/memory_version.rs
+23
View File
@@ -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",
}
}
}
codex-rs/state/src/runtime.rs
+45 -11
View File
@@ -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<sqlx::SqlitePool>,
thread_goals: GoalStore,
memories: MemoryStore,
memories_v2: Arc<tokio::sync::OnceCell<MemoryStore>>,
thread_queue: SqliteQueueStore,
thread_updated_at_millis: Arc<AtomicI64>,
thread_recency_at_millis: Arc<AtomicI64>,
@@ -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<bool> {
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)
}
}
codex-rs/state/src/runtime/memory_versions.rs
+56
View File
@@ -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<MemoryStore> {
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;
codex-rs/state/src/runtime/memory_versions_tests.rs
+113
View File
@@ -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<_>>(),
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(())
}
codex-rs/state/src/runtime/threads.rs
+1 -1
View File
@@ -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?;
}
codex-rs/state/src/sqlite.rs
+21 -1
View File
@@ -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<SqlitePool> {
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())