# 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

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

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

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

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

- [x] Task 5: Wire CopilotPanel to generation flow (AC: #1, #2, #3, #5)
  - [x] 5.1 `graphContext` serialized in `prepareSendMessagesRequest` via `useGraphStore.getState()` + `flowToGraph()`
  - [x] 5.2 Tool invocations detected via `part.type === "tool-generateDiagram"` (AI SDK v4 pattern)
  - [x] 5.3 `useGraphMutation.applyGraphPatch()` called on `output-available` state, tracked by `appliedToolCallIds` ref
  - [x] 5.4 "Generating diagram..." indicator shown during `input-streaming`/`input-available` states
  - [x] 5.5 Empty state shows "What are you designing today?" greeting

- [x] Task 6: Implement diagram type inference (AC: #2)
  - [x] 6.1 Type inference rules added to system prompt with keyword → diagram type mapping
  - [x] 6.2 System prompt receives `diagramType` param — AI uses existing type for consistency
  - [x] 6.3 AI sets `meta.diagramType` in the patch for empty diagrams
  - [x] 6.4 Extended `updateDiagramBodySchema` to accept `graphData`; persistence via fire-and-forget PATCH in `useGraphMutation`

- [x] Task 7: Write tests (AC: all)
  - [x] 7.1 Mutation schema validation tests: 30 tests covering all diagram types, node/edge type validation, referential integrity
  - [x] 7.2 System prompt tests: 19 tests covering generation instructions, type-specific schemas, few-shot examples, type inference rules
  - [x] 7.3 Graph patch application: tested indirectly via graph-converter (39 tests) and graph store (25 tests) — hook is glue code
  - [x] 7.4 Type inference prompt tests: covered in system-prompt tests (type inference rules, keyword presence)
  - [x] 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:**
```typescript
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:**
```typescript
// 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`

```typescript
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