gitmost/apps/server/src/core/comment/comment.service.ts

import {
  BadRequestException,
  ForbiddenException,
  Injectable,
  Logger,
  NotFoundException,
} from '@nestjs/common';
import { InjectQueue } from '@nestjs/bullmq';
import { Queue } from 'bullmq';
import { CreateCommentDto, yjsSelectionSchema } from './dto/create-comment.dto';
import { CollaborationGateway } from '../../collaboration/collaboration.gateway';
import { UpdateCommentDto } from './dto/update-comment.dto';
import { CommentRepo } from '@docmost/db/repos/comment/comment.repo';
import { Comment, Page, User } from '@docmost/db/types/entity.types';
import { PaginationOptions } from '@docmost/db/pagination/pagination-options';
import { PageRepo } from '@docmost/db/repos/page/page.repo';
import { CursorPaginationResult } from '@docmost/db/pagination/cursor-pagination';
import { QueueJob, QueueName } from '../../integrations/queue/constants';
import { extractUserMentionIdsFromJson } from '../../common/helpers/prosemirror/utils';
import {
  ICommentNotificationJob,
  ICommentResolvedNotificationJob,
} from '../../integrations/queue/constants/queue.interface';
import { WsService } from '../../ws/ws.service';
import {
  AuthProvenanceData,
  agentSourceFields,
} from '../../common/decorators/auth-provenance.decorator';

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

  constructor(
    private commentRepo: CommentRepo,
    private pageRepo: PageRepo,
    private wsService: WsService,
    private collaborationGateway: CollaborationGateway,
    @InjectQueue(QueueName.GENERAL_QUEUE)
    private generalQueue: Queue,
    @InjectQueue(QueueName.NOTIFICATION_QUEUE)
    private notificationQueue: Queue,
  ) {}

  async findById(commentId: string) {
    const comment = await this.commentRepo.findById(commentId, {
      includeCreator: true,
      includeResolvedBy: true,
    });
    if (!comment) {
      throw new NotFoundException('Comment not found');
    }
    return comment;
  }

  async create(
    opts: { page: Page; workspaceId: string; user: User },
    createCommentDto: CreateCommentDto,
    // Optional agent-edit provenance (from the signed access claim). When the
    // actor is 'agent', stamp created_source/ai_chat_id so an agent-authored
    // comment (incl. a reply) shows the AI marker (§15 C3). Normal user: default.
    provenance?: AuthProvenanceData,
  ) {
    const { page, workspaceId, user } = opts;
    const commentContent = JSON.parse(createCommentDto.content);

    if (createCommentDto.parentCommentId) {
      const parentComment = await this.commentRepo.findById(
        createCommentDto.parentCommentId,
      );

      if (!parentComment || parentComment.pageId !== page.id) {
        throw new BadRequestException('Parent comment not found');
      }

      if (parentComment.parentCommentId !== null) {
        throw new BadRequestException('You cannot reply to a reply');
      }
    }

    const inserted = await this.commentRepo.insertComment({
      pageId: page.id,
      content: commentContent,
      selection: createCommentDto?.selection?.substring(0, 250) ?? null,
      type: createCommentDto.type ?? 'page',
      parentCommentId: createCommentDto?.parentCommentId,
      creatorId: user.id,
      workspaceId: workspaceId,
      spaceId: page.spaceId,
      // Agent-edit provenance: the user stays creatorId; this only annotates the
      // source. Normal user requests leave the column default ('user').
      ...agentSourceFields(provenance, 'createdSource', 'aiChatId'),
    });

    if (createCommentDto.yjsSelection) {
      const parsed = yjsSelectionSchema.safeParse(createCommentDto.yjsSelection);
      if (!parsed.success) {
        this.logger.warn(
          `Invalid yjsSelection for comment ${inserted.id}: ${parsed.error.message}`,
        );
      } else {
        const documentName = `page.${page.id}`;
        try {
          await this.collaborationGateway.handleYjsEvent(
            'setCommentMark',
            documentName,
            {
              yjsSelection: parsed.data,
              commentId: inserted.id,
              resolved: false,
              user,
            },
          );
        } catch (error) {
          this.logger.warn(
            `Failed to apply comment mark for comment ${inserted.id}, comment saved without inline highlight`,
            error,
          );
        }
      }
    }

    const comment = await this.commentRepo.findById(inserted.id, {
      includeCreator: true,
      includeResolvedBy: true,
    });

    this.generalQueue
      .add(QueueJob.ADD_PAGE_WATCHERS, {
        userIds: [user.id],
        pageId: page.id,
        spaceId: page.spaceId,
        workspaceId,
      })
      .catch((err) =>
        this.logger.warn(`Failed to queue add-page-watchers: ${err.message}`),
      );

    const isReply = !!createCommentDto.parentCommentId;

    await this.queueCommentNotification(
      commentContent,
      [],
      comment.id,
      page.id,
      page.spaceId,
      workspaceId,
      user.id,
      !isReply,
      createCommentDto.parentCommentId,
    );

    this.wsService.emitCommentEvent(page.spaceId, page.id, {
      operation: 'commentCreated',
      pageId: page.id,
      comment,
    });

    return comment;
  }

  async findByPageId(
    pageId: string,
    pagination: PaginationOptions,
  ): Promise<CursorPaginationResult<Comment>> {
    const page = await this.pageRepo.findById(pageId);

    if (!page) {
      throw new BadRequestException('Page not found');
    }

    return this.commentRepo.findPageComments(pageId, pagination);
  }

  async update(
    comment: Comment,
    updateCommentDto: UpdateCommentDto,
    authUser: User,
  ): Promise<Comment> {
    const commentContent = JSON.parse(updateCommentDto.content);

    if (comment.creatorId !== authUser.id) {
      throw new ForbiddenException('You can only edit your own comments');
    }

    const oldMentionIds = extractUserMentionIdsFromJson(comment.content);

    const editedAt = new Date();

    await this.commentRepo.updateComment(
      {
        content: commentContent,
        editedAt: editedAt,
        updatedAt: editedAt,
      },
      comment.id,
    );

    await this.queueCommentNotification(
      commentContent,
      oldMentionIds,
      comment.id,
      comment.pageId,
      comment.spaceId,
      comment.workspaceId,
      authUser.id,
      false,
    );

    // Re-fetch the enriched comment before broadcasting, symmetric with
    // create()/resolveComment(). updateComment() above has already persisted the
    // new content/timestamps, so this single-row read reflects the edit AND
    // carries the same {agent,launcher} avatar stack (via includeCreator) as the
    // other two broadcasts. This deliberately does NOT reuse the caller's
    // pre-loaded `comment`: relying on the controller happening to load it with
    // includeCreator:true is exactly the fragile coupling that let the agent
    // stack silently vanish on edit once already (#300/#304) — a future caller
    // dropping that flag must not regress the broadcast.
    const updatedComment = await this.commentRepo.findById(comment.id, {
      includeCreator: true,
      includeResolvedBy: true,
    });

    this.wsService.emitCommentEvent(comment.spaceId, comment.pageId, {
      operation: 'commentUpdated',
      pageId: comment.pageId,
      comment: updatedComment,
    });

    return updatedComment;
  }

  async resolveComment(
    comment: Comment,
    resolved: boolean,
    authUser: User,
    // Optional agent-edit provenance (from the signed access claim). When the
    // actor is 'agent' and the thread is being resolved, stamp resolved_source
    // so the "resolved by" mark shows the AI marker (§15 C3). On unresolve the
    // source is cleared alongside resolvedAt/resolvedById.
    provenance?: AuthProvenanceData,
  ): Promise<Comment> {
    const resolvedAt = resolved ? new Date() : null;
    const resolvedById = resolved ? authUser.id : null;
    const isAgent = provenance?.actor === 'agent';
    // Set the agent marker only when resolving; on unresolve clear it back to
    // null so a reopened thread carries no stale source. A normal user resolve
    // leaves resolved_source null (no agent annotation).
    const resolvedSource = resolved && isAgent ? 'agent' : null;

    await this.commentRepo.updateComment(
      { resolvedAt, resolvedById, resolvedSource },
      comment.id,
    );

    // Reflect the resolved state on the inline comment mark in the
    // collaborative document so all connected clients stay in sync.
    const documentName = `page.${comment.pageId}`;
    try {
      await this.collaborationGateway.handleYjsEvent(
        'resolveCommentMark',
        documentName,
        { commentId: comment.id, resolved, user: authUser },
      );
    } catch (error) {
      this.logger.warn(
        `Failed to update comment mark for comment ${comment.id}`,
        error,
      );
    }

    // Notify the comment author when someone else resolves their comment.
    if (resolved && comment.creatorId !== authUser.id) {
      const jobData: ICommentResolvedNotificationJob = {
        commentId: comment.id,
        commentCreatorId: comment.creatorId,
        pageId: comment.pageId,
        spaceId: comment.spaceId,
        workspaceId: comment.workspaceId,
        actorId: authUser.id,
      };
      await this.notificationQueue.add(
        QueueJob.COMMENT_RESOLVED_NOTIFICATION,
        jobData,
      );
    }

    const updatedComment = await this.commentRepo.findById(comment.id, {
      includeCreator: true,
      includeResolvedBy: true,
    });

    this.wsService.emitCommentEvent(comment.spaceId, comment.pageId, {
      operation: 'commentResolved',
      pageId: comment.pageId,
      comment: updatedComment,
    });

    return updatedComment;
  }

  private async queueCommentNotification(
    content: any,
    oldMentionIds: string[],
    commentId: string,
    pageId: string,
    spaceId: string,
    workspaceId: string,
    actorId: string,
    notifyWatchers: boolean,
    parentCommentId?: string,
  ) {
    const mentionedUserIds = extractUserMentionIdsFromJson(content);
    const newMentionIds = mentionedUserIds.filter(
      (id) => id !== actorId && !oldMentionIds.includes(id),
    );

    if (newMentionIds.length === 0 && !notifyWatchers && !parentCommentId) return;

    const jobData: ICommentNotificationJob = {
      commentId,
      parentCommentId,
      pageId,
      spaceId,
      workspaceId,
      actorId,
      mentionedUserIds: newMentionIds,
      notifyWatchers,
    };

    await this.notificationQueue.add(
      QueueJob.COMMENT_NOTIFICATION,
      jobData,
    );
  }
}