From a1686025b6fcb58c7c795be11a44da3e791e68fe Mon Sep 17 00:00:00 2001 From: Joshen Lim Date: Wed, 9 Sep 2026 17:17:25 +0800 Subject: [PATCH] Joshenlim/fe 4291 keep unsaved notebooks accessible after page refresh (#49673) ## Context Changes here adds a "Draft" state for notebooks with a new `notebook-drafts` store - similar to how we handle query tabs in the explorer. This implies that if a user refreshes the tab while there's unsaved changes to notebooks, the changes can be persisted locally and the user will be able to continue from where they left off. This also implies that If you create a new notebook (OR open an existing notebook and make some changes) and refresh the browser, we no longer show the native browser confirmation dialog about discarding changes. We also reuse the existing confirmation dialog when saving a notebook if its draft has diverged from the server side content - just updated the language to be more generic rather than saying that the Assistant made changes image Also fixes an unrelated issue - renaming a notebook should mark the notebook as having unsaved changes (with the orange indicator) ## Summary by CodeRabbit ## Summary by CodeRabbit * **New Features** * Unsaved notebook edits are saved locally and restored when reopening Studio. * Drafts are scoped by project and protected from server changes through conflict detection. * Notebook tabs indicate unsaved changes, including drafts from unsaved notebooks. * **Bug Fixes** * Closing a tab with local edits prompts for confirmation and removes its saved draft. * Notebook save state reflects the server-confirmed update time. * Conflict messages clearly describe changes made on the server. * **Style** * Improved keyboard focus behavior for tab controls. --- .../Explorer/ExplorerNotebookTab.tsx | 10 +- .../ExplorerNotebookTabCoordinator.tsx | 49 ++--- ...kTab.assistant-cache-invalidation.test.tsx | 12 +- .../ExplorerNotebookTabCoordinator.test.tsx | 87 +++++++++ .../components/interfaces/Explorer/hooks.ts | 45 ++++- .../components/layouts/Tabs/SortableTab.tsx | 4 +- .../data/content/notebooks/notebook-cache.ts | 2 + .../state/notebooks/notebook-drafts.test.ts | 177 ++++++++++++++++++ .../studio/state/notebooks/notebook-drafts.ts | 152 +++++++++++++++ .../state/notebooks/notebooks-state.test.ts | 122 ++++++++++++ .../studio/state/notebooks/notebooks-state.ts | 85 ++++++++- apps/studio/state/notebooks/types.ts | 1 + packages/common/constants/local-storage.ts | 1 + 13 files changed, 702 insertions(+), 45 deletions(-) create mode 100644 apps/studio/state/notebooks/notebook-drafts.test.ts create mode 100644 apps/studio/state/notebooks/notebook-drafts.ts diff --git a/apps/studio/components/interfaces/Explorer/ExplorerNotebookTab.tsx b/apps/studio/components/interfaces/Explorer/ExplorerNotebookTab.tsx index 2d97f1576cd..b0fc1af2124 100644 --- a/apps/studio/components/interfaces/Explorer/ExplorerNotebookTab.tsx +++ b/apps/studio/components/interfaces/Explorer/ExplorerNotebookTab.tsx @@ -118,9 +118,9 @@ export const ExplorerNotebookTab = () => { const scrollContainerRef = useRef(null) const { mutate: updateNotebook, isPending: isUpdating } = useUpsertNotebookMutation({ - onSuccess: () => { + onSuccess: (data) => { if (id && content === savedContentRef.current) { - snap.markSaved({ id }) + snap.markSaved({ id, updatedAt: data?.updated_at }) toast.success('Successfully saved notebook!') if (isSaveBeforeAnalyzeOpen) { setIsSaveBeforeAnalyzeOpen(false) @@ -543,7 +543,7 @@ export const ExplorerNotebookTab = () => { { >

{id && snap.serverDivergedWhileDirty.get(id) === 'deleted' - ? 'An assistant deleted this notebook after your local changes. Saving will recreate it.' - : "An assistant updated this notebook after your local changes. Saving will overwrite the assistant's update."} + ? 'This notebook was deleted on the server after your local changes. Saving will recreate it.' + : 'This notebook changed on the server after your local changes. Saving will overwrite those changes.'}

