From fa0247304c6a8a055906e5998bbb1d73421baaa1 Mon Sep 17 00:00:00 2001 From: suharshit singh Date: Wed, 16 Sep 2026 00:49:24 +0530 Subject: [PATCH] feat(ai-agent): canvas-aware generation with roles, async edges and repair The design agent now sees what is already on the canvas, so follow-up requests extend the existing design instead of redrawing it, and generated diagrams match the starter templates' visual conventions. - lib/ai/canvas-summary.ts: compact canvas listing with refs (ex-N), kinds, used connection points, connections and notes, passed to the analyze, plan and generate steps - Generation schema: role + kicker instead of legacy dark colorId, sync/async delivery (async edges dashed); edges may reference existing components by ref - buildCanvasGraph: resolves refs, folds redrawn or repeated components into the existing/first node, drops unknown and duplicate links, and assigns connection points counting existing edges - lib/ai/graph-validation.ts + one repair call: duplicate/redrawn components, unknown refs, self loops, repeated links, >4 connections, unconnected components, missing plan components - Shared NODE_ROLES / getRoleFill and SHAPE_KICKERS in types/canvas.ts (starter templates use them) - Plan text no longer shows internal canvas refs Co-Authored-By: Claude Opus 5 --- components/editor/canvas-node.tsx | 14 +-- components/editor/starter-templates.ts | 22 +--- context/progress-tracker.md | 60 +++++++++- lib/ai/canvas-summary.ts | 132 +++++++++++++++++++++ lib/ai/design-agent-engine.ts | 94 +++++++++++++-- lib/ai/graph-validation.ts | 124 ++++++++++++++++++++ lib/ai/prompts.ts | 53 +++++++-- lib/design-generation.ts | 156 ++++++++++++++++++------- src/trigger/design-agent.ts | 55 +++++++-- types/canvas.ts | 40 +++++++ 10 files changed, 644 insertions(+), 106 deletions(-) create mode 100644 lib/ai/canvas-summary.ts create mode 100644 lib/ai/graph-validation.ts diff --git a/components/editor/canvas-node.tsx b/components/editor/canvas-node.tsx index a1b3931..07e87cf 100644 --- a/components/editor/canvas-node.tsx +++ b/components/editor/canvas-node.tsx @@ -20,7 +20,7 @@ import { } from "@xyflow/react"; import { getUsedSides } from "@/lib/canvas-connections"; -import { resolveNodeFill, TEXT_NODE_SHAPE } from "@/types/canvas"; +import { resolveNodeFill, SHAPE_KICKERS, TEXT_NODE_SHAPE } from "@/types/canvas"; import type { CanvasNodeData, CanvasNodeShape } from "@/types/canvas"; // --------------------------------------------------------------------------- @@ -92,16 +92,8 @@ const STROKE_WIDTH_ACTIVE = 2.5; const SHAPE_SHADOW_FILTER = "drop-shadow(var(--shadow-flat))"; export const DEFAULT_NODE_FONT_SIZE = 11; -/** Default mono kicker above the label, describing what the shape usually stands for. */ -export const SHAPE_KICKERS: Record = { - rectangle: "Service", - circle: "Event", - diamond: "Decision", - pill: "Queue", - cylinder: "Database", - hexagon: "External", - text: "Note", -}; +// Defined in `types/canvas.ts`; re-exported for the control bar, which imports it from here. +export { SHAPE_KICKERS }; // --------------------------------------------------------------------------- // Label block — mono kicker + Archivo label diff --git a/components/editor/starter-templates.ts b/components/editor/starter-templates.ts index bf7d994..1e86fab 100644 --- a/components/editor/starter-templates.ts +++ b/components/editor/starter-templates.ts @@ -2,14 +2,14 @@ import { assignEdgeHandles } from "@/lib/canvas-connections"; import { CANVAS_EDGE_TYPE, CANVAS_NODE_TYPE, - NODE_FILLS, + getRoleFill, SHAPE_DEFAULTS, type CanvasArrowDirection, type CanvasEdge, type CanvasEdgeData, type CanvasNode, type CanvasShape, - type NodeFill, + type NodeRole, } from "@/types/canvas"; export interface CanvasTemplate { @@ -46,20 +46,6 @@ export interface CanvasTemplate { * has more than four edges. Handles are assigned by `assignEdgeHandles`. */ -type TemplateRole = "entry" | "compute" | "messaging" | "data" | "output"; - -const ROLE_FILLS: Record = { - entry: "blue", - compute: "paper", - messaging: "coral", - data: "amber", - output: "sage", -}; - -function getFillValue(fillId: NodeFill["id"]) { - return (NODE_FILLS.find((fill) => fill.id === fillId) ?? NODE_FILLS[0]).value; -} - /** `cx` / `cy` are the node's centre, so shapes of different sizes on one row or column line up. */ export function createTemplateNode( id: string, @@ -67,7 +53,7 @@ export function createTemplateNode( cx: number, cy: number, shape: CanvasShape, - role: TemplateRole, + role: NodeRole, kicker: string, ): CanvasNode { const { width, height } = SHAPE_DEFAULTS[shape]; @@ -81,7 +67,7 @@ export function createTemplateNode( label, shape, kicker, - color: getFillValue(ROLE_FILLS[role]), + color: getRoleFill(role).value, }, }; } diff --git a/context/progress-tracker.md b/context/progress-tracker.md index ba686c9..d44846b 100644 --- a/context/progress-tracker.md +++ b/context/progress-tracker.md @@ -10,9 +10,9 @@ Update this file whenever the current phase, active feature, or implementation s - Persistent, multi-turn AI design sessions (plan: storage → wire sessions → turn engine → UI cards → generation quality → docs). Steps 1 (storage, PR #21), 2 (sessions in the sidebar, PR #22) and 3 (clarify → plan → generate - turn engine, PR #23) are done; step 4 (question/plan/result cards in the sidebar) is built on - `feat/ai-agent-cards` and awaiting a live run; next is step 5, generation quality (canvas context, kickers, - fills, edge styles, validate-and-repair). + turn engine, PR #23) and 4 (question/plan/result cards, PR #25) are done; step 5 (generation quality: canvas + context, roles and kickers, async edges, validate-and-repair) is built on `feat/ai-generation-quality` and + awaiting its live run; step 6 (docs) follows, then the Specs tab. - After that: the Specs tab (Generate Spec + automatic Markdown download), which remains inert. ## Completed @@ -1373,4 +1373,56 @@ Update this file whenever the current phase, active feature, or implementation s - DB turn checks (15) passed after the shared-reader refactor - Live flow on `gemini-3.5-flash-lite` through the Trigger worker passed (message → questions → skip → plan → generate; see the step 3 entry), including QUESTIONS and PLAN payloads that now store the agent's reply. - - Not yet verified: clicking through the cards in a signed-in browser (chips, typed answers, send, skip, draw). + - Browser check by the user (2026-09-16), both working: a detailed prompt with every question answered, and a + low-detail prompt that skipped the questions and generated a good design. +- AI sessions, step 5: generation quality — canvas context, roles and kickers, async edges, validate and repair + (2026-09-16, branch `feat/ai-generation-quality`): + - Shared node roles: `types/canvas.ts` now has `NODE_ROLES` (entry, compute, messaging, data, output), + `NODE_ROLE_FILLS`, and `getRoleFill`, taken from the starter templates' private role map, which now uses them. + `SHAPE_KICKERS` moved here from `canvas-node.tsx` (re-exported there for the control bar) so server code can read + it without importing React Flow. + - `lib/ai/canvas-summary.ts`: `summarizeCanvas` turns the room into components with refs (`ex-1`…), kind (kicker or + the shape default), and used connection points; connections (label, async); and text notes. Capped at 60 + components, 80 connections, 10 notes. `formatCanvasSummary` renders it for prompts. + - `lib/design-generation.ts`: the generation schema replaces the legacy dark `colorId` with `role` and `kicker`, and + edges gain `delivery: sync | async`. Edge endpoints may be a new node id or an existing ref. `buildCanvasGraph` + takes the existing canvas: refs resolve to existing node ids; a generated node that repeats an existing label is + not drawn and its edges attach to the existing node; repeated labels among new nodes fold into the first; + unknown endpoints, self loops, and repeated pairs (either direction) are dropped; fills come from the role; async + edges are dashed; connection points are assigned with existing edges counted. `MAX_GRAPH_NODES` / + `MAX_GRAPH_EDGES` are exported. + - `lib/ai/graph-validation.ts`: `validateDesignGraph` reports, in messages written for the model: over-limit counts, + duplicate ids, new ids shaped like refs, components already on the canvas, duplicate labels, unknown references, + self loops, repeated connections, components over 4 connections (existing ones included), unconnected new + components, and plan components missing from the diagram. + - `lib/ai/prompts.ts`: the canvas summary goes into analyze (treat requests as changes to the existing design), + plan (list only components to add, refer to existing ones by name), and generate (connect by ref, roles, + kickers, async delivery, every new component connected). New `buildRepairPrompt`. + - `lib/ai/design-agent-engine.ts`: `analyzeTurn`, `draftPlan`, and `generateDesignGraph` take the canvas summary. + `generateDesignGraph` validates the diagram and, when anything is wrong, makes one repair call with the issues + and the previous diagram, keeping whichever attempt has fewer issues (a failed repair keeps the first). It + returns `{ graph, issues, repairedIssues }`. + - `src/trigger/design-agent.ts`: reads the room read-only at the start of every turn and passes the summary to each + step. When writing, it uses the live room for layout and connection points and drops refs to components deleted + since the read. Repaired and remaining issues are logged. Comment updated: a turn makes at most three model calls. + - Validation checks: + - `pnpm typecheck` and `pnpm lint` passed + - Offline checks (36, no model calls): canvas summary (refs only for labelled drawn components, kinds, used points, + connections, notes, empty canvas, 60-component cap); every validation rule, including a clean graph with no + issues; `buildCanvasGraph` (existing duplicates and repeated labels not drawn, role fills, kickers, ref and + duplicate resolution, dropped unknown refs, self loops and repeated pairs, dashed async edges, handles on every + edge, new edges avoiding used sides, only an already-full node sharing a side, refs dropped without canvas + context); all 13 starter templates still on paper fills, entry still draft blue. + - Live extend run on `gemini-3.5-flash-lite` through the Trigger worker, on a throwaway room seeded with the + Microservices template (8 nodes, 8 edges): "Add a Redis cache in front of the user database, and a notification + service that consumes order events." → PLAN in 12.4s listing only the 3 new components (User Cache, Event + Broker, Notification Service) with 2 decisions → generate in 12.6s added 3 nodes (kickers Queue / Worker / + Cache, paper fills) and 4 edges: User Service → User Cache → User DB, and Order Service → Event Broker → + Notification Service, both order-event links dashed. Existing nodes untouched, no component redrawn, no + node forced to share a connection point, counts match the room. Test project, room, sessions, and TaskRuns + removed. 0 failures, 0 warnings. + - Follow-up from that run: the plan's flows read "User Service (ex-7) → User Cache → User DB (ex-8)" — the model + copied internal refs into user-facing plan text despite the prompt. `removeCanvasRefs` in + `design-agent-engine.ts` now strips parenthesised refs ("(ex-7)", "(ex-4, ex-5)") from the plan summary, + responsibilities, flows, decisions, and assumptions after drafting; other parentheses are left alone. Covered + by 3 more offline checks. diff --git a/lib/ai/canvas-summary.ts b/lib/ai/canvas-summary.ts new file mode 100644 index 0000000..7af00a2 --- /dev/null +++ b/lib/ai/canvas-summary.ts @@ -0,0 +1,132 @@ +import { getUsedSides } from "@/lib/canvas-connections"; +import { SHAPE_KICKERS, TEXT_NODE_SHAPE, type CanvasEdge, type CanvasNode } from "@/types/canvas"; + +// --------------------------------------------------------------------------- +// Canvas summary +// +// A compact, model-readable view of what is already on the canvas, so the +// design agent extends an existing design instead of redrawing it. Each drawn +// component gets a short ref ("ex-3") the model can use to connect new +// components to it. +// --------------------------------------------------------------------------- + +export const CANVAS_SUMMARY_LIMITS = { + components: 60, + connections: 80, + notes: 10, +} as const; + +/** Connection points per component; each holds one edge. */ +export const CONNECTION_POINTS_PER_COMPONENT = 4; + +export interface ExistingComponent { + /** Short reference the model uses to connect to this component, e.g. "ex-3". */ + ref: string; + nodeId: string; + label: string; + /** The node's kicker, or its shape's default kicker. */ + kind: string; + /** Connection points already holding an edge (0 to 4). */ + usedConnections: number; +} + +export interface ExistingConnection { + from: string; + to: string; + label?: string; + async: boolean; +} + +export interface CanvasSummary { + components: ExistingComponent[]; + connections: ExistingConnection[]; + /** Text of free text notes on the canvas. */ + notes: string[]; + /** Drawn components on the canvas, including any beyond the listing limit. */ + totalComponents: number; +} + +export const EXISTING_REF_PATTERN = /^ex-\d+$/; + +function labelOf(node: CanvasNode): string { + return typeof node.data.label === "string" ? node.data.label.trim() : ""; +} + +export function summarizeCanvas(nodes: readonly CanvasNode[], edges: readonly CanvasEdge[]): CanvasSummary { + const drawn = nodes.filter((node) => node.data.shape !== TEXT_NODE_SHAPE && labelOf(node).length > 0); + + const components: ExistingComponent[] = drawn + .slice(0, CANVAS_SUMMARY_LIMITS.components) + .map((node, index) => ({ + ref: `ex-${index + 1}`, + nodeId: node.id, + label: labelOf(node), + kind: node.data.kicker?.trim() || SHAPE_KICKERS[node.data.shape ?? "rectangle"], + usedConnections: getUsedSides(edges, node.id).size, + })); + + const refByNodeId = new Map(components.map((component) => [component.nodeId, component.ref])); + + const connections: ExistingConnection[] = edges + .flatMap((edge): ExistingConnection[] => { + const from = refByNodeId.get(edge.source); + const to = refByNodeId.get(edge.target); + if (!from || !to) { + return []; + } + const label = edge.data?.label?.trim(); + return [{ from, to, ...(label ? { label } : {}), async: edge.data?.edgeStyle === "dashed" }]; + }) + .slice(0, CANVAS_SUMMARY_LIMITS.connections); + + const notes = nodes + .filter((node) => node.data.shape === TEXT_NODE_SHAPE) + .map(labelOf) + .filter((note) => note.length > 0) + .slice(0, CANVAS_SUMMARY_LIMITS.notes); + + return { components, connections, notes, totalComponents: drawn.length }; +} + +/** The summary as prompt text. */ +export function formatCanvasSummary(summary: CanvasSummary | null): string { + if (!summary || summary.components.length === 0) { + return summary && summary.notes.length > 0 + ? `The canvas has no components yet. Notes on it:\n${summary.notes.map((note) => `- ${note}`).join("\n")}` + : "The canvas is empty."; + } + + const byRef = new Map(summary.components.map((component) => [component.ref, component])); + const sections = [ + `Components already on the canvas (ref: name, kind, connection points used):\n${summary.components + .map( + (component) => + `- ${component.ref}: ${component.label} (${component.kind}), ` + + `${component.usedConnections}/${CONNECTION_POINTS_PER_COMPONENT} used`, + ) + .join("\n")}`, + ]; + + if (summary.totalComponents > summary.components.length) { + sections.push(`...and ${summary.totalComponents - summary.components.length} more components not listed.`); + } + + if (summary.connections.length > 0) { + sections.push( + `Existing connections:\n${summary.connections + .map((connection) => { + const from = byRef.get(connection.from)?.label ?? connection.from; + const to = byRef.get(connection.to)?.label ?? connection.to; + const details = [connection.label, connection.async ? "async" : ""].filter(Boolean).join(", "); + return `- ${from} (${connection.from}) → ${to} (${connection.to})${details ? `: ${details}` : ""}`; + }) + .join("\n")}`, + ); + } + + if (summary.notes.length > 0) { + sections.push(`Notes on the canvas:\n${summary.notes.map((note) => `- ${note}`).join("\n")}`); + } + + return sections.join("\n\n"); +} diff --git a/lib/ai/design-agent-engine.ts b/lib/ai/design-agent-engine.ts index bb7d7b3..9ccac1a 100644 --- a/lib/ai/design-agent-engine.ts +++ b/lib/ai/design-agent-engine.ts @@ -14,11 +14,14 @@ import { type DesignPlan, type TurnAnalysis, } from "@/lib/ai/agent-schema"; +import { formatCanvasSummary, type CanvasSummary } from "@/lib/ai/canvas-summary"; +import { validateDesignGraph } from "@/lib/ai/graph-validation"; import { ANALYZE_SYSTEM_PROMPT, buildAnalyzeContext, buildGeneratePrompt, buildPlanPrompt, + buildRepairPrompt, GENERATE_SYSTEM_PROMPT, PLAN_SYSTEM_PROMPT, } from "@/lib/ai/prompts"; @@ -133,11 +136,15 @@ export function enforceDecision(analysis: TurnAnalysis, payload: DesignAgentPayl }; } -export async function analyzeTurn(model: LanguageModel, payload: DesignAgentPayload): Promise { +export async function analyzeTurn( + model: LanguageModel, + payload: DesignAgentPayload, + canvas: CanvasSummary | null = null, +): Promise { const { output } = await withSchemaRetry(() => generateText({ model, - system: `${ANALYZE_SYSTEM_PROMPT}\n\n${buildAnalyzeContext(payload)}`, + system: `${ANALYZE_SYSTEM_PROMPT}\n\n${buildAnalyzeContext(payload, formatCanvasSummary(canvas))}`, messages: toMessages(payload), output: Output.object({ schema: turnAnalysisSchema, name: "turn_analysis" }), providerOptions: thinking(model, "low"), @@ -151,45 +158,110 @@ export async function analyzeTurn(model: LanguageModel, payload: DesignAgentPayl ); } +/** Parenthesised canvas refs such as "(ex-7)" or "(ex-2, ex-3)". */ +const CANVAS_REF_MENTION = /\s*\((?:\s*ex-\d+\s*,?)+\)/g; + +function withoutRefs(text: string): string { + return text.replace(CANVAS_REF_MENTION, ""); +} + +/** + * Removes canvas refs the model copies from the canvas summary into plan text. + * Refs are internal handles for connecting to existing components; the plan + * is shown to the user, who only knows components by name. + */ +export function removeCanvasRefs(plan: DesignPlan): DesignPlan { + return { + summary: withoutRefs(plan.summary), + components: plan.components.map((component) => ({ + ...component, + responsibility: withoutRefs(component.responsibility), + })), + flows: plan.flows.map(withoutRefs), + decisions: plan.decisions.map((decision) => ({ + ...decision, + choice: withoutRefs(decision.choice), + rationale: withoutRefs(decision.rationale), + alternatives: decision.alternatives.map(withoutRefs), + })), + assumptions: plan.assumptions.map(withoutRefs), + }; +} + export async function draftPlan( model: LanguageModel, brief: DesignBrief, previousPlan: DesignPlan | null, latestInput: string, + canvas: CanvasSummary | null = null, ): Promise { const { output } = await withSchemaRetry(() => generateText({ model, system: PLAN_SYSTEM_PROMPT, - prompt: buildPlanPrompt(brief, previousPlan, latestInput), + prompt: buildPlanPrompt(brief, previousPlan, latestInput, formatCanvasSummary(canvas)), output: Output.object({ schema: designPlanSchema, name: "design_plan" }), providerOptions: thinking(model, "medium"), timeout: { totalMs: CALL_TIMEOUT_MS.plan }, }), ); - return clampPlan(output); + return removeCanvasRefs(clampPlan(output)); } -export async function generateDesignGraph( - model: LanguageModel, - brief: DesignBrief, - plan: DesignPlan, -): Promise { +export interface GeneratedDesign { + graph: DesignGraph; + /** Problems still present in `graph`; buildCanvasGraph drops what it cannot draw. */ + issues: string[]; + /** Problems found in the first attempt and sent back for repair, if a repair ran. */ + repairedIssues: string[]; +} + +async function requestGraph(model: LanguageModel, prompt: string): Promise { const { output } = await withSchemaRetry(() => generateText({ model, system: GENERATE_SYSTEM_PROMPT, - prompt: buildGeneratePrompt(brief, plan), + prompt, output: Output.object({ schema: designGraphSchema, name: "design_graph" }), providerOptions: thinking(model, "low"), timeout: { totalMs: CALL_TIMEOUT_MS.generate }, }), ); - return output; } +/** + * Generates the diagram for an approved plan, validates it, and when anything + * is wrong asks the model once to fix exactly those problems. The attempt with + * fewer remaining problems wins; a failed repair call keeps the first attempt. + */ +export async function generateDesignGraph( + model: LanguageModel, + brief: DesignBrief, + plan: DesignPlan, + canvas: CanvasSummary | null = null, +): Promise { + const prompt = buildGeneratePrompt(brief, plan, formatCanvasSummary(canvas)); + const first = await requestGraph(model, prompt); + const firstIssues = validateDesignGraph(first, { canvas, plan }); + + if (firstIssues.length === 0) { + return { graph: first, issues: [], repairedIssues: [] }; + } + + try { + const repaired = await requestGraph(model, buildRepairPrompt(prompt, first, firstIssues)); + const remaining = validateDesignGraph(repaired, { canvas, plan }); + return remaining.length <= firstIssues.length + ? { graph: repaired, issues: remaining, repairedIssues: firstIssues } + : { graph: first, issues: firstIssues, repairedIssues: firstIssues }; + } catch (error) { + console.warn("[design-agent] repair attempt failed; keeping the first diagram", error); + return { graph: first, issues: firstIssues, repairedIssues: [] }; + } +} + /** An empty brief, for a generate turn that somehow has none stored. */ export function emptyBrief(): DesignBrief { return designBriefSchema.parse({ diff --git a/lib/ai/graph-validation.ts b/lib/ai/graph-validation.ts new file mode 100644 index 0000000..fc501e2 --- /dev/null +++ b/lib/ai/graph-validation.ts @@ -0,0 +1,124 @@ +import type { DesignPlan } from "@/lib/ai/agent-schema"; +import { CONNECTION_POINTS_PER_COMPONENT, EXISTING_REF_PATTERN, type CanvasSummary } from "@/lib/ai/canvas-summary"; +import { MAX_GRAPH_EDGES, MAX_GRAPH_NODES, type DesignGraph } from "@/lib/design-generation"; + +// --------------------------------------------------------------------------- +// Generated graph validation +// +// Checks a generated diagram against the rules the canvas and the plan impose. +// The messages are written for the model: they are sent back verbatim in a +// single repair attempt. Whatever survives repair is still sanitised by +// buildCanvasGraph, so these are quality checks, not safety checks. +// --------------------------------------------------------------------------- + +interface ValidationContext { + canvas?: CanvasSummary | null; + plan?: DesignPlan | null; +} + +function normalizeLabel(label: string): string { + return label.trim().toLowerCase(); +} + +export function validateDesignGraph(graph: DesignGraph, { canvas, plan }: ValidationContext = {}): string[] { + const issues: string[] = []; + const existing = canvas?.components ?? []; + const existingByRef = new Map(existing.map((component) => [component.ref, component])); + const existingByLabel = new Map(existing.map((component) => [normalizeLabel(component.label), component])); + + if (graph.nodes.length > MAX_GRAPH_NODES) { + issues.push(`There are ${graph.nodes.length} components; draw at most ${MAX_GRAPH_NODES}.`); + } + if (graph.edges.length > MAX_GRAPH_EDGES) { + issues.push(`There are ${graph.edges.length} connections; draw at most ${MAX_GRAPH_EDGES}.`); + } + + // --- components + const labelById = new Map(); + const seenLabels = new Set(); + + for (const node of graph.nodes) { + if (labelById.has(node.id)) { + issues.push(`Component id "${node.id}" is used more than once; give every component a unique id.`); + continue; + } + if (EXISTING_REF_PATTERN.test(node.id)) { + issues.push( + `Component id "${node.id}" looks like a ref for an existing canvas component; choose a different id for new components.`, + ); + } + labelById.set(node.id, node.label); + + const label = normalizeLabel(node.label); + const onCanvas = existingByLabel.get(label); + if (onCanvas) { + issues.push( + `"${node.label}" is already on the canvas as ${onCanvas.ref}; do not draw it again, connect to ${onCanvas.ref} instead.`, + ); + } else if (seenLabels.has(label)) { + issues.push(`"${node.label}" appears more than once; each component needs a unique name.`); + } + seenLabels.add(label); + } + + const nameOf = (id: string) => labelById.get(id) ?? existingByRef.get(id)?.label ?? id; + + // --- connections + const connectionCounts = new Map(existing.map((component) => [component.ref, component.usedConnections])); + const seenPairs = new Set(); + + for (const edge of graph.edges) { + const unknown = [edge.source, edge.target].filter((id) => !labelById.has(id) && !existingByRef.has(id)); + if (unknown.length > 0) { + issues.push( + `Connection "${edge.id}" references ${unknown.map((id) => `"${id}"`).join(" and ")}, which is neither a new component id nor an existing ref.`, + ); + continue; + } + if (edge.source === edge.target) { + issues.push(`Connection "${edge.id}" connects "${nameOf(edge.source)}" to itself.`); + continue; + } + + const pair = [edge.source, edge.target].sort().join("|"); + if (seenPairs.has(pair)) { + issues.push(`"${nameOf(edge.source)}" and "${nameOf(edge.target)}" are connected more than once; keep one connection.`); + continue; + } + seenPairs.add(pair); + + for (const id of [edge.source, edge.target]) { + connectionCounts.set(id, (connectionCounts.get(id) ?? 0) + 1); + } + } + + for (const [id, count] of connectionCounts) { + if (count > CONNECTION_POINTS_PER_COMPONENT) { + const note = existingByRef.has(id) ? " (including connections it already has on the canvas)" : ""; + issues.push( + `"${nameOf(id)}" has ${count} connections${note}; each component has ${CONNECTION_POINTS_PER_COMPONENT} connection points, so keep it to ${CONNECTION_POINTS_PER_COMPONENT}.`, + ); + } + } + + // --- unconnected components + if (graph.nodes.length > 1 || existing.length > 0) { + for (const [id, label] of labelById) { + if (!connectionCounts.has(id)) { + issues.push(`"${label}" is not connected to anything.`); + } + } + } + + // --- plan coverage + if (plan) { + for (const component of plan.components) { + const name = normalizeLabel(component.name); + if (!seenLabels.has(name) && !existingByLabel.has(name)) { + issues.push(`Plan component "${component.name}" is missing from the diagram.`); + } + } + } + + return issues; +} diff --git a/lib/ai/prompts.ts b/lib/ai/prompts.ts index b57b763..b653e32 100644 --- a/lib/ai/prompts.ts +++ b/lib/ai/prompts.ts @@ -27,14 +27,19 @@ export const ANALYZE_SYSTEM_PROMPT = [ " without asking for changes.", " A detailed first message can go straight to a plan; a vague one should get questions.", "3. Write a short reply: one or two plain sentences, no markdown, no restating the questions.", + "", + "The canvas may already hold a design; its contents are listed below. Treat the user's request as", + "a change to that design unless they ask for a separate system: do not ask about anything the", + "canvas already answers, and plan only what needs to be added.", ].join("\n"); function describeBrief(brief: DesignBrief | null): string { return brief ? JSON.stringify(brief, null, 2) : "No brief yet — this is the start of the conversation."; } -export function buildAnalyzeContext(payload: DesignAgentPayload): string { +export function buildAnalyzeContext(payload: DesignAgentPayload, canvasText: string): string { const notes = [ + `Current canvas:\n${canvasText}`, `Current brief:\n${describeBrief(payload.brief)}`, payload.plan ? `A plan has been proposed:\n${describePlan(payload.plan)}` @@ -70,15 +75,23 @@ export const PLAN_SYSTEM_PROMPT = [ " to the brief and name realistic alternatives.", "- Assumptions: carry over the brief's assumptions and add any the plan depends on.", "When a previous plan exists and the user asked for changes, revise that plan rather than starting over.", + "When the canvas already has components, list only the components to add. Refer to existing", + "components by their exact names in flows and decisions, and never list them again as components.", ].join("\n"); function describePlan(plan: DesignPlan): string { return JSON.stringify(plan, null, 2); } -export function buildPlanPrompt(brief: DesignBrief, previousPlan: DesignPlan | null, latestInput: string): string { +export function buildPlanPrompt( + brief: DesignBrief, + previousPlan: DesignPlan | null, + latestInput: string, + canvasText: string, +): string { return [ `Requirements brief:\n${JSON.stringify(brief, null, 2)}`, + `Current canvas:\n${canvasText}`, previousPlan ? `Previous plan:\n${describePlan(previousPlan)}` : "", `The user's latest message:\n${latestInput}`, ] @@ -92,17 +105,35 @@ export function buildPlanPrompt(brief: DesignBrief, previousPlan: DesignPlan | n export const GENERATE_SYSTEM_PROMPT = [ "You are a system design architect. Turn the approved plan into a component diagram.", - "Draw exactly the plan's components, using each component's name verbatim as its label.", - "Connect components in the direction traffic actually flows, following the plan's flows, and label", - "a connection only when the protocol or payload is not obvious from the two components it joins.", + "Draw exactly the plan's components as new nodes, using each component's name verbatim as its label.", + "Components already on the canvas are listed with refs such as 'ex-2': never draw them again;", + "connect new components to them by using the ref as the connection's source or target.", + "Give every node a role (entry, compute, messaging, data, output) and a short kicker such as", + "'Gateway', 'Worker', 'Queue', or 'Database'.", + "Connect components in the direction traffic actually flows, following the plan's flows. Mark a", + "connection 'async' for events, queues, streams, replication, and background work; otherwise 'sync'.", + "Label a connection only when the protocol or payload is not obvious from the two components it joins.", "Lay the diagram out top to bottom: entry points at the lowest y values, datastores at the", "highest. Components that sit at the same level of the request path share a y value.", - "Every component has four connection points (top, right, bottom, left) and each point takes", - "one connection, so give each component at most four connections in total.", + "Every component has four connection points and each takes one connection, so give each component", + "at most four connections in total, counting the ones an existing component already uses.", + "Every new component must be connected to at least one other component.", ].join(" "); -export function buildGeneratePrompt(brief: DesignBrief, plan: DesignPlan): string { - return [`Approved plan:\n${describePlan(plan)}`, `Requirements brief:\n${JSON.stringify(brief, null, 2)}`].join( - "\n\n", - ); +export function buildGeneratePrompt(brief: DesignBrief, plan: DesignPlan, canvasText: string): string { + return [ + `Approved plan:\n${describePlan(plan)}`, + `Current canvas:\n${canvasText}`, + `Requirements brief:\n${JSON.stringify(brief, null, 2)}`, + ].join("\n\n"); +} + +/** The generate prompt plus the previous attempt and what was wrong with it. */ +export function buildRepairPrompt(generatePrompt: string, previousGraph: unknown, issues: readonly string[]): string { + return [ + generatePrompt, + `Your previous diagram had these problems:\n${issues.map((issue) => `- ${issue}`).join("\n")}`, + "Return the complete corrected diagram — every component and connection — fixing each problem and keeping everything else the same.", + `Previous diagram:\n${JSON.stringify(previousGraph, null, 2)}`, + ].join("\n\n"); } diff --git a/lib/design-generation.ts b/lib/design-generation.ts index a80bac5..dbf2466 100644 --- a/lib/design-generation.ts +++ b/lib/design-generation.ts @@ -1,28 +1,28 @@ import { z } from "zod"; +import type { CanvasSummary } from "@/lib/ai/canvas-summary"; import { assignEdgeHandles } from "@/lib/canvas-connections"; import { CANVAS_EDGE_TYPE, CANVAS_NODE_TYPE, CANVAS_SHAPES, EDGE_ARROW_DIRECTIONS, - NODE_COLOR_IDS, - NODE_COLOR_PALETTE, + getRoleFill, + NODE_ROLES, SHAPE_DEFAULTS, type CanvasEdge, type CanvasNode, - type NodeColorId, } from "@/types/canvas"; // --------------------------------------------------------------------------- // Generation schema // // This is the contract handed to the model. It is deliberately narrower than -// CanvasNode / CanvasEdge: the model only chooses semantics (what exists, how -// it connects, roughly where it sits). Everything the canvas needs in order to -// render — node type, edge type, pixel sizes, resolved hex colors — is filled -// in by buildCanvasGraph() below, so the model can never emit a value the -// canvas is unable to draw. +// CanvasNode / CanvasEdge: the model only chooses semantics (what exists, what +// role it plays, how it connects, roughly where it sits). Everything the canvas +// needs in order to render — node type, edge type, pixel sizes, fills, edge +// styles, connection points — is filled in by buildCanvasGraph() below, so the +// model can never emit a value the canvas is unable to draw. // --------------------------------------------------------------------------- // --------------------------------------------------------------------------- @@ -51,14 +51,17 @@ export function parseDesignAgentStage(value: unknown): DesignAgentStage | null { } /** Upper bound on generated components, to keep a single run legible. */ -const MAX_NODES = 24; +export const MAX_GRAPH_NODES = 24; /** Upper bound on generated connections. */ -const MAX_EDGES = 48; +export const MAX_GRAPH_EDGES = 48; /** A single generated component in the system design. */ export const designNodeSchema = z.object({ - id: z.string().min(1).describe("Short unique slug for this component, e.g. 'api-gateway'."), + id: z + .string() + .min(1) + .describe("Short unique slug for this new component, e.g. 'api-gateway'. Never of the form 'ex-N'."), label: z .string() .min(1) @@ -72,33 +75,49 @@ export const designNodeSchema = z.object({ shape: z .enum(CANVAS_SHAPES) .describe( - "Visual shape. Use 'cylinder' for datastores, 'pill' for gateways and entry points, " + - "'diamond' for routers and decisions, 'rectangle' for services, 'hexagon' for queues " + - "and brokers, 'circle' for clients and external actors.", + "Visual shape. 'pill' for clients and entry points (DNS, CDN, load balancers, gateways), " + + "'rectangle' for services, APIs and workers, 'hexagon' for queues, topics and streams, " + + "'cylinder' for databases, caches and storage, 'diamond' for routers and decisions, " + + "'circle' for external actors and events.", + ), + role: z + .enum(NODE_ROLES) + .describe( + "What the component does, which sets its colour: 'entry' (clients, DNS, CDN, load balancers, gateways), " + + "'compute' (services, APIs, workers), 'messaging' (queues, topics, streams, brokers), " + + "'data' (databases, caches, object storage), 'output' (external systems, consumers, delivery).", ), - colorId: z - .enum(NODE_COLOR_IDS) - .optional() - .describe("Palette entry used to tint the node. Give related components the same entry."), + kicker: z + .string() + .describe("One or two word kind shown above the name, e.g. 'Gateway', 'Worker', 'Queue', 'Database'."), }); /** A single generated connection between two components. */ export const designEdgeSchema = z.object({ id: z.string().min(1).describe("Short unique slug for this connection."), - source: z.string().min(1).describe("`id` of the node the connection starts at."), - target: z.string().min(1).describe("`id` of the node the connection ends at."), + source: z + .string() + .min(1) + .describe("Where the connection starts: a new component's `id`, or an existing component's ref such as 'ex-2'."), + target: z + .string() + .min(1) + .describe("Where the connection ends: a new component's `id`, or an existing component's ref such as 'ex-2'."), label: z.string().optional().describe("Short description of the traffic, e.g. 'writes'."), arrowDirection: z .enum(EDGE_ARROW_DIRECTIONS) .describe("Arrowheads to draw. Use 'forward' for one-way flow."), + delivery: z + .enum(["sync", "async"]) + .describe("'async' for events, queues, streams, replication and background work; otherwise 'sync'."), }); /** The full structured output requested from the model. */ // The limits are stated, not enforced as maxItems: models don't reliably honour // them, and one extra item would fail the run. buildCanvasGraph trims instead. export const designGraphSchema = z.object({ - nodes: z.array(designNodeSchema).min(1).describe(`At most ${MAX_NODES} components.`), - edges: z.array(designEdgeSchema).describe(`At most ${MAX_EDGES} connections.`), + nodes: z.array(designNodeSchema).min(1).describe(`New components to draw. At most ${MAX_GRAPH_NODES}.`), + edges: z.array(designEdgeSchema).describe(`Connections. At most ${MAX_GRAPH_EDGES}.`), }); export type DesignGraph = z.infer; @@ -124,6 +143,12 @@ interface LayoutOrigin { y: number; } +export interface ExistingCanvas { + summary: CanvasSummary; + nodes: readonly CanvasNode[]; + edges: readonly CanvasEdge[]; +} + export interface BuildCanvasGraphOptions { /** * Prefix applied to every generated id, so a run can never collide with @@ -132,10 +157,16 @@ export interface BuildCanvasGraphOptions { idPrefix: string; /** Top-left corner the generated graph is laid out from. */ origin?: LayoutOrigin; + /** + * What is already in the room. Lets edges reference existing components by + * ref, folds redrawn duplicates into the existing node, and keeps connection + * points that existing edges already use. + */ + existing?: ExistingCanvas; } -function resolveColor(colorId: NodeColorId | undefined) { - return NODE_COLOR_PALETTE.find((entry) => entry.id === colorId) ?? NODE_COLOR_PALETTE[0]; +function normalizeLabel(label: string): string { + return label.trim().toLowerCase(); } /** @@ -192,25 +223,55 @@ function assignGridCells(nodes: DesignGraph["nodes"]): Map [component.ref, component.nodeId])); + const existingNodeIdByLabel = new Map( + existingComponents.map((component) => [normalizeLabel(component.label), component.nodeId]), + ); + + /** Design ids that resolve to a node other than their own: existing duplicates and repeated labels. */ + const aliases = new Map(); + const firstIdByLabel = new Map(); const seenNodeIds = new Set(); + const uniqueNodes = design.nodes .filter((node) => { - if (seenNodeIds.has(node.id)) { + if (seenNodeIds.has(node.id) || existingNodeIdByRef.has(node.id)) { return false; } seenNodeIds.add(node.id); + + const label = normalizeLabel(node.label); + const onCanvas = existingNodeIdByLabel.get(label); + if (onCanvas) { + aliases.set(node.id, onCanvas); + return false; + } + + const first = firstIdByLabel.get(label); + if (first) { + aliases.set(node.id, `${idPrefix}-${first}`); + return false; + } + firstIdByLabel.set(label, node.id); return true; }) - .slice(0, MAX_NODES); + .slice(0, MAX_GRAPH_NODES); const cells = assignGridCells(uniqueNodes); const canvasNodeIds = new Map(); @@ -220,8 +281,8 @@ export function buildCanvasGraph( canvasNodeIds.set(node.id, canvasNodeId); const dimensions = SHAPE_DEFAULTS[node.shape]; - const color = resolveColor(node.colorId); const cell = cells.get(node.id) ?? { column: 0, row: 0 }; + const kicker = node.kicker.trim(); return { id: canvasNodeId, @@ -237,31 +298,40 @@ export function buildCanvasGraph( data: { label: node.label, shape: node.shape, - color: color.bg, - textColor: "var(--text-primary)", - strokeColor: "var(--text-primary)", + color: getRoleFill(node.role).value, + ...(kicker ? { kicker } : {}), }, }; }); + const drawnIds = new Set(canvasNodeIds.values()); + const resolveEndpoint = (id: string): string | undefined => { + const resolved = canvasNodeIds.get(id) ?? aliases.get(id) ?? existingNodeIdByRef.get(id); + // An alias to a repeated label points at the first node, which may have been trimmed away. + return resolved && (drawnIds.has(resolved) || !resolved.startsWith(`${idPrefix}-`)) ? resolved : undefined; + }; + const seenEdgeIds = new Set(); + const seenPairs = new Set(); const edges: CanvasEdge[] = []; for (const edge of design.edges) { - if (edges.length >= MAX_EDGES) { + if (edges.length >= MAX_GRAPH_EDGES) { break; } - const source = canvasNodeIds.get(edge.source); - const target = canvasNodeIds.get(edge.target); + const source = resolveEndpoint(edge.source); + const target = resolveEndpoint(edge.target); if (!source || !target || source === target) { continue; } + const pair = [source, target].sort().join("|"); const canvasEdgeId = `${idPrefix}-${edge.id}`; - if (seenEdgeIds.has(canvasEdgeId)) { + if (seenEdgeIds.has(canvasEdgeId) || seenPairs.has(pair)) { continue; } seenEdgeIds.add(canvasEdgeId); + seenPairs.add(pair); const label = edge.label?.trim(); @@ -272,11 +342,15 @@ export function buildCanvasGraph( target, data: { arrowDirection: edge.arrowDirection, + edgeStyle: edge.delivery === "async" ? "dashed" : "solid", ...(label ? { label } : {}), }, }); } - // One edge per node side, so connections don't pile up on a single point. - return { nodes, edges: assignEdgeHandles(nodes, edges) }; + // One edge per node side, counting the sides existing edges already hold. + const existingEdges = existing?.edges ?? []; + const assigned = assignEdgeHandles([...(existing?.nodes ?? []), ...nodes], [...existingEdges, ...edges]); + + return { nodes, edges: assigned.slice(existingEdges.length) }; } diff --git a/src/trigger/design-agent.ts b/src/trigger/design-agent.ts index e74541d..04354a2 100644 --- a/src/trigger/design-agent.ts +++ b/src/trigger/design-agent.ts @@ -3,6 +3,7 @@ import { logger, metadata, task } from "@trigger.dev/sdk/v3"; import type { LanguageModel } from "ai"; import type { DesignAgentPayload, DesignAgentResult, DesignBrief, DesignPlan } from "@/lib/ai/agent-schema"; +import { summarizeCanvas, type CanvasSummary } from "@/lib/ai/canvas-summary"; import { analyzeTurn, draftPlan, @@ -53,6 +54,17 @@ function resolveOrigin(existingNodes: readonly CanvasNode[]): { x: number; y: nu return { x: left, y: bottom + EXISTING_CONTENT_GAP }; } +/** Reads the room without changing it, so the agent knows what is already drawn. */ +async function readCanvas(roomId: string): Promise { + let snapshot: { nodes: CanvasNode[]; edges: CanvasEdge[] } = { nodes: [], edges: [] }; + + await mutateFlow({ client: getLiveblocksClient(), roomId }, (flow) => { + snapshot = { nodes: [...flow.nodes], edges: [...flow.edges] }; + }); + + return summarizeCanvas(snapshot.nodes, snapshot.edges); +} + /** Generates the diagram for an approved plan and writes it into the room. */ async function drawPlan( model: LanguageModel, @@ -60,18 +72,32 @@ async function drawPlan( runId: string, brief: DesignBrief, plan: DesignPlan, + canvas: CanvasSummary, ): Promise<{ nodeCount: number; edgeCount: number }> { setStage("generating"); - const design = await generateDesignGraph(model, brief, plan); - logger.log("Design generated", { nodeCount: design.nodes.length, edgeCount: design.edges.length }); + const { graph, issues, repairedIssues } = await generateDesignGraph(model, brief, plan, canvas); + logger.log("Design generated", { + nodeCount: graph.nodes.length, + edgeCount: graph.edges.length, + repairedIssues, + remainingIssues: issues, + }); setStage("writing"); let counts = { nodeCount: 0, edgeCount: 0 }; await mutateFlow({ client: getLiveblocksClient(), roomId }, (flow) => { - const { nodes, edges } = buildCanvasGraph(design, { + // Collaborators may have deleted components since the canvas was read; refs to them are dropped. + const liveNodeIds = new Set(flow.nodes.map((node) => node.id)); + const summary: CanvasSummary = { + ...canvas, + components: canvas.components.filter((component) => liveNodeIds.has(component.nodeId)), + }; + + const { nodes, edges } = buildCanvasGraph(graph, { idPrefix: runId, origin: resolveOrigin(flow.nodes), + existing: { summary, nodes: flow.nodes, edges: flow.edges }, }); flow.addNodes(nodes); @@ -87,7 +113,9 @@ async function drawPlan( /** * One turn of a design session. Reads the user's turn, then either asks * clarifying questions, proposes (or revises) a plan, or draws the approved - * plan on the canvas. The server stores the returned result on the session. + * plan on the canvas. Every step sees what is already on the canvas, so + * follow-up requests extend the existing design. The server stores the + * returned result on the session. */ export const designAgentTask = task({ id: "design-agent", @@ -95,23 +123,30 @@ export const designAgentTask = task({ // whole turn on top multiplied quota use (up to 9 calls per failure) and could // write a generated diagram twice; a failed turn is shown to the user instead. retry: { maxAttempts: 1 }, - // A turn makes at most two model calls, each bounded by its own timeout; this - // is the backstop so a stuck turn fails and settles instead of staying pending. + // A turn makes at most three model calls (analyze and plan, or generate and a + // repair), each bounded by its own timeout; this is the backstop so a stuck + // turn fails and settles instead of staying pending. maxDuration: 300, run: async (payload: DesignAgentPayload, { ctx }): Promise => { logger.log("Design agent turn", { roomId: payload.roomId, intent: payload.intent }); const model = resolveDesignModel(); + const canvas = await readCanvas(payload.roomId); + logger.log("Canvas read", { + components: canvas.totalComponents, + connections: canvas.connections.length, + notes: canvas.notes.length, + }); // An explicit approval skips analysis: the plan is already agreed. if (payload.intent === "generate" && payload.plan) { const brief = payload.brief ?? emptyBrief(); - const counts = await drawPlan(model, payload.roomId, ctx.run.id, brief, payload.plan); + const counts = await drawPlan(model, payload.roomId, ctx.run.id, brief, payload.plan, canvas); setStage("done"); return { action: "generated", reply: "", ...counts, decisions: payload.plan.decisions, brief }; } setStage("analyzing"); - const analysis = await analyzeTurn(model, payload); + const analysis = await analyzeTurn(model, payload, canvas); logger.log("Turn analyzed", { decision: analysis.decision, questions: analysis.questions.length }); if (analysis.decision === "ask") { @@ -120,7 +155,7 @@ export const designAgentTask = task({ } if (analysis.decision === "generate" && payload.plan) { - const counts = await drawPlan(model, payload.roomId, ctx.run.id, analysis.brief, payload.plan); + const counts = await drawPlan(model, payload.roomId, ctx.run.id, analysis.brief, payload.plan, canvas); setStage("done"); return { action: "generated", @@ -132,7 +167,7 @@ export const designAgentTask = task({ } setStage("planning"); - const plan = await draftPlan(model, analysis.brief, payload.plan, payload.input); + const plan = await draftPlan(model, analysis.brief, payload.plan, payload.input, canvas); logger.log("Plan drafted", { components: plan.components.length, decisions: plan.decisions.length }); setStage("done"); diff --git a/types/canvas.ts b/types/canvas.ts index bbb6f65..99b76fb 100644 --- a/types/canvas.ts +++ b/types/canvas.ts @@ -38,6 +38,20 @@ export const SHAPE_DEFAULTS: Record = { + rectangle: "Service", + circle: "Event", + diamond: "Decision", + pill: "Queue", + cylinder: "Database", + hexagon: "External", + text: "Note", +}; + // --------------------------------------------------------------------------- // Node color palette // --------------------------------------------------------------------------- @@ -124,6 +138,32 @@ export function resolveNodeFill(color: string | undefined): NodeFill { return NODE_FILLS.find((fill) => fill.id === fillId) ?? NODE_FILLS[0]; } +/** + * What a component does in the system. Its fill follows from the role, so starter templates and + * AI-generated designs colour components the same way: + * entry -> draft blue clients, DNS, CDNs, load balancers, gateways + * compute -> paper services, APIs, workers + * messaging -> scrap coral queues, topics, streams, brokers + * data -> marker amber databases, caches, object storage + * output -> cut sage external systems, consumers, delivery + */ +export const NODE_ROLES = ["entry", "compute", "messaging", "data", "output"] as const; + +export type NodeRole = (typeof NODE_ROLES)[number]; + +export const NODE_ROLE_FILLS: Record = { + entry: "blue", + compute: "paper", + messaging: "coral", + data: "amber", + output: "sage", +}; + +export function getRoleFill(role: NodeRole): NodeFill { + const fillId = NODE_ROLE_FILLS[role]; + return NODE_FILLS.find((fill) => fill.id === fillId) ?? NODE_FILLS[0]; +} + // --------------------------------------------------------------------------- // Node / edge types // ---------------------------------------------------------------------------