pi-tldraw 0.1.0

package/mcp-app/src/widget/persistence.ts

@@ -0,0 +1,337 @@
+import type { App } from '@modelcontextprotocol/ext-apps/react'
+import { Editor, structuredClone } from 'tldraw'
+import type { TLAsset, TLBindingCreate, TLShape } from 'tldraw'
+import { isPlainObject } from '../shared/utils'
+import { convertTldrawShapeToFocusedShape } from './focused/to-focused'
+
+export interface CanvasSnapshot {
+	shapes: TLShape[]
+	assets: TLAsset[]
+	bindings?: TLBindingCreate[]
+}
+
+// --- Canvas-scoped localStorage persistence ---
+
+let currentSessionId: string | null = null
+let currentCanvasId: string | null = null
+
+export function setCurrentSessionId(id: string): void {
+	currentSessionId = id
+}
+
+export function setCurrentCanvasId(id: string): void {
+	currentCanvasId = id
+}
+
+export function getCurrentCanvasId(): string | null {
+	return currentCanvasId
+}
+
+function localStorageKey(checkpointId: string): string {
+	if (currentCanvasId) return `tldraw:canvas:${currentCanvasId}:${checkpointId}`
+	if (currentSessionId) return `tldraw:${currentSessionId}:${checkpointId}`
+	return `tldraw:${checkpointId}`
+}
+
+function parseSnapshotData(
+	raw: string
+): { shapes: TLShape[]; assets: TLAsset[]; bindings: TLBindingCreate[] } | null {
+	const parsed = JSON.parse(raw)
+	// Backwards compat: old format was a plain array of shapes
+	if (Array.isArray(parsed)) {
+		const shapes = parsed.filter(
+			(s: unknown): s is TLShape =>
+				isPlainObject(s) && typeof s.id === 'string' && typeof s.type === 'string'
+		)
+		if (parsed.length > 0 && shapes.length === 0) return null
+		return { shapes, assets: [], bindings: [] }
+	}
+	if (!isPlainObject(parsed) || !Array.isArray(parsed.shapes)) return null
+	const shapes = parsed.shapes.filter(
+		(s: unknown): s is TLShape =>
+			isPlainObject(s) && typeof s.id === 'string' && typeof s.type === 'string'
+	)
+	if (parsed.shapes.length > 0 && shapes.length === 0) return null
+	const assets = (Array.isArray(parsed.assets) ? parsed.assets : []).filter(
+		(a: unknown): a is TLAsset => isPlainObject(a) && typeof a.id === 'string'
+	)
+	const bindings = (Array.isArray(parsed.bindings) ? parsed.bindings : []) as TLBindingCreate[]
+	return { shapes, assets, bindings }
+}
+
+export function loadLocalSnapshot(
+	checkpointId: string
+): { shapes: TLShape[]; assets: TLAsset[]; bindings: TLBindingCreate[] } | null {
+	try {
+		// eslint-disable-next-line tldraw/no-direct-storage
+		const raw = localStorage.getItem(localStorageKey(checkpointId))
+		if (!raw) return null
+		return parseSnapshotData(raw)
+	} catch {
+		return null
+	}
+}
+
+export function saveLocalSnapshot(
+	checkpointId: string,
+	shapes: TLShape[],
+	assets: TLAsset[] = [],
+	bindings: TLBindingCreate[] = []
+): void {
+	const scopeId = currentCanvasId ?? currentSessionId
+	if (!scopeId) return
+	try {
+		// eslint-disable-next-line tldraw/no-direct-storage
+		localStorage.setItem(
+			localStorageKey(checkpointId),
+			JSON.stringify({ shapes, assets, bindings })
+		)
+		const latestKey = currentCanvasId
+			? `tldraw:canvas:${currentCanvasId}:latest`
+			: `tldraw:${currentSessionId}:latest`
+		// eslint-disable-next-line tldraw/no-direct-storage
+		localStorage.setItem(latestKey, checkpointId)
+	} catch {
+		// localStorage may be full or unavailable.
+	}
+}
+
+export function getLatestCheckpointSnapshot(): {
+	shapes: TLShape[]
+	assets: TLAsset[]
+	bindings: TLBindingCreate[]
+} | null {
+	const scopeId = currentCanvasId ?? currentSessionId
+	if (!scopeId) return null
+	try {
+		const latestKey = currentCanvasId
+			? `tldraw:canvas:${currentCanvasId}:latest`
+			: `tldraw:${currentSessionId}:latest`
+		// eslint-disable-next-line tldraw/no-direct-storage
+		const latestId = localStorage.getItem(latestKey)
+		if (!latestId) return null
+		return loadLocalSnapshot(latestId)
+	} catch {
+		return null
+	}
+}
+
+// --- Embedded bootstrap ---
+
+declare global {
+	interface Window {
+		__TLDRAW_BOOTSTRAP__?: unknown
+	}
+}
+
+export function getEmbeddedBootstrap(): {
+	sessionId: string
+	canvasId?: string
+	checkpointId?: string
+	isDev: boolean
+	workerOrigin?: string
+	mcpSessionId?: string
+	snapshot?: CanvasSnapshot
+} | null {
+	const data = window.__TLDRAW_BOOTSTRAP__
+	if (!isPlainObject(data)) return null
+	const sessionId = typeof data.sessionId === 'string' ? data.sessionId : null
+	if (!sessionId) return null
+
+	const canvasId = typeof data.canvasId === 'string' ? data.canvasId : undefined
+	const checkpointId = typeof data.checkpointId === 'string' ? data.checkpointId : undefined
+	const isDev = data.isDev === true
+	const workerOrigin = typeof data.workerOrigin === 'string' ? data.workerOrigin : undefined
+	const mcpSessionId = typeof data.mcpSessionId === 'string' ? data.mcpSessionId : undefined
+
+	let snapshot: CanvasSnapshot | undefined
+	if (Array.isArray(data.shapes)) {
+		const shapes = data.shapes.filter(
+			(s: unknown): s is TLShape =>
+				isPlainObject(s) && typeof s.id === 'string' && typeof s.type === 'string'
+		)
+		const assets = (Array.isArray(data.assets) ? data.assets : []).filter(
+			(a: unknown): a is TLAsset => isPlainObject(a) && typeof a.id === 'string'
+		)
+		const bindings = (Array.isArray(data.bindings) ? data.bindings : []) as TLBindingCreate[]
+		if (data.shapes.length === 0 || shapes.length > 0) {
+			snapshot = { shapes, assets, bindings }
+		}
+	}
+
+	return { sessionId, canvasId, checkpointId, isDev, workerOrigin, mcpSessionId, snapshot }
+}
+
+// --- Tool result parsing ---
+
+function toSnapshotShapesFromRecords(value: unknown): TLShape[] | null {
+	if (!Array.isArray(value)) return null
+	return value.filter(
+		(shape): shape is TLShape =>
+			isPlainObject(shape) && typeof shape.id === 'string' && typeof shape.type === 'string'
+	)
+}
+
+function toAssetRecords(value: unknown): TLAsset[] {
+	if (!Array.isArray(value)) return []
+	return value.filter((a): a is TLAsset => isPlainObject(a) && typeof a.id === 'string')
+}
+
+interface CheckpointResult {
+	checkpointId: string
+	sessionId: string | null
+	shapes: TLShape[]
+	assets: TLAsset[]
+	bindings: TLBindingCreate[]
+	action: string | null
+	/** True if the server found base shapes to merge with (for create action). */
+	hadBaseShapes: boolean
+	/** True if the server started from a blank canvas (for create action). */
+	newBlankCanvas: boolean
+}
+
+function toBindingRecords(value: unknown): TLBindingCreate[] {
+	if (!Array.isArray(value)) return []
+	return value.filter(
+		(b): b is TLBindingCreate =>
+			isPlainObject(b) && typeof b.type === 'string' && typeof b.fromId === 'string'
+	)
+}
+
+export function parseCheckpointFromToolResult(result: unknown): CheckpointResult | null {
+	if (!isPlainObject(result)) return null
+	const sc = result.structuredContent
+	if (!isPlainObject(sc)) return null
+	const checkpointId = typeof sc.checkpointId === 'string' ? sc.checkpointId : null
+	if (!checkpointId) return null
+	const shapes = toSnapshotShapesFromRecords(sc.tldrawRecords)
+	if (!shapes) return null
+
+	const assets = toAssetRecords(sc.assets ?? sc.assetRecords)
+	const bindings = toBindingRecords(sc.bindings)
+	return {
+		checkpointId,
+		sessionId: typeof sc.sessionId === 'string' ? sc.sessionId : null,
+		shapes,
+		assets,
+		bindings,
+		action: typeof sc.action === 'string' ? sc.action : null,
+		hadBaseShapes: sc.hadBaseShapes === true,
+		newBlankCanvas: sc.newBlankCanvas === true,
+	}
+}
+
+// --- Snapshot capture + sync ---
+
+/**
+ * Capture the current editor state as a cloned snapshot.
+ */
+export function captureEditorSnapshot(editor: Editor): CanvasSnapshot {
+	return {
+		shapes: [...editor.getCurrentPageShapes()].map((s) => structuredClone(s)),
+		assets: [...editor.getAssets()].map((a) => structuredClone(a)),
+		bindings: getEditorBindings(editor),
+	}
+}
+
+/**
+ * Push model context and persist the editor state to localStorage + server.
+ * Returns the captured snapshot so the caller can store it.
+ */
+export function syncEditorState(
+	app: App,
+	editor: Editor,
+	checkpointId: string | null,
+	opts?: { message?: string }
+): CanvasSnapshot {
+	const snapshot = captureEditorSnapshot(editor)
+	pushCanvasContext(app, editor, opts)
+	if (checkpointId) {
+		saveLocalSnapshot(checkpointId, snapshot.shapes, snapshot.assets, snapshot.bindings ?? [])
+		void saveCheckpointToServer(app, checkpointId, editor)
+	}
+	return snapshot
+}
+
+// --- Canvas context push ---
+
+export function pushCanvasContext(app: App, editor: Editor, opts?: { message?: string }) {
+	const shapes = [...editor.getCurrentPageShapes()].map((shape) =>
+		convertTldrawShapeToFocusedShape(editor, shape)
+	)
+	const assets = [...editor.getAssets()].map((asset) => structuredClone(asset))
+	const bindings = getEditorBindings(editor)
+
+	const canvasStatus = shapes.length > 0 ? `Current canvas state is attached.` : 'Canvas is empty.'
+
+	const text = opts?.message ? `${opts.message}\n\n${canvasStatus}` : canvasStatus
+
+	void app.updateModelContext({
+		content: [{ type: 'text', text }],
+		structuredContent: {
+			shapes,
+			assets,
+			bindings,
+		},
+	})
+}
+
+export function clearCanvasContext(app: App, opts?: { message?: string }) {
+	const text = opts?.message ?? 'Canvas context was cleared.'
+	void app.updateModelContext({
+		content: [{ type: 'text', text }],
+		structuredContent: {
+			shapes: [],
+		},
+	})
+}
+
+export function getEditorBindings(editor: Editor): TLBindingCreate[] {
+	const bindings: TLBindingCreate[] = []
+	for (const record of editor.getCurrentPageShapes()) {
+		if (record.type !== 'arrow') continue
+		const arrowBindings = editor.getBindingsFromShape(record.id, 'arrow')
+		for (const binding of arrowBindings) {
+			bindings.push({
+				type: binding.type,
+				fromId: binding.fromId,
+				toId: binding.toId,
+				props: structuredClone(binding.props),
+				meta: structuredClone(binding.meta),
+			} as TLBindingCreate)
+		}
+	}
+	return bindings
+}
+
+export async function saveCheckpointToServer(
+	app: App,
+	checkpointId: string,
+	editor: Editor
+): Promise<boolean> {
+	const shapes = [...editor.getCurrentPageShapes()].map((s) => structuredClone(s))
+	const assets = [...editor.getAssets()].map((a) => structuredClone(a))
+	const bindings = getEditorBindings(editor)
+	const args: Record<string, string> = {
+		checkpointId,
+		shapesJson: JSON.stringify(shapes),
+	}
+	if (assets.length > 0) {
+		args.assetsJson = JSON.stringify(assets)
+	}
+	if (bindings.length > 0) {
+		args.bindingsJson = JSON.stringify(bindings)
+	}
+	if (currentCanvasId) {
+		args.canvasId = currentCanvasId
+	}
+	try {
+		const result = await app.callServerTool({
+			name: 'save_checkpoint',
+			arguments: args,
+		})
+		return result.isError !== true
+	} catch {
+		return false
+	}
+}

