turbostarter/_bmad-output/implementation-artifacts/3-2-ai-diagram-generation-from-natural-language.md

Story 3.2: AI Diagram Generation from Natural Language

Status: done

Story

As a user, I want to describe what I need in natural language and have the AI generate a diagram, so that I can go from idea to visual diagram without manual drawing.

Acceptance Criteria

1. Given I have a new empty diagram open, When the chat panel loads, Then the AI greets me with "What are you designing today?" and the canvas shows the dot grid background (no blank canvas paralysis)

2. Given I type a description like "Customer checkout flow with payment, address confirmation, and error handling", When the AI processes my message, Then the AI determines the best diagram type (BPMN for this example), generates a JSON graph patch with appropriate nodes and edges, and the diagram materializes on the canvas with nodes animating into position via ELK.js. Full mutation applied and rendered in < 3 seconds (NFR4)

3. Given no elements are selected on the canvas (FR3), When I send a message to the AI, Then the AI operates in whole-diagram scope and can create new diagrams, restructure existing ones, or analyze the entire graph

4. Given the AI generates a diagram, When the graph data is created, Then nodes use the correct diagram-type-specific node types (e.g., BPMN gateways, E-R entities), edges have proper routing metadata, and the graph data is saved to the database

5. Given the AI is generating a diagram, When streaming is in progress, Then the chat shows the AI reasoning/response streaming, and once complete, the diagram renders on canvas

6. Given the AI service fails or times out, When an error occurs during generation, Then the user sees a friendly error message, the chat remains functional, and no partial/corrupt graph data is applied

Tasks / Subtasks

