1
import { YServer } from "y-partyserver";
2
import type { Connection, ConnectionContext } from "partyserver";
3
import * as Y from "yjs";
4
import { eq } from "drizzle-orm";
5
import { drizzle, type DrizzleSqliteDODatabase } from "drizzle-orm/durable-sqlite";
6
import { migrate } from "drizzle-orm/durable-sqlite/migrator";
7
import { createSessionDb } from "@/worker/db/d1/client";
8
import { pages } from "@/worker/db/d1/schema";
9
import * as docSyncSchema from "@/worker/db/docsync-do/schema";
10
import { createLogger, errorContext, setLevel } from "@/worker/lib/logger";
11
import { DEFAULT_PAGE_TITLE } from "@/worker/lib/constants";
12
import { recordDocSyncPageSave } from "@/worker/lib/site-invalidation";
13
import { YJS_PAGE_TITLE } from "@/shared/constants";
14
import { parseDocMessage } from "@/shared/doc-messages";
15
import type { PageKind } from "@/shared/types";
16
import { extractCanvasPlaintext, extractPlaintext } from "@/shared/editor/yjs-text";
17
import docSyncMigrations from "../../../drizzle/docsync-do/migrations.js";
18
 
19
const MAX_CONNECTIONS_PER_DOC = 20;
20
const CHUNK_SIZE = 1.5 * 1024 * 1024; // 1.5MB per chunk, under 2MB SQLite row limit
21
const log = createLogger("doc-sync");
22
 
23
const READONLY_TAG = "readonly";
24
const MEMBER_EDIT_TAG = "member_edit";
25
 
26
type DocSyncDb = DrizzleSqliteDODatabase<typeof docSyncSchema>;
27
 
28
interface YpsConnectionState {
29
  __ypsAwarenessIds?: number[];
30
}
31
 
32
/** Read awareness client IDs stored on the connection by y-partyserver. */
33
function getAwarenessIds(conn: Connection): readonly number[] {
34
  try {
35
    return (conn as Connection<YpsConnectionState>).state?.__ypsAwarenessIds ?? [];
36
  } catch {
37
    return [];
38
  }
39
}
40
 
41
function chunkBuffer(buf: Uint8Array, chunkSize: number): Uint8Array[] {
42
  if (buf.byteLength <= chunkSize) return [buf];
43
  const chunks: Uint8Array[] = [];
44
  for (let offset = 0; offset < buf.byteLength; offset += chunkSize) {
45
    chunks.push(buf.slice(offset, offset + chunkSize));
46
  }
47
  return chunks;
48
}
49
 
50
interface ChunkRow {
51
  chunk_index: number;
52
  data: Uint8Array;
53
}
54
 
55
function reassembleChunks(rows: ChunkRow[]): Uint8Array | null {
56
  if (rows.length === 0) return null;
57
  if (rows.length === 1) return rows[0].data;
58
  const sorted = [...rows].sort((a, b) => a.chunk_index - b.chunk_index);
59
  let totalLength = 0;
60
  for (const r of sorted) totalLength += r.data.byteLength;
61
  const result = new Uint8Array(totalLength);
62
  let offset = 0;
63
  for (const row of sorted) {
64
    result.set(row.data, offset);
65
    offset += row.data.byteLength;
66
  }
67
  return result;
68
}
69
 
70
function toSnapshotChunkView(data: unknown): Uint8Array<ArrayBuffer> {
71
  if (!(data instanceof ArrayBuffer)) {
72
    throw new TypeError("Persisted snapshot chunk must be binary");
73
  }
74
  return new Uint8Array(data);
75
}
76
 
