🤖 fix: preserve SSH/Docker runtime selection

ammar-agent · ammar-agent · commit d25724d84c94 · 2025-12-11T13:22:19.000-06:00
Fixes a regression where editing SSH host / Docker image could reset the selection back to the project default runtime.\n\nAlso fixes SSHRuntime.resolvePath() so ~ and ~/path expand correctly on the remote host, and adds coverage for resolvePath in runtime integration tests.
diff --git a/src/browser/hooks/useDraftWorkspaceSettings.test.ts b/src/browser/hooks/useDraftWorkspaceSettings.test.ts
@@ -0,0 +1,116 @@
+import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test";
+import { renderHook, act, cleanup, waitFor } from "@testing-library/react";
+import { GlobalWindow } from "happy-dom";
+import { useState } from "react";
+import {
+  getLastDockerImageKey,
+  getLastSshHostKey,
+  getRuntimeKey,
+} from "@/common/constants/storage";
+import { useDraftWorkspaceSettings } from "./useDraftWorkspaceSettings";
+
+// A minimal in-memory persisted-state implementation.
+// We keep it here (rather than relying on real localStorage) so tests remain deterministic.
+const persisted = new Map<string, unknown>();
+
+void mock.module("@/browser/hooks/usePersistedState", () => {
+  return {
+    usePersistedState: <T>(key: string, defaultValue: T) => {
+      const [value, setValue] = useState<T>(() => {
+        return persisted.has(key) ? (persisted.get(key) as T) : defaultValue;
+      });
+
+      const setAndPersist = (next: T) => {
+        persisted.set(key, next);
+        setValue(next);
+      };
+
+      return [value, setAndPersist] as const;
+    },
+  };
+});
+
+void mock.module("@/browser/hooks/useModelLRU", () => ({
+  useModelLRU: () => ({ recentModels: ["test-model"] }),
+}));
+
+void mock.module("@/browser/hooks/useThinkingLevel", () => ({
+  useThinkingLevel: () => ["medium", () => undefined] as const,
+}));
+
+void mock.module("@/browser/contexts/ModeContext", () => ({
+  useMode: () => ["plan", () => undefined] as const,
+}));
+
+describe("useDraftWorkspaceSettings", () => {
+  beforeEach(() => {
+    persisted.clear();
+
+    globalThis.window = new GlobalWindow() as unknown as Window & typeof globalThis;
+    globalThis.document = globalThis.window.document;
+  });
+
+  afterEach(() => {
+    cleanup();
+    globalThis.window = undefined as unknown as Window & typeof globalThis;
+    globalThis.document = undefined as unknown as Document;
+  });
+
+  test("does not reset selected runtime to the default while editing SSH host", async () => {
+    const projectPath = "/tmp/project";
+
+    const { result } = renderHook(() => useDraftWorkspaceSettings(projectPath, ["main"], "main"));
+
+    act(() => {
+      result.current.setSelectedRuntime({ mode: "ssh", host: "dev@host" });
+    });
+
+    await waitFor(() => {
+      expect(result.current.settings.selectedRuntime).toEqual({ mode: "ssh", host: "dev@host" });
+    });
+  });
+
+  test("seeds SSH host from the remembered value when switching modes", async () => {
+    const projectPath = "/tmp/project";
+    persisted.set(getRuntimeKey(projectPath), undefined);
+    persisted.set(getLastSshHostKey(projectPath), "remembered@host");
+
+    const { result } = renderHook(() => useDraftWorkspaceSettings(projectPath, ["main"], "main"));
+
+    act(() => {
+      // Simulate UI switching into ssh mode with an empty field.
+      result.current.setSelectedRuntime({ mode: "ssh", host: "" });
+    });
+
+    await waitFor(() => {
+      expect(result.current.settings.selectedRuntime).toEqual({
+        mode: "ssh",
+        host: "remembered@host",
+      });
+    });
+
+    expect(persisted.get(getLastSshHostKey(projectPath))).toBe("remembered@host");
+  });
+
+  test("seeds Docker image from the remembered value when switching modes", async () => {
+    const projectPath = "/tmp/project";
+    persisted.set(getRuntimeKey(projectPath), undefined);
+    persisted.set(getLastDockerImageKey(projectPath), "ubuntu:22.04");
+
+    const { result } = renderHook(() => useDraftWorkspaceSettings(projectPath, ["main"], "main"));
+
+    act(() => {
+      // Simulate UI switching into docker mode with an empty field.
+      result.current.setSelectedRuntime({ mode: "docker", image: "" });
+    });
+
+    await waitFor(() => {
+      expect(result.current.settings.selectedRuntime).toEqual({
+        mode: "docker",
+        image: "ubuntu:22.04",
+      });
+    });
+
+    expect(persisted.get(getLastDockerImageKey(projectPath))).toBe("ubuntu:22.04");
+  });
+});
diff --git a/src/browser/hooks/useDraftWorkspaceSettings.ts b/src/browser/hooks/useDraftWorkspaceSettings.ts
@@ -1,8 +1,8 @@
-import { useState, useEffect } from "react";
-import { usePersistedState } from "./usePersistedState";
-import { useThinkingLevel } from "./useThinkingLevel";
+import { useState, useEffect, useRef } from "react";
+import { usePersistedState } from "@/browser/hooks/usePersistedState";
+import { useThinkingLevel } from "@/browser/hooks/useThinkingLevel";
 import { useMode } from "@/browser/contexts/ModeContext";
