🤖 feat: add default vs custom trunk branch toggle

Generated with

Author: ibetitsmike

src/browser/components/ChatInput/CreationControls.tsx

@@ -2,11 +2,14 @@ import React from "react";
 import { RUNTIME_MODE, type RuntimeMode } from "@/common/types/runtime";
 import { TooltipWrapper, Tooltip } from "../Tooltip";
 import { Select } from "../Select";
+import { DEFAULT_TRUNK_BRANCH, TRUNK_SELECTION, type TrunkSelection } from "@/common/constants/workspace";
 
 interface CreationControlsProps {
   branches: string[];
-  trunkBranch: string;
-  onTrunkBranchChange: (branch: string) => void;
+  trunkSelection: TrunkSelection;
+  customTrunkBranch: string;
+  onTrunkSelectionChange: (selection: TrunkSelection) => void;
+  onCustomTrunkBranchChange: (branch: string) => void;
   runtimeMode: RuntimeMode;
   sshHost: string;
   onRuntimeChange: (mode: RuntimeMode, host: string) => void;
@@ -19,24 +22,52 @@ interface CreationControlsProps {
  * - Runtime mode (local vs SSH)
  */
 export function CreationControls(props: CreationControlsProps) {
+  const defaultUnavailable =
+    props.branches.length > 0 && !props.branches.includes(DEFAULT_TRUNK_BRANCH);
+  const showCustomPicker = props.trunkSelection === TRUNK_SELECTION.CUSTOM;
+
+  const handleTrunkSelectionChange = (value: string) => {
+    const selection = value as TrunkSelection;
+    if (selection === TRUNK_SELECTION.DEFAULT && defaultUnavailable) {
+      return;
+    }
+    props.onTrunkSelectionChange(selection);
+  };
+
   return (
     <div className="flex flex-wrap items-center gap-x-3 gap-y-2">
       {/* Trunk Branch Selector */}
-      {props.branches.length > 0 && (
-        <div className="flex items-center gap-1" data-component="TrunkBranchGroup">
-          <label htmlFor="trunk-branch" className="text-muted text-xs">
-            From:
-          </label>
+      <div className="flex items-center gap-1" data-component="TrunkBranchGroup">
+        <label htmlFor="trunk-branch-mode" className="text-muted text-xs">
+          From:
+        </label>
+        <Select
+          id="trunk-branch-mode"
+          value={props.trunkSelection}
+          options={[
+            { value: TRUNK_SELECTION.DEFAULT, label: "Main" },
+            { value: TRUNK_SELECTION.CUSTOM, label: "Custom" },
+          ]}
+          onChange={handleTrunkSelectionChange}
+          disabled={props.disabled}
+          className="max-w-[110px]"
+          aria-label="Trunk branch selection"
+        />
+        {defaultUnavailable && (
+          <span className="text-muted text-[11px]">Main branch not found</span>
+        )}
+        {showCustomPicker && props.branches.length > 0 && (
           <Select
             id="trunk-branch"
-            value={props.trunkBranch}
+            value={props.customTrunkBranch}
             options={props.branches}
-            onChange={props.onTrunkBranchChange}
+            onChange={props.onCustomTrunkBranchChange}
             disabled={props.disabled}
-            className="max-w-[120px]"
+            className="max-w-[130px]"
+            aria-label="Custom trunk branch"
           />
-        </div>
-      )}
+        )}
+      </div>
 
       {/* Runtime Selector */}
       <div className="flex items-center gap-1" data-component="RuntimeSelectorGroup">

src/browser/components/ChatInput/index.tsx

@@ -1016,8 +1016,10 @@ export const ChatInput: React.FC<ChatInputProps> = (props) => {
             {variant === "creation" && (
               <CreationControls
                 branches={creationState.branches}
-                trunkBranch={creationState.trunkBranch}
-                onTrunkBranchChange={creationState.setTrunkBranch}
+                trunkSelection={creationState.trunkSelection}
+                customTrunkBranch={creationState.customTrunkBranch}
+                onTrunkSelectionChange={creationState.setTrunkSelection}
+                onCustomTrunkBranchChange={creationState.setTrunkBranch}
                 runtimeMode={creationState.runtimeMode}
                 sshHost={creationState.sshHost}
                 onRuntimeChange={creationState.setRuntimeOptions}

src/browser/components/ChatInput/useCreationWorkspace.ts

@@ -5,6 +5,7 @@ import type { UIMode } from "@/common/types/mode";
 import type { ThinkingLevel } from "@/common/types/thinking";
 import { parseRuntimeString } from "@/browser/utils/chatCommands";
 import { useDraftWorkspaceSettings } from "@/browser/hooks/useDraftWorkspaceSettings";
+import type { TrunkSelection } from "@/common/constants/workspace";
 import { readPersistedState, updatePersistedState } from "@/browser/hooks/usePersistedState";
 import { useSendMessageOptions } from "@/browser/hooks/useSendMessageOptions";
 import { getModeKey, getProjectScopeId, getThinkingLevelKey } from "@/common/constants/storage";
@@ -35,7 +36,10 @@ function syncCreationPreferences(projectPath: string, workspaceId: string): void
 interface UseCreationWorkspaceReturn {
   branches: string[];
   trunkBranch: string;
+  trunkSelection: TrunkSelection;
+  customTrunkBranch: string;
   setTrunkBranch: (branch: string) => void;
+  setTrunkSelection: (selection: TrunkSelection) => void;
   runtimeMode: RuntimeMode;
   sshHost: string;
   setRuntimeOptions: (mode: RuntimeMode, host: string) => void;
@@ -62,7 +66,7 @@ export function useCreationWorkspace({
   const [isSending, setIsSending] = useState(false);
 
   // Centralized draft workspace settings with automatic persistence
-  const { settings, setRuntimeOptions, setTrunkBranch, getRuntimeString } =
+  const { settings, setRuntimeOptions, setTrunkBranch, setTrunkSelection, getRuntimeString } =
     useDraftWorkspaceSettings(projectPath, branches, recommendedTrunk);
 
   // Get send options from shared hook (uses project-scoped storage key)
@@ -145,7 +149,10 @@ export function useCreationWorkspace({
   return {
     branches,
     trunkBranch: settings.trunkBranch,
+    trunkSelection: settings.trunkSelection,
+    customTrunkBranch: settings.customTrunkBranch,
     setTrunkBranch,
+    setTrunkSelection,
     runtimeMode: settings.runtimeMode,
     sshHost: settings.sshHost,
     setRuntimeOptions,

src/browser/hooks/useDraftWorkspaceSettings.ts

@@ -13,8 +13,10 @@ import {
   getModelKey,
   getRuntimeKey,
   getTrunkBranchKey,
+  getTrunkSelectionKey,
   getProjectScopeId,
 } from "@/common/constants/storage";
+import { DEFAULT_TRUNK_BRANCH, TRUNK_SELECTION, type TrunkSelection } from "@/common/constants/workspace";
 import type { UIMode } from "@/common/types/mode";
 import type { ThinkingLevel } from "@/common/types/thinking";
 
@@ -33,6 +35,8 @@ export interface DraftWorkspaceSettings {
   runtimeMode: RuntimeMode;
   sshHost: string;
   trunkBranch: string;
+  trunkSelection: TrunkSelection;
+  customTrunkBranch: string;
 }
 
 /**
@@ -52,6 +56,7 @@ export function useDraftWorkspaceSettings(
   settings: DraftWorkspaceSettings;
   setRuntimeOptions: (mode: RuntimeMode, host: string) => void;
   setTrunkBranch: (branch: string) => void;
+  setTrunkSelection: (selection: TrunkSelection) => void;
   getRuntimeString: () => string | undefined;
 } {
   // Global AI settings (read-only from global state)
@@ -75,22 +80,50 @@ export function useDraftWorkspaceSettings(
   );
 
   // Project-scoped trunk branch preference (persisted per project)
-  const [trunkBranch, setTrunkBranch] = usePersistedState<string>(
+  const [customTrunkBranch, setCustomTrunkBranch] = usePersistedState<string>(
     getTrunkBranchKey(projectPath),
     "",
     { listener: true }
   );
 
+  const [trunkSelection, setTrunkSelection] = usePersistedState<TrunkSelection>(
+    getTrunkSelectionKey(projectPath),
+    TRUNK_SELECTION.DEFAULT,
+    { listener: true }
+  );
+
   // Parse runtime string into mode and host
   const { mode: runtimeMode, host: sshHost } = parseRuntimeModeAndHost(runtimeString);
 
-  // Initialize trunk branch from backend recommendation or first branch
+  // Initialize custom trunk branch from backend recommendation or first branch
   useEffect(() => {
-    if (!trunkBranch && branches.length > 0) {
-      const defaultBranch = recommendedTrunk ?? branches[0];
-      setTrunkBranch(defaultBranch);
+    if (branches.length === 0) {
+      return;
     }
-  }, [branches, recommendedTrunk, trunkBranch, setTrunkBranch]);
+
+    const recommendedInList = recommendedTrunk && branches.includes(recommendedTrunk);
+    const currentIsValid = customTrunkBranch && branches.includes(customTrunkBranch);
+
+    if (currentIsValid) {
+      return;
+    }
+
+    const fallback = (recommendedInList ? recommendedTrunk : undefined) ?? branches[0];
+    setCustomTrunkBranch(fallback);
+  }, [branches, recommendedTrunk, customTrunkBranch, setCustomTrunkBranch]);
+
+  // Fall back to custom mode when default "main" is unavailable in the repo
+  useEffect(() => {
+    if (
+      branches.length === 0 ||
+      trunkSelection === TRUNK_SELECTION.CUSTOM ||
+      branches.includes(DEFAULT_TRUNK_BRANCH)
+    ) {
+      return;
+    }
+
+    setTrunkSelection(TRUNK_SELECTION.CUSTOM);
+  }, [branches, trunkSelection, setTrunkSelection]);
 
   // Setter for runtime options (updates persisted runtime string)
   const setRuntimeOptions = (newMode: RuntimeMode, newHost: string) => {
@@ -103,6 +136,17 @@ export function useDraftWorkspaceSettings(
     return buildRuntimeString(runtimeMode, sshHost);
   };
 
+  const resolvedCustomBranch =
+    customTrunkBranch ||
+    (recommendedTrunk ?? branches[0]) ||
+    DEFAULT_TRUNK_BRANCH;
+
+  const defaultAvailable = branches.length === 0 || branches.includes(DEFAULT_TRUNK_BRANCH);
+  const effectiveTrunkBranch =
+    trunkSelection === TRUNK_SELECTION.DEFAULT && defaultAvailable
+      ? DEFAULT_TRUNK_BRANCH
+      : resolvedCustomBranch;
+
   return {
     settings: {
       model,
@@ -111,10 +155,13 @@ export function useDraftWorkspaceSettings(
       use1M,
       runtimeMode,
       sshHost,
-      trunkBranch,
+      trunkBranch: effectiveTrunkBranch,
+      trunkSelection,
+      customTrunkBranch: resolvedCustomBranch,
     },
     setRuntimeOptions,
-    setTrunkBranch,
+    setTrunkBranch: setCustomTrunkBranch,
+    setTrunkSelection,
     getRuntimeString,
   };
 }

src/browser/hooks/useStartWorkspaceCreation.test.ts

@@ -12,8 +12,10 @@ import {
   getProjectScopeId,
   getRuntimeKey,
   getTrunkBranchKey,
+  getTrunkSelectionKey,
 } from "@/common/constants/storage";
 import type { ProjectConfig } from "@/node/config";
+import { TRUNK_SELECTION } from "@/common/constants/workspace";
 
 import type { updatePersistedState } from "@/browser/hooks/usePersistedState";
 
@@ -71,6 +73,7 @@ describe("persistWorkspaceCreationPrefill", () => {
     expect(callMap.get(getInputKey(getPendingScopeId(projectPath)))).toBe("Ship it");
     expect(callMap.get(getModelKey(getProjectScopeId(projectPath)))).toBe("provider/model");
     expect(callMap.get(getTrunkBranchKey(projectPath))).toBe("main");
+    expect(callMap.get(getTrunkSelectionKey(projectPath))).toBe(TRUNK_SELECTION.CUSTOM);
     expect(callMap.get(getRuntimeKey(projectPath))).toBe("ssh dev");
   });
 
@@ -90,6 +93,7 @@ describe("persistWorkspaceCreationPrefill", () => {
     }
 
     expect(callMap.get(getTrunkBranchKey(projectPath))).toBeUndefined();
+    expect(callMap.get(getTrunkSelectionKey(projectPath))).toBe(TRUNK_SELECTION.DEFAULT);
     expect(callMap.get(getRuntimeKey(projectPath))).toBeUndefined();
   });
 

src/browser/hooks/useStartWorkspaceCreation.ts

@@ -10,8 +10,10 @@ import {
   getProjectScopeId,
   getRuntimeKey,
   getTrunkBranchKey,
+  getTrunkSelectionKey,
 } from "@/common/constants/storage";
 import { RUNTIME_MODE, SSH_RUNTIME_PREFIX } from "@/common/types/runtime";
+import { TRUNK_SELECTION } from "@/common/constants/workspace";
 
 export type StartWorkspaceCreationDetail =
   CustomEventPayloads[typeof CUSTOM_EVENTS.START_WORKSPACE_CREATION];
@@ -69,9 +71,11 @@ export function persistWorkspaceCreationPrefill(
 
   if (detail.trunkBranch !== undefined) {
     const normalizedTrunk = detail.trunkBranch.trim();
+    const hasCustomTrunk = normalizedTrunk.length > 0;
+    persist(getTrunkBranchKey(projectPath), hasCustomTrunk ? normalizedTrunk : undefined);
     persist(
-      getTrunkBranchKey(projectPath),
-      normalizedTrunk.length > 0 ? normalizedTrunk : undefined
+      getTrunkSelectionKey(projectPath),
+      hasCustomTrunk ? TRUNK_SELECTION.CUSTOM : TRUNK_SELECTION.DEFAULT
     );
   }
 

src/common/constants/storage.ts

@@ -108,6 +108,14 @@ export function getTrunkBranchKey(projectPath: string): string {
   return `trunkBranch:${projectPath}`;
 }
 
+/**
+ * Get the localStorage key for trunk branch selection mode (default vs custom)
+ * Format: "trunkSelection:{projectPath}"
+ */
+export function getTrunkSelectionKey(projectPath: string): string {
+  return `trunkSelection:${projectPath}`;
+}
+
 /**
  * Get the localStorage key for the 1M context preference (global)
  * Format: "use1MContext"

src/common/constants/workspace.ts

@@ -8,3 +8,15 @@ export const DEFAULT_RUNTIME_CONFIG: RuntimeConfig = {
   type: "local",
   srcBaseDir: "~/.mux/src",
 } as const;
+
+/**
+ * Default trunk branch to fork from when creating workspaces
+ */
+export const DEFAULT_TRUNK_BRANCH = "main" as const;
+
+export const TRUNK_SELECTION = {
+  DEFAULT: "default",
+  CUSTOM: "custom",
+} as const;
+
+export type TrunkSelection = (typeof TRUNK_SELECTION)[keyof typeof TRUNK_SELECTION];