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
<img width="429" height="238" alt="image"
src="https://github.com/user-attachments/assets/5c392aed-1633-4428-8060-28f495a01f04"
/>

Also fixes an unrelated issue - renaming a notebook should mark the
notebook as having unsaved changes (with the orange indicator)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## 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.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

apps/studio/components/interfaces/Explorer/ExplorerNotebookTab.tsx

@@ -118,9 +118,9 @@ export const ExplorerNotebookTab = () => {
   const scrollContainerRef = useRef<HTMLDivElement>(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 = () => {
       <ConfirmationModal
         size="small"
         visible={isSaveConflictOpen}
-        title="Assistant changes detected"
+        title="Notebook changed on the server"
         additionalActionLabel="Discard changes"
         confirmLabel={
           id && snap.serverDivergedWhileDirty.get(id) === 'deleted' ? 'Recreate' : 'Save anyway'
@@ -554,8 +554,8 @@ export const ExplorerNotebookTab = () => {
       >
         <p className="text-sm">
           {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.'}
         </p>
       </ConfirmationModal>
 

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<StateNotebook> | 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 (
     <span
@@ -24,30 +44,11 @@ const NotebookTabStatusIndicator = ({ tab }: { tab: Tab }) => {
   )
 }
 
-/**
- * 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

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()

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(
+      <StatusIndicator
+        tab={{ id: tabId, type: 'notebook', metadata: { notebookId: NOTEBOOK_ID } }}
+      />
+    )
+
+    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(
+      <StatusIndicator
+        tab={{ id: tabId, type: 'notebook', metadata: { notebookId: NOTEBOOK_ID } }}
+      />
+    )
+
+    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(
+      <StatusIndicator
+        tab={{ id: tabId, type: 'notebook', metadata: { notebookId: NOTEBOOK_ID } }}
+      />
+    )
+    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()
+  })
 })

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 = () => {

apps/studio/components/layouts/Tabs/SortableTab.tsx

@@ -128,7 +128,7 @@ export const SortableTab = ({
             aria-hidden
           >
             {StatusIndicator && (
-              <span className="absolute inset-0 flex items-center justify-center group-hover/tab:opacity-0 group-focus-within/tab:opacity-0">
+              <span className="absolute inset-0 flex items-center justify-center group-hover/tab:opacity-0 group-focus-visible/tab:opacity-0">
                 <StatusIndicator tab={tab} />
               </span>
             )}
@@ -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'
             )}

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
 }

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<string, string>()
+
+  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()
+  })
+})

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<string, PersistedNotebookDraft>
+
+type StorageLike = Pick<typeof safeLocalStorage, 'getItem' | 'setItem' | 'removeItem'>
+
+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)
+}

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')
+  })
 })

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' })
+    }
   },
 
   /**

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 {

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) =>