From 8001cae1b4c2550fab5be59fef9b7082d38e5ed9 Mon Sep 17 00:00:00 2001
From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com>
Date: Fri, 15 May 2026 10:48:57 +0530
Subject: [PATCH] refactor: update MessageInfoDropdown to accept chatTurnId
 prop and enhance RevertTurnButton integration for improved functionality

---
 .../assistant-ui/assistant-message.tsx        |  8 ++--
 .../assistant-ui/revert-turn-button.tsx       | 40 ++++++++++++++-----
 2 files changed, 32 insertions(+), 16 deletions(-)

diff --git a/surfsense_web/components/assistant-ui/assistant-message.tsx b/surfsense_web/components/assistant-ui/assistant-message.tsx
index c0716f1eb..76396b397 100644
--- a/surfsense_web/components/assistant-ui/assistant-message.tsx
+++ b/surfsense_web/components/assistant-ui/assistant-message.tsx
@@ -274,7 +274,7 @@ function formatTurnCost(micros: number): string {
 	return "$0";
 }
 
-const MessageInfoDropdown: FC = () => {
+const MessageInfoDropdown: FC<{ chatTurnId: string | null | undefined }> = ({ chatTurnId }) => {
 	const messageId = useAuiState(({ message }) => message?.id);
 	const createdAt = useAuiState(({ message }) => message?.createdAt);
 	const usage = useTokenUsage(messageId);
@@ -359,6 +359,7 @@ const MessageInfoDropdown: FC = () => {
 						)}
 					</>
 				)}
+				<RevertTurnButton chatTurnId={chatTurnId} variant="menu-item" />
 			</ActionBarMorePrimitive.Content>
 		</ActionBarMorePrimitive.Root>
 	);
@@ -623,10 +624,7 @@ const AssistantActionBar: FC = () => {
 					<ClipboardPaste />
 				</TooltipIconButton>
 			)}
-			<MessageInfoDropdown />
-			<div className="ml-auto">
-				<RevertTurnButton chatTurnId={chatTurnId} />
-			</div>
+			<MessageInfoDropdown chatTurnId={chatTurnId} />
 		</ActionBarPrimitive.Root>
 	);
 };
diff --git a/surfsense_web/components/assistant-ui/revert-turn-button.tsx b/surfsense_web/components/assistant-ui/revert-turn-button.tsx
index 2cc151ee2..7b2725989 100644
--- a/surfsense_web/components/assistant-ui/revert-turn-button.tsx
+++ b/surfsense_web/components/assistant-ui/revert-turn-button.tsx
@@ -15,6 +15,7 @@
  * with their messages.
  */
 
+import { ActionBarMorePrimitive } from "@assistant-ui/react";
 import { useQueryClient } from "@tanstack/react-query";
 import { useAtomValue } from "jotai";
 import { CheckIcon, RotateCcw, XCircleIcon } from "lucide-react";
@@ -47,9 +48,10 @@ import { cn } from "@/lib/utils";
 
 interface RevertTurnButtonProps {
 	chatTurnId: string | null | undefined;
+	variant?: "button" | "menu-item";
 }
 
-export function RevertTurnButton({ chatTurnId }: RevertTurnButtonProps) {
+export function RevertTurnButton({ chatTurnId, variant = "button" }: RevertTurnButtonProps) {
 	const session = useAtomValue(chatSessionStateAtom);
 	const threadId = session?.threadId ?? null;
 	const queryClient = useQueryClient();
@@ -125,23 +127,39 @@ export function RevertTurnButton({ chatTurnId }: RevertTurnButtonProps) {
 	return (
 		<>
 			<AlertDialog open={confirmOpen} onOpenChange={setConfirmOpen}>
-				<AlertDialogTrigger asChild>
-					<Button
-						size="sm"
-						variant="ghost"
-						className="text-muted-foreground hover:text-accent-foreground gap-1.5"
-						onClick={(e) => {
-							e.stopPropagation();
+				{variant === "menu-item" ? (
+					<ActionBarMorePrimitive.Item
+						className="focus:bg-accent focus:text-accent-foreground relative flex cursor-default items-center gap-2 rounded-sm px-2 py-1.5 text-sm outline-hidden select-none"
+						onSelect={(e) => {
+							e.preventDefault();
 							setConfirmOpen(true);
 						}}
 					>
 						<RotateCcw className="size-3.5" />
 						<span>Revert turn</span>
-						<span className="text-xs tabular-nums opacity-70">
+						<span className="ml-auto text-xs tabular-nums opacity-70">
 							{reversibleCount}/{totalCount}
 						</span>
-					</Button>
-				</AlertDialogTrigger>
+					</ActionBarMorePrimitive.Item>
+				) : (
+					<AlertDialogTrigger asChild>
+						<Button
+							size="sm"
+							variant="ghost"
+							className="text-muted-foreground hover:text-accent-foreground gap-1.5"
+							onClick={(e) => {
+								e.stopPropagation();
+								setConfirmOpen(true);
+							}}
+						>
+							<RotateCcw className="size-3.5" />
+							<span>Revert turn</span>
+							<span className="text-xs tabular-nums opacity-70">
+								{reversibleCount}/{totalCount}
+							</span>
+						</Button>
+					</AlertDialogTrigger>
+				)}
 				<AlertDialogContent>
 					<AlertDialogHeader>
 						<AlertDialogTitle>Revert this turn?</AlertDialogTitle>