From ec3b1ca18f25386d7224c06a74bab6eeef407c6e Mon Sep 17 00:00:00 2001 From: Adam Van Ymeren Date: Mon, 15 Jun 2026 12:53:25 -0700 Subject: [PATCH] [app-server] allow custom in-process thread stores --- codex-rs/app-server/src/in_process.rs | 105 ++++++++++++++++-- codex-rs/app-server/src/lib.rs | 1 + codex-rs/app-server/src/message_processor.rs | 7 +- .../src/message_processor_tracing_tests.rs | 1 + 4 files changed, 105 insertions(+), 9 deletions(-) diff --git a/codex-rs/app-server/src/in_process.rs b/codex-rs/app-server/src/in_process.rs index 938d539a9f69..1123aed8a4c9 100644 --- a/codex-rs/app-server/src/in_process.rs +++ b/codex-rs/app-server/src/in_process.rs @@ -88,6 +88,7 @@ use codex_login::AuthManager; use codex_protocol::protocol::SessionSource; pub use codex_rollout::StateDbHandle; pub use codex_state::log_db::LogDbLayer; +use codex_thread_store::ThreadStore; use tokio::sync::mpsc; use tokio::sync::oneshot; use tokio::time::timeout; @@ -150,6 +151,23 @@ pub struct InProcessStartArgs { pub channel_capacity: usize, } +/// Optional dependencies and behavior overrides for an in-process runtime. +/// +/// Use [`InProcessStartOptions::default`] to preserve the standard app-server +/// startup behavior. +#[derive(Default)] +pub struct InProcessStartOptions { + thread_store: Option>, +} + +impl InProcessStartOptions { + /// Uses `thread_store` instead of the store derived from the runtime config. + pub fn with_thread_store(mut self, thread_store: Arc) -> Self { + self.thread_store = Some(thread_store); + self + } +} + /// Event emitted from the app-server to the in-process client. /// /// [`Lagged`](Self::Lagged) is a transport health marker, not an application @@ -350,8 +368,16 @@ impl InProcessClientHandle { /// the handle, so callers receive a ready-to-use runtime. If initialize fails, /// the runtime is shut down and an `InvalidData` error is returned. pub async fn start(args: InProcessStartArgs) -> IoResult { + start_with_options(args, InProcessStartOptions::default()).await +} + +/// Starts an in-process app-server runtime with explicit dependency overrides. +pub async fn start_with_options( + args: InProcessStartArgs, + options: InProcessStartOptions, +) -> IoResult { let initialize = args.initialize.clone(); - let client = start_uninitialized(args).await?; + let client = start_uninitialized(args, options).await?; let initialize_response = client .request(ClientRequest::Initialize { @@ -371,7 +397,10 @@ pub async fn start(args: InProcessStartArgs) -> IoResult Ok(client) } -async fn start_uninitialized(args: InProcessStartArgs) -> IoResult { +async fn start_uninitialized( + args: InProcessStartArgs, + options: InProcessStartOptions, +) -> IoResult { let channel_capacity = args.channel_capacity.max(1); let installation_id = resolve_installation_id(&args.config.codex_home).await?; let (client_tx, mut client_rx) = mpsc::channel::(channel_capacity); @@ -440,6 +469,7 @@ async fn start_uninitialized(args: InProcessStartArgs) -> IoResult InProcessClientHandle { + ) -> (TempDir, InProcessStartArgs) { let codex_home = TempDir::new().expect("temp dir"); let config = Arc::new(build_test_config(codex_home.path()).await); - let state_db = codex_rollout::state_db::try_init(config.as_ref()) - .await - .expect("state db should initialize for in-process test"); let args = InProcessStartArgs { arg0_paths: Arg0DispatchPaths::default(), config, @@ -779,7 +808,7 @@ mod tests { thread_config_loader: Arc::new(codex_config::NoopThreadConfigLoader), feedback: CodexFeedback::new(), log_db: None, - state_db: Some(state_db), + state_db: None, environment_manager: Arc::new(EnvironmentManager::default_for_tests()), config_warnings: Vec::new(), session_source, @@ -794,6 +823,27 @@ mod tests { }, channel_capacity, }; + (codex_home, args) + } + + async fn start_test_client_with_capacity_and_options( + session_source: SessionSource, + channel_capacity: usize, + options: InProcessStartOptions, + ) -> InProcessClientHandle { + let (codex_home, args) = build_test_start_args(session_source, channel_capacity).await; + let mut client = start_with_options(args, options) + .await + .expect("in-process runtime should start"); + client._test_codex_home = Some(codex_home); + client + } + + async fn start_test_client_with_capacity( + session_source: SessionSource, + channel_capacity: usize, + ) -> InProcessClientHandle { + let (codex_home, args) = build_test_start_args(session_source, channel_capacity).await; let mut client = start(args).await.expect("in-process runtime should start"); client._test_codex_home = Some(codex_home); client @@ -852,6 +902,45 @@ mod tests { } } + #[tokio::test] + async fn in_process_start_uses_injected_thread_store() { + let thread_store = Arc::new(InMemoryThreadStore::default()); + let client = start_test_client_with_capacity_and_options( + SessionSource::Cli, + DEFAULT_IN_PROCESS_CHANNEL_CAPACITY, + InProcessStartOptions::default().with_thread_store(thread_store.clone()), + ) + .await; + + client + .request(ClientRequest::ThreadList { + request_id: RequestId::Integer(3), + params: ThreadListParams { + cursor: None, + limit: Some(1), + sort_key: None, + sort_direction: None, + model_providers: Some(Vec::new()), + source_kinds: None, + archived: None, + cwd: None, + use_state_db_only: false, + search_term: None, + parent_thread_id: None, + ancestor_thread_id: None, + }, + }) + .await + .expect("request transport should work") + .expect("thread/list should succeed"); + + assert_eq!(thread_store.calls().await.list_threads, 1); + client + .shutdown() + .await + .expect("in-process runtime should shutdown cleanly"); + } + #[tokio::test] async fn in_process_start_clamps_zero_channel_capacity() { let client = diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index beed35848b31..a3265dc10254 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -901,6 +901,7 @@ pub async fn run_main_with_transport_options( rpc_transport: analytics_rpc_transport(&transport), remote_control_handle: Some(remote_control_handle.clone()), plugin_startup_tasks: runtime_options.plugin_startup_tasks, + thread_store: None, })); let mut thread_created_rx = processor.thread_created_receiver(); let mut running_turn_count_rx = processor.subscribe_running_assistant_turn_count(); diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 41ecbfe5613e..e21aa76cce71 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -84,6 +84,7 @@ use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::W3cTraceContext; use codex_rollout::StateDbHandle; use codex_state::log_db::LogDbLayer; +use codex_thread_store::ThreadStore; use tokio::sync::Mutex; use tokio::sync::Semaphore; use tokio::sync::broadcast; @@ -303,6 +304,7 @@ pub(crate) struct MessageProcessorArgs { pub(crate) rpc_transport: AppServerRpcTransport, pub(crate) remote_control_handle: Option, pub(crate) plugin_startup_tasks: crate::PluginStartupTasks, + pub(crate) thread_store: Option>, } impl MessageProcessor { @@ -326,6 +328,7 @@ impl MessageProcessor { rpc_transport, remote_control_handle, plugin_startup_tasks, + thread_store, } = args; auth_manager.set_external_auth(Arc::new(ExternalAuthRefreshBridge { outgoing: outgoing.clone(), @@ -334,7 +337,9 @@ impl MessageProcessor { // The thread store is intentionally process-scoped. Config reloads can // affect per-thread behavior, but they must not move newly started, // resumed, or forked threads to a different persistence backend/root. - let thread_store = codex_core::thread_store_from_config(config.as_ref(), state_db.clone()); + let thread_store = thread_store.unwrap_or_else(|| { + codex_core::thread_store_from_config(config.as_ref(), state_db.clone()) + }); let environment_manager_for_requests = Arc::clone(&environment_manager); let environment_manager_for_extensions = Arc::clone(&environment_manager); let restriction_product = session_source.restriction_product(); diff --git a/codex-rs/app-server/src/message_processor_tracing_tests.rs b/codex-rs/app-server/src/message_processor_tracing_tests.rs index 7c4ae8c98563..db81409a0c0e 100644 --- a/codex-rs/app-server/src/message_processor_tracing_tests.rs +++ b/codex-rs/app-server/src/message_processor_tracing_tests.rs @@ -267,6 +267,7 @@ async fn build_test_processor( rpc_transport: AppServerRpcTransport::Stdio, remote_control_handle: None, plugin_startup_tasks: crate::PluginStartupTasks::Start, + thread_store: None, })); (processor, outgoing_rx) }