chore: Use extraction from AST to infer mentions (#38845)

Co-authored-by: Kevin Aleman <kaleman960@gmail.com>

apps/meteor/app/lib/server/functions/extractMentionsFromMessageAST.ts

@@ -0,0 +1,63 @@
+import type { Root, Paragraph, Blocks, Inlines, UserMention, ChannelMention, Task, ListItem, BigEmoji } from '@rocket.chat/message-parser';
+
+type ExtractedMentions = {
+	mentions: string[];
+	channels: string[];
+};
+
+type MessageNode = Paragraph | Blocks | Inlines | Task | ListItem | BigEmoji;
+
+function isUserMention(node: MessageNode): node is UserMention {
+	return node.type === 'MENTION_USER';
+}
+
+function isChannelMention(node: MessageNode): node is ChannelMention {
+	return node.type === 'MENTION_CHANNEL';
+}
+
+function hasArrayValue(node: MessageNode): node is MessageNode & { value: MessageNode[] } {
+	return Array.isArray(node.value);
+}
+
+function hasObjectValue(node: MessageNode): node is MessageNode & { value: Record<string, MessageNode> } {
+	return typeof node.value === 'object' && node.value !== null && !Array.isArray(node.value);
+}
+
+function traverse(node: MessageNode, mentions: Set<string>, channels: Set<string>): void {
+	if (isUserMention(node)) {
+		mentions.add(node.value.value);
+		return;
+	}
+
+	if (isChannelMention(node)) {
+		channels.add(node.value.value);
+		return;
+	}
+
+	if (hasArrayValue(node)) {
+		for (const child of node.value) {
+			traverse(child, mentions, channels);
+		}
+		return;
+	}
+
+	if (hasObjectValue(node)) {
+		for (const key of Object.keys(node.value)) {
+			traverse(node.value[key], mentions, channels);
+		}
+	}
+}
+
+export function extractMentionsFromMessageAST(ast: Root): ExtractedMentions {
+	const mentions = new Set<string>();
+	const channels = new Set<string>();
+
+	for (const node of ast) {
+		traverse(node, mentions, channels);
+	}
+
+	return {
+		mentions: Array.from(mentions),
+		channels: Array.from(channels),
+	};
+}

apps/meteor/app/mentions/server/Mentions.ts

@@ -4,6 +4,7 @@
  */
 import { isE2EEMessage, type IMessage, type IRoom, type IUser } from '@rocket.chat/core-typings';
 
+import { extractMentionsFromMessageAST } from '../../lib/server/functions/extractMentionsFromMessageAST';
 import { type MentionsParserArgs, MentionsParser } from '../lib/MentionsParser';
 
 type MentionsServerArgs = MentionsParserArgs & {
@@ -50,13 +51,25 @@ export class MentionsServer extends MentionsParser {
 			isE2EEMessage(message) && e2eMentions?.e2eUserMentions && e2eMentions?.e2eUserMentions.length > 0
 				? e2eMentions?.e2eUserMentions
 				: this.getUserMentions(msg);
+
+		return this.convertMentionsToUsers(mentions, rid, sender);
+	}
+
+	async convertMentionsToUsers(mentions: string[], rid: string, sender: IMessage['u']): Promise<IMessage['mentions']> {
 		const mentionsAll: { _id: string; username: string }[] = [];
-		const userMentions = [];
+		const userMentions = new Set<string>();
 
 		for await (const m of mentions) {
-			const mention = m.includes(':') ? m.trim() : m.trim().substring(1);
+			let mention: string;
+			if (m.includes(':')) {
+				mention = m.trim();
+			} else if (m.startsWith('@')) {
+				mention = m.substring(1);
+			} else {
+				mention = m;
+			}
 			if (mention !== 'all' && mention !== 'here') {
-				userMentions.push(mention);
+				userMentions.add(mention);
 				continue;
 			}
 			if (this.messageMaxAll() > 0 && (await this.getTotalChannelMembers(rid)) > this.messageMaxAll()) {
@@ -69,7 +82,7 @@ export class MentionsServer extends MentionsParser {
 			});
 		}
 
-		return [...mentionsAll, ...(userMentions.length ? await this.getUsers(userMentions) : [])];
+		return [...mentionsAll, ...(userMentions.size ? await this.getUsers(Array.from(userMentions)) : [])];
 	}
 
 	async getChannelbyMentions(message: IMessage) {
@@ -79,15 +92,23 @@ export class MentionsServer extends MentionsParser {
 			isE2EEMessage(message) && e2eMentions?.e2eChannelMentions && e2eMentions?.e2eChannelMentions.length > 0
 				? e2eMentions?.e2eChannelMentions
 				: this.getChannelMentions(msg);
-		return this.getChannels(channels.map((c) => c.trim().substring(1)));
+		return this.convertMentionsToChannels(channels);
+	}
+
+	async convertMentionsToChannels(channels: string[]): Promise<Pick<IRoom, '_id' | 'name' | 'fname' | 'federated'>[]> {
+		return this.getChannels(channels.map((c) => (c.startsWith('#') ? c.substring(1) : c)));
 	}
 
 	async execute(message: IMessage) {
-		const mentionsAll = await this.getUsersByMentions(message);
-		const channels = await this.getChannelbyMentions(message);
+		if (message.md) {
+			const { mentions, channels } = extractMentionsFromMessageAST(message.md);
+			message.mentions = await this.convertMentionsToUsers(mentions, message.rid, message.u);
+			message.channels = await this.convertMentionsToChannels(channels);
+			return message;
+		}
 
-		message.mentions = mentionsAll;
-		message.channels = channels;
+		message.mentions = await this.getUsersByMentions(message);
+		message.channels = await this.getChannelbyMentions(message);
 
 		return message;
 	}

apps/meteor/server/services/messages/service.ts

@@ -234,10 +234,10 @@ export class MessageService extends ServiceClassInternal implements IMessageServ
 			throw new FederationMatrixInvalidConfigurationError('Unable to send message');
 		}
 
-		message = await mentionServer.execute(message);
 		message = await this.cannedResponse.replacePlaceholders({ message, room, user });
 		message = await this.badWords.filterBadWords({ message });
 		message = await this.markdownParser.parseMarkdown({ message, config: this.getMarkdownConfig() });
+		message = await mentionServer.execute(message);
 		if (parseUrls) {
 			message.urls = parseUrlsInMessage(message, previewUrls);
 		}

apps/meteor/tests/unit/app/lib/server/functions/extractMentionsFromMessageAST.spec.ts

@@ -0,0 +1,396 @@
+import type { Root } from '@rocket.chat/message-parser';
+import { expect } from 'chai';
+
+import { extractMentionsFromMessageAST } from '../../../../../../app/lib/server/functions/extractMentionsFromMessageAST';
+
+describe('extractMentionsFromMessageAST', () => {
+	it('should return empty arrays when AST has no mentions', () => {
+		const ast: Root = [
+			{
+				type: 'PARAGRAPH',
+				value: [
+					{
+						type: 'PLAIN_TEXT',
+						value: 'Hello world',
+					},
+				],
+			},
+		];
+
+		const result = extractMentionsFromMessageAST(ast);
+
+		expect(result.mentions).to.be.an('array').that.is.empty;
+		expect(result.channels).to.be.an('array').that.is.empty;
+	});
+
+	it('should extract user mentions from AST', () => {
+		const ast: Root = [
+			{
+				type: 'PARAGRAPH',
+				value: [
+					{
+						type: 'PLAIN_TEXT',
+						value: 'Hello ',
+					},
+					{
+						type: 'MENTION_USER',
+						value: {
+							type: 'PLAIN_TEXT',
+							value: 'john.doe',
+						},
+					},
+				],
+			},
+		];
+
+		const result = extractMentionsFromMessageAST(ast);
+
+		expect(result.mentions).to.deep.equal(['john.doe']);
+		expect(result.channels).to.be.an('array').that.is.empty;
+	});
+
+	it('should extract channel mentions from AST', () => {
+		const ast: Root = [
+			{
+				type: 'PARAGRAPH',
+				value: [
+					{
+						type: 'PLAIN_TEXT',
+						value: 'Check ',
+					},
+					{
+						type: 'MENTION_CHANNEL',
+						value: {
+							type: 'PLAIN_TEXT',
+							value: 'general',
+						},
+					},
+				],
+			},
+		];
+
+		const result = extractMentionsFromMessageAST(ast);
+
+		expect(result.mentions).to.be.an('array').that.is.empty;
+		expect(result.channels).to.deep.equal(['general']);
+	});
+
+	it('should extract both user and channel mentions', () => {
+		const ast: Root = [
+			{
+				type: 'PARAGRAPH',
+				value: [
+					{
+						type: 'MENTION_USER',
+						value: {
+							type: 'PLAIN_TEXT',
+							value: 'admin',
+						},
+					},
+					{
+						type: 'PLAIN_TEXT',
+						value: ' please check ',
+					},
+					{
+						type: 'MENTION_CHANNEL',
+						value: {
+							type: 'PLAIN_TEXT',
+							value: 'support',
+						},
+					},
+				],
+			},
+		];
+
+		const result = extractMentionsFromMessageAST(ast);
+
+		expect(result.mentions).to.deep.equal(['admin']);
+		expect(result.channels).to.deep.equal(['support']);
+	});
+
+	it('should extract multiple user mentions', () => {
+		const ast: Root = [
+			{
+				type: 'PARAGRAPH',
+				value: [
+					{
+						type: 'MENTION_USER',
+						value: {
+							type: 'PLAIN_TEXT',
+							value: 'user1',
+						},
+					},
+					{
+						type: 'PLAIN_TEXT',
+						value: ' and ',
+					},
+					{
+						type: 'MENTION_USER',
+						value: {
+							type: 'PLAIN_TEXT',
+							value: 'user2',
+						},
+					},
+				],
+			},
+		];
+
+		const result = extractMentionsFromMessageAST(ast);
+
+		expect(result.mentions).to.have.members(['user1', 'user2']);
+		expect(result.mentions).to.have.lengthOf(2);
+	});
+
+	it('should deduplicate repeated mentions', () => {
+		const ast: Root = [
+			{
+				type: 'PARAGRAPH',
+				value: [
+					{
+						type: 'MENTION_USER',
+						value: {
+							type: 'PLAIN_TEXT',
+							value: 'admin',
+						},
+					},
+					{
+						type: 'PLAIN_TEXT',
+						value: ' hello ',
+					},
+					{
+						type: 'MENTION_USER',
+						value: {
+							type: 'PLAIN_TEXT',
+							value: 'admin',
+						},
+					},
+				],
+			},
+		];
+
+		const result = extractMentionsFromMessageAST(ast);
+
+		expect(result.mentions).to.deep.equal(['admin']);
+	});
+
+	it('should extract mentions from nested structures like blockquotes', () => {
+		const ast: Root = [
+			{
+				type: 'QUOTE',
+				value: [
+					{
+						type: 'PARAGRAPH',
+						value: [
+							{
+								type: 'MENTION_USER',
+								value: {
+									type: 'PLAIN_TEXT',
+									value: 'quoted.user',
+								},
+							},
+						],
+					},
+				],
+			},
+		];
+
+		const result = extractMentionsFromMessageAST(ast);
+
+		expect(result.mentions).to.deep.equal(['quoted.user']);
+	});
+
+	it('should extract mentions from bold text', () => {
+		const ast: Root = [
+			{
+				type: 'PARAGRAPH',
+				value: [
+					{
+						type: 'BOLD',
+						value: [
+							{
+								type: 'MENTION_USER',
+								value: {
+									type: 'PLAIN_TEXT',
+									value: 'bold.user',
+								},
+							},
+						],
+					},
+				],
+			},
+		];
+
+		const result = extractMentionsFromMessageAST(ast);
+
+		expect(result.mentions).to.deep.equal(['bold.user']);
+	});
+
+	it('should extract mentions from italic text', () => {
+		const ast: Root = [
+			{
+				type: 'PARAGRAPH',
+				value: [
+					{
+						type: 'ITALIC',
+						value: [
+							{
+								type: 'MENTION_CHANNEL',
+								value: {
+									type: 'PLAIN_TEXT',
+									value: 'italic-channel',
+								},
+							},
+						],
+					},
+				],
+			},
+		];
+
+		const result = extractMentionsFromMessageAST(ast);
+
+		expect(result.channels).to.deep.equal(['italic-channel']);
+	});
+
+	it('should handle special mentions like all and here', () => {
+		const ast: Root = [
+			{
+				type: 'PARAGRAPH',
+				value: [
+					{
+						type: 'MENTION_USER',
+						value: {
+							type: 'PLAIN_TEXT',
+							value: 'all',
+						},
+					},
+					{
+						type: 'PLAIN_TEXT',
+						value: ' and ',
+					},
+					{
+						type: 'MENTION_USER',
+						value: {
+							type: 'PLAIN_TEXT',
+							value: 'here',
+						},
+					},
+				],
+			},
+		];
+
+		const result = extractMentionsFromMessageAST(ast);
+
+		expect(result.mentions).to.have.members(['all', 'here']);
+	});
+
+	it('should handle empty AST', () => {
+		const ast: Root = [];
+
+		const result = extractMentionsFromMessageAST(ast);
+
+		expect(result.mentions).to.be.an('array').that.is.empty;
+		expect(result.channels).to.be.an('array').that.is.empty;
+	});
+
+	it('should extract mentions from list items', () => {
+		const ast: Root = [
+			{
+				type: 'UNORDERED_LIST',
+				value: [
+					{
+						type: 'LIST_ITEM',
+						value: [
+							{
+								type: 'MENTION_USER',
+								value: {
+									type: 'PLAIN_TEXT',
+									value: 'list.user',
+								},
+							},
+						],
+					},
+				],
+			},
+		];
+
+		const result = extractMentionsFromMessageAST(ast);
+
+		expect(result.mentions).to.deep.equal(['list.user']);
+	});
+
+	it('should extract mentions from tasks', () => {
+		const ast: Root = [
+			{
+				type: 'TASKS',
+				value: [
+					{
+						type: 'TASK',
+						status: false,
+						value: [
+							{
+								type: 'MENTION_USER',
+								value: {
+									type: 'PLAIN_TEXT',
+									value: 'task.assignee',
+								},
+							},
+						],
+					},
+				],
+			},
+		];
+
+		const result = extractMentionsFromMessageAST(ast);
+
+		expect(result.mentions).to.deep.equal(['task.assignee']);
+	});
+
+	it('should handle BigEmoji AST structure', () => {
+		const ast: Root = [
+			{
+				type: 'BIG_EMOJI',
+				value: [
+					{
+						type: 'EMOJI',
+						value: {
+							type: 'PLAIN_TEXT',
+							value: 'smile',
+						},
+						shortCode: ':smile:',
+					},
+				],
+			},
+		];
+
+		const result = extractMentionsFromMessageAST(ast);
+
+		expect(result.mentions).to.be.an('array').that.is.empty;
+		expect(result.channels).to.be.an('array').that.is.empty;
+	});
+
+	it('should extract mentions from spoiler blocks', () => {
+		const ast: Root = [
+			{
+				type: 'SPOILER_BLOCK',
+				value: [
+					{
+						type: 'PARAGRAPH',
+						value: [
+							{
+								type: 'MENTION_USER',
+								value: {
+									type: 'PLAIN_TEXT',
+									value: 'hidden.user',
+								},
+							},
+						],
+					},
+				],
+			},
+		];
+
+		const result = extractMentionsFromMessageAST(ast);
+
+		expect(result.mentions).to.deep.equal(['hidden.user']);
+	});
+});