From 8c08fde96c042e5d1b38afc2dcae1274247b1d58 Mon Sep 17 00:00:00 2001 From: audichuang Date: Wed, 2 Sep 2026 22:11:50 +0800 Subject: [PATCH] fix(components): keep New Chat attachments across tab switches MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The landing draft was split across two lifetimes. Prompt text lived in `chatLandingSessionStateAtomFamily`, a module-level atom, so it survived the chat route unmounting. Attachments and the reserved draft session id were component `useState`/`useRef`, so visiting another tab destroyed them — and both hooks ran an unmount cleanup that made the loss irreversible even if the state had been kept: the image hook revoked every preview `blob:` URL and the file hook aborted every upload still in flight. The user came back to their own text with the picture gone, and nothing said so. Both attachment lists and the reserved session id now live in module-level atoms (`atoms/chat-landing-draft.ts`), so they outlive the route exactly as the text does. Deliberately not `atomWithStorage`: a `blob:` URL and an `AbortController` do not serialize, and losing a draft attachment when the app restarts is expected. With the state kept, the unmount cleanups had to go — `URL.revokeObjectURL` and `AbortController.abort()` now belong only to removing one attachment or clearing the whole draft (send accepted, draft reset). An upload in flight when the user leaves keeps running and settles into the atom, so returning shows the finished attachment. That costs nothing new for images: `uploadSessionImage` takes no signal, so those requests already outlived the unmount and only their result was discarded. The attachment key is workspace-scoped where the prompt-text key is not, because an `imageId`/`fileId` is addressable only inside the workspace it was uploaded to — carrying one into a session created elsewhere would attach a block pointing at another workspace's object. It scopes on the workspace slug rather than the resolved id: `useResolvedWorkspaceScope()` reports `null` until the workspace resolves, and a key that flipped mid-mount would strand whatever was added first. So text still follows the user across a workspace switch and attachments do not, which is the intended asymmetry. `useChatLandingDraftSession` reads and writes the store directly, which also removes the state/ref pair its synchronous `ensureSessionId` needed. Mobile is unaffected: `mobile-workspace-stack.tsx` keeps the landing mounted beneath the session drawer, so it never had this bug. The hoist also forces one guard to move. `chat-landing.tsx` tracked the applied `resetDraftKey` in a `useRef`, which is per-mount. That was harmless while the state it cleared was per-mount too, but a New chat URL keeps that key in its history entry, so navigating back re-applied the reset and cleared the draft this change preserves. The marker moves into `chatLandingAppliedResetKeyAtomFamily`, scoped exactly like the draft it guards. Closes #242 Model: claude-opus-5 --- .../src/atoms/chat-landing-draft.ts | 80 +++++ .../components/src/components/chat/AGENTS.md | 16 + .../src/components/chat/chat-landing.tsx | 25 +- .../hooks/use-chat-landing-draft-session.ts | 30 +- .../src/hooks/use-chat-landing-file-draft.ts | 66 ++-- .../src/hooks/use-chat-landing-image-draft.ts | 62 ++-- .../chat-landing-draft-persistence.test.tsx | 320 ++++++++++++++++++ 7 files changed, 515 insertions(+), 84 deletions(-) create mode 100644 packages/components/src/atoms/chat-landing-draft.ts create mode 100644 packages/components/tests/chat-landing-draft-persistence.test.tsx diff --git a/packages/components/src/atoms/chat-landing-draft.ts b/packages/components/src/atoms/chat-landing-draft.ts new file mode 100644 index 000000000..d7db962bb --- /dev/null +++ b/packages/components/src/atoms/chat-landing-draft.ts @@ -0,0 +1,80 @@ +import { atom } from 'jotai'; +import { atomFamily } from 'jotai/utils'; +import type { SessionFilePayload, SessionId, SessionImagePayload } from '@lody/shared'; +import type { SessionFileTransferPhase } from '@/lib/session-file-upload'; + +/** + * In-memory chat-landing draft state that must outlive the landing route's + * unmount. The prompt text already survives through + * `chatLandingSessionStateAtomFamily`; attachments and the reserved draft + * session id used to live in component state, so navigating to another tab and + * back silently dropped them (#242). + * + * Deliberately NOT `atomWithStorage`: a preview `blob:` URL and an in-flight + * upload's `AbortController` cannot be serialized. Surviving a route unmount is + * the whole requirement — losing a draft attachment when the app restarts is + * expected. + */ + +export type PendingImage = { + localId: string; + previewUrl: string; + file: File; + status: 'uploading' | 'uploaded' | 'failed'; + progress: number; + error?: string; + uploaded?: SessionImagePayload; +}; + +export type PendingFile = { + localId: string; + file: File; + status: SessionFileTransferPhase | 'uploaded' | 'failed'; + progress: number; + error?: string; + uploaded?: SessionFilePayload; + abort?: AbortController; +}; + +/** + * The one home for the landing draft's attachment scope. Attachments are + * uploaded into a specific workspace — an `imageId`/`fileId` from one workspace + * cannot be attached to a session in another — so unlike the prompt text (keyed + * by user alone) the attachment draft is workspace-scoped. + * + * Keyed on the workspace SLUG rather than the resolved id: the slug is a route + * param that is stable for the whole mount, while `useResolvedWorkspaceScope()` + * reports `null` until the workspace resolves. A key that flips mid-mount would + * strand whatever was added before it settled. + * + * `stateKey` is the prompt text's own key — the one passed to + * `chatLandingSessionStateAtomFamily`, which an alternate landing surface may + * suffix so it does not clobber the main draft. Pass that, not a raw user id, or + * a suffixed surface would share this scope with the main landing. + */ +export const buildChatLandingDraftKey = (stateKey: string | null, workspaceSlug: string): string => + `${stateKey ?? 'anonymous'}:${workspaceSlug}`; + +export const chatLandingPendingImagesAtomFamily = atomFamily((_draftKey: string) => + atom([]) +); + +export const chatLandingPendingFilesAtomFamily = atomFamily((_draftKey: string) => + atom([]) +); + +export const chatLandingDraftSessionIdAtomFamily = atomFamily((_draftKey: string) => + atom(null) +); + +/** + * The `resetDraftKey` this scope has already been cleared for. It lives beside + * the draft rather than in a mount-scoped ref because the draft now outlives the + * route: a `New chat` URL keeps its `resetDraftKey` in the history entry, so + * navigating back to it would otherwise re-apply the same reset and destroy the + * draft this module exists to preserve — revoking its preview URLs and aborting + * its uploads on the way out. + */ +export const chatLandingAppliedResetKeyAtomFamily = atomFamily((_draftKey: string) => + atom(null) +); diff --git a/packages/components/src/components/chat/AGENTS.md b/packages/components/src/components/chat/AGENTS.md index 38312a9a3..e4e40d179 100644 --- a/packages/components/src/components/chat/AGENTS.md +++ b/packages/components/src/components/chat/AGENTS.md @@ -93,6 +93,22 @@ that same identity. Attachment hooks never reset it independently; reset only after full draft clear. Submit blocks while either `hasBlockingImages` or `hasBlockingFiles`. +- The reserved session id and both attachment lists live in module-level atoms + (`atoms/chat-landing-draft.ts`), keyed by `buildChatLandingDraftKey`, so the + draft survives the landing route unmounting when the user visits another tab + — the prompt text always did, and an attachment that silently vanished was + sent-without-context waiting to happen (#242). Consequences the hooks must + keep: NO unmount cleanup (revoking a preview URL or aborting an upload on the + way out is what broke the returning draft), so `URL.revokeObjectURL` and + `AbortController.abort()` belong only to removing one attachment or clearing + the whole draft; an upload in flight at unmount keeps running and settles into + the atom. That key is workspace-scoped while the prompt-text key is not, + because an `imageId`/`fileId` is addressable only inside the workspace it was + uploaded to. It scopes on the workspace SLUG: `useResolvedWorkspaceScope()` + reports `null` until the workspace resolves, and a key that flipped mid-mount + would strand whatever was added first. Never persist these to localStorage — + a `blob:` URL and an `AbortController` do not serialize, and losing a draft + attachment on app restart is expected. - Submit immediately hides and disables the visible landing draft, but preserves its controlled text, attachment resources, and reserved session id until `startSession` accepts. Failure must reveal the unchanged draft; only acceptance diff --git a/packages/components/src/components/chat/chat-landing.tsx b/packages/components/src/components/chat/chat-landing.tsx index d372a0b52..889ea4f51 100644 --- a/packages/components/src/components/chat/chat-landing.tsx +++ b/packages/components/src/components/chat/chat-landing.tsx @@ -11,7 +11,7 @@ import { type ReactNode, } from 'react'; import { useTranslation } from 'react-i18next'; -import { useAtom, useAtomValue, useSetAtom } from 'jotai'; +import { useAtom, useAtomValue, useSetAtom, useStore } from 'jotai'; import { buildPendingUserHistoryEntry, buildSessionPreparationRunConfig, @@ -201,6 +201,10 @@ import { arePersistedMentionRangesEqual, toPersistedMentionRanges, } from '@/components/mentions/mention-persistence'; +import { + buildChatLandingDraftKey, + chatLandingAppliedResetKeyAtomFamily, +} from '@/atoms/chat-landing-draft'; import { useChatLandingImageDraft } from '@/hooks/use-chat-landing-image-draft'; import { useChatLandingFileDraft } from '@/hooks/use-chat-landing-file-draft'; import { useChatLandingDraftSession } from '@/hooks/use-chat-landing-draft-session'; @@ -955,6 +959,12 @@ function WorkspaceChatLanding({ const [sessionState, setSessionState] = useAtom( chatLandingSessionStateAtomFamily(chatLandingStateKey) ); + /** + * Scope for the attachment draft and the reserved session id. Unlike the + * prompt text this is workspace-scoped, because an uploaded image/file is + * addressable only inside the workspace it was uploaded to. + */ + const chatLandingDraftKey = buildChatLandingDraftKey(chatLandingStateKey, workspaceSlug); const prompt = sessionState.prompt; const [draftActivityRevision, setDraftActivityRevision] = useState(0); const pastedTextDrafts = useMemo( @@ -1279,7 +1289,7 @@ function WorkspaceChatLanding({ sessionId: draftSessionId, ensureSessionId: ensureDraftSessionId, resetSessionId: resetDraftSessionId, - } = useChatLandingDraftSession(); + } = useChatLandingDraftSession(chatLandingDraftKey); const attachmentInputRef = useRef(null); const { imageItems, @@ -1293,6 +1303,7 @@ function WorkspaceChatLanding({ clearPendingImages, buildInputBlocks, } = useChatLandingImageDraft({ + draftKey: chatLandingDraftKey, workspaceId: (workspaceId as WorkspaceId | null) ?? null, authToken, isMobile, @@ -1311,13 +1322,15 @@ function WorkspaceChatLanding({ clearPendingFiles, buildFileInputBlocks, } = useChatLandingFileDraft({ + draftKey: chatLandingDraftKey, workspaceId: (workspaceId as WorkspaceId | null) ?? null, authToken, machineId: selectedMachineId, sessionId: draftSessionId, ensureSessionId: ensureDraftSessionId, }); - const lastAppliedResetDraftKeyRef = useRef(null); + const draftStore = useStore(); + const appliedResetKeyAtom = chatLandingAppliedResetKeyAtomFamily(chatLandingDraftKey); useEffect(() => { if (!resetDraftKey) { @@ -1325,10 +1338,10 @@ function WorkspaceChatLanding({ } const scopedResetKey = `${chatLandingStateKey ?? 'anonymous'}:${resetDraftKey}`; - if (lastAppliedResetDraftKeyRef.current === scopedResetKey) { + if (draftStore.get(appliedResetKeyAtom) === scopedResetKey) { return; } - lastAppliedResetDraftKeyRef.current = scopedResetKey; + draftStore.set(appliedResetKeyAtom, scopedResetKey); if (resetDraftOnKeyChange) { setSessionState({ prompt: '', pastedTextDrafts: [] }); } @@ -1337,9 +1350,11 @@ function WorkspaceChatLanding({ clearPendingFiles(); resetDraftSessionId(); }, [ + appliedResetKeyAtom, chatLandingStateKey, clearPendingFiles, clearPendingImages, + draftStore, resetDraftSessionId, resetDraftKey, resetDraftOnKeyChange, diff --git a/packages/components/src/hooks/use-chat-landing-draft-session.ts b/packages/components/src/hooks/use-chat-landing-draft-session.ts index 748098fa7..a4aef197c 100644 --- a/packages/components/src/hooks/use-chat-landing-draft-session.ts +++ b/packages/components/src/hooks/use-chat-landing-draft-session.ts @@ -1,5 +1,7 @@ -import { useCallback, useRef, useState } from 'react'; +import { useCallback } from 'react'; +import { useAtomValue, useStore } from 'jotai'; import type { SessionId } from '@lody/shared'; +import { chatLandingDraftSessionIdAtomFamily } from '@/atoms/chat-landing-draft'; function createDraftSessionId(): SessionId { if (typeof crypto !== 'undefined' && typeof crypto.randomUUID === 'function') { @@ -8,22 +10,28 @@ function createDraftSessionId(): SessionId { return `${Date.now()}-${Math.random().toString(36).slice(2)}` as SessionId; } -export function useChatLandingDraftSession() { - const [sessionId, setSessionId] = useState(null); - const sessionIdRef = useRef(null); +/** + * The landing's reserved session id. Held in a module-level atom so the + * attachments uploaded against it stay addressable after the landing route + * unmounts and remounts; the store read also removes the state/ref pair the + * synchronous `ensureSessionId` used to need. + */ +export function useChatLandingDraftSession(draftKey: string) { + const sessionIdAtom = chatLandingDraftSessionIdAtomFamily(draftKey); + const store = useStore(); + const sessionId = useAtomValue(sessionIdAtom); const ensureSessionId = useCallback((): SessionId => { - if (sessionIdRef.current) return sessionIdRef.current; + const current = store.get(sessionIdAtom); + if (current) return current; const next = createDraftSessionId(); - sessionIdRef.current = next; - setSessionId(next); + store.set(sessionIdAtom, next); return next; - }, []); + }, [sessionIdAtom, store]); const resetSessionId = useCallback(() => { - sessionIdRef.current = null; - setSessionId(null); - }, []); + store.set(sessionIdAtom, null); + }, [sessionIdAtom, store]); return { sessionId, ensureSessionId, resetSessionId }; } diff --git a/packages/components/src/hooks/use-chat-landing-file-draft.ts b/packages/components/src/hooks/use-chat-landing-file-draft.ts index bc6a26ab3..e55ae051f 100644 --- a/packages/components/src/hooks/use-chat-landing-file-draft.ts +++ b/packages/components/src/hooks/use-chat-landing-file-draft.ts @@ -1,4 +1,4 @@ -import { useCallback, useEffect, useMemo, useState } from 'react'; +import { useCallback, useMemo } from 'react'; import { SESSION_FILE_MAX_COUNT, type MachineId, @@ -7,9 +7,10 @@ import { type SessionInputBlock, type WorkspaceId, } from '@lody/shared'; -import { useAtomValue } from 'jotai'; +import { useAtom, useAtomValue } from 'jotai'; import { toast } from 'sonner'; import { useTranslation } from 'react-i18next'; +import { chatLandingPendingFilesAtomFamily, type PendingFile } from '@/atoms/chat-landing-draft'; import { localMachineIdAtom } from '@/atoms/local-probe'; import { formatFileSize } from '@/lib/session-file-presentation'; import { @@ -28,16 +29,6 @@ import { type SessionFileUploadProgress, } from '@/lib/session-file-upload'; -type PendingFile = { - localId: string; - file: File; - status: SessionFileTransferPhase | 'uploaded' | 'failed'; - progress: number; - error?: string; - uploaded?: SessionFilePayload; - abort?: AbortController; -}; - export type ChatLandingFileDraftItem = { id: string; name: string; @@ -81,6 +72,8 @@ const createLocalFileId = (): string => { * support for removal mid-flight. */ export function useChatLandingFileDraft(args: { + /** Scope shared with the sibling image draft and the reserved session id. */ + draftKey: string; workspaceId: WorkspaceId | null; authToken: string | null; /** Selected machine for the eventual session; enables the local fast path. */ @@ -90,9 +83,16 @@ export function useChatLandingFileDraft(args: { ensureSessionId: () => SessionId; }) { const { t } = useTranslation(); - const { workspaceId, authToken, machineId, sessionId: draftSessionId, ensureSessionId } = args; + const { + draftKey, + workspaceId, + authToken, + machineId, + sessionId: draftSessionId, + ensureSessionId, + } = args; const localMachineId = useAtomValue(localMachineIdAtom); - const [pendingFiles, setPendingFiles] = useState([]); + const [pendingFiles, setPendingFiles] = useAtom(chatLandingPendingFilesAtomFamily(draftKey)); // Desktop local-transport fast path: available only when the selected machine // is this machine's local CLI and the Electron preload bridge exposes the @@ -116,18 +116,13 @@ export function useChatLandingFileDraft(args: { } return []; }); - }, []); + }, [setPendingFiles]); - useEffect(() => { - return () => { - setPendingFiles((prev) => { - for (const entry of prev) { - entry.abort?.abort(); - } - return []; - }); - }; - }, []); + // No unmount cleanup: the draft outlives the landing route (#242). An upload + // still in flight when the user switches tabs keeps running and settles into + // the atom, so returning shows the finished attachment rather than an entry + // aborted on the way out. Aborting stays tied to the user's own actions — + // removing one file, or clearing the draft (send accepted / draft reset). const updatePendingFile = useCallback( (localId: string, updater: (file: PendingFile) => PendingFile) => { @@ -135,7 +130,7 @@ export function useChatLandingFileDraft(args: { prev.map((entry) => (entry.localId === localId ? updater(entry) : entry)) ); }, - [] + [setPendingFiles] ); const startUpload = useCallback( @@ -310,16 +305,19 @@ export function useChatLandingFileDraft(args: { void startUpload(entry.localId, entry.file, sessionId); } }, - [ensureSessionId, pendingFiles.length, startUpload, t] + [ensureSessionId, pendingFiles.length, setPendingFiles, startUpload, t] ); - const handleRemoveFile = useCallback((localId: string) => { - setPendingFiles((prev) => { - const target = prev.find((entry) => entry.localId === localId); - target?.abort?.abort(); - return prev.filter((entry) => entry.localId !== localId); - }); - }, []); + const handleRemoveFile = useCallback( + (localId: string) => { + setPendingFiles((prev) => { + const target = prev.find((entry) => entry.localId === localId); + target?.abort?.abort(); + return prev.filter((entry) => entry.localId !== localId); + }); + }, + [setPendingFiles] + ); const handleRetryFile = useCallback( (localId: string) => { diff --git a/packages/components/src/hooks/use-chat-landing-image-draft.ts b/packages/components/src/hooks/use-chat-landing-image-draft.ts index 27ac25dd8..8976d0c56 100644 --- a/packages/components/src/hooks/use-chat-landing-image-draft.ts +++ b/packages/components/src/hooks/use-chat-landing-image-draft.ts @@ -1,4 +1,4 @@ -import { useCallback, useEffect, useMemo, useState, type ClipboardEvent } from 'react'; +import { useCallback, useMemo, type ClipboardEvent } from 'react'; import type { MessageTextSpan } from '@lody/shared'; import { SESSION_IMAGE_MAX_COUNT, @@ -7,20 +7,13 @@ import { type SessionInputBlock, type WorkspaceId, } from '@lody/shared'; +import { useAtom } from 'jotai'; import { toast } from 'sonner'; import { useTranslation } from 'react-i18next'; import { usePostHog } from '@posthog/react'; +import { chatLandingPendingImagesAtomFamily, type PendingImage } from '@/atoms/chat-landing-draft'; import { capturePostHogEvent } from '@/lib/posthog-analytics'; import { uploadSessionImage, validateSessionImageFile } from '@/lib/session-image-upload'; -type PendingImage = { - localId: string; - previewUrl: string; - file: File; - status: 'uploading' | 'uploaded' | 'failed'; - progress: number; - error?: string; - uploaded?: SessionImagePayload; -}; export type ChatLandingImageDraftItem = { id: string; @@ -49,6 +42,8 @@ const createLocalImageId = (): string => { }; export function useChatLandingImageDraft(args: { + /** Scope shared with the sibling file draft and the reserved session id. */ + draftKey: string; workspaceId: WorkspaceId | null; authToken: string | null; isMobile: boolean; @@ -58,6 +53,7 @@ export function useChatLandingImageDraft(args: { }) { const { t } = useTranslation(); const { + draftKey, workspaceId, authToken, isMobile, @@ -66,7 +62,7 @@ export function useChatLandingImageDraft(args: { ensureSessionId, } = args; const postHog = usePostHog(); - const [pendingImages, setPendingImages] = useState([]); + const [pendingImages, setPendingImages] = useAtom(chatLandingPendingImagesAtomFamily(draftKey)); const imageUploadFailedLabel = t('sessions.imageUploadFailed', 'Image upload failed'); const imageUploadMissingAuthLabel = t( 'sessions.imageUploadMissingAuth', @@ -111,18 +107,12 @@ export function useChatLandingImageDraft(args: { } return []; }); - }, []); + }, [setPendingImages]); - useEffect(() => { - return () => { - setPendingImages((prev) => { - for (const image of prev) { - URL.revokeObjectURL(image.previewUrl); - } - return []; - }); - }; - }, []); + // No unmount cleanup: the draft outlives the landing route (#242), so a + // preview URL is revoked only when its image is removed or the whole draft + // is cleared (send accepted / draft reset), never because the user navigated + // to another tab. const updatePendingImage = useCallback( (localId: string, updater: (image: PendingImage) => PendingImage) => { @@ -130,7 +120,7 @@ export function useChatLandingImageDraft(args: { prev.map((image) => (image.localId === localId ? updater(image) : image)) ); }, - [] + [setPendingImages] ); const startUpload = useCallback( @@ -284,6 +274,7 @@ export function useChatLandingImageDraft(args: { ensureSessionId, imageCountLimitLabel, pendingImages.length, + setPendingImages, showImageSelectionIssues, startUpload, ] @@ -309,17 +300,20 @@ export function useChatLandingImageDraft(args: { [handleAddFiles, isMobile] ); - const handleRemoveImage = useCallback((localId: string) => { - // The landing owns the shared draft session id, so removing the last image - // cannot orphan file attachments or an in-flight ACP preparation. - setPendingImages((prev) => { - const target = prev.find((item) => item.localId === localId); - if (target) { - URL.revokeObjectURL(target.previewUrl); - } - return prev.filter((item) => item.localId !== localId); - }); - }, []); + const handleRemoveImage = useCallback( + (localId: string) => { + // The landing owns the shared draft session id, so removing the last image + // cannot orphan file attachments or an in-flight ACP preparation. + setPendingImages((prev) => { + const target = prev.find((item) => item.localId === localId); + if (target) { + URL.revokeObjectURL(target.previewUrl); + } + return prev.filter((item) => item.localId !== localId); + }); + }, + [setPendingImages] + ); const handleRetryImage = useCallback( (localId: string) => { diff --git a/packages/components/tests/chat-landing-draft-persistence.test.tsx b/packages/components/tests/chat-landing-draft-persistence.test.tsx new file mode 100644 index 000000000..a2303c5ad --- /dev/null +++ b/packages/components/tests/chat-landing-draft-persistence.test.tsx @@ -0,0 +1,320 @@ +// @vitest-environment jsdom + +import { act, createElement } from 'react'; +import { createRoot, type Root } from 'react-dom/client'; +import { Provider, createStore } from 'jotai'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import type { SessionId, WorkspaceId } from '@lody/shared'; + +import { useChatLandingDraftSession } from '../src/hooks/use-chat-landing-draft-session'; +import { + useChatLandingImageDraft, + type ChatLandingImageDraftItem, +} from '../src/hooks/use-chat-landing-image-draft'; +import { + useChatLandingFileDraft, + type ChatLandingFileDraftItem, +} from '../src/hooks/use-chat-landing-file-draft'; + +type Deferred = { promise: Promise; resolve: (value: T) => void }; + +function deferred(): Deferred { + let resolve!: (value: T) => void; + const promise = new Promise((done) => { + resolve = done; + }); + return { promise, resolve }; +} + +const uploadMocks = vi.hoisted(() => ({ + imageUpload: null as Deferred | null, + fileUpload: null as Deferred | null, + /** Resolves the moment the hook reaches `uploadSessionFile`, so the test + * waits on that call rather than on a guessed number of microtasks. */ + fileUploadStarted: null as Deferred | null, + fileUploadSignals: [] as (AbortSignal | undefined)[], +})); + +vi.mock('react-i18next', () => ({ + useTranslation: () => ({ + t: (_key: string, fallback?: string) => fallback ?? _key, + i18n: { language: 'en' }, + }), +})); + +vi.mock('sonner', () => ({ toast: { error: vi.fn() } })); + +vi.mock('@posthog/react', () => ({ usePostHog: () => null })); + +vi.mock('../src/lib/posthog-analytics', () => ({ capturePostHogEvent: vi.fn() })); + +vi.mock('../src/lib/session-image-upload', () => ({ + validateSessionImageFile: () => null, + uploadSessionImage: () => { + uploadMocks.imageUpload = deferred(); + return uploadMocks.imageUpload.promise; + }, +})); + +vi.mock('../src/lib/session-file-upload', () => ({ + SESSION_FILE_MAX_SIZE_MB: 20, + validateSessionFile: () => null, + computeSha256Hex: async () => 'sha256', + computeTextPreviewable: async () => undefined, + isUploadAbortedError: () => false, + isSessionFileTransferPhase: (status: string) => status === 'preparing' || status === 'uploading', + uploadSessionFile: ({ signal }: { signal?: AbortSignal }) => { + uploadMocks.fileUploadSignals.push(signal); + uploadMocks.fileUpload = deferred(); + uploadMocks.fileUploadStarted?.resolve(); + return uploadMocks.fileUpload.promise; + }, +})); + +vi.mock('../src/lib/electron-session-file-sender', () => ({ + canUseElectronLocalFileSend: () => false, + sendSessionFileToLocalRuntime: async () => null, +})); + +const WORKSPACE_A_KEY = 'user-1:workspace-a'; +const WORKSPACE_B_KEY = 'user-1:workspace-b'; + +type Harness = { + imageItems: ChatLandingImageDraftItem[]; + fileItems: ChatLandingFileDraftItem[]; + sessionId: SessionId | null; + addImages: (files: File[]) => void; + addFiles: (files: File[]) => void; + removeImage: (localId: string) => void; + clearDraft: () => void; +}; + +let harness: Harness | null = null; + +function DraftHarness({ draftKey }: { draftKey: string }) { + const { sessionId, ensureSessionId } = useChatLandingDraftSession(draftKey); + const imageDraft = useChatLandingImageDraft({ + draftKey, + workspaceId: 'workspace-a' as WorkspaceId, + authToken: 'token', + isMobile: false, + projectKind: null, + sessionId, + ensureSessionId, + }); + const fileDraft = useChatLandingFileDraft({ + draftKey, + workspaceId: 'workspace-a' as WorkspaceId, + authToken: 'token', + machineId: null, + sessionId, + ensureSessionId, + }); + harness = { + imageItems: imageDraft.imageItems, + fileItems: fileDraft.fileItems, + sessionId, + addImages: imageDraft.addFiles, + addFiles: fileDraft.addFiles, + removeImage: imageDraft.handleRemoveImage, + // What submit-accepted and `resetDraftKey` call in `chat-landing.tsx`. + clearDraft: () => { + imageDraft.clearPendingImages(); + fileDraft.clearPendingFiles(); + }, + }; + return null; +} + +let store = createStore(); +let root: Root | null = null; +let container: HTMLDivElement | null = null; +let objectUrlSeq = 0; +let revokedUrls: string[] = []; + +function mountLanding(draftKey: string): void { + container = document.createElement('div'); + document.body.appendChild(container); + root = createRoot(container); + act(() => { + root?.render(createElement(Provider, { store }, createElement(DraftHarness, { draftKey }))); + }); +} + +function unmountLanding(): void { + act(() => { + root?.unmount(); + }); + container?.remove(); + root = null; + container = null; +} + +function readHarness(): Harness { + if (!harness) throw new Error('landing harness is not mounted'); + return harness; +} + +function pngFile(name: string): File { + return new File([new Uint8Array([1, 2, 3])], name, { type: 'image/png' }); +} + +function textFile(name: string): File { + return new File([new Uint8Array([4, 5, 6])], name, { type: 'text/plain' }); +} + +beforeEach(() => { + store = createStore(); + harness = null; + objectUrlSeq = 0; + revokedUrls = []; + uploadMocks.imageUpload = null; + uploadMocks.fileUpload = null; + uploadMocks.fileUploadStarted = deferred(); + uploadMocks.fileUploadSignals = []; + URL.createObjectURL = () => `blob:preview/${(objectUrlSeq += 1)}`; + URL.revokeObjectURL = (url: string) => { + revokedUrls.push(url); + }; +}); + +afterEach(() => { + if (root) unmountLanding(); +}); + +describe('chat landing draft attachments across a route unmount', () => { + it('restores images, their preview URLs, and the reserved session id', () => { + mountLanding(WORKSPACE_A_KEY); + act(() => { + readHarness().addImages([pngFile('shot.png')]); + }); + + const added = readHarness().imageItems; + expect(added).toHaveLength(1); + const previewUrl = added[0]!.previewUrl; + const reservedSessionId = readHarness().sessionId; + expect(reservedSessionId).not.toBeNull(); + + unmountLanding(); + expect(revokedUrls).toEqual([]); + + mountLanding(WORKSPACE_A_KEY); + const restored = readHarness().imageItems; + expect(restored).toHaveLength(1); + expect(restored[0]!.id).toBe(added[0]!.id); + expect(restored[0]!.previewUrl).toBe(previewUrl); + expect(readHarness().sessionId).toBe(reservedSessionId); + }); + + it('revokes a preview URL when the user removes that image', () => { + mountLanding(WORKSPACE_A_KEY); + act(() => { + readHarness().addImages([pngFile('shot.png')]); + }); + const [item] = readHarness().imageItems; + + act(() => { + readHarness().removeImage(item!.id); + }); + + expect(revokedUrls).toEqual([item!.previewUrl]); + expect(readHarness().imageItems).toEqual([]); + }); + + it('lets an image upload that was in flight at unmount finish into the restored draft', async () => { + mountLanding(WORKSPACE_A_KEY); + act(() => { + readHarness().addImages([pngFile('shot.png')]); + }); + expect(readHarness().imageItems[0]!.status).toBe('uploading'); + + unmountLanding(); + await act(async () => { + uploadMocks.imageUpload?.resolve({ + imageId: 'image-1', + mimeType: 'image/png', + fileName: 'shot.png', + sizeBytes: 3, + }); + await Promise.resolve(); + }); + + mountLanding(WORKSPACE_A_KEY); + expect(readHarness().imageItems[0]!.status).toBe('uploaded'); + }); + + it('does not abort a file upload when the landing unmounts', async () => { + mountLanding(WORKSPACE_A_KEY); + await act(async () => { + readHarness().addFiles([textFile('notes.txt')]); + await uploadMocks.fileUploadStarted?.promise; + }); + expect(uploadMocks.fileUploadSignals).toHaveLength(1); + + unmountLanding(); + expect(uploadMocks.fileUploadSignals[0]?.aborted).toBe(false); + + await act(async () => { + uploadMocks.fileUpload?.resolve({ + fileId: 'file-1', + fileName: 'notes.txt', + mimeType: 'text/plain', + sizeBytes: 3, + sha256: 'sha256', + transport: 'cloud', + uploadedAt: 0, + }); + await Promise.resolve(); + }); + + mountLanding(WORKSPACE_A_KEY); + const restored = readHarness().fileItems; + expect(restored).toHaveLength(1); + expect(restored[0]!.status).toBe('uploaded'); + }); + + it('clears the draft for good once submit or a draft reset releases it', async () => { + mountLanding(WORKSPACE_A_KEY); + act(() => { + readHarness().addImages([pngFile('shot.png')]); + }); + await act(async () => { + readHarness().addFiles([textFile('notes.txt')]); + await uploadMocks.fileUploadStarted?.promise; + }); + const [image] = readHarness().imageItems; + + act(() => { + readHarness().clearDraft(); + }); + + expect(revokedUrls).toEqual([image!.previewUrl]); + expect(uploadMocks.fileUploadSignals[0]?.aborted).toBe(true); + expect(readHarness().imageItems).toEqual([]); + expect(readHarness().fileItems).toEqual([]); + + unmountLanding(); + mountLanding(WORKSPACE_A_KEY); + expect(readHarness().imageItems).toEqual([]); + expect(readHarness().fileItems).toEqual([]); + }); + + it('keeps drafts in different workspaces apart', () => { + mountLanding(WORKSPACE_A_KEY); + act(() => { + readHarness().addImages([pngFile('from-a.png')]); + }); + const workspaceAItems = readHarness().imageItems; + const workspaceASessionId = readHarness().sessionId; + unmountLanding(); + + mountLanding(WORKSPACE_B_KEY); + expect(readHarness().imageItems).toEqual([]); + expect(readHarness().sessionId).toBeNull(); + unmountLanding(); + + mountLanding(WORKSPACE_A_KEY); + expect(readHarness().imageItems).toEqual(workspaceAItems); + expect(readHarness().sessionId).toBe(workspaceASessionId); + }); +});