From c1a28d54c2c3cb0041fd9d7016d5ce372b0c2e76 Mon Sep 17 00:00:00 2001
From: majdyz <zamil.majdy@agpt.co>
Date: Wed, 8 Apr 2026 18:41:58 +0700
Subject: [PATCH] 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
---
 .../BuilderChatPanel/BuilderChatPanel.tsx     |  34 +++++-
 .../__tests__/BuilderChatPanel.test.tsx       |  31 ++++-
 .../__tests__/helpers.test.ts                 | 111 ++++++++++++++++++
 .../components/BuilderChatPanel/helpers.ts    |  43 +++++--
 .../BuilderChatPanel/useBuilderChatPanel.ts   |  67 +++++------
 5 files changed, 233 insertions(+), 53 deletions(-)
 create mode 100644 autogpt_platform/frontend/src/app/(platform)/build/components/BuilderChatPanel/__tests__/helpers.test.ts

diff --git a/autogpt_platform/frontend/src/app/(platform)/build/components/BuilderChatPanel/BuilderChatPanel.tsx b/autogpt_platform/frontend/src/app/(platform)/build/components/BuilderChatPanel/BuilderChatPanel.tsx
index cd763c5756..c38772a145 100644
--- a/autogpt_platform/frontend/src/app/(platform)/build/components/BuilderChatPanel/BuilderChatPanel.tsx
+++ b/autogpt_platform/frontend/src/app/(platform)/build/components/BuilderChatPanel/BuilderChatPanel.tsx
@@ -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>
   );
 }
diff --git a/autogpt_platform/frontend/src/app/(platform)/build/components/BuilderChatPanel/__tests__/BuilderChatPanel.test.tsx b/autogpt_platform/frontend/src/app/(platform)/build/components/BuilderChatPanel/__tests__/BuilderChatPanel.test.tsx
index 5c0401c410..64f01fb8e7 100644
--- a/autogpt_platform/frontend/src/app/(platform)/build/components/BuilderChatPanel/__tests__/BuilderChatPanel.test.tsx
+++ b/autogpt_platform/frontend/src/app/(platform)/build/components/BuilderChatPanel/__tests__/BuilderChatPanel.test.tsx
@@ -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", () => {
diff --git a/autogpt_platform/frontend/src/app/(platform)/build/components/BuilderChatPanel/__tests__/helpers.test.ts b/autogpt_platform/frontend/src/app/(platform)/build/components/BuilderChatPanel/__tests__/helpers.test.ts
new file mode 100644
index 0000000000..373971cd5b
--- /dev/null
+++ b/autogpt_platform/frontend/src/app/(platform)/build/components/BuilderChatPanel/__tests__/helpers.test.ts
@@ -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;");
+  });
+});
diff --git a/autogpt_platform/frontend/src/app/(platform)/build/components/BuilderChatPanel/helpers.ts b/autogpt_platform/frontend/src/app/(platform)/build/components/BuilderChatPanel/helpers.ts
index 7fd30108a7..d1ab5dae9f 100644
--- a/autogpt_platform/frontend/src/app/(platform)/build/components/BuilderChatPanel/helpers.ts
+++ b/autogpt_platform/frontend/src/app/(platform)/build/components/BuilderChatPanel/helpers.ts
@@ -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");
 }
diff --git a/autogpt_platform/frontend/src/app/(platform)/build/components/BuilderChatPanel/useBuilderChatPanel.ts b/autogpt_platform/frontend/src/app/(platform)/build/components/BuilderChatPanel/useBuilderChatPanel.ts
index 7e6a96ad1c..02c9521748 100644
--- a/autogpt_platform/frontend/src/app/(platform)/build/components/BuilderChatPanel/useBuilderChatPanel.ts
+++ b/autogpt_platform/frontend/src/app/(platform)/build/components/BuilderChatPanel/useBuilderChatPanel.ts
@@ -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,
   };