From ea7bcebcd05c7fb4cc2452d5b265934bd64db707 Mon Sep 17 00:00:00 2001
From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com>
Date: Mon, 13 Apr 2026 20:19:23 +0530
Subject: [PATCH] refactor: integrate HITL approval UI for interrupt results

Enhanced the NewChatPage to utilize the new GenericHitlApprovalToolUI for handling interrupt results. Updated the ToolFallback component to conditionally render the approval UI based on the result type. Additionally, introduced a new GenericHitlApprovalToolUI component to manage user approvals and parameter editing for tool actions.
---
 .../new-chat/[[...chat_id]]/page.tsx          |   4 +-
 .../components/assistant-ui/tool-fallback.tsx |  11 +-
 .../tool-ui/generic-hitl-approval.tsx         | 268 ++++++++++++++++++
 3 files changed, 280 insertions(+), 3 deletions(-)
 create mode 100644 surfsense_web/components/tool-ui/generic-hitl-approval.tsx

diff --git a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx
index 0b1369340..58eb58f4b 100644
--- a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx
+++ b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx
@@ -798,7 +798,7 @@ export default function NewChatPage() {
 									});
 								} else {
 									const tcId = `interrupt-${action.name}`;
-									addToolCall(contentPartsState, TOOLS_WITH_UI, tcId, action.name, action.args);
+									addToolCall(contentPartsState, TOOLS_WITH_UI, tcId, action.name, action.args, true);
 									updateToolCall(contentPartsState, tcId, {
 										result: { __interrupt__: true, ...interruptData },
 									});
@@ -1125,7 +1125,7 @@ export default function NewChatPage() {
 									});
 								} else {
 									const tcId = `interrupt-${action.name}`;
-									addToolCall(contentPartsState, TOOLS_WITH_UI, tcId, action.name, action.args);
+									addToolCall(contentPartsState, TOOLS_WITH_UI, tcId, action.name, action.args, true);
 									updateToolCall(contentPartsState, tcId, {
 										result: {
 											__interrupt__: true,
diff --git a/surfsense_web/components/assistant-ui/tool-fallback.tsx b/surfsense_web/components/assistant-ui/tool-fallback.tsx
index b658dba6d..d9833b387 100644
--- a/surfsense_web/components/assistant-ui/tool-fallback.tsx
+++ b/surfsense_web/components/assistant-ui/tool-fallback.tsx
@@ -1,14 +1,16 @@
 import type { ToolCallMessagePartComponent } from "@assistant-ui/react";
 import { CheckIcon, ChevronDownIcon, ChevronUpIcon, XCircleIcon } from "lucide-react";
 import { useMemo, useState } from "react";
+import { GenericHitlApprovalToolUI } from "@/components/tool-ui/generic-hitl-approval";
 import { getToolIcon } from "@/contracts/enums/toolIcons";
+import { isInterruptResult } from "@/lib/hitl";
 import { cn } from "@/lib/utils";
 
 function formatToolName(name: string): string {
 	return name.replace(/_/g, " ").replace(/\b\w/g, (c) => c.toUpperCase());
 }
 
-export const ToolFallback: ToolCallMessagePartComponent = ({
+const DefaultToolFallbackInner: ToolCallMessagePartComponent = ({
 	toolName,
 	argsText,
 	result,
@@ -145,3 +147,10 @@ export const ToolFallback: ToolCallMessagePartComponent = ({
 		</div>
 	);
 };
+
+export const ToolFallback: ToolCallMessagePartComponent = (props) => {
+	if (isInterruptResult(props.result)) {
+		return <GenericHitlApprovalToolUI {...props} />;
+	}
+	return <DefaultToolFallbackInner {...props} />;
+};
diff --git a/surfsense_web/components/tool-ui/generic-hitl-approval.tsx b/surfsense_web/components/tool-ui/generic-hitl-approval.tsx
new file mode 100644
index 000000000..48d0d3764
--- /dev/null
+++ b/surfsense_web/components/tool-ui/generic-hitl-approval.tsx
@@ -0,0 +1,268 @@
+"use client";
+
+import type { ToolCallMessagePartComponent } from "@assistant-ui/react";
+import { CornerDownLeftIcon, ShieldAlertIcon, ShieldCheckIcon } from "lucide-react";
+import { useCallback, useEffect, useMemo, useState } from "react";
+import { TextShimmerLoader } from "@/components/prompt-kit/loader";
+import { Button } from "@/components/ui/button";
+import { Input } from "@/components/ui/input";
+import { Textarea } from "@/components/ui/textarea";
+import { useHitlPhase } from "@/hooks/use-hitl-phase";
+import { connectorsApiService } from "@/lib/apis/connectors-api.service";
+import { isInterruptResult, useHitlDecision } from "@/lib/hitl";
+import type { HitlDecision, InterruptResult } from "@/lib/hitl";
+
+function ParamEditor({
+	params,
+	onChange,
+	disabled,
+}: {
+	params: Record<string, unknown>;
+	onChange: (updated: Record<string, unknown>) => void;
+	disabled: boolean;
+}) {
+	const entries = Object.entries(params);
+	if (entries.length === 0) return null;
+
+	return (
+		<div className="space-y-2">
+			{entries.map(([key, value]) => {
+				const strValue = value == null ? "" : String(value);
+				const isLong = strValue.length > 120;
+				const fieldId = `hitl-param-${key}`;
+
+				return (
+					<div key={key} className="space-y-1">
+						<label htmlFor={fieldId} className="text-xs font-medium text-muted-foreground">
+							{key}
+						</label>
+						{isLong ? (
+							<Textarea
+								id={fieldId}
+								value={strValue}
+								disabled={disabled}
+								rows={3}
+								onChange={(e) => onChange({ ...params, [key]: e.target.value })}
+								className="text-xs"
+							/>
+						) : (
+							<Input
+								id={fieldId}
+								value={strValue}
+								disabled={disabled}
+								onChange={(e) => onChange({ ...params, [key]: e.target.value })}
+								className="text-xs"
+							/>
+						)}
+					</div>
+				);
+			})}
+		</div>
+	);
+}
+
+function GenericApprovalCard({
+	toolName,
+	args,
+	interruptData,
+	onDecision,
+}: {
+	toolName: string;
+	args: Record<string, unknown>;
+	interruptData: InterruptResult;
+	onDecision: (decision: HitlDecision) => void;
+}) {
+	const { phase, setProcessing, setRejected } = useHitlPhase(interruptData);
+	const [editedParams, setEditedParams] = useState<Record<string, unknown>>(args);
+	const [isEditing, setIsEditing] = useState(false);
+
+	const displayName = toolName.replace(/_/g, " ").replace(/\b\w/g, (c) => c.toUpperCase());
+
+	const mcpServer = interruptData.context?.mcp_server as string | undefined;
+	const toolDescription = interruptData.context?.tool_description as string | undefined;
+	const mcpConnectorId = interruptData.context?.mcp_connector_id as number | undefined;
+	const isMCPTool = mcpConnectorId != null;
+
+	const reviewConfig = interruptData.review_configs?.[0];
+	const allowedDecisions = reviewConfig?.allowed_decisions ?? ["approve", "reject"];
+	const canEdit = allowedDecisions.includes("edit");
+
+	const hasChanged = useMemo(() => {
+		return JSON.stringify(editedParams) !== JSON.stringify(args);
+	}, [editedParams, args]);
+
+	const handleApprove = useCallback(() => {
+		if (phase !== "pending") return;
+		const isEdited = isEditing && hasChanged;
+		setProcessing();
+		onDecision({
+			type: isEdited ? "edit" : "approve",
+			edited_action: isEdited
+				? { name: interruptData.action_requests[0]?.name ?? toolName, args: editedParams }
+				: undefined,
+		});
+	}, [
+		phase,
+		setProcessing,
+		isEditing,
+		hasChanged,
+		onDecision,
+		interruptData,
+		toolName,
+		editedParams,
+	]);
+
+	const handleAlwaysAllow = useCallback(() => {
+		if (phase !== "pending" || !isMCPTool) return;
+		setProcessing();
+		onDecision({ type: "approve" });
+		connectorsApiService.trustMCPTool(mcpConnectorId, toolName).catch((err) => {
+			console.error("Failed to trust MCP tool:", err);
+		});
+	}, [phase, setProcessing, onDecision, isMCPTool, mcpConnectorId, toolName]);
+
+	useEffect(() => {
+		const handler = (e: KeyboardEvent) => {
+			if (e.key === "Enter" && !e.shiftKey && !e.ctrlKey && !e.metaKey && phase === "pending") {
+				handleApprove();
+			}
+		};
+		window.addEventListener("keydown", handler);
+		return () => window.removeEventListener("keydown", handler);
+	}, [handleApprove, phase]);
+
+	return (
+		<div className="my-4 max-w-lg overflow-hidden rounded-2xl border bg-muted/30 transition-[box-shadow] duration-300">
+			{/* Header */}
+			<div className="flex items-start gap-3 px-5 pt-5 pb-4 select-none">
+				<div className="flex size-8 shrink-0 items-center justify-center rounded-lg bg-amber-500/10">
+					<ShieldAlertIcon className="size-4 text-amber-600" />
+				</div>
+				<div className="flex-1 min-w-0">
+					<p className="text-sm font-semibold text-foreground">
+						{phase === "rejected"
+							? `${displayName} — Rejected`
+							: phase === "processing" || phase === "complete"
+								? `${displayName} — Approved`
+								: `Approve: ${displayName}`}
+					</p>
+					{phase === "processing" ? (
+						<TextShimmerLoader text="Executing..." size="sm" />
+					) : phase === "complete" ? (
+						<p className="text-xs text-muted-foreground mt-0.5">Action completed</p>
+					) : phase === "rejected" ? (
+						<p className="text-xs text-muted-foreground mt-0.5">Action was cancelled</p>
+					) : (
+						<p className="text-xs text-muted-foreground mt-0.5">
+							Requires your approval to proceed
+						</p>
+					)}
+					{mcpServer && (
+						<p className="text-[10px] text-muted-foreground/70 mt-1">
+							via <span className="font-medium">{mcpServer}</span>
+						</p>
+					)}
+				</div>
+			</div>
+
+			{/* Description */}
+			{toolDescription && phase === "pending" && (
+				<>
+					<div className="mx-5 h-px bg-border/50" />
+					<div className="px-5 py-2">
+						<p className="text-xs text-muted-foreground">{toolDescription}</p>
+					</div>
+				</>
+			)}
+
+			{/* Parameters */}
+			{Object.keys(args).length > 0 && (
+				<>
+					<div className="mx-5 h-px bg-border/50" />
+					<div className="px-5 py-3 space-y-2">
+						<p className="text-xs font-medium text-muted-foreground">Parameters</p>
+						{phase === "pending" && isEditing ? (
+							<ParamEditor
+								params={editedParams}
+								onChange={setEditedParams}
+								disabled={phase !== "pending"}
+							/>
+						) : (
+							<pre className="text-xs text-foreground/80 whitespace-pre-wrap break-all bg-muted/50 rounded-lg p-3">
+								{JSON.stringify(args, null, 2)}
+							</pre>
+						)}
+					</div>
+				</>
+			)}
+
+			{/* Action buttons */}
+			{phase === "pending" && (
+				<>
+					<div className="mx-5 h-px bg-border/50" />
+					<div className="px-5 py-3 flex items-center gap-2 select-none flex-wrap">
+						{allowedDecisions.includes("approve") && (
+							<Button size="sm" className="rounded-lg gap-1.5" onClick={handleApprove}>
+								{isEditing && hasChanged ? "Approve with edits" : "Approve"}
+								<CornerDownLeftIcon className="size-3 opacity-60" />
+							</Button>
+						)}
+						{isMCPTool && (
+							<Button
+								size="sm"
+								variant="outline"
+								className="rounded-lg gap-1.5"
+								onClick={handleAlwaysAllow}
+							>
+								<ShieldCheckIcon className="size-3" />
+								Always Allow
+							</Button>
+						)}
+						{canEdit && !isEditing && (
+							<Button
+								size="sm"
+								variant="outline"
+								className="rounded-lg"
+								onClick={() => setIsEditing(true)}
+							>
+								Edit
+							</Button>
+						)}
+						{allowedDecisions.includes("reject") && (
+							<Button
+								size="sm"
+								variant="ghost"
+								className="rounded-lg text-muted-foreground"
+								onClick={() => {
+									setRejected();
+									onDecision({ type: "reject", message: "User rejected the action." });
+								}}
+							>
+								Reject
+							</Button>
+						)}
+					</div>
+				</>
+			)}
+		</div>
+	);
+}
+
+export const GenericHitlApprovalToolUI: ToolCallMessagePartComponent = ({
+	toolName,
+	args,
+	result,
+}) => {
+	const { dispatch } = useHitlDecision();
+
+	if (!result || !isInterruptResult(result)) return null;
+
+	return (
+		<GenericApprovalCard
+			toolName={toolName}
+			args={args as Record<string, unknown>}
+			interruptData={result}
+			onDecision={(decision) => dispatch([decision])}
+		/>
+	);
+};