openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-05 15:18:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

[thread-store] derive in-memory metadata

Browse Source
This commit is contained in:
Michael Zeng
2026-07-06 17:30:59 -07:00
parent bbd91ab14e
commit 0001ed71ea
3 changed files with 100 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
codex-rs/app-server/tests/suite/v2/remote_thread_store.rs
View File

@@ -32,6 +32,8 @@ use codex_app_server_protocol::ThreadDeleteParams;
use codex_app_server_protocol::ThreadDeleteResponse;
use codex_app_server_protocol::ThreadListParams;
use codex_app_server_protocol::ThreadListResponse;
use codex_app_server_protocol::ThreadReadParams;
use codex_app_server_protocol::ThreadReadResponse;
use codex_app_server_protocol::ThreadResumeParams;
use codex_app_server_protocol::ThreadStartParams;
use codex_app_server_protocol::ThreadStartResponse;
@@ -143,8 +145,27 @@ async fn thread_delete_with_non_local_thread_store_does_not_create_local_persist
assert_eq!(data.len(), 1);
assert_eq!(data[0].id, thread.id);
assert_eq!(data[0].path, None);
assert_eq!(data[0].preview, "Hello");
assert_eq!(data[0].model_provider, "mock_provider");
delete_thread(&client, /*request_id*/ 4, thread.id.clone()).await?;
let response = client
.request(ClientRequest::ThreadRead {
request_id: RequestId::Integer(4),
params: ThreadReadParams {
thread_id: thread.id.clone(),
include_turns: false,
},
})
.await?
.expect("thread/read should succeed");
let ThreadReadResponse {
thread: read_thread,
} = serde_json::from_value(response).expect("thread/read response should parse");
assert_eq!(read_thread.preview, "Hello");
assert_eq!(read_thread.model_provider, "mock_provider");
assert_eq!(read_thread.path, None);
delete_thread(&client, /*request_id*/ 5, thread.id.clone()).await?;
let unloaded_thread_id = ThreadId::from_string(&Uuid::new_v4().to_string())?;
thread_store
.create_thread(StoreCreateThreadParams {
@@ -171,7 +192,7 @@ async fn thread_delete_with_non_local_thread_store_does_not_create_local_persist
.await?;
delete_thread(
&client,
/*request_id*/ 5,
/*request_id*/ 6,
unloaded_thread_id.to_string(),
)
.await?;

87
codex-rs/thread-store/src/in_memory.rs
View File

@@ -38,6 +38,8 @@ use crate::ThreadStoreFuture;
use crate::ThreadStoreResult;
use crate::UpdateThreadMetadataParams;
use crate::error::reject_paginated_history_mode;
use crate::thread_metadata_sync::PendingThreadMetadataPatch;
use crate::thread_metadata_sync::ThreadMetadataSync;
use crate::types::canonical_history_mode_from_rollout_items;
static IN_MEMORY_THREAD_STORES: OnceLock<Mutex<HashMap<String, Arc<InMemoryThreadStore>>>> =
@@ -393,11 +395,47 @@ struct InMemoryThreadStoreState {
calls: InMemoryThreadStoreCalls,
created_threads: HashMap<ThreadId, CreateThreadParams>,
histories: HashMap<ThreadId, Vec<RolloutItem>>,
metadata_syncs: HashMap<ThreadId, ThreadMetadataSync>,
metadata_updates: HashMap<ThreadId, ThreadMetadataPatch>,
names: HashMap<ThreadId, Option<String>>,
rollout_paths: HashMap<PathBuf, ThreadId>,
}
impl InMemoryThreadStoreState {
fn apply_pending_metadata_update(
&mut self,
thread_id: ThreadId,
update: Option<PendingThreadMetadataPatch>,
) {
let Some(update) = update else {
return;
};
self.metadata_updates
.entry(thread_id)
.or_default()
.merge(update.patch.clone());
if let Some(metadata_sync) = self.metadata_syncs.get_mut(&thread_id) {
metadata_sync.mark_pending_update_applied(&update);
}
}
fn flush_pending_metadata_update(&mut self, thread_id: ThreadId) {
let update = self
.metadata_syncs
.get(&thread_id)
.and_then(ThreadMetadataSync::take_pending_update);
self.apply_pending_metadata_update(thread_id, update);
}
fn flush_pending_metadata_update_for_existing_history(&mut self, thread_id: ThreadId) {
let update = self
.metadata_syncs
.get(&thread_id)
.and_then(ThreadMetadataSync::take_pending_update_for_existing_history);
self.apply_pending_metadata_update(thread_id, update);
}
}
impl InMemoryThreadStore {
/// Returns the store associated with `id`, creating it if needed.
pub fn for_id(id: impl Into<String>) -> Arc<Self> {
@@ -421,6 +459,7 @@ impl InMemoryThreadStore {
async fn create_thread(&self, params: CreateThreadParams) -> ThreadStoreResult<()> {
reject_paginated_history_mode(params.history_mode)?;
let metadata_sync = ThreadMetadataSync::for_create(&params).await;
let mut state = self.state.lock().await;
state.calls.create_thread += 1;
let session_meta = SessionMeta {
@@ -454,30 +493,36 @@ impl InMemoryThreadStore {
meta: session_meta,
git: None,
}));
state.created_threads.insert(params.thread_id, params);
let thread_id = params.thread_id;
state.created_threads.insert(thread_id, params);
state.metadata_syncs.insert(thread_id, metadata_sync);
Ok(())
}
async fn resume_thread(&self, params: ResumeThreadParams) -> ThreadStoreResult<()> {
let mut metadata_sync = ThreadMetadataSync::for_resume(&params);
let mut state = self.state.lock().await;
state.calls.resume_thread += 1;
let thread_id = params.thread_id;
let history_mode = params
.history
.as_deref()
.map(Vec::as_slice)
.map(canonical_history_mode_from_rollout_items)
.unwrap_or_else(|| history_mode_from_state(&state, params.thread_id));
.unwrap_or_else(|| history_mode_from_state(&state, thread_id));
reject_paginated_history_mode(history_mode)?;
if let Some(history) = params.history {
state
.histories
.insert(params.thread_id, Arc::unwrap_or_clone(history));
.insert(thread_id, Arc::unwrap_or_clone(history));
} else {
state.histories.entry(params.thread_id).or_default();
let history = state.histories.entry(thread_id).or_default();
metadata_sync.record_resume_history(history);
}
if let Some(rollout_path) = params.rollout_path {
state.rollout_paths.insert(rollout_path, params.thread_id);
state.rollout_paths.insert(rollout_path, thread_id);
}
state.metadata_syncs.insert(thread_id, metadata_sync);
Ok(())
}
@@ -488,11 +533,16 @@ impl InMemoryThreadStore {
}
let mut state = self.state.lock().await;
state.calls.append_items += 1;
let update = state
.metadata_syncs
.get_mut(&params.thread_id)
.and_then(|metadata_sync| metadata_sync.observe_appended_items(&canonical_items));
state
.histories
.entry(params.thread_id)
.or_default()
.extend(canonical_items);
state.apply_pending_metadata_update(params.thread_id, update);
Ok(())
}
@@ -572,6 +622,7 @@ impl InMemoryThreadStore {
) -> ThreadStoreResult<StoredThread> {
let mut state = self.state.lock().await;
state.calls.update_thread_metadata += 1;
state.flush_pending_metadata_update(params.thread_id);
if let Some(name) = params.patch.name.clone() {
state.names.insert(params.thread_id, name);
}
@@ -588,6 +639,7 @@ impl InMemoryThreadStore {
state.calls.delete_thread += 1;
let existed = state.histories.remove(&params.thread_id).is_some();
state.created_threads.remove(&params.thread_id);
state.metadata_syncs.remove(&params.thread_id);
state.names.remove(&params.thread_id);
state.metadata_updates.remove(&params.thread_id);
state
@@ -620,30 +672,39 @@ impl ThreadStore for InMemoryThreadStore {
Box::pin(InMemoryThreadStore::append_items(self, params))
}
fn persist_thread(&self, _thread_id: ThreadId) -> ThreadStoreFuture<'_, ()> {
fn persist_thread(&self, thread_id: ThreadId) -> ThreadStoreFuture<'_, ()> {
Box::pin(async move {
self.state.lock().await.calls.persist_thread += 1;
let mut state = self.state.lock().await;
state.calls.persist_thread += 1;
state.flush_pending_metadata_update(thread_id);
Ok(())
})
}
fn flush_thread(&self, _thread_id: ThreadId) -> ThreadStoreFuture<'_, ()> {
fn flush_thread(&self, thread_id: ThreadId) -> ThreadStoreFuture<'_, ()> {
Box::pin(async move {
self.state.lock().await.calls.flush_thread += 1;
let mut state = self.state.lock().await;
state.calls.flush_thread += 1;
state.flush_pending_metadata_update_for_existing_history(thread_id);
Ok(())
})
}
fn shutdown_thread(&self, _thread_id: ThreadId) -> ThreadStoreFuture<'_, ()> {
fn shutdown_thread(&self, thread_id: ThreadId) -> ThreadStoreFuture<'_, ()> {
Box::pin(async move {
self.state.lock().await.calls.shutdown_thread += 1;
let mut state = self.state.lock().await;
state.calls.shutdown_thread += 1;
state.flush_pending_metadata_update_for_existing_history(thread_id);
state.metadata_syncs.remove(&thread_id);
Ok(())
})
}
fn discard_thread(&self, _thread_id: ThreadId) -> ThreadStoreFuture<'_, ()> {
fn discard_thread(&self, thread_id: ThreadId) -> ThreadStoreFuture<'_, ()> {
Box::pin(async move {
self.state.lock().await.calls.discard_thread += 1;
let mut state = self.state.lock().await;
state.calls.discard_thread += 1;
state.metadata_syncs.remove(&thread_id);
Ok(())
})
}

6
codex-rs/thread-store/src/thread_metadata_sync.rs
View File

@@ -27,10 +27,10 @@ const THREAD_UPDATED_AT_TOUCH_INTERVAL: Duration = Duration::from_secs(5);
#[cfg(test)]
const THREAD_UPDATED_AT_TOUCH_INTERVAL: Duration = Duration::from_millis(50);
/// Local-store helper that derives SQLite metadata from canonical rollout items.
/// Store implementation helper that derives query metadata from canonical rollout items.
///
/// This is private implementation state for [`crate::LocalThreadStore`]. Other stores receive raw
/// history and may derive their own metadata, or none at all.
/// Stores may retain this as private per-thread state when they need the local metadata projection.
/// Other stores receive raw history and may derive different metadata, or none at all.
pub(crate) struct ThreadMetadataSync {
thread_id: ThreadId,
cwd_seen: bool,