From 4ab5fca4ca9d8fb12b1442ed46107e15c4275bf2 Mon Sep 17 00:00:00 2001 From: Julius Marminge Date: Sun, 1 Mar 2026 22:26:54 -0800 Subject: [PATCH] Preserve draft env mode across thread and branch context updates - Track and persist `envMode` in composer draft thread state - Keep worktree mode for local drafts even when `worktreePath` is null - Update branch toolbar/chat/sidebar to use effective draft env mode - Add tests for env mode persistence and undefined option merges --- apps/web/src/components/BranchToolbar.tsx | 21 ++++++---- apps/web/src/components/ChatView.tsx | 23 ++++++++--- apps/web/src/components/Sidebar.tsx | 16 ++++++-- apps/web/src/composerDraftStore.test.ts | 30 ++++++++++++++ apps/web/src/composerDraftStore.ts | 48 +++++++++++++++++------ 5 files changed, 111 insertions(+), 27 deletions(-) diff --git a/apps/web/src/components/BranchToolbar.tsx b/apps/web/src/components/BranchToolbar.tsx index bbe0447a7f7923da82cf7c34d07968c4301f1e5b..ce095c3b9c57ace35bd2acb35d02fa5bd98ae5e0 100644 --- a/apps/web/src/components/BranchToolbar.tsx +++ b/apps/web/src/components/BranchToolbar.tsx @@ -59,6 +59,10 @@ export default function BranchToolbar({ const activeWorktreePath = serverThread?.worktreePath ?? draftThread?.worktreePath ?? null; const branchCwd = activeWorktreePath ?? activeProject?.cwd ?? null; const hasServerThread = serverThread !== undefined; + const effectiveEnvMode: "local" | "worktree" = + activeWorktreePath || (!hasServerThread && draftThread?.envMode === "worktree") + ? "worktree" + : envMode; // ── Queries ─────────────────────────────────────────────────────────── @@ -68,7 +72,7 @@ export default function BranchToolbar({ const branchNames = branches.map((branch) => branch.name); const branchByName = new Map(branches.map((branch) => [branch.name, branch])); const trimmedBranchQuery = branchQuery.trim(); - const canCreateBranch = envMode === "local" && trimmedBranchQuery.length > 0; + const canCreateBranch = effectiveEnvMode === "local" && trimmedBranchQuery.length > 0; const hasExactBranchMatch = branchByName.has(trimmedBranchQuery); const createBranchItemValue = canCreateBranch ? `__create_new_branch__:${trimmedBranchQuery}` @@ -93,7 +97,7 @@ export default function BranchToolbar({ const syncedBranch = deriveSyncedLocalBranch({ activeThreadId, activeWorktreePath, - envMode, + envMode: effectiveEnvMode, activeThreadBranch, queryBranches, }); @@ -127,7 +131,7 @@ export default function BranchToolbar({ activeWorktreePath, activeThreadBranch, queryBranches, - envMode, + effectiveEnvMode, dispatch, hasServerThread, setDraftThreadContext, @@ -178,6 +182,7 @@ export default function BranchToolbar({ setDraftThreadContext(threadId, { branch, worktreePath, + envMode: effectiveEnvMode, }); }; @@ -186,7 +191,7 @@ export default function BranchToolbar({ if (!api || !activeThreadId || !branchCwd) return; // For new worktree mode, selecting a branch picks the base branch. - if (envMode === "worktree" && !envLocked && !activeWorktreePath) { + if (effectiveEnvMode === "worktree" && !envLocked && !activeWorktreePath) { setThreadError(null); setThreadBranch(branch.name, null); setIsBranchMenuOpen(false); @@ -265,9 +270,11 @@ export default function BranchToolbar({ variant="ghost" className="text-muted-foreground/70 hover:text-foreground/80" size="xs" - onClick={() => onEnvModeChange(envMode === "local" ? "worktree" : "local")} + onClick={() => + onEnvModeChange(effectiveEnvMode === "local" ? "worktree" : "local") + } > - {envMode === "worktree" ? "New worktree" : "Local"} + {effectiveEnvMode === "worktree" ? "New worktree" : "Local"} )} @@ -286,7 +293,7 @@ export default function BranchToolbar({ > {activeThreadBranch - ? envMode === "worktree" && !activeWorktreePath + ? effectiveEnvMode === "worktree" && !activeWorktreePath ? `From ${activeThreadBranch}` : activeThreadBranch : "Select branch"} diff --git a/apps/web/src/components/ChatView.tsx b/apps/web/src/components/ChatView.tsx index af023f741b333a325a2e9c34f6ffc7a8e54ccd02..5017fb3f47df4a0b4392c8097aa844e20cbadf1c 100644 --- a/apps/web/src/components/ChatView.tsx +++ b/apps/web/src/components/ChatView.tsx @@ -114,6 +114,7 @@ import { readNativeApi } from "~/nativeApi"; import { getAppModelOptions, useAppSettings } from "../appSettings"; import { type ComposerImageAttachment, + type DraftThreadEnvMode, type DraftThreadState, type PersistedComposerImageAttachment, useComposerDraftStore, @@ -465,6 +466,7 @@ export default function ChatView({ threadId }: ChatViewProps) { ); const clearComposerDraftContent = useComposerDraftStore((store) => store.clearComposerContent); const clearDraftThread = useComposerDraftStore((store) => store.clearDraftThread); + const setDraftThreadContext = useComposerDraftStore((store) => store.setDraftThreadContext); const draftThread = useComposerDraftStore((store) => store.draftThreadsByThreadId[threadId] ?? null); const promptRef = useRef(prompt); const [isDragOverComposer, setIsDragOverComposer] = useState(false); @@ -476,7 +478,7 @@ export default function ChatView({ threadId }: ChatViewProps) { const [sendPhase, setSendPhase] = useState("idle"); const [isConnecting, _setIsConnecting] = useState(false); const [isRevertingCheckpoint, setIsRevertingCheckpoint] = useState(false); - const [envMode, setEnvMode] = useState<"local" | "worktree">("local"); + const [envMode, setEnvMode] = useState("local"); const [isSwitchingRuntimeMode, setIsSwitchingRuntimeMode] = useState(false); const [respondingRequestIds, setRespondingRequestIds] = useState([]); const [expandedWorkGroups, setExpandedWorkGroups] = useState>({}); @@ -1386,8 +1388,16 @@ export default function ChatView({ threadId }: ChatViewProps) { useEffect(() => { if (!activeThread?.id) return; - setEnvMode(activeWorktreePath ? "worktree" : "local"); - }, [activeThread?.id, activeWorktreePath]); + if (activeWorktreePath) { + setEnvMode("worktree"); + return; + } + if (isLocalDraftThread) { + setEnvMode(draftThread?.envMode ?? "local"); + return; + } + setEnvMode("local"); + }, [activeThread?.id, activeWorktreePath, draftThread?.envMode, isLocalDraftThread]); // Auto-resize textarea useEffect(() => { @@ -1973,11 +1983,14 @@ export default function ChatView({ threadId }: ChatViewProps) { [scheduleComposerFocus, setComposerDraftEffort, threadId], ); const onEnvModeChange = useCallback( - (mode: "local" | "worktree") => { + (mode: DraftThreadEnvMode) => { setEnvMode(mode); + if (isLocalDraftThread) { + setDraftThreadContext(threadId, { envMode: mode }); + } scheduleComposerFocus(); }, - [scheduleComposerFocus], + [isLocalDraftThread, scheduleComposerFocus, setDraftThreadContext, threadId], ); const applyPromptReplacement = useCallback( diff --git a/apps/web/src/components/Sidebar.tsx b/apps/web/src/components/Sidebar.tsx index 9e38eb9dfc808846be8daad09d485dfd74b02a05..f6e477417a408991bce4098b8a6d7539dde90d11 100644 --- a/apps/web/src/components/Sidebar.tsx +++ b/apps/web/src/components/Sidebar.tsx @@ -20,7 +20,7 @@ import { derivePendingApprovals } from "../session-logic"; import { gitRemoveWorktreeMutationOptions, gitStatusQueryOptions } from "../lib/gitReactQuery"; import { serverConfigQueryOptions } from "../lib/serverReactQuery"; import { readNativeApi } from "../nativeApi"; -import { useComposerDraftStore } from "../composerDraftStore"; +import { type DraftThreadEnvMode, useComposerDraftStore } from "../composerDraftStore"; import { toastManager } from "./ui/toast"; import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "./ui/collapsible"; import { Tooltip, TooltipPopup, TooltipTrigger } from "./ui/tooltip"; @@ -345,17 +345,20 @@ export default function Sidebar() { options?: { branch?: string | null; worktreePath?: string | null; + envMode?: DraftThreadEnvMode; }, ): Promise => { const hasBranchOption = options?.branch !== undefined; const hasWorktreePathOption = options?.worktreePath !== undefined; + const hasEnvModeOption = options?.envMode !== undefined; const storedDraftThread = getDraftThreadByProjectId(projectId); if (storedDraftThread) { return (async () => { - if (hasBranchOption || hasWorktreePathOption) { + if (hasBranchOption || hasWorktreePathOption || hasEnvModeOption) { setDraftThreadContext(storedDraftThread.threadId, { ...(hasBranchOption ? { branch: options?.branch ?? null } : {}), ...(hasWorktreePathOption ? { worktreePath: options?.worktreePath ?? null } : {}), + ...(hasEnvModeOption ? { envMode: options?.envMode } : {}), }); } setProjectDraftThreadId(projectId, storedDraftThread.threadId); @@ -372,10 +375,11 @@ export default function Sidebar() { const activeDraftThread = routeThreadId ? getDraftThread(routeThreadId) : null; if (activeDraftThread && routeThreadId && activeDraftThread.projectId === projectId) { - if (hasBranchOption || hasWorktreePathOption) { + if (hasBranchOption || hasWorktreePathOption || hasEnvModeOption) { setDraftThreadContext(routeThreadId, { ...(hasBranchOption ? { branch: options?.branch ?? null } : {}), ...(hasWorktreePathOption ? { worktreePath: options?.worktreePath ?? null } : {}), + ...(hasEnvModeOption ? { envMode: options?.envMode } : {}), }); } setProjectDraftThreadId(projectId, routeThreadId); @@ -388,6 +392,7 @@ export default function Sidebar() { createdAt, branch: options?.branch ?? null, worktreePath: options?.worktreePath ?? null, + envMode: options?.envMode ?? "local", }); await navigate({ @@ -684,6 +689,9 @@ export default function Sidebar() { void handleNewThread(projectId, { branch: activeThread?.branch ?? activeDraftThread?.branch ?? null, worktreePath: activeThread?.worktreePath ?? activeDraftThread?.worktreePath ?? null, + envMode: + activeDraftThread?.envMode ?? + (activeThread?.worktreePath ? "worktree" : "local"), }); }; @@ -691,7 +699,7 @@ export default function Sidebar() { return () => { window.removeEventListener("keydown", onWindowKeyDown); }; - }, [handleNewThread, keybindings, routeThreadId, state.projects, state.threads]); + }, [getDraftThread, handleNewThread, keybindings, routeThreadId, state.projects, state.threads]); const onCreateThreadClick = () => { if (state.projects.length === 0) { diff --git a/apps/web/src/composerDraftStore.test.ts b/apps/web/src/composerDraftStore.test.ts index c994fcb79f78fbee0fcc09efd2568ad2bca124d6..12c46a21b16cdb951df62c5f00a4b493dea92206 100644 --- a/apps/web/src/composerDraftStore.test.ts +++ b/apps/web/src/composerDraftStore.test.ts @@ -183,12 +183,14 @@ describe("composerDraftStore project draft thread mapping", () => { projectId, branch: "feature/test", worktreePath: "/tmp/worktree-test", + envMode: "worktree", createdAt: "2026-01-01T00:00:00.000Z", }); expect(useComposerDraftStore.getState().getDraftThread(threadId)).toEqual({ projectId, branch: "feature/test", worktreePath: "/tmp/worktree-test", + envMode: "worktree", createdAt: "2026-01-01T00:00:00.000Z", }); }); @@ -273,6 +275,7 @@ describe("composerDraftStore project draft thread mapping", () => { projectId, branch: "feature/next", worktreePath: "/tmp/feature-next", + envMode: "worktree", }); }); @@ -295,6 +298,33 @@ describe("composerDraftStore project draft thread mapping", () => { projectId, branch: "main", worktreePath: "/tmp/main-worktree", + envMode: "worktree", + }); + }); + + it("preserves worktree env mode without a worktree path", () => { + const store = useComposerDraftStore.getState(); + store.setProjectDraftThreadId(projectId, threadId, { + branch: "feature/base", + worktreePath: null, + envMode: "worktree", + }); + const runtimeUndefinedOptions = { + branch: undefined, + worktreePath: undefined, + envMode: undefined, + } as unknown as { + branch?: string | null; + worktreePath?: string | null; + envMode?: "local" | "worktree"; + }; + store.setProjectDraftThreadId(projectId, threadId, runtimeUndefinedOptions); + + expect(useComposerDraftStore.getState().getDraftThread(threadId)).toMatchObject({ + projectId, + branch: "feature/base", + worktreePath: null, + envMode: "worktree", }); }); }); diff --git a/apps/web/src/composerDraftStore.ts b/apps/web/src/composerDraftStore.ts index 1606f2aeb1af30d84d44c8f0897686afa0a5b836..47fdd719221d3daa03dde04f76ff51c597d91f46 100644 --- a/apps/web/src/composerDraftStore.ts +++ b/apps/web/src/composerDraftStore.ts @@ -11,6 +11,7 @@ import { create } from "zustand"; import { createJSONStorage, persist } from "zustand/middleware"; export const COMPOSER_DRAFT_STORAGE_KEY = "t3code:composer-drafts:v1"; +export type DraftThreadEnvMode = "local" | "worktree"; export interface PersistedComposerImageAttachment { id: string; @@ -38,6 +39,7 @@ interface PersistedDraftThreadState { createdAt: string; branch: string | null; worktreePath: string | null; + envMode: DraftThreadEnvMode; } interface PersistedComposerDraftStoreState { @@ -61,6 +63,7 @@ export interface DraftThreadState { createdAt: string; branch: string | null; worktreePath: string | null; + envMode: DraftThreadEnvMode; } interface ProjectDraftThread extends DraftThreadState { @@ -80,6 +83,7 @@ interface ComposerDraftStoreState { branch?: string | null; worktreePath?: string | null; createdAt?: string; + envMode?: DraftThreadEnvMode; }, ) => void; setDraftThreadContext: ( @@ -89,6 +93,7 @@ interface ComposerDraftStoreState { worktreePath?: string | null; projectId?: ProjectId; createdAt?: string; + envMode?: DraftThreadEnvMode; }, ) => void; clearProjectDraftThreadId: (projectId: ProjectId) => void; @@ -203,6 +208,16 @@ function normalizePersistedAttachment(value: unknown): PersistedComposerImageAtt }; } +function normalizeDraftThreadEnvMode( + value: unknown, + fallbackWorktreePath: string | null, +): DraftThreadEnvMode { + if (value === "local" || value === "worktree") { + return value; + } + return fallbackWorktreePath ? "worktree" : "local"; +} + function normalizePersistedComposerDraftState(value: unknown): PersistedComposerDraftStoreState { if (!value || typeof value !== "object") { return EMPTY_PERSISTED_DRAFT_STORE_STATE; @@ -227,6 +242,7 @@ function normalizePersistedComposerDraftState(value: unknown): PersistedComposer const createdAt = candidateDraftThread.createdAt; const branch = candidateDraftThread.branch; const worktreePath = candidateDraftThread.worktreePath; + const normalizedWorktreePath = typeof worktreePath === "string" ? worktreePath : null; if (typeof projectId !== "string" || projectId.length === 0) { continue; } @@ -237,7 +253,8 @@ function normalizePersistedComposerDraftState(value: unknown): PersistedComposer ? createdAt : new Date().toISOString(), branch: typeof branch === "string" ? branch : null, - worktreePath: typeof worktreePath === "string" ? worktreePath : null, + worktreePath: normalizedWorktreePath, + envMode: normalizeDraftThreadEnvMode(candidateDraftThread.envMode, normalizedWorktreePath), }; } } @@ -263,6 +280,7 @@ function normalizePersistedComposerDraftState(value: unknown): PersistedComposer createdAt: new Date().toISOString(), branch: null, worktreePath: null, + envMode: "local", }; } else if (draftThreadsByThreadId[threadId as ThreadId]?.projectId !== projectId) { draftThreadsByThreadId[threadId as ThreadId] = { @@ -456,6 +474,10 @@ export const useComposerDraftStore = create()( set((state) => { const existingThread = state.draftThreadsByThreadId[threadId]; const previousThreadIdForProject = state.projectDraftThreadIdByProjectId[projectId]; + const nextWorktreePath = + options?.worktreePath === undefined + ? (existingThread?.worktreePath ?? null) + : (options.worktreePath ?? null); const nextDraftThread: DraftThreadState = { projectId, createdAt: options?.createdAt ?? existingThread?.createdAt ?? new Date().toISOString(), @@ -463,10 +485,10 @@ export const useComposerDraftStore = create()( options?.branch === undefined ? (existingThread?.branch ?? null) : (options.branch ?? null), - worktreePath: - options?.worktreePath === undefined - ? (existingThread?.worktreePath ?? null) - : (options.worktreePath ?? null), + worktreePath: nextWorktreePath, + envMode: + options?.envMode ?? + (nextWorktreePath ? "worktree" : (existingThread?.envMode ?? "local")), }; const hasSameProjectMapping = previousThreadIdForProject === threadId; const hasSameDraftThread = @@ -474,7 +496,8 @@ export const useComposerDraftStore = create()( existingThread.projectId === nextDraftThread.projectId && existingThread.createdAt === nextDraftThread.createdAt && existingThread.branch === nextDraftThread.branch && - existingThread.worktreePath === nextDraftThread.worktreePath; + existingThread.worktreePath === nextDraftThread.worktreePath && + existingThread.envMode === nextDraftThread.envMode; if (hasSameProjectMapping && hasSameDraftThread) { return state; } @@ -518,6 +541,8 @@ export const useComposerDraftStore = create()( if (nextProjectId.length === 0) { return state; } + const nextWorktreePath = + options.worktreePath === undefined ? existing.worktreePath : (options.worktreePath ?? null); const nextDraftThread: DraftThreadState = { projectId: nextProjectId, createdAt: @@ -525,16 +550,17 @@ export const useComposerDraftStore = create()( ? existing.createdAt : options.createdAt || existing.createdAt, branch: options.branch === undefined ? existing.branch : (options.branch ?? null), - worktreePath: - options.worktreePath === undefined - ? existing.worktreePath - : (options.worktreePath ?? null), + worktreePath: nextWorktreePath, + envMode: + options.envMode ?? + (nextWorktreePath ? "worktree" : (existing.envMode ?? "local")), }; const isUnchanged = nextDraftThread.projectId === existing.projectId && nextDraftThread.createdAt === existing.createdAt && nextDraftThread.branch === existing.branch && - nextDraftThread.worktreePath === existing.worktreePath; + nextDraftThread.worktreePath === existing.worktreePath && + nextDraftThread.envMode === existing.envMode; if (isUnchanged) { return state; } -- 2.54.0