SurfSense/surfsense_web/components/assistant-ui/revert-turn-button.tsx

"use client";

/**
 * "Revert turn" button rendered at the bottom of every completed
 * assistant turn that has at least one reversible action.
 *
 * The button reads from the unified ``useAgentActionsQuery`` cache
 * (the SAME react-query cache the agent-actions sheet and the inline
 * Revert button consume) filtered by ``chat_turn_id``. It shows a
 * confirmation dialog summarising "N reversible / M total" and, on
 * confirm, calls ``POST /threads/{id}/revert-turn/{chat_turn_id}``.
 *
 * The route returns a per-action result list and never collapses the
 * batch into a 4xx — so we render any failed/not_reversible rows inline
 * with their messages.
 */

import { useQueryClient } from "@tanstack/react-query";
import { useAtomValue } from "jotai";
import { CheckIcon, RotateCcw, XCircleIcon } from "lucide-react";
import { useMemo, useState } from "react";
import { toast } from "sonner";
import { chatSessionStateAtom } from "@/atoms/chat/chat-session-state.atom";
import {
	AlertDialog,
	AlertDialogAction,
	AlertDialogCancel,
	AlertDialogContent,
	AlertDialogDescription,
	AlertDialogFooter,
	AlertDialogHeader,
	AlertDialogTitle,
	AlertDialogTrigger,
} from "@/components/ui/alert-dialog";
import { Button } from "@/components/ui/button";
import { getToolDisplayName } from "@/contracts/enums/toolIcons";
import {
	applyRevertTurnResultsToCache,
	useAgentActionsQuery,
} from "@/hooks/use-agent-actions-query";
import {
	agentActionsApiService,
	type RevertTurnActionResult,
} from "@/lib/apis/agent-actions-api.service";
import { AppError } from "@/lib/error";
import { cn } from "@/lib/utils";

interface RevertTurnButtonProps {
	chatTurnId: string | null | undefined;
}

export function RevertTurnButton({ chatTurnId }: RevertTurnButtonProps) {
	const session = useAtomValue(chatSessionStateAtom);
	const threadId = session?.threadId ?? null;
	const queryClient = useQueryClient();
	const { findByChatTurnId } = useAgentActionsQuery(threadId);
	const [isReverting, setIsReverting] = useState(false);
	const [confirmOpen, setConfirmOpen] = useState(false);
	const [resultsOpen, setResultsOpen] = useState(false);
	const [results, setResults] = useState<RevertTurnActionResult[]>([]);

	const actions = useMemo(() => findByChatTurnId(chatTurnId), [findByChatTurnId, chatTurnId]);

	const reversibleCount = useMemo(
		() =>
			actions.filter(
				(a) =>
					a.reversible &&
					(a.reverted_by_action_id === null || a.reverted_by_action_id === undefined) &&
					!a.is_revert_action &&
					(a.error === null || a.error === undefined)
			).length,
		[actions]
	);
	const totalCount = useMemo(() => actions.filter((a) => !a.is_revert_action).length, [actions]);

	if (!chatTurnId) return null;
	if (reversibleCount === 0) return null;
	if (!threadId) return null;

	const handleRevertTurn = async () => {
		setIsReverting(true);
		try {
			const response = await agentActionsApiService.revertTurn(threadId, chatTurnId);
			setResults(response.results);
			const revertedEntries = response.results
				.filter((r) => r.status === "reverted" || r.status === "already_reverted")
				.map((r) => ({ id: r.action_id, newActionId: r.new_action_id ?? null }));
			if (revertedEntries.length > 0) {
				applyRevertTurnResultsToCache(queryClient, threadId, revertedEntries);
			}
			if (response.status === "ok") {
				toast.success(
					response.reverted === 1 ? "Reverted 1 action." : `Reverted ${response.reverted} actions.`
				);
			} else {
				// Every "not undone" bucket counts as a failure for the
				// user-facing summary. ``skipped`` rows are batch
				// artefacts (revert rows themselves) and intentionally
				// excluded from the failure tally.
				const failureCount =
					response.failed + response.not_reversible + (response.permission_denied ?? 0);
				toast.warning(
					`Reverted ${response.reverted} of ${response.total}. ${failureCount} could not be undone.`
				);
				setResultsOpen(true);
			}
		} catch (err) {
			if (err instanceof AppError && err.status === 503) {
				return;
			}
			const message =
				err instanceof AppError
					? err.message
					: err instanceof Error
						? err.message
						: "Failed to revert turn.";
			toast.error(message);
		} finally {
			setIsReverting(false);
			setConfirmOpen(false);
		}
	};

	return (
		<>
			<AlertDialog open={confirmOpen} onOpenChange={setConfirmOpen}>
				<AlertDialogTrigger asChild>
					<Button
						size="sm"
						variant="ghost"
						className="text-muted-foreground hover:text-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>
						<AlertDialogDescription>
							This will undo {reversibleCount} of {totalCount} action
							{totalCount === 1 ? "" : "s"} from this turn in reverse order. The chat history and
							any read-only actions are preserved. Some rows may not be reversible — partial success
							is normal.
						</AlertDialogDescription>
					</AlertDialogHeader>
					<AlertDialogFooter>
						<AlertDialogCancel disabled={isReverting}>Cancel</AlertDialogCancel>
						<AlertDialogAction
							onClick={(e) => {
								e.preventDefault();
								handleRevertTurn();
							}}
							disabled={isReverting}
						>
							{isReverting ? "Reverting…" : "Revert turn"}
						</AlertDialogAction>
					</AlertDialogFooter>
				</AlertDialogContent>
			</AlertDialog>

			<AlertDialog open={resultsOpen} onOpenChange={setResultsOpen}>
				<AlertDialogContent>
					<AlertDialogHeader>
						<AlertDialogTitle>Revert results</AlertDialogTitle>
						<AlertDialogDescription>
							Some actions could not be reverted. Review per-row outcomes below.
						</AlertDialogDescription>
					</AlertDialogHeader>
					<ul className="max-h-72 overflow-y-auto space-y-2 text-sm">
						{results.map((r) => (
							<RevertResultRow key={r.action_id} result={r} />
						))}
					</ul>
					<AlertDialogFooter>
						<AlertDialogAction onClick={() => setResultsOpen(false)}>Close</AlertDialogAction>
					</AlertDialogFooter>
				</AlertDialogContent>
			</AlertDialog>
		</>
	);
}

function RevertResultRow({ result }: { result: RevertTurnActionResult }) {
	const isOk = result.status === "reverted" || result.status === "already_reverted";
	const Icon = isOk ? CheckIcon : XCircleIcon;
	return (
		<li className="flex items-start gap-2 rounded-md border bg-muted/30 px-3 py-2">
			<Icon
				className={cn("size-4 mt-0.5 shrink-0", isOk ? "text-emerald-500" : "text-destructive")}
			/>
			<div className="min-w-0 flex-1">
				<p className="font-medium truncate">
					{getToolDisplayName(result.tool_name)}{" "}
					<span className="ml-1 text-xs text-muted-foreground">
						{result.status.replace(/_/g, " ")}
					</span>
				</p>
				{(result.message || result.error) && (
					<p className="text-xs text-muted-foreground mt-0.5">{result.error ?? result.message}</p>
				)}
			</div>
		</li>
	);
}