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?