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

View File

@@ -14,9 +14,25 @@ import { ScrollArea } from "@turbostarter/ui-web/scroll-area";
import { api } from "~/lib/api/client";
import { MemoizedMarkdown } from "~/modules/common/markdown/memoized-markdown";
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 type { DiagramType } from "~/modules/diagram/types/graph";
// Type helper for tool invocation parts from AI SDK
interface ToolPart {
type: string;
toolCallId: string;
state: "input-streaming" | "input-available" | "output-available" | "output-error";
output?: unknown;
input?: unknown;
}
function isGenerateDiagramTool(part: { type: string }): part is ToolPart {
return part.type === "tool-generateDiagram";
}
interface CopilotPanelProps {
diagramId: string;
diagramType: DiagramType;
@@ -28,6 +44,9 @@ export function CopilotPanel({ diagramId, diagramType }: CopilotPanelProps) {
const scrollRef = useRef<HTMLDivElement>(null);
const inputRef = useRef<HTMLTextAreaElement>(null);
const userScrolledRef = useRef(false);
const appliedToolCallIds = useRef(new Set<string>());
const { applyGraphPatch } = useGraphMutation(diagramId, diagramType);
// Fetch existing chat history on mount (H1 fix)
const { data: initialMessages } = useQuery({
@@ -49,12 +68,28 @@ export function CopilotPanel({ diagramId, diagramType }: CopilotPanelProps) {
api: api.ai.copilot.$url().toString(),
prepareSendMessagesRequest: ({ messages, id }) => {
const lastMessage = messages.at(-1);
// Serialize current graph state for AI context
const currentNodes = useGraphStore.getState().nodes;
const currentEdges = useGraphStore.getState().edges;
const graphContext =
currentNodes.length > 0
? JSON.stringify(
flowToGraph(currentNodes, currentEdges, {
version: "1",
title: "",
diagramType,
}),
)
: undefined;
return {
body: {
...lastMessage,
chatId: id,
diagramId,
diagramType,
graphContext,
},
};
},
@@ -83,8 +118,45 @@ export function CopilotPanel({ diagramId, diagramType }: CopilotPanelProps) {
}
}, [initialMessages, messages.length, setMessages]);
// Detect and apply graph patches from tool invocations
useEffect(() => {
for (const message of messages) {
if (message.role !== "assistant") continue;
for (const part of message.parts) {
if (
isGenerateDiagramTool(part) &&
part.state === "output-available" &&
!appliedToolCallIds.current.has(part.toolCallId)
) {
appliedToolCallIds.current.add(part.toolCallId);
const result = part.output as
| { success: true; data: Parameters<typeof applyGraphPatch>[0] }
| { success: false; errors: string[] };
if (result.success) {
applyGraphPatch(result.data);
} else {
toast.error("Diagram generation failed: invalid graph structure");
console.error("[copilot] Graph validation errors:", result.errors);
}
}
}
}
}, [messages, applyGraphPatch]);
const isSubmitting = status === "submitted" || status === "streaming";
// Check if currently generating a diagram (tool call in progress)
const isGeneratingDiagram = useMemo(() => {
const lastMessage = messages.at(-1);
if (!lastMessage || lastMessage.role !== "assistant") return false;
return lastMessage.parts.some(
(p) =>
isGenerateDiagramTool(p) &&
(p.state === "input-streaming" || p.state === "input-available"),
);
}, [messages]);
// Auto-scroll on new content, but pause if user scrolled up
useEffect(() => {
if (userScrolledRef.current) return;
@@ -187,6 +259,14 @@ export function CopilotPanel({ diagramId, diagramType }: CopilotPanelProps) {
</div>
)}
{/* Diagram generation indicator */}
{isGeneratingDiagram && (
<div className="flex items-center gap-2 rounded-lg bg-muted/50 px-3 py-2 text-xs text-muted-foreground">
<Icons.Loader2 className="size-3 animate-spin" />
Generating diagram...
</div>
)}
{/* Error display */}
{error && (
<div className="rounded-lg bg-destructive/10 px-3 py-2 text-xs text-destructive">
@@ -204,7 +284,7 @@ export function CopilotPanel({ diagramId, diagramType }: CopilotPanelProps) {
value={input}
onChange={(e) => setInput(e.target.value)}
onKeyDown={handleKeyDown}
placeholder="Ask about your diagram..."
placeholder="Describe what you want to build..."
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"
/>
@@ -252,9 +332,11 @@ function EmptyState() {
return (
<div className="flex flex-col items-center justify-center py-12 text-center">
<Icons.Sparkles className="mb-3 size-8 text-muted-foreground/30" />
<p className="text-sm font-medium text-muted-foreground">AI Copilot</p>
<p className="text-sm font-medium text-muted-foreground">
What are you designing today?
</p>
<p className="mt-1 text-xs text-muted-foreground/60">
Ask questions about your diagram or describe what you want to build
Describe your diagram and I'll generate it for you
</p>
</div>
);
@@ -276,24 +358,36 @@ const UserBubble = memo<{ message: { id: string; parts: Array<{ type: string; te
UserBubble.displayName = "UserBubble";
const AssistantBubble = memo<{
message: { id: string; parts: Array<{ type: string; text?: string }> };
message: { id: string; parts: Array<{ type: string; text?: string; state?: string }> };
isStreaming: boolean;
}>(({ message, isStreaming }) => (
<div className="max-w-[95%]">
<Prose className="text-sm">
{message.parts.map((part, i) =>
part.type === "text" && part.text ? (
<MemoizedMarkdown
key={`${message.id}-${i}`}
content={part.text}
id={`copilot-${message.id}-${i}`}
/>
) : null,
}>(({ message, isStreaming }) => {
const hasToolResult = message.parts.some(
(p) => isGenerateDiagramTool(p) && p.state === "output-available",
);
return (
<div className="max-w-[95%]">
<Prose className="text-sm">
{message.parts.map((part, i) =>
part.type === "text" && part.text ? (
<MemoizedMarkdown
key={`${message.id}-${i}`}
content={part.text}
id={`copilot-${message.id}-${i}`}
/>
) : null,
)}
{isStreaming && message.parts.length === 0 && (
<span className="inline-block size-2 animate-pulse rounded-full bg-muted-foreground/50" />
)}
</Prose>
{hasToolResult && (
<div className="mt-1 flex items-center gap-1.5 text-xs text-muted-foreground">
<Icons.Check className="size-3 text-green-500" />
Diagram updated
</div>
)}
{isStreaming && message.parts.length === 0 && (
<span className="inline-block size-2 animate-pulse rounded-full bg-muted-foreground/50" />
)}
</Prose>
</div>
));
</div>
);
});
AssistantBubble.displayName = "AssistantBubble";

View File

@@ -0,0 +1,93 @@
"use client";
import { useCallback } from "react";
import { toast } from "sonner";
import { api } from "~/lib/api/client";
import { useGraphStore } from "~/modules/diagram/stores/useGraphStore";
import { graphToFlow } from "~/modules/diagram/lib/graph-converter";
import type { DiagramType, GraphData } from "~/modules/diagram/types/graph";
interface GraphPatchData {
meta: {
diagramType: string;
title: string;
version?: string;
layoutDirection?: "DOWN" | "RIGHT" | "LEFT" | "UP";
edgeRouting?: "ORTHOGONAL" | "SPLINES" | "POLYLINE";
};
nodes: GraphData["nodes"];
edges: GraphData["edges"];
pools?: GraphData["pools"];
groups?: GraphData["groups"];
}
export function useGraphMutation(diagramId: string, diagramType: DiagramType) {
const setNodes = useGraphStore((s) => s.setNodes);
const setEdges = useGraphStore((s) => s.setEdges);
const setLayoutDirection = useGraphStore((s) => s.setLayoutDirection);
const setEdgeRouting = useGraphStore((s) => s.setEdgeRouting);
const requestLayout = useGraphStore((s) => s.requestLayout);
const applyGraphPatch = useCallback(
(patch: GraphPatchData) => {
const effectiveDiagramType =
(patch.meta.diagramType as DiagramType) ?? diagramType;
const graphData: GraphData = {
meta: {
version: patch.meta.version ?? "1",
title: patch.meta.title,
diagramType: effectiveDiagramType,
layoutDirection: patch.meta.layoutDirection,
edgeRouting: patch.meta.edgeRouting,
},
nodes: patch.nodes,
edges: patch.edges,
pools: patch.pools,
groups: patch.groups,
};
const { nodes, edges } = graphToFlow(graphData);
setNodes(nodes);
setEdges(edges);
if (graphData.meta?.layoutDirection) {
setLayoutDirection(graphData.meta.layoutDirection);
}
if (graphData.meta?.edgeRouting) {
setEdgeRouting(graphData.meta.edgeRouting);
}
requestLayout();
// Persist graphData to database (fire-and-forget with error reporting)
api.diagrams[":id"]
.$patch({
param: { id: diagramId },
json: { graphData: graphData as unknown as Record<string, unknown> },
})
.then((res) => {
if (!res.ok) {
toast.error("Failed to save diagram — changes may be lost on reload");
}
})
.catch(() => {
toast.error("Failed to save diagram — changes may be lost on reload");
});
},
[
diagramId,
diagramType,
setNodes,
setEdges,
setLayoutDirection,
setEdgeRouting,
requestLayout,
],
);
return { applyGraphPatch };
}

View File

@@ -139,6 +139,15 @@ export function useAutoLayout() {
// eslint-disable-next-line react-hooks/exhaustive-deps
}, [layoutDirection, edgeRouting]);
// Re-layout on explicit request (e.g., AI-generated graph patch)
const layoutRequestId = useGraphStore((s) => s.layoutRequestId);
useEffect(() => {
if (layoutRequestId > 0 && nodeCount > 0) {
runLayout();
}
// eslint-disable-next-line react-hooks/exhaustive-deps
}, [layoutRequestId]);
// Cleanup worker on unmount
useEffect(() => {
return () => {

View File

@@ -24,6 +24,7 @@ interface GraphState {
isLayouting: boolean;
highlightedNodeId: string | null;
selectedNodeIds: string[];
layoutRequestId: number;
setNodes: (nodes: Node[]) => void;
setEdges: (edges: Edge[]) => void;
onNodesChange: OnNodesChange;
@@ -34,6 +35,7 @@ interface GraphState {
setIsLayouting: (isLayouting: boolean) => void;
setHighlightedNodeId: (id: string | null) => void;
setSelectedNodeIds: (ids: string[]) => void;
requestLayout: () => void;
initializeFromGraphData: (nodes: Node[], edges: Edge[]) => void;
reset: () => void;
}
@@ -49,6 +51,7 @@ export const useGraphStore = create<GraphState>((set, get) => ({
isLayouting: false,
highlightedNodeId: null,
selectedNodeIds: [],
layoutRequestId: 0,
setNodes: (nodes) => set({ nodes, nodeCount: nodes.length }),
setEdges: (edges) => set({ edges }),
@@ -71,6 +74,7 @@ export const useGraphStore = create<GraphState>((set, get) => ({
setIsLayouting: (isLayouting) => set({ isLayouting }),
setHighlightedNodeId: (highlightedNodeId) => set({ highlightedNodeId }),
setSelectedNodeIds: (selectedNodeIds) => set({ selectedNodeIds }),
requestLayout: () => set((s) => ({ layoutRequestId: s.layoutRequestId + 1 })),
initializeFromGraphData: (nodes, edges) => {
set({ nodes, edges, nodeCount: nodes.length });
@@ -88,6 +92,7 @@ export const useGraphStore = create<GraphState>((set, get) => ({
isLayouting: false,
highlightedNodeId: null,
selectedNodeIds: [],
layoutRequestId: 0,
});
},
}));

View File

@@ -3,7 +3,9 @@ import {
createUIMessageStream,
createUIMessageStreamResponse,
smoothStream,
stepCountIs,
streamText,
tool,
} from "ai";
import { eq } from "@turbostarter/db";
@@ -13,6 +15,13 @@ import { db } from "@turbostarter/db/server";
import { modelStrategies } from "../chat/strategies";
import { Model, Role } from "../chat/types";
import {
graphPatchSchema,
validateEdgeReferences,
validateEdgeTypes,
validateNodeTypes,
validateUniqueIds,
} from "./mutation-schema";
import { buildCopilotSystemPrompt } from "./system-prompt";
import type { CopilotMessagePayload } from "./schema";
@@ -25,6 +34,24 @@ import type {
const DEFAULT_MODEL = Model.CLAUDE_4_SONNET;
const generateDiagramTool = tool({
description:
"Generate or modify a diagram with the given graph data. Call this when the user asks to create, restructure, or modify a diagram.",
inputSchema: graphPatchSchema,
execute: async (patch) => {
const errors = [
...validateUniqueIds(patch),
...validateNodeTypes(patch),
...validateEdgeTypes(patch),
...validateEdgeReferences(patch),
];
if (errors.length > 0) {
return { success: false as const, errors };
}
return { success: true as const, data: patch };
},
});
const createCopilotChat = async (data: InsertChat) =>
db
.insert(chat)
@@ -60,6 +87,15 @@ const toChatMessage = (msg: Awaited<ReturnType<typeof getChatMessages>>[number])
role: msg.role as "user" | "assistant",
parts: msg.part.map((p) => {
const details = p.details as Record<string, unknown>;
// Preserve tool invocation parts so the AI retains generation context
if (p.type.startsWith("tool-")) {
return {
type: p.type as "text",
...(details as Record<string, unknown>),
};
}
return {
type: "text" as const,
text: (details.text as string) ?? "",
@@ -76,6 +112,7 @@ export const streamCopilot = async ({
chatId,
diagramId,
diagramType,
graphContext,
userId,
signal,
...msg
@@ -102,7 +139,9 @@ export const streamCopilot = async ({
})),
);
const systemPrompt = buildCopilotSystemPrompt(diagramType as DiagramType);
const systemPrompt = buildCopilotSystemPrompt(diagramType as DiagramType, {
graphContext,
});
const stream = createUIMessageStream({
execute: ({ writer }) => {
@@ -117,6 +156,8 @@ export const streamCopilot = async ({
},
]),
system: systemPrompt,
tools: { generateDiagram: generateDiagramTool },
stopWhen: stepCountIs(2),
abortSignal: signal,
experimental_transform: smoothStream({
chunking: "word",

View File

@@ -0,0 +1,464 @@
import { describe, expect, it } from "vitest";
import {
graphPatchSchema,
validateEdgeReferences,
validateEdgeTypes,
validateGraphPatch,
validateNodeTypes,
validateUniqueIds,
} from "./mutation-schema";
import type { GraphPatch } from "./mutation-schema";
function createTestPatch(overrides: Partial<GraphPatch> = {}): GraphPatch {
return {
meta: { diagramType: "flowchart", title: "Test Diagram" },
nodes: [
{ id: "n1", type: "terminal", label: "Start" },
{ id: "n2", type: "process", label: "Do something" },
{ id: "n3", type: "terminal", label: "End" },
],
edges: [
{ id: "e1", from: "n1", to: "n2" },
{ id: "e2", from: "n2", to: "n3" },
],
...overrides,
};
}
describe("graphPatchSchema", () => {
it("should accept a valid flowchart patch", () => {
const patch = createTestPatch();
const result = graphPatchSchema.safeParse(patch);
expect(result.success).toBe(true);
});
it("should accept a valid BPMN patch with pools", () => {
const patch = createTestPatch({
meta: { diagramType: "bpmn", title: "BPMN Process" },
nodes: [
{ id: "n1", type: "start-event", label: "Start", lane: "lane1" },
{ id: "n2", type: "activity", label: "Task", lane: "lane1" },
],
edges: [{ id: "e1", from: "n1", to: "n2", type: "sequence" }],
pools: [
{
id: "pool1",
label: "Customer",
lanes: [{ id: "lane1", label: "Web" }],
},
],
});
const result = graphPatchSchema.safeParse(patch);
expect(result.success).toBe(true);
});
it("should accept a valid E-R patch with columns", () => {
const patch = createTestPatch({
meta: { diagramType: "er", title: "Schema" },
nodes: [
{
id: "n1",
type: "entity",
label: "User",
columns: [
{ name: "id", type: "uuid", isPrimaryKey: true },
{ name: "email", type: "varchar", isUnique: true },
],
},
],
edges: [],
});
const result = graphPatchSchema.safeParse(patch);
expect(result.success).toBe(true);
});
it("should accept a valid sequence patch with lifeline", () => {
const patch = createTestPatch({
meta: { diagramType: "sequence", title: "Flow" },
nodes: [
{ id: "n1", type: "participant", label: "Client", lifeline: true },
{ id: "n2", type: "participant", label: "Server", lifeline: true },
],
edges: [{ id: "e1", from: "n1", to: "n2", type: "sync", label: "request" }],
});
const result = graphPatchSchema.safeParse(patch);
expect(result.success).toBe(true);
});
it("should reject patch without meta", () => {
const result = graphPatchSchema.safeParse({
nodes: [{ id: "n1", type: "process", label: "A" }],
edges: [],
});
expect(result.success).toBe(false);
});
it("should reject patch without nodes", () => {
const result = graphPatchSchema.safeParse({
meta: { diagramType: "flowchart", title: "Test" },
edges: [],
});
expect(result.success).toBe(false);
});
it("should reject node without required fields", () => {
const result = graphPatchSchema.safeParse({
meta: { diagramType: "flowchart", title: "Test" },
nodes: [{ id: "n1" }],
edges: [],
});
expect(result.success).toBe(false);
});
it("should reject edge without from/to", () => {
const result = graphPatchSchema.safeParse({
meta: { diagramType: "flowchart", title: "Test" },
nodes: [{ id: "n1", type: "process", label: "A" }],
edges: [{ id: "e1", from: "n1" }],
});
expect(result.success).toBe(false);
});
it("should reject more than 200 nodes", () => {
const nodes = Array.from({ length: 201 }, (_, i) => ({
id: `n${i}`,
type: "process",
label: `Node ${i}`,
}));
const result = graphPatchSchema.safeParse({
meta: { diagramType: "flowchart", title: "Too Many" },
nodes,
edges: [],
});
expect(result.success).toBe(false);
});
it("should accept optional layoutDirection", () => {
const patch = createTestPatch({
meta: { diagramType: "flowchart", title: "Test", layoutDirection: "RIGHT" },
});
const result = graphPatchSchema.safeParse(patch);
expect(result.success).toBe(true);
});
it("should accept groups", () => {
const patch = createTestPatch({
groups: [{ id: "g1", label: "Group A", color: "#ff0000" }],
nodes: [
{ id: "n1", type: "process", label: "Task", group: "g1" },
{ id: "n2", type: "terminal", label: "End" },
],
edges: [{ id: "e1", from: "n1", to: "n2" }],
});
const result = graphPatchSchema.safeParse(patch);
expect(result.success).toBe(true);
});
});
describe("validateNodeTypes", () => {
it("should accept valid BPMN node types", () => {
const patch = createTestPatch({
meta: { diagramType: "bpmn", title: "Test" },
nodes: [
{ id: "n1", type: "activity", label: "Task" },
{ id: "n2", type: "gateway-exclusive", label: "Decision" },
{ id: "n3", type: "start-event", label: "Start" },
],
});
expect(validateNodeTypes(patch)).toEqual([]);
});
it("should reject invalid BPMN node types", () => {
const patch = createTestPatch({
meta: { diagramType: "bpmn", title: "Test" },
nodes: [{ id: "n1", type: "invalid-type", label: "Bad" }],
});
const errors = validateNodeTypes(patch);
expect(errors).toHaveLength(1);
expect(errors[0]).toContain("invalid-type");
});
it("should strip prefix before validation", () => {
const patch = createTestPatch({
meta: { diagramType: "bpmn", title: "Test" },
nodes: [{ id: "n1", type: "bpmn:activity", label: "Task" }],
});
expect(validateNodeTypes(patch)).toEqual([]);
});
it("should accept valid architecture node types", () => {
const patch = createTestPatch({
meta: { diagramType: "architecture", title: "Test" },
nodes: [
{ id: "n1", type: "service", label: "API" },
{ id: "n2", type: "database", label: "DB" },
{ id: "n3", type: "queue", label: "Queue" },
{ id: "n4", type: "loadbalancer", label: "LB" },
{ id: "n5", type: "external", label: "Ext" },
],
});
expect(validateNodeTypes(patch)).toEqual([]);
});
it("should accept valid flowchart node types", () => {
const patch = createTestPatch({
meta: { diagramType: "flowchart", title: "Test" },
nodes: [
{ id: "n1", type: "process", label: "Do" },
{ id: "n2", type: "decision", label: "?" },
{ id: "n3", type: "terminal", label: "End" },
{ id: "n4", type: "io", label: "Read" },
{ id: "n5", type: "subprocess", label: "Sub" },
],
});
expect(validateNodeTypes(patch)).toEqual([]);
});
it("should accept valid sequence node types", () => {
const patch = createTestPatch({
meta: { diagramType: "sequence", title: "Test" },
nodes: [
{ id: "n1", type: "participant", label: "Client" },
{ id: "n2", type: "fragment", label: "alt" },
],
});
expect(validateNodeTypes(patch)).toEqual([]);
});
});
describe("validateEdgeTypes", () => {
it("should accept valid BPMN edge types", () => {
const patch = createTestPatch({
meta: { diagramType: "bpmn", title: "Test" },
nodes: [
{ id: "n1", type: "activity", label: "A" },
{ id: "n2", type: "activity", label: "B" },
],
edges: [
{ id: "e1", from: "n1", to: "n2", type: "sequence" },
{ id: "e2", from: "n1", to: "n2", type: "message" },
],
});
expect(validateEdgeTypes(patch)).toEqual([]);
});
it("should reject invalid BPMN edge types", () => {
const patch = createTestPatch({
meta: { diagramType: "bpmn", title: "Test" },
nodes: [
{ id: "n1", type: "activity", label: "A" },
{ id: "n2", type: "activity", label: "B" },
],
edges: [{ id: "e1", from: "n1", to: "n2", type: "invalid" }],
});
const errors = validateEdgeTypes(patch);
expect(errors).toHaveLength(1);
expect(errors[0]).toContain("invalid");
});
it("should accept valid sequence edge types", () => {
const patch = createTestPatch({
meta: { diagramType: "sequence", title: "Test" },
nodes: [
{ id: "n1", type: "participant", label: "A" },
{ id: "n2", type: "participant", label: "B" },
],
edges: [
{ id: "e1", from: "n1", to: "n2", type: "sync" },
{ id: "e2", from: "n2", to: "n1", type: "return" },
{ id: "e3", from: "n1", to: "n2", type: "async" },
],
});
expect(validateEdgeTypes(patch)).toEqual([]);
});
it("should skip validation for types with no edge type constraints", () => {
const patch = createTestPatch({
meta: { diagramType: "er", title: "Test" },
nodes: [{ id: "n1", type: "entity", label: "A" }],
edges: [{ id: "e1", from: "n1", to: "n1", type: "anything" }],
});
expect(validateEdgeTypes(patch)).toEqual([]);
});
it("should allow edges without type for constrained diagram types", () => {
const patch = createTestPatch({
meta: { diagramType: "bpmn", title: "Test" },
nodes: [
{ id: "n1", type: "activity", label: "A" },
{ id: "n2", type: "activity", label: "B" },
],
edges: [{ id: "e1", from: "n1", to: "n2" }],
});
expect(validateEdgeTypes(patch)).toEqual([]);
});
});
describe("validateEdgeReferences", () => {
it("should accept edges referencing existing nodes", () => {
const patch = createTestPatch();
expect(validateEdgeReferences(patch)).toEqual([]);
});
it("should reject edges referencing non-existent source", () => {
const patch = createTestPatch({
edges: [{ id: "e1", from: "missing", to: "n1" }],
});
const errors = validateEdgeReferences(patch);
expect(errors).toHaveLength(1);
expect(errors[0]).toContain("missing");
});
it("should reject edges referencing non-existent target", () => {
const patch = createTestPatch({
edges: [{ id: "e1", from: "n1", to: "missing" }],
});
const errors = validateEdgeReferences(patch);
expect(errors).toHaveLength(1);
expect(errors[0]).toContain("missing");
});
});
describe("validateGraphPatch", () => {
it("should return success for valid patch", () => {
const result = validateGraphPatch(createTestPatch());
expect(result.success).toBe(true);
if (result.success) {
expect(result.data.nodes).toHaveLength(3);
expect(result.data.edges).toHaveLength(2);
}
});
it("should return errors for invalid schema", () => {
const result = validateGraphPatch({ invalid: true });
expect(result.success).toBe(false);
if (!result.success) {
expect(result.errors.length).toBeGreaterThan(0);
}
});
it("should return errors for invalid node types", () => {
const result = validateGraphPatch({
meta: { diagramType: "bpmn", title: "Test" },
nodes: [{ id: "n1", type: "invalid", label: "Bad" }],
edges: [],
});
expect(result.success).toBe(false);
});
it("should return errors for broken edge references", () => {
const result = validateGraphPatch({
meta: { diagramType: "flowchart", title: "Test" },
nodes: [{ id: "n1", type: "process", label: "A" }],
edges: [{ id: "e1", from: "n1", to: "missing" }],
});
expect(result.success).toBe(false);
});
it("should combine node type and edge reference errors", () => {
const result = validateGraphPatch({
meta: { diagramType: "bpmn", title: "Test" },
nodes: [{ id: "n1", type: "invalid", label: "Bad" }],
edges: [{ id: "e1", from: "n1", to: "missing" }],
});
expect(result.success).toBe(false);
if (!result.success) {
expect(result.errors.length).toBe(2);
}
});
it("should reject invalid diagramType", () => {
const result = validateGraphPatch({
meta: { diagramType: "unknown", title: "Test" },
nodes: [{ id: "n1", type: "process", label: "A" }],
edges: [],
});
expect(result.success).toBe(false);
});
it("should reject empty title", () => {
const result = validateGraphPatch({
meta: { diagramType: "flowchart", title: "" },
nodes: [{ id: "n1", type: "process", label: "A" }],
edges: [],
});
expect(result.success).toBe(false);
});
it("should reject duplicate node IDs", () => {
const result = validateGraphPatch({
meta: { diagramType: "flowchart", title: "Test" },
nodes: [
{ id: "n1", type: "process", label: "A" },
{ id: "n1", type: "decision", label: "B" },
],
edges: [],
});
expect(result.success).toBe(false);
if (!result.success) {
expect(result.errors[0]).toContain("Duplicate node ID");
}
});
it("should reject duplicate edge IDs", () => {
const result = validateGraphPatch({
meta: { diagramType: "flowchart", title: "Test" },
nodes: [
{ id: "n1", type: "process", label: "A" },
{ id: "n2", type: "process", label: "B" },
],
edges: [
{ id: "e1", from: "n1", to: "n2" },
{ id: "e1", from: "n2", to: "n1" },
],
});
expect(result.success).toBe(false);
if (!result.success) {
expect(result.errors[0]).toContain("Duplicate edge ID");
}
});
});
describe("validateUniqueIds", () => {
it("should accept unique IDs", () => {
const patch = createTestPatch();
expect(validateUniqueIds(patch)).toEqual([]);
});
it("should detect duplicate node IDs", () => {
const patch = createTestPatch({
nodes: [
{ id: "n1", type: "process", label: "A" },
{ id: "n1", type: "process", label: "B" },
],
});
const errors = validateUniqueIds(patch);
expect(errors).toHaveLength(1);
expect(errors[0]).toContain("n1");
});
it("should detect duplicate edge IDs", () => {
const patch = createTestPatch({
edges: [
{ id: "e1", from: "n1", to: "n2" },
{ id: "e1", from: "n2", to: "n3" },
],
});
const errors = validateUniqueIds(patch);
expect(errors).toHaveLength(1);
expect(errors[0]).toContain("e1");
});
it("should allow same ID across nodes and edges", () => {
const patch = createTestPatch({
nodes: [
{ id: "x1", type: "process", label: "A" },
{ id: "x2", type: "process", label: "B" },
],
edges: [{ id: "x1", from: "x1", to: "x2" }],
});
expect(validateUniqueIds(patch)).toEqual([]);
});
});

View File

@@ -0,0 +1,229 @@
import * as z from "zod";
const VALID_DIAGRAM_TYPES = [
"bpmn",
"er",
"orgchart",
"architecture",
"sequence",
"flowchart",
] as const;
const VALID_NODE_TYPES: Record<string, string[]> = {
bpmn: [
"activity",
"subprocess",
"start-event",
"end-event",
"event-timer",
"event-message",
"gateway-exclusive",
"gateway-parallel",
"gateway-inclusive",
"data-object",
"annotation",
],
er: ["entity"],
orgchart: ["person"],
architecture: ["service", "database", "queue", "loadbalancer", "external"],
sequence: ["participant", "fragment"],
flowchart: ["process", "decision", "terminal", "io", "subprocess"],
};
const VALID_EDGE_TYPES: Record<string, string[] | null> = {
bpmn: ["sequence", "message", "association"],
er: null,
orgchart: null,
architecture: null,
sequence: ["sync", "async", "return"],
flowchart: null,
};
const MAX_NODES = 200;
const columnSchema = z.object({
name: z.string(),
type: z.string(),
isPrimaryKey: z.boolean().optional(),
isForeignKey: z.boolean().optional(),
isNullable: z.boolean().optional(),
isUnique: z.boolean().optional(),
references: z.string().optional(),
});
const nodeSchema = z.object({
id: z.string(),
type: z.string(),
label: z.string(),
tag: z.string().optional(),
icon: z.string().optional(),
color: z.string().optional(),
w: z.number().optional(),
position: z
.object({ x: z.number(), y: z.number() })
.optional(),
lane: z.string().optional(),
group: z.string().optional(),
columns: z.array(columnSchema).optional(),
lifeline: z.boolean().optional(),
parentId: z.string().optional(),
});
const edgeSchema = z.object({
id: z.string(),
from: z.string(),
to: z.string(),
label: z.string().optional(),
color: z.string().optional(),
type: z.string().optional(),
cardinality: z.string().optional(),
});
const metaSchema = z.object({
diagramType: z.enum(VALID_DIAGRAM_TYPES),
title: z.string().min(1),
description: z.string().optional(),
version: z.string().optional(),
layoutDirection: z.enum(["DOWN", "RIGHT", "LEFT", "UP"]).optional(),
edgeRouting: z.enum(["ORTHOGONAL", "SPLINES", "POLYLINE"]).optional(),
});
const poolSchema = z.object({
id: z.string(),
label: z.string(),
lanes: z.array(
z.object({
id: z.string(),
label: z.string(),
}),
),
});
const groupSchema = z.object({
id: z.string(),
label: z.string(),
color: z.string().optional(),
});
export const graphPatchSchema = z.object({
meta: metaSchema,
nodes: z.array(nodeSchema).max(MAX_NODES),
edges: z.array(edgeSchema),
pools: z.array(poolSchema).optional(),
groups: z.array(groupSchema).optional(),
});
export type GraphPatch = z.infer<typeof graphPatchSchema>;
/**
* Validate that node types match the diagram type's allowed types.
* Returns an array of error messages (empty if valid).
*/
export function validateNodeTypes(patch: GraphPatch): string[] {
const diagramType = patch.meta.diagramType;
const validTypes = VALID_NODE_TYPES[diagramType];
if (!validTypes) return [];
const errors: string[] = [];
for (const node of patch.nodes) {
const bare = node.type.includes(":") ? node.type.split(":")[1]! : node.type;
if (!validTypes.includes(bare)) {
errors.push(
`Invalid node type "${node.type}" for ${diagramType}. Valid types: ${validTypes.join(", ")}`,
);
}
}
return errors;
}
/**
* Validate that edge types match the diagram type's allowed types.
* Returns an array of error messages (empty if valid).
*/
export function validateEdgeTypes(patch: GraphPatch): string[] {
const diagramType = patch.meta.diagramType;
const validTypes = VALID_EDGE_TYPES[diagramType];
if (!validTypes) return [];
const errors: string[] = [];
for (const edge of patch.edges) {
if (edge.type && !validTypes.includes(edge.type)) {
errors.push(
`Invalid edge type "${edge.type}" for ${diagramType}. Valid types: ${validTypes.join(", ")}`,
);
}
}
return errors;
}
/**
* Validate that all edge from/to references point to existing node IDs.
*/
export function validateEdgeReferences(patch: GraphPatch): string[] {
const nodeIds = new Set(patch.nodes.map((n) => n.id));
const errors: string[] = [];
for (const edge of patch.edges) {
if (!nodeIds.has(edge.from)) {
errors.push(`Edge "${edge.id}" references non-existent source node "${edge.from}"`);
}
if (!nodeIds.has(edge.to)) {
errors.push(`Edge "${edge.id}" references non-existent target node "${edge.to}"`);
}
}
return errors;
}
/**
* Validate that all node and edge IDs are unique.
*/
export function validateUniqueIds(patch: GraphPatch): string[] {
const errors: string[] = [];
const nodeIds = new Set<string>();
for (const node of patch.nodes) {
if (nodeIds.has(node.id)) {
errors.push(`Duplicate node ID "${node.id}"`);
}
nodeIds.add(node.id);
}
const edgeIds = new Set<string>();
for (const edge of patch.edges) {
if (edgeIds.has(edge.id)) {
errors.push(`Duplicate edge ID "${edge.id}"`);
}
edgeIds.add(edge.id);
}
return errors;
}
/**
* Full validation: schema + type-specific + referential integrity + unique IDs.
* Returns { success: true, data } or { success: false, errors }.
*/
export function validateGraphPatch(
input: unknown,
): { success: true; data: GraphPatch } | { success: false; errors: string[] } {
const result = graphPatchSchema.safeParse(input);
if (!result.success) {
return {
success: false,
errors: result.error.issues.map((i) => `${i.path.join(".")}: ${i.message}`),
};
}
const patch = result.data;
const typeErrors = [
...validateUniqueIds(patch),
...validateNodeTypes(patch),
...validateEdgeTypes(patch),
...validateEdgeReferences(patch),
];
if (typeErrors.length > 0) {
return { success: false, errors: typeErrors };
}
return { success: true, data: patch };
}

View File

@@ -25,14 +25,15 @@ describe("buildCopilotSystemPrompt", () => {
const prompt = buildCopilotSystemPrompt(type);
expect(prompt).toContain(`Type: ${type.toUpperCase()}`);
expect(prompt).toContain("domaingraph AI copilot");
expect(prompt).toContain("CHAT-ONLY mode");
}
});
it("should mention chat-only constraint", () => {
const prompt = buildCopilotSystemPrompt("er");
expect(prompt).toContain("CHAT-ONLY mode");
expect(prompt).toContain("cannot modify the diagram directly");
it("should include generation instructions with generateDiagram tool", () => {
const prompt = buildCopilotSystemPrompt("bpmn");
expect(prompt).toContain("generateDiagram");
expect(prompt).toContain("Generate complete diagrams");
expect(prompt).not.toContain("CHAT-ONLY");
expect(prompt).not.toContain("cannot modify the diagram directly");
});
it("should include today's date", () => {
@@ -52,4 +53,124 @@ describe("buildCopilotSystemPrompt", () => {
expect(prompt).toContain("services");
expect(prompt).toContain("databases");
});
it("should include BPMN-specific node types", () => {
const prompt = buildCopilotSystemPrompt("bpmn");
expect(prompt).toContain("activity");
expect(prompt).toContain("gateway-exclusive");
expect(prompt).toContain("gateway-parallel");
expect(prompt).toContain("gateway-inclusive");
expect(prompt).toContain("start-event");
expect(prompt).toContain("end-event");
expect(prompt).toContain("event-timer");
expect(prompt).toContain("event-message");
expect(prompt).toContain("subprocess");
expect(prompt).toContain("data-object");
expect(prompt).toContain("annotation");
});
it("should include BPMN-specific edge types", () => {
const prompt = buildCopilotSystemPrompt("bpmn");
expect(prompt).toContain("sequence");
expect(prompt).toContain("message");
expect(prompt).toContain("association");
});
it("should include BPMN pools and lanes guidance", () => {
const prompt = buildCopilotSystemPrompt("bpmn");
expect(prompt).toContain("pools");
expect(prompt).toContain("lanes");
});
it("should include node types for each diagram type", () => {
const typeNodeMap: Record<DiagramType, string[]> = {
bpmn: ["activity", "gateway-exclusive", "start-event", "end-event"],
er: ["entity"],
orgchart: ["person"],
architecture: ["service", "database", "queue", "loadbalancer", "external"],
sequence: ["participant", "fragment"],
flowchart: ["process", "decision", "terminal", "io"],
};
for (const [type, nodes] of Object.entries(typeNodeMap)) {
const prompt = buildCopilotSystemPrompt(type as DiagramType);
for (const node of nodes) {
expect(prompt).toContain(node);
}
}
});
it("should include sequence-specific edge types", () => {
const prompt = buildCopilotSystemPrompt("sequence");
expect(prompt).toContain("sync");
expect(prompt).toContain("async");
expect(prompt).toContain("return");
});
it("should include ER column format with key fields", () => {
const prompt = buildCopilotSystemPrompt("er");
expect(prompt).toContain("columns");
expect(prompt).toContain("isPrimaryKey");
expect(prompt).toContain("isForeignKey");
expect(prompt).toContain("isNullable");
});
it("should include type inference instructions for all types", () => {
const prompt = buildCopilotSystemPrompt("bpmn");
expect(prompt).toContain("infer");
expect(prompt).toContain("bpmn");
expect(prompt).toContain("er");
expect(prompt).toContain("orgchart");
expect(prompt).toContain("architecture");
expect(prompt).toContain("sequence");
expect(prompt).toContain("flowchart");
});
it("should include 200-node generation cap", () => {
const prompt = buildCopilotSystemPrompt("bpmn");
expect(prompt).toContain("200");
});
it("should include graph context when provided", () => {
const graphContext = '{"nodes":[],"edges":[]}';
const prompt = buildCopilotSystemPrompt("bpmn", { graphContext });
expect(prompt).toContain(graphContext);
});
it("should show empty canvas when no graph context", () => {
const prompt = buildCopilotSystemPrompt("bpmn");
expect(prompt).toContain("Empty canvas");
});
it("should include a JSON example for each diagram type", () => {
const types: DiagramType[] = [
"bpmn",
"er",
"orgchart",
"architecture",
"sequence",
"flowchart",
];
for (const type of types) {
const prompt = buildCopilotSystemPrompt(type);
expect(prompt).toContain("```json");
expect(prompt).toContain(`"diagramType":"${type}"`);
}
});
it("should include GraphData format specification", () => {
const prompt = buildCopilotSystemPrompt("bpmn");
expect(prompt).toContain("meta");
expect(prompt).toContain("nodes");
expect(prompt).toContain("edges");
expect(prompt).toContain("diagramType");
expect(prompt).toContain("layoutDirection");
});
it("should include ID generation convention", () => {
const prompt = buildCopilotSystemPrompt("flowchart");
expect(prompt).toContain('"n1"');
expect(prompt).toContain('"e1"');
});
});

View File

@@ -9,25 +9,137 @@ const DIAGRAM_DESCRIPTIONS: Record<DiagramType, string> = {
flowchart: "Flowchart — decision flows with processes, decisions, terminals, I/O, and subprocesses",
};
export function buildCopilotSystemPrompt(diagramType: DiagramType): string {
const description = DIAGRAM_DESCRIPTIONS[diagramType];
const NODE_TYPE_REFERENCE: Record<DiagramType, string> = {
bpmn: `- activity: Task or action step
- subprocess: Collapsed subprocess
- start-event: Process start
- end-event: Process end
- event-timer: Timer event
- event-message: Message event
- gateway-exclusive: XOR decision (one path)
- gateway-parallel: AND split/join (all paths)
- gateway-inclusive: OR split/join (one or more)
- data-object: Data artifact
- annotation: Text note`,
er: `- entity: Database table (use columns[] for attributes)`,
orgchart: `- person: Team member or role`,
architecture: `- service: Application or microservice
- database: Data store
- queue: Message queue or event bus
- loadbalancer: Load balancer or API gateway
- external: External system`,
sequence: `- participant: Actor or system (set lifeline: true)
- fragment: Combined fragment (loop, alt, opt)`,
flowchart: `- process: Action or operation
- decision: Yes/no branch
- terminal: Start or end point
- io: Input/output step
- subprocess: Nested process`,
};
return `You are the domaingraph AI copilot — a diagram design assistant.
const EDGE_TYPE_REFERENCE: Record<DiagramType, string> = {
bpmn: `- sequence (default): Flow between activities
- message: Cross-pool message flow
- association: Link to annotations/data`,
er: `- (default): Relationship. Set cardinality: "1:1", "1:N", "N:1", or "M:N"`,
orgchart: `- (default): Reports-to relationship`,
architecture: `- (default): Connection. Use label for protocol/description`,
sequence: `- sync (default): Synchronous call (solid arrow)
- async: Asynchronous message (open arrow)
- return: Response (dashed arrow)`,
flowchart: `- (default): Flow. Use label for conditions ("Yes", "No")`,
};
const TYPE_INFERENCE_RULES = `If the diagram type is not established, infer from the user's description:
- Business processes, workflows, approvals, order handling → bpmn
- Database schemas, tables, entities, data models → er
- Team structures, org hierarchies, reporting lines → orgchart
- System design, microservices, infrastructure, APIs → architecture
- Interactions between actors over time, API calls, request/response → sequence
- Decision logic, algorithms, if/else flows → flowchart`;
const DIAGRAM_EXAMPLES: Record<DiagramType, string> = {
bpmn: `{"meta":{"diagramType":"bpmn","title":"Order Process","layoutDirection":"RIGHT"},"nodes":[{"id":"n1","type":"start-event","label":"Order received"},{"id":"n2","type":"activity","label":"Validate order"},{"id":"n3","type":"gateway-exclusive","label":"Valid?"},{"id":"n4","type":"activity","label":"Process payment"},{"id":"n5","type":"end-event","label":"Complete"}],"edges":[{"id":"e1","from":"n1","to":"n2"},{"id":"e2","from":"n2","to":"n3"},{"id":"e3","from":"n3","to":"n4","label":"Yes"},{"id":"e4","from":"n4","to":"n5"}]}`,
er: `{"meta":{"diagramType":"er","title":"Blog Schema"},"nodes":[{"id":"n1","type":"entity","label":"User","columns":[{"name":"id","type":"uuid","isPrimaryKey":true},{"name":"email","type":"varchar","isUnique":true}]},{"id":"n2","type":"entity","label":"Post","columns":[{"name":"id","type":"uuid","isPrimaryKey":true},{"name":"author_id","type":"uuid","isForeignKey":true,"references":"User.id"},{"name":"title","type":"varchar"}]}],"edges":[{"id":"e1","from":"n1","to":"n2","cardinality":"1:N","label":"writes"}]}`,
orgchart: `{"meta":{"diagramType":"orgchart","title":"Engineering Team","layoutDirection":"DOWN"},"nodes":[{"id":"n1","type":"person","label":"VP Engineering"},{"id":"n2","type":"person","label":"Frontend Lead"},{"id":"n3","type":"person","label":"Backend Lead"}],"edges":[{"id":"e1","from":"n1","to":"n2"},{"id":"e2","from":"n1","to":"n3"}]}`,
architecture: `{"meta":{"diagramType":"architecture","title":"Web Stack"},"nodes":[{"id":"n1","type":"loadbalancer","label":"API Gateway"},{"id":"n2","type":"service","label":"Auth Service"},{"id":"n3","type":"service","label":"User Service"},{"id":"n4","type":"database","label":"PostgreSQL"}],"edges":[{"id":"e1","from":"n1","to":"n2","label":"REST"},{"id":"e2","from":"n1","to":"n3","label":"REST"},{"id":"e3","from":"n2","to":"n4"},{"id":"e4","from":"n3","to":"n4"}]}`,
sequence: `{"meta":{"diagramType":"sequence","title":"Login Flow"},"nodes":[{"id":"n1","type":"participant","label":"Browser","lifeline":true},{"id":"n2","type":"participant","label":"Auth API","lifeline":true},{"id":"n3","type":"participant","label":"Database","lifeline":true}],"edges":[{"id":"e1","from":"n1","to":"n2","type":"sync","label":"POST /login"},{"id":"e2","from":"n2","to":"n3","type":"sync","label":"SELECT user"},{"id":"e3","from":"n3","to":"n2","type":"return","label":"user record"},{"id":"e4","from":"n2","to":"n1","type":"return","label":"JWT token"}]}`,
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"}]}`,
};
export function buildCopilotSystemPrompt(
diagramType: DiagramType,
options?: { graphContext?: string },
): string {
const description = DIAGRAM_DESCRIPTIONS[diagramType];
const graphContext = options?.graphContext;
const date = new Date().toLocaleDateString("en-US", {
year: "numeric",
month: "short",
day: "2-digit",
weekday: "short",
});
const typeSpecificSections: string[] = [];
if (diagramType === "er") {
typeSpecificSections.push(`### E-R column format
Each entity node has a \`columns\` array: [{ name, type, isPrimaryKey?, isForeignKey?, isNullable?, isUnique?, references? }]`);
}
if (diagramType === "bpmn") {
typeSpecificSections.push(`### BPMN pools and lanes
Use \`pools\` to define swim lanes. Assign nodes to lanes via the \`lane\` field matching a lane ID.`);
}
return `You are the domaingraph AI copilot — a diagram design assistant that can discuss, generate, and modify diagrams.
## Current diagram
Type: ${diagramType.toUpperCase()}${description}
## Your role
- Help users think through their diagram design
- Explain diagram concepts and best practices for ${diagramType.toUpperCase()} diagrams
- Suggest improvements, missing elements, or structural changes
- Answer questions about the current diagram or diagram type
- Keep responses concise and diagram-focused
## Capabilities
- Answer questions about diagram design and best practices
- Generate complete diagrams from natural language descriptions
- Modify existing diagrams based on user requests
- Analyze the current diagram and suggest improvements
## Important constraints
- You are in CHAT-ONLY mode: you can discuss and advise, but you cannot modify the diagram directly yet
- When users ask you to add or change elements, explain what you would do and tell them this capability is coming soon
- Use markdown formatting for clarity (bold, lists, code blocks)
- Do not use h1 headings in responses
- Today's date is ${new Date().toLocaleDateString("en-US", { year: "numeric", month: "short", day: "2-digit", weekday: "short" })}`;
## Generating and modifying diagrams
When the user describes what they want to build, asks you to create a diagram, or requests changes, use the \`generateDiagram\` tool.
1. Briefly explain what you'll create or change (1-2 sentences)
2. Call the \`generateDiagram\` tool with the complete GraphData object
When modifying an existing diagram, include ALL nodes and edges in the output (not just changes). The tool output replaces the entire graph.
### GraphData format
The tool accepts:
- \`meta\`: { diagramType, title, description?, layoutDirection?: "DOWN"|"RIGHT"|"LEFT"|"UP" }
- \`nodes\`: [{ id, type, label, tag?, columns?, lane?, group?, lifeline? }]
- \`edges\`: [{ id, from, to, label?, type?, cardinality? }]
- \`pools\`: (BPMN only) [{ id, label, lanes: [{ id, label }] }]
- \`groups\`: [{ id, label, color? }]
### Node types for ${diagramType.toUpperCase()}
${NODE_TYPE_REFERENCE[diagramType]}
### Edge types for ${diagramType.toUpperCase()}
${EDGE_TYPE_REFERENCE[diagramType]}
${typeSpecificSections.length > 0 ? typeSpecificSections.join("\n\n") + "\n" : ""}## Example
\`\`\`json
${DIAGRAM_EXAMPLES[diagramType]}
\`\`\`
## Type inference
${TYPE_INFERENCE_RULES}
## Current graph state
${graphContext ? `The diagram currently contains:\n\`\`\`json\n${graphContext}\n\`\`\`` : "Empty canvas — no nodes or edges yet."}
## Constraints
- Maximum 200 nodes per generation
- 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}`;
}

View File

@@ -3,3 +3,6 @@ import { diagramTypeEnum } from "@turbostarter/db/schema/diagram";
export type DiagramType = (typeof diagramTypeEnum.enumValues)[number];
export const DIAGRAM_TYPES = diagramTypeEnum.enumValues;
export type { GraphPatch } from "./mutation-schema";
export { graphPatchSchema, validateGraphPatch, validateUniqueIds } from "./mutation-schema";

View File

@@ -32,12 +32,22 @@ export const updateDiagramBodySchema = z
title: z.string().min(1).max(255).optional(),
projectId: z.string().nullable().optional(),
sortOrder: z.number().int().min(0).optional(),
graphData: z
.object({
meta: z.record(z.string(), z.unknown()).optional(),
nodes: z.array(z.record(z.string(), z.unknown())),
edges: z.array(z.record(z.string(), z.unknown())),
pools: z.array(z.record(z.string(), z.unknown())).optional(),
groups: z.array(z.record(z.string(), z.unknown())).optional(),
})
.optional(),
})
.refine(
(data) =>
data.title !== undefined ||
data.projectId !== undefined ||
data.sortOrder !== undefined,
data.sortOrder !== undefined ||
data.graphData !== undefined,
{ message: "At least one field must be provided" },
);