77
export class DocSync extends YServer<Cloudflare.Env> {
78
  static options = { hibernate: true };
79
 
80
  static callbackOptions = {
81
    debounceWait: 2000,
82
    debounceMaxWait: 10000,
83
  };
84
 
85
  private readonly doDb: DocSyncDb;
86
 
87
  constructor(ctx: DurableObjectState, env: Cloudflare.Env) {
88
    super(ctx, env);
89
    this.doDb = drizzle(ctx.storage, { schema: docSyncSchema });
90
 
91
    ctx.blockConcurrencyWhile(async () => {
92
      await migrate(this.doDb, docSyncMigrations);
93
    });
94
  }
95
 
96
  private get d1Db() {
97
    return createSessionDb(this.env.DB, "first-primary").db;
98
  }
99
 
100
  onClose(connection: Connection, code: number, reason: string, wasClean: boolean) {
101
    setLevel(this.env.LOG_LEVEL);
102
    log.debug("connection_closed", { pageId: this.name });
103
    // After hibernation wake-up, awareness state is empty (in-memory only).
104
    // Ensure stub entries exist so removeAwarenessStates in super.onClose
105
    // can find them, triggering the broadcast to remaining clients.
106
    const ids = getAwarenessIds(connection);
107
    const awareness = this.document.awareness;
108
    for (const id of ids) {
109
      if (!awareness.states.has(id)) {
110
        awareness.states.set(id, {});
111
      }
112
      if (!awareness.meta.has(id)) {
113
        awareness.meta.set(id, { clock: Date.now(), lastUpdated: 0 });
114
      }
115
    }
116
    super.onClose(connection, code, reason, wasClean);
117
  }
118
 
119
  getConnectionTags(_connection: Connection, ctx: ConnectionContext): string[] {
120
    const url = new URL(ctx.request.url);
121
    const tags: string[] = [];
122
    if (url.searchParams.get("readOnly") === "1") {
123
      log.debug("connection_readonly", { pageId: this.name });
124
      tags.push(READONLY_TAG);
125
    }
126
    if (url.searchParams.get("authType") === "member_edit") {
127
      tags.push(MEMBER_EDIT_TAG);
128
    }
129
    return tags;
130
  }
131
 
132
  async onConnect(connection: Connection, ctx: ConnectionContext): Promise<void> {
133
    const isMemberEdit = connection.tags.includes(MEMBER_EDIT_TAG);
134
 
135
    // Member-edit connections are guaranteed — only cap headroom connections
136
    if (!isMemberEdit) {
137
      let total = 0;
138
      let memberEditCount = 0;
139
      for (const conn of this.getConnections()) {
140
        total++;
141
        if (conn.tags.includes(MEMBER_EDIT_TAG)) memberEditCount++;
142
      }
143
      const headroomCount = total - memberEditCount;
144
      log.debug("connection_attempt", { pageId: this.name, total, headroomCount, isMemberEdit });
145
      if (headroomCount >= MAX_CONNECTIONS_PER_DOC) {
146
        log.info("connection_rejected", { pageId: this.name, reason: "headroom_full", headroomCount });
147
        connection.close(4029, "Too many concurrent connections");
148
        return;
149
      }
150
    }
151
 
152
    return super.onConnect(connection, ctx);
153
  }
154
 
155
  /** y-partyserver calls this in readSyncMessage to gate syncStep2/update. */
156
  isReadOnly(connection: Connection): boolean {
157
    return connection.tags.includes(READONLY_TAG);
158
  }
159
 
160
  async onCustomMessage(connection: Connection, message: string): Promise<void> {
161
    const msg = parseDocMessage(message);
162
    if (!msg || msg.type !== "page-metadata-refresh") return;
163
 
164
    // Don't allow readonly connections to trigger D1 reads
165
    if (connection.tags.includes(READONLY_TAG)) return;
166
 
167
    try {
168
      const row = await this.d1Db
169
        .select({ icon: pages.icon, cover_url: pages.cover_url })
170
        .from(pages)
171
        .where(eq(pages.id, this.name))
172
        .get();
173
      if (!row) return;
174
 
175
      this.broadcastCustomMessage(
176
        JSON.stringify({
177
          type: "page-metadata-updated",
178
          pageId: this.name,
179
          icon: row.icon,
180
          cover_url: row.cover_url,
181
        }),
182
        connection,
183
      );
184
    } catch (e) {
185
      log.error("metadata_refresh_failed", errorContext(e));
186
    }
187
  }
188
 
189
  async onLoad(): Promise<Y.Doc | void> {
190
    setLevel(this.env.LOG_LEVEL);
191
 
192
    const chunkRows = await this.doDb
193
      .select({ chunk_index: docSyncSchema.snapshotChunks.chunk_index, data: docSyncSchema.snapshotChunks.data })
194
      .from(docSyncSchema.snapshotChunks)
195
      .orderBy(docSyncSchema.snapshotChunks.chunk_index);
196
 
197
    const state = reassembleChunks(chunkRows);
198
    log.debug("snapshot_loaded", { pageId: this.name, hasSnapshot: !!state });
199
 
200
    if (state) {
201
      const doc = new Y.Doc();
202
      Y.applyUpdate(doc, state);
203
      return doc;
204
    }
205
  }
206
 
207
  async onSave(): Promise<void> {
208
    const state = Y.encodeStateAsUpdate(this.document);
209
    const dl = log.child({ pageId: this.name });
210
 
211
    const chunks = chunkBuffer(state, CHUNK_SIZE);
212
    const now = new Date().toISOString();
213
 
214
    // Persist to DO-local SQLite in a single transaction
215
    try {
216
      this.doDb.transaction((tx) => {
217
        tx.delete(docSyncSchema.snapshotChunks).run();
218
        for (let i = 0; i < chunks.length; i++) {
219
          tx.insert(docSyncSchema.snapshotChunks).values({ chunk_index: i, data: chunks[i] }).run();
220
        }
221
        tx.insert(docSyncSchema.snapshotMeta)
222
          .values({ id: 1, chunk_count: chunks.length, total_bytes: state.byteLength, snapshot_at: now })
223
          .onConflictDoUpdate({
224
            target: docSyncSchema.snapshotMeta.id,
225
            set: { chunk_count: chunks.length, total_bytes: state.byteLength, snapshot_at: now },
226
          })
227
          .run();
228
      });
229
      dl.debug("snapshot_saved", { sizeBytes: state.byteLength, chunks: chunks.length });
230
    } catch (e) {
231
      dl.error("snapshot_save_failed", errorContext(e));
232
    }
233
 
234
    // Mirror DocSync save metadata into D1 for page lists and public Sites freshness.
235
    const title = this.document.getText(YJS_PAGE_TITLE).toString() || DEFAULT_PAGE_TITLE;
236
    try {
237
      await recordDocSyncPageSave(this.d1Db, this.name, title, now);
238
    } catch (e) {
239
      dl.error("page_save_sync_failed", errorContext(e));
240
    }
241
 
242
    // Derived indexing/projection must not break snapshot persistence (spec S7)
243
    try {
244
      await this.env.TASKS_QUEUE.sendBatch([
245
        { body: { type: "index-page", pageId: this.name } },
246
        { body: { type: "page-projection", pageId: this.name } },
247
      ]);
248
    } catch (e) {
249
      dl.error("queue_send_failed", errorContext(e));
250
    }
251
  }
252
 
253
  /**
254
   * RPC method for Worker to extract indexable text from the persisted snapshot.
255
   * Does NOT use this.document or this.name (RPC bypasses partyserver init).
256
   */
257
  async getIndexPayload(
258
    pageId: string,
259
    pageKind: PageKind = "doc",
260
  ): Promise<{ kind: "found"; title: string; bodyText: string } | { kind: "missing" }> {
261
    const chunkRows = await this.doDb
262
      .select({ chunk_index: docSyncSchema.snapshotChunks.chunk_index, data: docSyncSchema.snapshotChunks.data })
263
      .from(docSyncSchema.snapshotChunks)
264
      .orderBy(docSyncSchema.snapshotChunks.chunk_index);
265
 
266
    const state = reassembleChunks(chunkRows);
267
    if (!state) {
268
      log.debug("index_payload_missing", { pageId });
269
      return { kind: "missing" };
270
    }
271
 
272
    const ydoc = new Y.Doc();
273
    try {
274
      Y.applyUpdate(ydoc, state);
275
      const { title, bodyText } = pageKind === "canvas" ? extractCanvasPlaintext(ydoc) : extractPlaintext(ydoc);
276
      return { kind: "found", title, bodyText };
277
    } finally {
278
      ydoc.destroy();
279
    }
280
  }
281
 
282
  async getSnapshotResponse(pageId: string): Promise<{ kind: "found"; response: Response } | { kind: "missing" }> {
283
    // This RPC exists for cold editor bootstrap: uncached clients need the
284
    // persisted Yjs body before mounting a live editor, otherwise an empty
285
    // local Y.Doc can race with authoritative content on first sync. Stream
286
    // snapshot rows directly from DO SQLite so the Worker<->DO hop never has
287
    // to materialize the whole Yjs blob in memory. Drizzle does not expose an
288
    // iterator over durable-sqlite query results yet, so this path uses the
289
    // underlying sql.exec cursor on purpose.
290
    const cursor = this.ctx.storage.sql.exec<{ data: ArrayBuffer }>(
291
      "SELECT data FROM snapshot_chunks ORDER BY chunk_index",
292
    );
293
    const iterator = cursor[Symbol.iterator]();
294
    const first = iterator.next();
295
 
296
    if (first.done) {
297
      log.debug("snapshot_response_missing", { pageId });
298
      return { kind: "missing" };
299
    }
300
 
301
    let nextChunk: Uint8Array<ArrayBuffer> | null = toSnapshotChunkView(first.value.data);
302
 
303
    return {
304
      kind: "found",
305
      response: new Response(
306
        new ReadableStream({
307
          type: "bytes",
308
          pull(controller) {
309
            if (nextChunk) {
310
              controller.enqueue(nextChunk);
311
              nextChunk = null;
312
              return;
313
            }
314
 
315
            const next = iterator.next();
316
            if (next.done) {
317
              controller.close();
318
              return;
319
            }
320
 
321
            controller.enqueue(toSnapshotChunkView(next.value.data));
322
          },
323
        }),
324
        {
325
          headers: {
326
            "Content-Type": "application/octet-stream",
327
          },
328
        },
329
      ),
330
    };
331
  }
332
}