🤖 refactor: use discriminated union for runtime settings

Refactor DraftWorkspaceSettings to use ParsedRuntime discriminated union
instead of separate runtimeMode/sshHost/dockerImage fields.

This makes invalid states unrepresentable:
- Cannot have SSH mode without a host
- Cannot have Docker mode without an image
- Runtime-specific args are nested within their runtime type

Changes:
- DraftWorkspaceSettings.selectedRuntime: ParsedRuntime | null
- DraftWorkspaceSettings.defaultRuntime: ParsedRuntime | null
- Remove sshHost/dockerImage as separate top-level fields
- Update CreationControls to derive mode from selectedRuntime
- Add Docker image input field to CreationControls
- Update test harness to use new discriminated union API

_Generated with mux_

Author: ammar-agent

src/browser/components/ChatInput/CreationControls.tsx

@@ -1,5 +1,5 @@
-import React, { useCallback, useEffect } from "react";
-import { RUNTIME_MODE, type RuntimeMode } from "@/common/types/runtime";
+import React, { useCallback, useEffect, useState } from "react";
+import { RUNTIME_MODE, type RuntimeMode, type ParsedRuntime } from "@/common/types/runtime";
 import { Select } from "../Select";
 import { RuntimeIconSelector } from "../RuntimeIconSelector";
 import { Loader2, Wand2 } from "lucide-react";
@@ -13,40 +13,150 @@ interface CreationControlsProps {
   branchesLoaded: boolean;
   trunkBranch: string;
   onTrunkBranchChange: (branch: string) => void;
-  runtimeMode: RuntimeMode;
-  defaultRuntimeMode: RuntimeMode;
-  sshHost: string;
-  onRuntimeModeChange: (mode: RuntimeMode) => void;
-  onSetDefaultRuntime: (mode: RuntimeMode) => void;
-  onSshHostChange: (host: string) => void;
+  /**
+   * Currently selected runtime (discriminated union).
+   * null when mode requires additional input (e.g., SSH without host).
+   */
+  selectedRuntime: ParsedRuntime | null;
+  /** Persisted default runtime for this project */
+  defaultRuntime: ParsedRuntime | null;
+  /** Set the currently selected runtime */
+  onRuntimeChange: (runtime: ParsedRuntime | null) => void;
+  /** Set the default runtime for this project (persists via checkbox) */
+  onSetDefaultRuntime: (runtime: ParsedRuntime | null) => void;
   disabled: boolean;
   /** Workspace name/title generation state and actions */
   nameState: WorkspaceNameState;
 }
 
