refactor: update MessageInfoDropdown to accept chatTurnId prop and enhance RevertTurnButton integration for improved functionality

This commit is contained in:
Anish Sarkar 2026-05-15 10:48:57 +05:30
parent 5092bd3e8c
commit 8001cae1b4
2 changed files with 32 additions and 16 deletions

View file

@ -274,7 +274,7 @@ function formatTurnCost(micros: number): string {
return "$0"; return "$0";
} }
const MessageInfoDropdown: FC = () => { const MessageInfoDropdown: FC<{ chatTurnId: string | null | undefined }> = ({ chatTurnId }) => {
const messageId = useAuiState(({ message }) => message?.id); const messageId = useAuiState(({ message }) => message?.id);
const createdAt = useAuiState(({ message }) => message?.createdAt); const createdAt = useAuiState(({ message }) => message?.createdAt);
const usage = useTokenUsage(messageId); const usage = useTokenUsage(messageId);
@ -359,6 +359,7 @@ const MessageInfoDropdown: FC = () => {
)} )}
</> </>
)} )}
<RevertTurnButton chatTurnId={chatTurnId} variant="menu-item" />
</ActionBarMorePrimitive.Content> </ActionBarMorePrimitive.Content>
</ActionBarMorePrimitive.Root> </ActionBarMorePrimitive.Root>
); );
@ -623,10 +624,7 @@ const AssistantActionBar: FC = () => {
<ClipboardPaste /> <ClipboardPaste />
</TooltipIconButton> </TooltipIconButton>
)} )}
<MessageInfoDropdown /> <MessageInfoDropdown chatTurnId={chatTurnId} />
<div className="ml-auto">
<RevertTurnButton chatTurnId={chatTurnId} />
</div>
</ActionBarPrimitive.Root> </ActionBarPrimitive.Root>
); );
}; };

View file

@ -15,6 +15,7 @@
* with their messages. * with their messages.
*/ */
import { ActionBarMorePrimitive } from "@assistant-ui/react";
import { useQueryClient } from "@tanstack/react-query"; import { useQueryClient } from "@tanstack/react-query";
import { useAtomValue } from "jotai"; import { useAtomValue } from "jotai";
import { CheckIcon, RotateCcw, XCircleIcon } from "lucide-react"; import { CheckIcon, RotateCcw, XCircleIcon } from "lucide-react";
@ -47,9 +48,10 @@ import { cn } from "@/lib/utils";
interface RevertTurnButtonProps { interface RevertTurnButtonProps {
chatTurnId: string | null | undefined; 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 session = useAtomValue(chatSessionStateAtom);
const threadId = session?.threadId ?? null; const threadId = session?.threadId ?? null;
const queryClient = useQueryClient(); const queryClient = useQueryClient();
@ -125,23 +127,39 @@ export function RevertTurnButton({ chatTurnId }: RevertTurnButtonProps) {
return ( return (
<> <>
<AlertDialog open={confirmOpen} onOpenChange={setConfirmOpen}> <AlertDialog open={confirmOpen} onOpenChange={setConfirmOpen}>
<AlertDialogTrigger asChild> {variant === "menu-item" ? (
<Button <ActionBarMorePrimitive.Item
size="sm" 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"
variant="ghost" onSelect={(e) => {
className="text-muted-foreground hover:text-accent-foreground gap-1.5" e.preventDefault();
onClick={(e) => {
e.stopPropagation();
setConfirmOpen(true); setConfirmOpen(true);
}} }}
> >
<RotateCcw className="size-3.5" /> <RotateCcw className="size-3.5" />
<span>Revert turn</span> <span>Revert turn</span>
<span className="text-xs tabular-nums opacity-70"> <span className="ml-auto text-xs tabular-nums opacity-70">
{reversibleCount}/{totalCount} {reversibleCount}/{totalCount}
</span> </span>
</Button> </ActionBarMorePrimitive.Item>
</AlertDialogTrigger> ) : (
<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> <AlertDialogContent>
<AlertDialogHeader> <AlertDialogHeader>
<AlertDialogTitle>Revert this turn?</AlertDialogTitle> <AlertDialogTitle>Revert this turn?</AlertDialogTitle>