refactor: use discriminated unions for runtime parsing

Refactored parseRuntimeModeAndHost to return discriminated unions:
- SSH: { mode: 'ssh', host: string }
- Docker: { mode: 'docker', image: string }
- Local/Worktree: { mode: 'local' | 'worktree' }

Key changes:
- Returns null instead of empty strings for invalid ssh/docker without args
- buildRuntimeString now takes ParsedRuntime instead of (mode, host)
- Added dockerImage to DraftWorkspaceSettings (persisted per project)
- Updated CLI to use new API with proper error handling
- Updated tests for new behavior

Author: ammar-agent

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

@@ -521,6 +521,7 @@ function createDraftSettingsHarness(
   initial?: Partial<{
     runtimeMode: RuntimeMode;
     sshHost: string;
+    dockerImage: string;
     trunkBranch: string;
     runtimeString?: string | undefined;
     defaultRuntimeMode?: RuntimeMode;
@@ -530,12 +531,14 @@ function createDraftSettingsHarness(
     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;
     trunkBranch: string;
     runtimeString: string | undefined;
   };
@@ -563,18 +566,24 @@ function createDraftSettingsHarness(
     state.sshHost = host;
   });
 
+  const setDockerImage = mock((image: string) => {
+    state.dockerImage = image;
+  });
+
   return {
     state,
     setRuntimeMode,
     setDefaultRuntimeMode,
     setSshHost,
+    setDockerImage,
     setTrunkBranch,
     getRuntimeString,
     snapshot(): {
       settings: DraftWorkspaceSettings;
       setRuntimeMode: typeof setRuntimeMode;
       setDefaultRuntimeMode: typeof setDefaultRuntimeMode;
       setSshHost: typeof setSshHost;
+      setDockerImage: typeof setDockerImage;
       setTrunkBranch: typeof setTrunkBranch;
       getRuntimeString: typeof getRuntimeString;
     } {
@@ -585,13 +594,15 @@ function createDraftSettingsHarness(
         runtimeMode: state.runtimeMode,
         defaultRuntimeMode: state.defaultRuntimeMode,
         sshHost: state.sshHost,
+        dockerImage: state.dockerImage ?? "",
         trunkBranch: state.trunkBranch,
       };
       return {
         settings,
         setRuntimeMode,
         setDefaultRuntimeMode,
         setSshHost,
+        setDockerImage,
         setTrunkBranch,
         getRuntimeString,
       };

src/browser/hooks/useDraftWorkspaceSettings.ts

@@ -5,14 +5,17 @@ import { useMode } from "@/browser/contexts/ModeContext";
 import { useModelLRU } from "./useModelLRU";
 import {
   type RuntimeMode,
+  type ParsedRuntime,
   parseRuntimeModeAndHost,
   buildRuntimeString,
+  RUNTIME_MODE,
 } from "@/common/types/runtime";
 import {
   getModelKey,
   getRuntimeKey,
   getTrunkBranchKey,
   getLastSshHostKey,
+  getLastDockerImageKey,
   getProjectScopeId,
 } from "@/common/constants/storage";
 import type { UIMode } from "@/common/types/mode";
@@ -33,7 +36,10 @@ export interface DraftWorkspaceSettings {
   runtimeMode: RuntimeMode;
   /** Persisted default runtime for this project (used to initialize selection) */
   defaultRuntimeMode: RuntimeMode;
+  /** SSH host (persisted separately from mode) */
   sshHost: string;
+  /** Docker image (persisted separately from mode) */
+  dockerImage: string;
   trunkBranch: string;
 }
 
@@ -57,6 +63,7 @@ export function useDraftWorkspaceSettings(
   /** Set the default runtime mode for this project (persists via checkbox) */
   setDefaultRuntimeMode: (mode: RuntimeMode) => void;
   setSshHost: (host: string) => void;
+  setDockerImage: (image: string) => void;
   setTrunkBranch: (branch: string) => void;
   getRuntimeString: () => string | undefined;
 } {
@@ -79,8 +86,9 @@ export function useDraftWorkspaceSettings(
     { listener: true }
   );
 
-  // Parse default runtime string into mode (worktree when undefined)
-  const { mode: defaultRuntimeMode } = parseRuntimeModeAndHost(defaultRuntimeString);
+  // Parse default runtime string into mode (worktree when undefined or invalid)
+  const parsedDefault = parseRuntimeModeAndHost(defaultRuntimeString);
+  const defaultRuntimeMode: RuntimeMode = parsedDefault?.mode ?? RUNTIME_MODE.WORKTREE;
 
   // Currently selected runtime mode for this session (initialized from default)
   // This allows user to select a different runtime without changing the default
@@ -106,6 +114,13 @@ export function useDraftWorkspaceSettings(
     { listener: true }
   );
 
+  // Project-scoped Docker image preference (persisted separately from runtime mode)
+  const [lastDockerImage, setLastDockerImage] = usePersistedState<string>(
+    getLastDockerImageKey(projectPath),
+    "",
+    { listener: true }
+  );
+
   // Initialize trunk branch from backend recommendation or first branch
   useEffect(() => {
     if (!trunkBranch && branches.length > 0) {
@@ -114,14 +129,31 @@ export function useDraftWorkspaceSettings(
     }
   }, [branches, recommendedTrunk, trunkBranch, setTrunkBranch]);
 
+  // Build ParsedRuntime from mode + stored host/image
+  const buildParsedRuntime = (mode: RuntimeMode): ParsedRuntime | null => {
+    switch (mode) {
+      case RUNTIME_MODE.LOCAL:
+        return { mode: "local" };
+      case RUNTIME_MODE.WORKTREE:
+        return { mode: "worktree" };
+      case RUNTIME_MODE.SSH:
+        return lastSshHost ? { mode: "ssh", host: lastSshHost } : null;
+      case RUNTIME_MODE.DOCKER:
+        return lastDockerImage ? { mode: "docker", image: lastDockerImage } : null;
+      default:
+        return null;
+    }
+  };
+
   // Setter for selected runtime mode (changes current selection, does not persist)
   const setRuntimeMode = (newMode: RuntimeMode) => {
     setSelectedRuntimeMode(newMode);
   };
 
   // Setter for default runtime mode (persists via checkbox in tooltip)
   const setDefaultRuntimeMode = (newMode: RuntimeMode) => {
-    const newRuntimeString = buildRuntimeString(newMode, lastSshHost);
+    const parsed = buildParsedRuntime(newMode);
+    const newRuntimeString = parsed ? buildRuntimeString(parsed) : undefined;
     setDefaultRuntimeString(newRuntimeString);
     // Also update selection to match new default
     setSelectedRuntimeMode(newMode);
@@ -132,9 +164,15 @@ export function useDraftWorkspaceSettings(
     setLastSshHost(newHost);
   };
 
+  // Setter for Docker image (persisted separately so it's remembered across mode switches)
+  const setDockerImage = (newImage: string) => {
+    setLastDockerImage(newImage);
+  };
+
   // Helper to get runtime string for IPC calls (uses selected mode, not default)
   const getRuntimeString = (): string | undefined => {
-    return buildRuntimeString(selectedRuntimeMode, lastSshHost);
+    const parsed = buildParsedRuntime(selectedRuntimeMode);
+    return parsed ? buildRuntimeString(parsed) : undefined;
   };
 
   return {
@@ -145,11 +183,13 @@ export function useDraftWorkspaceSettings(
       runtimeMode: selectedRuntimeMode,
       defaultRuntimeMode,
       sshHost: lastSshHost,
+      dockerImage: lastDockerImage,
       trunkBranch,
     },
     setRuntimeMode,
     setDefaultRuntimeMode,
     setSshHost,
+    setDockerImage,
     setTrunkBranch,
     getRuntimeString,
   };

src/cli/run.ts

@@ -52,18 +52,22 @@ function parseRuntimeConfig(value: string | undefined, srcBaseDir: string): Runt
     return { type: "local" };
   }
 
-  const { mode, host } = parseRuntimeModeAndHost(value);
+  const parsed = parseRuntimeModeAndHost(value);
+  if (!parsed) {
+    throw new Error(
+      `Invalid runtime: '${value}'. Use 'local', 'worktree', 'ssh <host>', or 'docker <image>'`
+    );
+  }
 
-  switch (mode) {
+  switch (parsed.mode) {
     case RUNTIME_MODE.LOCAL:
       return { type: "local" };
     case RUNTIME_MODE.WORKTREE:
       return { type: "worktree", srcBaseDir };
     case RUNTIME_MODE.SSH:
-      if (!host.trim()) {
-        throw new Error("SSH runtime requires a host (e.g., --runtime 'ssh user@host')");
-      }
-      return { type: "ssh", host: host.trim(), srcBaseDir };
+      return { type: "ssh", host: parsed.host, srcBaseDir };
+    case RUNTIME_MODE.DOCKER:
+      return { type: "docker", image: parsed.image };
     default:
       return { type: "local" };
   }

src/common/constants/storage.ts

@@ -145,6 +145,16 @@ export function getLastSshHostKey(projectPath: string): string {
   return `lastSshHost:${projectPath}`;
 }
 
+/**
+ * Get the localStorage key for the last entered Docker image for a project
+ * Stores the last entered Docker image separately from runtime mode
+ * so it persists when switching between runtime modes
+ * Format: "lastDockerImage:{projectPath}"
+ */
+export function getLastDockerImageKey(projectPath: string): string {
+  return `lastDockerImage:${projectPath}`;
+}
+
 /**
  * Get the localStorage key for the preferred compaction model (global)
  * Format: "preferredCompactionModel"

src/common/types/runtime.test.ts

@@ -9,69 +9,96 @@ describe("parseRuntimeModeAndHost", () => {
     });
   });
 
-  it("parses SSH mode without host", () => {
-    expect(parseRuntimeModeAndHost("ssh")).toEqual({
-      mode: "ssh",
-      host: "",
+  it("returns null for SSH mode without host", () => {
+    expect(parseRuntimeModeAndHost("ssh")).toBeNull();
+  });
+
+  it("returns null for SSH with trailing space but no host", () => {
+    expect(parseRuntimeModeAndHost("ssh ")).toBeNull();
+  });
+
+  it("parses Docker mode with image", () => {
+    expect(parseRuntimeModeAndHost("docker ubuntu:22.04")).toEqual({
+      mode: "docker",
+      image: "ubuntu:22.04",
     });
   });
 
+  it("returns null for Docker mode without image", () => {
+    expect(parseRuntimeModeAndHost("docker")).toBeNull();
+  });
+
   it("parses local mode", () => {
     expect(parseRuntimeModeAndHost("local")).toEqual({
       mode: "local",
-      host: "",
+    });
+  });
+
+  it("parses worktree mode", () => {
+    expect(parseRuntimeModeAndHost("worktree")).toEqual({
+      mode: "worktree",
     });
   });
 
   it("defaults to worktree for undefined", () => {
     expect(parseRuntimeModeAndHost(undefined)).toEqual({
       mode: "worktree",
-      host: "",
     });
   });
 
   it("defaults to worktree for null", () => {
     expect(parseRuntimeModeAndHost(null)).toEqual({
       mode: "worktree",
-      host: "",
     });
   });
+
+  it("returns null for unrecognized runtime", () => {
+    expect(parseRuntimeModeAndHost("unknown")).toBeNull();
+  });
 });
 
 describe("buildRuntimeString", () => {
   it("builds SSH string with host", () => {
-    expect(buildRuntimeString("ssh", "user@host")).toBe("ssh user@host");
+    expect(buildRuntimeString({ mode: "ssh", host: "user@host" })).toBe("ssh user@host");
   });
 
-  it("builds SSH string without host (persists SSH mode)", () => {
-    expect(buildRuntimeString("ssh", "")).toBe("ssh");
+  it("builds Docker string with image", () => {
+    expect(buildRuntimeString({ mode: "docker", image: "ubuntu:22.04" })).toBe(
+      "docker ubuntu:22.04"
+    );
   });
 
   it("returns 'local' for local mode", () => {
-    expect(buildRuntimeString("local", "")).toBe("local");
+    expect(buildRuntimeString({ mode: "local" })).toBe("local");
   });
 
   it("returns undefined for worktree mode (default)", () => {
-    expect(buildRuntimeString("worktree", "")).toBeUndefined();
-  });
-
-  it("trims whitespace from host", () => {
-    expect(buildRuntimeString("ssh", "  user@host  ")).toBe("ssh user@host");
+    expect(buildRuntimeString({ mode: "worktree" })).toBeUndefined();
   });
 });
 
 describe("round-trip parsing and building", () => {
-  it("preserves SSH mode without host", () => {
-    const built = buildRuntimeString("ssh", "");
+  it("preserves SSH mode with host", () => {
+    const built = buildRuntimeString({ mode: "ssh", host: "user@host" });
     const parsed = parseRuntimeModeAndHost(built);
-    expect(parsed.mode).toBe("ssh");
-    expect(parsed.host).toBe("");
+    expect(parsed).toEqual({ mode: "ssh", host: "user@host" });
   });
 
-  it("preserves SSH mode with host", () => {
-    const built = buildRuntimeString("ssh", "user@host");
+  it("preserves Docker mode with image", () => {
+    const built = buildRuntimeString({ mode: "docker", image: "node:20" });
+    const parsed = parseRuntimeModeAndHost(built);
+    expect(parsed).toEqual({ mode: "docker", image: "node:20" });
+  });
+
+  it("preserves local mode", () => {
+    const built = buildRuntimeString({ mode: "local" });
+    const parsed = parseRuntimeModeAndHost(built);
+    expect(parsed).toEqual({ mode: "local" });
+  });
+
+  it("preserves worktree mode", () => {
+    const built = buildRuntimeString({ mode: "worktree" });
     const parsed = parseRuntimeModeAndHost(built);
-    expect(parsed.mode).toBe("ssh");
-    expect(parsed.host).toBe("user@host");
+    expect(parsed).toEqual({ mode: "worktree" });
   });
 });