mirror of
https://github.com/open-webui/open-webui.git
synced 2025-12-15 05:45:19 +00:00
337 lines
9.2 KiB
TypeScript
337 lines
9.2 KiB
TypeScript
import * as Y from 'yjs';
|
|
import {
|
|
ySyncPlugin,
|
|
yCursorPlugin,
|
|
yUndoPlugin,
|
|
undo,
|
|
redo,
|
|
prosemirrorJSONToYDoc
|
|
} from 'y-prosemirror';
|
|
import type { Socket } from 'socket.io-client';
|
|
import type { Awareness } from 'y-protocols/awareness';
|
|
import type { SessionUser } from '$lib/stores';
|
|
import { Editor, Extension } from '@tiptap/core';
|
|
import { keymap } from 'prosemirror-keymap';
|
|
import { tick } from 'svelte';
|
|
|
|
const USER_COLORS = [
|
|
'#FF6B6B',
|
|
'#4ECDC4',
|
|
'#45B7D1',
|
|
'#96CEB4',
|
|
'#FFEAA7',
|
|
'#DDA0DD',
|
|
'#98D8C8',
|
|
'#F7DC6F',
|
|
'#BB8FCE',
|
|
'#85C1E9'
|
|
];
|
|
const generateUserColor = () => {
|
|
return USER_COLORS[Math.floor(Math.random() * USER_COLORS.length)];
|
|
};
|
|
|
|
export type EditorContentGetter = () => {
|
|
md: string;
|
|
html: string;
|
|
json: string;
|
|
};
|
|
|
|
// Custom Yjs Socket.IO provider
|
|
export class SocketIOCollaborationProvider {
|
|
private readonly doc = new Y.Doc();
|
|
private readonly awareness = new SimpleAwareness(this.doc);
|
|
private isConnected = false;
|
|
private synced = false;
|
|
private editor: Editor | null = null;
|
|
private editorContentGetter: EditorContentGetter | null = null;
|
|
|
|
constructor(
|
|
private readonly documentId: string,
|
|
private readonly socket: Socket,
|
|
private readonly user: SessionUser,
|
|
private readonly initialContent: string | null = null
|
|
) {
|
|
this.setupEventListeners();
|
|
}
|
|
|
|
public getEditorExtension() {
|
|
return Extension.create({
|
|
name: 'yjsCollaboration',
|
|
|
|
addProseMirrorPlugins: () => {
|
|
const yXmlFragment = this.doc.getXmlFragment('prosemirror');
|
|
if (!yXmlFragment) return [];
|
|
|
|
const plugins = [
|
|
ySyncPlugin(yXmlFragment),
|
|
yUndoPlugin(),
|
|
keymap({
|
|
'Mod-z': undo,
|
|
'Mod-y': redo,
|
|
'Mod-Shift-z': redo
|
|
})
|
|
];
|
|
|
|
plugins.push(yCursorPlugin(this.awareness as unknown as Awareness));
|
|
|
|
return plugins;
|
|
}
|
|
});
|
|
}
|
|
|
|
public setEditor(editor: Editor, editorContentGetter: EditorContentGetter) {
|
|
this.editor = editor;
|
|
this.editorContentGetter = editorContentGetter;
|
|
}
|
|
|
|
private joinDocument() {
|
|
const userColor = generateUserColor();
|
|
this.socket.emit('ydoc:document:join', {
|
|
document_id: this.documentId,
|
|
user_id: this.user?.id,
|
|
user_name: this.user?.name,
|
|
user_color: userColor
|
|
});
|
|
|
|
// Set user awareness info
|
|
if (this.user) {
|
|
this.awareness.setLocalStateField('user', {
|
|
name: `${this.user.name}`,
|
|
color: userColor,
|
|
id: this.socket.id
|
|
});
|
|
}
|
|
}
|
|
|
|
private setupEventListeners() {
|
|
// Listen for document updates from server
|
|
this.socket.on('ydoc:document:update', (data) => {
|
|
if (data.document_id === this.documentId && data.socket_id !== this.socket.id) {
|
|
try {
|
|
const update = new Uint8Array(data.update);
|
|
Y.applyUpdate(this.doc, update);
|
|
} catch (error) {
|
|
console.error('Error applying Yjs update:', error);
|
|
}
|
|
}
|
|
});
|
|
|
|
// Listen for document state from server
|
|
this.socket.on('ydoc:document:state', async (data) => {
|
|
if (data.document_id === this.documentId) {
|
|
try {
|
|
if (data.state) {
|
|
const state = new Uint8Array(data.state);
|
|
|
|
if (state.length === 2 && state[0] === 0 && state[1] === 0) {
|
|
// Empty state, check if we have content to initialize
|
|
// check if editor empty as well
|
|
// const editor = await getEditorInstance();
|
|
|
|
const isEmptyEditor = !this.editor?.getText().trim();
|
|
if (isEmptyEditor && this.editor) {
|
|
if (this.initialContent && (data?.sessions ?? ['']).length === 1) {
|
|
const editorYdoc = prosemirrorJSONToYDoc(this.editor.schema, this.initialContent);
|
|
if (editorYdoc) {
|
|
Y.applyUpdate(this.doc, Y.encodeStateAsUpdate(editorYdoc));
|
|
}
|
|
}
|
|
} else {
|
|
// If the editor already has content, we don't need to send an empty state
|
|
if (this.doc.getXmlFragment('prosemirror').length > 0) {
|
|
this.socket.emit('ydoc:document:update', {
|
|
document_id: this.documentId,
|
|
user_id: this.user?.id,
|
|
socket_id: this.socket.id,
|
|
update: Y.encodeStateAsUpdate(this.doc)
|
|
});
|
|
} else {
|
|
console.warn('Yjs document is empty, not sending state.');
|
|
}
|
|
}
|
|
} else {
|
|
Y.applyUpdate(this.doc, state, 'server');
|
|
}
|
|
}
|
|
this.synced = true;
|
|
} catch (error) {
|
|
console.error('Error applying Yjs state:', error);
|
|
|
|
this.synced = false;
|
|
this.socket.emit('ydoc:document:state', {
|
|
document_id: this.documentId
|
|
});
|
|
}
|
|
}
|
|
});
|
|
|
|
// Listen for awareness updates
|
|
this.socket.on('ydoc:awareness:update', (data) => {
|
|
if (data.document_id === this.documentId) {
|
|
try {
|
|
const awarenessUpdate = new Uint8Array(data.update);
|
|
this.awareness.applyUpdate(awarenessUpdate, 'server');
|
|
} catch (error) {
|
|
console.error('Error applying awareness update:', error);
|
|
}
|
|
}
|
|
});
|
|
|
|
// Handle connection events
|
|
this.socket.on('connect', this.onConnect);
|
|
this.socket.on('disconnect', this.onDisconnect);
|
|
|
|
// Listen for document updates from Yjs
|
|
this.doc.on('update', async (update, origin) => {
|
|
if (this.editor && origin !== 'server' && this.isConnected) {
|
|
await tick(); // Ensure the DOM is updated before sending
|
|
this.socket.emit('ydoc:document:update', {
|
|
document_id: this.documentId,
|
|
user_id: this.user?.id,
|
|
socket_id: this.socket.id,
|
|
update: Array.from(update),
|
|
data: {
|
|
content: this.editorContentGetter?.() ?? {
|
|
md: '',
|
|
html: '',
|
|
json: ''
|
|
}
|
|
}
|
|
});
|
|
}
|
|
});
|
|
|
|
// Listen for awareness updates from Yjs
|
|
this.awareness.on(
|
|
'change',
|
|
(
|
|
{ added, updated, removed }: { added: number[]; updated: number[]; removed: number[] },
|
|
origin: string
|
|
) => {
|
|
if (origin !== 'server' && this.isConnected) {
|
|
const changedClients = added.concat(updated).concat(removed);
|
|
const awarenessUpdate = this.awareness.encodeUpdate(changedClients);
|
|
this.socket.emit('ydoc:awareness:update', {
|
|
document_id: this.documentId,
|
|
user_id: this.socket.id,
|
|
update: Array.from(awarenessUpdate)
|
|
});
|
|
}
|
|
}
|
|
);
|
|
|
|
if (this.socket.connected) {
|
|
this.isConnected = true;
|
|
this.joinDocument();
|
|
}
|
|
}
|
|
|
|
private readonly onConnect = () => {
|
|
this.isConnected = true;
|
|
this.joinDocument();
|
|
};
|
|
|
|
private readonly onDisconnect = () => {
|
|
this.isConnected = false;
|
|
this.synced = false;
|
|
};
|
|
|
|
public destroy() {
|
|
this.socket.off('ydoc:document:update');
|
|
this.socket.off('ydoc:document:state');
|
|
this.socket.off('ydoc:awareness:update');
|
|
this.socket.off('connect', this.onConnect);
|
|
this.socket.off('disconnect', this.onDisconnect);
|
|
|
|
if (this.isConnected) {
|
|
this.socket.emit('ydoc:document:leave', {
|
|
document_id: this.documentId,
|
|
user_id: this.user?.id
|
|
});
|
|
}
|
|
|
|
this.editor = null;
|
|
this.editorContentGetter = null;
|
|
}
|
|
}
|
|
|
|
// Simple awareness implementation
|
|
class SimpleAwareness {
|
|
public readonly clientID: number;
|
|
private readonly _states: Map<number, any>;
|
|
private readonly _updateHandlers: any[];
|
|
private readonly _localState: any;
|
|
|
|
public constructor(public readonly doc: Y.Doc) {
|
|
// Yjs awareness expects clientID (not clientId) property
|
|
this.clientID = doc.clientID ? doc.clientID : Math.floor(Math.random() * 0xffffffff);
|
|
// Map from clientID (number) to state (object)
|
|
this._states = new Map(); // _states, not states; will make getStates() for compat
|
|
this._updateHandlers = [];
|
|
this._localState = {};
|
|
// As in Yjs Awareness, add our local state to the states map from the start:
|
|
this._states.set(this.clientID, this._localState);
|
|
}
|
|
|
|
public on(event: string, handler: any) {
|
|
if (event === 'change') this._updateHandlers.push(handler);
|
|
}
|
|
|
|
public off(event: string, handler: any) {
|
|
if (event === 'change') {
|
|
const i = this._updateHandlers.indexOf(handler);
|
|
if (i !== -1) this._updateHandlers.splice(i, 1);
|
|
}
|
|
}
|
|
|
|
public getLocalState() {
|
|
return this._states.get(this.clientID) || null;
|
|
}
|
|
|
|
public getStates() {
|
|
// Yjs returns a Map (clientID->state)
|
|
return this._states;
|
|
}
|
|
|
|
public setLocalStateField(field: string, value: any) {
|
|
let localState = this._states.get(this.clientID);
|
|
if (!localState) {
|
|
localState = {};
|
|
this._states.set(this.clientID, localState);
|
|
}
|
|
localState[field] = value;
|
|
// After updating, fire 'update' event to all handlers
|
|
for (const cb of this._updateHandlers) {
|
|
// Follows Yjs Awareness ({ added, updated, removed }, origin)
|
|
cb({ added: [], updated: [this.clientID], removed: [] }, 'local');
|
|
}
|
|
}
|
|
|
|
public applyUpdate(update: Uint8Array, origin: string) {
|
|
// Very simple: Accepts a serialized JSON state for now as Uint8Array
|
|
try {
|
|
const str = new TextDecoder().decode(update);
|
|
const obj = JSON.parse(str);
|
|
// Should be a plain object: { clientID: state, ... }
|
|
for (const [k, v] of Object.entries(obj)) {
|
|
this._states.set(+k, v);
|
|
}
|
|
for (const cb of this._updateHandlers) {
|
|
cb({ added: [], updated: Array.from(Object.keys(obj)).map(Number), removed: [] }, origin);
|
|
}
|
|
} catch (e) {
|
|
console.warn('SimpleAwareness: Could not decode update:', e);
|
|
}
|
|
}
|
|
|
|
public encodeUpdate(clients: number[]) {
|
|
// Encodes the states for the given clientIDs as Uint8Array (JSON)
|
|
const obj: Record<number, any> = {};
|
|
for (const id of clients || Array.from(this._states.keys())) {
|
|
const st = this._states.get(id);
|
|
if (st) obj[id] = st;
|
|
}
|
|
const json = JSON.stringify(obj);
|
|
return new TextEncoder().encode(json);
|
|
}
|
|
}
|