feat: implement Story 3.2 — AI diagram generation from natural language

Add complete AI-powered diagram generation pipeline: natural language input
→ type inference → graph patch generation → validated canvas render with
ELK.js layout animation. Includes adversarial code review fixes for
diagramType enum validation, duplicate ID detection, tool part history
preservation, PATCH error handling, and graphData structural validation.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Alejandro Gutiérrez
2026-02-28 13:34:46 +00:00
parent 26215d9060
commit 6dcb4dcd6f
13 changed files with 1577 additions and 44 deletions

View File

@@ -0,0 +1,352 @@
# 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

View File

@@ -64,7 +64,7 @@ development_status:
# ── Epic 3: AI Copilot & Chat (Phase 2) ──
epic-3: in-progress
3-1-chat-panel-ui-with-streaming-ai-responses: done
3-2-ai-diagram-generation-from-natural-language: backlog
3-2-ai-diagram-generation-from-natural-language: done
3-3-badge-based-element-referencing-for-targeted-modifications: backlog
3-4-ai-semantic-suggestions-and-accept-reject-workflow: backlog
3-5-new-diagram-wizard-with-ai-type-inference-and-chat-first-onboarding: backlog