fix(tldraw): fix color picker crash by persisting the new styles (tldraw#8980)

Fixing the app example page for custom colors: the example did not take
reloading into account and when a shape is created with a custom style
unfortunately it crashes the app because we don't persist the new
additions to palette and fonts

### Change type

- [x] `bugfix`
- [ ] `improvement`
- [ ] `feature`
- [ ] `api`
- [ ] `other`

### Test plan

1. go to example `/color-picker`
2. add a colour in the palette and some fonts
3. create a shape with new colour and font
4. reload 

- [ ] Unit tests
- [ ] End to end tests

### Release notes

- Fixed a bug in color-picker example that occurs when reloading the
canvas without persisting the added styles

Co-authored-by: Guillaume <guillaume@tldraw.com>

Author: kaneel

apps/examples/src/examples/ui/color-picker/ColorPickerExample.tsx

@@ -1,27 +1,32 @@
-import { useEffect, useRef, useState } from 'react'
+import { useEffect, useMemo, useRef, useState } from 'react'
+import 'tldraw/tldraw.css'
+import './color-picker.css'
 import {
 	DEFAULT_THEME,
 	DefaultColorStyle,
 	DefaultFontStyle,
 	Editor,
+	JsonObject,
+	TLDOCUMENT_ID,
 	TLDefaultColor,
+	TLShape,
+	TLShapePartial,
 	TLTheme,
 	TLThemeFont,
 	TLThemes,
 	TLUiOverrides,
 	Tldraw,
-	registerColorsFromThemes,
-	registerFontsFromThemes,
+	useValue,
 } from 'tldraw'
-import 'tldraw/tldraw.css'
-import './color-picker.css'
 
 // [1] Pre-declare slot names for custom colors and fonts. The runtime only
 // fills the slots the user has actually added, but TypeScript needs concrete
 // keys to compute `TLDefaultColorStyle` / `TLDefaultFontStyle`.
 const MAX_CUSTOM_COLORS = 20
 const MAX_CUSTOM_FONTS = 10
 
+const PERSISTENCE_KEY = 'color-picker-example'
+
 type CustomColorKey =
 	| 'custom-1'
 	| 'custom-2'
@@ -171,69 +176,200 @@ function ensureGoogleFontsLoaded() {
 	document.head.appendChild(link)
 }
 
+// [8] The palette (which hex/family each `custom-*` / `gf-*` slot maps to) lives
+// on the tldraw document's `meta`. Shapes only persist the slot *key* (`custom-2`,
+// `gf-1`); this is where the slot's actual value is kept. Storing it in the
+// document — instead of a separate React/localStorage copy — means the editor
+// persists it (via `persistenceKey`), syncs it across tabs, and includes it in
+// undo/redo for free, always atomically with the shapes that reference it. There
+// is therefore nothing to hand-reconcile: no cross-tab listener, no load-time
+// recovery, no save-failure handling.
+const PALETTE_META_KEY = 'colorPickerPalette'
+
+interface Palette {
+	hexes: string[]
+	fonts: GoogleFont[]
+}
+
+function isGoogleFont(v: unknown): v is GoogleFont {
+	return (
+		typeof v === 'object' &&
+		v !== null &&
+		typeof (v as GoogleFont).family === 'string' &&
+		typeof (v as GoogleFont).name === 'string'
+	)
+}
+
+// `meta` is free-form JSON, so normalise whatever we read into a known shape.
+// This also guards against a corrupt or hand-edited document.
+function parsePalette(raw: unknown): Palette {
+	const obj = (raw ?? {}) as { hexes?: unknown; fonts?: unknown }
+	return {
+		hexes: Array.isArray(obj.hexes)
+			? obj.hexes.filter((h): h is string => typeof h === 'string').slice(0, MAX_CUSTOM_COLORS)
+			: [],
+		fonts: Array.isArray(obj.fonts)
+			? obj.fonts.filter(isGoogleFont).slice(0, MAX_CUSTOM_FONTS)
+			: [],
+	}
+}
+
+function readRawPalette(editor: Editor): unknown {
+	return editor.getDocumentSettings().meta[PALETTE_META_KEY]
+}
+
+// Write the palette into the document `meta`. Call this inside `editor.run` so it
+// batches into one undo step with any shape repair that goes with it.
+function writePalette(editor: Editor, palette: Palette) {
+	editor.store.update(TLDOCUMENT_ID, (doc) => ({
+		...doc,
+		meta: { ...doc.meta, [PALETTE_META_KEY]: palette as unknown as JsonObject },
+	}))
+}
+
+const isCustomColor = (v: unknown): v is string => typeof v === 'string' && /^custom-\d+$/.test(v)
+const isCustomFont = (v: unknown): v is string => typeof v === 'string' && /^gf-\d+$/.test(v)
+
+const DEFAULT_COLOR = DefaultColorStyle.defaultValue
+const DEFAULT_FONT = DefaultFontStyle.defaultValue
+
+// [10] When the palette is cleared, the slots leave the *live* theme, so every
+// shape and the per-tab "next shape" style that still names a `custom-*` / `gf-*`
+// slot is remapped back to a built-in value (`black` / `draw`). The enum keeps
+// every slot registered (see `buildCompleteTheme`), so this isn't needed to keep
+// records valid — it just stops those shapes from rendering with an empty slot.
+// Run inside the same `editor.run` as the palette write so the two move together
+// through undo/redo.
+function repairShapesUsingCustomStyles(editor: Editor) {
+	const updates: TLShapePartial[] = []
+	for (const record of editor.store.allRecords()) {
+		if (record.typeName !== 'shape') continue
+		const shape = record as TLShape
+		const props = shape.props as Record<string, unknown>
+		const next: Record<string, unknown> = {}
+		if (isCustomColor(props.color)) next.color = DEFAULT_COLOR
+		if (isCustomColor(props.labelColor)) next.labelColor = DEFAULT_COLOR
+		if (isCustomFont(props.font)) next.font = DEFAULT_FONT
+		if (Object.keys(next).length > 0) {
+			updates.push({ id: shape.id, type: shape.type, props: next } as TLShapePartial)
+		}
+	}
+	if (updates.length > 0) editor.updateShapes(updates)
+
+	// `stylesForNextShape` lives on the per-tab `instance` record; reset any entry
+	// that points at a slot we're removing so the next created shape doesn't
+	// re-apply it.
+	const stylesForNextShape = editor.getInstanceState().stylesForNextShape
+	const fixed: Record<string, unknown> = {}
+	for (const [id, value] of Object.entries(stylesForNextShape)) {
+		if (isCustomColor(value)) fixed[id] = DEFAULT_COLOR
+		else if (isCustomFont(value)) fixed[id] = DEFAULT_FONT
+	}
+	if (Object.keys(fixed).length > 0) {
+		editor.updateInstanceState({ stylesForNextShape: { ...stylesForNextShape, ...fixed } })
+	}
+}
+
 export default function ColorPickerExample() {
-	const [customHexes, setCustomHexes] = useState<string[]>([])
-	const [customFonts, setCustomFonts] = useState<GoogleFont[]>([])
 	const [editor, setEditor] = useState<Editor | null>(null)
-	const prevColorCountRef = useRef(0)
-	const prevFontCountRef = useRef(0)
-
-	// [4] Imperative theme updates. Mutating the `themes` prop on <Tldraw>
-	// re-renders the whole editor tree and causes a visible flash. Calling
-	// `registerColorsFromThemes` + `editor.updateTheme` directly updates the
-	// style enum and theme atom in place — the style panel and affected shapes
-	// repaint reactively without remounting anything.
+	// [9] Register *every* possible slot up front (with placeholder values). Passed
+	// as the stable `themes` prop, this makes `createTLStore` register the full
+	// style enum before the persisted document is validated on load, so shapes
+	// referencing any slot pass validation even before the palette is read. The
+	// visible colors/fonts are pushed separately via `editor.updateTheme` below.
+	const [initialThemes] = useState<Partial<TLThemes>>(() => ({ default: buildCompleteTheme() }))
+
+	// [8b] Read the palette straight from the document `meta`, reactively. Because
+	// the store drives local edits, undo/redo, and cross-tab sync alike, this one
+	// read covers all three with no extra plumbing. `useValue` returns the stored
+	// reference (stable until the document record changes, so unrelated edits like
+	// moving a shape don't churn it); `useMemo` normalises it.
+	const rawPalette = useValue('palette', () => (editor ? readRawPalette(editor) : undefined), [
+		editor,
+	])
+	const palette = useMemo(() => parsePalette(rawPalette), [rawPalette])
+
+	// [4] Push the palette's actual colors/fonts into the live theme. The style
+	// panel and canvas read this, so only slots the palette contains are shown; the
+	// complete prop's placeholder slots stay hidden. Runs on load, on every palette
+	// change, and on undo/redo (all surface through `palette` above).
 	useEffect(() => {
 		if (!editor) return
-		const theme = buildTheme(customHexes, customFonts)
-		const themes = { default: theme } as TLThemes
-		registerColorsFromThemes(themes)
-		registerFontsFromThemes(themes)
-		editor.updateTheme(theme)
-
-		// [5] If this render added a new slot and the user has shapes selected,
-		// push the new slot onto the selection. We do this inside the effect
-		// (rather than in the click handler) so the enum is already updated by
-		// `registerColorsFromThemes` — setting an unknown value on a shape
-		// would throw.
-		const hasSelection = editor.getSelectedShapeIds().length > 0
-		if (hasSelection && customHexes.length > prevColorCountRef.current) {
-			const key = `custom-${customHexes.length}` as CustomColorKey
-			editor.setStyleForSelectedShapes(DefaultColorStyle, key)
-		}
-		if (hasSelection && customFonts.length > prevFontCountRef.current) {
-			const key = `gf-${customFonts.length}` as CustomFontKey
-			// The static type of DefaultFontStyle is narrowed to the four
-			// built-in fonts; we registered the runtime value ourselves above.
-			editor.setStyleForSelectedShapes(DefaultFontStyle, key as 'draw')
-		}
-
-		prevColorCountRef.current = customHexes.length
-		prevFontCountRef.current = customFonts.length
-	}, [editor, customHexes, customFonts])
+		editor.updateTheme(buildTheme(palette.hexes, palette.fonts))
+		if (palette.fonts.length > 0) ensureGoogleFontsLoaded()
+	}, [editor, palette])
 
 	const addColor = (hex: string) => {
-		setCustomHexes((prev) => (prev.length >= MAX_CUSTOM_COLORS ? prev : [...prev, hex]))
+		if (!editor || palette.hexes.length >= MAX_CUSTOM_COLORS) return
+		const hexes = [...palette.hexes, hex]
+		editor.run(() => {
+			writePalette(editor, { hexes, fonts: palette.fonts })
+			// [5] Put the new color into the live theme before applying it, so a
+			// selected shape repaints with it immediately. The enum already knows the
+			// slot (complete `themes` prop), so the style set never trips validation.
+			editor.updateTheme(buildTheme(hexes, palette.fonts))
+			if (editor.getSelectedShapeIds().length > 0) {
+				const key = `custom-${hexes.length}` as CustomColorKey
+				editor.setStyleForSelectedShapes(DefaultColorStyle, key)
+			}
+		})
 	}
 
 	const addFont = (font: GoogleFont) => {
-		setCustomFonts((prev) => {
-			if (prev.length >= MAX_CUSTOM_FONTS) return prev
-			if (prev.some((f) => f.family === font.family)) return prev
-			return [...prev, font]
+		if (!editor || palette.fonts.length >= MAX_CUSTOM_FONTS) return
+		if (palette.fonts.some((f) => f.family === font.family)) return
+		const fonts = [...palette.fonts, font]
+		ensureGoogleFontsLoaded()
+		editor.run(() => {
+			writePalette(editor, { hexes: palette.hexes, fonts })
+			editor.updateTheme(buildTheme(palette.hexes, fonts))
+			if (editor.getSelectedShapeIds().length > 0) {
+				// DefaultFontStyle's static type is the four built-ins; the runtime
+				// enum (complete `themes` prop) knows this slot.
+				const key = `gf-${fonts.length}` as CustomFontKey
+				editor.setStyleForSelectedShapes(DefaultFontStyle, key as 'draw')
+			}
+		})
+	}
+
+	// [11] Remove every custom color and font. We clear the palette and remap every
+	// shape / next-shape style off the custom slots in a single `editor.run`, so
+	// it's one undo step and the palette can never drift from the shapes that
+	// reference it — undo restores both together, redo clears both together.
+	const clearCustomStyles = () => {
+		if (!editor) return
+		editor.run(() => {
+			repairShapesUsingCustomStyles(editor)
+			writePalette(editor, { hexes: [], fonts: [] })
 		})
 	}
 
+	const hasCustomStyles = palette.hexes.length > 0 || palette.fonts.length > 0
+
 	return (
 		<div className="tldraw__editor">
-			<Tldraw persistenceKey="color-picker-example" overrides={uiOverrides} onMount={setEditor}>
+			<Tldraw
+				persistenceKey={PERSISTENCE_KEY}
+				themes={initialThemes}
+				overrides={uiOverrides}
+				onMount={(editor) => setEditor(editor)}
+			>
 				<div className="color-picker-toolbar" onPointerDown={(e) => e.stopPropagation()}>
-					<AddColorButton addColor={addColor} isFull={customHexes.length >= MAX_CUSTOM_COLORS} />
+					<AddColorButton addColor={addColor} isFull={palette.hexes.length >= MAX_CUSTOM_COLORS} />
 					<AddFontButton
-						addedFamilies={customFonts.map((f) => f.family)}
+						addedFamilies={palette.fonts.map((f) => f.family)}
 						addFont={addFont}
-						isFull={customFonts.length >= MAX_CUSTOM_FONTS}
+						isFull={palette.fonts.length >= MAX_CUSTOM_FONTS}
 					/>
+					{hasCustomStyles && (
+						<button
+							className="color-picker-button"
+							onClick={clearCustomStyles}
+							title="Remove all custom colors and fonts, resetting any shapes that use them"
+						>
+							Clear custom styles
+						</button>
+					)}
 				</div>
 			</Tldraw>
 		</div>
@@ -264,6 +400,26 @@ function buildTheme(hexes: string[], fonts: GoogleFont[]): TLTheme {
 	}
 }
 
+// A stable theme that registers *every* possible custom slot (`custom-1..N`,
+// `gf-1..N`). It's passed as the `themes` prop, which `createTLStore` registers
+// into the style enum *before* the persisted document is validated on load, and
+// which the editor re-registers on every render. Registering the full set means:
+//   - persisted or cross-tab shapes can never reference a slot the enum doesn't
+//     know, so they always pass validation (no load crash, no per-render drop);
+//   - the enum is never narrowed, so adding/removing palette entries can't
+//     invalidate existing shapes.
+// The placeholder values here are never shown: the style panel and canvas read
+// the *live* theme, which we set with `editor.updateTheme(buildTheme(...))` to
+// contain only the slots actually in the palette.
+function buildCompleteTheme(): TLTheme {
+	const hexes = Array.from({ length: MAX_CUSTOM_COLORS }, () => '#000000')
+	const fonts: GoogleFont[] = Array.from({ length: MAX_CUSTOM_FONTS }, (_, i) => ({
+		name: `placeholder-${i + 1}`,
+		family: 'sans-serif',
+	}))
+	return buildTheme(hexes, fonts)
+}
+
 // [6] Color add flow — the native picker opens immediately on "+ Add color".
 // The picker fires onChange continuously while dragging, so we stage a preview
 // and only commit on the explicit Add button.
@@ -418,29 +574,27 @@ Module augmentation needs concrete property names, so we pre-declare every
 possible slot up front. Only the ones the user actually adds land in the
 theme at runtime.
 
-[3]
-The link tag is lazily added on the first "+ Add font" open. Subsequent opens
-reuse the already-loaded CSS. We use a module-level flag instead of React
-state because the tag only needs to exist once in the document — remounting
-the example shouldn't add it again.
+[8]
+The palette lives on the document's `meta`, not in a separate React/localStorage
+copy. That single decision removes the cross-tab listener, the load-time
+recovery pass, and the save-failure handling that a side store would need: the
+editor already persists `meta` (via `persistenceKey`), syncs it to other tabs,
+and records it in undo/redo — always atomically with the shapes that reference
+the slots, so the two can never disagree.
 
 [4]
-Before: we passed custom colors and fonts through the `themes` prop, which
-meant every state change rebuilt the prop, re-rendered `TldrawEditor`, and
-caused a visible flash. Now the `themes` prop is omitted entirely; we
-imperatively register + update the theme from an effect. The style panel
-enum and the theme atom both update in place.
-
-[5]
-When the user adds a style while a shape is selected, we push the new slot
-onto the selection immediately via `queueMicrotask`. The microtask runs after
-the React commit and the effect in [4], so by the time we call
-`setStyleForSelectedShapes`, `DefaultColorStyle` / `DefaultFontStyle` already
-contains the new slot (otherwise the enum validator would reject it).
-
-Removal of custom slots is intentionally not supported:
-`registerColorsFromThemes` / `registerFontsFromThemes` strip entries that are
-absent from the theme, which would invalidate any shape still naming the
-removed slot.
+The style enum is registered once, from a *complete* theme passed via the
+`themes` prop (`buildCompleteTheme`), not from the palette. `createTLStore`
+registers the prop before the persisted document is validated on load, and the
+editor re-registers it on every render — so a shape can never reference a slot
+the enum doesn't know. The palette's actual colors/fonts are pushed into the
+live theme with `editor.updateTheme`; the style panel reads that live theme, so
+only slots the palette actually contains are shown.
+
+[11]
+"Clear custom styles" clears the palette and remaps every shape and the
+next-shape style off the custom slots back to the built-in defaults, both in one
+`editor.run`. Because the palette and the shapes are now both in the document,
+that single batch keeps them consistent through undo and redo.
 
 */