Skip to content

Commit b7b1c96

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 b7b1c96

File tree

8 files changed

+333
-263
lines changed

8 files changed

+333
-263
lines changed

‎docs/system-prompt.mdx‎

Lines changed: 47 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -47,59 +47,59 @@ When the user asks you to remember something:
4747
/**
4848
* Build environment context XML block describing the workspace.
4949
* @param workspacePath - Workspace directory path
50-
* @param runtimeType - Runtime type: "local", "worktree", "ssh", or "docker"
50+
* @param runtimeType - Runtime type (local, worktree, ssh, docker)
5151
*/
52-
function buildEnvironmentContext(
53-
workspacePath: string,
54-
runtimeType: "local" | "worktree" | "ssh" | "docker"
55-
): string {
56-
if (runtimeType === "local") {
57-
// Local runtime works directly in project directory - may or may not be git
58-
return `
59-
<environment>
60-
You are working in a directory at ${workspacePath}
61-
62-
- Tools run here automatically
63-
- You are meant to do your work isolated from the user and other agents
64-
</environment>
65-
`;
66-
}
67-
68-
if (runtimeType === "ssh") {
69-
// SSH runtime clones the repository on a remote host
70-
return `
71-
<environment>
72-
You are in a clone of a git repository at ${workspacePath}
73-
74-
- This IS a git repository - run git commands directly (no cd needed)
75-
- Tools run here automatically
76-
- You are meant to do your work isolated from the user and other agents
77-
</environment>
78-
`;
79-
}
80-
81-
if (runtimeType === "docker") {
82-
// Docker runtime runs in an isolated container
83-
return `
84-
<environment>
85-
You are in a clone of a git repository at ${workspacePath} inside a Docker container
86-
87-
- This IS a git repository - run git commands directly (no cd needed)
88-
- Tools run here automatically inside the container
89-
- You are meant to do your work isolated from the user and other agents
90-
</environment>
91-
`;
52+
function buildEnvironmentContext(workspacePath: string, runtimeType: RuntimeMode): string {
53+
// Common lines shared across git-based runtimes
54+
const gitCommonLines = [
55+
"- This IS a git repository - run git commands directly (no cd needed)",
56+
"- Tools run here automatically",
57+
"- You are meant to do your work isolated from the user and other agents",
58+
];
59+
60+
let description: string;
61+
let lines: string[];
62+
63+
switch (runtimeType) {
64+
case RUNTIME_MODE.LOCAL:
65+
// Local runtime works directly in project directory - may or may not be git
66+
description = `You are working in a directory at ${workspacePath}`;
67+
lines = [
68+
"- Tools run here automatically",
69+
"- You are meant to do your work isolated from the user and other agents",
70+
];
71+
break;
72+
73+
case RUNTIME_MODE.WORKTREE:
74+
// Worktree runtime creates a git worktree locally
75+
description = `You are in a git worktree at ${workspacePath}`;
76+
lines = [
77+
...gitCommonLines,
78+
"- Do not modify or visit other worktrees (especially the main project) without explicit user intent",
79+
];
80+
break;
81+
82+
case RUNTIME_MODE.SSH:
83+
// SSH runtime clones the repository on a remote host
84+
description = `You are in a clone of a git repository at ${workspacePath}`;
85+
lines = gitCommonLines;
86+
break;
87+
88+
case RUNTIME_MODE.DOCKER:
89+
// Docker runtime runs in an isolated container
90+
description = `You are in a clone of a git repository at ${workspacePath} inside a Docker container`;
91+
lines = gitCommonLines;
92+
break;
93+
94+
default:
95+
assertNever(runtimeType, `Unknown runtime type: ${String(runtimeType)}`);
9296
}
9397

94-
// Worktree runtime creates a git worktree locally
9598
return `
9699
<environment>
97-
You are in a git worktree at ${workspacePath}
100+
${description}
98101
99-
- This IS a git repository - run git commands directly (no cd needed)
100-
- Tools run here automatically
101-
- Do not modify or visit other worktrees (especially the main project) without explicit user intent
102-
- You are meant to do your work isolated from the user and other agents
102+
${lines.join("\n")}
103103
</environment>
104104
`;
105105
}

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

Lines changed: 44 additions & 21 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,20 @@ 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)
170-
const showTrunkBranchSelector =
171-
props.branches.length > 0 && props.runtimeMode !== RUNTIME_MODE.LOCAL;
171+
const showTrunkBranchSelector = props.branches.length > 0 && runtimeMode !== RUNTIME_MODE.LOCAL;
172172

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

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

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

312335
{/* SSH Host Input */}
313-
{props.runtimeMode === RUNTIME_MODE.SSH && (
336+
{selectedRuntime.mode === "ssh" && (
314337
<div className="flex items-center gap-2">
315338
<label className="text-muted-foreground text-xs">host</label>
316339
<input
317340
type="text"
318-
value={props.sshHost}
319-
onChange={(e) => props.onSshHostChange(e.target.value)}
341+
value={selectedRuntime.host}
342+
onChange={(e) => onSelectedRuntimeChange({ mode: "ssh", host: e.target.value })}
320343
placeholder="user@host"
321344
disabled={props.disabled}
322345
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 +348,13 @@ export function CreationControls(props: CreationControlsProps) {
325348
)}
326349

327350
{/* Docker Image Input */}
328-
{props.runtimeMode === RUNTIME_MODE.DOCKER && (
351+
{selectedRuntime.mode === "docker" && (
329352
<div className="flex items-center gap-2">
330353
<label className="text-muted-foreground text-xs">image</label>
331354
<input
332355
type="text"
333-
value={props.dockerImage}
334-
onChange={(e) => props.onDockerImageChange(e.target.value)}
356+
value={selectedRuntime.image}
357+
onChange={(e) => onSelectedRuntimeChange({ mode: "docker", image: e.target.value })}
335358
placeholder="ubuntu:22.04"
336359
disabled={props.disabled}
337360
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
};

0 commit comments

Comments
 (0)