• Task 1: Extend system prompt for diagram generation (AC: #2, #3, #4)

  • 1.1 Update buildCopilotSystemPrompt() in packages/ai/src/modules/copilot/system-prompt.ts to include mutation instructions
  • 1.2 Add JSON output format specification matching GraphData interface (nodes, edges, meta)
  • 1.3 Add diagram-type-specific schemas and best practices for each of 6 types (BPMN, E-R, orgchart, architecture, sequence, flowchart)
  • 1.4 Add type inference instructions so AI determines diagram type from natural language
  • 1.5 Add few-shot examples for each diagram type showing expected JSON output
  • 1.6 Update system prompt tests in system-prompt.test.ts (19 tests)

• Task 2: Create mutation schema and types (AC: #4, #6)

  • 2.1 Create packages/ai/src/modules/copilot/mutation-schema.ts with Zod schema for AI-generated GraphData patches
  • 2.2 Schema must validate: nodes[] (id, type, label required), edges[] (id, from, to required), meta (diagramType required)
  • 2.3 Add type-specific node type validation (e.g., BPMN nodes must use bpmn:* prefix types)
  • 2.4 Create mutation-schema.test.ts with valid/invalid patch tests (30 tests)
  • 2.5 Export types from packages/ai/src/modules/copilot/types.ts

• Task 3: Implement AI diagram generation handler (AC: #2, #5, #6)

  • 3.1 Added generateDiagramTool using AI SDK tool() with inputSchema: graphPatchSchema and server-side validation
  • 3.2 graphContext field already existed in copilotMessageSchema (from Story 3.1 M3 fix)
  • 3.3 Tool execute function validates node/edge types and referential integrity before returning
  • 3.4 Streaming handled via createUIMessageStream — text streams first, tool results emit as structured parts
  • 3.5 stopWhen: stepCountIs(2) prevents infinite loops; validation errors returned as { success: false, errors }

• Task 4: Implement graph patch application on canvas (AC: #2, #4)

  • 4.1 Create apps/web/src/modules/copilot/hooks/useGraphMutation.ts hook
  • 4.2 Tool invocation detection via isGenerateDiagramTool type guard in CopilotPanel
  • 4.3 Patch applied to Zustand store via setNodes/setEdges after graphToFlow() conversion
  • 4.4 Layout triggered via requestLayout() → layoutRequestId counter → useAutoLayout watcher
  • 4.5 Animation uses existing CSS transition on .react-flow__node.layouting class from Stories 2.x

• Task 5: Wire CopilotPanel to generation flow (AC: #1, #2, #3, #5)

  • 5.1 graphContext serialized in prepareSendMessagesRequest via useGraphStore.getState() + flowToGraph()
  • 5.2 Tool invocations detected via part.type === "tool-generateDiagram" (AI SDK v4 pattern)
  • 5.3 useGraphMutation.applyGraphPatch() called on output-available state, tracked by appliedToolCallIds ref
  • 5.4 "Generating diagram..." indicator shown during input-streaming/input-available states
  • 5.5 Empty state shows "What are you designing today?" greeting

• Task 6: Implement diagram type inference (AC: #2)

  • 6.1 Type inference rules added to system prompt with keyword → diagram type mapping
  • 6.2 System prompt receives diagramType param — AI uses existing type for consistency
  • 6.3 AI sets meta.diagramType in the patch for empty diagrams
  • 6.4 Extended updateDiagramBodySchema to accept graphData; persistence via fire-and-forget PATCH in useGraphMutation

• Task 7: Write tests (AC: all)

  • 7.1 Mutation schema validation tests: 30 tests covering all diagram types, node/edge type validation, referential integrity
  • 7.2 System prompt tests: 19 tests covering generation instructions, type-specific schemas, few-shot examples, type inference rules
  • 7.3 Graph patch application: tested indirectly via graph-converter (39 tests) and graph store (25 tests) — hook is glue code
  • 7.4 Type inference prompt tests: covered in system-prompt tests (type inference rules, keyword presence)
  • 7.5 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). Do NOT create side-channel API calls that directly modify the graph.
• ELK.js in Web Worker: Layout runs in Web Worker. Use the existing layout trigger mechanism from Stories 2.x. Do NOT run ELK on the main thread.
• 200-node soft cap: AI should not generate graphs with more than ~200 nodes per request. Add instruction in system prompt.
• Unified graph model: All diagram types use the same GraphData interface with type discriminators. The AI output MUST match this exact interface.

Critical Implementation Patterns (from Story 3.1)

Transport pattern for sending graph context:

const transport = useMemo(
  () => new DefaultChatTransport({
    api: api.ai.copilot.$url().toString(),
    prepareSendMessagesRequest: ({ messages, id }) => {
      const lastMessage = messages.at(-1);
      return {
        body: {
          ...lastMessage,
          chatId: id,
          diagramId,
          diagramType,
          graphContext: JSON.stringify(currentGraphData), // NEW: serialize graph state
        },
      };
    },
  }),
  [diagramId, diagramType, currentGraphData],
);

Server-side streaming pattern:

// In copilot router handler:
const result = streamText({
  model: registry.languageModel(modelId),
  system: getDiagramSystemPrompt(diagramType, { mode: 'generate' }),
  messages: convertToModelMessages(messages),
  // ...
});
return result.toUIMessageStreamResponse();

graphContext schema field already exists in copilotMessageSchema (added during Story 3.1 code review fix M3). Use this to pass current graph state to the AI.

Unified Graph Model (Source of Truth)

Location: apps/web/src/modules/diagram/types/graph.ts

type DiagramType = "bpmn" | "er" | "orgchart" | "architecture" | "sequence" | "flowchart";

interface DiagramNode {
  id: string;      // Required
  type: string;    // Required - diagram-type-specific (e.g., "activity", "gateway", "entity")
  tag?: string;    // Optional subtype
  label: string;   // Required - display text
  icon?: string;
  color?: string;
  w?: number;      // Width override
  position?: { x: number; y: number };
  manuallyPositioned?: boolean;
  lane?: string;   // BPMN lane assignment
  group?: string;  // Group containment
  columns?: Column[]; // E-R entity columns
  lifeline?: boolean; // Sequence diagram lifeline
  parentId?: string;  // Nesting / subprocess containment
}

interface DiagramEdge {
  id: string;      // Required
  from: string;    // Required - source node ID
  to: string;      // Required - target node ID
  label?: string;
  color?: string;
  type?: string;   // Edge type (e.g., "association", "dependency")
  cardinality?: string; // E-R cardinality (e.g., "1:N")
}

interface GraphData {
  meta?: DiagramMeta;           // { version, title, description, diagramType, layoutDirection, edgeRouting }
  nodes: DiagramNode[];
  edges: DiagramEdge[];
  pools?: Array<{ id, label, lanes: Array<{ id, label }> }>;  // BPMN pools/lanes
  groups?: Array<{ id, label, color? }>;                       // Grouping containers
}

Diagram Type Node Types (from codebase)

Each diagram type has specific node/edge types defined in apps/web/src/modules/diagram/types/<type>/constants.ts:

Diagram Type	Node Types	Edge Types
BPMN	activity, gateway, timerEvent, dataObject, pool, subprocess, group	messageFlow, flow
E-R	entity (with columns[])	relationship (with cardinality)
Orgchart	person	hierarchyEdge
Architecture	service, database, queue, loadBalancer, external	connection
Sequence	actor (with lifeline)	message
Flowchart	process, decision, terminal, io, subprocess	flow

The AI system prompt MUST include these type mappings so the AI generates valid node types.

File Structure

Files to CREATE:

packages/ai/src/modules/copilot/
  ├── mutation-schema.ts          # Zod schema for AI-generated GraphData patches
  ├── mutation-schema.test.ts     # Mutation schema validation tests

apps/web/src/modules/copilot/hooks/
  └── useGraphMutation.ts         # Hook to apply AI patches to canvas

Files to MODIFY:

packages/ai/src/modules/copilot/
  ├── system-prompt.ts            # Add generation mode instructions, type schemas, JSON format
  ├── system-prompt.test.ts       # Add generation prompt tests
  ├── schema.ts                   # Extend copilotMessageSchema if needed
  ├── types.ts                    # Add mutation-related types

packages/api/src/modules/ai/copilot/
  └── router.ts                   # Extend handler for generation mode

apps/web/src/modules/copilot/components/
  └── CopilotPanel.tsx            # Wire graphContext, detect patches, apply to canvas

Files to REFERENCE (read-only):

apps/web/src/modules/diagram/types/graph.ts              # GraphData interface (source of truth)
apps/web/src/modules/diagram/types/*/constants.ts         # Node/edge type definitions per diagram type
apps/web/src/modules/diagram/components/editor/*.tsx      # DiagramEditor, RightPanel (wiring context)
apps/web/src/modules/diagram/hooks/useLayout.ts           # ELK.js layout trigger (if exists)

Project Structure Notes

• Copilot AI logic lives in packages/ai/src/modules/copilot/ — do NOT create a separate module
• Copilot API routes in packages/api/src/modules/ai/copilot/ — extend existing router
• Copilot UI in apps/web/src/modules/copilot/ — add hooks subdirectory for new hooks
• Feature code in ~/modules/<feature>/ — do NOT co-locate in route directories
• Hooks go in feature module, not in apps/web/src/hooks/

Anti-Patterns to AVOID

• Do NOT create a new API endpoint for generation. Extend the existing copilot POST endpoint. The AI response should include graph patches as structured parts alongside chat text.
• 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 use uuid() column type — always text().primaryKey().$defaultFn(generateId)
• Do NOT create a separate "generation" mode that bypasses chat. Generation IS chat — the AI responds to messages with both text explanations AND graph patches.
• Do NOT modify graph data outside Zustand store — mutations flow through the store
• Do NOT run ELK.js on main thread — use existing Web Worker
• Do NOT generate more than 200 nodes per AI request

Performance Requirements

• First token < 1 second (NFR3) — already achieved in Story 3.1
• Full mutation applied and rendered in < 3 seconds (NFR4)
• ELK.js layout < 500ms (already benchmarked in Story 2.2)
• Zero-to-diagram < 30 seconds (UX spec: greeting + description + generation)
• Rate limit: 30 requests/min per user on AI endpoints

Security Requirements

• enforceAuth middleware on all endpoints (already in place)
• deductCredits middleware before AI generation calls (already in place)
• Validate AI output against mutation schema before applying to graph (prevent injection of malformed data)
• Do NOT expose raw LLM output to graph store — always validate through Zod schema first

Testing Standards

• Test runner: Vitest with explicit imports (import { describe, it, expect } from 'vitest')
• Test location: co-located with source files (e.g., mutation-schema.test.ts next to mutation-schema.ts)
• Factory pattern for test data (e.g., createTestGraphPatch(overrides))
• Component rendering tests deferred to E2E per project standards
• Workspace command: pnpm --filter @turbostarter/ai test and pnpm --filter @turbostarter/web test
• Expected new tests: ~15-20 (mutation schema validation + generation prompt tests + graph patch application)

Previous Story Intelligence (Story 3.1)

What was built:

• CopilotPanel UI with streaming, auto-scroll, typing indicator, stop button, markdown rendering
• Copilot API route (POST for chat, GET for history) with SSE streaming
• System prompt generator (getDiagramSystemPrompt) — needs extension for generation mode
• Chat persistence via existing chat.chat, chat.message, chat.part tables with diagramId column
• graphContext field in copilotMessageSchema (ready for graph state)
• DiagramType sourced from DB diagramTypeEnum (single source of truth)

Key learnings:

• Use prepareSendMessagesRequest in DefaultChatTransport for custom body params (NOT body option on useChat)
• DB text() columns need "text" as const casting for literal types in UIMessagePart
• Wrap DefaultChatTransport in useMemo to prevent recreation on every render
• AI SDK uses message.parts: Part[] (NOT message.content) for structured responses
• streamText() → convertToModelMessages() → result.toUIMessageStreamResponse() for server-side streaming

Files created in 3.1 (do NOT recreate):

• packages/ai/src/modules/copilot/types.ts — extend with mutation types
• packages/ai/src/modules/copilot/schema.ts — already has graphContext field
• packages/ai/src/modules/copilot/system-prompt.ts — extend for generation
• packages/ai/src/modules/copilot/api.ts — getCopilotHistory query
• packages/api/src/modules/ai/copilot/router.ts — extend handler
• apps/web/src/modules/copilot/components/CopilotPanel.tsx — extend for patches

No new dependencies needed — everything required is already in the workspace (AI SDK, Zod, React Query, Motion, etc.)

References

• [Source: _bmad-output/planning-artifacts/epics.md#Story 3.2] — Acceptance criteria and technical notes
• [Source: _bmad-output/planning-artifacts/architecture.md] — AI mutation architecture, CRDT decisions, ELK.js Web Worker
• [Source: _bmad-output/planning-artifacts/ux-design-specification.md] — Chat-first UX, conversational design loop, performance requirements
• [Source: _bmad-output/project-context.md] — Framework rules, coding standards, anti-patterns
• [Source: apps/web/src/modules/diagram/types/graph.ts] — Unified graph data model (GraphData, DiagramNode, DiagramEdge)
• [Source: _bmad-output/implementation-artifacts/3-1-chat-panel-ui-with-streaming-ai-responses.md] — Previous story learnings and established patterns

Dev Agent Record

Agent Model Used

Claude Opus 4.6

Debug Log References

• AI SDK v4 uses tool-${toolName} type pattern (e.g., tool-generateDiagram), not tool-invocation with toolName property
• AI SDK v4 tool part states: input-streaming, input-available, output-available, output-error — not partial-call/call/result
• AI SDK v4 uses part.output not part.result for tool invocation output
• z.record(z.unknown()) fails in this Zod version — must use z.record(z.string(), z.unknown())
• Cross-component layout triggering: CopilotPanel is outside ReactFlowProvider, solved via layoutRequestId counter in Zustand store

Completion Notes List

• 49 new tests added (19 system-prompt + 30 mutation-schema), exceeding 15-20 target
• All 261 tests pass (75 AI + 186 web)
• TypeScript compiles clean (pre-existing errors in bpmn-layout.ts/bpmn constants.ts are unrelated)
• graphContext serialized at send-time via useGraphStore.getState() in transport callback — avoids transport re-creation
• Graph persistence is fire-and-forget PATCH call after store update — will be replaced by Liveblocks CRDT in Epic 4
• stopWhen: stepCountIs(2) limits AI to one tool call + one follow-up response

Senior Developer Review (AI)

Reviewer: Mou on 2026-02-28 Issues Found: 2 High, 3 Medium, 3 Low — All Fixed

Fixes Applied:

• H1 (FIXED): diagramType in metaSchema changed from z.string() to z.enum(VALID_DIAGRAM_TYPES) — prevents unknown diagram types from bypassing node/edge type validation
• H2 (FIXED): toChatMessage() now preserves tool invocation parts (type tool-*) instead of converting all to empty text — AI retains diagram generation context across reloads
• M1 (FIXED): Fire-and-forget PATCH in useGraphMutation now has .then()/.catch() error handling with toast.error() — prevents silent data loss
• M2 (FIXED): Added validateUniqueIds() function detecting duplicate node/edge IDs — prevents React Flow silent node dropping
• M3 (FIXED): graphData in updateDiagramBodySchema now validates structural shape (meta, nodes[], edges[], pools?, groups?) instead of accepting any record
• L1 (FIXED): Fixed missing blank line before ## Example in system prompt template when typeSpecificSections is empty
• L2 (FIXED): meta.title now requires z.string().min(1) — prevents untitled diagrams
• L3 (FIXED): Added description? to system prompt GraphData format spec

Tests added: 8 new tests (validateUniqueIds: 4, validateGraphPatch: 4 for diagramType enum, empty title, duplicate node/edge IDs) Total tests after review: 269 (83 AI + 186 web) — all passing

File List

Created:

• packages/ai/src/modules/copilot/mutation-schema.ts — Zod schema + validation for AI-generated graph patches
• packages/ai/src/modules/copilot/mutation-schema.test.ts — 38 tests for mutation schema validation (30 original + 8 review fixes)
• apps/web/src/modules/copilot/hooks/useGraphMutation.ts — Hook to apply AI patches to Zustand store + persist to DB

Modified:

• packages/ai/src/modules/copilot/system-prompt.ts — Full rewrite: generation instructions, node/edge type references, type inference rules, few-shot examples per diagram type, 200-node cap, graphContext section
• packages/ai/src/modules/copilot/system-prompt.test.ts — Rewritten: 19 tests (up from 6), covers generation mode, type-specific schemas, type inference
• packages/ai/src/modules/copilot/api.ts — Added generateDiagramTool with graphPatchSchema, validation in execute, tools + stopWhen in streamText, graphContext passed to system prompt. [Review: preserved tool parts in toChatMessage, added validateUniqueIds to tool execute]
• packages/ai/src/modules/copilot/types.ts — Re-exports GraphPatch, graphPatchSchema, validateGraphPatch, validateUniqueIds
• packages/api/src/modules/diagram/router.ts — Extended updateDiagramBodySchema with structurally validated graphData field
• apps/web/src/modules/diagram/stores/useGraphStore.ts — Added layoutRequestId counter + requestLayout() action
• apps/web/src/modules/diagram/hooks/useAutoLayout.ts — Added layoutRequestId watcher to trigger layout from outside ReactFlowProvider
• apps/web/src/modules/copilot/components/CopilotPanel.tsx — Rewired: graphContext serialization, tool invocation detection, graph patch application, "Generating diagram..." indicator, "Diagram updated" confirmation, updated greeting/placeholder