Merge pull request #292 from dev-kittens/fix/edit-replies

Fix handling replies in edited messages

PluralKit.Bot/Handlers/MessageCreated.cs

@@ -53,7 +53,7 @@ namespace PluralKit.Bot
 
         private bool IsDuplicateMessage(Message msg) =>
             // We consider a message duplicate if it has the same ID as the previous message that hit the gateway
-            _lastMessageCache.GetLastMessage(msg.ChannelId) == msg.Id;
+            _lastMessageCache.GetLastMessage(msg.ChannelId)?.mid == msg.Id;
 
         public async Task Handle(Shard shard, MessageCreateEvent evt)
         {
@@ -66,7 +66,7 @@ namespace PluralKit.Bot
             
             // Log metrics and message info
             _metrics.Measure.Meter.Mark(BotMetrics.MessagesReceived);
-            _lastMessageCache.AddMessage(evt.ChannelId, evt.Id);
+            _lastMessageCache.AddMessage(evt);
             
             // Get message context from DB (tracking w/ metrics)
             MessageContext ctx;

PluralKit.Bot/Handlers/MessageEdited.cs

@@ -6,6 +6,7 @@ using App.Metrics;
 using Myriad.Cache;
 using Myriad.Extensions;
 using Myriad.Gateway;
+using Myriad.Rest;
 using Myriad.Types;
 
 using PluralKit.Core;
@@ -23,8 +24,10 @@ namespace PluralKit.Bot
         private readonly Cluster _client;
         private readonly IDiscordCache _cache;
         private readonly Bot _bot;
+        private readonly DiscordApiClient _rest;
         
-        public MessageEdited(LastMessageCacheService lastMessageCache, ProxyService proxy, IDatabase db, IMetrics metrics, ModelRepository repo, Cluster client, IDiscordCache cache, Bot bot)
+
+        public MessageEdited(LastMessageCacheService lastMessageCache, ProxyService proxy, IDatabase db, IMetrics metrics, ModelRepository repo, Cluster client, IDiscordCache cache, Bot bot, DiscordApiClient rest)
         {
             _lastMessageCache = lastMessageCache;
             _proxy = proxy;
@@ -34,6 +37,7 @@ namespace PluralKit.Bot
             _client = client;
             _cache = cache;
             _bot = bot;
+            _rest = rest;
         }
 
         public async Task Handle(Shard shard, MessageUpdateEvent evt)
@@ -48,9 +52,10 @@ namespace PluralKit.Bot
             if (channel.Type != Channel.ChannelType.GuildText)
                 return;
             var guild = _cache.GetGuild(channel.GuildId!.Value);
+            var lastMessage = _lastMessageCache.GetLastMessage(evt.ChannelId);
 
             // Only react to the last message in the channel
-            if (_lastMessageCache.GetLastMessage(evt.ChannelId) != evt.Id)
+            if (lastMessage?.mid != evt.Id)
                 return;
             
             // Just run the normal message handling code, with a flag to disable autoproxying
@@ -59,6 +64,8 @@ namespace PluralKit.Bot
             using (_metrics.Measure.Timer.Time(BotMetrics.MessageContextQueryTime))
                 ctx = await _repo.GetMessageContext(conn, evt.Author.Value!.Id, channel.GuildId!.Value, evt.ChannelId);
 
+            Message referencedMessage = (lastMessage.referenced_message != null) ? await _rest.GetMessage(evt.ChannelId, lastMessage.referenced_message.Value) : null;
+
             // TODO: is this missing anything?
             var equivalentEvt = new MessageCreateEvent
             {
@@ -68,7 +75,10 @@ namespace PluralKit.Bot
                 Author = evt.Author.Value,
                 Member = evt.Member.Value,
                 Content = evt.Content.Value,
-                Attachments = evt.Attachments.Value ?? Array.Empty<Message.Attachment>()
+                Attachments = evt.Attachments.Value ?? Array.Empty<Message.Attachment>(),
+                MessageReference = (lastMessage.referenced_message != null) ? new (channel.GuildId, evt.ChannelId, lastMessage.referenced_message.Value) : null,
+                ReferencedMessage = referencedMessage,
+                Type = (lastMessage.referenced_message != null) ? Message.MessageType.Reply : Message.MessageType.Default,
             };
             var botPermissions = _bot.PermissionsIn(channel.Id);
             await _proxy.HandleIncomingMessage(shard, equivalentEvt, ctx, allowAutoproxy: false, guild: guild, channel: channel, botPermissions: botPermissions);

PluralKit.Bot/Services/LastMessageCacheService.cs

@@ -1,6 +1,8 @@
 using System.Collections.Concurrent;
 using System.Collections.Generic;
 
+using Myriad.Types;
+
 namespace PluralKit.Bot
 {
     // Doing things like this instead of enabling D.NET's message cache because the message cache is, let's face it,
@@ -10,17 +12,30 @@ namespace PluralKit.Bot
     // TODO: is this still needed after the D#+ migration?
     public class LastMessageCacheService
     {
-        private readonly IDictionary<ulong, ulong> _cache = new ConcurrentDictionary<ulong, ulong>();
+        private readonly IDictionary<ulong, CachedMessage> _cache = new ConcurrentDictionary<ulong, CachedMessage>();
 
-        public void AddMessage(ulong channel, ulong message)
+        public void AddMessage(Message msg)
         {
-            _cache[channel] = message;
+            _cache[msg.ChannelId] = new CachedMessage(msg);
         }
 
-        public ulong? GetLastMessage(ulong channel)
+        public CachedMessage GetLastMessage(ulong channel)
         {
             if (_cache.TryGetValue(channel, out var message)) return message;
             return null;
         }
     }
+
+    public class CachedMessage
+    {
+        public ulong mid;
+        public ulong? referenced_message;
+
+        public CachedMessage(Message msg)
+        {
+            mid = msg.Id;
+            if (msg.ReferencedMessage.HasValue)
+                referenced_message = msg.ReferencedMessage.Value.Id;
+        }
+    }
 }