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 = () => {
)}
>
)}
+
);
@@ -623,10 +624,7 @@ const AssistantActionBar: FC = () => {
)}
-
-
-
-
+
);
};
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 (
<>
-
-
-
+
+ ) : (
+
+
+
+ )}
Revert this turn?