add multi-host ssh exec registry

codex-rs/core/src/codex.rs

@@ -675,6 +675,7 @@ impl Codex {
             mcp_manager.clone(),
             skills_watcher,
             agent_control,
+            environment_manager,
             environment,
         )
         .await
@@ -1603,6 +1604,7 @@ impl Session {
         mcp_manager: Arc<McpManager>,
         skills_watcher: Arc<SkillsWatcher>,
         agent_control: AgentControl,
+        environment_manager: Arc<EnvironmentManager>,
         environment: Option<Arc<Environment>>,
     ) -> anyhow::Result<Arc<Self>> {
         debug!(
@@ -2040,6 +2042,7 @@ impl Session {
             code_mode_service: crate::tools::code_mode::CodeModeService::new(
                 config.js_repl_node_path.clone(),
             ),
+            environment_manager,
             environment,
         };
         services

codex-rs/core/src/codex_delegate.rs

@@ -4,7 +4,6 @@ use std::sync::Arc;
 use async_channel::Receiver;
 use async_channel::Sender;
 use codex_async_utils::OrCancelExt;
-use codex_exec_server::EnvironmentManager;
 use codex_protocol::protocol::ApplyPatchApprovalRequestEvent;
 use codex_protocol::protocol::Event;
 use codex_protocol::protocol::EventMsg;
@@ -78,9 +77,7 @@ pub(crate) async fn run_codex_thread_interactive(
         config,
         auth_manager,
         models_manager,
-        environment_manager: Arc::new(EnvironmentManager::from_environment(
-            parent_ctx.environment.as_deref(),
-        )),
+        environment_manager: Arc::clone(&parent_session.services.environment_manager),
         skills_manager: Arc::clone(&parent_session.services.skills_manager),
         plugins_manager: Arc::clone(&parent_session.services.plugins_manager),
         mcp_manager: Arc::clone(&parent_session.services.mcp_manager),

codex-rs/core/src/codex_tests.rs

@@ -2743,6 +2743,9 @@ async fn session_new_fails_when_zsh_fork_enabled_without_zsh_path() {
         mcp_manager,
         Arc::new(SkillsWatcher::noop()),
         AgentControl::default(),
+        Arc::new(codex_exec_server::EnvironmentManager::new(
+            /*exec_server_url*/ None,
+        )),
         Some(Arc::new(
             codex_exec_server::Environment::create(/*exec_server_url*/ None)
                 .await
@@ -2901,6 +2904,9 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
         code_mode_service: crate::tools::code_mode::CodeModeService::new(
             config.js_repl_node_path.clone(),
         ),
+        environment_manager: Arc::new(codex_exec_server::EnvironmentManager::new(
+            /*exec_server_url*/ None,
+        )),
         environment: Some(Arc::clone(&environment)),
     };
     let js_repl = Arc::new(JsReplHandle::with_node_path(
@@ -3746,6 +3752,9 @@ pub(crate) async fn make_session_and_context_with_dynamic_tools_and_rx(
         code_mode_service: crate::tools::code_mode::CodeModeService::new(
             config.js_repl_node_path.clone(),
         ),
+        environment_manager: Arc::new(codex_exec_server::EnvironmentManager::new(
+            /*exec_server_url*/ None,
+        )),
         environment: Some(Arc::clone(&environment)),
     };
     let js_repl = Arc::new(JsReplHandle::with_node_path(

codex-rs/core/src/state/service.rs

@@ -16,6 +16,7 @@ use crate::tools::sandboxing::ApprovalStore;
 use crate::unified_exec::UnifiedExecProcessManager;
 use codex_analytics::AnalyticsEventsClient;
 use codex_exec_server::Environment;
+use codex_exec_server::EnvironmentManager;
 use codex_hooks::Hooks;
 use codex_login::AuthManager;
 use codex_mcp::McpConnectionManager;
@@ -59,5 +60,6 @@ pub(crate) struct SessionServices {
     /// Session-scoped model client shared across turns.
     pub(crate) model_client: ModelClient,
     pub(crate) code_mode_service: CodeModeService,
+    pub(crate) environment_manager: Arc<EnvironmentManager>,
     pub(crate) environment: Option<Arc<Environment>>,
 }

codex-rs/core/src/tools/handlers/unified_exec.rs

@@ -44,6 +44,8 @@ pub(crate) struct ExecCommandArgs {
     #[serde(default)]
     shell: Option<String>,
     #[serde(default)]
+    host_id: Option<String>,
+    #[serde(default)]
     login: Option<bool>,
     #[serde(default = "default_tty")]
     tty: bool,
@@ -176,13 +178,6 @@ impl ToolHandler for UnifiedExecHandler {
             }
         };
 
-        let Some(environment) = turn.environment.as_ref() else {
-            return Err(FunctionCallError::RespondToModel(
-                "unified exec is unavailable in this session".to_string(),
-            ));
-        };
-        let fs = environment.get_filesystem();
-
         let manager: &UnifiedExecProcessManager = &session.services.unified_exec_manager;
         let context = UnifiedExecContext::new(session.clone(), turn.clone(), call_id.clone());
 
@@ -190,7 +185,32 @@ impl ToolHandler for UnifiedExecHandler {
             "exec_command" => {
                 let cwd = resolve_workdir_base_path(&arguments, &context.turn.cwd)?;
                 let args: ExecCommandArgs = parse_arguments_with_base_path(&arguments, &cwd)?;
-                let workdir = context.turn.resolve_path(args.workdir.clone());
+                let environment = resolve_exec_environment(
+                    session.as_ref(),
+                    turn.as_ref(),
+                    args.host_id.as_deref(),
+                )
+                .await?;
+                let default_host_cwd = if environment.is_remote() && args.workdir.is_none() {
+                    environment
+                        .default_cwd()
+                        .map(|cwd| {
+                            codex_utils_absolute_path::AbsolutePathBuf::from_absolute_path(cwd)
+                                .map_err(|err| {
+                                    FunctionCallError::RespondToModel(format!(
+                                        "exec host default cwd `{}` is not usable on this client: {err}",
+                                        cwd.display()
+                                    ))
+                                })
+                        })
+                        .transpose()?
+                } else {
+                    None
+                };
+                let fs = environment.get_filesystem();
+                let workdir = default_host_cwd
+                    .clone()
+                    .unwrap_or_else(|| context.turn.resolve_path(args.workdir.clone()));
                 maybe_emit_implicit_skill_invocation(
                     session.as_ref(),
                     context.turn.as_ref(),
@@ -199,14 +219,20 @@ impl ToolHandler for UnifiedExecHandler {
                 )
                 .await;
                 let process_id = manager.allocate_process_id().await;
-                let command = get_command(
+                let command = get_command_with_shell_override(
                     &args,
                     session.user_shell(),
                     &turn.tools_config.unified_exec_shell_mode,
                     turn.tools_config.allow_login_shell,
+                    if environment.is_remote() {
+                        environment.default_shell()
+                    } else {
+                        None
+                    },
                 )
                 .map_err(FunctionCallError::RespondToModel)?;
                 let command_for_display = codex_shell_command::parse_command::shlex_join(&command);
+                let request_host_id = args.host_id.clone();
 
                 let ExecCommandArgs {
                     workdir,
@@ -253,7 +279,9 @@ impl ToolHandler for UnifiedExecHandler {
 
                 let workdir = workdir.filter(|value| !value.is_empty());
 
-                let workdir = workdir.map(|dir| context.turn.resolve_path(Some(dir)));
+                let workdir = workdir
+                    .map(|dir| context.turn.resolve_path(Some(dir)))
+                    .or(default_host_cwd);
                 let cwd = workdir.clone().unwrap_or(cwd);
                 let normalized_additional_permissions = match implicit_granted_permissions(
                     sandbox_permissions,
@@ -312,6 +340,8 @@ impl ToolHandler for UnifiedExecHandler {
                     .exec_command(
                         ExecCommandRequest {
                             command,
+                            environment,
+                            host_id: request_host_id,
                             process_id,
                             yield_time_ms,
                             max_output_tokens,
@@ -379,11 +409,51 @@ fn emit_unified_exec_tty_metric(session_telemetry: &SessionTelemetry, tty: bool)
     );
 }
 
+async fn resolve_exec_environment(
+    session: &crate::codex::Session,
+    turn: &crate::codex::TurnContext,
+    host_id: Option<&str>,
+) -> Result<Arc<codex_exec_server::Environment>, FunctionCallError> {
+    let environment = match host_id {
+        Some(host_id) => session
+            .services
+            .environment_manager
+            .current_for_host(host_id)
+            .await
+            .map_err(|err| {
+                FunctionCallError::RespondToModel(format!(
+                    "failed to resolve exec host `{host_id}`: {err}"
+                ))
+            })?,
+        None => turn.environment.clone(),
+    };
+
+    environment.ok_or_else(|| {
+        FunctionCallError::RespondToModel("unified exec is unavailable in this session".to_string())
+    })
+}
+
 pub(crate) fn get_command(
     args: &ExecCommandArgs,
     session_shell: Arc<Shell>,
     shell_mode: &UnifiedExecShellMode,
     allow_login_shell: bool,
+) -> Result<Vec<String>, String> {
+    get_command_with_shell_override(
+        args,
+        session_shell,
+        shell_mode,
+        allow_login_shell,
+        /*shell_override*/ None,
+    )
+}
+
+fn get_command_with_shell_override(
+    args: &ExecCommandArgs,
+    session_shell: Arc<Shell>,
+    shell_mode: &UnifiedExecShellMode,
+    allow_login_shell: bool,
+    shell_override: Option<&str>,
 ) -> Result<Vec<String>, String> {
     let use_login_shell = match args.login {
         Some(true) if !allow_login_shell => {
@@ -397,7 +467,8 @@ pub(crate) fn get_command(
 
     match shell_mode {
         UnifiedExecShellMode::Direct => {
-            let model_shell = args.shell.as_ref().map(|shell_str| {
+            let selected_shell = args.shell.as_deref().or(shell_override);
+            let model_shell = selected_shell.map(|shell_str| {
                 let mut shell = get_shell_by_model_provided_path(&PathBuf::from(shell_str));
                 shell.shell_snapshot = crate::shell::empty_shell_snapshot_receiver();
                 shell

codex-rs/core/src/tools/runtimes/unified_exec.rs

@@ -44,12 +44,15 @@ use codex_tools::UnifiedExecShellMode;
 use codex_utils_absolute_path::AbsolutePathBuf;
 use futures::future::BoxFuture;
 use std::collections::HashMap;
+use std::sync::Arc;
 
 /// Request payload used by the unified-exec runtime after approvals and
 /// sandbox preferences have been resolved for the current turn.
 #[derive(Clone, Debug)]
 pub struct UnifiedExecRequest {
     pub command: Vec<String>,
+    pub environment: Arc<codex_exec_server::Environment>,
+    pub host_id: Option<String>,
     pub process_id: i32,
     pub cwd: AbsolutePathBuf,
     pub env: HashMap<String, String>,
@@ -69,6 +72,7 @@ pub struct UnifiedExecRequest {
 #[derive(serde::Serialize, Clone, Debug, Eq, PartialEq, Hash)]
 pub struct UnifiedExecApprovalKey {
     pub command: Vec<String>,
+    pub host_id: Option<String>,
     pub cwd: AbsolutePathBuf,
     pub tty: bool,
     pub sandbox_permissions: SandboxPermissions,
@@ -108,6 +112,7 @@ impl Approvable<UnifiedExecRequest> for UnifiedExecRuntime<'_> {
     fn approval_keys(&self, req: &UnifiedExecRequest) -> Vec<Self::ApprovalKey> {
         vec![UnifiedExecApprovalKey {
             command: canonicalize_command_for_approval(&req.command),
+            host_id: req.host_id.clone(),
             cwd: req.cwd.clone(),
             tty: req.tty,
             sandbox_permissions: req.sandbox_permissions,
@@ -202,11 +207,7 @@ impl<'a> ToolRuntime<UnifiedExecRequest, UnifiedExecProcess> for UnifiedExecRunt
     ) -> Result<UnifiedExecProcess, ToolError> {
         let base_command = &req.command;
         let session_shell = ctx.session.user_shell();
-        let environment_is_remote = ctx
-            .turn
-            .environment
-            .as_ref()
-            .is_some_and(|environment| environment.is_remote());
+        let environment_is_remote = req.environment.is_remote();
         let command = if environment_is_remote {
             base_command.to_vec()
         } else {
@@ -249,12 +250,7 @@ impl<'a> ToolRuntime<UnifiedExecRequest, UnifiedExecProcess> for UnifiedExecRunt
             .await?
             {
                 Some(prepared) => {
-                    let Some(environment) = ctx.turn.environment.as_ref() else {
-                        return Err(ToolError::Rejected(
-                            "exec_command is unavailable in this session".to_string(),
-                        ));
-                    };
-                    if environment.is_remote() {
+                    if req.environment.is_remote() {
                         return Err(ToolError::Rejected(
                             "unified_exec zsh-fork is not supported when exec_server_url is configured".to_string(),
                         ));
@@ -266,7 +262,7 @@ impl<'a> ToolRuntime<UnifiedExecRequest, UnifiedExecProcess> for UnifiedExecRunt
                             &prepared.exec_request,
                             req.tty,
                             prepared.spawn_lifecycle,
-                            environment.as_ref(),
+                            req.environment.as_ref(),
                         )
                         .await
                         .map_err(|err| match err {
@@ -296,18 +292,13 @@ impl<'a> ToolRuntime<UnifiedExecRequest, UnifiedExecProcess> for UnifiedExecRunt
         let exec_env = attempt
             .env_for(command, options, req.network.as_ref())
             .map_err(|err| ToolError::Codex(err.into()))?;
-        let Some(environment) = ctx.turn.environment.as_ref() else {
-            return Err(ToolError::Rejected(
-                "exec_command is unavailable in this session".to_string(),
-            ));
-        };
         self.manager
             .open_session_with_exec_env(
                 req.process_id,
                 &exec_env,
                 req.tty,
                 Box::new(NoopSpawnLifecycle),
-                environment.as_ref(),
+                req.environment.as_ref(),
             )
             .await
             .map_err(|err| match err {

codex-rs/core/src/unified_exec/mod.rs

@@ -27,6 +27,7 @@ use std::collections::HashSet;
 use std::sync::Arc;
 use std::sync::Weak;
 
+use codex_exec_server::Environment;
 use codex_network_proxy::NetworkProxy;
 use codex_protocol::models::PermissionProfile;
 use codex_utils_absolute_path::AbsolutePathBuf;
@@ -88,6 +89,8 @@ impl UnifiedExecContext {
 #[derive(Debug)]
 pub(crate) struct ExecCommandRequest {
     pub command: Vec<String>,
+    pub environment: Arc<Environment>,
+    pub host_id: Option<String>,
     pub process_id: i32,
     pub yield_time_ms: u64,
     pub max_output_tokens: Option<usize>,

codex-rs/core/src/unified_exec/process_manager.rs

@@ -677,6 +677,8 @@ impl UnifiedExecProcessManager {
             .await;
         let req = UnifiedExecToolRequest {
             command: request.command.clone(),
+            environment: Arc::clone(&request.environment),
+            host_id: request.host_id.clone(),
             process_id: request.process_id,
             cwd,
             env,

codex-rs/exec-server/src/client.rs

@@ -5,10 +5,13 @@ use std::time::Duration;
 use arc_swap::ArcSwap;
 use codex_app_server_protocol::JSONRPCNotification;
 use serde_json::Value;
+use tokio::io::AsyncRead;
+use tokio::io::AsyncWrite;
 use tokio::sync::Mutex;
 use tokio::sync::watch;
 
 use tokio::time::timeout;
+use tokio_tungstenite::client_async;
 use tokio_tungstenite::connect_async;
 use tracing::debug;
 
@@ -187,6 +190,31 @@ impl ExecServerClient {
         .await
     }
 
+    pub async fn connect_websocket_stream<S>(
+        websocket_url: String,
+        stream: S,
+        options: ExecServerClientConnectOptions,
+    ) -> Result<Self, ExecServerError>
+    where
+        S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
+    {
+        let (stream, _) = client_async(websocket_url.as_str(), stream)
+            .await
+            .map_err(|source| ExecServerError::WebSocketConnect {
+                url: websocket_url.clone(),
+                source,
+            })?;
+
+        Self::connect(
+            JsonRpcConnection::from_websocket(
+                stream,
+                format!("exec-server websocket stream {websocket_url}"),
+            ),
+            options,
+        )
+        .await
+    }
+
     pub async fn initialize(
         &self,
         options: ExecServerClientConnectOptions,
@@ -209,7 +237,7 @@ impl ExecServerClient {
                     },
                 )
                 .await?;
-            {
+            if !response.session_id.is_empty() {
                 let mut session_id = self
                     .inner
                     .session_id

codex-rs/exec-server/src/lib.rs

@@ -17,9 +17,15 @@ pub use client::ExecServerClient;
 pub use client::ExecServerError;
 pub use client_api::ExecServerClientConnectOptions;
 pub use client_api::RemoteExecServerConnectArgs;
+pub use environment::CODEX_EXEC_SERVER_DEFAULT_HOST_ENV_VAR;
+pub use environment::CODEX_EXEC_SERVER_SSH_HOSTS_ENV_VAR;
 pub use environment::CODEX_EXEC_SERVER_URL_ENV_VAR;
+pub use environment::DEFAULT_HOST_ALIAS;
 pub use environment::Environment;
 pub use environment::EnvironmentManager;
+pub use environment::HostConfig;
+pub use environment::HostConnection;
+pub use environment::LOCAL_HOST_ID;
 pub use file_system::CopyOptions;
 pub use file_system::CreateDirectoryOptions;
 pub use file_system::ExecutorFileSystem;

codex-rs/exec-server/src/protocol.rs

@@ -53,6 +53,7 @@ pub struct InitializeParams {
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct InitializeResponse {
+    #[serde(default, alias = "session_id")]
     pub session_id: String,
 }
 
@@ -293,3 +294,15 @@ mod base64_bytes {
             .map_err(serde::de::Error::custom)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::InitializeResponse;
+
+    #[test]
+    fn initialize_response_accepts_legacy_empty_payload() {
+        let response: InitializeResponse =
+            serde_json::from_str("{}").expect("legacy initialize response should deserialize");
+        assert!(response.session_id.is_empty());
+    }
+}

codex-rs/tools/src/local_tool.rs

@@ -35,6 +35,13 @@ pub fn create_exec_command_tool(options: CommandToolOptions) -> ToolSpec {
                 "Shell binary to launch. Defaults to the user's default shell.".to_string(),
             )),
         ),
+        (
+            "host_id".to_string(),
+            JsonSchema::string(Some(
+                "Optional registered exec host to run the command on; omit for the default host."
+                    .to_string(),
+            )),
+        ),
         (
             "tty".to_string(),
             JsonSchema::boolean(Some(

codex-rs/tools/src/local_tool_tests.rs

@@ -126,6 +126,13 @@ fn exec_command_tool_matches_expected_spec() {
                     "Shell binary to launch. Defaults to the user's default shell.".to_string(),
                 )),
         ),
+        (
+            "host_id".to_string(),
+            JsonSchema::string(Some(
+                    "Optional registered exec host to run the command on; omit for the default host."
+                        .to_string(),
+                )),
+        ),
         (
             "tty".to_string(),
             JsonSchema::boolean(Some(