feat: improved agent streaming

surfsense_web/components/assistant-ui/tool-fallback.tsx

@@ -1,12 +1,33 @@
 import type { ToolCallMessagePartComponent } from "@assistant-ui/react";
-import { CheckIcon, ChevronDownIcon, ChevronUpIcon, XCircleIcon } from "lucide-react";
+import { useAtomValue, useSetAtom } from "jotai";
+import { CheckIcon, ChevronDownIcon, ChevronUpIcon, RotateCcw, XCircleIcon } from "lucide-react";
 import { useMemo, useState } from "react";
+import { toast } from "sonner";
+import {
+	agentActionByToolCallIdAtom,
+	markAgentActionRevertedAtom,
+} from "@/atoms/chat/agent-actions.atom";
+import { chatSessionStateAtom } from "@/atoms/chat/chat-session-state.atom";
 import {
 	DoomLoopApprovalToolUI,
 	isDoomLoopInterrupt,
 } from "@/components/tool-ui/doom-loop-approval";
 import { GenericHitlApprovalToolUI } from "@/components/tool-ui/generic-hitl-approval";
+import {
+	AlertDialog,
+	AlertDialogAction,
+	AlertDialogCancel,
+	AlertDialogContent,
+	AlertDialogDescription,
+	AlertDialogFooter,
+	AlertDialogHeader,
+	AlertDialogTitle,
+	AlertDialogTrigger,
+} from "@/components/ui/alert-dialog";
+import { Button } from "@/components/ui/button";
 import { getToolIcon } from "@/contracts/enums/toolIcons";
+import { agentActionsApiService } from "@/lib/apis/agent-actions-api.service";
+import { AppError } from "@/lib/error";
 import { isInterruptResult } from "@/lib/hitl";
 import { cn } from "@/lib/utils";
 
@@ -14,7 +35,99 @@ function formatToolName(name: string): string {
 	return name.replace(/_/g, " ").replace(/\b\w/g, (c) => c.toUpperCase());
 }
 
+/**
+ * Inline Revert button rendered on a tool card when the matching
+ * ``AgentActionLog`` row is reversible and hasn't been reverted yet.
+ * Reads from the SSE side-channel atom keyed by the synthetic
+ * ``toolCallId`` so it lights up even when ``GET /threads/.../actions``
+ * is gated behind ``SURFSENSE_ENABLE_ACTION_LOG=False`` (503).
+ */
+function ToolCardRevertButton({ toolCallId }: { toolCallId: string }) {
+	const session = useAtomValue(chatSessionStateAtom);
+	const actionMap = useAtomValue(agentActionByToolCallIdAtom);
+	const markReverted = useSetAtom(markAgentActionRevertedAtom);
+	const action = actionMap.get(toolCallId);
+	const [isReverting, setIsReverting] = useState(false);
+	const [confirmOpen, setConfirmOpen] = useState(false);
+
+	if (!action) return null;
+	if (!action.reversible) return null;
+	if (action.revertedByActionId !== null) return null;
+	if (action.isRevertAction) return null;
+	if (action.error) return null;
+	const threadId = session?.threadId;
+	if (!threadId) return null;
+
+	const handleRevert = async () => {
+		setIsReverting(true);
+		try {
+			const response = await agentActionsApiService.revert(threadId, action.id);
+			markReverted({ id: action.id, newActionId: response.new_action_id ?? null });
+			toast.success(response.message || "Action reverted.");
+		} catch (err) {
+			// 503 means revert is gated off on this deployment — hide the
+			// button silently rather than nagging the user. Any other error
+			// is surfaced as a toast so the operator can investigate.
+			if (err instanceof AppError && err.status === 503) {
+				return;
+			}
+			const message =
+				err instanceof AppError
+					? err.message
+					: err instanceof Error
+						? err.message
+						: "Failed to revert action.";
+			toast.error(message);
+		} finally {
+			setIsReverting(false);
+			setConfirmOpen(false);
+		}
+	};
+
+	return (
+		<AlertDialog open={confirmOpen} onOpenChange={setConfirmOpen}>
+			<AlertDialogTrigger asChild>
+				<Button
+					size="sm"
+					variant="outline"
+					className="gap-1.5"
+					onClick={(e) => {
+						e.stopPropagation();
+						setConfirmOpen(true);
+					}}
+				>
+					<RotateCcw className="size-3.5" />
+					Revert
+				</Button>
+			</AlertDialogTrigger>
+			<AlertDialogContent>
+				<AlertDialogHeader>
+					<AlertDialogTitle>Revert this action?</AlertDialogTitle>
+					<AlertDialogDescription>
+						This will undo <span className="font-medium">{formatToolName(action.toolName)}</span>{" "}
+						and append a new audit entry. Chat history is preserved — only the tool's effects on
+						your knowledge base or connectors will be reversed where possible.
+					</AlertDialogDescription>
+				</AlertDialogHeader>
+				<AlertDialogFooter>
+					<AlertDialogCancel disabled={isReverting}>Cancel</AlertDialogCancel>
+					<AlertDialogAction
+						onClick={(e) => {
+							e.preventDefault();
+							handleRevert();
+						}}
+						disabled={isReverting}
+					>
+						{isReverting ? "Reverting…" : "Revert"}
+					</AlertDialogAction>
+				</AlertDialogFooter>
+			</AlertDialogContent>
+		</AlertDialog>
+	);
+}
+
 const DefaultToolFallbackInner: ToolCallMessagePartComponent = ({
+	toolCallId,
 	toolName,
 	argsText,
 	result,
@@ -145,6 +258,9 @@ const DefaultToolFallbackInner: ToolCallMessagePartComponent = ({
 								</div>
 							</>
 						)}
+						<div className="flex justify-end">
+							<ToolCardRevertButton toolCallId={toolCallId} />
+						</div>
 					</div>
 				</>
 			)}