package/mcp-app/src/widget/snapshot.ts

@@ -0,0 +1,157 @@
+import { Box, Editor, structuredClone } from 'tldraw'
+import type { TLBindingCreate, TLShape, TLShapeId } from 'tldraw'
+import type { CanvasSnapshot } from './persistence'
+
+/**
+ * Re-run each ShapeUtil's onBeforeCreate to recalculate auto-sizing dimensions
+ * (e.g. growY for geo/note shapes). This compensates for shapes arriving from
+ * the server with static dimensions (growY: 0) that get applied via updateShapes,
+ * which doesn't trigger the sizing recalculation in onBeforeUpdate unless
+ * text/font/size actually changed.
+ */
+function forceAutoSize(editor: Editor) {
+	const updates: TLShape[] = []
+	for (const shape of editor.getCurrentPageShapes()) {
+		const util = editor.getShapeUtil(shape)
+		const adjusted = util.onBeforeCreate?.(shape as any)
+		if (adjusted) {
+			updates.push(adjusted as TLShape)
+		}
+	}
+	if (updates.length > 0) {
+		editor.updateShapes(updates)
+	}
+}
+
+/**
+ * Create bindings on the editor from a list of TLBindingCreate objects.
+ * Only creates bindings where the target shape exists on the page.
+ * Removes existing arrow bindings for affected arrows to prevent duplicates.
+ */
+function applyBindings(editor: Editor, bindings: TLBindingCreate[]) {
+	if (bindings.length === 0) return
+
+	// Collect arrow shape IDs that have bindings to dedup
+	const arrowIds = new Set(bindings.map((b) => b.fromId))
+	for (const arrowId of arrowIds) {
+		const existing = editor.getBindingsFromShape(arrowId, 'arrow')
+		if (existing.length > 0) {
+			editor.deleteBindings(existing)
+		}
+	}
+
+	for (const binding of bindings) {
+		if (editor.getShape(binding.toId)) {
+			editor.createBinding(binding)
+		}
+	}
+}
+
+export const CAMERA_ANIM_MS = 300
+let cameraAnimEndTime = 0
+
+/**
+ * If any of the given shape IDs extend outside the current viewport,
+ * pan (and only zoom out if necessary) to keep them visible.
+ * Never zooms in beyond the current zoom level.
+ * Skips the call if a previous animation is still playing.
+ */
+export function zoomToFitRequestShapes(editor: Editor, shapeIds: Set<TLShapeId>) {
+	if (shapeIds.size === 0) return
+
+	// Don't interrupt an in-progress animation — wait for it to finish
+	if (Date.now() < cameraAnimEndTime) return
+
+	const shapeBounds: Box[] = []
+	for (const id of shapeIds) {
+		const bounds = editor.getShapePageBounds(id)
+		if (bounds) shapeBounds.push(bounds)
+	}
+	if (shapeBounds.length === 0) return
+
+	const commonBounds = Box.Common(shapeBounds)
+	const viewportBounds = editor.getViewportPageBounds()
+
+	// All request shapes already visible — nothing to do
+	const contained =
+		commonBounds.x >= viewportBounds.x &&
+		commonBounds.y >= viewportBounds.y &&
+		commonBounds.x + commonBounds.w <= viewportBounds.x + viewportBounds.w &&
+		commonBounds.y + commonBounds.h <= viewportBounds.y + viewportBounds.h
+	if (contained) return
+
+	const currentZoom = editor.getZoomLevel()
+	const screenBounds = editor.getViewportScreenBounds()
+	const inset = 100
+
+	// The zoom level needed to fit the shapes with padding
+	const fitZoomX =
+		commonBounds.w > 0 ? (screenBounds.w - inset) / commonBounds.w : Number.POSITIVE_INFINITY
+	const fitZoomY =
+		commonBounds.h > 0 ? (screenBounds.h - inset) / commonBounds.h : Number.POSITIVE_INFINITY
+	const fitZoom = Math.min(fitZoomX, fitZoomY)
+
+	// Never zoom in past the current level — only zoom out if shapes don't fit
+	const zoom = Math.min(currentZoom, fitZoom)
+	if (!Number.isFinite(zoom) || zoom <= 0) return
+
+	// Center the shapes in the viewport at the chosen zoom
+	const cx = commonBounds.x + commonBounds.w / 2
+	const cy = commonBounds.y + commonBounds.h / 2
+	const cameraX = -cx + screenBounds.w / zoom / 2
+	const cameraY = -cy + screenBounds.h / zoom / 2
+	if (!Number.isFinite(cameraX) || !Number.isFinite(cameraY)) return
+
+	cameraAnimEndTime = Date.now() + CAMERA_ANIM_MS
+
+	editor.setCamera(
+		{
+			x: cameraX,
+			y: cameraY,
+			z: zoom,
+		},
+		{ animation: { duration: CAMERA_ANIM_MS } }
+	)
+}
+
+export function applySnapshot(editor: Editor, snapshot: CanvasSnapshot) {
+	const nextShapes = snapshot.shapes.map((shape) => structuredClone(shape))
+	const nextAssets = (snapshot.assets ?? []).map((asset) => structuredClone(asset))
+	const nextBindings = (snapshot.bindings ?? []) as TLBindingCreate[]
+
+	editor.store.mergeRemoteChanges(() => {
+		editor.run(
+			() => {
+				// Restore asset records first so image shapes can resolve them
+				if (nextAssets.length > 0) {
+					const existingAssetIds = new Set(editor.getAssets().map((a) => a.id))
+					for (const asset of nextAssets) {
+						if (!existingAssetIds.has(asset.id)) {
+							editor.createAssets([asset])
+						}
+					}
+				}
+
+				const existingIds = [...editor.getCurrentPageShapeIds()]
+				if (existingIds.length > 0) {
+					editor.deleteShapes(existingIds)
+				}
+				if (nextShapes.length <= 0) return
+
+				// Preserve parent relationships while allowing the editor to assign fresh indices.
+				const createInputs = nextShapes.map((shape) => {
+					const { index: _index, ...partial } = shape
+					return partial
+				})
+				editor.createShapes(createInputs)
+
+				// Re-run auto-sizing so growY / fontSizeAdjustment are correct
+				forceAutoSize(editor)
+
+				// Create bindings after all shapes are on the page
+				applyBindings(editor, nextBindings)
+			},
+			{ history: 'ignore' }
+		)
+	})
+}