pajlads · iProdigy · Sep 27, 2023
diff --git a/src/main/java/dinkplugin/message/DiscordMessageHandler.java b/src/main/java/dinkplugin/message/DiscordMessageHandler.java
@@ -93,7 +93,7 @@ public DiscordMessageHandler(Gson gson, Client client, DrawManager drawManager,
             .build();
     }
 
-    public void createMessage(String webhookUrl, boolean sendImage, @NonNull NotificationBody<?> inputBody) {
+    public void createMessage(String webhookUrl, boolean sendImage, @NonNull NotificationBody<?> inputBody, @Nullable CompletableFuture<Image> screenshot) {
         if (StringUtils.isBlank(webhookUrl)) return;
 
         Collection<HttpUrl> urlList = Arrays.stream(StringUtils.split(webhookUrl, '\n'))
@@ -108,7 +108,7 @@ public void createMessage(String webhookUrl, boolean sendImage, @NonNull Notific
             boolean chatHidden = hideWidget(config.screenshotHideChat(), client, WidgetInfo.CHATBOX);
             boolean whispersHidden = hideWidget(config.screenshotHideChat(), client, WidgetInfo.PRIVATE_CHAT_MESSAGE);
 
-            captureScreenshot(drawManager, config.screenshotScale() / 100.0)
+            captureScreenshot(screenshot, drawManager, config.screenshotScale() / 100.0)
                 .thenApply(image ->
                     RequestBody.create(MediaType.parse("image/" + image.getKey()), image.getValue())
                 )
@@ -267,15 +267,17 @@ private MultipartBody createBody(NotificationBody<?> mBody, @Nullable RequestBod
      * Captures the next frame and applies the specified rescaling
      * while abiding by {@link Embed#MAX_IMAGE_SIZE}.
      *
+     * @param future       a screenshot that was already taken (optional)
      * @param drawManager  {@link DrawManager}
      * @param scalePercent {@link DinkPluginConfig#screenshotScale()} divided by 100.0
      * @return future of the image byte array by the image format name
      * @apiNote scalePercent should be in (0, 1]
      * @implNote the image format is either "png" (lossless) or "jpeg" (lossy), both of which can be used in MIME type
      */
-    private static CompletableFuture<Map.Entry<String, byte[]>> captureScreenshot(DrawManager drawManager, double scalePercent) {
-        CompletableFuture<Image> future = new CompletableFuture<>();
-        drawManager.requestNextFrameListener(future::complete);
+    private static CompletableFuture<Map.Entry<String, byte[]>> captureScreenshot(CompletableFuture<Image> future, DrawManager drawManager, double scalePercent) {
+        if (future == null) {
+            future = Utils.captureScreenshot(drawManager);
+        }
         return future.thenApply(ImageUtil::bufferedImageFromImage)
             .thenApply(input -> Utils.rescale(input, scalePercent))
             .thenApply(image -> {

diff --git a/src/main/java/dinkplugin/notifiers/BaseNotifier.java b/src/main/java/dinkplugin/notifiers/BaseNotifier.java
@@ -9,6 +9,8 @@
 import org.apache.commons.lang3.StringUtils;
 
 import javax.inject.Inject;
+import java.awt.Image;
+import java.util.concurrent.CompletableFuture;
 
 public abstract class BaseNotifier {
 
@@ -31,9 +33,13 @@ public boolean isEnabled() {
     protected abstract String getWebhookUrl();
 
     protected final void createMessage(boolean sendImage, NotificationBody<?> body) {
+        this.createMessage(sendImage, null, body);
+    }
+
+    protected final void createMessage(boolean sendImage, CompletableFuture<Image> screenshot, NotificationBody<?> body) {
         String overrideUrl = getWebhookUrl();
         String url = StringUtils.isNotBlank(overrideUrl) ? overrideUrl : config.primaryWebhook();
-        messageHandler.createMessage(url, sendImage, body);
+        messageHandler.createMessage(url, sendImage, body, screenshot);
     }
 
 }
diff --git a/src/main/java/dinkplugin/notifiers/DeathNotifier.java b/src/main/java/dinkplugin/notifiers/DeathNotifier.java
@@ -1,6 +1,7 @@
 package dinkplugin.notifiers;
 
 import dinkplugin.domain.AccountType;
+import dinkplugin.message.DiscordMessageHandler;
 import dinkplugin.message.Embed;
 import dinkplugin.message.templating.Replacements;
 import dinkplugin.message.templating.Template;
@@ -26,6 +27,7 @@
 import net.runelite.api.events.InteractingChanged;
 import net.runelite.client.game.ItemManager;
 import net.runelite.client.game.NPCManager;
+import net.runelite.client.ui.DrawManager;
 import org.apache.commons.lang3.ArrayUtils;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.commons.lang3.tuple.Pair;
@@ -35,6 +37,7 @@
 
 import javax.inject.Inject;
 import javax.inject.Singleton;
+import java.awt.Image;
 import java.lang.ref.WeakReference;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -43,6 +46,10 @@
 import java.util.Comparator;
 import java.util.List;
 import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiFunction;
 import java.util.function.BiPredicate;
 import java.util.function.Function;
@@ -55,7 +62,7 @@ public class DeathNotifier extends BaseNotifier {
 
     private static final String ATTACK_OPTION = "Attack";
 
-    private static final String TOA_DEATH_MSG = "You failed to survive the Tombs of Amascut";
+    static final @VisibleForTesting String TOA_DEATH_MSG = "You failed to survive the Tombs of Amascut.";
 
     /**
      * Checks whether the actor is alive and interacting with the specified player.
@@ -83,6 +90,19 @@ public class DeathNotifier extends BaseNotifier {
     @Inject
     private NPCManager npcManager;
 
+    @Inject
+    private DrawManager drawManager;
+
+    @Inject
+    private ScheduledExecutorService executor;
+
+    /**
+     * Stores a screenshot that is captured at an earlier moment than {@link DiscordMessageHandler} would.
+     * <p>
+     * This is relevant for TOA since notifications rely upon a game message that occurs well after the actual death.
+     */
+    private final AtomicReference<CompletableFuture<Image>> screenshot = new AtomicReference<>();
+
     /**
      * Tracks the last {@link Actor} our local player interacted with,
      * for the purposes of attributing deaths to particular {@link Player}'s.
@@ -116,7 +136,8 @@ public void onActorDeath(ActorDeath actor) {
     }
 
     public void onGameMessage(String message) {
-        if (config.deathIgnoreSafe() && !Utils.getAccountType(client).isHardcore() && message.contains(TOA_DEATH_MSG) && isEnabled()) {
+        boolean shouldNotify = screenshot.get() != null || (config.deathIgnoreSafe() && !Utils.getAccountType(client).isHardcore() && isEnabled());
+        if (message.equals(TOA_DEATH_MSG) && shouldNotify) {
             // https://github.com/pajlads/DinkPlugin/issues/316
             // though, hardcore (group) ironmen just use the normal ActorDeath trigger for TOA
             handleNotify(Danger.DANGEROUS);
@@ -130,9 +151,23 @@ public void onInteraction(InteractingChanged event) {
     }
 
     private void handleNotify(Danger dangerOverride) {
+        // Check if a screenshot for this death was already captured (relevant for TOA)
+        CompletableFuture<Image> screenshotFuture = screenshot.getAndSet(null);
+
+        // Ignore safe deaths, depending on config
         Danger danger = dangerOverride != null ? dangerOverride : getDangerLevel(client);
-        if (danger == Danger.SAFE && config.deathIgnoreSafe())
+        if (danger == Danger.SAFE && config.deathIgnoreSafe()) {
+            // Capture screenshot now if the death was in TOA since the game message comes late
+            if (config.deathSendImage() && WorldUtils.isAmascutTombs(WorldUtils.getLocation(client).getRegionID())) {
+                screenshot.set(Utils.captureScreenshot(drawManager));
+                executor.schedule(() -> {
+                    // release memory just in case onGameMessage branch doesn't get called
+                    screenshot.set(null);
+                }, 15L, TimeUnit.SECONDS);
+            }
+
             return;
+        }
 
         Collection<Item> items = ItemUtils.getItems(client);
         List<Pair<Item, Long>> itemsByPrice = getPricedItems(itemManager, items);
@@ -188,7 +223,7 @@ private void handleNotify(Danger dangerOverride) {
             lostStacks
         );
 
-        createMessage(config.deathSendImage(), NotificationBody.builder()
+        createMessage(config.deathSendImage(), screenshotFuture, NotificationBody.builder()
             .text(notifyMessage)
             .extra(extra)
             .embeds(keptItemEmbeds)

diff --git a/src/main/java/dinkplugin/util/Utils.java b/src/main/java/dinkplugin/util/Utils.java
@@ -8,6 +8,7 @@
 import lombok.extern.slf4j.Slf4j;
 import net.runelite.api.Client;
 import net.runelite.api.Varbits;
+import net.runelite.client.ui.DrawManager;
 import net.runelite.client.util.ColorUtil;
 import net.runelite.client.util.Text;
 import okhttp3.Call;
@@ -23,6 +24,7 @@
 import javax.imageio.ImageIO;
 import javax.swing.SwingUtilities;
 import java.awt.Color;
+import java.awt.Image;
 import java.awt.Toolkit;
 import java.awt.datatransfer.Clipboard;
 import java.awt.datatransfer.DataFlavor;
@@ -198,6 +200,12 @@ public boolean hasImage(@NotNull MultipartBody body) {
         });
     }
 
+    public CompletableFuture<Image> captureScreenshot(@NotNull DrawManager drawManager) {
+        CompletableFuture<Image> future = new CompletableFuture<>();
+        drawManager.requestNextFrameListener(future::complete);
+        return future;
+    }
+
     public CompletableFuture<String> readClipboard() {
         CompletableFuture<String> future = new CompletableFuture<>();
         SwingUtilities.invokeLater(() -> {

diff --git a/src/test/java/dinkplugin/notifiers/ClueNotifierTest.java b/src/test/java/dinkplugin/notifiers/ClueNotifierTest.java
@@ -88,7 +88,8 @@ void testNotify() {
                 )
                 .extra(new ClueNotificationData("medium", 1312, Collections.singletonList(new SerializedItemStack(ItemID.RUBY, 1, RUBY_PRICE, "Ruby"))))
                 .type(NotificationType.CLUE)
-                .build()
+                .build(),
+            null
         );
     }
 
@@ -114,7 +115,7 @@ void testIgnoreTier() {
         plugin.onWidgetLoaded(event);
 
         // ensure no notification was fired
-        verify(messageHandler, never()).createMessage(any(), anyBoolean(), any());
+        verify(messageHandler, never()).createMessage(any(), anyBoolean(), any(), any());
     }
 
     @Test
@@ -139,7 +140,7 @@ void testIgnoreLoot() {
         plugin.onWidgetLoaded(event);
 
         // ensure no notification was fired
-        verify(messageHandler, never()).createMessage(any(), anyBoolean(), any());
+        verify(messageHandler, never()).createMessage(any(), anyBoolean(), any(), any());
     }
 
     @Test
@@ -167,7 +168,7 @@ void testDisabled() {
         plugin.onWidgetLoaded(event);
 
         // ensure no notification was fired
-        verify(messageHandler, never()).createMessage(any(), anyBoolean(), any());
+        verify(messageHandler, never()).createMessage(any(), anyBoolean(), any(), any());
     }
 
 }
diff --git a/src/test/java/dinkplugin/notifiers/CollectionNotifierTest.java b/src/test/java/dinkplugin/notifiers/CollectionNotifierTest.java
@@ -77,7 +77,8 @@ void testNotify() {
                 )
                 .extra(new CollectionNotificationData(item, ItemID.SEERCULL, (long) price, 1, TOTAL_ENTRIES))
                 .type(NotificationType.COLLECTION)
-                .build()
+                .build(),
+            null
         );
     }
 
@@ -117,7 +118,8 @@ void testNotifyPopup() {
                 )
                 .extra(new CollectionNotificationData(item, ItemID.SEERCULL, (long) price, 1, TOTAL_ENTRIES))
                 .type(NotificationType.COLLECTION)
-                .build()
+                .build(),
+            null
         );
     }
 
@@ -127,7 +129,7 @@ void testIgnore() {
         notifier.onChatMessage("New item added to your backpack: weed");
 
         // ensure no notification occurred
-        verify(messageHandler, never()).createMessage(any(), anyBoolean(), any());
+        verify(messageHandler, never()).createMessage(any(), anyBoolean(), any(), any());
     }
 
     @Test
@@ -140,7 +142,7 @@ void testDisabled() {
         notifier.onChatMessage("New item added to your collection log: " + item);
 
         // ensure no notification occurred
-        verify(messageHandler, never()).createMessage(any(), anyBoolean(), any());
+        verify(messageHandler, never()).createMessage(any(), anyBoolean(), any(), any());
     }
 
 }
diff --git a/src/test/java/dinkplugin/notifiers/CombatTaskNotifierTest.java b/src/test/java/dinkplugin/notifiers/CombatTaskNotifierTest.java
@@ -76,7 +76,8 @@ void testNotify() {
                 .extra(new CombatAchievementData(CombatAchievementTier.HARD, "Whack-a-Mole", 3, 200, 85, 189, null))
                 .playerName(PLAYER_NAME)
                 .type(NotificationType.COMBAT_ACHIEVEMENT)
-                .build()
+                .build(),
+            null
         );
     }
 
@@ -103,7 +104,8 @@ void testNotifyUnlock() {
                 .extra(new CombatAchievementData(CombatAchievementTier.GRANDMASTER, "No Pressure", 6, 1466, 1466 - 1465, 2005 - 1465, CombatAchievementTier.MASTER))
                 .playerName(PLAYER_NAME)
                 .type(NotificationType.COMBAT_ACHIEVEMENT)
-                .build()
+                .build(),
+            null
         );
     }
 
@@ -130,7 +132,8 @@ void testNotifyUnlockGrand() {
                 .extra(new CombatAchievementData(CombatAchievementTier.GRANDMASTER, "No Pressure", 6, 2005, null, null, CombatAchievementTier.GRANDMASTER))
                 .playerName(PLAYER_NAME)
                 .type(NotificationType.COMBAT_ACHIEVEMENT)
-                .build()
+                .build(),
+            null
         );
     }
 
@@ -140,7 +143,7 @@ void testSkipped() {
         notifier.onGameMessage("Congratulations, you've completed an easy combat task: A Slow Death.");
 
         // ensure no notification occurred
-        verify(messageHandler, never()).createMessage(any(), anyBoolean(), any());
+        verify(messageHandler, never()).createMessage(any(), anyBoolean(), any(), any());
     }
 
     @Test
@@ -149,7 +152,7 @@ void testIgnore() {
         notifier.onGameMessage("Congratulations, you've completed a gachi combat task: Swordfight with the homies.");
 
         // ensure no notification occurred
-        verify(messageHandler, never()).createMessage(any(), anyBoolean(), any());
+        verify(messageHandler, never()).createMessage(any(), anyBoolean(), any(), any());
     }
 
     @Test
@@ -161,7 +164,7 @@ void testDisabled() {
         notifier.onGameMessage("Congratulations, you've completed a hard combat task: Whack-a-Mole.");
 
         // ensure no notification occurred
-        verify(messageHandler, never()).createMessage(any(), anyBoolean(), any());
+        verify(messageHandler, never()).createMessage(any(), anyBoolean(), any(), any());
     }
 
     private static Template buildUnlockTemplate(String tier, String task) {