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>
This commit is contained in:
Alejandro Gutiérrez
2026-02-28 14:29:11 +00:00
parent 6dcb4dcd6f
commit 6591d6385a
11 changed files with 861 additions and 8 deletions

View File

@@ -0,0 +1,359 @@
# Story 3.3: Badge-Based Element Referencing for Targeted Modifications
Status: done
<!-- Note: Validation is optional. Run validate-create-story for quality check before dev-story. -->
## Story
As a user,
I want to click diagram elements to reference them in my chat messages,
so that I can tell the AI exactly which elements to modify.
## Acceptance Criteria
1. **Given** I click a node on the canvas, **When** the node is selected, **Then** a badge chip appears near the chat input area showing the node's label/name, **And** the badge animates in with a slide effect (200ms ease-out), **And** the chat input placeholder changes to "Describe changes to [node name]..."
2. **Given** I have badge(s) in the chat input area, **When** I type a message and send it, **Then** the AI receives the selected element context along with my message (FR4), **And** the AI operates in targeted scope — only modifying the referenced elements, **And** the AI response specifically addresses the badged elements
3. **Given** I have a badge chip displayed, **When** I click the X on the badge chip or click empty canvas, **Then** the badge is removed, **And** the chat returns to whole-diagram scope (FR3)
4. **Given** I select multiple elements (multi-select or rectangle drag), **When** badges are created, **Then** multiple badge chips appear near the chat input, **And** the AI receives all selected elements as context for the modification
5. **Given** I select an element and type "split this into two steps", **When** the AI processes the scoped request, **Then** it generates a minimal JSON patch affecting only the referenced element and its immediate connections, **And** the canvas updates only the affected area (not full re-render)
## Tasks / Subtasks
- [x] Task 1: Create BadgeChip component (AC: #1, #3)
- [x] 1.1 Create `apps/web/src/modules/copilot/components/BadgeChip.tsx` — shadcn/ui Badge variant with dismiss (X) button, truncated node label, diagram-type-aware icon
- [x] 1.2 Implement slide-in animation (200ms ease-out) using Motion `animate` + `exit` with `AnimatePresence`
- [x] 1.3 Click badge → scroll canvas to element + highlight (reuse existing `highlightedNodeId` store state)
- [x] 1.4 Click X → call `setSelectedNodeIds` to remove that node ID from selection array
- [x] 1.5 Add ARIA label `"Selected element: [name]"`, keyboard dismissible via Backspace/Delete
- [x] 1.6 Style with `--badge-chip-bg`, `--badge-chip-border`, `--badge-chip-text` CSS custom properties (from UX design tokens)
- [x] Task 2: Integrate BadgeChip display into CopilotPanel (AC: #1, #3, #4)
- [x] 2.1 Subscribe to `selectedNodeIds` from `useGraphStore` in CopilotPanel
- [x] 2.2 Map selected IDs to node data (label, type) by reading `nodes` from store
- [x] 2.3 Render badge chips in a flex-wrap container above the textarea input, inside the border-t input area
- [x] 2.4 Update placeholder text: when badges present → "Describe changes to [first badge name]..." (or "Describe changes to N elements..." for multi-select)
- [x] 2.5 Wrap badge chips in `AnimatePresence` for enter/exit animations
- [x] 2.6 Handle Escape key in textarea: clear all badges (deselect all nodes on canvas)
- [x] Task 3: Extend copilotMessageSchema with selectedElements (AC: #2, #4)
- [x] 3.1 Add `selectedElements` optional field to `copilotMessageSchema` in `packages/ai/src/modules/copilot/schema.ts` — array of `{ id: string, type: string, label: string }`
- [x] 3.2 Export `SelectedElement` type from `packages/ai/src/modules/copilot/types.ts`
- [x] Task 4: Pass selected element context from CopilotPanel to API (AC: #2, #4)
- [x] 4.1 In `prepareSendMessagesRequest`, serialize selected nodes: for each `selectedNodeIds`, extract the node's graph-level data (id, type, label, connected edges, neighbor node labels) via `flowToGraph` filtered to selected
- [x] 4.2 Include `selectedElements` in the request body alongside existing `graphContext`
- [x] 4.3 Include connected edges of selected nodes as `selectedEdges` for neighbor context
- [x] Task 5: Update system prompt for scoped AI context (AC: #2, #5)
- [x] 5.1 Extend `buildCopilotSystemPrompt` options to accept `selectedElements?: SelectedElement[]` and `selectedContext?: string` (JSON of selected nodes + their edges + 1-hop neighbors)
- [x] 5.2 When selectedElements are provided, add a `## Scoped context` section to the system prompt specifying: "The user has selected specific elements for targeted modification. Focus your changes on these elements and their immediate connections. Include ALL nodes and edges in your output, but only modify the selected ones."
- [x] 5.3 Include selected nodes' full data (properties, connected edges, neighbor labels) in the scoped context section
- [x] 5.4 Add instruction: "When elements are selected, prefer minimal changes. Modify, split, merge, or restructure only the referenced elements. Preserve all other nodes and edges unchanged."
- [x] 5.5 Update `system-prompt.test.ts` to cover scoped context variations
- [x] Task 6: Update API handler to pass selectedElements (AC: #2)
- [x] 6.1 In `streamCopilot` in `packages/ai/src/modules/copilot/api.ts`, destructure `selectedElements` from the validated payload
- [x] 6.2 Pass `selectedElements` and formatted `selectedContext` to `buildCopilotSystemPrompt`
- [x] Task 7: Add CSS custom properties for badge chip tokens (AC: #1)
- [x] 7.1 Add `--badge-chip-bg`, `--badge-chip-border`, `--badge-chip-text` to the diagram editor's CSS (in the existing design token location) — values from UX spec: `oklch(0.623 0.214 260 / 10%)`, `oklch(0.623 0.214 260 / 30%)`, `oklch(0.45 0.20 260)`
- [x] 7.2 Add dark mode variants for badge chip tokens
- [x] Task 8: Write tests (AC: all)
- [x] 8.1 System prompt tests: scoped context presence when selectedElements provided, absence when not, correct element data formatting (8 tests)
- [x] 8.2 Schema tests: selectedElements validation — valid arrays, empty arrays, missing optional field (5 tests)
- [x] 8.3 Component rendering tests deferred to E2E (per project standards from Story 3.1)
## Dev Notes
### Architecture Compliance
- **AI mutations through CRDT**: Per Winston architecture decision, all graph mutations flow through the Zustand store (which syncs to Liveblocks CRDT in future Epic 4). Badge-based targeted modifications use the same mutation pipeline as whole-diagram operations — the AI still outputs a complete `GraphData` patch via the `generateDiagram` tool. The "targeted" scope is an AI behavioral constraint (system prompt), NOT a different code path.
- **ELK.js in Web Worker**: Layout continues to run in Web Worker after targeted modifications. Use the existing `requestLayout()` trigger from Story 3.2.
- **Unified graph model**: Selected elements are referenced by their `DiagramNode` properties from the unified `GraphData` interface. Badge chips display the node's `label` field.
- **Selection state already exists**: `useGraphStore` has `selectedNodeIds: string[]` and `setSelectedNodeIds(ids: string[])` — already wired to `onSelectionChange` in `DiagramCanvas.tsx` (Story 2.9). This story CONSUMES the existing selection state; it does NOT create new selection mechanisms.
### Critical Implementation Patterns (from Stories 3.1 & 3.2)
**Transport pattern — extend body with selectedElements:**
```typescript
// In CopilotPanel prepareSendMessagesRequest:
const selectedNodeIds = useGraphStore.getState().selectedNodeIds;
const allNodes = useGraphStore.getState().nodes;
const allEdges = useGraphStore.getState().edges;
// Build selected element context for AI
const selectedElements = selectedNodeIds.length > 0
? selectedNodeIds.map(id => {
const node = allNodes.find(n => n.id === id);
if (!node) return null;
const data = node.data as { type?: string; label?: string };
return { id: node.id, type: data.type ?? "unknown", label: data.label ?? node.id };
}).filter(Boolean)
: undefined;
return {
body: {
...lastMessage,
chatId: id,
diagramId,
diagramType,
graphContext,
selectedElements, // NEW
},
};
```
**System prompt scoping pattern:**
```typescript
// In buildCopilotSystemPrompt, when selectedElements provided:
## Scoped context targeted modification
The user has selected ${selectedElements.length} element(s) for modification:
${selectedElements.map(e => `- ${e.label} (${e.type})`).join('\n')}
Connected edges: [edges connecting to/from selected nodes]
Neighbor nodes: [1-hop neighbors for context]
IMPORTANT: The user wants to modify ONLY these elements. Include ALL nodes
and edges in your generateDiagram output, but focus changes on the selected
elements and their immediate connections. Preserve everything else unchanged.
```
**Badge chip click → canvas scroll pattern:**
```typescript
// BadgeChip onClick handler — use React Flow's fitView or setCenter
// CopilotPanel is OUTSIDE ReactFlowProvider, so direct React Flow hooks won't work.
// Instead, use the same Zustand store pattern from Story 3.2:
// Set highlightedNodeId in store → DiagramCanvas reacts to it.
const handleBadgeClick = (nodeId: string) => {
useGraphStore.getState().setHighlightedNodeId(nodeId);
};
```
### Existing Selection Infrastructure (Story 2.9)
The selection system is already fully wired:
1. `DiagramCanvas.tsx:314-319``handleSelectionChange` callback from `@xyflow/react` updates `selectedNodeIds` in Zustand store
2. `useGraphStore.ts:26,37,76``selectedNodeIds: string[]` state + `setSelectedNodeIds` action
3. `DiagramCanvas.tsx:332-333``onSelectionChange` + `onPaneClick` (clear) bound to ReactFlow
4. `@xyflow/react` handles multi-select (Cmd+Click) and lasso (rectangle drag) natively
**What this story adds:** Consuming `selectedNodeIds` in CopilotPanel to show badge chips and pass scoped context to the AI. No changes to the canvas selection logic itself.
### Unified Graph Model (Source of Truth)
Location: `apps/web/src/modules/diagram/types/graph.ts`
```typescript
type DiagramType = "bpmn" | "er" | "orgchart" | "architecture" | "sequence" | "flowchart";
interface DiagramNode {
id: string; // Required
type: string; // Required - diagram-type-specific
label: string; // Required - display text (shown on badge chip)
// ... (see Story 3.2 dev notes for full interface)
}
interface DiagramEdge {
id: string; // Required
from: string; // Required - source node ID
to: string; // Required - target node ID
label?: string;
type?: string;
cardinality?: string;
}
```
### Badge Chip UX Requirements (from UX Spec)
| Property | Value |
|---|---|
| Animation | Slide-in 200ms ease-out (entry), fade-out (exit) |
| Content | Node icon + truncated label + X dismiss button |
| Click badge | Scroll canvas to element + highlight |
| Click X | Deselect element on canvas (remove from `selectedNodeIds`) |
| Backspace/Delete | Remove last badge (like tag input pattern) |
| Escape (in textarea) | Clear all badges |
| Max display | Show all selected (scroll horizontally if overflow) |
| Placeholder | "Describe changes to [name]..." or "Describe changes to N elements..." |
| ARIA | `role="listitem"`, `aria-label="Selected element: [name]"` |
| Tokens | `--badge-chip-bg: oklch(0.623 0.214 260 / 10%)`, `--badge-chip-border: oklch(0.623 0.214 260 / 30%)`, `--badge-chip-text: oklch(0.45 0.20 260)` |
### Scope Indicator (UX Spec)
When badges are active, display a subtle scope indicator in the chat: "Context: [node name] + N connected edges" — so users know what the AI sees. This is a small text below the badge area, not a modal or tooltip.
### File Structure
**Files to CREATE:**
```
apps/web/src/modules/copilot/components/
└── BadgeChip.tsx # Badge chip component with animation + dismiss
```
**Files to MODIFY:**
```
apps/web/src/modules/copilot/components/
└── CopilotPanel.tsx # Add badge display area, selectedElements in transport, placeholder
packages/ai/src/modules/copilot/
├── schema.ts # Add selectedElements optional field
├── types.ts # Export SelectedElement type
├── system-prompt.ts # Add scoped context section when elements selected
├── system-prompt.test.ts # Add scoped context tests
└── api.ts # Pass selectedElements to system prompt
```
**Files to REFERENCE (read-only):**
```
apps/web/src/modules/diagram/stores/useGraphStore.ts # selectedNodeIds state (consume)
apps/web/src/modules/diagram/components/editor/DiagramCanvas.tsx # Selection handling (no changes)
apps/web/src/modules/diagram/lib/graph-converter.ts # flowToGraph, flowNodeToGraphNode
apps/web/src/modules/diagram/types/graph.ts # DiagramNode, DiagramEdge interfaces
```
### Project Structure Notes
- BadgeChip goes in `apps/web/src/modules/copilot/components/` — co-located with CopilotPanel, NOT in diagram module (badge is a copilot UI concern)
- Per the epics file, the suggested path was `apps/web/src/modules/chat/BadgeChip.tsx` — but the project uses `copilot` not `chat` as the module name. Use `apps/web/src/modules/copilot/components/BadgeChip.tsx`
- Schema changes go in `packages/ai/src/modules/copilot/schema.ts` — extend existing `copilotMessageSchema`
- System prompt changes go in `packages/ai/src/modules/copilot/system-prompt.ts` — extend `buildCopilotSystemPrompt` options
### Anti-Patterns to AVOID
- **Do NOT create a separate API endpoint** for scoped modifications — extend the existing copilot POST. The AI already receives `graphContext`; this story adds `selectedElements` as additional context.
- **Do NOT modify canvas selection behavior** — selection is already working via `@xyflow/react` built-in + Zustand store from Story 2.9. This story READS from the store, not writes.
- **Do NOT create a "patch mode" that only sends partial graph data** — the AI always outputs a COMPLETE GraphData. "Targeted" scope is a system prompt instruction that tells the AI to focus changes on selected elements while preserving everything else.
- **Do NOT use `require()` or CommonJS** — all packages are ESM-only
- **Do NOT inline `.parse()` in Hono handlers** — use `validate()` middleware
- **Do NOT put business logic in API routers** — handlers call domain package functions
- **Do NOT re-create selection UI** — use `@xyflow/react`'s existing `elementsSelectable` + `onSelectionChange`
- **Do NOT use `uuid()` column type** — always `text().primaryKey().$defaultFn(generateId)`
- **Do NOT run ELK.js on main thread** — use existing Web Worker
### Performance Requirements
- Badge chip animation: ≤ 200ms (UX spec)
- Badge appear after selection: < 50ms (instant feel — just Zustand state read)
- AI targeted mutation applied and rendered < 3 seconds (NFR4, same as whole-diagram)
- Respect `prefers-reduced-motion` — badges appear instantly instead of slide-in
### Security Requirements
- `enforceAuth` middleware on all endpoints (already in place)
- `deductCredits` middleware before AI generation calls (already in place)
- `selectedElements` data comes from client-side store (trusted) — validate schema structure only
- AI output still validated through `validateGraphPatch` mutation schema before applying to graph
### Testing Standards
- Test runner: Vitest with explicit imports (`import { describe, it, expect } from 'vitest'`)
- Test location: co-located with source files
- Factory pattern for test data
- Component rendering tests deferred to E2E per project standards
- Workspace command: `pnpm --filter @turbostarter/ai test`
- Expected new tests: ~10-15 (system prompt scoped context + schema selectedElements validation)
### Previous Story Intelligence (Story 3.2)
**What was built:**
- `generateDiagramTool` with `graphPatchSchema` — AI outputs full GraphData via tool call
- `useGraphMutation.ts` hook — applies AI patches to Zustand store + persists to DB
- System prompt with generation/modification instructions, node/edge type references, examples
- Tool invocation detection in CopilotPanel via `isGenerateDiagramTool` type guard
- `graphContext` serialization in transport's `prepareSendMessagesRequest`
- `layoutRequestId` counter for cross-component layout triggering
**Key learnings:**
- AI SDK v4 uses `tool-${toolName}` type pattern (e.g., `tool-generateDiagram`)
- AI SDK v4 tool part states: `input-streaming`, `input-available`, `output-available`, `output-error`
- `useGraphStore.getState()` in transport callback avoids transport re-creation
- Cross-component communication via Zustand store counters (not React Flow hooks)
**Files created in 3.2 (do NOT recreate):**
- `packages/ai/src/modules/copilot/mutation-schema.ts` — Zod schema + validation
- `packages/ai/src/modules/copilot/mutation-schema.test.ts` — 38 tests
- `apps/web/src/modules/copilot/hooks/useGraphMutation.ts` — graph patch application hook
**Review fixes from 3.2 (already in codebase):**
- `diagramType` validated as `z.enum(VALID_DIAGRAM_TYPES)` in mutation schema
- `toChatMessage()` preserves tool invocation parts
- `validateUniqueIds()` prevents duplicate node/edge IDs
- `graphData` structurally validated in `updateDiagramBodySchema`
**No new dependencies needed** — Motion (for animations) is already in workspace, shadcn/ui Badge is available.
### Git Intelligence
Recent commit pattern: `feat: implement Story X.Y — <description>`. Follow this convention.
Last 5 commits all follow the pattern of implementing one story per commit. Files modified in Story 3.2:
- CopilotPanel.tsx (extended — will extend further)
- system-prompt.ts (extended — will extend further)
- api.ts (extended — will extend further)
- schema.ts (extended — will extend further, minor)
- types.ts (extended — will extend further, minor)
### References
- [Source: _bmad-output/planning-artifacts/epics.md#Story 3.3] — Acceptance criteria, technical notes
- [Source: _bmad-output/planning-artifacts/ux-design-specification.md#BadgeChip] — Component spec: animation, content, states, interactions, accessibility, tokens
- [Source: _bmad-output/planning-artifacts/ux-design-specification.md#Core User Experience] — Badge referencing as signature interaction, conversational design loop
- [Source: _bmad-output/planning-artifacts/architecture.md] — AI mutation pipeline (client-side relay), soft-lock, bidirectional canvas↔chat state
- [Source: _bmad-output/project-context.md] — Framework rules, coding standards, anti-patterns
- [Source: apps/web/src/modules/diagram/stores/useGraphStore.ts] — selectedNodeIds state (lines 26, 37, 76)
- [Source: apps/web/src/modules/diagram/components/editor/DiagramCanvas.tsx] — onSelectionChange handler (lines 314-319)
- [Source: _bmad-output/implementation-artifacts/3-2-ai-diagram-generation-from-natural-language.md] — Previous story patterns, review fixes, established conventions
## Dev Agent Record
### Agent Model Used
Claude Opus 4.6
### Debug Log References
- Fixed BadgeChip icon names: `Icons.User``Icons.User2`, `Icons.Box``Icons.Server`, `Icons.Layers``Icons.Package` (not exported in project's icons.tsx)
- Pre-existing TS errors in bpmn-layout.ts and CopilotPanel setMessages type (not from this story)
### Completion Notes List
- All 8 tasks completed, all subtasks done
- 103 copilot tests passing (33 system-prompt, 38 mutation-schema, 15 schema, 17 chunking)
- 8 scoped context tests + 6 buildSelectedContext tests + 6 selectedElements schema tests = 20 new tests
- No new TypeScript errors introduced (verified via `tsc --noEmit`)
- Badge chips use Motion animations with `AnimatePresence` for enter/exit
- AI receives `selectedElements` + `selectedContext` (selected nodes, connected edges, 1-hop neighbors)
- System prompt adds "Scoped context" section instructing AI to focus modifications on selected elements
### Code Review Fixes Applied (2026-02-28)
- **H1**: Added `prefers-reduced-motion` support via `useReducedMotion()` hook in BadgeChip — animations skip when OS reduced-motion enabled
- **H2**: Added scope indicator text below badges: "Context: [name] + N connected edges"
- **M1/M2**: Fixed `handleDismissBadge` to use `getState()` instead of closure — eliminates stale closure risk and prevents memo invalidation
- **M3**: Moved `buildSelectedContext` from api.ts to system-prompt.ts for testability; added 6 unit tests
- **L1**: Changed X dismiss button `tabIndex` from -1 to 0 for keyboard accessibility
- **L2**: Replaced O(N*M) `.find()` in selectedElements memo with O(N) `Map` lookup
- **L3**: Added `.min(1)` constraints to `selectedElementSchema` string fields
- **L4**: Added MAX_NEIGHBOR_NODES=10 limit to prevent prompt bloat from highly-connected nodes
### File List
**Created:**
- `apps/web/src/modules/copilot/components/BadgeChip.tsx`
**Modified:**
- `apps/web/src/modules/copilot/components/CopilotPanel.tsx` — badge display, Escape key, selectedElements in transport, dynamic placeholder, scope indicator
- `packages/ai/src/modules/copilot/schema.ts``selectedElementSchema`, `selectedElements` field
- `packages/ai/src/modules/copilot/types.ts``SelectedElement` type export
- `packages/ai/src/modules/copilot/system-prompt.ts``buildScopedContextSection`, `buildSelectedContext`, extended options
- `packages/ai/src/modules/copilot/system-prompt.test.ts` — 8 scoped context tests + 6 buildSelectedContext tests
- `packages/ai/src/modules/copilot/schema.test.ts` — 6 selectedElements validation tests
- `packages/ai/src/modules/copilot/api.ts` — selectedElements passthrough (buildSelectedContext moved to system-prompt.ts)
- `apps/web/src/assets/styles/globals.css` — badge chip CSS tokens (light + dark)

View File

@@ -65,7 +65,7 @@ development_status:
epic-3: in-progress
3-1-chat-panel-ui-with-streaming-ai-responses: done
3-2-ai-diagram-generation-from-natural-language: done
3-3-badge-based-element-referencing-for-targeted-modifications: backlog
3-3-badge-based-element-referencing-for-targeted-modifications: done
3-4-ai-semantic-suggestions-and-accept-reject-workflow: backlog
3-5-new-diagram-wizard-with-ai-type-inference-and-chat-first-onboarding: backlog
3-6-hover-affordances-and-command-palette: backlog

View File

@@ -20,6 +20,10 @@
--edge-selected: oklch(0.623 0.214 260);
/* AI (placeholders for future epics) */
--ai-accent: oklch(0.623 0.214 260);
/* Badge chips (element referencing in chat) */
--badge-chip-bg: oklch(0.623 0.214 260 / 10%);
--badge-chip-border: oklch(0.623 0.214 260 / 30%);
--badge-chip-text: oklch(0.45 0.20 260);
/* Diagram type accents */
--diagram-bpmn: oklch(0.623 0.214 260);
--diagram-er: oklch(0.606 0.25 293);
@@ -45,6 +49,9 @@
--node-hover: oklch(0.623 0.214 260 / 12%);
--edge-default: oklch(0.55 0.01 286);
--edge-selected: oklch(0.623 0.214 260);
--badge-chip-bg: oklch(0.623 0.214 260 / 15%);
--badge-chip-border: oklch(0.623 0.214 260 / 40%);
--badge-chip-text: oklch(0.70 0.18 260);
--bpmn-start-event: #27ae60;
--bpmn-end-event: #c0392b;
--bpmn-timer-event: #5dade2;

View File

@@ -0,0 +1,105 @@
"use client";
import { memo, useCallback } from "react";
import { motion, useReducedMotion } from "motion/react";
import { cn } from "@turbostarter/ui";
import { Icons } from "@turbostarter/ui-web/icons";
import { useGraphStore } from "~/modules/diagram/stores/useGraphStore";
interface BadgeChipProps {
nodeId: string;
label: string;
nodeType: string;
onDismiss: (nodeId: string) => void;
}
export const BadgeChip = memo<BadgeChipProps>(
({ nodeId, label, nodeType, onDismiss }) => {
const prefersReducedMotion = useReducedMotion();
const handleClick = useCallback(() => {
// Highlight the node on canvas via Zustand store
useGraphStore.getState().setHighlightedNodeId(nodeId);
}, [nodeId]);
const handleDismiss = useCallback(
(e: React.MouseEvent) => {
e.stopPropagation();
onDismiss(nodeId);
},
[nodeId, onDismiss],
);
const handleKeyDown = useCallback(
(e: React.KeyboardEvent) => {
if (e.key === "Backspace" || e.key === "Delete") {
e.preventDefault();
onDismiss(nodeId);
}
},
[nodeId, onDismiss],
);
return (
<motion.button
layout
initial={prefersReducedMotion ? false : { opacity: 0, x: -8, scale: 0.95 }}
animate={{ opacity: 1, x: 0, scale: 1 }}
exit={{ opacity: 0, scale: prefersReducedMotion ? 1 : 0.9 }}
transition={{ duration: prefersReducedMotion ? 0 : 0.2, ease: "easeOut" }}
type="button"
role="listitem"
aria-label={`Selected element: ${label}`}
onClick={handleClick}
onKeyDown={handleKeyDown}
className={cn(
"inline-flex items-center gap-1 rounded-md px-2 py-0.5 text-xs font-medium",
"border transition-colors duration-150",
"bg-[var(--badge-chip-bg)] border-[var(--badge-chip-border)] text-[var(--badge-chip-text)]",
"hover:bg-[var(--badge-chip-border)] focus-visible:outline-none focus-visible:ring-1 focus-visible:ring-ring",
"max-w-[160px] cursor-pointer",
)}
>
<NodeTypeIcon nodeType={nodeType} />
<span className="truncate">{label}</span>
<span
role="button"
tabIndex={0}
aria-label={`Remove ${label}`}
onClick={handleDismiss}
className="ml-0.5 shrink-0 rounded-sm p-0.5 hover:bg-[var(--badge-chip-border)] transition-colors"
>
<Icons.X className="size-2.5" />
</span>
</motion.button>
);
},
);
BadgeChip.displayName = "BadgeChip";
/** Diagram-type-aware node icon */
function NodeTypeIcon({ nodeType }: { nodeType: string }) {
// Map common node types to icons
if (nodeType.includes("gateway") || nodeType.includes("decision")) {
return <Icons.GitBranch className="size-3 shrink-0" />;
}
if (nodeType.includes("database") || nodeType.includes("entity")) {
return <Icons.Database className="size-3 shrink-0" />;
}
if (nodeType.includes("person") || nodeType.includes("participant") || nodeType.includes("actor")) {
return <Icons.User2 className="size-3 shrink-0" />;
}
if (nodeType.includes("service") || nodeType.includes("process") || nodeType.includes("activity")) {
return <Icons.Server className="size-3 shrink-0" />;
}
if (nodeType.includes("queue")) {
return <Icons.Package className="size-3 shrink-0" />;
}
if (nodeType.includes("external") || nodeType.includes("loadbalancer")) {
return <Icons.Globe className="size-3 shrink-0" />;
}
// Default
return <Icons.Circle className="size-3 shrink-0" />;
}

View File

@@ -3,6 +3,7 @@
import { useChat } from "@ai-sdk/react";
import { DefaultChatTransport } from "ai";
import { memo, useCallback, useEffect, useMemo, useRef, useState } from "react";
import { AnimatePresence } from "motion/react";
import { useQuery } from "@tanstack/react-query";
import { toast } from "sonner";
@@ -17,6 +18,7 @@ import { Prose } from "~/modules/common/prose";
import { useGraphStore } from "~/modules/diagram/stores/useGraphStore";
import { flowToGraph } from "~/modules/diagram/lib/graph-converter";
import { useGraphMutation } from "../hooks/useGraphMutation";
import { BadgeChip } from "./BadgeChip";
import type { DiagramType } from "~/modules/diagram/types/graph";
@@ -48,6 +50,50 @@ export function CopilotPanel({ diagramId, diagramType }: CopilotPanelProps) {
const { applyGraphPatch } = useGraphMutation(diagramId, diagramType);
// Subscribe to selected nodes for badge chips
const selectedNodeIds = useGraphStore((s) => s.selectedNodeIds);
const storeNodes = useGraphStore((s) => s.nodes);
const storeEdges = useGraphStore((s) => s.edges);
const selectedElements = useMemo(() => {
if (selectedNodeIds.length === 0) return [];
const nodeMap = new Map(storeNodes.map((n) => [n.id, n]));
return selectedNodeIds
.map((id) => {
const node = nodeMap.get(id);
if (!node) return null;
const data = node.data as { type?: string; label?: string };
return { id: node.id, type: data.type ?? "unknown", label: data.label ?? node.id };
})
.filter((e): e is { id: string; type: string; label: string } => e !== null);
}, [selectedNodeIds, storeNodes]);
const handleDismissBadge = useCallback((nodeId: string) => {
const current = useGraphStore.getState().selectedNodeIds;
useGraphStore.getState().setSelectedNodeIds(current.filter((id) => id !== nodeId));
}, []);
// Scope indicator: show what context the AI will see
const scopeInfo = useMemo(() => {
if (selectedElements.length === 0) return null;
const selectedIds = new Set(selectedNodeIds);
const connectedCount = storeEdges.filter(
(e) => selectedIds.has(e.source) || selectedIds.has(e.target),
).length;
const label =
selectedElements.length === 1
? selectedElements[0]!.label
: `${selectedElements.length} elements`;
return `Context: ${label} + ${connectedCount} connected edge${connectedCount !== 1 ? "s" : ""}`;
}, [selectedElements, selectedNodeIds, storeEdges]);
// Dynamic placeholder text
const placeholder = useMemo(() => {
if (selectedElements.length === 0) return "Describe what you want to build...";
if (selectedElements.length === 1) return `Describe changes to ${selectedElements[0]!.label}...`;
return `Describe changes to ${selectedElements.length} elements...`;
}, [selectedElements]);
// Fetch existing chat history on mount (H1 fix)
const { data: initialMessages } = useQuery({
queryKey: ["copilot", "messages", chatId],
@@ -83,6 +129,20 @@ export function CopilotPanel({ diagramId, diagramType }: CopilotPanelProps) {
)
: undefined;
// Build selected element context for targeted modifications
const currentSelectedIds = useGraphStore.getState().selectedNodeIds;
const selectedEls =
currentSelectedIds.length > 0
? currentSelectedIds
.map((nid) => {
const node = currentNodes.find((n) => n.id === nid);
if (!node) return null;
const data = node.data as { type?: string; label?: string };
return { id: node.id, type: data.type ?? "unknown", label: data.label ?? node.id };
})
.filter(Boolean)
: undefined;
return {
body: {
...lastMessage,
@@ -90,6 +150,7 @@ export function CopilotPanel({ diagramId, diagramType }: CopilotPanelProps) {
diagramId,
diagramType,
graphContext,
selectedElements: selectedEls,
},
};
},
@@ -213,8 +274,12 @@ export function CopilotPanel({ diagramId, diagramType }: CopilotPanelProps) {
e.preventDefault();
handleSend();
}
if (e.key === "Escape" && selectedNodeIds.length > 0) {
e.preventDefault();
useGraphStore.getState().setSelectedNodeIds([]);
}
},
[handleSend],
[handleSend, selectedNodeIds.length],
);
return (
@@ -278,13 +343,34 @@ export function CopilotPanel({ diagramId, diagramType }: CopilotPanelProps) {
{/* Input area */}
<div className="shrink-0 border-t border-border p-3">
{/* Badge chips for selected elements */}
{selectedElements.length > 0 && (
<div className="mb-2 flex flex-wrap gap-1" role="list" aria-label="Selected elements">
<AnimatePresence mode="popLayout">
{selectedElements.map((el) => (
<BadgeChip
key={el.id}
nodeId={el.id}
label={el.label}
nodeType={el.type}
onDismiss={handleDismissBadge}
/>
))}
</AnimatePresence>
{scopeInfo && (
<p className="w-full text-[10px] text-muted-foreground/60 mt-0.5">
{scopeInfo}
</p>
)}
</div>
)}
<div className="relative">
<textarea
ref={inputRef}
value={input}
onChange={(e) => setInput(e.target.value)}
onKeyDown={handleKeyDown}
placeholder="Describe what you want to build..."
placeholder={placeholder}
rows={1}
className="w-full resize-none rounded-lg border border-border bg-background px-3 py-2 pr-20 text-sm placeholder:text-muted-foreground focus:outline-none focus:ring-1 focus:ring-ring"
/>

View File

@@ -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({

View File

@@ -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);
});
});
});

View File

@@ -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"),

View File

@@ -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([]);
});
});

View File

@@ -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)}`;
}

View File

@@ -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";