gitmost/apps/server/src/collaboration/collaboration.handler.ts

import { Injectable, Logger } from '@nestjs/common';
import { Hocuspocus, Document } from '@hocuspocus/server';
import { TiptapTransformer } from '@hocuspocus/transformer';
import {
  prosemirrorNodeToYElement,
  tiptapExtensions,
} from './collaboration.util';
import { setYjsMark, updateYjsMarkAttribute, YjsSelection } from './yjs.util';
import * as Y from 'yjs';
import { User } from '@docmost/db/types/entity.types';
import {
  hasHtmlEmbedNode,
  htmlEmbedAllowed,
  isHtmlEmbedFeatureEnabled,
  stripHtmlEmbedNodes,
} from '../common/helpers/prosemirror/html-embed.util';
import { WorkspaceRepo } from '@docmost/db/repos/workspace/workspace.repo';

export type CollabEventHandlers = ReturnType<
  CollaborationHandler['getHandlers']
>;

@Injectable()
export class CollaborationHandler {
  private readonly logger = new Logger(CollaborationHandler.name);

  constructor(private readonly workspaceRepo: WorkspaceRepo) {}

  getHandlers(hocuspocus: Hocuspocus) {
    return {
      alterState: async (documentName: string, payload: { pageId: string }) => {
        // dummy
        // this.logger.log('Processing', documentName, payload);
        // await this.withYdocConnection(hocuspocus, documentName, {}, (doc) => {
        //   const fragment = doc.getXmlFragment('default');
        //});
      },
      setCommentMark: async (
        documentName: string,
        payload: {
          yjsSelection: YjsSelection;
          commentId: string;
          resolved: boolean;
          user: User;
        },
      ) => {
        const { yjsSelection, commentId, resolved, user } = payload;
        await this.withYdocConnection(
          hocuspocus,
          documentName,
          { user },
          (doc) => {
            const fragment = doc.getXmlFragment('default');
            setYjsMark(doc, fragment, yjsSelection, 'comment', {
              commentId,
              resolved,
            });
          },
        );
      },
      resolveCommentMark: async (
        documentName: string,
        payload: {
          commentId: string;
          resolved: boolean;
          user: User;
        },
      ) => {
        const { commentId, resolved, user } = payload;
        await this.withYdocConnection(
          hocuspocus,
          documentName,
          { user },
          (doc) => {
            const fragment = doc.getXmlFragment('default');
            updateYjsMarkAttribute(
              fragment,
              'comment',
              { name: 'commentId', value: commentId },
              { resolved },
            );
          },
        );
      },
      updatePageContent: async (
        documentName: string,
        payload: {
          prosemirrorJson: any;
          operation: string;
          user: User;
        },
      ) => {
        const { operation, user } = payload;
        let { prosemirrorJson } = payload;
        this.logger.debug('Updating page content via yjs', documentName);

        // SECURITY (Variant C admin gate, REST/MCP/AI write path):
        // updatePageContent is the server-side entrypoint used by the REST page
        // update endpoint and by the MCP/AI agent. Raw `htmlEmbed` nodes execute
        // arbitrary JS in every reader's browser, so a NON-admin caller must not
        // be able to persist them here. If the editing user is not a workspace
        // admin/owner, strip every htmlEmbed node before it reaches the ydoc.
        // Toggle-AND-admin gate: htmlEmbed survives only when the workspace
        // feature toggle is ON and the editing user is an admin/owner. OFF
        // (default) => stripped for everyone.
        const htmlEmbedEnabled = isHtmlEmbedFeatureEnabled(
          (await this.workspaceRepo.findById(user?.workspaceId))?.settings,
        );
        if (!htmlEmbedAllowed(htmlEmbedEnabled, user?.role)) {
          if (hasHtmlEmbedNode(prosemirrorJson)) {
            this.logger.warn(
              `Stripping htmlEmbed node(s) from update by user ${user?.id} on ${documentName}`,
            );
            prosemirrorJson = stripHtmlEmbedNodes(prosemirrorJson);
          }
        }

        await this.withYdocConnection(
          hocuspocus,
          documentName,
          { user },
          (doc) => {
            const fragment = doc.getXmlFragment('default');

            if (operation === 'replace') {
              if (fragment.length > 0) {
                fragment.delete(0, fragment.length);
              }

              const newDoc = TiptapTransformer.toYdoc(
                prosemirrorJson,
                'default',
                tiptapExtensions,
              );
              Y.applyUpdate(doc, Y.encodeStateAsUpdate(newDoc));
            } else {
              const newContent = prosemirrorJson.content || [];
              const yElements = newContent.map(prosemirrorNodeToYElement);
              const position = operation === 'prepend' ? 0 : fragment.length;
              fragment.insert(position, yElements);
            }
          },
        );
      },
    };
  }

  async withYdocConnection(
    hocuspocus: Hocuspocus,
    documentName: string,
    context: any = {},
    fn: (doc: Document) => void,
  ): Promise<void> {
    const connection = await hocuspocus.openDirectConnection(
      documentName,
      context,
    );
    try {
      await connection.transact(fn);
    } finally {
      await connection.disconnect();
    }
  }
}