mirror of
https://github.com/Significant-Gravitas/AutoGPT.git
synced 2026-04-30 03:00:41 -04:00
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:
@@ -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>
|
||||
);
|
||||
}
|
||||
|
||||
@@ -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", () => {
|
||||
|
||||
@@ -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("<script>");
|
||||
});
|
||||
|
||||
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("<injection>");
|
||||
});
|
||||
});
|
||||
@@ -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, "<").replace(/>/g, ">");
|
||||
}
|
||||
|
||||
/**
|
||||
* 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");
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
};
|
||||
|
||||
Reference in New Issue
Block a user