mirror of
https://github.com/openai/codex.git
synced 2026-11-13 22:47:20 +00:00
[codex] Speed up external agent session imports (#26637)
## Why Importing large external-agent session histories currently starts a full live Codex thread for every imported session. This initializes unrelated runtime systems and repeats expensive transcript, metadata, hashing, and ledger work. On a 50-session, 238 MiB fixture, the existing path took roughly 70 seconds to complete the import and 77 seconds end to end. ## What changed - Persist imported sessions directly through `ThreadStore` instead of starting full live threads. - Process imports through a bounded five-session pipeline. - Parse, extract, and hash each source file in one pass. - Move blocking source preparation onto the blocking thread pool. - Reuse prepared content hashes and update the import ledger once per batch. - Avoid metadata readback for newly written rollouts. - Preserve imported conversation history and visible thread metadata. - Keep the implementation out of `codex-core` and avoid changes to the public `ThreadStore` trait. ## Performance For the same 50-session, 238 MiB fixture: | Path | Import completion | End to end | | --- | ---: | ---: | | Existing import | 69.61s | 76.62s | | This change | 5.95s | 6.58s | All 50 sessions imported successfully with no warnings or contention signals. ## Validation - `just test -p codex-external-agent-sessions` - `just test -p codex-app-server external_agent_config_import` - Verified imports do not initialize unrelated required MCP servers. - Verified previously imported source versions are skipped and changed sources can be imported again. - Verified imported rollouts remain readable through thread listing and history APIs.
This commit is contained in:
stefanstokic-oai authored and GitHub committed
1 parent
f1c18df9ae
commit
6d8e12ac42
10 files changed
+793
-394
No files matched your search
@@ -26,53 +26,47 @@ use codex_app_server_protocol::MigrationDetails;
|
||||
use codex_app_server_protocol::PluginsMigration;
|
||||
use codex_app_server_protocol::ServerNotification;
|
||||
use codex_arg0::Arg0DispatchPaths;
|
||||
use codex_core::StartThreadOptions;
|
||||
use codex_core::ThreadManager;
|
||||
use codex_core::config::ConfigOverrides;
|
||||
use codex_external_agent_sessions::ExternalAgentSessionMigration as CoreSessionMigration;
|
||||
use codex_external_agent_sessions::ImportedExternalAgentSession;
|
||||
use codex_external_agent_sessions::PendingSessionImport;
|
||||
use codex_external_agent_sessions::prepare_validated_session_imports;
|
||||
use codex_external_agent_sessions::record_imported_session;
|
||||
use codex_protocol::ThreadId;
|
||||
use codex_protocol::protocol::InitialHistory;
|
||||
use codex_thread_store::ThreadMetadataPatch;
|
||||
use codex_thread_store::ThreadStore;
|
||||
use std::collections::HashSet;
|
||||
use std::path::PathBuf;
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
use super::ConfigRequestProcessor;
|
||||
use super::external_agent_session_import::ExternalAgentSessionImporter;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct ExternalAgentConfigRequestProcessor {
|
||||
outgoing: Arc<OutgoingMessageSender>,
|
||||
codex_home: PathBuf,
|
||||
migration_service: ExternalAgentConfigService,
|
||||
session_import_permits: Arc<Semaphore>,
|
||||
session_importer: ExternalAgentSessionImporter,
|
||||
thread_manager: Arc<ThreadManager>,
|
||||
config_manager: ConfigManager,
|
||||
config_processor: ConfigRequestProcessor,
|
||||
arg0_paths: Arg0DispatchPaths,
|
||||
}
|
||||
|
||||
impl ExternalAgentConfigRequestProcessor {
|
||||
pub(crate) fn new(
|
||||
outgoing: Arc<OutgoingMessageSender>,
|
||||
thread_manager: Arc<ThreadManager>,
|
||||
thread_store: Arc<dyn ThreadStore>,
|
||||
config_manager: ConfigManager,
|
||||
config_processor: ConfigRequestProcessor,
|
||||
arg0_paths: Arg0DispatchPaths,
|
||||
codex_home: PathBuf,
|
||||
) -> Self {
|
||||
let session_importer = ExternalAgentSessionImporter::new(
|
||||
codex_home.clone(),
|
||||
Arc::clone(&thread_manager),
|
||||
thread_store,
|
||||
config_manager,
|
||||
arg0_paths,
|
||||
);
|
||||
Self {
|
||||
outgoing,
|
||||
migration_service: ExternalAgentConfigService::new(codex_home.clone()),
|
||||
codex_home,
|
||||
session_import_permits: Arc::new(Semaphore::new(1)),
|
||||
migration_service: ExternalAgentConfigService::new(codex_home),
|
||||
session_importer,
|
||||
thread_manager,
|
||||
config_manager,
|
||||
config_processor,
|
||||
arg0_paths,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -207,42 +201,12 @@ impl ExternalAgentConfigRequestProcessor {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let session_import_permits = Arc::clone(&self.session_import_permits);
|
||||
let session_processor = self.clone();
|
||||
let session_importer = self.session_importer.clone();
|
||||
let plugin_processor = self.clone();
|
||||
let outgoing = Arc::clone(&self.outgoing);
|
||||
let thread_manager = Arc::clone(&self.thread_manager);
|
||||
tokio::spawn(async move {
|
||||
let session_imports = async move {
|
||||
if !pending_session_imports.is_empty() {
|
||||
let Ok(_session_import_permit) = session_import_permits.acquire_owned().await
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let pending_session_imports = session_processor
|
||||
.prepare_validated_session_imports(pending_session_imports);
|
||||
for pending_session_import in pending_session_imports {
|
||||
match session_processor
|
||||
.import_external_agent_session(pending_session_import.session)
|
||||
.await
|
||||
{
|
||||
Ok(imported_thread_id) => {
|
||||
session_processor.record_imported_session(
|
||||
&pending_session_import.source_path,
|
||||
imported_thread_id,
|
||||
);
|
||||
}
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
error = %error.message,
|
||||
path = %pending_session_import.source_path.display(),
|
||||
"external agent session import failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
let session_imports = session_importer.import_sessions(pending_session_imports);
|
||||
let plugin_imports = async move {
|
||||
for pending_plugin_import in pending_plugin_imports {
|
||||
match plugin_processor
|
||||
@@ -274,65 +238,6 @@ impl ExternalAgentConfigRequestProcessor {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn import_external_agent_session(
|
||||
&self,
|
||||
session: ImportedExternalAgentSession,
|
||||
) -> Result<ThreadId, JSONRPCErrorError> {
|
||||
let ImportedExternalAgentSession {
|
||||
cwd,
|
||||
title,
|
||||
rollout_items,
|
||||
} = session;
|
||||
let config = self
|
||||
.config_manager
|
||||
.load_with_overrides(
|
||||
/*request_overrides*/ None,
|
||||
ConfigOverrides {
|
||||
cwd: Some(PathBuf::from(cwd.to_string_lossy().into_owned())),
|
||||
codex_linux_sandbox_exe: self.arg0_paths.codex_linux_sandbox_exe.clone(),
|
||||
main_execve_wrapper_exe: self.arg0_paths.main_execve_wrapper_exe.clone(),
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await
|
||||
.map_err(|err| {
|
||||
internal_error(format!("failed to load imported session config: {err}"))
|
||||
})?;
|
||||
let environments = self
|
||||
.thread_manager
|
||||
.default_environment_selections(&config.cwd);
|
||||
let imported_thread = self
|
||||
.thread_manager
|
||||
.start_thread_with_options(StartThreadOptions {
|
||||
config,
|
||||
initial_history: InitialHistory::Forked(rollout_items),
|
||||
session_source: None,
|
||||
thread_source: None,
|
||||
dynamic_tools: Vec::new(),
|
||||
metrics_service_name: None,
|
||||
parent_trace: None,
|
||||
environments,
|
||||
})
|
||||
.await
|
||||
.map_err(|err| internal_error(format!("failed to import session: {err}")))?;
|
||||
if let Some(title) = title
|
||||
&& let Some(name) = codex_core::util::normalize_thread_name(&title)
|
||||
{
|
||||
imported_thread
|
||||
.thread
|
||||
.update_thread_metadata(
|
||||
ThreadMetadataPatch {
|
||||
name: Some(Some(name)),
|
||||
..Default::default()
|
||||
},
|
||||
/*include_archived*/ false,
|
||||
)
|
||||
.await
|
||||
.map_err(|err| internal_error(format!("failed to name imported session: {err}")))?;
|
||||
}
|
||||
Ok(imported_thread.thread_id)
|
||||
}
|
||||
|
||||
fn validate_pending_session_imports(
|
||||
&self,
|
||||
params: &ExternalAgentConfigImportParams,
|
||||
@@ -371,24 +276,6 @@ impl ExternalAgentConfigRequestProcessor {
|
||||
Ok(selected_sessions)
|
||||
}
|
||||
|
||||
fn prepare_validated_session_imports(
|
||||
&self,
|
||||
sessions: Vec<CoreSessionMigration>,
|
||||
) -> Vec<PendingSessionImport> {
|
||||
prepare_validated_session_imports(&self.codex_home, sessions)
|
||||
}
|
||||
|
||||
fn record_imported_session(&self, source_path: &std::path::Path, imported_thread_id: ThreadId) {
|
||||
if let Err(err) = record_imported_session(&self.codex_home, source_path, imported_thread_id)
|
||||
{
|
||||
tracing::warn!(
|
||||
error = %err,
|
||||
path = %source_path.display(),
|
||||
"external agent session import ledger update failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
async fn import_external_agent_config(
|
||||
&self,
|
||||
params: ExternalAgentConfigImportParams,
|
||||
|
||||
@@ -0,0 +1,260 @@
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
|
||||
use chrono::Utc;
|
||||
use codex_arg0::Arg0DispatchPaths;
|
||||
use codex_core::ThreadManager;
|
||||
use codex_core::config::ConfigOverrides;
|
||||
use codex_external_agent_sessions::CompletedExternalAgentSessionImport;
|
||||
use codex_external_agent_sessions::ExternalAgentSessionMigration;
|
||||
use codex_external_agent_sessions::ImportedExternalAgentSession;
|
||||
use codex_external_agent_sessions::PendingSessionImport;
|
||||
use codex_external_agent_sessions::prepare_validated_session_import;
|
||||
use codex_external_agent_sessions::record_completed_session_imports;
|
||||
use codex_models_manager::manager::RefreshStrategy;
|
||||
use codex_protocol::ThreadId;
|
||||
use codex_protocol::models::BaseInstructions;
|
||||
use codex_protocol::protocol::MultiAgentVersion;
|
||||
use codex_protocol::protocol::ThreadMemoryMode;
|
||||
use codex_rollout::is_persisted_rollout_item;
|
||||
use codex_thread_store::AppendThreadItemsParams;
|
||||
use codex_thread_store::CreateThreadParams;
|
||||
use codex_thread_store::ThreadMetadataPatch;
|
||||
use codex_thread_store::ThreadPersistenceMetadata;
|
||||
use codex_thread_store::ThreadStore;
|
||||
use codex_thread_store::UpdateThreadMetadataParams;
|
||||
use futures::StreamExt;
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
use crate::config_manager::ConfigManager;
|
||||
|
||||
const SESSION_IMPORT_CONCURRENCY: usize = 5;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(super) struct ExternalAgentSessionImporter {
|
||||
codex_home: PathBuf,
|
||||
permits: Arc<Semaphore>,
|
||||
thread_manager: Arc<ThreadManager>,
|
||||
thread_store: Arc<dyn ThreadStore>,
|
||||
config_manager: ConfigManager,
|
||||
arg0_paths: Arg0DispatchPaths,
|
||||
}
|
||||
|
||||
impl ExternalAgentSessionImporter {
|
||||
pub(super) fn new(
|
||||
codex_home: PathBuf,
|
||||
thread_manager: Arc<ThreadManager>,
|
||||
thread_store: Arc<dyn ThreadStore>,
|
||||
config_manager: ConfigManager,
|
||||
arg0_paths: Arg0DispatchPaths,
|
||||
) -> Self {
|
||||
Self {
|
||||
codex_home,
|
||||
permits: Arc::new(Semaphore::new(1)),
|
||||
thread_manager,
|
||||
thread_store,
|
||||
config_manager,
|
||||
arg0_paths,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn import_sessions(&self, sessions: Vec<ExternalAgentSessionMigration>) {
|
||||
if sessions.is_empty() {
|
||||
return;
|
||||
}
|
||||
let Ok(_permit) = self.permits.acquire().await else {
|
||||
return;
|
||||
};
|
||||
let import_results = futures::stream::iter(sessions)
|
||||
.map(|session| {
|
||||
let importer = self.clone();
|
||||
async move { importer.import_requested_session(session).await }
|
||||
})
|
||||
.buffer_unordered(SESSION_IMPORT_CONCURRENCY);
|
||||
futures::pin_mut!(import_results);
|
||||
|
||||
let mut completed_imports = Vec::new();
|
||||
while let Some(result) = import_results.next().await {
|
||||
match result {
|
||||
Ok(Some(completed_import)) => completed_imports.push(completed_import),
|
||||
Ok(None) => {}
|
||||
Err(failure) => {
|
||||
tracing::warn!(
|
||||
error = %failure.message,
|
||||
path = %failure.source_path.display(),
|
||||
"external agent session import failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
if let Err(err) = record_completed_session_imports(&self.codex_home, completed_imports) {
|
||||
tracing::warn!(
|
||||
error = %err,
|
||||
"external agent session import ledger update failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
async fn import_requested_session(
|
||||
&self,
|
||||
session: ExternalAgentSessionMigration,
|
||||
) -> Result<Option<CompletedExternalAgentSessionImport>, SessionImportFailure> {
|
||||
let source_path = session.path.clone();
|
||||
let Some(pending_import) =
|
||||
self.prepare_session_import(session)
|
||||
.await
|
||||
.map_err(|message| SessionImportFailure {
|
||||
source_path: source_path.clone(),
|
||||
message,
|
||||
})?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let imported_thread_id =
|
||||
self.persist_session(pending_import.session)
|
||||
.await
|
||||
.map_err(|message| SessionImportFailure {
|
||||
source_path: pending_import.source_path.clone(),
|
||||
message,
|
||||
})?;
|
||||
Ok(Some(CompletedExternalAgentSessionImport {
|
||||
source_path: pending_import.source_path,
|
||||
source_content_sha256: pending_import.source_content_sha256,
|
||||
imported_thread_id,
|
||||
}))
|
||||
}
|
||||
|
||||
async fn prepare_session_import(
|
||||
&self,
|
||||
session: ExternalAgentSessionMigration,
|
||||
) -> Result<Option<PendingSessionImport>, String> {
|
||||
let codex_home = self.codex_home.clone();
|
||||
tokio::task::spawn_blocking(move || prepare_validated_session_import(&codex_home, session))
|
||||
.await
|
||||
.map_err(|err| format!("external agent session preparation task failed: {err}"))?
|
||||
.map_err(|err| format!("failed to prepare external agent session: {err}"))
|
||||
}
|
||||
|
||||
async fn persist_session(
|
||||
&self,
|
||||
session: ImportedExternalAgentSession,
|
||||
) -> Result<ThreadId, String> {
|
||||
let ImportedExternalAgentSession {
|
||||
cwd,
|
||||
title,
|
||||
first_user_message,
|
||||
mut rollout_items,
|
||||
} = session;
|
||||
let config = self
|
||||
.config_manager
|
||||
.load_with_overrides(
|
||||
/*request_overrides*/ None,
|
||||
ConfigOverrides {
|
||||
cwd: Some(cwd),
|
||||
codex_linux_sandbox_exe: self.arg0_paths.codex_linux_sandbox_exe.clone(),
|
||||
main_execve_wrapper_exe: self.arg0_paths.main_execve_wrapper_exe.clone(),
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await
|
||||
.map_err(|err| format!("failed to load imported session config: {err}"))?;
|
||||
let models_manager = self.thread_manager.get_models_manager();
|
||||
let model = models_manager
|
||||
.get_default_model(&config.model, RefreshStrategy::Offline)
|
||||
.await;
|
||||
let model_info = models_manager
|
||||
.get_model_info(model.as_str(), &config.to_models_manager_config())
|
||||
.await;
|
||||
let thread_id = ThreadId::new();
|
||||
let source = self.thread_manager.session_source();
|
||||
let cwd = config.cwd.to_path_buf();
|
||||
let model_provider = config.model_provider_id.clone();
|
||||
let memory_mode = if config.memories.generate_memories {
|
||||
ThreadMemoryMode::Enabled
|
||||
} else {
|
||||
ThreadMemoryMode::Disabled
|
||||
};
|
||||
let now = Utc::now();
|
||||
let create_params = CreateThreadParams {
|
||||
thread_id,
|
||||
forked_from_id: None,
|
||||
parent_thread_id: None,
|
||||
source: source.clone(),
|
||||
thread_source: None,
|
||||
base_instructions: BaseInstructions {
|
||||
text: config
|
||||
.base_instructions
|
||||
.clone()
|
||||
.unwrap_or_else(|| model_info.get_model_instructions(config.personality)),
|
||||
},
|
||||
dynamic_tools: Vec::new(),
|
||||
multi_agent_version: Some(MultiAgentVersion::V1),
|
||||
metadata: ThreadPersistenceMetadata {
|
||||
cwd: Some(cwd.clone()),
|
||||
model_provider: model_provider.clone(),
|
||||
memory_mode,
|
||||
},
|
||||
};
|
||||
rollout_items.retain(is_persisted_rollout_item);
|
||||
let title = title
|
||||
.as_deref()
|
||||
.and_then(codex_core::util::normalize_thread_name);
|
||||
let metadata = ThreadMetadataPatch {
|
||||
title,
|
||||
preview: first_user_message.clone(),
|
||||
model_provider: Some(model_provider),
|
||||
created_at: Some(now),
|
||||
updated_at: Some(now),
|
||||
source: Some(source.clone()),
|
||||
thread_source: Some(None),
|
||||
agent_nickname: Some(source.get_nickname()),
|
||||
agent_role: Some(source.get_agent_role()),
|
||||
agent_path: Some(source.get_agent_path().map(Into::into)),
|
||||
cwd: Some(cwd),
|
||||
cli_version: Some(env!("CARGO_PKG_VERSION").to_string()),
|
||||
first_user_message,
|
||||
memory_mode: Some(memory_mode),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
self.thread_store
|
||||
.create_thread(create_params)
|
||||
.await
|
||||
.map_err(|err| format!("failed to import session: {err}"))?;
|
||||
if !rollout_items.is_empty()
|
||||
&& let Err(err) = self
|
||||
.thread_store
|
||||
.append_items(AppendThreadItemsParams {
|
||||
thread_id,
|
||||
items: rollout_items,
|
||||
})
|
||||
.await
|
||||
{
|
||||
let _ = self.thread_store.discard_thread(thread_id).await;
|
||||
return Err(format!("failed to import session: {err}"));
|
||||
}
|
||||
|
||||
self.thread_store
|
||||
.update_thread_metadata(UpdateThreadMetadataParams {
|
||||
thread_id,
|
||||
patch: metadata,
|
||||
include_archived: false,
|
||||
})
|
||||
.await
|
||||
.map_err(|err| format!("failed to update imported session: {err}"))?;
|
||||
self.thread_store
|
||||
.persist_thread(thread_id)
|
||||
.await
|
||||
.map_err(|err| format!("failed to persist imported session: {err}"))?;
|
||||
self.thread_store
|
||||
.shutdown_thread(thread_id)
|
||||
.await
|
||||
.map_err(|err| format!("failed to shutdown imported session: {err}"))?;
|
||||
Ok(thread_id)
|
||||
}
|
||||
}
|
||||
|
||||
struct SessionImportFailure {
|
||||
source_path: PathBuf,
|
||||
message: String,
|
||||
}
|
||||
Reference in new issue
Block a user