mirror of
https://github.com/openai/codex.git
synced 2026-04-30 01:16:54 +00:00
state: add memory consolidation lock primitives (#11199)
## Summary - add a migration for memory_consolidation_locks - add acquire/release lock primitives to codex-state runtime - add core/state_db wrappers and cwd normalization for memory queries and lock keys ## Testing - cargo test -p codex-state memory_consolidation_lock_ - cargo test -p codex-core --lib state_db::
This commit is contained in:
@@ -1,5 +1,6 @@
|
||||
use crate::config::Config;
|
||||
use crate::features::Feature;
|
||||
use crate::path_utils::normalize_for_path_comparison;
|
||||
use crate::rollout::list::Cursor;
|
||||
use crate::rollout::list::ThreadSortKey;
|
||||
use crate::rollout::metadata;
|
||||
@@ -156,6 +157,10 @@ fn cursor_to_anchor(cursor: Option<&Cursor>) -> Option<codex_state::Anchor> {
|
||||
Some(codex_state::Anchor { ts, id })
|
||||
}
|
||||
|
||||
fn normalize_cwd_for_state_db(cwd: &Path) -> PathBuf {
|
||||
normalize_for_path_comparison(cwd).unwrap_or_else(|_| cwd.to_path_buf())
|
||||
}
|
||||
|
||||
/// List thread ids from SQLite for parity checks without rollout scanning.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub async fn list_thread_ids_db(
|
||||
@@ -355,7 +360,11 @@ pub async fn get_last_n_thread_memories_for_cwd(
|
||||
stage: &str,
|
||||
) -> Option<Vec<codex_state::ThreadMemory>> {
|
||||
let ctx = context?;
|
||||
match ctx.get_last_n_thread_memories_for_cwd(cwd, n).await {
|
||||
let normalized_cwd = normalize_cwd_for_state_db(cwd);
|
||||
match ctx
|
||||
.get_last_n_thread_memories_for_cwd(&normalized_cwd, n)
|
||||
.await
|
||||
{
|
||||
Ok(memories) => Some(memories),
|
||||
Err(err) => {
|
||||
warn!("state db get_last_n_thread_memories_for_cwd failed during {stage}: {err}");
|
||||
@@ -364,6 +373,49 @@ pub async fn get_last_n_thread_memories_for_cwd(
|
||||
}
|
||||
}
|
||||
|
||||
/// Try to acquire or renew a per-cwd memory consolidation lock.
|
||||
pub async fn try_acquire_memory_consolidation_lock(
|
||||
context: Option<&codex_state::StateRuntime>,
|
||||
cwd: &Path,
|
||||
working_thread_id: ThreadId,
|
||||
lease_seconds: i64,
|
||||
stage: &str,
|
||||
) -> Option<bool> {
|
||||
let ctx = context?;
|
||||
let normalized_cwd = normalize_cwd_for_state_db(cwd);
|
||||
match ctx
|
||||
.try_acquire_memory_consolidation_lock(&normalized_cwd, working_thread_id, lease_seconds)
|
||||
.await
|
||||
{
|
||||
Ok(acquired) => Some(acquired),
|
||||
Err(err) => {
|
||||
warn!("state db try_acquire_memory_consolidation_lock failed during {stage}: {err}");
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Release a per-cwd memory consolidation lock if held by `working_thread_id`.
|
||||
pub async fn release_memory_consolidation_lock(
|
||||
context: Option<&codex_state::StateRuntime>,
|
||||
cwd: &Path,
|
||||
working_thread_id: ThreadId,
|
||||
stage: &str,
|
||||
) -> Option<bool> {
|
||||
let ctx = context?;
|
||||
let normalized_cwd = normalize_cwd_for_state_db(cwd);
|
||||
match ctx
|
||||
.release_memory_consolidation_lock(&normalized_cwd, working_thread_id)
|
||||
.await
|
||||
{
|
||||
Ok(released) => Some(released),
|
||||
Err(err) => {
|
||||
warn!("state db release_memory_consolidation_lock failed during {stage}: {err}");
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Reconcile rollout items into SQLite, falling back to scanning the rollout file.
|
||||
pub async fn reconcile_rollout(
|
||||
context: Option<&codex_state::StateRuntime>,
|
||||
@@ -400,6 +452,7 @@ pub async fn reconcile_rollout(
|
||||
}
|
||||
};
|
||||
let mut metadata = outcome.metadata;
|
||||
metadata.cwd = normalize_cwd_for_state_db(&metadata.cwd);
|
||||
match archived_only {
|
||||
Some(true) if metadata.archived_at.is_none() => {
|
||||
metadata.archived_at = Some(metadata.updated_at);
|
||||
@@ -447,6 +500,7 @@ pub async fn read_repair_rollout_path(
|
||||
&& let Ok(Some(mut metadata)) = ctx.get_thread(thread_id).await
|
||||
{
|
||||
metadata.rollout_path = rollout_path.to_path_buf();
|
||||
metadata.cwd = normalize_cwd_for_state_db(&metadata.cwd);
|
||||
match archived_only {
|
||||
Some(true) if metadata.archived_at.is_none() => {
|
||||
metadata.archived_at = Some(metadata.updated_at);
|
||||
@@ -509,6 +563,7 @@ pub async fn apply_rollout_items(
|
||||
},
|
||||
};
|
||||
builder.rollout_path = rollout_path.to_path_buf();
|
||||
builder.cwd = normalize_cwd_for_state_db(&builder.cwd);
|
||||
if let Err(err) = ctx.apply_rollout_items(&builder, items, None).await {
|
||||
warn!(
|
||||
"state db apply_rollout_items failed during {stage} for {}: {err}",
|
||||
|
||||
Reference in New Issue
Block a user