fix(frontend/builder): require manual action confirmation and prevent prompt injection

- Replace auto-apply with per-action Apply buttons; users must explicitly
  confirm each AI suggestion before the graph is mutated
- Accumulate parsedActions across all assistant messages so multi-turn
  suggestions remain visible rather than disappearing after the next turn
- Escape < and > in node names/descriptions before embedding in XML prompt
  context to prevent AI prompt injection via crafted node labels
- Add MAX_EDGES cap (200) in serializeGraphForChat to mirror the MAX_NODES
  limit and prevent token overruns on dense graphs
- Add Escape key handler in the hook to close the chat panel
- Add helpers.test.ts with unit tests for buildSeedPrompt,
  extractTextFromParts, and XML sanitization
This commit is contained in:
majdyz
2026-04-08 18:41:58 +07:00
parent b8cb9a506f
commit c1a28d54c2
5 changed files with 233 additions and 53 deletions

View File

@@ -34,6 +34,8 @@ export function BuilderChatPanel({ className, isGraphLoaded }: Props) {
sessionId,
nodes,
parsedActions,
appliedActionKeys,
handleApplyAction,
seedMessageId,
} = useBuilderChatPanel({ isGraphLoaded });
@@ -91,6 +93,8 @@ export function BuilderChatPanel({ className, isGraphLoaded }: Props) {
streamError={error}
nodes={nodes}
parsedActions={parsedActions}
appliedActionKeys={appliedActionKeys}
onApplyAction={handleApplyAction}
seedMessageId={seedMessageId}
messagesEndRef={messagesEndRef}
/>
@@ -147,6 +151,8 @@ interface MessageListProps {
streamError: Error | undefined;
nodes: CustomNode[];
parsedActions: GraphAction[];
appliedActionKeys: Set<string>;
onApplyAction: (action: GraphAction) => void;
seedMessageId: string | null;
messagesEndRef: React.RefObject<HTMLDivElement>;
}
@@ -158,6 +164,8 @@ function MessageList({
streamError,
nodes,
parsedActions,
appliedActionKeys,
onApplyAction,
seedMessageId,
messagesEndRef,
}: MessageListProps) {
@@ -246,14 +254,17 @@ function MessageList({
{parsedActions.length > 0 && (
<div className="space-y-2 rounded-lg border border-violet-100 bg-violet-50 p-3">
<p className="text-xs font-medium text-violet-700">
AI applied these changes
Suggested changes
</p>
{parsedActions.map((action) => {
const key = getActionKey(action);
return (
<ActionItem
key={getActionKey(action)}
key={key}
action={action}
nodes={nodes}
isApplied={appliedActionKeys.has(key)}
onApply={onApplyAction}
/>
);
})}
@@ -268,9 +279,13 @@ function MessageList({
function ActionItem({
action,
nodes,
isApplied,
onApply,
}: {
action: GraphAction;
nodes: CustomNode[];
isApplied: boolean;
onApply: (action: GraphAction) => void;
}) {
const nodeName = (id: string) =>
nodes.find((n) => n.id === id)?.data.metadata?.customized_name ||
@@ -285,9 +300,18 @@ function ActionItem({
return (
<div className="flex items-start justify-between gap-2 rounded bg-white p-2 text-xs shadow-sm">
<span className="leading-tight text-slate-700">{label}</span>
<span className="shrink-0 rounded bg-green-100 px-2 py-0.5 text-xs font-medium text-green-700">
Applied
</span>
{isApplied ? (
<span className="shrink-0 rounded bg-green-100 px-2 py-0.5 text-xs font-medium text-green-700">
Applied
</span>
) : (
<button
onClick={() => onApply(action)}
className="shrink-0 rounded bg-violet-100 px-2 py-0.5 text-xs font-medium text-violet-700 hover:bg-violet-200"
>
Apply
</button>
)}
</div>
);
}

View File

@@ -39,6 +39,7 @@ function makeMockHook(
sessionId: null,
nodes: [],
parsedActions: [],
appliedActionKeys: new Set<string>(),
handleApplyAction: vi.fn(),
seedMessageId: null,
...overrides,
@@ -119,7 +120,7 @@ describe("BuilderChatPanel", () => {
expect(screen.getByText("This agent searches the web.")).toBeDefined();
});
it("renders applied actions section when parsedActions are present", () => {
it("renders suggested changes section when parsedActions are present", () => {
mockUseBuilderChatPanel.mockReturnValue(
makeMockHook({
isOpen: true,
@@ -134,11 +135,11 @@ describe("BuilderChatPanel", () => {
}),
);
render(<BuilderChatPanel />);
expect(screen.getByText("AI applied these changes")).toBeDefined();
expect(screen.getByText("Applied")).toBeDefined();
expect(screen.getByText("Suggested changes")).toBeDefined();
expect(screen.getByText("Apply")).toBeDefined();
});
it("shows applied badge for actions", () => {
it("shows Apply button for unapplied actions and Applied badge for applied actions", () => {
const action = {
type: "update_node_input" as const,
nodeId: "1",
@@ -149,10 +150,32 @@ describe("BuilderChatPanel", () => {
makeMockHook({
isOpen: true,
parsedActions: [action],
appliedActionKeys: new Set(["1:query"]),
}),
);
render(<BuilderChatPanel />);
expect(screen.getByText("Applied")).toBeDefined();
expect(screen.queryByText("Apply")).toBeNull();
});
it("calls handleApplyAction when Apply button is clicked", () => {
const handleApplyAction = vi.fn();
const action = {
type: "update_node_input" as const,
nodeId: "1",
key: "query",
value: "AI news",
};
mockUseBuilderChatPanel.mockReturnValue(
makeMockHook({
isOpen: true,
parsedActions: [action],
handleApplyAction,
}),
);
render(<BuilderChatPanel />);
fireEvent.click(screen.getByText("Apply"));
expect(handleApplyAction).toHaveBeenCalledWith(action);
});
it("calls sendMessage when the user submits a message", () => {

View File

@@ -0,0 +1,111 @@
import { describe, expect, it } from "vitest";
import {
buildSeedPrompt,
extractTextFromParts,
serializeGraphForChat,
} from "../helpers";
import type { CustomNode } from "../../FlowEditor/nodes/CustomNode/CustomNode";
describe("extractTextFromParts", () => {
it("returns empty string for empty array", () => {
expect(extractTextFromParts([])).toBe("");
});
it("concatenates text parts in order", () => {
const parts = [
{ type: "text", text: "Hello, " },
{ type: "text", text: "world!" },
];
expect(extractTextFromParts(parts)).toBe("Hello, world!");
});
it("ignores non-text parts", () => {
const parts = [
{ type: "text", text: "visible" },
{ type: "tool-call", text: "ignored" },
{ type: "text", text: " text" },
];
expect(extractTextFromParts(parts)).toBe("visible text");
});
it("returns empty string when all parts are non-text", () => {
const parts = [{ type: "tool-result" }, { type: "image" }];
expect(extractTextFromParts(parts)).toBe("");
});
});
describe("buildSeedPrompt", () => {
it("wraps the summary in <graph_context> tags", () => {
const result = buildSeedPrompt("some graph summary");
expect(result).toContain(
"<graph_context>\nsome graph summary\n</graph_context>",
);
});
it("includes instructions for update_node_input format", () => {
const result = buildSeedPrompt("");
expect(result).toContain('"action": "update_node_input"');
});
it("includes instructions for connect_nodes format", () => {
const result = buildSeedPrompt("");
expect(result).toContain('"action": "connect_nodes"');
});
it("ends with a question to prompt AI response", () => {
const result = buildSeedPrompt("");
expect(result.trim().endsWith("What does this agent do?")).toBe(true);
});
});
describe("serializeGraphForChat XML injection prevention", () => {
it("escapes < and > in node names before embedding in prompt", () => {
const nodes = [
{
id: "1",
data: {
title: "<script>alert(1)</script>",
description: "",
hardcodedValues: {},
inputSchema: {},
outputSchema: {},
uiType: 1,
block_id: "b1",
costs: [],
categories: [],
},
type: "custom" as const,
position: { x: 0, y: 0 },
},
] as unknown as CustomNode[];
const result = serializeGraphForChat(nodes, []);
expect(result).not.toContain("<script>");
expect(result).toContain("&lt;script&gt;");
});
it("escapes < and > in node descriptions", () => {
const nodes = [
{
id: "1",
data: {
title: "Node",
description: "desc with <injection>",
hardcodedValues: {},
inputSchema: {},
outputSchema: {},
uiType: 1,
block_id: "b1",
costs: [],
categories: [],
},
type: "custom" as const,
position: { x: 0, y: 0 },
},
] as unknown as CustomNode[];
const result = serializeGraphForChat(nodes, []);
expect(result).not.toContain("<injection>");
expect(result).toContain("&lt;injection&gt;");
});
});

View File

@@ -3,6 +3,13 @@ import type { CustomEdge } from "../FlowEditor/edges/CustomEdge";
/** Maximum nodes serialized into the AI context to prevent token overruns. */
const MAX_NODES = 100;
/** Maximum edges serialized into the AI context to prevent token overruns. */
const MAX_EDGES = 200;
/** Escapes XML special characters in user-controlled strings before embedding in prompts. */
function sanitizeForXml(s: string): string {
return s.replace(/</g, "&lt;").replace(/>/g, "&gt;");
}
/**
* Action emitted by the AI to edit the agent graph.
@@ -45,8 +52,14 @@ export function serializeGraphForChat(
const visibleNodes = nodes.slice(0, MAX_NODES);
const nodeLines = visibleNodes.map((n) => {
const name = n.data.metadata?.customized_name || n.data.title;
const desc = n.data.description ? `${n.data.description}` : "";
const name = sanitizeForXml(
(n.data.metadata?.customized_name as string | undefined) ||
n.data.title ||
"",
);
const desc = n.data.description
? `${sanitizeForXml(n.data.description)}`
: "";
return `- Node ${n.id}: "${name}"${desc}`;
});
@@ -55,21 +68,35 @@ export function serializeGraphForChat(
? `\n(${nodes.length - MAX_NODES} additional nodes not shown)`
: "";
const edgeLines = edges.map((e) => {
const visibleEdges = edges.slice(0, MAX_EDGES);
const edgeLines = visibleEdges.map((e) => {
const src = nodes.find((n) => n.id === e.source);
const tgt = nodes.find((n) => n.id === e.target);
const srcName =
src?.data.metadata?.customized_name || src?.data.title || e.source;
const tgtName =
tgt?.data.metadata?.customized_name || tgt?.data.title || e.target;
const srcName = sanitizeForXml(
(src?.data.metadata?.customized_name as string | undefined) ||
src?.data.title ||
e.source,
);
const tgtName = sanitizeForXml(
(tgt?.data.metadata?.customized_name as string | undefined) ||
tgt?.data.title ||
e.target,
);
return `- "${srcName}" (${e.sourceHandle}) → "${tgtName}" (${e.targetHandle})`;
});
const edgeTruncationNote =
edges.length > MAX_EDGES
? `\n(${edges.length - MAX_EDGES} additional connections not shown)`
: "";
const parts = [
`Blocks (${nodes.length}):\n${nodeLines.join("\n")}${truncationNote}`,
];
if (edgeLines.length > 0) {
parts.push(`Connections (${edges.length}):\n${edgeLines.join("\n")}`);
parts.push(
`Connections (${edges.length}):\n${edgeLines.join("\n")}${edgeTruncationNote}`,
);
}
return parts.join("\n\n");
}

View File

@@ -32,13 +32,12 @@ export function useBuilderChatPanel({
const [sessionId, setSessionId] = useState<string | null>(null);
const [isCreatingSession, setIsCreatingSession] = useState(false);
const [sessionError, setSessionError] = useState(false);
const [appliedActionKeys, setAppliedActionKeys] = useState<Set<string>>(
new Set(),
);
// Guards whether the seed message has been sent for this session.
const hasSentSeedMessageRef = useRef(false);
const sendMessageRef = useRef<SendMessageFn | null>(null);
const handleApplyActionRef = useRef<((action: GraphAction) => void) | null>(
null,
);
const prevStatusRef = useRef<string>("ready");
const [{ flowID }] = useQueryStates({ flowID: parseAsString });
const queryClient = useQueryClient();
@@ -53,6 +52,7 @@ export function useBuilderChatPanel({
useEffect(() => {
setSessionId(null);
setSessionError(false);
setAppliedActionKeys(new Set());
hasSentSeedMessageRef.current = false;
}, [flowID]);
@@ -125,7 +125,6 @@ export function useBuilderChatPanel({
// Keep a stable ref so the initialization effect can call sendMessage
// without including it in the deps array (avoids re-triggering the effect).
sendMessageRef.current = sendMessage;
handleApplyActionRef.current = handleApplyAction;
// ID of the seed message sent on panel open. It contains prompt-engineering
// instructions that should not be shown to the user.
@@ -134,44 +133,31 @@ export function useBuilderChatPanel({
return messages.find((m) => m.role === "user")?.id ?? null;
}, [messages]);
// Parsed actions from the last assistant message. Gated on `status ===
// "ready"` so the expensive regex parse only runs once per completed AI turn,
// not on every streaming chunk.
// Parsed actions from all assistant messages, accumulated across turns.
// Gated on `status === "ready"` so parsing only runs on completed turns.
const parsedActions = useMemo(() => {
if (status !== "ready") return [];
const assistantMessages = messages.filter((m) => m.role === "assistant");
const last = assistantMessages[assistantMessages.length - 1];
if (!last) return [];
const text = extractTextFromParts(last.parts);
const parsed = parseGraphActions(text);
const seen = new Set<string>();
return parsed.filter((action) => {
const key = getActionKey(action);
if (seen.has(key)) return false;
seen.add(key);
return true;
});
return messages
.filter((m) => m.role === "assistant")
.flatMap((msg) => parseGraphActions(extractTextFromParts(msg.parts)))
.filter((action) => {
const key = getActionKey(action);
if (seen.has(key)) return false;
seen.add(key);
return true;
});
}, [messages, status]);
// After each AI turn: apply parsed actions to the local graph and refresh
// the canvas from the server. Gating on parsedActions.length > 0 avoids an
// unnecessary refetch after read-only turns (e.g. the initial description).
// Close the panel on Escape so keyboard users can dismiss it quickly.
useEffect(() => {
const prev = prevStatusRef.current;
prevStatusRef.current = status;
if (
status === "ready" &&
(prev === "streaming" || prev === "submitted") &&
parsedActions.length > 0
) {
parsedActions.forEach((a) => handleApplyActionRef.current?.(a));
if (flowID) {
queryClient.invalidateQueries({
queryKey: getGetV1GetSpecificGraphQueryKey(flowID),
});
}
if (!isOpen) return;
function onKeyDown(e: KeyboardEvent) {
if (e.key === "Escape") setIsOpen(false);
}
}, [status, flowID, queryClient, parsedActions]);
document.addEventListener("keydown", onKeyDown);
return () => document.removeEventListener("keydown", onKeyDown);
}, [isOpen]);
// Send the seed message once per session. `nodes` and `edges` are included in
// the dep array so this effect always has fresh data; the hasSentSeedMessageRef
@@ -228,6 +214,14 @@ export function useBuilderChatPanel({
targetHandle: action.targetHandle,
type: "custom",
});
} else {
return;
}
setAppliedActionKeys((prev) => new Set([...prev, getActionKey(action)]));
if (flowID) {
queryClient.invalidateQueries({
queryKey: getGetV1GetSpecificGraphQueryKey(flowID),
});
}
}
@@ -244,6 +238,7 @@ export function useBuilderChatPanel({
sessionId,
nodes,
parsedActions,
appliedActionKeys,
handleApplyAction,
seedMessageId,
};