Skip to content

Commit 1f7673c

Browse files
ammar-agentammar-agent
committed
🤖 refactor: use discriminated union for runtime selection
Replace separate runtimeMode/sshHost/dockerImage props with single ParsedRuntime discriminated union in useDraftWorkspaceSettings and CreationControls. This makes the API type-safe and ensures runtime parameters (host for SSH, image for Docker) are always bundled with their mode.
1 parent f7f721a commit 1f7673c

File tree

5 files changed

+137
-143
lines changed

5 files changed

+137
-143
lines changed

‎src/browser/components/ChatInput/CreationControls.tsx‎

Lines changed: 44 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import React, { useCallback, useEffect } from "react";
2-
import { RUNTIME_MODE, type RuntimeMode } from "@/common/types/runtime";
2+
import { RUNTIME_MODE, type RuntimeMode, type ParsedRuntime } from "@/common/types/runtime";
33
import { Select } from "../Select";
44
import { Loader2, Wand2 } from "lucide-react";
55
import { cn } from "@/common/lib/utils";
@@ -14,14 +14,12 @@ interface CreationControlsProps {
1414
branchesLoaded: boolean;
1515
trunkBranch: string;
1616
onTrunkBranchChange: (branch: string) => void;
17-
runtimeMode: RuntimeMode;
17+
/** Currently selected runtime (discriminated union: SSH has host, Docker has image) */
18+
selectedRuntime: ParsedRuntime;
1819
defaultRuntimeMode: RuntimeMode;
19-
sshHost: string;
20-
dockerImage: string;
21-
onRuntimeModeChange: (mode: RuntimeMode) => void;
20+
/** Set the currently selected runtime (discriminated union) */
21+
onSelectedRuntimeChange: (runtime: ParsedRuntime) => void;
2222
onSetDefaultRuntime: (mode: RuntimeMode) => void;
23-
onSshHostChange: (host: string) => void;
24-
onDockerImageChange: (image: string) => void;
2523
disabled: boolean;
2624
/** Project name to display as header */
2725
projectName: string;
@@ -166,18 +164,21 @@ export function CreationControls(props: CreationControlsProps) {
166164
// Don't check until branches have loaded to avoid prematurely switching runtime
167165
const isNonGitRepo = props.branchesLoaded && props.branches.length === 0;
168166

167+
// Extract mode from discriminated union for convenience
168+
const runtimeMode = props.selectedRuntime.mode;
169+
169170
// Local runtime doesn't need a trunk branch selector (uses project dir as-is)
170171
const showTrunkBranchSelector =
171-
props.branches.length > 0 && props.runtimeMode !== RUNTIME_MODE.LOCAL;
172+
props.branches.length > 0 && runtimeMode !== RUNTIME_MODE.LOCAL;
172173

173-
const { runtimeMode, onRuntimeModeChange } = props;
174+
const { selectedRuntime, onSelectedRuntimeChange } = props;
174175

175176
// Force local runtime for non-git directories (only after branches loaded)
176177
useEffect(() => {
177-
if (isNonGitRepo && runtimeMode !== RUNTIME_MODE.LOCAL) {
178-
onRuntimeModeChange(RUNTIME_MODE.LOCAL);
178+
if (isNonGitRepo && selectedRuntime.mode !== RUNTIME_MODE.LOCAL) {
179+
onSelectedRuntimeChange({ mode: "local" });
179180
}
180-
}, [isNonGitRepo, runtimeMode, onRuntimeModeChange]);
181+
}, [isNonGitRepo, selectedRuntime.mode, onSelectedRuntimeChange]);
181182

182183
const handleNameChange = useCallback(
183184
(e: React.ChangeEvent<HTMLInputElement>) => {
@@ -276,8 +277,31 @@ export function CreationControls(props: CreationControlsProps) {
276277
<label className="text-muted-foreground text-xs font-medium">Workspace Type</label>
277278
<div className="flex flex-wrap items-center gap-3">
278279
<RuntimeButtonGroup
279-
value={props.runtimeMode}
280-
onChange={props.onRuntimeModeChange}
280+
value={runtimeMode}
281+
onChange={(mode) => {
282+
// Convert mode to ParsedRuntime with appropriate defaults
283+
switch (mode) {
284+
case RUNTIME_MODE.SSH:
285+
onSelectedRuntimeChange({
286+
mode: "ssh",
287+
host: selectedRuntime.mode === "ssh" ? selectedRuntime.host : "",
288+
});
289+
break;
290+
case RUNTIME_MODE.DOCKER:
291+
onSelectedRuntimeChange({
292+
mode: "docker",
293+
image: selectedRuntime.mode === "docker" ? selectedRuntime.image : "",
294+
});
295+
break;
296+
case RUNTIME_MODE.LOCAL:
297+
onSelectedRuntimeChange({ mode: "local" });
298+
break;
299+
case RUNTIME_MODE.WORKTREE:
300+
default:
301+
onSelectedRuntimeChange({ mode: "worktree" });
302+
break;
303+
}
304+
}}
281305
defaultMode={props.defaultRuntimeMode}
282306
onSetDefault={props.onSetDefaultRuntime}
283307
disabled={props.disabled}
@@ -310,13 +334,13 @@ export function CreationControls(props: CreationControlsProps) {
310334
)}
311335

312336
{/* SSH Host Input */}
313-
{props.runtimeMode === RUNTIME_MODE.SSH && (
337+
{selectedRuntime.mode === "ssh" && (
314338
<div className="flex items-center gap-2">
315339
<label className="text-muted-foreground text-xs">host</label>
316340
<input
317341
type="text"
318-
value={props.sshHost}
319-
onChange={(e) => props.onSshHostChange(e.target.value)}
342+
value={selectedRuntime.host}
343+
onChange={(e) => onSelectedRuntimeChange({ mode: "ssh", host: e.target.value })}
320344
placeholder="user@host"
321345
disabled={props.disabled}
322346
className="bg-bg-dark text-foreground border-border-medium focus:border-accent h-7 w-36 rounded-md border px-2 text-sm focus:outline-none disabled:opacity-50"
@@ -325,13 +349,13 @@ export function CreationControls(props: CreationControlsProps) {
325349
)}
326350

327351
{/* Docker Image Input */}
328-
{props.runtimeMode === RUNTIME_MODE.DOCKER && (
352+
{selectedRuntime.mode === "docker" && (
329353
<div className="flex items-center gap-2">
330354
<label className="text-muted-foreground text-xs">image</label>
331355
<input
332356
type="text"
333-
value={props.dockerImage}
334-
onChange={(e) => props.onDockerImageChange(e.target.value)}
357+
value={selectedRuntime.image}
358+
onChange={(e) => onSelectedRuntimeChange({ mode: "docker", image: e.target.value })}
335359
placeholder="ubuntu:22.04"
336360
disabled={props.disabled}
337361
className="bg-bg-dark text-foreground border-border-medium focus:border-accent h-7 w-36 rounded-md border px-2 text-sm focus:outline-none disabled:opacity-50"

‎src/browser/components/ChatInput/index.tsx‎

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1438,14 +1438,10 @@ export const ChatInput: React.FC<ChatInputProps> = (props) => {
14381438
branchesLoaded={creationState.branchesLoaded}
14391439
trunkBranch={creationState.trunkBranch}
14401440
onTrunkBranchChange={creationState.setTrunkBranch}
1441-
runtimeMode={creationState.runtimeMode}
1441+
selectedRuntime={creationState.selectedRuntime}
14421442
defaultRuntimeMode={creationState.defaultRuntimeMode}
1443-
sshHost={creationState.sshHost}
1444-
dockerImage={creationState.dockerImage}
1445-
onRuntimeModeChange={creationState.setRuntimeMode}
1443+
onSelectedRuntimeChange={creationState.setSelectedRuntime}
14461444
onSetDefaultRuntime={creationState.setDefaultRuntimeMode}
1447-
onSshHostChange={creationState.setSshHost}
1448-
onDockerImageChange={creationState.setDockerImage}
14491445
disabled={creationState.isSending || isSending}
14501446
projectName={props.projectName}
14511447
nameState={creationState.nameState}

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

Lines changed: 34 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ import {
1010
} from "@/common/constants/storage";
1111
import type { SendMessageError as _SendMessageError } from "@/common/types/errors";
1212
import type { WorkspaceChatMessage } from "@/common/orpc/types";
13-
import type { RuntimeMode } from "@/common/types/runtime";
13+
import type { RuntimeMode, ParsedRuntime } from "@/common/types/runtime";
1414
import type {
1515
FrontendWorkspaceMetadata,
1616
WorkspaceActivitySnapshot,
@@ -409,8 +409,7 @@ describe("useCreationWorkspace", () => {
409409
persistedPreferences[getModelKey(getProjectScopeId(TEST_PROJECT_PATH))] = "gpt-4";
410410

411411
draftSettingsState = createDraftSettingsHarness({
412-
runtimeMode: "ssh",
413-
sshHost: "example.com",
412+
selectedRuntime: { mode: "ssh", host: "example.com" },
414413
runtimeString: "ssh example.com",
415414
trunkBranch: "dev",
416415
});
@@ -519,26 +518,20 @@ type DraftSettingsHarness = ReturnType<typeof createDraftSettingsHarness>;
519518

520519
function createDraftSettingsHarness(
521520
initial?: Partial<{
522-
runtimeMode: RuntimeMode;
523-
sshHost: string;
524-
dockerImage: string;
521+
selectedRuntime: ParsedRuntime;
525522
trunkBranch: string;
526523
runtimeString?: string | undefined;
527524
defaultRuntimeMode?: RuntimeMode;
528525
}>
529526
) {
530527
const state = {
531-
runtimeMode: initial?.runtimeMode ?? "local",
528+
selectedRuntime: initial?.selectedRuntime ?? { mode: "local" as const },
532529
defaultRuntimeMode: initial?.defaultRuntimeMode ?? "worktree",
533-
sshHost: initial?.sshHost ?? "",
534-
dockerImage: initial?.dockerImage ?? "",
535530
trunkBranch: initial?.trunkBranch ?? "main",
536531
runtimeString: initial?.runtimeString,
537532
} satisfies {
538-
runtimeMode: RuntimeMode;
533+
selectedRuntime: ParsedRuntime;
539534
defaultRuntimeMode: RuntimeMode;
540-
sshHost: string;
541-
dockerImage: string;
542535
trunkBranch: string;
543536
runtimeString: string | undefined;
544537
};
@@ -549,60 +542,62 @@ function createDraftSettingsHarness(
549542

550543
const getRuntimeString = mock(() => state.runtimeString);
551544

552-
const setRuntimeMode = mock((mode: RuntimeMode) => {
553-
state.runtimeMode = mode;
554-
const trimmedHost = state.sshHost.trim();
555-
state.runtimeString = mode === "ssh" ? (trimmedHost ? `ssh ${trimmedHost}` : "ssh") : undefined;
545+
const setSelectedRuntime = mock((runtime: ParsedRuntime) => {
546+
state.selectedRuntime = runtime;
547+
if (runtime.mode === "ssh") {
548+
state.runtimeString = runtime.host ? `ssh ${runtime.host}` : "ssh";
549+
} else if (runtime.mode === "docker") {
550+
state.runtimeString = runtime.image ? `docker ${runtime.image}` : "docker";
551+
} else {
552+
state.runtimeString = undefined;
553+
}
556554
});
557555

558556
const setDefaultRuntimeMode = mock((mode: RuntimeMode) => {
559557
state.defaultRuntimeMode = mode;
560-
state.runtimeMode = mode;
561-
const trimmedHost = state.sshHost.trim();
562-
state.runtimeString = mode === "ssh" ? (trimmedHost ? `ssh ${trimmedHost}` : "ssh") : undefined;
563-
});
564-
565-
const setSshHost = mock((host: string) => {
566-
state.sshHost = host;
567-
});
568-
569-
const setDockerImage = mock((image: string) => {
570-
state.dockerImage = image;
558+
// Update selected runtime to match new default
559+
if (mode === "ssh") {
560+
const host = state.selectedRuntime.mode === "ssh" ? state.selectedRuntime.host : "";
561+
state.selectedRuntime = { mode: "ssh", host };
562+
state.runtimeString = host ? `ssh ${host}` : "ssh";
563+
} else if (mode === "docker") {
564+
const image = state.selectedRuntime.mode === "docker" ? state.selectedRuntime.image : "";
565+
state.selectedRuntime = { mode: "docker", image };
566+
state.runtimeString = image ? `docker ${image}` : "docker";
567+
} else if (mode === "local") {
568+
state.selectedRuntime = { mode: "local" };
569+
state.runtimeString = undefined;
570+
} else {
571+
state.selectedRuntime = { mode: "worktree" };
572+
state.runtimeString = undefined;
573+
}
571574
});
572575

573576
return {
574577
state,
575-
setRuntimeMode,
578+
setSelectedRuntime,
576579
setDefaultRuntimeMode,
577-
setSshHost,
578-
setDockerImage,
579580
setTrunkBranch,
580581
getRuntimeString,
581582
snapshot(): {
582583
settings: DraftWorkspaceSettings;
583-
setRuntimeMode: typeof setRuntimeMode;
584+
setSelectedRuntime: typeof setSelectedRuntime;
584585
setDefaultRuntimeMode: typeof setDefaultRuntimeMode;
585-
setSshHost: typeof setSshHost;
586-
setDockerImage: typeof setDockerImage;
587586
setTrunkBranch: typeof setTrunkBranch;
588587
getRuntimeString: typeof getRuntimeString;
589588
} {
590589
const settings: DraftWorkspaceSettings = {
591590
model: "gpt-4",
592591
thinkingLevel: "medium",
593592
mode: "exec",
594-
runtimeMode: state.runtimeMode,
593+
selectedRuntime: state.selectedRuntime,
595594
defaultRuntimeMode: state.defaultRuntimeMode,
596-
sshHost: state.sshHost,
597-
dockerImage: state.dockerImage ?? "",
598595
trunkBranch: state.trunkBranch,
599596
};
600597
return {
601598
settings,
602-
setRuntimeMode,
599+
setSelectedRuntime,
603600
setDefaultRuntimeMode,
604-
setSshHost,
605-
setDockerImage,
606601
setTrunkBranch,
607602
getRuntimeString,
608603
};

‎src/browser/components/ChatInput/useCreationWorkspace.ts‎

Lines changed: 8 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import { useState, useEffect, useCallback } from "react";
22
import type { FrontendWorkspaceMetadata } from "@/common/types/workspace";
3-
import type { RuntimeConfig, RuntimeMode } from "@/common/types/runtime";
3+
import type { RuntimeConfig, RuntimeMode, ParsedRuntime } from "@/common/types/runtime";
44
import type { UIMode } from "@/common/types/mode";
55
import type { ThinkingLevel } from "@/common/types/thinking";
66
import { parseRuntimeString } from "@/browser/utils/chatCommands";
@@ -62,18 +62,13 @@ interface UseCreationWorkspaceReturn {
6262
branchesLoaded: boolean;
6363
trunkBranch: string;
6464
setTrunkBranch: (branch: string) => void;
65-
runtimeMode: RuntimeMode;
65+
/** Currently selected runtime (discriminated union: SSH has host, Docker has image) */
66+
selectedRuntime: ParsedRuntime;
6667
defaultRuntimeMode: RuntimeMode;
67-
sshHost: string;
68-
dockerImage: string;
69-
/** Set the currently selected runtime mode (does not persist) */
70-
setRuntimeMode: (mode: RuntimeMode) => void;
68+
/** Set the currently selected runtime (discriminated union) */
69+
setSelectedRuntime: (runtime: ParsedRuntime) => void;
7170
/** Set the default runtime mode for this project (persists via checkbox) */
7271
setDefaultRuntimeMode: (mode: RuntimeMode) => void;
73-
/** Set the SSH host (persisted separately from runtime mode) */
74-
setSshHost: (host: string) => void;
75-
/** Set the Docker image (persisted separately from runtime mode) */
76-
setDockerImage: (image: string) => void;
7772
toast: Toast | null;
7873
setToast: (toast: Toast | null) => void;
7974
isSending: boolean;
@@ -109,10 +104,8 @@ export function useCreationWorkspace({
109104
// Centralized draft workspace settings with automatic persistence
110105
const {
111106
settings,
112-
setRuntimeMode,
107+
setSelectedRuntime,
113108
setDefaultRuntimeMode,
114-
setSshHost,
115-
setDockerImage,
116109
setTrunkBranch,
117110
getRuntimeString,
118111
} = useDraftWorkspaceSettings(projectPath, branches, recommendedTrunk);
@@ -262,14 +255,10 @@ export function useCreationWorkspace({
262255
branchesLoaded,
263256
trunkBranch: settings.trunkBranch,
264257
setTrunkBranch,
265-
runtimeMode: settings.runtimeMode,
258+
selectedRuntime: settings.selectedRuntime,
266259
defaultRuntimeMode: settings.defaultRuntimeMode,
267-
sshHost: settings.sshHost,
268-
dockerImage: settings.dockerImage,
269-
setRuntimeMode,
260+
setSelectedRuntime,
270261
setDefaultRuntimeMode,
271-
setSshHost,
272-
setDockerImage,
273262
toast,
274263
setToast,
275264
isSending,

0 commit comments

Comments
 (0)