diff --git a/codex-rs/app-server-client/src/lib.rs b/codex-rs/app-server-client/src/lib.rs index c75915629c..ffe22f11a5 100644 --- a/codex-rs/app-server-client/src/lib.rs +++ b/codex-rs/app-server-client/src/lib.rs @@ -25,6 +25,7 @@ use std::io::Result as IoResult; use std::sync::Arc; use std::time::Duration; +pub use codex_app_server::AppServerRuntimeStorageDeps; pub use codex_app_server::app_server_control_socket_path; pub use codex_app_server::in_process::DEFAULT_IN_PROCESS_CHANNEL_CAPACITY; pub use codex_app_server::in_process::InProcessServerEvent; @@ -489,9 +490,20 @@ impl InProcessAppServerClient { /// internal event queue is saturated later, server requests are rejected /// with overload error instead of being silently dropped. pub async fn start(args: InProcessClientStartArgs) -> IoResult { + Self::start_with_runtime_storage_deps(args, AppServerRuntimeStorageDeps::default()).await + } + + /// Starts the in-process runtime with explicit thread storage backends. + pub async fn start_with_runtime_storage_deps( + args: InProcessClientStartArgs, + runtime_storage_deps: AppServerRuntimeStorageDeps, + ) -> IoResult { let channel_capacity = args.channel_capacity.max(1); - let mut handle = - codex_app_server::in_process::start(args.into_runtime_start_args()).await?; + let mut handle = codex_app_server::in_process::start_with_runtime_storage_deps( + args.into_runtime_start_args(), + runtime_storage_deps, + ) + .await?; let request_sender = handle.sender(); let (command_tx, mut command_rx) = mpsc::channel::(channel_capacity); let (event_tx, event_rx) = mpsc::channel::(channel_capacity); diff --git a/codex-rs/app-server/src/in_process.rs b/codex-rs/app-server/src/in_process.rs index 448a389851..05319e4d33 100644 --- a/codex-rs/app-server/src/in_process.rs +++ b/codex-rs/app-server/src/in_process.rs @@ -50,6 +50,7 @@ use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; use std::time::Duration; +use crate::AppServerRuntimeStorageDeps; use crate::analytics_utils::analytics_events_client_from_config; use crate::config_manager::ConfigManager; use crate::error_code::OVERLOADED_ERROR_CODE; @@ -348,8 +349,19 @@ 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_runtime_storage_deps(args, AppServerRuntimeStorageDeps::default()).await +} + +/// Starts an in-process app-server runtime with explicit thread storage backends. +/// +/// This is the embedded-host path for supplying executor-visible artifact and +/// shell snapshot stores without changing the existing start-argument shape. +pub async fn start_with_runtime_storage_deps( + args: InProcessStartArgs, + runtime_storage_deps: AppServerRuntimeStorageDeps, +) -> IoResult { let initialize = args.initialize.clone(); - let client = start_uninitialized(args).await?; + let client = start_uninitialized(args, runtime_storage_deps).await?; let initialize_response = client .request(ClientRequest::Initialize { @@ -369,7 +381,10 @@ pub async fn start(args: InProcessStartArgs) -> IoResult Ok(client) } -async fn start_uninitialized(args: InProcessStartArgs) -> IoResult { +async fn start_uninitialized( + args: InProcessStartArgs, + runtime_storage_deps: AppServerRuntimeStorageDeps, +) -> 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); @@ -438,6 +453,7 @@ async fn start_uninitialized(args: InProcessStartArgs) -> IoResult, +} + #[allow(clippy::too_many_arguments)] pub async fn run_main_with_transport_options( arg0_paths: Arg0DispatchPaths, @@ -427,6 +438,35 @@ pub async fn run_main_with_transport_options( session_source: SessionSource, auth: AppServerWebsocketAuthSettings, runtime_options: AppServerRuntimeOptions, +) -> IoResult<()> { + run_main_with_transport_options_and_storage_deps( + arg0_paths, + cli_config_overrides, + loader_overrides, + strict_config, + default_analytics_enabled, + transport, + session_source, + auth, + runtime_options, + AppServerRuntimeStorageDeps::default(), + ) + .await +} + +/// Runs app-server with explicit storage backends for managed threads. +#[allow(clippy::too_many_arguments)] +pub async fn run_main_with_transport_options_and_storage_deps( + arg0_paths: Arg0DispatchPaths, + cli_config_overrides: CliConfigOverrides, + loader_overrides: LoaderOverrides, + strict_config: bool, + default_analytics_enabled: bool, + transport: AppServerTransport, + session_source: SessionSource, + auth: AppServerWebsocketAuthSettings, + runtime_options: AppServerRuntimeOptions, + runtime_storage_deps: AppServerRuntimeStorageDeps, ) -> IoResult<()> { let (transport_event_tx, mut transport_event_rx) = mpsc::channel::(CHANNEL_CAPACITY); @@ -815,6 +855,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, + runtime_storage_deps, })); 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 84f22809ee..c63aa2728f 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -272,6 +272,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) runtime_storage_deps: crate::AppServerRuntimeStorageDeps, } impl MessageProcessor { @@ -295,6 +296,7 @@ impl MessageProcessor { rpc_transport, remote_control_handle, plugin_startup_tasks, + runtime_storage_deps, } = args; auth_manager.set_external_auth(Arc::new(ExternalAuthRefreshBridge { outgoing: outgoing.clone(), @@ -305,26 +307,45 @@ impl MessageProcessor { // 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 environment_manager_for_requests = Arc::clone(&environment_manager); + let thread_manager_storage_deps = runtime_storage_deps.thread_manager; let thread_manager = Arc::new_cyclic(|thread_manager| { - ThreadManager::new( - config.as_ref(), + let thread_manager_storage_deps = thread_manager_storage_deps.clone(); + let extensions = thread_extensions( + guardian_agent_spawner(thread_manager.clone()), + app_server_extension_event_sink(outgoing.clone()), auth_manager.clone(), - session_source, - environment_manager, - thread_extensions( - guardian_agent_spawner(thread_manager.clone()), - app_server_extension_event_sink(outgoing.clone()), + ); + let attestation_provider = Some(app_server_attestation_provider( + outgoing.clone(), + thread_state_manager.clone(), + )); + match thread_manager_storage_deps { + Some(storage_deps) => ThreadManager::new_with_storage_deps( + config.as_ref(), auth_manager.clone(), + session_source, + environment_manager, + extensions, + Some(analytics_events_client.clone()), + Arc::clone(&thread_store), + storage_deps, + state_db.clone(), + installation_id, + attestation_provider, ), - Some(analytics_events_client.clone()), - Arc::clone(&thread_store), - state_db.clone(), - installation_id, - Some(app_server_attestation_provider( - outgoing.clone(), - thread_state_manager.clone(), - )), - ) + None => ThreadManager::new( + config.as_ref(), + auth_manager.clone(), + session_source, + environment_manager, + extensions, + Some(analytics_events_client.clone()), + Arc::clone(&thread_store), + state_db.clone(), + installation_id, + attestation_provider, + ), + } }); thread_manager .plugins_manager() 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 771b1fd128..41b49db98f 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, + runtime_storage_deps: crate::AppServerRuntimeStorageDeps::default(), })); (processor, outgoing_rx) } diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 4e3082d9b0..cf72520edd 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -113,6 +113,7 @@ pub use thread_manager::ForkSnapshot; pub use thread_manager::NewThread; pub use thread_manager::StartThreadOptions; pub use thread_manager::ThreadManager; +pub use thread_manager::ThreadManagerStorageDeps; pub use thread_manager::ThreadShutdownReport; pub use thread_manager::build_models_manager; pub use thread_manager::thread_store_from_config; diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 1ffc79ec4c..1039fb0c39 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -177,6 +177,30 @@ pub struct ThreadManager { _test_codex_home_guard: Option, } +/// Storage backends used by [`ThreadManager`] for executor-visible artifacts. +/// +/// Cloud and embedded hosts can provide implementations here without asking +/// core to treat `codex_home` as a remotely mounted filesystem. +#[derive(Clone)] +pub struct ThreadManagerStorageDeps { + /// Stores generated artifacts and materializes executor-readable paths. + pub artifact_store: Arc, + /// Stores executor-local shell snapshots used by shell startup. + pub shell_snapshot_store: Arc, +} + +impl ThreadManagerStorageDeps { + /// Builds the standard local backends rooted under the configured Codex home. + pub fn local(config: &Config) -> Self { + Self { + artifact_store: Arc::new(LocalArtifactStore::from_codex_home(&config.codex_home)), + shell_snapshot_store: Arc::new(LocalShellSnapshotStore::from_codex_home( + &config.codex_home, + )), + } + } +} + pub struct StartThreadOptions { pub config: Config, pub initial_history: InitialHistory, @@ -269,7 +293,7 @@ impl ThreadManager { installation_id: String, attestation_provider: Option>, ) -> Self { - Self::new_with_artifact_store( + Self::new_with_storage_deps( config, auth_manager, session_source, @@ -277,7 +301,7 @@ impl ThreadManager { extensions, analytics_events_client, thread_store, - Arc::new(LocalArtifactStore::from_codex_home(&config.codex_home)), + ThreadManagerStorageDeps::local(config), state_db, installation_id, attestation_provider, @@ -299,7 +323,9 @@ impl ThreadManager { installation_id: String, attestation_provider: Option>, ) -> Self { - Self::new_with_storage_backends( + let mut storage_deps = ThreadManagerStorageDeps::local(config); + storage_deps.artifact_store = artifact_store; + Self::new_with_storage_deps( config, auth_manager, session_source, @@ -307,8 +333,7 @@ impl ThreadManager { extensions, analytics_events_client, thread_store, - artifact_store, - Arc::new(LocalShellSnapshotStore::from_codex_home(&config.codex_home)), + storage_deps, state_db, installation_id, attestation_provider, @@ -330,7 +355,9 @@ impl ThreadManager { installation_id: String, attestation_provider: Option>, ) -> Self { - Self::new_with_storage_backends( + let mut storage_deps = ThreadManagerStorageDeps::local(config); + storage_deps.shell_snapshot_store = shell_snapshot_store; + Self::new_with_storage_deps( config, auth_manager, session_source, @@ -338,8 +365,7 @@ impl ThreadManager { extensions, analytics_events_client, thread_store, - Arc::new(LocalArtifactStore::from_codex_home(&config.codex_home)), - shell_snapshot_store, + storage_deps, state_db, installation_id, attestation_provider, @@ -347,7 +373,8 @@ impl ThreadManager { } #[allow(clippy::too_many_arguments)] - fn new_with_storage_backends( + /// Constructs a thread manager with explicit executor-visible storage backends. + pub fn new_with_storage_deps( config: &Config, auth_manager: Arc, session_source: SessionSource, @@ -355,12 +382,15 @@ impl ThreadManager { extensions: Arc>, analytics_events_client: Option, thread_store: Arc, - artifact_store: Arc, - shell_snapshot_store: Arc, + storage_deps: ThreadManagerStorageDeps, state_db: Option, installation_id: String, attestation_provider: Option>, ) -> Self { + let ThreadManagerStorageDeps { + artifact_store, + shell_snapshot_store, + } = storage_deps; let codex_home = config.codex_home.clone(); let restriction_product = session_source.restriction_product(); let (thread_created_tx, _) = broadcast::channel(THREAD_CREATED_CHANNEL_CAPACITY); diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index a7d27ad56f..99efd9bdb0 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -7,6 +7,8 @@ use crate::installation_id::INSTALLATION_ID_FILENAME; use crate::rollout::RolloutRecorder; use crate::session::session::SessionSettingsUpdate; use crate::session::tests::make_session_and_context; +use crate::shell_snapshot::ShellSnapshotPaths; +use crate::shell_snapshot::ShellSnapshotStoreFuture; use crate::tasks::InterruptedTurnHistoryMarker; use crate::tasks::interrupted_turn_history_marker; use codex_extension_api::empty_extension_registry; @@ -52,6 +54,26 @@ impl ArtifactStore for FakeArtifactStore { } } +struct FakeShellSnapshotStore; + +impl ShellSnapshotStore for FakeShellSnapshotStore { + fn snapshot_paths( + &self, + _session_id: ThreadId, + _shell_type: crate::shell::ShellType, + ) -> ShellSnapshotPaths { + unreachable!("shell snapshot paths should not be requested") + } + + fn cleanup_stale_snapshots( + &self, + _active_session_id: ThreadId, + _state_db: Option, + ) -> ShellSnapshotStoreFuture<'_> { + unreachable!("shell snapshot cleanup should not be requested") + } +} + fn user_msg(text: &str) -> ResponseItem { ResponseItem::Message { id: None, @@ -84,10 +106,11 @@ fn developer_interrupted_marker() -> ResponseItem { } #[tokio::test] -async fn start_thread_uses_injected_artifact_store() { +async fn start_thread_uses_injected_storage_deps() { let config = test_config().await; let artifact_store: Arc = Arc::new(FakeArtifactStore); - let manager = ThreadManager::new_with_artifact_store( + let shell_snapshot_store: Arc = Arc::new(FakeShellSnapshotStore); + let manager = ThreadManager::new_with_storage_deps( &config, AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()), SessionSource::Exec, @@ -95,7 +118,10 @@ async fn start_thread_uses_injected_artifact_store() { empty_extension_registry(), /*analytics_events_client*/ None, thread_store_from_config(&config, /*state_db*/ None), - Arc::clone(&artifact_store), + ThreadManagerStorageDeps { + artifact_store: Arc::clone(&artifact_store), + shell_snapshot_store: Arc::clone(&shell_snapshot_store), + }, /*state_db*/ None, TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, @@ -106,6 +132,10 @@ async fn start_thread_uses_injected_artifact_store() { &thread.thread.codex.session.services.artifact_store, &artifact_store )); + assert!(Arc::ptr_eq( + &thread.thread.codex.session.services.shell_snapshot_store, + &shell_snapshot_store + )); } #[test]