From e79f9782ab43b44243352e41f812b43232726a26 Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Thu, 29 Aug 2024 08:13:16 +1000
Subject: [PATCH] fix(ui): queue count badge positioning

---
 .../components/QueueActionsMenuButton.tsx     | 36 ++++++++++++++-----
 .../queue/components/QueueControls.tsx        |  7 ++--
 2 files changed, 32 insertions(+), 11 deletions(-)

diff --git a/invokeai/frontend/web/src/features/queue/components/QueueActionsMenuButton.tsx b/invokeai/frontend/web/src/features/queue/components/QueueActionsMenuButton.tsx
index ce96322606..8332aabd4b 100644
--- a/invokeai/frontend/web/src/features/queue/components/QueueActionsMenuButton.tsx
+++ b/invokeai/frontend/web/src/features/queue/components/QueueActionsMenuButton.tsx
@@ -1,6 +1,5 @@
 import {
   Badge,
-  Box,
   IconButton,
   Menu,
   MenuButton,
@@ -18,13 +17,18 @@ import { usePauseProcessor } from 'features/queue/hooks/usePauseProcessor';
 import { useResumeProcessor } from 'features/queue/hooks/useResumeProcessor';
 import { useFeatureStatus } from 'features/system/hooks/useFeatureStatus';
 import { setActiveTab } from 'features/ui/store/uiSlice';
+import type { RefObject } from 'react';
 import { memo, useCallback, useEffect, useRef, useState } from 'react';
 import { useTranslation } from 'react-i18next';
 import { PiPauseFill, PiPlayFill, PiTrashSimpleBold } from 'react-icons/pi';
 import { RiListCheck, RiPlayList2Fill } from 'react-icons/ri';
 import { useGetQueueStatusQuery } from 'services/api/endpoints/queue';
 
-export const QueueActionsMenuButton = memo(() => {
+type Props = {
+  containerRef: RefObject<HTMLDivElement>;
+};
+
+export const QueueActionsMenuButton = memo(({ containerRef }: Props) => {
   const { isOpen, onOpen, onClose } = useDisclosure();
   const dispatch = useAppDispatch();
   const { t } = useTranslation();
@@ -54,14 +58,30 @@ export const QueueActionsMenuButton = memo(() => {
   }, [dispatch]);
 
   useEffect(() => {
-    if (menuButtonRef.current) {
-      const { x, y } = menuButtonRef.current.getBoundingClientRect();
-      setBadgePos({ x: x - 10, y: y - 10 });
+    if (!containerRef.current || !menuButtonRef.current) {
+      return;
     }
-  }, []);
+
+    const container = containerRef.current;
+    const menuButton = menuButtonRef.current;
+
+    const cb = () => {
+      const { x, y } = menuButton.getBoundingClientRect();
+      setBadgePos({ x: x - 10, y: y - 10 });
+    };
+
+    // // update badge position on resize
+    const resizeObserver = new ResizeObserver(cb);
+    resizeObserver.observe(container);
+    cb();
+
+    return () => {
+      resizeObserver.disconnect();
+    };
+  }, [containerRef]);
 
   return (
-    <Box pos="relative">
+    <>
       <Menu isOpen={isOpen} onOpen={onOpen} onClose={onClose} placement="bottom-end">
         <MenuButton ref={menuButtonRef} as={IconButton} aria-label="Queue Actions Menu" icon={<RiListCheck />} />
         <MenuList>
@@ -113,7 +133,7 @@ export const QueueActionsMenuButton = memo(() => {
           </Badge>
         </Portal>
       )}
-    </Box>
+    </>
   );
 });
 
diff --git a/invokeai/frontend/web/src/features/queue/components/QueueControls.tsx b/invokeai/frontend/web/src/features/queue/components/QueueControls.tsx
index 1acbd6b697..34b481fc6b 100644
--- a/invokeai/frontend/web/src/features/queue/components/QueueControls.tsx
+++ b/invokeai/frontend/web/src/features/queue/components/QueueControls.tsx
@@ -3,20 +3,21 @@ import { ClearQueueIconButton } from 'features/queue/components/ClearQueueIconBu
 import QueueFrontButton from 'features/queue/components/QueueFrontButton';
 import ProgressBar from 'features/system/components/ProgressBar';
 import { useFeatureStatus } from 'features/system/hooks/useFeatureStatus';
-import { memo } from 'react';
+import { memo, useRef } from 'react';
 
 import { InvokeQueueBackButton } from './InvokeQueueBackButton';
 import { QueueActionsMenuButton } from './QueueActionsMenuButton';
 
 const QueueControls = () => {
   const isPrependEnabled = useFeatureStatus('prependQueue');
+  const containerRef = useRef<HTMLDivElement>(null);
   return (
-    <Flex w="full" position="relative" borderRadius="base" gap={2} flexDir="column">
+    <Flex ref={containerRef} w="full" position="relative" borderRadius="base" gap={2} flexDir="column">
       <ButtonGroup size="lg" isAttached={false}>
         {isPrependEnabled && <QueueFrontButton />}
         <InvokeQueueBackButton />
         <Spacer />
-        <QueueActionsMenuButton />
+        <QueueActionsMenuButton containerRef={containerRef} />
         <ClearQueueIconButton />
       </ButtonGroup>
       <ProgressBar />