-import { useModelLRU } from "./useModelLRU";
+import { useModelLRU } from "@/browser/hooks/useModelLRU";
 import {
   type RuntimeMode,
   type ParsedRuntime,
@@ -109,6 +109,37 @@ export function useDraftWorkspaceSettings(
     { listener: true }
   );
 
+  // If the default runtime string contains a host/image (e.g. older persisted values like "ssh devbox"),
+  // prefer it as the initial remembered value.
+  useEffect(() => {
+    if (
+      parsedDefault?.mode === RUNTIME_MODE.SSH &&
+      !lastSshHost.trim() &&
+      parsedDefault.host.trim()
+    ) {
+      setLastSshHost(parsedDefault.host);
+    }
+    if (
+      parsedDefault?.mode === RUNTIME_MODE.DOCKER &&
+      !lastDockerImage.trim() &&
+      parsedDefault.image.trim()
+    ) {
+      setLastDockerImage(parsedDefault.image);
+    }
+  }, [
+    projectPath,
+    parsedDefault,
+    lastSshHost,
+    lastDockerImage,
+    setLastSshHost,
+    setLastDockerImage,
+  ]);
+
+  const defaultSshHost =
+    parsedDefault?.mode === RUNTIME_MODE.SSH ? parsedDefault.host : lastSshHost;
+  const defaultDockerImage =
+    parsedDefault?.mode === RUNTIME_MODE.DOCKER ? parsedDefault.image : lastDockerImage;
+
   // Build ParsedRuntime from mode + stored host/image
   // Defined as a function so it can be used in both useState init and useEffect
   const buildRuntimeForMode = (
@@ -132,13 +163,49 @@ export function useDraftWorkspaceSettings(
   // Currently selected runtime for this session (initialized from default)
   // Uses discriminated union: SSH has host, Docker has image
   const [selectedRuntime, setSelectedRuntimeState] = useState<ParsedRuntime>(() =>
-    buildRuntimeForMode(defaultRuntimeMode, lastSshHost, lastDockerImage)
+    buildRuntimeForMode(defaultRuntimeMode, defaultSshHost, defaultDockerImage)
   );
 
-  // Sync selected runtime when default changes (e.g., from checkbox or project switch)
+  const prevProjectPathRef = useRef<string | null>(null);
+  const prevDefaultRuntimeModeRef = useRef<RuntimeMode | null>(null);
+
+  // When switching projects or changing the persisted default mode, reset the selection.
+  // Importantly: do NOT reset selection when lastSshHost/lastDockerImage changes while typing.
   useEffect(() => {
-    setSelectedRuntimeState(buildRuntimeForMode(defaultRuntimeMode, lastSshHost, lastDockerImage));
-  }, [defaultRuntimeMode, lastSshHost, lastDockerImage]);
+    const projectChanged = prevProjectPathRef.current !== projectPath;
+    const defaultModeChanged = prevDefaultRuntimeModeRef.current !== defaultRuntimeMode;
+
+    if (projectChanged || defaultModeChanged) {
+      setSelectedRuntimeState(
+        buildRuntimeForMode(defaultRuntimeMode, defaultSshHost, defaultDockerImage)
+      );
+    }
+
+    prevProjectPathRef.current = projectPath;
+    prevDefaultRuntimeModeRef.current = defaultRuntimeMode;
+  }, [projectPath, defaultRuntimeMode, defaultSshHost, defaultDockerImage]);
+
+  // When the user switches into SSH/Docker mode, seed the field with the remembered host/image.
+  // This avoids clearing the last host/image when the UI switches modes with an empty field.
+  const prevSelectedRuntimeModeRef = useRef<RuntimeMode | null>(null);
+  useEffect(() => {
+    const prevMode = prevSelectedRuntimeModeRef.current;
+    if (prevMode !== selectedRuntime.mode) {
+      if (selectedRuntime.mode === RUNTIME_MODE.SSH) {
+        if (!selectedRuntime.host.trim() && lastSshHost.trim()) {
+          setSelectedRuntimeState({ mode: RUNTIME_MODE.SSH, host: lastSshHost });
+        }
+      }
+
+      if (selectedRuntime.mode === RUNTIME_MODE.DOCKER) {
+        if (!selectedRuntime.image.trim() && lastDockerImage.trim()) {
+          setSelectedRuntimeState({ mode: RUNTIME_MODE.DOCKER, image: lastDockerImage });
+        }
+      }
+    }
+
+    prevSelectedRuntimeModeRef.current = selectedRuntime.mode;
+  }, [selectedRuntime, lastSshHost, lastDockerImage]);
 
   // Initialize trunk branch from backend recommendation or first branch
   useEffect(() => {
@@ -151,11 +218,17 @@ export function useDraftWorkspaceSettings(
   // Setter for selected runtime (also persists host/image for future mode switches)
   const setSelectedRuntime = (runtime: ParsedRuntime) => {
     setSelectedRuntimeState(runtime);
-    // Persist host/image so they're remembered when switching modes
-    if (runtime.mode === "ssh") {
-      setLastSshHost(runtime.host);
-    } else if (runtime.mode === "docker") {
-      setLastDockerImage(runtime.image);
+
+    // Persist host/image so they're remembered when switching modes.
+    // Avoid wiping the remembered value when the UI switches modes with an empty field.
+    if (runtime.mode === RUNTIME_MODE.SSH) {
+      if (runtime.host.trim()) {
+        setLastSshHost(runtime.host);
+      }
+    } else if (runtime.mode === RUNTIME_MODE.DOCKER) {
+      if (runtime.image.trim()) {
+        setLastDockerImage(runtime.image);
+      }
     }
   };
 
diff --git a/src/common/types/runtime.ts b/src/common/types/runtime.ts
@@ -40,7 +40,8 @@ export type ParsedRuntime =
  * Format: "ssh <host>" -> { mode: "ssh", host: "<host>" }
  *         "docker <image>" -> { mode: "docker", image: "<image>" }
  *         "worktree" -> { mode: "worktree" }
- *         "local" or undefined -> { mode: "local" }
+ *         "local" -> { mode: "local" }
+ *         undefined/null -> { mode: "worktree" } (default)
  *
  * Note: "ssh" or "docker" without arguments returns null (invalid).
  * Use this for UI state management (localStorage, form inputs).
diff --git a/src/node/runtime/Runtime.ts b/src/node/runtime/Runtime.ts
@@ -15,7 +15,7 @@
  * srcBaseDir (base directory for all workspaces):
  *   - Where mux stores ALL workspace directories
  *   - Local: ~/.mux/src (tilde expanded to full path by LocalRuntime)
- *   - SSH: /home/user/workspace (must be absolute path, no tilde allowed)
+ *   - SSH: /home/user/workspace (tilde paths are allowed and are resolved before use)
  *
  * Workspace Path Computation:
  *   {srcBaseDir}/{projectName}/{workspaceName}
diff --git a/src/node/runtime/SSHRuntime.ts b/src/node/runtime/SSHRuntime.ts
@@ -291,10 +291,24 @@ export class SSHRuntime extends RemoteRuntime {
   // ===== Runtime interface implementations =====
 
   async resolvePath(filePath: string): Promise<string> {
-    // Use shell to expand tildes on remote system
-    // Bash will expand ~ automatically when we echo the unquoted variable
-    // This works with BusyBox (doesn't require GNU coreutils)
-    const command = `bash -c 'p=${shescape.quote(filePath)}; echo $p'`;
+    // Expand ~ on the remote host.
+    // Note: `p='~/x'; echo "$p"` does NOT expand ~ (tilde expansion happens before assignment),
+    // so we do explicit expansion via shell logic.
+    const script = [
+      `p=${shescape.quote(filePath)}`,
+      'if [ "$p" = "~" ]; then',
+      '  echo "$HOME"',
+      'elif [ "${p#\\~/}" != "$p" ]; then',
+      '  echo "$HOME/${p#\\~/}"',
+      'elif [ "${p#/}" != "$p" ]; then',
+      '  echo "$p"',
+      "else",
+      '  echo "$PWD/$p"',
+      "fi",
+    ].join("\n");
+
+    const command = `bash -lc ${shescape.quote(script)}`;
+
     // Use 10 second timeout for path resolution to allow for slower SSH connections
     return this.execSSHCommand(command, 10000);
   }
diff --git a/tests/runtime/runtime.test.ts b/tests/runtime/runtime.test.ts
@@ -193,6 +193,29 @@ describeIntegration("Runtime integration tests", () => {
         ); // 15 second timeout for test (includes workspace creation overhead)
       });
 
+      describe("resolvePath() - Path resolution", () => {
+        test.concurrent("expands ~ to the home directory", async () => {
+          const runtime = createRuntime();
+
+          const resolved = await runtime.resolvePath("~");
+
+          if (type === "ssh") {
+            expect(resolved).toBe("/home/testuser");
+          } else {
+            expect(resolved).toBe(os.homedir());
+          }
+        });
+
+        test.concurrent("expands ~/path by prefixing the home directory", async () => {
+          const runtime = createRuntime();
+
+          const home = await runtime.resolvePath("~");
+          const resolved = await runtime.resolvePath("~/mux");
+
+          expect(resolved).toBe(`${home}/mux`);
+        });
+      });
+
       describe("readFile() - File reading", () => {
         test.concurrent("reads file contents", async () => {
           const runtime = createRuntime();

Original file line number	Diff line number	Diff line change
`@@ -15,7 +15,7 @@`
`15`	`15`	`* srcBaseDir (base directory for all workspaces):`
`16`	`16`	`* - Where mux stores ALL workspace directories`
`17`	`17`	`* - Local: ~/.mux/src (tilde expanded to full path by LocalRuntime)`
`18`		`- * - SSH: /home/user/workspace (must be absolute path, no tilde allowed)`
	`18`	`+ * - SSH: /home/user/workspace (tilde paths are allowed and are resolved before use)`
`19`	`19`	`*`
`20`	`20`	`* Workspace Path Computation:`
`21`	`21`	`* {srcBaseDir}/{projectName}/{workspaceName}`