Skip to content
490 linesCodeBlameRaw

Pick any line to see why it is the way it is: the commit, the pull request and issue it came from, and what the agent was thinking.

Docs: a workspace knowledge base people and agents write together1/**
2 * One page's live room, as a Durable Object named by the page id.
3 *
4 * It owns the page's Yjs document: every editor's socket syncs with it
5 * (the y-protocols sync and awareness messages, as binary frames), and
6 * every change made on the server (an agent's edit, an accepted
7 * suggestion, a restore, a comment) is applied here, so all of them merge
8 * as one CRDT. It holds sockets with the WebSocket Hibernation API, so a
9 * page people have open but aren't typing in costs nothing between
10 * keystrokes; the document is read back from storage when it wakes.
11 *
12 * Storage (the object's own SQLite): the document as a snapshot plus the
13 * updates since, compacted every so often. A few seconds after a burst of
14 * edits (an alarm), the room saves the Markdown rendition, the search
15 * index, backlinks and history to D1 (src/persist.ts).
16 *
17 * The room authorizes nothing about who may open the page: the Worker
18 * checks the viewer's role before forwarding a socket (src/index.ts,
19 * `live`) and puts it in ROOM_MEMBER_HEADER. The room enforces that role:
20 * a socket that may only view or comment never changes the document.
21 */
22import { DurableObject } from "cloudflare:workers";
23
24import type { DocEditTarget, DocRole, DocThreadAction, DocVersionKind, DocsLiveEvent, MemberProfile, ServiceBinding } from "@g1t/contracts";
25import * as decoding from "lib0/decoding";
26import * as encoding from "lib0/encoding";
27import * as syncProtocol from "y-protocols/sync";
28import * as Y from "yjs";
29
30import { atLeast } from "./access.ts";
31import { seed } from "./blocks.ts";
32import { anchorThread, applyEdit, findTarget, rangeIds, rangeMarkdown, restoreFrom, unanchorThread } from "./edits.ts";
33import { documentMarkdown, mentionedIds, outline, type Outline } from "./markdown.ts";
34import { save } from "./persist.ts";
35import { applyThreadAction, listThreads, setQuote, type ThreadResult } from "./threads.ts";
36
37/** Header the Worker sets on a socket it forwards: who it is, as JSON (`RoomMember`). */
38export const ROOM_MEMBER_HEADER = "x-g1t-docs-member";
39
40export type RoomMember = {
41 page_id: string;
42 workspace_slug: string;
43 /** `user:<id>`. */
44 key: string;
45 member: MemberProfile;
46 role: DocRole;
47};
48
49/** Who made a change on the server, for history. */
50export type Origin = { key: string; kind: DocVersionKind; note: string | null; authors?: string[] };
51
52type Attachment = RoomMember & { clients: number[] };
53
54type Env = { DB: D1Database; NOTIFY?: ServiceBinding };
55
56const FRAGMENT = "document-store";
57const MESSAGE_SYNC = 0;
58const MESSAGE_AWARENESS = 1;
59const MESSAGE_QUERY_AWARENESS = 3;
60/** Save this long after the last change. */
61const SAVE_AFTER_MS = 4_000;
62/** Compact the stored updates into one snapshot past this many. */
63const COMPACT_AT = 300;
64
65export class PageRoom extends DurableObject<Env> {
66 private doc: Y.Doc | null = null;
67 /** The last awareness update each client sent, so newcomers see everyone at once. Lost on hibernation; clients resend every 15 s. */
68 private awareness = new Map<number, Uint8Array>();
69 /** Each client's last awareness clock, to mark it gone with the next one. */
70 private clocks = new Map<number, number>();
71
72 constructor(ctx: DurableObjectState, env: Env) {
73 super(ctx, env);
74 ctx.setWebSocketAutoResponse(new WebSocketRequestResponsePair("ping", "pong"));
75 ctx.blockConcurrencyWhile(async () => {
76 this.ctx.storage.sql.exec("CREATE TABLE IF NOT EXISTS updates (seq INTEGER PRIMARY KEY AUTOINCREMENT, data BLOB NOT NULL)");
77 this.ctx.storage.sql.exec("CREATE TABLE IF NOT EXISTS snapshot (id INTEGER PRIMARY KEY CHECK (id = 1), data BLOB NOT NULL)");
78 this.ctx.storage.sql.exec("CREATE TABLE IF NOT EXISTS meta (key TEXT PRIMARY KEY, value TEXT NOT NULL)");
79 });
80 }
81
82 // ── State ──────────────────────────────────────────────────────────────
83
84 private meta<T>(key: string, fallback: T): T {
85 const row = this.ctx.storage.sql.exec<{ value: string }>("SELECT value FROM meta WHERE key = ?", key).toArray()[0];
86 if (!row) return fallback;
87 try {
88 return JSON.parse(row.value) as T;
89 } catch {
90 return fallback;
91 }
92 }
93
94 private setMeta(key: string, value: unknown): void {
95 this.ctx.storage.sql.exec("INSERT INTO meta (key, value) VALUES (?, ?) ON CONFLICT (key) DO UPDATE SET value = excluded.value", key, JSON.stringify(value));
96 }
97
98 /** The document, read from storage the first time it is needed after waking. */
99 private load(): Y.Doc {
100 if (this.doc) return this.doc;
101 const doc = new Y.Doc({ gc: true });
102 const snapshot = this.ctx.storage.sql.exec<{ data: ArrayBuffer }>("SELECT data FROM snapshot WHERE id = 1").toArray()[0];
103 if (snapshot) Y.applyUpdate(doc, new Uint8Array(snapshot.data));
104 for (const row of this.ctx.storage.sql.exec<{ data: ArrayBuffer }>("SELECT data FROM updates ORDER BY seq")) {
105 Y.applyUpdate(doc, new Uint8Array(row.data));
106 }
107 doc.on("update", (update: Uint8Array, origin: unknown) => this.onUpdate(update, origin));
108 this.doc = doc;
109 return doc;
110 }
111
112 private fragment(): Y.XmlFragment {
113 return this.load().getXmlFragment(FRAGMENT);
114 }
115
116 private onUpdate(update: Uint8Array, origin: unknown): void {
117 this.ctx.storage.sql.exec("INSERT INTO updates (data) VALUES (?)", update);
118 const count = this.ctx.storage.sql.exec<{ n: number }>("SELECT COUNT(*) AS n FROM updates").one().n;
119 if (count >= COMPACT_AT) this.compact();
120 // Who changed it: a socket's member, or a change made here.
121 const member = origin instanceof WebSocket ? (origin.deserializeAttachment() as Attachment | null) : null;
122 const key = member ? member.key : (origin as Origin | null)?.key;
123 if (member?.member.kind === "user") {
124 const names = this.meta<string[]>("editor_names", []);
125 const name = member.member.name.toLowerCase();
126 if (!names.includes(name)) this.setMeta("editor_names", [...names, name]);
127 }
128 if (key) {
129 const editors = this.meta<string[]>("editors", []).filter((k) => k !== key);
130 editors.push(key);
131 this.setMeta("editors", editors);
132 const pending = this.meta<string[]>("pending_authors", []);
133 if (!pending.includes(key)) this.setMeta("pending_authors", [...pending, key]);
134 }
135 // Everyone else sees it at once.
136 const encoder = encoding.createEncoder();
137 encoding.writeVarUint(encoder, MESSAGE_SYNC);
138 syncProtocol.writeUpdate(encoder, update);
139 this.send(encoding.toUint8Array(encoder), origin instanceof WebSocket ? origin : null);
140 this.scheduleSave();
141 }
142
143 private compact(): void {
144 const doc = this.load();
145 this.ctx.storage.transactionSync(() => {
146 this.ctx.storage.sql.exec("INSERT INTO snapshot (id, data) VALUES (1, ?) ON CONFLICT (id) DO UPDATE SET data = excluded.data", Y.encodeStateAsUpdate(doc));
147 this.ctx.storage.sql.exec("DELETE FROM updates");
148 });
149 }
150
151 private scheduleSave(): void {
152 void this.ctx.storage.getAlarm().then((at) => {
153 if (at == null) return this.ctx.storage.setAlarm(Date.now() + SAVE_AFTER_MS);
154 });
155 }
156
157 private send(message: Uint8Array | string, except: WebSocket | null = null): void {
158 for (const socket of this.ctx.getWebSockets()) {
159 if (socket === except) continue;
160 try {
161 socket.send(message);
162 } catch {
163 // Closing already.
164 }
165 }
166 }
167
168 /** Saves to D1 now: the Markdown, search, links, mentions, and a version if one is due or asked for. */
169 private async persist(version: Origin | null = null): Promise<string | null> {
170 const pageId = this.meta<string | null>("page_id", null);
171 if (!pageId) return null;
172 const fragment = this.fragment();
173 const editors = this.meta<string[]>("editors", []);
174 const pending = this.meta<string[]>("pending_authors", []);
175 const result = await save(this.env, {
176 page_id: pageId,
177 markdown: documentMarkdown(fragment),
178 editors,
179 mentioned: mentionedIds(fragment).users,
180 editor_names: this.meta<string[]>("editor_names", []),
181 state: Y.encodeStateAsUpdate(this.load()),
182 version: version ? { kind: version.kind, note: version.note, authors: version.authors ?? [version.key] } : null,
183 pending_authors: pending,
184 last_version_at: this.meta<number>("last_version_at", 0),
185 workspace_slug: this.meta<string | null>("workspace_slug", null),
186 });
187 this.setMeta("editors", []);
188 this.setMeta("editor_names", []);
189 if (result.version_id) {
190 this.setMeta("last_version_at", Date.now());
191 this.setMeta("pending_authors", []);
192 }
193 return result.version_id;
194 }
195
196 override async alarm(): Promise<void> {
197 await this.persist();
198 }
199
200 // ── Calls from the Worker ──────────────────────────────────────────────
201
202 /**
203 * Names the page this room is for, and fills an empty document from
204 * Markdown (a template, an agent's new page, or a page made before its
205 * room existed) or from a Yjs state (a duplicate).
206 */
207 async ensure(init: { page_id: string; workspace_slug: string; markdown?: string | null; state?: Uint8Array | null }): Promise<void> {
208 this.setMeta("page_id", init.page_id);
209 if (init.workspace_slug) this.setMeta("workspace_slug", init.workspace_slug);
210 const doc = this.load();
211 const fragment = this.fragment();
212 if (fragment.length > 0) return;
213 if (init.state) Y.applyUpdate(doc, init.state, { key: "system", kind: "created", note: null } satisfies Origin);
214 else seed(doc, fragment, init.markdown ?? "");
215 }
216
217 /** The document as Markdown, and its top-level blocks. */
218 async read(): Promise<{ markdown: string; blocks: Outline[] }> {
219 const fragment = this.fragment();
220 return { markdown: documentMarkdown(fragment), blocks: outline(fragment) };
221 }
222
223 /** The whole document's state, for a duplicate. */
224 async state(): Promise<Uint8Array> {
225 return Y.encodeStateAsUpdate(this.load());
226 }
227
228 /** A target's current Markdown and blocks, or null when it is gone. */
229 async target(target: DocEditTarget): Promise<{ markdown: string; block_ids: string[] } | null> {
230 const fragment = this.fragment();
231 const range = findTarget(fragment, target);
232 if (!range) return null;
233 return { markdown: rangeMarkdown(fragment, range), block_ids: rangeIds(fragment, range) };
234 }
235
236 /** Where each target is now, for marking open suggestions in the editor. */
237 async targets(targets: DocEditTarget[]): Promise<(string[] | null)[]> {
238 const fragment = this.fragment();
239 return targets.map((t) => {
240 const range = findTarget(fragment, t);
241 return range ? rangeIds(fragment, range) : null;
242 });
243 }
244
245 /** Applies an edit and records a version for it. False when the target is gone. */
246 async edit(target: DocEditTarget, markdown: string, origin: Origin): Promise<{ applied: boolean; version_id: string | null }> {
247 const doc = this.load();
248 const applied = applyEdit(doc, this.fragment(), target, markdown, origin);
249 if (!applied) return { applied: false, version_id: null };
250 const version_id = await this.persist(origin);
251 return { applied, version_id };
252 }
253
254 /** Makes the document what a version's was, as a new version. */
255 async restore(input: { state: Uint8Array | null; markdown: string }, origin: Origin): Promise<string | null> {
256 const doc = this.load();
257 if (input.state) {
258 const old = new Y.Doc();
259 Y.applyUpdate(old, input.state);
260 restoreFrom(doc, this.fragment(), old.getXmlFragment(FRAGMENT), origin);
261 } else applyEdit(doc, this.fragment(), { kind: "document" }, input.markdown, origin);
262 return this.persist(origin);
263 }
264
265 /** A comment operation from `actor` with `role`. */
266 async thread(actor: string, role: DocRole, action: DocThreadAction): Promise<ThreadResult> {
267 const doc = this.load();
268 if (action.op === "anchor") {
269 if (!atLeast(role, "comment")) return { ok: false, code: "forbidden", message: "You can read this page but not comment on it." };
270 const thread = doc.getMap<Y.Map<unknown>>("threads").get(action.thread_id);
271 if (!thread) return { ok: false, code: "not_found", message: "That thread is gone." };
272 const quote = anchorThread(doc, this.fragment(), action.anchor, action.head, action.thread_id);
273 if (quote) setQuote(doc, action.thread_id, quote);
274 return { ok: true, value: { quote } };
275 }
276 const result = applyThreadAction(doc, actor, role, action);
277 if (result.ok && action.op === "delete_thread") unanchorThread(doc, this.fragment(), action.thread_id);
278 return result;
279 }
280
281 /** The page's threads, authors as member keys. */
282 async threads(): Promise<ReturnType<typeof listThreads>> {
283 return listThreads(this.load());
284 }
285
286 /**
287 * Shows an agent on the page for a little while (its face in everyone's
288 * presence row) when it edits or suggests: an awareness entry of its
289 * own, which editors drop after 30 s without renewal, as for anyone.
290 */
291 async announce(key: string, name: string): Promise<void> {
292 let id = 0;
293 for (let i = 0; i < key.length; i++) id = (Math.imul(id, 31) + key.charCodeAt(i)) | 0;
294 const client = (id & 0x3fffffff) + 1;
295 const clock = Math.floor(Date.now() / 1000);
296 const state = JSON.stringify({ user: { name, color: "#b8a6ff", key, kind: "agent", avatar: "" } });
297 const update = encoding.createEncoder();
298 encoding.writeVarUint(update, 1);
299 encoding.writeVarUint(update, client);
300 encoding.writeVarUint(update, clock);
301 encoding.writeVarString(update, state);
302 const bytes = encoding.toUint8Array(update);
303 this.awareness.set(client, bytes);
304 this.clocks.set(client, clock);
305 const message = encoding.createEncoder();
306 encoding.writeVarUint(message, MESSAGE_AWARENESS);
307 encoding.writeVarUint8Array(message, bytes);
308 this.send(encoding.toUint8Array(message));
309 }
310
311 /** Tells everyone with the page open. */
312 async notice(event: DocsLiveEvent): Promise<void> {
313 this.send(JSON.stringify(event));
314 }
315
316 /** A member's role changed (null: they can no longer read it): their sockets follow. */
317 async setRole(key: string, role: DocRole | null): Promise<void> {
318 for (const socket of this.ctx.getWebSockets(key)) {
319 const who = socket.deserializeAttachment() as Attachment | null;
320 if (!who) continue;
321 if (!role) {
322 try {
323 socket.close(4403, "No longer allowed");
324 } catch {
325 // Already closed.
326 }
327 continue;
328 }
329 socket.serializeAttachment({ ...who, role });
330 try {
331 socket.send(JSON.stringify({ type: "access", role } satisfies DocsLiveEvent));
332 } catch {
333 // Closing.
334 }
335 }
336 }
337
338 /** Saves now, as before the page is archived or exported. */
339 async flush(): Promise<void> {
340 if (this.doc) await this.persist();
341 }
342
343 /** Closes every socket: the page went to the trash. */
344 async closeAll(reason: string): Promise<void> {
345 for (const socket of this.ctx.getWebSockets()) {
346 try {
347 socket.send(JSON.stringify({ type: "page.archived", page_id: this.meta<string>("page_id", "") } satisfies DocsLiveEvent));
348 socket.close(4410, reason);
349 } catch {
350 // Already closed.
351 }
352 }
353 }
354
355 // ── Sockets ────────────────────────────────────────────────────────────
356
357 override async fetch(request: Request): Promise<Response> {
358 if (request.headers.get("upgrade")?.toLowerCase() !== "websocket") {
359 return new Response("Expected a WebSocket upgrade\n", { status: 426 });
360 }
361 let who: RoomMember;
362 try {
363 who = JSON.parse(request.headers.get(ROOM_MEMBER_HEADER) ?? "") as RoomMember;
364 } catch {
365 return new Response("Missing member\n", { status: 400 });
366 }
367 if (who.workspace_slug) this.setMeta("workspace_slug", who.workspace_slug);
368 if (who.page_id) this.setMeta("page_id", who.page_id);
369 const pair = new WebSocketPair();
370 const [client, server] = Object.values(pair) as [WebSocket, WebSocket];
371 this.ctx.acceptWebSocket(server, [who.key]);
372 server.serializeAttachment({ ...who, clients: [] } satisfies Attachment);
373 // Start syncing: our state vector, and everyone's presence so far.
374 const doc = this.load();
375 const encoder = encoding.createEncoder();
376 encoding.writeVarUint(encoder, MESSAGE_SYNC);
377 syncProtocol.writeSyncStep1(encoder, doc);
378 server.send(encoding.toUint8Array(encoder));
379 for (const update of this.awareness.values()) {
380 const e = encoding.createEncoder();
381 encoding.writeVarUint(e, MESSAGE_AWARENESS);
382 encoding.writeVarUint8Array(e, update);
383 server.send(encoding.toUint8Array(e));
384 }
385 return new Response(null, { status: 101, webSocket: client });
386 }
387
388 override async webSocketMessage(socket: WebSocket, message: string | ArrayBuffer): Promise<void> {
389 // Text frames are for the Worker's notices only; clients speak binary.
390 if (typeof message === "string") return;
391 const who = socket.deserializeAttachment() as Attachment | null;
392 if (!who) return;
393 const data = new Uint8Array(message);
394 const decoder = decoding.createDecoder(data);
395 const type = decoding.readVarUint(decoder);
396 if (type === MESSAGE_SYNC) {
397 // Peek at the sync message: viewers and commenters may ask for the
398 // document (step 1) but never change it (step 2, update).
399 const peek = decoding.createDecoder(data);
400 decoding.readVarUint(peek);
401 const step = decoding.readVarUint(peek);
402 if (step !== syncProtocol.messageYjsSyncStep1 && !atLeast(who.role, "edit")) return;
403 const doc = this.load();
404 const encoder = encoding.createEncoder();
405 encoding.writeVarUint(encoder, MESSAGE_SYNC);
406 syncProtocol.readSyncMessage(decoder, encoder, doc, socket);
407 if (encoding.length(encoder) > 1) socket.send(encoding.toUint8Array(encoder));
408 return;
409 }
410 if (type === MESSAGE_AWARENESS) {
411 const update = decoding.readVarUint8Array(decoder);
412 // Remember which clients this socket speaks for, to clear them when it closes.
413 const entries = awarenessEntries(update);
414 const clients = entries.map((e) => e.id);
415 for (const e of entries) {
416 this.clocks.set(e.id, e.clock);
417 if (e.gone) this.awareness.delete(e.id);
418 else this.awareness.set(e.id, update);
419 }
420 const known = new Set(who.clients);
421 if (clients.some((id) => !known.has(id))) socket.serializeAttachment({ ...who, clients: [...new Set([...who.clients, ...clients])] });
422 const encoder = encoding.createEncoder();
423 encoding.writeVarUint(encoder, MESSAGE_AWARENESS);
424 encoding.writeVarUint8Array(encoder, update);
425 this.send(encoding.toUint8Array(encoder), socket);
426 return;
427 }
428 if (type === MESSAGE_QUERY_AWARENESS) {
429 for (const update of this.awareness.values()) {
430 const encoder = encoding.createEncoder();
431 encoding.writeVarUint(encoder, MESSAGE_AWARENESS);
432 encoding.writeVarUint8Array(encoder, update);
433 socket.send(encoding.toUint8Array(encoder));
434 }
435 }
436 }
437
438 override async webSocketClose(socket: WebSocket, code: number, reason: string): Promise<void> {
439 this.leave(socket);
440 try {
441 socket.close(code, reason);
442 } catch {
443 // Already closed.
444 }
445 }
446
447 override async webSocketError(socket: WebSocket): Promise<void> {
448 this.leave(socket);
449 }
450
451 /** Tells everyone a closed socket's people left: their cursors go. */
452 private leave(socket: WebSocket): void {
453 const who = socket.deserializeAttachment() as Attachment | null;
454 if (!who?.clients.length) return;
455 for (const id of who.clients) this.awareness.delete(id);
456 // Marks each client gone (state null, the next clock). One whose clock
457 // was lost to hibernation is left to time out in the others' editors
458 // (30 s), as y-protocols does for any client that stops renewing.
459 const gone = who.clients.filter((id) => this.clocks.has(id));
460 if (!gone.length) return;
461 const encoder = encoding.createEncoder();
462 encoding.writeVarUint(encoder, gone.length);
463 for (const id of gone) {
464 encoding.writeVarUint(encoder, id);
465 encoding.writeVarUint(encoder, this.clocks.get(id)! + 1);
466 encoding.writeVarString(encoder, "null");
467 this.clocks.delete(id);
468 }
469 const update = encoding.toUint8Array(encoder);
470 const message = encoding.createEncoder();
471 encoding.writeVarUint(message, MESSAGE_AWARENESS);
472 encoding.writeVarUint8Array(message, update);
473 this.send(encoding.toUint8Array(message), socket);
474 }
475}
476
477/** The clients in an awareness update (its format: count, then id, clock, JSON state). */
478export function awarenessEntries(update: Uint8Array): { id: number; clock: number; gone: boolean }[] {
479 const decoder = decoding.createDecoder(update);
480 const count = decoding.readVarUint(decoder);
481 const out: { id: number; clock: number; gone: boolean }[] = [];
482 for (let i = 0; i < count; i++) {
483 const id = decoding.readVarUint(decoder);
484 const clock = decoding.readVarUint(decoder);
485 const state = decoding.readVarString(decoder);
486 out.push({ id, clock, gone: state === "null" });
487 }
488 return out;
489}
490