diff --git a/apps/web/src/components/ChatView.tsx b/apps/web/src/components/ChatView.tsx index c9c074b045e255b8aa74968ac220e912a3a1e304..7c600faa6e27698d06c0b72e4bba9cce9e58cee9 100644 --- a/apps/web/src/components/ChatView.tsx +++ b/apps/web/src/components/ChatView.tsx @@ -53,6 +53,7 @@ import { DEFAULT_THREAD_TERMINAL_ID, MAX_THREAD_TERMINAL_COUNT, type ChatMessage, + type Thread, type TurnDiffSummary, } from "../types"; import { basenameOfPath, getVscodeIconUrlForEntry } from "../vscode-icons"; @@ -113,6 +114,7 @@ import { getAppModelOptions, useAppSettings } from "../appSettings"; import { COMPOSER_DRAFT_STORAGE_KEY, type ComposerImageAttachment, + type DraftThreadState, type PersistedComposerImageAttachment, readPersistedComposerDraftAttachments, useComposerDraftStore, @@ -195,6 +197,42 @@ function buildExpandedImagePreview( }; } +function buildLocalDraftThread( + threadId: ThreadId, + draftThread: DraftThreadState, + fallbackModel: string, +): Thread { + return { + id: threadId, + codexThreadId: null, + projectId: draftThread.projectId, + title: "New thread", + model: fallbackModel, + terminalOpen: false, + terminalHeight: 280, + terminalIds: [DEFAULT_THREAD_TERMINAL_ID], + runningTerminalIds: [], + activeTerminalId: DEFAULT_THREAD_TERMINAL_ID, + terminalGroups: [ + { + id: `group-${DEFAULT_THREAD_TERMINAL_ID}`, + terminalIds: [DEFAULT_THREAD_TERMINAL_ID], + }, + ], + activeTerminalGroupId: `group-${DEFAULT_THREAD_TERMINAL_ID}`, + session: null, + messages: [], + error: null, + createdAt: draftThread.createdAt, + latestTurn: null, + lastVisitedAt: draftThread.createdAt, + branch: draftThread.branch, + worktreePath: draftThread.worktreePath, + turnDiffSummaries: [], + activities: [], + }; +} + type ComposerCommandItem = | { id: string; @@ -428,9 +466,8 @@ export default function ChatView({ threadId }: ChatViewProps) { (store) => store.syncPersistedAttachments, ); const clearComposerDraftContent = useComposerDraftStore((store) => store.clearComposerContent); - const clearProjectDraftThreadById = useComposerDraftStore( - (store) => store.clearProjectDraftThreadById, - ); + const clearDraftThread = useComposerDraftStore((store) => store.clearDraftThread); + const draftThread = useComposerDraftStore((store) => store.draftThreadsByThreadId[threadId] ?? null); const promptRef = useRef(prompt); const [isDragOverComposer, setIsDragOverComposer] = useState(false); const [expandedImage, setExpandedImage] = useState(null); @@ -488,7 +525,18 @@ export default function ChatView({ threadId }: ChatViewProps) { [removeComposerDraftImage, threadId], ); - const activeThread = state.threads.find((t) => t.id === threadId); + const serverThread = state.threads.find((t) => t.id === threadId); + const fallbackDraftProject = state.projects.find((project) => project.id === draftThread?.projectId); + const localDraftThread = useMemo( + () => + draftThread + ? buildLocalDraftThread(threadId, draftThread, fallbackDraftProject?.model ?? DEFAULT_MODEL) + : undefined, + [draftThread, fallbackDraftProject?.model, threadId], + ); + const activeThread = serverThread ?? localDraftThread; + const isServerThread = serverThread !== undefined; + const isLocalDraftThread = !isServerThread && localDraftThread !== undefined; const diffSearch = useMemo( () => parseDiffRouteSearch(rawSearch as Record), [rawSearch], @@ -876,7 +924,7 @@ export default function ChatView({ threadId }: ChatViewProps) { }, ) => { const api = readNativeApi(); - if (!api || !activeThreadId || !activeProject || !activeThread) return; + if (!api || !activeThreadId || !activeProject || !activeThread || !isServerThread) return; if (options?.rememberAsLastInvoked !== false) { setLastInvokedScriptByProjectId((current) => { if (current[activeProject.id] === script.id) return current; @@ -948,7 +996,7 @@ export default function ChatView({ threadId }: ChatViewProps) { }); } }, - [activeProject, activeThread, activeThreadId, dispatch, gitCwd], + [activeProject, activeThread, activeThreadId, dispatch, gitCwd, isServerThread], ); const persistProjectScripts = useCallback( async (input: { @@ -1634,7 +1682,7 @@ export default function ChatView({ threadId }: ChatViewProps) { if (!trimmed && composerImages.length === 0) return; if (!activeProject) return; const threadIdForSend = activeThread.id; - const isFirstMessage = activeThread.messages.length === 0; + const isFirstMessage = !isServerThread || activeThread.messages.length === 0; const baseBranchForWorktree = isFirstMessage && envMode === "worktree" && !activeThread.worktreePath ? activeThread.branch @@ -1693,7 +1741,10 @@ export default function ChatView({ threadId }: ChatViewProps) { clearComposerDraftContent(threadIdForSend); setComposerHighlightedItemId(null); - let attemptedTurnStart = false; + let createdServerThreadForLocalDraft = false; + let turnStartSucceeded = false; + let nextThreadBranch = activeThread.branch; + let nextThreadWorktreePath = activeThread.worktreePath; try { // On first message: lock in branch + create worktree if needed. if (baseBranchForWorktree) { @@ -1704,23 +1755,27 @@ export default function ChatView({ threadId }: ChatViewProps) { branch: baseBranchForWorktree, newBranch, }); - await api.orchestration.dispatchCommand({ - type: "thread.meta.update", - commandId: newCommandId(), - threadId: threadIdForSend, - branch: result.worktree.branch, - worktreePath: result.worktree.path, - }); - // Keep local thread state in sync immediately so terminal drawer opens - // with the worktree cwd/env instead of briefly using the project root. - dispatch({ - type: "SET_THREAD_BRANCH", - threadId: threadIdForSend, - branch: result.worktree.branch, - worktreePath: result.worktree.path, - }); + nextThreadBranch = result.worktree.branch; + nextThreadWorktreePath = result.worktree.path; + if (isServerThread) { + await api.orchestration.dispatchCommand({ + type: "thread.meta.update", + commandId: newCommandId(), + threadId: threadIdForSend, + branch: result.worktree.branch, + worktreePath: result.worktree.path, + }); + // Keep local thread state in sync immediately so terminal drawer opens + // with the worktree cwd/env instead of briefly using the project root. + dispatch({ + type: "SET_THREAD_BRANCH", + threadId: threadIdForSend, + branch: result.worktree.branch, + worktreePath: result.worktree.path, + }); + } const setupScript = setupProjectScript(activeProject.scripts); - if (setupScript) { + if (setupScript && isServerThread) { await runProjectScript(setupScript, { cwd: result.worktree.path, worktreePath: result.worktree.path, @@ -1729,14 +1784,30 @@ export default function ChatView({ threadId }: ChatViewProps) { } } + const titleSeed = + trimmed || + (composerImagesSnapshot.length > 0 + ? `Image: ${composerImagesSnapshot[0]?.name ?? "attachment"}` + : "New thread"); + const title = truncateTitle(titleSeed); + + if (isLocalDraftThread) { + await api.orchestration.dispatchCommand({ + type: "thread.create", + commandId: newCommandId(), + threadId: threadIdForSend, + projectId: activeProject.id, + title, + model: selectedModel ?? activeProject.model ?? DEFAULT_MODEL, + branch: nextThreadBranch, + worktreePath: nextThreadWorktreePath, + createdAt: activeThread.createdAt, + }); + createdServerThreadForLocalDraft = true; + } + // Auto-title from first message - if (isFirstMessage) { - const titleSeed = - trimmed || - (composerImagesSnapshot.length > 0 - ? `Image: ${composerImagesSnapshot[0]?.name ?? "attachment"}` - : "New thread"); - const title = truncateTitle(titleSeed); + if (isFirstMessage && isServerThread) { await api.orchestration.dispatchCommand({ type: "thread.meta.update", commandId: newCommandId(), @@ -1747,7 +1818,6 @@ export default function ChatView({ threadId }: ChatViewProps) { setSendPhase("sending-turn"); const turnAttachments = await turnAttachmentsPromise; - attemptedTurnStart = true; const approvalPolicy = state.runtimeMode === "full-access" ? "never" : "on-request"; const sandboxMode = state.runtimeMode === "full-access" ? "danger-full-access" : "workspace-write"; @@ -1768,12 +1838,22 @@ export default function ChatView({ threadId }: ChatViewProps) { sandboxMode, createdAt: messageCreatedAt, }); + turnStartSucceeded = true; if (isFirstMessage) { - clearProjectDraftThreadById(activeProject.id, threadIdForSend); + clearDraftThread(threadIdForSend); } } catch (err) { + if (createdServerThreadForLocalDraft && !turnStartSucceeded) { + await api.orchestration + .dispatchCommand({ + type: "thread.delete", + commandId: newCommandId(), + threadId: threadIdForSend, + }) + .catch(() => undefined); + } if ( - !attemptedTurnStart && + !turnStartSucceeded && promptRef.current.length === 0 && composerImagesRef.current.length === 0 ) { @@ -1841,7 +1921,7 @@ export default function ChatView({ threadId }: ChatViewProps) { const normalizedModel = resolveModelSlug(model); setComposerDraftModel(threadId, normalizedModel); const api = readNativeApi(); - if (api && activeThread) { + if (api && isServerThread && activeThread) { void api.orchestration.dispatchCommand({ type: "thread.meta.update", commandId: newCommandId(), @@ -1851,7 +1931,7 @@ export default function ChatView({ threadId }: ChatViewProps) { } scheduleComposerFocus(); }, - [activeThread, scheduleComposerFocus, setComposerDraftModel, threadId], + [activeThread, isServerThread, scheduleComposerFocus, setComposerDraftModel, threadId], ); const onEffortSelect = useCallback( (effort: ReasoningEffort) => { diff --git a/apps/web/src/components/Sidebar.tsx b/apps/web/src/components/Sidebar.tsx index f3863bff7770200cfe654291316d94b4d1fc0c13..495ed4491961b3b46799aa0062848029ed3d2ab8 100644 --- a/apps/web/src/components/Sidebar.tsx +++ b/apps/web/src/components/Sidebar.tsx @@ -75,10 +75,6 @@ interface PrStatusIndicator { type ThreadPr = GitStatusResult["pr"]; -function shouldShowThreadInSidebar(thread: Thread): boolean { - return thread.messages.some((message) => message.role === "user"); -} - function hasUnseenCompletion(thread: Thread): boolean { if (!thread.latestTurn?.completedAt) return false; const completedAt = Date.parse(thread.latestTurn.completedAt); @@ -234,8 +230,10 @@ function ProjectFavicon({ cwd }: { cwd: string }) { export default function Sidebar() { const { state, dispatch } = useStore(); const clearComposerDraftForThread = useComposerDraftStore((store) => store.clearThreadDraft); + const readDraftThread = useComposerDraftStore((store) => store.readDraftThread); const readProjectDraftThreadId = useComposerDraftStore((store) => store.readProjectDraftThreadId); const setProjectDraftThreadId = useComposerDraftStore((store) => store.setProjectDraftThreadId); + const setDraftThreadContext = useComposerDraftStore((store) => store.setDraftThreadContext); const clearProjectDraftThreadId = useComposerDraftStore((store) => store.clearProjectDraftThreadId); const clearProjectDraftThreadById = useComposerDraftStore( (store) => store.clearProjectDraftThreadById, @@ -346,55 +344,25 @@ export default function Sidebar() { worktreePath?: string | null; }, ): Promise => { - const api = readNativeApi(); - if (!api) return Promise.resolve(); const hasBranchOption = options?.branch !== undefined; const hasWorktreePathOption = options?.worktreePath !== undefined; - const applyThreadOptions = async (thread: Thread): Promise => { - if (!hasBranchOption && !hasWorktreePathOption) { - return; - } - const nextBranch = hasBranchOption ? (options?.branch ?? null) : thread.branch; - const nextWorktreePath = hasWorktreePathOption - ? (options?.worktreePath ?? null) - : thread.worktreePath; - if (thread.branch === nextBranch && thread.worktreePath === nextWorktreePath) { - return; - } - await api.orchestration.dispatchCommand({ - type: "thread.meta.update", - commandId: newCommandId(), - threadId: thread.id, - ...(hasBranchOption ? { branch: nextBranch } : {}), - ...(hasWorktreePathOption ? { worktreePath: nextWorktreePath } : {}), - }); - dispatch({ - type: "SET_THREAD_BRANCH", - threadId: thread.id, - branch: nextBranch, - worktreePath: nextWorktreePath, - }); - }; const storedDraftThreadId = readProjectDraftThreadId(projectId); - const storedDraftThread = storedDraftThreadId - ? state.threads.find((thread) => thread.id === storedDraftThreadId) - : undefined; - const reusableStoredDraftThread = - storedDraftThread && - storedDraftThread.projectId === projectId && - !shouldShowThreadInSidebar(storedDraftThread) - ? storedDraftThread - : null; - if (reusableStoredDraftThread) { + const storedDraftThread = storedDraftThreadId ? readDraftThread(storedDraftThreadId) : null; + if (storedDraftThread && storedDraftThread.projectId === projectId && storedDraftThreadId) { return (async () => { - await applyThreadOptions(reusableStoredDraftThread); - setProjectDraftThreadId(projectId, reusableStoredDraftThread.id); - if (routeThreadId === reusableStoredDraftThread.id) { + if (hasBranchOption || hasWorktreePathOption) { + setDraftThreadContext(storedDraftThreadId, { + ...(hasBranchOption ? { branch: options?.branch ?? null } : {}), + ...(hasWorktreePathOption ? { worktreePath: options?.worktreePath ?? null } : {}), + }); + } + setProjectDraftThreadId(projectId, storedDraftThreadId); + if (routeThreadId === storedDraftThreadId) { return; } await navigate({ to: "/$threadId", - params: { threadId: reusableStoredDraftThread.id }, + params: { threadId: storedDraftThreadId }, }); })(); } @@ -402,53 +370,25 @@ export default function Sidebar() { clearProjectDraftThreadId(projectId); } - const activeThread = routeThreadId - ? state.threads.find((thread) => thread.id === routeThreadId) - : undefined; - const reusableDraftThread = - activeThread && - activeThread.projectId === projectId && - !shouldShowThreadInSidebar(activeThread) - ? activeThread - : null; - if (reusableDraftThread) { - return (async () => { - await applyThreadOptions(reusableDraftThread); - setProjectDraftThreadId(projectId, reusableDraftThread.id); - if (routeThreadId === reusableDraftThread.id) { - return; - } - await navigate({ - to: "/$threadId", - params: { threadId: reusableDraftThread.id }, + const activeDraftThread = routeThreadId ? readDraftThread(routeThreadId) : null; + if (activeDraftThread && routeThreadId && activeDraftThread.projectId === projectId) { + if (hasBranchOption || hasWorktreePathOption) { + setDraftThreadContext(routeThreadId, { + ...(hasBranchOption ? { branch: options?.branch ?? null } : {}), + ...(hasWorktreePathOption ? { worktreePath: options?.worktreePath ?? null } : {}), }); - })(); + } + setProjectDraftThreadId(projectId, routeThreadId); + return Promise.resolve(); } const threadId = newThreadId(); const createdAt = new Date().toISOString(); - const model = - state.projects.find((project) => project.id === projectId)?.model ?? DEFAULT_MODEL; return (async () => { - await api.orchestration.dispatchCommand({ - type: "thread.create", - commandId: newCommandId(), - threadId, - projectId, - title: "New thread", - model, + setProjectDraftThreadId(projectId, threadId, { + createdAt, branch: options?.branch ?? null, worktreePath: options?.worktreePath ?? null, - createdAt, }); - setProjectDraftThreadId(projectId, threadId); - - // Ensure route guards can see the new thread before navigating. - try { - const snapshot = await api.orchestration.getSnapshot(); - dispatch({ type: "SYNC_SERVER_READ_MODEL", readModel: snapshot }); - } catch { - // Event stream can still hydrate the thread shortly after dispatch. - } await navigate({ to: "/$threadId", @@ -458,13 +398,12 @@ export default function Sidebar() { }, [ clearProjectDraftThreadId, - dispatch, navigate, + readDraftThread, readProjectDraftThreadId, routeThreadId, + setDraftThreadContext, setProjectDraftThreadId, - state.projects, - state.threads, ], ); @@ -678,9 +617,7 @@ export default function Sidebar() { if (!project) return; const projectThreads = state.threads.filter((thread) => thread.projectId === projectId); - const visibleThreads = projectThreads.filter(shouldShowThreadInSidebar); - const draftThreads = projectThreads.filter((thread) => !shouldShowThreadInSidebar(thread)); - if (visibleThreads.length > 0) { + if (projectThreads.length > 0) { toastManager.add({ type: "warning", title: "Project is not empty", @@ -695,13 +632,9 @@ export default function Sidebar() { if (!confirmed) return; try { - for (const draftThread of draftThreads) { - await api.orchestration.dispatchCommand({ - type: "thread.delete", - commandId: newCommandId(), - threadId: draftThread.id, - }); - clearComposerDraftForThread(draftThread.id); + const projectDraftThreadId = readProjectDraftThreadId(projectId); + if (projectDraftThreadId) { + clearComposerDraftForThread(projectDraftThreadId); } clearProjectDraftThreadId(projectId); await api.orchestration.dispatchCommand({ @@ -719,7 +652,13 @@ export default function Sidebar() { }); } }, - [clearComposerDraftForThread, clearProjectDraftThreadId, state.projects, state.threads], + [ + clearComposerDraftForThread, + clearProjectDraftThreadId, + readProjectDraftThreadId, + state.projects, + state.threads, + ], ); useEffect(() => { @@ -727,8 +666,9 @@ export default function Sidebar() { const activeThread = routeThreadId ? state.threads.find((thread) => thread.id === routeThreadId) : undefined; + const activeDraftThread = routeThreadId ? readDraftThread(routeThreadId) : null; if (isChatNewLocalShortcut(event, keybindings)) { - const projectId = activeThread?.projectId ?? state.projects[0]?.id; + const projectId = activeThread?.projectId ?? activeDraftThread?.projectId ?? state.projects[0]?.id; if (!projectId) return; event.preventDefault(); void handleNewThread(projectId); @@ -736,12 +676,12 @@ export default function Sidebar() { } if (!isChatNewShortcut(event, keybindings)) return; - const projectId = activeThread?.projectId ?? state.projects[0]?.id; + const projectId = activeThread?.projectId ?? activeDraftThread?.projectId ?? state.projects[0]?.id; if (!projectId) return; event.preventDefault(); void handleNewThread(projectId, { - branch: activeThread?.branch ?? null, - worktreePath: activeThread?.worktreePath ?? null, + branch: activeThread?.branch ?? activeDraftThread?.branch ?? null, + worktreePath: activeThread?.worktreePath ?? activeDraftThread?.worktreePath ?? null, }); }; @@ -749,7 +689,7 @@ export default function Sidebar() { return () => { window.removeEventListener("keydown", onWindowKeyDown); }; - }, [handleNewThread, keybindings, routeThreadId, state.projects, state.threads]); + }, [handleNewThread, keybindings, readDraftThread, routeThreadId, state.projects, state.threads]); const onCreateThreadClick = () => { if (state.projects.length === 0) { @@ -808,7 +748,6 @@ export default function Sidebar() { {state.projects.map((project) => { const visibleThreads = state.threads .filter((thread) => thread.projectId === project.id) - .filter(shouldShowThreadInSidebar) .toSorted((a, b) => { const byDate = new Date(b.createdAt).getTime() - new Date(a.createdAt).getTime(); if (byDate !== 0) return byDate; diff --git a/apps/web/src/composerDraftStore.test.ts b/apps/web/src/composerDraftStore.test.ts index a54832dcd4ac33d6e1afb3da424a74d3ecd43831..0563d7b1245505427e2225adf1e4dd5f73e5d05d 100644 --- a/apps/web/src/composerDraftStore.test.ts +++ b/apps/web/src/composerDraftStore.test.ts @@ -33,7 +33,11 @@ describe("composerDraftStore addImages", () => { let revokeSpy: ReturnType void>>; beforeEach(() => { - useComposerDraftStore.setState({ draftsByThreadId: {}, projectDraftThreadIdByProjectId: {} }); + useComposerDraftStore.setState({ + draftsByThreadId: {}, + draftThreadsByThreadId: {}, + projectDraftThreadIdByProjectId: {}, + }); originalRevokeObjectUrl = URL.revokeObjectURL; revokeSpy = vi.fn(); URL.revokeObjectURL = revokeSpy; @@ -118,7 +122,11 @@ describe("composerDraftStore clearComposerContent", () => { let revokeSpy: ReturnType void>>; beforeEach(() => { - useComposerDraftStore.setState({ draftsByThreadId: {}, projectDraftThreadIdByProjectId: {} }); + useComposerDraftStore.setState({ + draftsByThreadId: {}, + draftThreadsByThreadId: {}, + projectDraftThreadIdByProjectId: {}, + }); originalRevokeObjectUrl = URL.revokeObjectURL; revokeSpy = vi.fn(); URL.revokeObjectURL = revokeSpy; @@ -149,15 +157,30 @@ describe("composerDraftStore project draft thread mapping", () => { const otherThreadId = ThreadId.makeUnsafe("thread-b"); beforeEach(() => { - useComposerDraftStore.setState({ draftsByThreadId: {}, projectDraftThreadIdByProjectId: {} }); + useComposerDraftStore.setState({ + draftsByThreadId: {}, + draftThreadsByThreadId: {}, + projectDraftThreadIdByProjectId: {}, + }); }); it("stores and reads project draft thread ids via actions", () => { const store = useComposerDraftStore.getState(); expect(store.readProjectDraftThreadId(projectId)).toBeNull(); + expect(store.readDraftThread(threadId)).toBeNull(); - store.setProjectDraftThreadId(projectId, threadId); + store.setProjectDraftThreadId(projectId, threadId, { + branch: "feature/test", + worktreePath: "/tmp/worktree-test", + createdAt: "2026-01-01T00:00:00.000Z", + }); expect(useComposerDraftStore.getState().readProjectDraftThreadId(projectId)).toBe(threadId); + expect(useComposerDraftStore.getState().readDraftThread(threadId)).toEqual({ + projectId, + branch: "feature/test", + worktreePath: "/tmp/worktree-test", + createdAt: "2026-01-01T00:00:00.000Z", + }); }); it("clears only matching project draft mapping entries", () => { @@ -169,6 +192,7 @@ describe("composerDraftStore project draft thread mapping", () => { store.clearProjectDraftThreadById(projectId, threadId); expect(useComposerDraftStore.getState().readProjectDraftThreadId(projectId)).toBeNull(); + expect(useComposerDraftStore.getState().readDraftThread(threadId)).toBeNull(); }); it("clears project draft mapping by project id", () => { @@ -176,5 +200,32 @@ describe("composerDraftStore project draft thread mapping", () => { store.setProjectDraftThreadId(projectId, threadId); store.clearProjectDraftThreadId(projectId); expect(useComposerDraftStore.getState().readProjectDraftThreadId(projectId)).toBeNull(); + expect(useComposerDraftStore.getState().readDraftThread(threadId)).toBeNull(); + }); + + it("clears draft registration independently", () => { + const store = useComposerDraftStore.getState(); + store.setProjectDraftThreadId(projectId, threadId); + store.clearDraftThread(threadId); + expect(useComposerDraftStore.getState().readProjectDraftThreadId(projectId)).toBeNull(); + expect(useComposerDraftStore.getState().readDraftThread(threadId)).toBeNull(); + }); + + it("updates branch context on an existing draft thread", () => { + const store = useComposerDraftStore.getState(); + store.setProjectDraftThreadId(projectId, threadId, { + branch: "main", + worktreePath: null, + }); + store.setDraftThreadContext(threadId, { + branch: "feature/next", + worktreePath: "/tmp/feature-next", + }); + expect(useComposerDraftStore.getState().readProjectDraftThreadId(projectId)).toBe(threadId); + expect(useComposerDraftStore.getState().readDraftThread(threadId)).toMatchObject({ + projectId, + branch: "feature/next", + worktreePath: "/tmp/feature-next", + }); }); }); diff --git a/apps/web/src/composerDraftStore.ts b/apps/web/src/composerDraftStore.ts index 49fba01777539afd6e1040cfd9782cba66018d75..40064e5eb021e8fba54479eda5b3a1a888419523 100644 --- a/apps/web/src/composerDraftStore.ts +++ b/apps/web/src/composerDraftStore.ts @@ -66,8 +66,16 @@ interface PersistedComposerThreadDraftState { effort?: ReasoningEffort | null; } +interface PersistedDraftThreadState { + projectId: ProjectId; + createdAt: string; + branch: string | null; + worktreePath: string | null; +} + interface PersistedComposerDraftStoreState { draftsByThreadId: Record; + draftThreadsByThreadId: Record; projectDraftThreadIdByProjectId: Record; } @@ -81,13 +89,41 @@ export interface ComposerThreadDraftState { effort: ReasoningEffort | null; } +export interface DraftThreadState { + projectId: ProjectId; + createdAt: string; + branch: string | null; + worktreePath: string | null; +} + interface ComposerDraftStoreState { draftsByThreadId: Record; + draftThreadsByThreadId: Record; projectDraftThreadIdByProjectId: Record; + readDraftThread: (threadId: ThreadId) => DraftThreadState | null; + isDraftThread: (threadId: ThreadId) => boolean; readProjectDraftThreadId: (projectId: ProjectId) => ThreadId | null; - setProjectDraftThreadId: (projectId: ProjectId, threadId: ThreadId) => void; + setProjectDraftThreadId: ( + projectId: ProjectId, + threadId: ThreadId, + options?: { + branch?: string | null; + worktreePath?: string | null; + createdAt?: string; + }, + ) => void; + setDraftThreadContext: ( + threadId: ThreadId, + options: { + branch?: string | null; + worktreePath?: string | null; + projectId?: ProjectId; + createdAt?: string; + }, + ) => void; clearProjectDraftThreadId: (projectId: ProjectId) => void; clearProjectDraftThreadById: (projectId: ProjectId, threadId: ThreadId) => void; + clearDraftThread: (threadId: ThreadId) => void; setPrompt: (threadId: ThreadId, prompt: string) => void; setCursor: (threadId: ThreadId, cursor: number) => void; setModel: (threadId: ThreadId, model: string | null | undefined) => void; @@ -187,11 +223,39 @@ function normalizePersistedAttachment(value: unknown): PersistedComposerImageAtt function normalizePersistedComposerDraftState(value: unknown): PersistedComposerDraftStoreState { if (!value || typeof value !== "object") { - return { draftsByThreadId: {}, projectDraftThreadIdByProjectId: {} }; + return { draftsByThreadId: {}, draftThreadsByThreadId: {}, projectDraftThreadIdByProjectId: {} }; } const candidate = value as Record; const rawDraftMap = candidate.draftsByThreadId; + const rawDraftThreadsByThreadId = candidate.draftThreadsByThreadId; const rawProjectDraftThreadIdByProjectId = candidate.projectDraftThreadIdByProjectId; + const draftThreadsByThreadId: PersistedComposerDraftStoreState["draftThreadsByThreadId"] = {}; + if (rawDraftThreadsByThreadId && typeof rawDraftThreadsByThreadId === "object") { + for (const [threadId, rawDraftThread] of Object.entries( + rawDraftThreadsByThreadId as Record, + )) { + if (typeof threadId !== "string" || threadId.length === 0) { + continue; + } + if (!rawDraftThread || typeof rawDraftThread !== "object") { + continue; + } + const candidateDraftThread = rawDraftThread as Record; + const projectId = candidateDraftThread.projectId; + const createdAt = candidateDraftThread.createdAt; + const branch = candidateDraftThread.branch; + const worktreePath = candidateDraftThread.worktreePath; + if (typeof projectId !== "string" || projectId.length === 0) { + continue; + } + draftThreadsByThreadId[threadId as ThreadId] = { + projectId: projectId as ProjectId, + createdAt: typeof createdAt === "string" && createdAt.length > 0 ? createdAt : new Date().toISOString(), + branch: typeof branch === "string" ? branch : null, + worktreePath: typeof worktreePath === "string" ? worktreePath : null, + }; + } + } const projectDraftThreadIdByProjectId: PersistedComposerDraftStoreState["projectDraftThreadIdByProjectId"] = {}; if (rawProjectDraftThreadIdByProjectId && typeof rawProjectDraftThreadIdByProjectId === "object") { @@ -205,11 +269,24 @@ function normalizePersistedComposerDraftState(value: unknown): PersistedComposer threadId.length > 0 ) { projectDraftThreadIdByProjectId[projectId as ProjectId] = threadId as ThreadId; + if (!draftThreadsByThreadId[threadId as ThreadId]) { + draftThreadsByThreadId[threadId as ThreadId] = { + projectId: projectId as ProjectId, + createdAt: new Date().toISOString(), + branch: null, + worktreePath: null, + }; + } else if (draftThreadsByThreadId[threadId as ThreadId]?.projectId !== projectId) { + draftThreadsByThreadId[threadId as ThreadId] = { + ...draftThreadsByThreadId[threadId as ThreadId]!, + projectId: projectId as ProjectId, + }; + } } } } if (!rawDraftMap || typeof rawDraftMap !== "object") { - return { draftsByThreadId: {}, projectDraftThreadIdByProjectId }; + return { draftsByThreadId: {}, draftThreadsByThreadId, projectDraftThreadIdByProjectId }; } const nextDraftsByThreadId: PersistedComposerDraftStoreState["draftsByThreadId"] = {}; for (const [threadId, draftValue] of Object.entries(rawDraftMap as Record)) { @@ -250,7 +327,7 @@ function normalizePersistedComposerDraftState(value: unknown): PersistedComposer ...(effort ? { effort } : {}), }; } - return { draftsByThreadId: nextDraftsByThreadId, projectDraftThreadIdByProjectId }; + return { draftsByThreadId: nextDraftsByThreadId, draftThreadsByThreadId, projectDraftThreadIdByProjectId }; } function fileFromDataUrl(dataUrl: string, name: string, mimeType: string): File | null { @@ -304,12 +381,12 @@ function hydrateImagesFromPersisted( function readPersistedDraftStateFromLocalStorage(): PersistedComposerDraftStoreState { if (typeof window === "undefined") { - return { draftsByThreadId: {}, projectDraftThreadIdByProjectId: {} }; + return { draftsByThreadId: {}, draftThreadsByThreadId: {}, projectDraftThreadIdByProjectId: {} }; } try { const raw = window.localStorage.getItem(COMPOSER_DRAFT_STORAGE_KEY); if (!raw) { - return { draftsByThreadId: {}, projectDraftThreadIdByProjectId: {} }; + return { draftsByThreadId: {}, draftThreadsByThreadId: {}, projectDraftThreadIdByProjectId: {} }; } const parsed = JSON.parse(raw) as unknown; if (parsed && typeof parsed === "object" && "state" in parsed) { @@ -317,7 +394,7 @@ function readPersistedDraftStateFromLocalStorage(): PersistedComposerDraftStoreS } return normalizePersistedComposerDraftState(parsed); } catch { - return { draftsByThreadId: {}, projectDraftThreadIdByProjectId: {} }; + return { draftsByThreadId: {}, draftThreadsByThreadId: {}, projectDraftThreadIdByProjectId: {} }; } } @@ -339,26 +416,122 @@ export const useComposerDraftStore = create()( persist( (set, get) => ({ draftsByThreadId: {}, + draftThreadsByThreadId: {}, projectDraftThreadIdByProjectId: {}, + readDraftThread: (threadId) => { + if (threadId.length === 0) { + return null; + } + return get().draftThreadsByThreadId[threadId] ?? null; + }, + isDraftThread: (threadId) => { + if (threadId.length === 0) { + return false; + } + return get().draftThreadsByThreadId[threadId] !== undefined; + }, readProjectDraftThreadId: (projectId) => { if (projectId.length === 0) { return null; } return get().projectDraftThreadIdByProjectId[projectId] ?? null; }, - setProjectDraftThreadId: (projectId, threadId) => { + setProjectDraftThreadId: (projectId, threadId, options) => { if (projectId.length === 0 || threadId.length === 0) { return; } set((state) => { - if (state.projectDraftThreadIdByProjectId[projectId] === threadId) { + const existingThread = state.draftThreadsByThreadId[threadId]; + const previousThreadIdForProject = state.projectDraftThreadIdByProjectId[projectId]; + const nextDraftThread: DraftThreadState = { + projectId, + createdAt: options?.createdAt ?? existingThread?.createdAt ?? new Date().toISOString(), + branch: + options && "branch" in options + ? (options.branch ?? null) + : (existingThread?.branch ?? null), + worktreePath: + options && "worktreePath" in options + ? (options.worktreePath ?? null) + : (existingThread?.worktreePath ?? null), + }; + const hasSameProjectMapping = previousThreadIdForProject === threadId; + const hasSameDraftThread = + existingThread && + existingThread.projectId === nextDraftThread.projectId && + existingThread.createdAt === nextDraftThread.createdAt && + existingThread.branch === nextDraftThread.branch && + existingThread.worktreePath === nextDraftThread.worktreePath; + if (hasSameProjectMapping && hasSameDraftThread) { return state; } + const nextProjectDraftThreadIdByProjectId: Record = { + ...state.projectDraftThreadIdByProjectId, + [projectId]: threadId, + }; + const nextDraftThreadsByThreadId: Record = { + ...state.draftThreadsByThreadId, + [threadId]: nextDraftThread, + }; + if ( + previousThreadIdForProject && + previousThreadIdForProject !== threadId && + !Object.values(nextProjectDraftThreadIdByProjectId).includes(previousThreadIdForProject) + ) { + delete nextDraftThreadsByThreadId[previousThreadIdForProject]; + } + return { + draftThreadsByThreadId: nextDraftThreadsByThreadId, + projectDraftThreadIdByProjectId: nextProjectDraftThreadIdByProjectId, + }; + }); + }, + setDraftThreadContext: (threadId, options) => { + if (threadId.length === 0) { + return; + } + set((state) => { + const existing = state.draftThreadsByThreadId[threadId]; + if (!existing) { + return state; + } + const nextProjectId = options.projectId ?? existing.projectId; + if (nextProjectId.length === 0) { + return state; + } + const nextDraftThread: DraftThreadState = { + projectId: nextProjectId, + createdAt: + options.createdAt === undefined ? existing.createdAt : (options.createdAt || existing.createdAt), + branch: options.branch === undefined ? existing.branch : (options.branch ?? null), + worktreePath: + options.worktreePath === undefined + ? existing.worktreePath + : (options.worktreePath ?? null), + }; + const isUnchanged = + nextDraftThread.projectId === existing.projectId && + nextDraftThread.createdAt === existing.createdAt && + nextDraftThread.branch === existing.branch && + nextDraftThread.worktreePath === existing.worktreePath; + if (isUnchanged) { + return state; + } + const nextProjectDraftThreadIdByProjectId: Record = { + ...state.projectDraftThreadIdByProjectId, + [nextProjectId]: threadId, + }; + if (existing.projectId !== nextProjectId) { + if (nextProjectDraftThreadIdByProjectId[existing.projectId] === threadId) { + delete nextProjectDraftThreadIdByProjectId[existing.projectId]; + } + } return { - projectDraftThreadIdByProjectId: { - ...state.projectDraftThreadIdByProjectId, - [projectId]: threadId, + draftThreadsByThreadId: { + ...state.draftThreadsByThreadId, + [threadId]: nextDraftThread, }, + projectDraftThreadIdByProjectId: nextProjectDraftThreadIdByProjectId, }; }); }, @@ -367,12 +540,21 @@ export const useComposerDraftStore = create()( return; } set((state) => { - if (state.projectDraftThreadIdByProjectId[projectId] === undefined) { + const threadId = state.projectDraftThreadIdByProjectId[projectId]; + if (threadId === undefined) { return state; } - const { [projectId]: _removed, ...rest } = state.projectDraftThreadIdByProjectId; + const { [projectId]: _removed, ...restProjectMappingsRaw } = state.projectDraftThreadIdByProjectId; + const restProjectMappings = restProjectMappingsRaw as Record; + const nextDraftThreadsByThreadId: Record = { + ...state.draftThreadsByThreadId, + }; + if (!Object.values(restProjectMappings).includes(threadId)) { + delete nextDraftThreadsByThreadId[threadId]; + } return { - projectDraftThreadIdByProjectId: rest, + draftThreadsByThreadId: nextDraftThreadsByThreadId, + projectDraftThreadIdByProjectId: restProjectMappings, }; }); }, @@ -384,9 +566,40 @@ export const useComposerDraftStore = create()( if (state.projectDraftThreadIdByProjectId[projectId] !== threadId) { return state; } - const { [projectId]: _removed, ...rest } = state.projectDraftThreadIdByProjectId; + const { [projectId]: _removed, ...restProjectMappingsRaw } = state.projectDraftThreadIdByProjectId; + const restProjectMappings = restProjectMappingsRaw as Record; + const nextDraftThreadsByThreadId: Record = { + ...state.draftThreadsByThreadId, + }; + if (!Object.values(restProjectMappings).includes(threadId)) { + delete nextDraftThreadsByThreadId[threadId]; + } return { - projectDraftThreadIdByProjectId: rest, + draftThreadsByThreadId: nextDraftThreadsByThreadId, + projectDraftThreadIdByProjectId: restProjectMappings, + }; + }); + }, + clearDraftThread: (threadId) => { + if (threadId.length === 0) { + return; + } + set((state) => { + const hasDraftThread = state.draftThreadsByThreadId[threadId] !== undefined; + const hasProjectMapping = Object.values(state.projectDraftThreadIdByProjectId).includes(threadId); + if (!hasDraftThread && !hasProjectMapping) { + return state; + } + const nextProjectDraftThreadIdByProjectId = Object.fromEntries( + Object.entries(state.projectDraftThreadIdByProjectId).filter( + ([, draftThreadId]) => draftThreadId !== threadId, + ), + ) as Record; + const { [threadId]: _removedDraftThread, ...restDraftThreadsByThreadId } = + state.draftThreadsByThreadId; + return { + draftThreadsByThreadId: restDraftThreadsByThreadId, + projectDraftThreadIdByProjectId: nextProjectDraftThreadIdByProjectId, }; }); }, @@ -688,11 +901,26 @@ export const useComposerDraftStore = create()( } } set((state) => { - if (state.draftsByThreadId[threadId] === undefined) { + const hasComposerDraft = state.draftsByThreadId[threadId] !== undefined; + const hasDraftThread = state.draftThreadsByThreadId[threadId] !== undefined; + const hasProjectMapping = Object.values(state.projectDraftThreadIdByProjectId).includes(threadId); + if (!hasComposerDraft && !hasDraftThread && !hasProjectMapping) { return state; } - const { [threadId]: _removed, ...rest } = state.draftsByThreadId; - return { draftsByThreadId: rest }; + const { [threadId]: _removedComposerDraft, ...restComposerDraftsByThreadId } = + state.draftsByThreadId; + const { [threadId]: _removedDraftThread, ...restDraftThreadsByThreadId } = + state.draftThreadsByThreadId; + const nextProjectDraftThreadIdByProjectId = Object.fromEntries( + Object.entries(state.projectDraftThreadIdByProjectId).filter( + ([, draftThreadId]) => draftThreadId !== threadId, + ), + ) as Record; + return { + draftsByThreadId: restComposerDraftsByThreadId, + draftThreadsByThreadId: restDraftThreadsByThreadId, + projectDraftThreadIdByProjectId: nextProjectDraftThreadIdByProjectId, + }; }); }, }), @@ -729,6 +957,7 @@ export const useComposerDraftStore = create()( } return { draftsByThreadId: persistedDraftsByThreadId, + draftThreadsByThreadId: state.draftThreadsByThreadId, projectDraftThreadIdByProjectId: state.projectDraftThreadIdByProjectId, }; }, @@ -743,6 +972,7 @@ export const useComposerDraftStore = create()( return { ...currentState, draftsByThreadId, + draftThreadsByThreadId: normalizedPersisted.draftThreadsByThreadId, projectDraftThreadIdByProjectId: normalizedPersisted.projectDraftThreadIdByProjectId, }; }, diff --git a/apps/web/src/routes/_chat.$threadId.tsx b/apps/web/src/routes/_chat.$threadId.tsx index 41472f5ecc321f1372aa204b3a0dd219dfc6b90d..03c5bb080496ffa33be8dff46adb50d68be59787 100644 --- a/apps/web/src/routes/_chat.$threadId.tsx +++ b/apps/web/src/routes/_chat.$threadId.tsx @@ -3,6 +3,7 @@ import { createFileRoute, useNavigate } from "@tanstack/react-router"; import { Suspense, lazy, type ReactNode, useCallback, useEffect } from "react"; import ChatView from "../components/ChatView"; +import { useComposerDraftStore } from "../composerDraftStore"; import { parseDiffRouteSearch, stripDiffSearchParams } from "../diffRouteSearch"; import { useMediaQuery } from "../hooks/useMediaQuery"; import { useStore } from "../store"; @@ -139,6 +140,10 @@ function ChatThreadRouteView() { }); const search = Route.useSearch(); const threadExists = state.threads.some((thread) => thread.id === threadId); + const draftThreadExists = useComposerDraftStore( + (store) => store.draftThreadsByThreadId[threadId] !== undefined, + ); + const routeThreadExists = threadExists || draftThreadExists; const diffOpen = search.diff === "1"; const shouldUseDiffSheet = useMediaQuery(DIFF_INLINE_LAYOUT_MEDIA_QUERY); const closeDiff = useCallback(() => { @@ -166,13 +171,13 @@ function ChatThreadRouteView() { return; } - if (!threadExists) { + if (!routeThreadExists) { void navigate({ to: "/", replace: true }); return; } - }, [navigate, state.threadsHydrated, threadExists, threadId]); + }, [navigate, routeThreadExists, state.threadsHydrated, threadId]); - if (!state.threadsHydrated || !threadExists) { + if (!state.threadsHydrated || !routeThreadExists) { return null; }