From 4f3db8982f80ad9bb78936e033bc19be26c47830 Mon Sep 17 00:00:00 2001 From: Guinness Chen Date: Thu, 26 Mar 2026 11:11:41 -0700 Subject: [PATCH] Refactor thread metadata projection --- .../app-server/src/bespoke_event_handling.rs | 6 +- .../app-server/src/codex_message_processor.rs | 245 +++++++++++++----- codex-rs/state/src/runtime/threads.rs | 53 ++++ 3 files changed, 235 insertions(+), 69 deletions(-) diff --git a/codex-rs/app-server/src/bespoke_event_handling.rs b/codex-rs/app-server/src/bespoke_event_handling.rs index 6700ffdf12..54a5432c9f 100644 --- a/codex-rs/app-server/src/bespoke_event_handling.rs +++ b/codex-rs/app-server/src/bespoke_event_handling.rs @@ -1,8 +1,8 @@ use crate::codex_message_processor::ApiVersion; use crate::codex_message_processor::merge_thread_metadata_from_state_db_context; +use crate::codex_message_processor::project_thread_from_summary; use crate::codex_message_processor::read_rollout_items_from_rollout; use crate::codex_message_processor::read_summary_from_rollout; -use crate::codex_message_processor::summary_to_thread; use crate::error_code::INTERNAL_ERROR_CODE; use crate::error_code::INVALID_REQUEST_ERROR_CODE; use crate::outgoing_message::ClientRequestResult; @@ -1810,7 +1810,9 @@ pub(crate) async fn apply_bespoke_event_handling( .await { Ok(summary) => { - let mut thread = summary_to_thread(summary); + let mut thread = project_thread_from_summary( + summary, /*metadata*/ None, /*persisted_metadata*/ None, + ); merge_thread_metadata_from_state_db_context( &mut thread, conversation.state_db().as_ref(), diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index 62ebfdebb6..fdf320e802 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -2102,20 +2102,12 @@ impl CodexMessageProcessor { otel.name = "app_server.thread_start.config_snapshot", )) .await; - let mut thread = build_thread_from_snapshot( + let started_thread_state_db = started_thread.state_db(); + if let Err(message) = persist_materialized_thread_metadata( + started_thread_state_db.as_ref(), + session_configured.rollout_path.as_deref(), thread_id, &config_snapshot, - session_configured.rollout_path.clone(), - ); - thread.metadata = thread_metadata.clone(); - if let (Some(state_db_ctx), Some(rollout_path)) = ( - started_thread.state_db(), - session_configured.rollout_path.as_ref(), - ) && let Err(message) = persist_thread_metadata_from_snapshot( - state_db_ctx.as_ref(), - thread_id, - rollout_path.as_path(), - &config_snapshot, &thread_metadata, ) .await @@ -2133,6 +2125,13 @@ impl CodexMessageProcessor { .await; return; } + let mut thread = project_thread_from_snapshot( + thread_id, + &config_snapshot, + session_configured.rollout_path.clone(), + Some(&thread_metadata), + /*persisted_metadata*/ None, + ); // Auto-attach a thread listener when starting a thread. Self::log_listener_attach_result( @@ -2669,8 +2668,7 @@ impl CodexMessageProcessor { return; }; - let mut thread = summary_to_thread(summary_from_thread_metadata(&persisted_metadata)); - merge_persisted_thread_metadata(&mut thread, &persisted_metadata); + let mut thread = project_thread_from_persisted_metadata(&persisted_metadata); self.attach_thread_name(thread_uuid, &mut thread).await; thread.status = resolve_thread_status( self.thread_watch_manager @@ -2970,15 +2968,17 @@ impl CodexMessageProcessor { .mark_unarchived(thread_id, restored_path.as_path()) .await; } - let mut thread = + let summary = read_summary_from_rollout(restored_path.as_path(), fallback_provider.as_str()) .await - .map(summary_to_thread) .map_err(|err| JSONRPCErrorError { code: INTERNAL_ERROR_CODE, message: format!("failed to read unarchived thread: {err}"), data: None, })?; + let mut thread = project_thread_from_summary( + summary, /*metadata*/ None, /*persisted_metadata*/ None, + ); merge_thread_metadata_from_state_db_context( &mut thread, state_db_ctx.as_ref(), @@ -3235,25 +3235,31 @@ impl CodexMessageProcessor { return; } }; - let mut threads = Vec::with_capacity(summaries.len()); let mut thread_ids = HashSet::with_capacity(summaries.len()); - let mut status_ids = Vec::with_capacity(summaries.len()); + let mut summary_items = Vec::with_capacity(summaries.len()); let state_db_ctx = get_state_db(&self.config).await; for summary in summaries { let conversation_id = summary.conversation_id; thread_ids.insert(conversation_id); - - let mut thread = summary_to_thread(summary); - if let Some(persisted_metadata) = - read_thread_metadata_from_state_db_context_by_thread_id( - state_db_ctx.as_ref(), - conversation_id, - ) - .await - { - merge_persisted_thread_metadata(&mut thread, &persisted_metadata); - } + summary_items.push((conversation_id, summary)); + } + let persisted_metadata_by_thread_id = read_thread_metadata_from_state_db_context_by_ids( + state_db_ctx.as_ref(), + &summary_items + .iter() + .map(|(conversation_id, _)| *conversation_id) + .collect::>(), + ) + .await; + let mut threads = Vec::with_capacity(summary_items.len()); + let mut status_ids = Vec::with_capacity(summary_items.len()); + for (conversation_id, summary) in summary_items { + let thread = project_thread_from_summary( + summary, + /*metadata*/ None, + persisted_metadata_by_thread_id.get(&conversation_id), + ); status_ids.push(thread.id.clone()); threads.push((conversation_id, thread)); } @@ -3416,13 +3422,13 @@ impl CodexMessageProcessor { } let mut thread = if let Some(persisted_metadata) = db_metadata.as_ref() { - let mut thread = summary_to_thread(summary_from_thread_metadata(persisted_metadata)); - merge_persisted_thread_metadata(&mut thread, persisted_metadata); - thread + project_thread_from_persisted_metadata(persisted_metadata) } else if let Some(rollout_path) = rollout_path.as_ref() { let fallback_provider = self.config.model_provider_id.as_str(); match read_summary_from_rollout(rollout_path, fallback_provider).await { - Ok(summary) => summary_to_thread(summary), + Ok(summary) => project_thread_from_summary( + summary, /*metadata*/ None, /*persisted_metadata*/ None, + ), Err(err) => { self.send_internal_error( request_id, @@ -3457,7 +3463,13 @@ impl CodexMessageProcessor { if include_turns { rollout_path = loaded_rollout_path.clone(); } - build_thread_from_snapshot(thread_uuid, &config_snapshot, loaded_rollout_path) + project_thread_from_snapshot( + thread_uuid, + &config_snapshot, + loaded_rollout_path, + /*metadata*/ None, + /*persisted_metadata*/ None, + ) }; self.attach_thread_name(thread_uuid, &mut thread).await; @@ -3541,8 +3553,13 @@ impl CodexMessageProcessor { ) { if let Ok(thread) = self.thread_manager.get_thread(thread_id).await { let config_snapshot = thread.config_snapshot().await; - let loaded_thread = - build_thread_from_snapshot(thread_id, &config_snapshot, thread.rollout_path()); + let loaded_thread = project_thread_from_snapshot( + thread_id, + &config_snapshot, + thread.rollout_path(), + /*metadata*/ None, + /*persisted_metadata*/ None, + ); self.thread_watch_manager.upsert_thread(loaded_thread).await; } @@ -4054,10 +4071,12 @@ impl CodexMessageProcessor { } InitialHistory::Forked(items) => { let config_snapshot = thread.config_snapshot().await; - let mut thread = build_thread_from_snapshot( + let mut thread = project_thread_from_snapshot( thread_id, &config_snapshot, Some(rollout_path.into()), + /*metadata*/ None, + persisted_resume_metadata, ); thread.preview = preview_from_rollout_items(items); Ok(thread) @@ -4264,6 +4283,19 @@ impl CodexMessageProcessor { } }; let fork_config_snapshot = forked_thread.config_snapshot().await; + let forked_thread_state_db = forked_thread.state_db(); + if let Err(message) = persist_materialized_thread_metadata( + forked_thread_state_db.as_ref(), + session_configured.rollout_path.as_deref(), + thread_id, + &fork_config_snapshot, + &thread_metadata, + ) + .await + { + self.send_internal_error(request_id, message).await; + return; + } // Auto-attach a conversation listener when forking a thread. Self::log_listener_attach_result( @@ -4288,7 +4320,11 @@ impl CodexMessageProcessor { ) .await { - Ok(summary) => summary_to_thread(summary), + Ok(summary) => project_thread_from_summary( + summary, + Some(&thread_metadata), + /*persisted_metadata*/ None, + ), Err(err) => { self.send_internal_error( request_id, @@ -4303,8 +4339,13 @@ impl CodexMessageProcessor { } } else { // forked thread names do not inherit the source thread name - let mut thread = - build_thread_from_snapshot(thread_id, &fork_config_snapshot, /*path*/ None); + let mut thread = project_thread_from_snapshot( + thread_id, + &fork_config_snapshot, + /*path*/ None, + Some(&thread_metadata), + /*persisted_metadata*/ None, + ); let history_items = match read_rollout_items_from_rollout(rollout_path.as_path()).await { Ok(items) => items, @@ -4333,7 +4374,6 @@ impl CodexMessageProcessor { } thread }; - thread.metadata = thread_metadata.clone(); if let Some(fork_rollout_path) = session_configured.rollout_path.as_ref() && let Err(message) = populate_thread_turns( @@ -4346,20 +4386,6 @@ impl CodexMessageProcessor { self.send_internal_error(request_id, message).await; return; } - if let Some(fork_rollout_path) = session_configured.rollout_path.as_ref() - && let Some(state_db_ctx) = forked_thread.state_db() - && let Err(message) = persist_thread_metadata_from_snapshot( - state_db_ctx.as_ref(), - thread_id, - fork_rollout_path.as_path(), - &fork_config_snapshot, - &thread_metadata, - ) - .await - { - self.send_internal_error(request_id, message).await; - return; - } self.thread_watch_manager .upsert_thread_silently(thread.clone()) @@ -6761,7 +6787,9 @@ impl CodexMessageProcessor { if let Some(rollout_path) = review_thread.rollout_path() { match read_summary_from_rollout(rollout_path.as_path(), fallback_provider).await { Ok(summary) => { - let mut thread = summary_to_thread(summary); + let mut thread = project_thread_from_summary( + summary, /*metadata*/ None, /*persisted_metadata*/ None, + ); self.thread_watch_manager .upsert_thread_silently(thread.clone()) .await; @@ -8151,6 +8179,41 @@ async fn read_thread_metadata_from_state_db_context_by_thread_id( } } +async fn read_thread_metadata_from_state_db_context_by_ids( + state_db_ctx: Option<&StateDbHandle>, + thread_ids: &[ThreadId], +) -> HashMap { + let Some(state_db_ctx) = state_db_ctx else { + return HashMap::new(); + }; + + state_db_ctx + .as_ref() + .get_threads_by_ids(thread_ids) + .await + .unwrap_or_default() +} + +async fn persist_materialized_thread_metadata( + state_db_ctx: Option<&StateDbHandle>, + rollout_path: Option<&Path>, + thread_id: ThreadId, + config_snapshot: &ThreadConfigSnapshot, + thread_metadata: &BTreeMap, +) -> Result<(), String> { + let (Some(state_db_ctx), Some(rollout_path)) = (state_db_ctx, rollout_path) else { + return Ok(()); + }; + persist_thread_metadata_from_snapshot( + state_db_ctx.as_ref(), + thread_id, + rollout_path, + config_snapshot, + thread_metadata, + ) + .await +} + pub(crate) async fn merge_thread_metadata_from_state_db_context( thread: &mut Thread, state_db_ctx: Option<&StateDbHandle>, @@ -8159,10 +8222,51 @@ pub(crate) async fn merge_thread_metadata_from_state_db_context( if let Some(persisted_metadata) = read_thread_metadata_from_state_db_context_by_thread_id(state_db_ctx, thread_id).await { - merge_persisted_thread_metadata(thread, &persisted_metadata); + apply_persisted_mutable_thread_metadata(thread, Some(&persisted_metadata)); } } +pub(crate) fn project_thread_from_summary( + summary: ConversationSummary, + metadata: Option<&BTreeMap>, + persisted_metadata: Option<&ThreadMetadata>, +) -> Thread { + let mut thread = summary_to_thread(summary); + apply_thread_projection_metadata(&mut thread, metadata, persisted_metadata); + thread +} + +fn project_thread_from_snapshot( + thread_id: ThreadId, + config_snapshot: &ThreadConfigSnapshot, + path: Option, + metadata: Option<&BTreeMap>, + persisted_metadata: Option<&ThreadMetadata>, +) -> Thread { + let mut thread = build_thread_from_snapshot(thread_id, config_snapshot, path); + apply_thread_projection_metadata(&mut thread, metadata, persisted_metadata); + thread +} + +fn project_thread_from_persisted_metadata(persisted_metadata: &ThreadMetadata) -> Thread { + project_thread_from_summary( + summary_from_thread_metadata(persisted_metadata), + /*metadata*/ None, + Some(persisted_metadata), + ) +} + +fn apply_thread_projection_metadata( + thread: &mut Thread, + metadata: Option<&BTreeMap>, + persisted_metadata: Option<&ThreadMetadata>, +) { + if let Some(metadata) = metadata { + thread.metadata = metadata.clone(); + } + apply_persisted_mutable_thread_metadata(thread, persisted_metadata); +} + async fn read_summary_from_state_db_by_thread_id( config: &Config, thread_id: ThreadId, @@ -8461,26 +8565,33 @@ async fn load_thread_summary_for_rollout( fallback_provider: &str, persisted_metadata: Option<&ThreadMetadata>, ) -> std::result::Result { - let mut thread = read_summary_from_rollout(rollout_path, fallback_provider) + let summary = read_summary_from_rollout(rollout_path, fallback_provider) .await - .map(summary_to_thread) .map_err(|err| { format!( "failed to load rollout `{}` for thread {thread_id}: {err}", rollout_path.display() ) })?; - if let Some(persisted_metadata) = persisted_metadata { - merge_persisted_thread_metadata(&mut thread, persisted_metadata); - } else if let Some(persisted_metadata) = + let loaded_persisted_metadata = if persisted_metadata.is_some() { + None + } else { read_thread_metadata_from_state_db_by_thread_id(config, thread_id).await - { - merge_persisted_thread_metadata(&mut thread, &persisted_metadata); - } - Ok(thread) + }; + Ok(project_thread_from_summary( + summary, + /*metadata*/ None, + persisted_metadata.or(loaded_persisted_metadata.as_ref()), + )) } -fn merge_persisted_thread_metadata(thread: &mut Thread, persisted_metadata: &ThreadMetadata) { +fn apply_persisted_mutable_thread_metadata( + thread: &mut Thread, + persisted_metadata: Option<&ThreadMetadata>, +) { + let Some(persisted_metadata) = persisted_metadata else { + return; + }; thread.git_info = if persisted_metadata.git_sha.is_none() && persisted_metadata.git_branch.is_none() && persisted_metadata.git_origin_url.is_none() diff --git a/codex-rs/state/src/runtime/threads.rs b/codex-rs/state/src/runtime/threads.rs index 28c31a7c88..2484204239 100644 --- a/codex-rs/state/src/runtime/threads.rs +++ b/codex-rs/state/src/runtime/threads.rs @@ -40,6 +40,59 @@ WHERE id = ? .transpose() } + pub async fn get_threads_by_ids( + &self, + ids: &[ThreadId], + ) -> anyhow::Result> { + if ids.is_empty() { + return Ok(std::collections::HashMap::new()); + } + + let mut builder = QueryBuilder::::new( + r#" +SELECT + id, + rollout_path, + created_at, + updated_at, + source, + agent_nickname, + agent_role, + agent_path, + model_provider, + model, + reasoning_effort, + cwd, + cli_version, + title, + sandbox_policy, + approval_mode, + tokens_used, + first_user_message, + archived_at, + git_sha, + git_branch, + git_origin_url, + metadata_json +FROM threads +WHERE id IN ( + "#, + ); + let mut separated = builder.separated(", "); + for id in ids { + separated.push_bind(id.to_string()); + } + separated.push_unseparated(")"); + + let rows = builder.build().fetch_all(self.pool.as_ref()).await?; + rows.into_iter() + .map(|row| { + let metadata = ThreadRow::try_from_row(&row).and_then(ThreadMetadata::try_from)?; + Ok((metadata.id, metadata)) + }) + .collect() + } + pub async fn get_thread_memory_mode(&self, id: ThreadId) -> anyhow::Result> { let row = sqlx::query("SELECT memory_mode FROM threads WHERE id = ?") .bind(id.to_string())