feat: implement Story 3.3 — badge-based element referencing for targeted modifications

Adds badge chips in the copilot chat input that reference selected diagram
elements, enabling scoped AI modifications. Includes code review fixes for
reduced-motion support, scope indicator, callback stability, schema validation,
neighbor limits, and buildSelectedContext test coverage (103 tests passing).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

packages/ai/src/modules/copilot/api.ts

@@ -22,7 +22,7 @@ import {
   validateNodeTypes,
   validateUniqueIds,
 } from "./mutation-schema";
-import { buildCopilotSystemPrompt } from "./system-prompt";
+import { buildCopilotSystemPrompt, buildSelectedContext } from "./system-prompt";
 
 import type { CopilotMessagePayload } from "./schema";
 import type { DiagramType } from "./types";
@@ -113,6 +113,7 @@ export const streamCopilot = async ({
   diagramId,
   diagramType,
   graphContext,
+  selectedElements,
   userId,
   signal,
   ...msg
@@ -141,6 +142,10 @@ export const streamCopilot = async ({
 
   const systemPrompt = buildCopilotSystemPrompt(diagramType as DiagramType, {
     graphContext,
+    selectedElements,
+    selectedContext: selectedElements?.length
+      ? buildSelectedContext(selectedElements, graphContext)
+      : undefined,
   });
 
   const stream = createUIMessageStream({

packages/ai/src/modules/copilot/schema.test.ts

@@ -84,4 +84,64 @@ describe("copilotMessageSchema", () => {
     });
     expect(result.success).toBe(false);
   });
+
+  describe("selectedElements", () => {
+    it("should accept valid selectedElements array", () => {
+      const result = copilotMessageSchema.safeParse({
+        ...validPayload,
+        selectedElements: [
+          { id: "n1", type: "activity", label: "Process order" },
+          { id: "n2", type: "gateway-exclusive", label: "Valid?" },
+        ],
+      });
+      expect(result.success).toBe(true);
+      if (result.success) {
+        expect(result.data.selectedElements).toHaveLength(2);
+        expect(result.data.selectedElements![0]!.id).toBe("n1");
+      }
+    });
+
+    it("should accept empty selectedElements array", () => {
+      const result = copilotMessageSchema.safeParse({
+        ...validPayload,
+        selectedElements: [],
+      });
+      expect(result.success).toBe(true);
+      if (result.success) {
+        expect(result.data.selectedElements).toHaveLength(0);
+      }
+    });
+
+    it("should accept missing selectedElements (optional)", () => {
+      const result = copilotMessageSchema.safeParse(validPayload);
+      expect(result.success).toBe(true);
+      if (result.success) {
+        expect(result.data.selectedElements).toBeUndefined();
+      }
+    });
+
+    it("should reject selectedElements with missing required fields", () => {
+      const result = copilotMessageSchema.safeParse({
+        ...validPayload,
+        selectedElements: [{ id: "n1" }],
+      });
+      expect(result.success).toBe(false);
+    });
+
+    it("should reject non-array selectedElements", () => {
+      const result = copilotMessageSchema.safeParse({
+        ...validPayload,
+        selectedElements: "n1",
+      });
+      expect(result.success).toBe(false);
+    });
+
+    it("should reject selectedElements with empty string fields", () => {
+      const result = copilotMessageSchema.safeParse({
+        ...validPayload,
+        selectedElements: [{ id: "", type: "activity", label: "Test" }],
+      });
+      expect(result.success).toBe(false);
+    });
+  });
 });

packages/ai/src/modules/copilot/schema.ts

@@ -2,12 +2,19 @@ import * as z from "zod";
 
 import { DIAGRAM_TYPES } from "./types";
 
+export const selectedElementSchema = z.object({
+  id: z.string().min(1),
+  type: z.string().min(1),
+  label: z.string().min(1),
+});
+
 export const copilotMessageSchema = z.object({
   id: z.string(),
   chatId: z.string(),
   diagramId: z.string(),
   diagramType: z.enum(DIAGRAM_TYPES as [string, ...string[]]),
   graphContext: z.string().optional(),
+  selectedElements: z.array(selectedElementSchema).optional(),
   parts: z.array(
     z.object({
       type: z.literal("text"),

packages/ai/src/modules/copilot/system-prompt.test.ts

@@ -1,6 +1,6 @@
 import { describe, expect, it } from "vitest";
 
-import { buildCopilotSystemPrompt } from "./system-prompt";
+import { buildCopilotSystemPrompt, buildSelectedContext } from "./system-prompt";
 
 import type { DiagramType } from "./types";
 
@@ -173,4 +173,153 @@ describe("buildCopilotSystemPrompt", () => {
     expect(prompt).toContain('"n1"');
     expect(prompt).toContain('"e1"');
   });
+
+  describe("scoped context — selectedElements", () => {
+    const selectedElements = [
+      { id: "n1", type: "activity", label: "Validate order" },
+      { id: "n2", type: "gateway-exclusive", label: "Valid?" },
+    ];
+
+    it("should include scoped context section when selectedElements provided", () => {
+      const prompt = buildCopilotSystemPrompt("bpmn", { selectedElements });
+      expect(prompt).toContain("Scoped context");
+      expect(prompt).toContain("targeted modification");
+    });
+
+    it("should list all selected elements by label and type", () => {
+      const prompt = buildCopilotSystemPrompt("bpmn", { selectedElements });
+      expect(prompt).toContain("**Validate order** (activity)");
+      expect(prompt).toContain("**Valid?** (gateway-exclusive)");
+    });
+
+    it("should show element count", () => {
+      const prompt = buildCopilotSystemPrompt("bpmn", { selectedElements });
+      expect(prompt).toContain("2 element(s)");
+    });
+
+    it("should include selectedContext JSON when provided", () => {
+      const selectedContext = JSON.stringify({
+        selectedNodes: selectedElements,
+        connectedEdges: [{ id: "e1", from: "n1", to: "n2" }],
+      });
+      const prompt = buildCopilotSystemPrompt("bpmn", {
+        selectedElements,
+        selectedContext,
+      });
+      expect(prompt).toContain("Selected element details");
+      expect(prompt).toContain(selectedContext);
+    });
+
+    it("should include preservation instructions for scoped mode", () => {
+      const prompt = buildCopilotSystemPrompt("bpmn", { selectedElements });
+      expect(prompt).toContain("Include ALL nodes and edges");
+      expect(prompt).toContain("Preserve all other nodes and edges unchanged");
+      expect(prompt).toContain("minimal changes");
+    });
+
+    it("should NOT include scoped context when selectedElements is empty", () => {
+      const prompt = buildCopilotSystemPrompt("bpmn", { selectedElements: [] });
+      expect(prompt).not.toContain("Scoped context");
+    });
+
+    it("should NOT include scoped context when selectedElements is undefined", () => {
+      const prompt = buildCopilotSystemPrompt("bpmn", {});
+      expect(prompt).not.toContain("Scoped context");
+    });
+
+    it("should include scoped context alongside graph context", () => {
+      const graphContext = '{"nodes":[{"id":"n1"}],"edges":[]}';
+      const prompt = buildCopilotSystemPrompt("bpmn", {
+        graphContext,
+        selectedElements,
+      });
+      expect(prompt).toContain(graphContext);
+      expect(prompt).toContain("Scoped context");
+    });
+  });
+});
+
+interface ContextResult {
+  selectedNodes: Array<{ id: string; [key: string]: unknown }>;
+  connectedEdges?: Array<{ id: string; from: string; to: string }>;
+  neighborNodes?: Array<{ id: string; label: string }>;
+}
+
+const parseContext = (raw: string): ContextResult => JSON.parse(raw) as ContextResult;
+
+describe("buildSelectedContext", () => {
+  const selectedElements = [
+    { id: "n1", type: "activity", label: "Validate order" },
+    { id: "n2", type: "gateway-exclusive", label: "Valid?" },
+  ];
+
+  it("should return selectedNodes-only JSON when no graphContext", () => {
+    const result = parseContext(buildSelectedContext(selectedElements));
+    expect(result.selectedNodes).toEqual(selectedElements);
+    expect(result.connectedEdges).toBeUndefined();
+    expect(result.neighborNodes).toBeUndefined();
+  });
+
+  it("should extract connected edges from graphContext", () => {
+    const graphContext = JSON.stringify({
+      nodes: [
+        { id: "n1", label: "Validate order" },
+        { id: "n2", label: "Valid?" },
+        { id: "n3", label: "Reject" },
+      ],
+      edges: [
+        { id: "e1", from: "n1", to: "n2" },
+        { id: "e2", from: "n2", to: "n3" },
+        { id: "e3", from: "n3", to: "n3" }, // self-loop on non-selected, not connected
+      ],
+    });
+    const result = parseContext(buildSelectedContext(selectedElements, graphContext));
+    expect(result.connectedEdges).toHaveLength(2);
+    expect(result.connectedEdges!.map((e) => e.id)).toEqual(["e1", "e2"]);
+  });
+
+  it("should extract 1-hop neighbor nodes", () => {
+    const graphContext = JSON.stringify({
+      nodes: [
+        { id: "n1", label: "Validate order" },
+        { id: "n2", label: "Valid?" },
+        { id: "n3", label: "Reject" },
+        { id: "n4", label: "Unconnected" },
+      ],
+      edges: [{ id: "e1", from: "n2", to: "n3" }],
+    });
+    const result = parseContext(buildSelectedContext(selectedElements, graphContext));
+    expect(result.neighborNodes).toHaveLength(1);
+    expect(result.neighborNodes![0]!.id).toBe("n3");
+    expect(result.neighborNodes![0]!.label).toBe("Reject");
+  });
+
+  it("should limit neighbor nodes to 10", () => {
+    const nodes = [
+      { id: "n1", label: "Selected" },
+      ...Array.from({ length: 15 }, (_, i) => ({ id: `nb${i}`, label: `Neighbor ${i}` })),
+    ];
+    const edges = Array.from({ length: 15 }, (_, i) => ({
+      id: `e${i}`,
+      from: "n1",
+      to: `nb${i}`,
+    }));
+    const graphContext = JSON.stringify({ nodes, edges });
+    const result = parseContext(
+      buildSelectedContext([{ id: "n1", type: "activity", label: "Selected" }], graphContext),
+    );
+    expect(result.neighborNodes).toHaveLength(10);
+  });
+
+  it("should handle malformed graphContext gracefully", () => {
+    const result = parseContext(buildSelectedContext(selectedElements, "not valid json"));
+    expect(result.selectedNodes).toEqual(selectedElements);
+  });
+
+  it("should handle graphContext with missing nodes/edges arrays", () => {
+    const result = parseContext(buildSelectedContext(selectedElements, "{}"));
+    expect(result.selectedNodes).toEqual([]);
+    expect(result.connectedEdges).toEqual([]);
+    expect(result.neighborNodes).toEqual([]);
+  });
 });

packages/ai/src/modules/copilot/system-prompt.ts

@@ -1,4 +1,4 @@
-import type { DiagramType } from "./types";
+import type { DiagramType, SelectedElement } from "./types";
 
 const DIAGRAM_DESCRIPTIONS: Record<DiagramType, string> = {
   bpmn: "BPMN (Business Process Model and Notation) — processes with activities, gateways, events, pools, and lanes",
@@ -67,12 +67,82 @@ const DIAGRAM_EXAMPLES: Record<DiagramType, string> = {
   flowchart: `{"meta":{"diagramType":"flowchart","title":"Validation Flow","layoutDirection":"DOWN"},"nodes":[{"id":"n1","type":"terminal","label":"Start"},{"id":"n2","type":"io","label":"Read input"},{"id":"n3","type":"decision","label":"Valid?"},{"id":"n4","type":"process","label":"Process data"},{"id":"n5","type":"terminal","label":"End"}],"edges":[{"id":"e1","from":"n1","to":"n2"},{"id":"e2","from":"n2","to":"n3"},{"id":"e3","from":"n3","to":"n4","label":"Yes"},{"id":"e4","from":"n3","to":"n2","label":"No"},{"id":"e5","from":"n4","to":"n5"}]}`,
 };
 
+const MAX_NEIGHBOR_NODES = 10;
+
+/** Extract selected nodes + their connected edges from the full graph context */
+export function buildSelectedContext(
+  selectedElements: SelectedElement[],
+  graphContext?: string,
+): string {
+  if (!graphContext) {
+    return JSON.stringify({ selectedNodes: selectedElements });
+  }
+
+  try {
+    const graph = JSON.parse(graphContext) as {
+      nodes?: Array<{ id: string; [key: string]: unknown }>;
+      edges?: Array<{ id: string; from: string; to: string; [key: string]: unknown }>;
+    };
+    const selectedIds = new Set(selectedElements.map((e) => e.id));
+
+    const selectedNodes = (graph.nodes ?? []).filter((n) => selectedIds.has(n.id));
+    const connectedEdges = (graph.edges ?? []).filter(
+      (e) => selectedIds.has(e.from) || selectedIds.has(e.to),
+    );
+
+    // 1-hop neighbor IDs (nodes connected to selected nodes but not selected themselves)
+    const neighborIds = new Set<string>();
+    for (const edge of connectedEdges) {
+      if (!selectedIds.has(edge.from)) neighborIds.add(edge.from);
+      if (!selectedIds.has(edge.to)) neighborIds.add(edge.to);
+    }
+    const neighborNodes = (graph.nodes ?? [])
+      .filter((n) => neighborIds.has(n.id))
+      .slice(0, MAX_NEIGHBOR_NODES)
+      .map((n) => ({ id: n.id, label: (n as { label?: string }).label ?? n.id }));
+
+    return JSON.stringify({ selectedNodes, connectedEdges, neighborNodes });
+  } catch {
+    return JSON.stringify({ selectedNodes: selectedElements });
+  }
+}
+
+function buildScopedContextSection(
+  selectedElements?: SelectedElement[],
+  selectedContext?: string,
+): string {
+  if (!selectedElements || selectedElements.length === 0) return "";
+
+  const elementList = selectedElements
+    .map((e) => `- **${e.label}** (${e.type})`)
+    .join("\n");
+
+  return `
+
+## Scoped context — targeted modification
+The user has selected ${selectedElements.length} element(s) for modification:
+${elementList}
+${selectedContext ? `\nSelected element details:\n\`\`\`json\n${selectedContext}\n\`\`\`` : ""}
+
+**IMPORTANT:** The user wants to modify ONLY these elements. When using the \`generateDiagram\` tool:
+- Include ALL nodes and edges in your output (the tool replaces the entire graph)
+- Focus changes on the selected elements and their immediate connections
+- Preserve all other nodes and edges unchanged
+- Prefer minimal changes: modify, split, merge, or restructure only the referenced elements`;
+}
+
 export function buildCopilotSystemPrompt(
   diagramType: DiagramType,
-  options?: { graphContext?: string },
+  options?: {
+    graphContext?: string;
+    selectedElements?: SelectedElement[];
+    selectedContext?: string;
+  },
 ): string {
   const description = DIAGRAM_DESCRIPTIONS[diagramType];
   const graphContext = options?.graphContext;
+  const selectedElements = options?.selectedElements;
+  const selectedContext = options?.selectedContext;
   const date = new Date().toLocaleDateString("en-US", {
     year: "numeric",
     month: "short",
@@ -141,5 +211,5 @@ ${graphContext ? `The diagram currently contains:\n\`\`\`json\n${graphContext}\n
 - Generate IDs as "n1", "n2", ... for nodes and "e1", "e2", ... for edges
 - Keep text responses concise and diagram-focused
 - Use markdown formatting (bold, lists, code blocks) — no h1 headings
-- Today's date is ${date}`;
+- Today's date is ${date}${buildScopedContextSection(selectedElements, selectedContext)}`;
 }

packages/ai/src/modules/copilot/types.ts

@@ -1,8 +1,13 @@
 import { diagramTypeEnum } from "@turbostarter/db/schema/diagram";
 
+import type { selectedElementSchema } from "./schema";
+import type * as z from "zod";
+
 export type DiagramType = (typeof diagramTypeEnum.enumValues)[number];
 
 export const DIAGRAM_TYPES = diagramTypeEnum.enumValues;
 
+export type SelectedElement = z.infer<typeof selectedElementSchema>;
+
 export type { GraphPatch } from "./mutation-schema";
 export { graphPatchSchema, validateGraphPatch, validateUniqueIds } from "./mutation-schema";