+/**
+ * Helper to get the mode from a ParsedRuntime (for UI display)
+ */
+function getRuntimeMode(runtime: ParsedRuntime | null): RuntimeMode {
+  return runtime?.mode ?? RUNTIME_MODE.WORKTREE;
+}
+
+/**
+ * Helper to get SSH host from runtime (empty string if not SSH)
+ */
+function getSshHost(runtime: ParsedRuntime | null): string {
+  return runtime?.mode === "ssh" ? runtime.host : "";
+}
+
+/**
+ * Helper to get Docker image from runtime (empty string if not Docker)
+ */
+function getDockerImage(runtime: ParsedRuntime | null): string {
+  return runtime?.mode === "docker" ? runtime.image : "";
+}
+
 /**
  * Additional controls shown only during workspace creation
  * - Trunk branch selector (which branch to fork from) - hidden for Local runtime
- * - Runtime mode (Local, Worktree, SSH) - only Local available for non-git directories
+ * - Runtime mode (Local, Worktree, SSH, Docker) - only Local available for non-git directories
  * - Workspace name (auto-generated with manual override)
  */
 export function CreationControls(props: CreationControlsProps) {
+  const { selectedRuntime, onRuntimeChange, nameState } = props;
+
+  // Derive the current mode for UI display
+  const runtimeMode = getRuntimeMode(selectedRuntime);
+
+  // Local state for SSH host and Docker image input (before user commits)
+  // This allows typing without immediately updating the runtime
+  const [sshHostInput, setSshHostInput] = useState(() => getSshHost(selectedRuntime));
+  const [dockerImageInput, setDockerImageInput] = useState(() => getDockerImage(selectedRuntime));
+
+  // Sync local input state when runtime changes externally
+  useEffect(() => {
+    setSshHostInput(getSshHost(selectedRuntime));
+    setDockerImageInput(getDockerImage(selectedRuntime));
+  }, [selectedRuntime]);
+
   // Non-git directories (empty branches after loading completes) can only use local runtime
   // Don't check until branches have loaded to avoid prematurely switching runtime
   const isNonGitRepo = props.branchesLoaded && props.branches.length === 0;
 
   // Local runtime doesn't need a trunk branch selector (uses project dir as-is)
-  const showTrunkBranchSelector =
-    props.branches.length > 0 && props.runtimeMode !== RUNTIME_MODE.LOCAL;
-
-  const { runtimeMode, onRuntimeModeChange, nameState } = props;
+  const showTrunkBranchSelector = props.branches.length > 0 && runtimeMode !== RUNTIME_MODE.LOCAL;
 
   // Force local runtime for non-git directories (only after branches loaded)
   useEffect(() => {
     if (isNonGitRepo && runtimeMode !== RUNTIME_MODE.LOCAL) {
-      onRuntimeModeChange(RUNTIME_MODE.LOCAL);
+      onRuntimeChange({ mode: "local" });
     }
-  }, [isNonGitRepo, runtimeMode, onRuntimeModeChange]);
+  }, [isNonGitRepo, runtimeMode, onRuntimeChange]);
+
+  // Handle mode change from RuntimeIconSelector
+  const handleModeChange = useCallback(
+    (mode: RuntimeMode) => {
+      switch (mode) {
+        case RUNTIME_MODE.LOCAL:
+          onRuntimeChange({ mode: "local" });
+          break;
+        case RUNTIME_MODE.WORKTREE:
+          onRuntimeChange({ mode: "worktree" });
+          break;
+        case RUNTIME_MODE.SSH:
+          // Use saved host if available, otherwise null (user needs to enter)
+          onRuntimeChange(sshHostInput ? { mode: "ssh", host: sshHostInput } : null);
+          break;
+        case RUNTIME_MODE.DOCKER:
+          // Use saved image if available, otherwise null (user needs to enter)
+          onRuntimeChange(dockerImageInput ? { mode: "docker", image: dockerImageInput } : null);
+          break;
+      }
+    },
+    [onRuntimeChange, sshHostInput, dockerImageInput]
+  );
+
+  // Handle setting default from RuntimeIconSelector
+  const { onSetDefaultRuntime } = props;
+  const handleSetDefault = useCallback(
+    (mode: RuntimeMode) => {
+      switch (mode) {
+        case RUNTIME_MODE.LOCAL:
+          onSetDefaultRuntime({ mode: "local" });
+          break;
+        case RUNTIME_MODE.WORKTREE:
+          onSetDefaultRuntime({ mode: "worktree" });
+          break;
+        case RUNTIME_MODE.SSH:
+          // Only set default if we have a valid host
+          if (sshHostInput) {
+            onSetDefaultRuntime({ mode: "ssh", host: sshHostInput });
+          }
+          break;
+        case RUNTIME_MODE.DOCKER:
+          // Only set default if we have a valid image
+          if (dockerImageInput) {
+            onSetDefaultRuntime({ mode: "docker", image: dockerImageInput });
+          }
+          break;
+      }
+    },
+    [onSetDefaultRuntime, sshHostInput, dockerImageInput]
+  );
+
+  // Handle SSH host input change
+  const handleSshHostChange = useCallback(
+    (host: string) => {
+      setSshHostInput(host);
+      // Update runtime with new host (or null if empty)
+      onRuntimeChange(host.trim() ? { mode: "ssh", host: host.trim() } : null);
+    },
+    [onRuntimeChange]
+  );
+
+  // Handle Docker image input change
+  const handleDockerImageChange = useCallback(
+    (image: string) => {
+      setDockerImageInput(image);
+      // Update runtime with new image (or null if empty)
+      onRuntimeChange(image.trim() ? { mode: "docker", image: image.trim() } : null);
+    },
+    [onRuntimeChange]
+  );
 
   const handleNameChange = useCallback(
     (e: React.ChangeEvent<HTMLInputElement>) => {
@@ -125,16 +235,20 @@ export function CreationControls(props: CreationControlsProps) {
         {nameState.error && <span className="text-xs text-red-500">{nameState.error}</span>}
       </div>
 
-      {/* Second row: Runtime, Branch, SSH */}
+      {/* Second row: Runtime, Branch, SSH/Docker input */}
       <div className="flex flex-wrap items-center gap-x-3 gap-y-2">
         {/* Runtime Selector - icon-based with tooltips */}
         <RuntimeIconSelector
-          value={props.runtimeMode}
-          onChange={props.onRuntimeModeChange}
-          defaultMode={props.defaultRuntimeMode}
-          onSetDefault={props.onSetDefaultRuntime}
+          value={runtimeMode}
+          onChange={handleModeChange}
+          defaultMode={getRuntimeMode(props.defaultRuntime)}
+          onSetDefault={handleSetDefault}
           disabled={props.disabled}
-          disabledModes={isNonGitRepo ? [RUNTIME_MODE.WORKTREE, RUNTIME_MODE.SSH] : undefined}
+          disabledModes={
+            isNonGitRepo
+              ? [RUNTIME_MODE.WORKTREE, RUNTIME_MODE.SSH, RUNTIME_MODE.DOCKER]
+              : undefined
+          }
         />
 
         {/* Trunk Branch Selector - hidden for Local runtime */}
@@ -158,17 +272,29 @@ export function CreationControls(props: CreationControlsProps) {
           </div>
         )}
 
-        {/* SSH Host Input - after From selector */}
-        {props.runtimeMode === RUNTIME_MODE.SSH && (
+        {/* SSH Host Input - shown when SSH mode selected */}
+        {runtimeMode === RUNTIME_MODE.SSH && (
           <input
             type="text"
-            value={props.sshHost}
-            onChange={(e) => props.onSshHostChange(e.target.value)}
+            value={sshHostInput}
+            onChange={(e) => handleSshHostChange(e.target.value)}
             placeholder="user@host"
             disabled={props.disabled}
             className="bg-separator text-foreground border-border-medium focus:border-accent h-6 w-32 rounded border px-1 text-xs focus:outline-none disabled:opacity-50"
           />
         )}
+
+        {/* Docker Image Input - shown when Docker mode selected */}
+        {runtimeMode === RUNTIME_MODE.DOCKER && (
+          <input
+            type="text"
+            value={dockerImageInput}
+            onChange={(e) => handleDockerImageChange(e.target.value)}
+            placeholder="ubuntu:22.04"
+            disabled={props.disabled}
+            className="bg-separator text-foreground border-border-medium focus:border-accent h-6 w-32 rounded border px-1 text-xs focus:outline-none disabled:opacity-50"
+          />
+        )}
       </div>
     </div>
   );

src/browser/components/ChatInput/index.tsx

@@ -1608,12 +1608,10 @@ export const ChatInput: React.FC<ChatInputProps> = (props) => {
                 branchesLoaded={creationState.branchesLoaded}
                 trunkBranch={creationState.trunkBranch}
                 onTrunkBranchChange={creationState.setTrunkBranch}
-                runtimeMode={creationState.runtimeMode}
-                defaultRuntimeMode={creationState.defaultRuntimeMode}
-                sshHost={creationState.sshHost}
-                onRuntimeModeChange={creationState.setRuntimeMode}
-                onSetDefaultRuntime={creationState.setDefaultRuntimeMode}
-                onSshHostChange={creationState.setSshHost}
+                selectedRuntime={creationState.selectedRuntime}
+                defaultRuntime={creationState.defaultRuntime}
+                onRuntimeChange={creationState.setSelectedRuntime}
+                onSetDefaultRuntime={creationState.setDefaultRuntime}
                 disabled={creationState.isSending || isSending}
                 nameState={creationState.nameState}
               />

src/browser/components/ChatInput/useCreationWorkspace.test.tsx

@@ -10,7 +10,8 @@ import {
 } from "@/common/constants/storage";
 import type { SendMessageError as _SendMessageError } from "@/common/types/errors";
 import type { WorkspaceChatMessage } from "@/common/orpc/types";
-import type { RuntimeMode } from "@/common/types/runtime";
+import type { ParsedRuntime } from "@/common/types/runtime";
+import { buildRuntimeString } from "@/common/types/runtime";
 import type {
   FrontendWorkspaceMetadata,
   WorkspaceActivitySnapshot,
@@ -409,9 +410,7 @@ describe("useCreationWorkspace", () => {
     persistedPreferences[getModelKey(getProjectScopeId(TEST_PROJECT_PATH))] = "gpt-4";
 
     draftSettingsState = createDraftSettingsHarness({
-      runtimeMode: "ssh",
-      sshHost: "example.com",
-      runtimeString: "ssh example.com",
+      selectedRuntime: { mode: "ssh", host: "example.com" },
       trunkBranch: "dev",
     });
     const onWorkspaceCreated = mock((metadata: FrontendWorkspaceMetadata) => metadata);
@@ -519,90 +518,65 @@ type DraftSettingsHarness = ReturnType<typeof createDraftSettingsHarness>;
 
 function createDraftSettingsHarness(
   initial?: Partial<{
-    runtimeMode: RuntimeMode;
-    sshHost: string;
-    dockerImage: string;
+    selectedRuntime: ParsedRuntime | null;
+    defaultRuntime: ParsedRuntime | null;
     trunkBranch: string;
-    runtimeString?: string | undefined;
-    defaultRuntimeMode?: RuntimeMode;
   }>
 ) {
-  const state = {
-    runtimeMode: initial?.runtimeMode ?? "local",
-    defaultRuntimeMode: initial?.defaultRuntimeMode ?? "worktree",
-    sshHost: initial?.sshHost ?? "",
-    dockerImage: initial?.dockerImage ?? "",
-    trunkBranch: initial?.trunkBranch ?? "main",
-    runtimeString: initial?.runtimeString,
-  } satisfies {
-    runtimeMode: RuntimeMode;
-    defaultRuntimeMode: RuntimeMode;
-    sshHost: string;
-    dockerImage: string;
+  // Use discriminated union - makes invalid states unrepresentable
+  // Explicitly type state to allow null assignment in setters
+  const state: {
+    selectedRuntime: ParsedRuntime | null;
+    defaultRuntime: ParsedRuntime | null;
     trunkBranch: string;
-    runtimeString: string | undefined;
+  } = {
+    selectedRuntime: initial?.selectedRuntime ?? { mode: "local" },
+    defaultRuntime: initial?.defaultRuntime ?? { mode: "worktree" },
+    trunkBranch: initial?.trunkBranch ?? "main",
   };
 
   const setTrunkBranch = mock((branch: string) => {
     state.trunkBranch = branch;
   });
 
-  const getRuntimeString = mock(() => state.runtimeString);
-
-  const setRuntimeMode = mock((mode: RuntimeMode) => {
-    state.runtimeMode = mode;
-    const trimmedHost = state.sshHost.trim();
-    state.runtimeString = mode === "ssh" ? (trimmedHost ? `ssh ${trimmedHost}` : "ssh") : undefined;
+  const getRuntimeString = mock(() => {
+    return state.selectedRuntime ? buildRuntimeString(state.selectedRuntime) : undefined;
   });
 
-  const setDefaultRuntimeMode = mock((mode: RuntimeMode) => {
-    state.defaultRuntimeMode = mode;
-    state.runtimeMode = mode;
-    const trimmedHost = state.sshHost.trim();
-    state.runtimeString = mode === "ssh" ? (trimmedHost ? `ssh ${trimmedHost}` : "ssh") : undefined;
-  });
+  const setSelectedRuntime = mock((runtime: ParsedRuntime | null) => {
+    state.selectedRuntime = runtime;
+  }) as (runtime: ParsedRuntime | null) => void;
 
-  const setSshHost = mock((host: string) => {
-    state.sshHost = host;
-  });
-
-  const setDockerImage = mock((image: string) => {
-    state.dockerImage = image;
-  });
+  const setDefaultRuntime = mock((runtime: ParsedRuntime | null) => {
+    state.defaultRuntime = runtime;
+    state.selectedRuntime = runtime;
+  }) as (runtime: ParsedRuntime | null) => void;
 
   return {
     state,
-    setRuntimeMode,
-    setDefaultRuntimeMode,
-    setSshHost,
-    setDockerImage,
+    setSelectedRuntime,
+    setDefaultRuntime,
     setTrunkBranch,
     getRuntimeString,
     snapshot(): {
       settings: DraftWorkspaceSettings;
-      setRuntimeMode: typeof setRuntimeMode;
-      setDefaultRuntimeMode: typeof setDefaultRuntimeMode;
-      setSshHost: typeof setSshHost;
-      setDockerImage: typeof setDockerImage;
+      setSelectedRuntime: typeof setSelectedRuntime;
+      setDefaultRuntime: typeof setDefaultRuntime;
       setTrunkBranch: typeof setTrunkBranch;
       getRuntimeString: typeof getRuntimeString;
     } {
       const settings: DraftWorkspaceSettings = {
         model: "gpt-4",
         thinkingLevel: "medium",
         mode: "exec",
-        runtimeMode: state.runtimeMode,
-        defaultRuntimeMode: state.defaultRuntimeMode,
-        sshHost: state.sshHost,
-        dockerImage: state.dockerImage ?? "",
+        selectedRuntime: state.selectedRuntime,
+        defaultRuntime: state.defaultRuntime,
         trunkBranch: state.trunkBranch,
       };
       return {
         settings,
-        setRuntimeMode,
-        setDefaultRuntimeMode,
-        setSshHost,
-        setDockerImage,
+        setSelectedRuntime,
+        setDefaultRuntime,
         setTrunkBranch,
         getRuntimeString,
       };