diff --git a/apps/studio/components/interfaces/Explorer/ExplorerNotebookTabCoordinator.tsx b/apps/studio/components/interfaces/Explorer/ExplorerNotebookTabCoordinator.tsx index 21f9af24ffa..482a48f6c20 100644 --- a/apps/studio/components/interfaces/Explorer/ExplorerNotebookTabCoordinator.tsx +++ b/apps/studio/components/interfaces/Explorer/ExplorerNotebookTabCoordinator.tsx @@ -1,19 +1,39 @@ import { useQueryClient } from '@tanstack/react-query' import { useParams } from 'common' import { useContext, useEffect } from 'react' +import type { Snapshot } from 'valtio' import { evictNotebookFromCaches, hasDiscardableChanges, } from '@/data/content/notebooks/notebook-cache' +import { hasNotebookDraft } from '@/state/notebooks/notebook-drafts' import { notebooksState, useNotebooksStateSnapshot } from '@/state/notebooks/notebooks-state' +import type { StateNotebook } from '@/state/notebooks/types' import { TabsStateContext, type Tab } from '@/state/tabs' +function isNotebookTabDirty({ + stateNotebook, + ref, + notebookId, +}: { + stateNotebook: StateNotebook | Snapshot | undefined + ref: string | undefined + notebookId: string | undefined +}): boolean { + if (!notebookId) return false + if (stateNotebook) return hasDiscardableChanges(stateNotebook) + + return !!ref && hasNotebookDraft({ projectRef: ref, id: notebookId }) +} + const NotebookTabStatusIndicator = ({ tab }: { tab: Tab }) => { + const { ref } = useParams() const notebooksSnap = useNotebooksStateSnapshot() const notebookId = tab.metadata?.notebookId const stateNotebook = notebookId ? notebooksSnap.notebooks[notebookId] : undefined - if (!hasDiscardableChanges(stateNotebook)) return null + + if (!isNotebookTabDirty({ stateNotebook, ref, notebookId })) return null return ( { ) } -/** - * Evicts a notebook's content from the valtio store and the React Query cache - * when its tab closes, so reopening it always refetches instead of showing - * whatever was last loaded. Unsaved edits are discarded the same way — safe - * because `confirmClose` below has already asked the user to confirm. - */ export const ExplorerNotebookTabCoordinator = () => { const { ref } = useParams() const queryClient = useQueryClient() const tabs = useContext(TabsStateContext) - useEffect(() => { - const handleBeforeUnload = (event: BeforeUnloadEvent) => { - const hasUnsaved = Object.values(notebooksState.notebooks).some(hasDiscardableChanges) - if (hasUnsaved) { - event.preventDefault() - event.returnValue = true - } - } - - window.addEventListener('beforeunload', handleBeforeUnload) - return () => window.removeEventListener('beforeunload', handleBeforeUnload) - }, []) - useEffect(() => { return tabs.registerTabTypeHandler('notebook', { onClose: (tab) => { @@ -59,9 +60,11 @@ export const ExplorerNotebookTabCoordinator = () => { confirmClose: (notebookTabs) => { const dirtyCount = notebookTabs.filter((tab) => { const notebookId = tab.metadata?.notebookId - if (!notebookId) return false - - return hasDiscardableChanges(notebooksState.notebooks[notebookId]) + return isNotebookTabDirty({ + stateNotebook: notebookId ? notebooksState.notebooks[notebookId] : undefined, + ref, + notebookId, + }) }).length if (dirtyCount === 0) return null diff --git a/apps/studio/components/interfaces/Explorer/__tests__/ExplorerNotebookTab.assistant-cache-invalidation.test.tsx b/apps/studio/components/interfaces/Explorer/__tests__/ExplorerNotebookTab.assistant-cache-invalidation.test.tsx index c12a5f123d3..b3d3b31dbe7 100644 --- a/apps/studio/components/interfaces/Explorer/__tests__/ExplorerNotebookTab.assistant-cache-invalidation.test.tsx +++ b/apps/studio/components/interfaces/Explorer/__tests__/ExplorerNotebookTab.assistant-cache-invalidation.test.tsx @@ -231,7 +231,7 @@ describe('ExplorerNotebookTab — assistant cache invalidation', () => { await waitFor(() => expect(mutationCount).toBe(1)) expect( - screen.queryByRole('dialog', { name: 'Assistant changes detected' }) + screen.queryByRole('dialog', { name: 'Notebook changed on the server' }) ).not.toBeInTheDocument() }) @@ -239,13 +239,13 @@ describe('ExplorerNotebookTab — assistant cache invalidation', () => { { type: 'updated' as const, description: - "An assistant updated this notebook after your local changes. Saving will overwrite the assistant's update.", + 'This notebook changed on the server after your local changes. Saving will overwrite those changes.', saveLabel: 'Save anyway', }, { type: 'deleted' as const, description: - 'An assistant deleted this notebook after your local changes. Saving will recreate it.', + 'This notebook was deleted on the server after your local changes. Saving will recreate it.', saveLabel: 'Recreate', }, ])('shows the $type conflict copy before saving', async ({ type, description, saveLabel }) => { @@ -261,7 +261,7 @@ describe('ExplorerNotebookTab — assistant cache invalidation', () => { await userEvent.click(await screen.findByRole('button', { name: 'Save changes' })) - const dialog = await screen.findByRole('dialog', { name: 'Assistant changes detected' }) + const dialog = await screen.findByRole('dialog', { name: 'Notebook changed on the server' }) expect(dialog).toHaveTextContent(description) expect(dialog).toHaveTextContent(saveLabel) expect(dialog).toHaveTextContent('Discard changes') @@ -425,7 +425,7 @@ describe('ExplorerNotebookTab — assistant cache invalidation', () => { renderNotebookTab(new QueryClient()) await userEvent.click(await screen.findByRole('button', { name: 'Save changes' })) - const dialog = await screen.findByRole('dialog', { name: 'Assistant changes detected' }) + const dialog = await screen.findByRole('dialog', { name: 'Notebook changed on the server' }) if (dismissal === 'the close button') { await userEvent.click(screen.getByRole('button', { name: 'Close' })) } else { @@ -434,7 +434,7 @@ describe('ExplorerNotebookTab — assistant cache invalidation', () => { await waitFor(() => expect( - screen.queryByRole('dialog', { name: 'Assistant changes detected' }) + screen.queryByRole('dialog', { name: 'Notebook changed on the server' }) ).not.toBeInTheDocument() ) expect(notebooksState.notebooks[NOTEBOOK_ID]).toBeDefined() diff --git a/apps/studio/components/interfaces/Explorer/__tests__/ExplorerNotebookTabCoordinator.test.tsx b/apps/studio/components/interfaces/Explorer/__tests__/ExplorerNotebookTabCoordinator.test.tsx index 26e9377f308..2a167aa22e9 100644 --- a/apps/studio/components/interfaces/Explorer/__tests__/ExplorerNotebookTabCoordinator.test.tsx +++ b/apps/studio/components/interfaces/Explorer/__tests__/ExplorerNotebookTabCoordinator.test.tsx @@ -1,9 +1,11 @@ import { QueryClient } from '@tanstack/react-query' +import { screen } from '@testing-library/react' import { afterEach, describe, expect, it, vi } from 'vitest' import { ExplorerNotebookTabCoordinator } from '../ExplorerNotebookTabCoordinator' import { createQueryCellSkeleton } from '../utils' import { contentKeys } from '@/data/content/keys' +import { persistNotebookDraft, readNotebookDraft } from '@/state/notebooks/notebook-drafts' import { notebooksState } from '@/state/notebooks/notebooks-state' import type { Notebook } from '@/state/notebooks/types' import { createTabId, createTabsState, TabsStateContext } from '@/state/tabs' @@ -54,6 +56,7 @@ const renderCoordinator = (queryClient: QueryClient) => { afterEach(() => { delete notebooksState.notebooks[NOTEBOOK_ID] notebooksState.needsSaving.clear() + localStorage.clear() }) describe('ExplorerNotebookTabCoordinator', () => { @@ -104,4 +107,88 @@ describe('ExplorerNotebookTabCoordinator', () => { expect(tabsState.getCloseConfirmation([tabId])).toBeNull() }) + + it('asks for confirmation before closing a notebook with only a local draft, never loaded this session', () => { + persistNotebookDraft({ + projectRef: 'default', + id: NOTEBOOK_ID, + name: 'Test notebook', + content: { schema_version: 1, cells: [createQueryCellSkeleton()] }, + baseUpdatedAt: null, + }) + const { tabsState, tabId } = renderCoordinator(new QueryClient()) + + expect(tabsState.getCloseConfirmation([tabId])).toEqual({ + title: 'Unsaved changes', + description: 'You have unsaved changes in this notebook. Closing it will discard them.', + }) + }) + + it('removes a notebook local draft on close, even if it was never loaded into the store', () => { + persistNotebookDraft({ + projectRef: 'default', + id: NOTEBOOK_ID, + name: 'Test notebook', + content: { schema_version: 1, cells: [createQueryCellSkeleton()] }, + baseUpdatedAt: null, + }) + const { tabsState, tabId } = renderCoordinator(new QueryClient()) + + tabsState.closeTabs([tabId]) + + expect(readNotebookDraft({ projectRef: 'default', id: NOTEBOOK_ID })).toBeUndefined() + }) + + it('shows the status dot for a notebook with only a local draft, never loaded this session', () => { + persistNotebookDraft({ + projectRef: 'default', + id: NOTEBOOK_ID, + name: 'Test notebook', + content: { schema_version: 1, cells: [createQueryCellSkeleton()] }, + baseUpdatedAt: null, + }) + const { tabsState, tabId } = renderCoordinator(new QueryClient()) + const StatusIndicator = tabsState.getTabStatusIndicator('notebook')! + + customRender( + + ) + + expect(screen.getByRole('img', { name: 'Unsaved changes' })).toBeInTheDocument() + }) + + it('does not show the status dot when there is no draft and the notebook is not loaded', () => { + const { tabsState, tabId } = renderCoordinator(new QueryClient()) + const StatusIndicator = tabsState.getTabStatusIndicator('notebook')! + + customRender( + + ) + + expect(screen.queryByRole('img', { name: 'Unsaved changes' })).not.toBeInTheDocument() + }) + + it('updates the status dot live when a loaded notebook is edited, without remounting', async () => { + seedNotebook('saved') + const { tabsState, tabId } = renderCoordinator(new QueryClient()) + const StatusIndicator = tabsState.getTabStatusIndicator('notebook')! + + customRender( + + ) + expect(screen.queryByRole('img', { name: 'Unsaved changes' })).not.toBeInTheDocument() + + notebooksState.updateCells({ + id: NOTEBOOK_ID, + cells: [{ _tag: 'markdown_cell', _id: 'cell-1', text: 'hello' }], + }) + + expect(await screen.findByRole('img', { name: 'Unsaved changes' })).toBeInTheDocument() + }) }) diff --git a/apps/studio/components/interfaces/Explorer/hooks.ts b/apps/studio/components/interfaces/Explorer/hooks.ts index 6d152d9ffb2..77791bf9cf5 100644 --- a/apps/studio/components/interfaces/Explorer/hooks.ts +++ b/apps/studio/components/interfaces/Explorer/hooks.ts @@ -9,6 +9,7 @@ import { useProfile } from '@/lib/profile' import type { AssistantModel } from '@/state/ai-assistant-state' import { useAiAssistantState, whenAiAssistantInitialized } from '@/state/ai-assistant-state' import { useExplorerQueryStateSnapshot } from '@/state/explorer-query' +import { readNotebookDraft } from '@/state/notebooks/notebook-drafts' import { useNotebooksStateSnapshot } from '@/state/notebooks/notebooks-state' import { type Notebook } from '@/state/notebooks/types' import { Notebooks } from '@/types' @@ -17,9 +18,16 @@ import { Notebooks } from '@/types' * Fetches a notebook's content by id and merges it into the valtio store, so landing on * a notebook any way other than creating it in this session (direct link, hard refresh, * clicking it from the nav list) still hydrates `notebooksState`. + * + * A notebook that isn't in the store yet is also checked for a locally-persisted draft of unsaved edits + * If the notebook loaded from the server, the draft is restored on top of it, otherwise if + * the notebook was never saved at all, the draft — if present — is the only copy + * that ever existed, so it's restored as a new local-only notebook instead. */ export const useLoadNotebook = ({ id, projectRef }: { id?: string; projectRef?: string }) => { const notebooksSnap = useNotebooksStateSnapshot() + const { profile } = useProfile() + const { data: project } = useSelectedProjectQuery() const currentNotebook = id ? notebooksSnap.notebooks[id] : undefined const isCurrentProjectNotebook = currentNotebook?.projectRef === projectRef @@ -35,15 +43,46 @@ export const useLoadNotebook = ({ id, projectRef }: { id?: string; projectRef?: } ) + const isNotFound = isError && error.code === 404 + const mergeNotebook = useEffectEvent(() => { - if (projectRef && data) notebooksSnap.setNotebook({ projectRef, notebook: data }) + if (!projectRef || !id) return + + if (data) { + const isFreshLoad = !isCurrentProjectNotebook + notebooksSnap.setNotebook({ projectRef, notebook: data }) + if (isFreshLoad) { + notebooksSnap.restoreDraft({ projectRef, id, baseUpdatedAt: data.updated_at }) + } + return + } + + if (isNotFound && !isCurrentProjectNotebook && profile && project) { + const draft = readNotebookDraft({ projectRef, id }) + if (!draft) return + + notebooksSnap.addNotebook({ + projectRef, + notebook: { + id, + type: 'notebook', + name: draft.name, + description: '', + visibility: 'project', + favorite: false, + content: draft.content, + owner_id: profile.id, + project_id: project.id, + }, + }) + } }) useEffect(() => { mergeNotebook() - }, [projectRef, data]) + }, [projectRef, id, data, isNotFound, !!profile, !!project]) - return { isNotFound: isError && error.code === 404 } + return { isNotFound: isNotFound && !isCurrentProjectNotebook } } export const useCreateNotebook = () => { diff --git a/apps/studio/components/layouts/Tabs/SortableTab.tsx b/apps/studio/components/layouts/Tabs/SortableTab.tsx index 28ad583c34f..a59cfab95a0 100644 --- a/apps/studio/components/layouts/Tabs/SortableTab.tsx +++ b/apps/studio/components/layouts/Tabs/SortableTab.tsx @@ -128,7 +128,7 @@ export const SortableTab = ({ aria-hidden > {StatusIndicator && ( - + )} @@ -159,7 +159,7 @@ export const SortableTab = ({ className={cn( 'absolute top-1/2 right-2.5 z-10 -translate-y-1/2', 'flex size-5 items-center justify-center rounded-xs', - 'opacity-0 group-hover/tab:opacity-100 group-focus-within/tab:opacity-100 focus-visible:opacity-100', + 'opacity-0 group-hover/tab:opacity-100 group-focus-visible/tab:opacity-100 focus-visible:opacity-100', 'hover:bg-200 focus-ring', 'cursor-pointer' )} diff --git a/apps/studio/data/content/notebooks/notebook-cache.ts b/apps/studio/data/content/notebooks/notebook-cache.ts index 10ce4711450..79372b2833d 100644 --- a/apps/studio/data/content/notebooks/notebook-cache.ts +++ b/apps/studio/data/content/notebooks/notebook-cache.ts @@ -2,6 +2,7 @@ import type { QueryClient } from '@tanstack/react-query' import type { Snapshot } from 'valtio' import { contentKeys } from '@/data/content/keys' +import { removeNotebookDraft } from '@/state/notebooks/notebook-drafts' import { notebooksState } from '@/state/notebooks/notebooks-state' import type { StateNotebook } from '@/state/notebooks/types' import { hasUnsavedChanges } from '@/state/sql-editor/sql-editor-lifecycle' @@ -43,6 +44,7 @@ export function evictNotebookFromCaches({ }): boolean { notebooksState.removeNotebook({ id }) queryClient.removeQueries({ queryKey: contentKeys.resource(projectRef, id) }) + removeNotebookDraft({ projectRef, id }) return true } diff --git a/apps/studio/state/notebooks/notebook-drafts.test.ts b/apps/studio/state/notebooks/notebook-drafts.test.ts new file mode 100644 index 00000000000..9433901b81a --- /dev/null +++ b/apps/studio/state/notebooks/notebook-drafts.test.ts @@ -0,0 +1,177 @@ +import { untrustedSql } from '@supabase/pg-meta' +import { LOCAL_STORAGE_KEYS } from 'common' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +import { + hasNotebookDraft, + MAX_PERSISTED_NOTEBOOK_DRAFTS, + persistNotebookDraft, + readNotebookDraft, + removeNotebookDraft, +} from './notebook-drafts' +import type { NotebookContent } from '@/data/content/notebooks/notebook-schema' + +const createMemoryStorage = () => { + const values = new Map() + + return { + getItem: (key: string) => values.get(key) ?? null, + setItem: vi.fn((key: string, value: string) => values.set(key, value)), + removeItem: (key: string) => values.delete(key), + } +} + +const makeContent = (sql = 'select 1'): NotebookContent => ({ + schema_version: 1, + cells: [ + { _tag: 'markdown_cell', _id: 'cell-1', text: 'hello' }, + { + _tag: 'database_cell', + _id: 'cell-2', + unchecked_sql: untrustedSql(sql), + row_limit: 100, + view: 'table', + }, + ], +}) + +describe('notebook drafts', () => { + beforeEach(() => vi.useFakeTimers()) + afterEach(() => vi.useRealTimers()) + + it('persists and reads back a draft, re-branding SQL through the domain parse', () => { + const storage = createMemoryStorage() + + persistNotebookDraft({ + storage, + projectRef: 'project-a', + id: 'notebook-1', + name: 'My notebook', + content: makeContent('select * from users'), + baseUpdatedAt: '2024-01-01T00:00:00.000Z', + }) + + const draft = readNotebookDraft({ storage, projectRef: 'project-a', id: 'notebook-1' }) + + expect(draft?.name).toBe('My notebook') + expect(draft?.baseUpdatedAt).toBe('2024-01-01T00:00:00.000Z') + expect(draft?.content.cells).toMatchObject([ + { _tag: 'markdown_cell', text: 'hello' }, + { _tag: 'database_cell', unchecked_sql: 'select * from users' }, + ]) + }) + + it('scopes drafts by project ref', () => { + const storage = createMemoryStorage() + persistNotebookDraft({ + storage, + projectRef: 'project-a', + id: 'notebook-1', + name: 'My notebook', + content: makeContent(), + baseUpdatedAt: null, + }) + + expect( + readNotebookDraft({ storage, projectRef: 'project-b', id: 'notebook-1' }) + ).toBeUndefined() + }) + + it('reports whether a draft exists without needing to parse its content', () => { + const storage = createMemoryStorage() + expect(hasNotebookDraft({ storage, projectRef: 'project-a', id: 'notebook-1' })).toBe(false) + + persistNotebookDraft({ + storage, + projectRef: 'project-a', + id: 'notebook-1', + name: 'My notebook', + content: makeContent(), + baseUpdatedAt: null, + }) + + expect(hasNotebookDraft({ storage, projectRef: 'project-a', id: 'notebook-1' })).toBe(true) + }) + + it('removes a draft, clearing the storage key entirely once empty', () => { + const storage = createMemoryStorage() + const key = LOCAL_STORAGE_KEYS.NOTEBOOK_DRAFTS('project-a') + + persistNotebookDraft({ + storage, + projectRef: 'project-a', + id: 'notebook-1', + name: 'My notebook', + content: makeContent(), + baseUpdatedAt: null, + }) + expect(storage.getItem(key)).not.toBeNull() + + removeNotebookDraft({ storage, projectRef: 'project-a', id: 'notebook-1' }) + + expect( + readNotebookDraft({ storage, projectRef: 'project-a', id: 'notebook-1' }) + ).toBeUndefined() + expect(storage.getItem(key)).toBeNull() + }) + + it('falls back to no drafts when storage holds malformed JSON', () => { + const storage = createMemoryStorage() + storage.setItem(LOCAL_STORAGE_KEYS.NOTEBOOK_DRAFTS('project-a'), '{not json') + + expect( + readNotebookDraft({ storage, projectRef: 'project-a', id: 'notebook-1' }) + ).toBeUndefined() + }) + + it('drops an individual draft that fails schema validation without discarding the rest', () => { + const storage = createMemoryStorage() + persistNotebookDraft({ + storage, + projectRef: 'project-a', + id: 'notebook-good', + name: 'Good notebook', + content: makeContent(), + baseUpdatedAt: null, + }) + + const key = LOCAL_STORAGE_KEYS.NOTEBOOK_DRAFTS('project-a') + const raw = JSON.parse(storage.getItem(key)!) + raw['notebook-bad'] = { name: 'Bad notebook' } // missing required fields + storage.setItem(key, JSON.stringify(raw)) + + expect( + readNotebookDraft({ storage, projectRef: 'project-a', id: 'notebook-bad' }) + ).toBeUndefined() + expect(readNotebookDraft({ storage, projectRef: 'project-a', id: 'notebook-good' })?.name).toBe( + 'Good notebook' + ) + }) + + it('caps the number of persisted drafts, evicting the least recently updated', () => { + const storage = createMemoryStorage() + + for (let i = 0; i < MAX_PERSISTED_NOTEBOOK_DRAFTS + 1; i++) { + vi.setSystemTime(i) + persistNotebookDraft({ + storage, + projectRef: 'project-a', + id: `notebook-${i}`, + name: `Notebook ${i}`, + content: makeContent(), + baseUpdatedAt: null, + }) + } + + expect( + readNotebookDraft({ storage, projectRef: 'project-a', id: 'notebook-0' }) + ).toBeUndefined() + expect( + readNotebookDraft({ + storage, + projectRef: 'project-a', + id: `notebook-${MAX_PERSISTED_NOTEBOOK_DRAFTS}`, + }) + ).toBeDefined() + }) +}) diff --git a/apps/studio/state/notebooks/notebook-drafts.ts b/apps/studio/state/notebooks/notebook-drafts.ts new file mode 100644 index 00000000000..13c4b796bca --- /dev/null +++ b/apps/studio/state/notebooks/notebook-drafts.ts @@ -0,0 +1,152 @@ +import { LOCAL_STORAGE_KEYS, safeLocalStorage } from 'common' +import { z } from 'zod' + +import { + notebookDomainSchema, + notebookSchema, + toWireNotebook, + type NotebookContent, + type NotebookWire, +} from '@/data/content/notebooks/notebook-schema' + +/** + * A notebook's unsaved cell content, persisted locally so it survives a browser refresh. + * `baseUpdatedAt` snapshots the server's `updated_at` this draft branched from (or `null` + * for a notebook that's never been saved) — compared against the server's current value + * on restore to tell "nothing else changed" apart from "the server moved on since", e.g. + * an assistant edit landed while this draft was pending. + */ +export type PersistedNotebookDraft = { + name: string + content: NotebookWire + baseUpdatedAt: string | null + updatedAt: number +} + +type PersistedNotebookDrafts = Record + +type StorageLike = Pick + +export const MAX_PERSISTED_NOTEBOOK_DRAFTS = 50 + +const persistedDraftsSchema = z.record(z.string(), z.unknown()) +const persistedDraftSchema = z.object({ + name: z.string(), + content: notebookSchema, + baseUpdatedAt: z.string().nullable(), + updatedAt: z.number(), +}) + +function readPersistedDrafts(storage: StorageLike, projectRef: string): PersistedNotebookDrafts { + const raw = storage.getItem(LOCAL_STORAGE_KEYS.NOTEBOOK_DRAFTS(projectRef)) + if (!raw) return {} + + try { + const parsed = persistedDraftsSchema.safeParse(JSON.parse(raw)) + if (!parsed.success) return {} + + return Object.fromEntries( + Object.entries(parsed.data).flatMap(([id, value]) => { + const draft = persistedDraftSchema.safeParse(value) + return draft.success ? [[id, draft.data]] : [] + }) + ) + } catch { + return {} + } +} + +function writePersistedDrafts( + storage: StorageLike, + projectRef: string, + drafts: PersistedNotebookDrafts +) { + const key = LOCAL_STORAGE_KEYS.NOTEBOOK_DRAFTS(projectRef) + const retainedDrafts = Object.fromEntries( + Object.entries(drafts) + .sort(([, a], [, b]) => b.updatedAt - a.updatedAt) + .slice(0, MAX_PERSISTED_NOTEBOOK_DRAFTS) + ) + + if (Object.keys(retainedDrafts).length === 0) storage.removeItem(key) + else storage.setItem(key, JSON.stringify(retainedDrafts)) +} + +/** + * Writes (or overwrites) a notebook's local draft. Every cell-level edit already goes + * through a single choke point (`notebooksState.updateCells`), so this is called there + * directly rather than debounced — unlike ad-hoc query drafts, there's no raw keystroke + * stream reaching this layer (cell editors only commit on blur/run). + */ +export function persistNotebookDraft({ + storage = safeLocalStorage, + projectRef, + id, + name, + content, + baseUpdatedAt, +}: { + storage?: StorageLike + projectRef: string + id: string + name: string + content: NotebookContent + baseUpdatedAt: string | null +}) { + const drafts = readPersistedDrafts(storage, projectRef) + drafts[id] = { name, content: toWireNotebook(content), baseUpdatedAt, updatedAt: Date.now() } + writePersistedDrafts(storage, projectRef, drafts) +} + +/** + * Reads back a notebook's local draft, re-branding its SQL through the same wire→domain + * parse a server response goes through (`notebookDomainSchema`). + */ +export function readNotebookDraft({ + storage = safeLocalStorage, + projectRef, + id, +}: { + storage?: StorageLike + projectRef: string + id: string +}): { name: string; content: NotebookContent; baseUpdatedAt: string | null } | undefined { + const persisted = readPersistedDrafts(storage, projectRef)[id] + if (!persisted) return undefined + + const content = notebookDomainSchema.safeParse(persisted.content) + if (!content.success) return undefined + + return { name: persisted.name, content: content.data, baseUpdatedAt: persisted.baseUpdatedAt } +} + +/** + * Whether a draft exists for a notebook that isn't currently loaded into `notebooksState` + * (e.g. a background tab from a previous session, never opened this one) — used as a + * close-confirmation backstop, since `hasDiscardableChanges` can only see loaded notebooks. + */ +export function hasNotebookDraft({ + storage = safeLocalStorage, + projectRef, + id, +}: { + storage?: StorageLike + projectRef: string + id: string +}): boolean { + return readPersistedDrafts(storage, projectRef)[id] !== undefined +} + +export function removeNotebookDraft({ + storage = safeLocalStorage, + projectRef, + id, +}: { + storage?: StorageLike + projectRef: string + id: string +}) { + const drafts = readPersistedDrafts(storage, projectRef) + delete drafts[id] + writePersistedDrafts(storage, projectRef, drafts) +} diff --git a/apps/studio/state/notebooks/notebooks-state.test.ts b/apps/studio/state/notebooks/notebooks-state.test.ts index 41d9309a558..68c4f2f23d6 100644 --- a/apps/studio/state/notebooks/notebooks-state.test.ts +++ b/apps/studio/state/notebooks/notebooks-state.test.ts @@ -1,5 +1,6 @@ import { beforeEach, describe, expect, it } from 'vitest' +import { persistNotebookDraft, readNotebookDraft } from './notebook-drafts' import { notebooksState } from './notebooks-state' import type { Notebook } from './types' import type { Notebooks } from '@/types' @@ -28,6 +29,7 @@ describe('notebooksState', () => { notebooksState.needsSaving.clear() notebooksState.cellLocalState.clear() notebooksState.serverDivergedWhileDirty.clear() + localStorage.clear() }) it('addNotebook marks a locally-created notebook as new', () => { @@ -36,6 +38,14 @@ describe('notebooksState', () => { expect(notebooksState.notebooks['notebook-1'].status).toBe('new') }) + it('addNotebook persists a local draft immediately, even for a still-empty notebook', () => { + notebooksState.addNotebook({ projectRef: 'ref', notebook: makeNotebook('notebook-1') }) + + const draft = readNotebookDraft({ projectRef: 'ref', id: 'notebook-1' }) + expect(draft?.name).toBe('My Notebook') + expect(draft?.content.cells).toEqual([]) + }) + it('setNotebook marks a notebook not yet in the store as saved', () => { notebooksState.setNotebook({ projectRef: 'ref', notebook: makeNotebook('notebook-1') }) @@ -78,6 +88,22 @@ describe('notebooksState', () => { expect(notebooksState.needsSaving.get('notebook-1')).toBe(false) }) + it('renaming a loaded (saved) notebook transitions it to unsaved and persists a draft', () => { + notebooksState.setNotebook({ + projectRef: 'ref', + notebook: makeNotebook('notebook-1', { updated_at: '2024-01-01T00:00:00.000Z' }), + }) + + notebooksState.renameNotebook({ id: 'notebook-1', name: 'Renamed notebook' }) + + expect(notebooksState.notebooks['notebook-1'].status).toBe('unsaved') + expect(notebooksState.notebooks['notebook-1'].notebook.name).toBe('Renamed notebook') + + const draft = readNotebookDraft({ projectRef: 'ref', id: 'notebook-1' }) + expect(draft?.name).toBe('Renamed notebook') + expect(draft?.baseUpdatedAt).toBe('2024-01-01T00:00:00.000Z') + }) + it('editing a notebook that has never been saved keeps it as new', () => { notebooksState.addNotebook({ projectRef: 'ref', notebook: makeNotebook('notebook-1') }) @@ -133,4 +159,100 @@ describe('notebooksState', () => { expect(notebooksState.notebooks['notebook-1']).toBeUndefined() expect(notebooksState.serverDivergedWhileDirty.has('notebook-1')).toBe(false) }) + + it('persists a local draft of every cell edit', () => { + notebooksState.setNotebook({ + projectRef: 'ref', + notebook: makeNotebook('notebook-1', { updated_at: '2024-01-01T00:00:00.000Z' }), + }) + + notebooksState.updateCells({ + id: 'notebook-1', + cells: [{ _tag: 'markdown_cell', _id: 'cell-1', text: 'hello' }], + }) + + const draft = readNotebookDraft({ projectRef: 'ref', id: 'notebook-1' }) + expect(draft?.content.cells).toMatchObject([{ _tag: 'markdown_cell', text: 'hello' }]) + expect(draft?.baseUpdatedAt).toBe('2024-01-01T00:00:00.000Z') + }) + + it('clears a notebook local draft once it is saved', () => { + notebooksState.setNotebook({ projectRef: 'ref', notebook: makeNotebook('notebook-1') }) + notebooksState.updateCells({ + id: 'notebook-1', + cells: [{ _tag: 'markdown_cell', _id: 'cell-1', text: 'hello' }], + }) + expect(readNotebookDraft({ projectRef: 'ref', id: 'notebook-1' })).toBeDefined() + + notebooksState.markSaved({ id: 'notebook-1', updatedAt: '2024-02-02T00:00:00.000Z' }) + + expect(readNotebookDraft({ projectRef: 'ref', id: 'notebook-1' })).toBeUndefined() + expect(notebooksState.notebooks['notebook-1'].notebook.updated_at).toBe( + '2024-02-02T00:00:00.000Z' + ) + }) + + it('clears a notebook local draft when the notebook is removed', () => { + notebooksState.setNotebook({ projectRef: 'ref', notebook: makeNotebook('notebook-1') }) + notebooksState.updateCells({ + id: 'notebook-1', + cells: [{ _tag: 'markdown_cell', _id: 'cell-1', text: 'hello' }], + }) + + notebooksState.removeNotebook({ id: 'notebook-1' }) + + expect(readNotebookDraft({ projectRef: 'ref', id: 'notebook-1' })).toBeUndefined() + }) + + it('restores a local draft onto a freshly-loaded notebook', () => { + persistNotebookDraft({ + projectRef: 'ref', + id: 'notebook-1', + name: 'Restored name', + content: { + schema_version: 1, + cells: [{ _tag: 'markdown_cell', _id: 'cell-1', text: 'draft' }], + }, + baseUpdatedAt: '2024-01-01T00:00:00.000Z', + }) + notebooksState.setNotebook({ + projectRef: 'ref', + notebook: makeNotebook('notebook-1', { updated_at: '2024-01-01T00:00:00.000Z' }), + }) + + notebooksState.restoreDraft({ + projectRef: 'ref', + id: 'notebook-1', + baseUpdatedAt: '2024-01-01T00:00:00.000Z', + }) + + expect(notebooksState.notebooks['notebook-1'].notebook.name).toBe('Restored name') + expect(notebooksState.notebooks['notebook-1'].notebook.content?.cells).toMatchObject([ + { _tag: 'markdown_cell', text: 'draft' }, + ]) + expect(notebooksState.notebooks['notebook-1'].status).toBe('unsaved') + expect(notebooksState.serverDivergedWhileDirty.has('notebook-1')).toBe(false) + }) + + it('flags a server-diverged conflict when the draft branched from a stale server version', () => { + persistNotebookDraft({ + projectRef: 'ref', + id: 'notebook-1', + name: 'Restored name', + content: { schema_version: 1, cells: [] }, + baseUpdatedAt: '2024-01-01T00:00:00.000Z', + }) + notebooksState.setNotebook({ + projectRef: 'ref', + notebook: makeNotebook('notebook-1', { updated_at: '2024-06-01T00:00:00.000Z' }), + }) + + notebooksState.restoreDraft({ + projectRef: 'ref', + id: 'notebook-1', + baseUpdatedAt: '2024-06-01T00:00:00.000Z', + }) + + expect(notebooksState.serverDivergedWhileDirty.get('notebook-1')).toBe('updated') + }) }) diff --git a/apps/studio/state/notebooks/notebooks-state.ts b/apps/studio/state/notebooks/notebooks-state.ts index 96a7b716151..0634aedc7a7 100644 --- a/apps/studio/state/notebooks/notebooks-state.ts +++ b/apps/studio/state/notebooks/notebooks-state.ts @@ -5,6 +5,7 @@ import { useMemo } from 'react' import { proxy, snapshot, useSnapshot, type Snapshot } from 'valtio' import { proxyMap } from 'valtio/utils' +import { persistNotebookDraft, readNotebookDraft, removeNotebookDraft } from './notebook-drafts' import type { Notebook, StateNotebook } from './types' import { isQueryCell } from '@/data/content/notebooks/notebook-schema' import type { SnippetStatus } from '@/data/content/snippet-status' @@ -39,6 +40,16 @@ export const notebooksState = proxy({ addNotebook: ({ projectRef, notebook }: { projectRef: string; notebook: Notebook }) => { if (notebooksState.notebooks[notebook.id]) return notebooksState.notebooks[notebook.id] = { projectRef, notebook, status: 'new' } + + if (notebook.content) { + persistNotebookDraft({ + projectRef, + id: notebook.id, + name: notebook.name, + content: notebook.content, + baseUpdatedAt: null, + }) + } }, /** @@ -67,10 +78,18 @@ export const notebooksState = proxy({ * 'unsaved' -> ...), but the one-time 'new' -> 'saved' transition has no * other trigger — the resource query that would otherwise pick it up is * disabled while the notebook is still 'new'. + * + * `updatedAt` is the server's confirmed timestamp for this save, so the + * next locally-persisted draft (if any) branches from an accurate base + * rather than the notebook's stale initial-load timestamp. */ - markSaved: ({ id }: { id: string }) => { + markSaved: ({ id, updatedAt }: { id: string; updatedAt?: string }) => { const stateNotebook = notebooksState.notebooks[id] - if (stateNotebook) stateNotebook.status = 'saved' + if (stateNotebook) { + stateNotebook.status = 'saved' + if (updatedAt) stateNotebook.notebook.updated_at = updatedAt + removeNotebookDraft({ projectRef: stateNotebook.projectRef, id }) + } notebooksState.clearServerDivergence({ id }) }, @@ -81,13 +100,25 @@ export const notebooksState = proxy({ notebooksState.serverDivergedWhileDirty.delete(id), /** - * Rename follows its own async save directly at the call site rather than going - * through needsSaving/the debounced scheduler. + * Rename is bundled into the same "Save changes" action as cell edits, rather than its + * own immediate save — so it needs the same dirty-tracking and draft persistence as + * `updateCells`, or a rename with no cell changes would look clean and never get saved. */ renameNotebook: ({ id, name }: { id: string; name: string }) => { const stateNotebook = notebooksState.notebooks[id] - if (stateNotebook) { - stateNotebook.notebook.name = name + if (!stateNotebook) return + + stateNotebook.notebook.name = name + stateNotebook.status = statusOnEdit(stateNotebook.status) + + if (stateNotebook.notebook.content) { + persistNotebookDraft({ + projectRef: stateNotebook.projectRef, + id, + name, + content: stateNotebook.notebook.content, + baseUpdatedAt: stateNotebook.notebook.updated_at ?? null, + }) } }, @@ -103,6 +134,7 @@ export const notebooksState = proxy({ notebooksState.notebooks = otherNotebooks if (!skipSave) notebooksState.needsSaving.delete(id) notebooksState.clearServerDivergence({ id }) + if (notebook) removeNotebookDraft({ projectRef: notebook.projectRef, id }) }, /** @@ -126,6 +158,47 @@ export const notebooksState = proxy({ stateNotebook.notebook.content.cells = cells as Notebooks.Cell[] stateNotebook.status = statusOnEdit(stateNotebook.status) if (!skipSave) notebooksState.needsSaving.set(id, false) + + persistNotebookDraft({ + projectRef: stateNotebook.projectRef, + id, + name: stateNotebook.notebook.name, + content: stateNotebook.notebook.content, + baseUpdatedAt: stateNotebook.notebook.updated_at ?? null, + }) + }, + + /** + * Applies a locally-persisted draft on top of a freshly-loaded notebook — restoring + * edits that were never saved before the browser refreshed. `baseUpdatedAt` is the + * server's current `updated_at` for this notebook; if the draft branched from a + * different value, the server moved on while the draft was pending (e.g. an assistant + * edit), so the existing "assistant changes detected" conflict is raised rather than + * silently restoring over it — the user still sees their draft, but saving it requires + * the same confirmation an in-session conflict would. + */ + restoreDraft: ({ + projectRef, + id, + baseUpdatedAt, + }: { + projectRef: string + id: string + baseUpdatedAt: string + }) => { + const stateNotebook = notebooksState.notebooks[id] + if (!stateNotebook) return + + const draft = readNotebookDraft({ projectRef, id }) + if (!draft) return + + stateNotebook.notebook.name = draft.name + stateNotebook.notebook.content = draft.content + stateNotebook.status = statusOnEdit('saved') + + if (draft.baseUpdatedAt !== null && draft.baseUpdatedAt !== baseUpdatedAt) { + notebooksState.markServerDivergence({ id, type: 'updated' }) + } }, /** diff --git a/apps/studio/state/notebooks/types.ts b/apps/studio/state/notebooks/types.ts index ca3fe7eb4e6..cd989982a73 100644 --- a/apps/studio/state/notebooks/types.ts +++ b/apps/studio/state/notebooks/types.ts @@ -11,6 +11,7 @@ export interface Notebook { owner_id: number project_id: number content?: Notebooks.Content // Undefined until loaded + updated_at?: string // Absent for a notebook that's never been saved } export interface StateNotebook { diff --git a/packages/common/constants/local-storage.ts b/packages/common/constants/local-storage.ts index d215f27c315..0ffda24c5ca 100644 --- a/packages/common/constants/local-storage.ts +++ b/packages/common/constants/local-storage.ts @@ -56,6 +56,7 @@ export const LOCAL_STORAGE_KEYS = { SQL_EDITOR_TEMPORARY_FROM_EXPLORER: (ref: string) => `sql-editor-temporary-from-explorer-${ref}`, EXPLORER_QUERY_DRAFTS: (ref: string) => `explorer-query-drafts-${ref}`, + NOTEBOOK_DRAFTS: (ref: string) => `notebook-drafts-${ref}`, LOG_EXPLORER_SPLIT_SIZE: 'supabase_log-explorer-split-size', GRAPHQL_INTROSPECTION_NOTICE_COLLAPSED: (ref: string) =>