worker/TldrawDurableObject.ts

import { RoomSnapshot, TLSocketRoom } from '@tldraw/sync-core'
import {
	TLRecord,
	createTLSchema,
	defaultShapeSchemas,
	RecordProps,
} from '@tldraw/tlschema'
import { T } from '@tldraw/validate'
import { DurableObject } from 'cloudflare:workers'
import { AutoRouter, IRequest, error } from 'itty-router'
import throttle from 'lodash.throttle'

// Define equation shape props using the correct v4 API
const equationShapeProps: RecordProps<any> = {
	w: T.number,
	h: T.number,
	type: T.string,
	expression: T.string,
	xExpression: T.string,
	yExpression: T.string,
	xMin: T.number,
	xMax: T.number,
	tMin: T.number,
	tMax: T.number,
	scaleX: T.number,
	scaleY: T.number,
	offsetX: T.number.optional(),
	offsetY: T.number.optional(),
	showAxes: T.boolean,
	showGrid: T.boolean,
	showNumbers: T.boolean,
	axisColor: T.string,
	gridColor: T.string,
	strokeColor: T.string,
	fontSize: T.number.optional(),
	fontFamily: T.string.optional(),
	lockAspectRatio: T.boolean.optional(),
}

// Create schema with custom equation shape
const schema = createTLSchema({
	shapes: { 
		...defaultShapeSchemas,
		equation: {
			props: equationShapeProps,
		}
	},
})
// each whiteboard room is hosted in a DurableObject:
// https://developers.cloudflare.com/durable-objects/

// there's only ever one durable object instance per room. it keeps all the room state in memory and
// handles websocket connections. periodically, it persists the room state to the R2 bucket.
export class TldrawDurableObject extends DurableObject {
	private r2: R2Bucket
	// the room ID will be missing while the room is being initialized
	private roomId: string | null = null
	// when we load the room from the R2 bucket, we keep it here. it's a promise so we only ever
	// load it once.
	private roomPromise: Promise<TLSocketRoom<TLRecord, void>> | null = null

    private lastPersistedSnapshot: string = ''

	constructor(ctx: DurableObjectState, env: Env) {
		super(ctx, env)
		this.r2 = env.TLDRAW_BUCKET

		ctx.blockConcurrencyWhile(async () => {
			this.roomId = ((await this.ctx.storage.get('roomId')) ?? null) as string | null
		})
	}

	private readonly router = AutoRouter({
		catch: (e) => {
			console.log(e)
			return error(e)
		},
	})
		// when we get a connection request, we stash the room id if needed and handle the connection
		.get('/api/connect/:roomId', async (request) => {
			if (!this.roomId) {
				await this.ctx.blockConcurrencyWhile(async () => {
					await this.ctx.storage.put('roomId', request.params.roomId)
					this.roomId = request.params.roomId
				})
			}
            console.log(`[DO] Connect request for room: ${this.roomId}`)
			return this.handleConnect(request)
		})
        // Get room metadata (name)
        .get('/api/meta/:roomId', async () => {
            const name = await this.ctx.storage.get('meta_name') as string | undefined
            console.log(`[DO] GET meta for ${this.roomId}: ${name || 'Untitled Board'}`)
            return { name: name || 'Untitled Board' }
        })
        // Update room metadata (name)
        .post('/api/meta/:roomId', async (request) => {
            try {
                const body = await request.json() as { name?: string }
                if (body.name) {
                    console.log(`[DO] UPDATE meta for ${this.roomId}: "${body.name}"`)
                    await this.ctx.storage.put('meta_name', body.name)
                    return { name: body.name }
                }
                return error(400, 'Missing name')
            } catch (e) {
                console.error(`[DO] Error updating meta:`, e)
                return error(400, 'Invalid JSON')
            }
        })
        .post('/internal/nuke', async () => {
            await this.ctx.storage.deleteAll();
            // Force reload on next connect
            this.roomPromise = null;
            this.roomId = null;
            this.lastPersistedSnapshot = '';
            return new Response(JSON.stringify({ success: true }));
        })

	// `fetch` is the entry point for all requests to the Durable Object
	fetch(request: Request): Response | Promise<Response> {
		return this.router.fetch(request)
	}

	// what happens when someone tries to connect to this room?
	async handleConnect(request: IRequest) {
		// extract query params from request
		const sessionId = request.query.sessionId as string
		if (!sessionId) return error(400, 'Missing sessionId')

		// Create the websocket pair for the client
		const { 0: clientWebSocket, 1: serverWebSocket } = new WebSocketPair()
		serverWebSocket.accept()

		// load the room, or retrieve it if it's already loaded
		const room = await this.getRoom()

		// connect the client to the room
		room.handleSocketConnect({ sessionId, socket: serverWebSocket })

		// return the websocket connection to the client
		return new Response(null, { status: 101, webSocket: clientWebSocket })
	}

	getRoom() {
		const roomId = this.roomId
		if (!roomId) throw new Error('Missing roomId')

		if (!this.roomPromise) {
			this.roomPromise = (async () => {
				// fetch the room from R2
				const roomFromBucket = await this.r2.get(`rooms/${roomId}`)

				// if it doesn't exist, we'll just create a new empty room
				const initialSnapshot = roomFromBucket
					? ((await roomFromBucket.json()) as RoomSnapshot)
					: undefined

				// create a new TLSocketRoom. This handles all the sync protocol & websocket connections.
				// it's up to us to persist the room state to R2 when needed though.
				return new TLSocketRoom<TLRecord, void>({
					schema,
					initialSnapshot,
					onDataChange: () => {
						// and persist whenever the data in the room changes
						this.schedulePersistToR2()
					},
				})
			})()
		}

		return this.roomPromise
	}

	// we throttle persistance so it only happens every 5 seconds (reduced from 10)
    // because we are now using non-blocking waitUntil
	schedulePersistToR2 = throttle(async () => {
		if (!this.roomPromise || !this.roomId) return
		const room = await this.getRoom()

		// convert the room to JSON
		const snapshot = JSON.stringify(room.getCurrentSnapshot())

        // Optimization: Skip if nothing changed
        if (snapshot === this.lastPersistedSnapshot) return
        this.lastPersistedSnapshot = snapshot

        const name = await this.ctx.storage.get('meta_name') as string || 'Untitled Board'
        
        // Optimization: Use waitUntil to allow the DO to continue processing while R2 writes
		this.ctx.waitUntil(
            this.r2.put(`rooms/${this.roomId}`, snapshot, {
                customMetadata: { name }
            })
        )
	}, 5_000)
}