mirror of
https://github.com/openai/codex.git
synced 2026-11-13 14:46:14 +00:00
feat: mem v2 - PR4 (#11369)
# Memories migration plan (simplified global workflow) ## Target behavior - One shared memory root only: `~/.codex/memories/`. - No per-cwd memory buckets, no cwd hash handling. - Phase 1 candidate rules: - Not currently being processed unless the job lease is stale. - Rollout updated within the max-age window (currently 30 days). - Rollout idle for at least 12 hours (new constant). - Global cap: at most 64 stage-1 jobs in `running` state at any time (new invariant). - Stage-1 model output shape (new): - `rollout_slug` (accepted but ignored for now). - `rollout_summary`. - `raw_memory`. - Phase-1 artifacts written under the shared root: - `rollout_summaries/<thread_id>.md` for each rollout summary. - `raw_memories.md` containing appended/merged raw memory paragraphs. - Phase 2 runs one consolidation agent for the shared `memories/` directory. - Phase-2 lock is DB-backed with 1 hour lease and heartbeat/expiry. ## Current code map - Core startup pipeline: `core/src/memories/startup/mod.rs`. - Stage-1 request+parse: `core/src/memories/startup/extract.rs`, `core/src/memories/stage_one.rs`, templates in `core/templates/memories/`. - File materialization: `core/src/memories/storage.rs`, `core/src/memories/layout.rs`. - Scope routing (cwd/user): `core/src/memories/scope.rs`, `core/src/memories/startup/mod.rs`. - DB job lifecycle and scope queueing: `state/src/runtime/memory.rs`. ## PR plan ## PR 1: Correct phase-1 selection invariants (no behavior-breaking layout changes yet) - Add `PHASE_ONE_MIN_ROLLOUT_IDLE_HOURS: i64 = 12` in `core/src/memories/mod.rs`. - Thread this into `state::claim_stage1_jobs_for_startup(...)`. - Enforce idle-time filter in DB selection logic (not only in-memory filtering after `scan_limit`) so eligible threads are not starved by very recent threads. - Enforce global running cap of 64 at claim time in DB logic: - Count fresh `memory_stage1` running jobs. - Only allow new claims while count < cap. - Keep stale-lease takeover behavior intact. - Add/adjust tests in `state/src/runtime.rs`: - Idle filter inclusion/exclusion around 12h boundary. - Global running-cap guarantee. - Existing stale/fresh ownership behavior still passes. Acceptance criteria: - Startup never creates more than 64 fresh `memory_stage1` running jobs. - Threads updated <12h ago are skipped. - Threads older than 30d are skipped. ## PR 2: Stage-1 output contract + storage artifacts (forward-compatible) - Update parser/types to accept the new structured output while keeping backward compatibility: - Add `rollout_slug` (optional for now). - Add `rollout_summary`. - Keep alias support for legacy `summary` and `rawMemory` until prompt swap completes. - Update stage-1 schema generator in `core/src/memories/stage_one.rs` to include the new keys. - Update prompt templates: - `core/templates/memories/stage_one_system.md`. - `core/templates/memories/stage_one_input.md`. - Replace storage model in `core/src/memories/storage.rs`: - Introduce `rollout_summaries/` directory writer (`<thread_id>.md` files). - Introduce `raw_memories.md` aggregator writer from DB rows. - Keep deterministic rebuild behavior from DB outputs so files can always be regenerated. - Update consolidation prompt template to reference `rollout_summaries/` + `raw_memories.md` inputs. Acceptance criteria: - Stage-1 accepts both old and new output keys during migration. - Phase-1 artifacts are generated in new format from DB state. - No dependence on per-thread files in `raw_memories/`. ## PR 3: Remove per-cwd memories and move to one global memory root - Simplify layout in `core/src/memories/layout.rs`: - Single root: `codex_home/memories`. - Remove cwd-hash bucket helpers and normalization logic used only for memory pathing. - Remove scope branching from startup phase-2 dispatch path: - No cwd/user mapping in `core/src/memories/startup/mod.rs`. - One target root for consolidation. - In `state/src/runtime/memory.rs`, stop enqueueing/handling cwd consolidation scope. - Keep one logical consolidation scope/job key (global/user) to avoid a risky schema rewrite in same PR. - Add one-time migration helper (core side) to preserve current shared memory output: - If `~/.codex/memories/user/memory` exists and new root is empty, move/copy contents into `~/.codex/memories`. - Leave old hashed cwd buckets untouched for now (safe/no-destructive migration). Acceptance criteria: - New runs only read/write `~/.codex/memories`. - No new cwd-scoped consolidation jobs are enqueued. - Existing user-shared memory content is preserved. ## PR 4: Phase-2 global lock simplification and cleanup - Replace multi-scope dispatch with a single global consolidation claim path: - Either reuse jobs table with one fixed key, or add a tiny dedicated lock helper; keep 1h lease. - Ensure at most one consolidation agent can run at once. - Keep heartbeat + stale lock recovery semantics in `core/src/memories/startup/watch.rs`. - Remove dead scope code and legacy constants no longer used. - Update tests: - One-agent-at-a-time behavior. - Lock expiry allows takeover after stale lease. Acceptance criteria: - Exactly one phase-2 consolidation agent can be active cluster-wide (per local DB). - Stale lock recovers automatically. ## PR 5: Final cleanup and docs - Remove legacy artifacts and references: - `raw_memories/` and `memory_summary.md` assumptions from prompts/comments/tests. - Scope constants for cwd memory pathing in core/state if fully unused. - Update docs under `docs/` for memory workflow and directory layout. - Add a brief operator note for rollout: compatibility window for old stage-1 JSON keys and when to remove aliases. Acceptance criteria: - Code and docs reflect only the simplified global workflow. - No stale references to per-cwd memory buckets. ## Notes on sequencing - PR 1 is safest first because it improves correctness without changing external artifact layout. - PR 2 keeps parser compatibility so prompt deployment can happen independently. - PR 3 and PR 4 split filesystem/scope simplification from locking simplification to reduce blast radius. - PR 5 is intentionally cleanup-only.
This commit is contained in:
1 parent
d8f9bb65e2
commit
623d3f4071
10 files changed
+384
-566
No files matched your search
+95
-115
@@ -39,9 +39,6 @@ use uuid::Uuid;
|
||||
pub const STATE_DB_FILENAME: &str = "state";
|
||||
pub const STATE_DB_VERSION: u32 = 4;
|
||||
|
||||
const MEMORY_SCOPE_KIND_USER: &str = "user";
|
||||
const MEMORY_SCOPE_KEY_USER: &str = "user";
|
||||
|
||||
const METRIC_DB_INIT: &str = "codex.db.init";
|
||||
|
||||
mod memory;
|
||||
@@ -86,27 +83,18 @@ pub struct Stage1StartupClaimParams<'a> {
|
||||
pub lease_seconds: i64,
|
||||
}
|
||||
|
||||
/// Scope row used to queue phase-2 consolidation work.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct PendingScopeConsolidation {
|
||||
/// Scope family (`user`).
|
||||
pub scope_kind: String,
|
||||
/// Scope identifier keyed by `scope_kind`.
|
||||
pub scope_key: String,
|
||||
}
|
||||
|
||||
/// Result of trying to claim a phase-2 consolidation job.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum Phase2JobClaimOutcome {
|
||||
/// The caller owns the scope and should spawn consolidation.
|
||||
/// The caller owns the global lock and should spawn consolidation.
|
||||
Claimed {
|
||||
ownership_token: String,
|
||||
/// Snapshot of `input_watermark` at claim time.
|
||||
input_watermark: i64,
|
||||
},
|
||||
/// The scope is not pending consolidation (or is already up to date).
|
||||
/// The global job is not pending consolidation (or is already up to date).
|
||||
SkippedNotDirty,
|
||||
/// Another worker currently owns a fresh lease for this scope.
|
||||
/// Another worker currently owns a fresh global consolidation lease.
|
||||
SkippedRunning,
|
||||
}
|
||||
|
||||
@@ -926,7 +914,6 @@ fn push_thread_order_and_limit(
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::PendingScopeConsolidation;
|
||||
use super::Phase2JobClaimOutcome;
|
||||
use super::STATE_DB_FILENAME;
|
||||
use super::STATE_DB_VERSION;
|
||||
@@ -1463,7 +1450,7 @@ WHERE kind = 'memory_stage1'
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn phase2_consolidation_jobs_rerun_when_watermark_advances() {
|
||||
async fn phase2_global_consolidation_reruns_when_watermark_advances() {
|
||||
let codex_home = unique_temp_dir();
|
||||
let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string(), None)
|
||||
.await
|
||||
@@ -1472,24 +1459,12 @@ WHERE kind = 'memory_stage1'
|
||||
let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id");
|
||||
|
||||
runtime
|
||||
.enqueue_scope_consolidation("user", "user", 100)
|
||||
.enqueue_global_consolidation(100)
|
||||
.await
|
||||
.expect("enqueue scope");
|
||||
|
||||
let scopes = runtime
|
||||
.list_pending_scope_consolidations(10)
|
||||
.await
|
||||
.expect("list pending");
|
||||
assert_eq!(
|
||||
scopes,
|
||||
vec![PendingScopeConsolidation {
|
||||
scope_kind: "user".to_string(),
|
||||
scope_key: "user".to_string(),
|
||||
}]
|
||||
);
|
||||
.expect("enqueue global consolidation");
|
||||
|
||||
let claim = runtime
|
||||
.try_claim_phase2_job("user", "user", owner, 3600)
|
||||
.try_claim_global_phase2_job(owner, 3600)
|
||||
.await
|
||||
.expect("claim phase2");
|
||||
let (ownership_token, input_watermark) = match claim {
|
||||
@@ -1501,30 +1476,25 @@ WHERE kind = 'memory_stage1'
|
||||
};
|
||||
assert!(
|
||||
runtime
|
||||
.mark_phase2_job_succeeded(
|
||||
"user",
|
||||
"user",
|
||||
ownership_token.as_str(),
|
||||
input_watermark,
|
||||
)
|
||||
.mark_global_phase2_job_succeeded(ownership_token.as_str(), input_watermark)
|
||||
.await
|
||||
.expect("mark phase2 succeeded"),
|
||||
"phase2 success should finalize for current token"
|
||||
);
|
||||
|
||||
let claim_up_to_date = runtime
|
||||
.try_claim_phase2_job("user", "user", owner, 3600)
|
||||
.try_claim_global_phase2_job(owner, 3600)
|
||||
.await
|
||||
.expect("claim phase2 up-to-date");
|
||||
assert_eq!(claim_up_to_date, Phase2JobClaimOutcome::SkippedNotDirty);
|
||||
|
||||
runtime
|
||||
.enqueue_scope_consolidation("user", "user", 101)
|
||||
.enqueue_global_consolidation(101)
|
||||
.await
|
||||
.expect("enqueue scope again");
|
||||
.expect("enqueue global consolidation again");
|
||||
|
||||
let claim_rerun = runtime
|
||||
.try_claim_phase2_job("user", "user", owner, 3600)
|
||||
.try_claim_global_phase2_job(owner, 3600)
|
||||
.await
|
||||
.expect("claim phase2 rerun");
|
||||
assert!(
|
||||
@@ -1536,7 +1506,7 @@ WHERE kind = 'memory_stage1'
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn list_stage1_outputs_for_user_scope_returns_latest_outputs() {
|
||||
async fn list_stage1_outputs_for_global_returns_latest_outputs() {
|
||||
let codex_home = unique_temp_dir();
|
||||
let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string(), None)
|
||||
.await
|
||||
@@ -1607,9 +1577,9 @@ WHERE kind = 'memory_stage1'
|
||||
);
|
||||
|
||||
let outputs = runtime
|
||||
.list_stage1_outputs_for_scope("user", "user", 10)
|
||||
.list_stage1_outputs_for_global(10)
|
||||
.await
|
||||
.expect("list stage1 outputs for user scope");
|
||||
.expect("list stage1 outputs for global");
|
||||
assert_eq!(outputs.len(), 2);
|
||||
assert_eq!(outputs[0].thread_id, thread_id_b);
|
||||
assert_eq!(outputs[0].summary, "summary b");
|
||||
@@ -1620,7 +1590,7 @@ WHERE kind = 'memory_stage1'
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn mark_stage1_job_succeeded_enqueues_single_user_scope() {
|
||||
async fn mark_stage1_job_succeeded_enqueues_global_consolidation() {
|
||||
let codex_home = unique_temp_dir();
|
||||
let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string(), None)
|
||||
.await
|
||||
@@ -1679,106 +1649,116 @@ WHERE kind = 'memory_stage1'
|
||||
"stage1 success should persist output for thread b"
|
||||
);
|
||||
|
||||
let pending_scopes = runtime
|
||||
.list_pending_scope_consolidations(10)
|
||||
let claim = runtime
|
||||
.try_claim_global_phase2_job(owner, 3600)
|
||||
.await
|
||||
.expect("list pending scopes");
|
||||
assert_eq!(
|
||||
pending_scopes,
|
||||
vec![PendingScopeConsolidation {
|
||||
scope_kind: "user".to_string(),
|
||||
scope_key: "user".to_string(),
|
||||
}]
|
||||
);
|
||||
.expect("claim global consolidation");
|
||||
let input_watermark = match claim {
|
||||
Phase2JobClaimOutcome::Claimed {
|
||||
input_watermark, ..
|
||||
} => input_watermark,
|
||||
other => panic!("unexpected global consolidation claim outcome: {other:?}"),
|
||||
};
|
||||
assert_eq!(input_watermark, 101);
|
||||
|
||||
let _ = tokio::fs::remove_dir_all(codex_home).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn list_pending_scope_consolidations_omits_unclaimable_jobs() {
|
||||
async fn phase2_global_lock_allows_only_one_fresh_runner() {
|
||||
let codex_home = unique_temp_dir();
|
||||
let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string(), None)
|
||||
.await
|
||||
.expect("initialize runtime");
|
||||
let owner = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner id");
|
||||
|
||||
runtime
|
||||
.enqueue_scope_consolidation("user", "scope-running", 200)
|
||||
.enqueue_global_consolidation(200)
|
||||
.await
|
||||
.expect("enqueue running scope");
|
||||
runtime
|
||||
.enqueue_scope_consolidation("user", "scope-backoff", 199)
|
||||
.await
|
||||
.expect("enqueue backoff scope");
|
||||
runtime
|
||||
.enqueue_scope_consolidation("user", "scope-exhausted", 198)
|
||||
.await
|
||||
.expect("enqueue exhausted scope");
|
||||
runtime
|
||||
.enqueue_scope_consolidation("user", "scope-claimable-a", 90)
|
||||
.await
|
||||
.expect("enqueue claimable scope a");
|
||||
runtime
|
||||
.enqueue_scope_consolidation("user", "scope-claimable-b", 89)
|
||||
.await
|
||||
.expect("enqueue claimable scope b");
|
||||
.expect("enqueue global consolidation");
|
||||
|
||||
let owner_a = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner a");
|
||||
let owner_b = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner b");
|
||||
|
||||
let running_claim = runtime
|
||||
.try_claim_phase2_job("user", "scope-running", owner, 3600)
|
||||
.try_claim_global_phase2_job(owner_a, 3600)
|
||||
.await
|
||||
.expect("claim running scope");
|
||||
.expect("claim global lock");
|
||||
assert!(
|
||||
matches!(running_claim, Phase2JobClaimOutcome::Claimed { .. }),
|
||||
"scope-running should be claimed"
|
||||
"first owner should claim global lock"
|
||||
);
|
||||
|
||||
let backoff_claim = runtime
|
||||
.try_claim_phase2_job("user", "scope-backoff", owner, 3600)
|
||||
let second_claim = runtime
|
||||
.try_claim_global_phase2_job(owner_b, 3600)
|
||||
.await
|
||||
.expect("claim backoff scope");
|
||||
let backoff_token = match backoff_claim {
|
||||
.expect("claim global lock from second owner");
|
||||
assert_eq!(second_claim, Phase2JobClaimOutcome::SkippedRunning);
|
||||
|
||||
let _ = tokio::fs::remove_dir_all(codex_home).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn phase2_global_lock_stale_lease_allows_takeover() {
|
||||
let codex_home = unique_temp_dir();
|
||||
let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string(), None)
|
||||
.await
|
||||
.expect("initialize runtime");
|
||||
|
||||
runtime
|
||||
.enqueue_global_consolidation(300)
|
||||
.await
|
||||
.expect("enqueue global consolidation");
|
||||
|
||||
let owner_a = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner a");
|
||||
let owner_b = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("owner b");
|
||||
|
||||
let initial_claim = runtime
|
||||
.try_claim_global_phase2_job(owner_a, 3600)
|
||||
.await
|
||||
.expect("claim initial global lock");
|
||||
let token_a = match initial_claim {
|
||||
Phase2JobClaimOutcome::Claimed {
|
||||
ownership_token, ..
|
||||
} => ownership_token,
|
||||
other => panic!("unexpected backoff claim outcome: {other:?}"),
|
||||
other => panic!("unexpected initial claim outcome: {other:?}"),
|
||||
};
|
||||
assert!(
|
||||
runtime
|
||||
.mark_phase2_job_failed(
|
||||
"user",
|
||||
"scope-backoff",
|
||||
backoff_token.as_str(),
|
||||
"temporary failure",
|
||||
3600,
|
||||
)
|
||||
.await
|
||||
.expect("mark backoff scope failed"),
|
||||
"backoff scope should transition to retry backoff"
|
||||
);
|
||||
|
||||
sqlx::query("UPDATE jobs SET retry_remaining = 0 WHERE kind = ? AND job_key = ?")
|
||||
.bind("memory_consolidate_user")
|
||||
.bind("scope-exhausted")
|
||||
sqlx::query("UPDATE jobs SET lease_until = ? WHERE kind = ? AND job_key = ?")
|
||||
.bind(Utc::now().timestamp() - 1)
|
||||
.bind("memory_consolidate_global")
|
||||
.bind("global")
|
||||
.execute(runtime.pool.as_ref())
|
||||
.await
|
||||
.expect("set exhausted scope retries to zero");
|
||||
.expect("expire global consolidation lease");
|
||||
|
||||
let pending = runtime
|
||||
.list_pending_scope_consolidations(2)
|
||||
let takeover_claim = runtime
|
||||
.try_claim_global_phase2_job(owner_b, 3600)
|
||||
.await
|
||||
.expect("list pending scopes");
|
||||
.expect("claim stale global lock");
|
||||
let (token_b, input_watermark) = match takeover_claim {
|
||||
Phase2JobClaimOutcome::Claimed {
|
||||
ownership_token,
|
||||
input_watermark,
|
||||
} => (ownership_token, input_watermark),
|
||||
other => panic!("unexpected takeover claim outcome: {other:?}"),
|
||||
};
|
||||
assert_ne!(token_a, token_b);
|
||||
assert_eq!(input_watermark, 300);
|
||||
|
||||
assert_eq!(
|
||||
pending,
|
||||
vec![
|
||||
PendingScopeConsolidation {
|
||||
scope_kind: "user".to_string(),
|
||||
scope_key: "scope-claimable-a".to_string(),
|
||||
},
|
||||
PendingScopeConsolidation {
|
||||
scope_kind: "user".to_string(),
|
||||
scope_key: "scope-claimable-b".to_string(),
|
||||
},
|
||||
]
|
||||
runtime
|
||||
.mark_global_phase2_job_succeeded(token_a.as_str(), 300)
|
||||
.await
|
||||
.expect("mark stale owner success result"),
|
||||
false,
|
||||
"stale owner should lose finalization ownership after takeover"
|
||||
);
|
||||
assert!(
|
||||
runtime
|
||||
.mark_global_phase2_job_succeeded(token_b.as_str(), 300)
|
||||
.await
|
||||
.expect("mark takeover owner success"),
|
||||
"takeover owner should finalize consolidation"
|
||||
);
|
||||
|
||||
let _ = tokio::fs::remove_dir_all(codex_home).await;
|
||||
|
||||
Reference in new issue
Block a user