diff --git a/components/editor/ai-sidebar.tsx b/components/editor/ai-sidebar.tsx index d194c99..ec73d1c 100644 --- a/components/editor/ai-sidebar.tsx +++ b/components/editor/ai-sidebar.tsx @@ -1,24 +1,21 @@ "use client"; -import { FormEvent, KeyboardEvent, useMemo, useRef, useState } from "react"; -import { Bot, Download, FileText, Send, X } from "lucide-react"; +import { FormEvent, KeyboardEvent, useMemo, useState } from "react"; +import { Bot, Download, FileText, Loader2, Send, X } from "lucide-react"; import { Button } from "@/components/ui/button"; import { Card, CardContent } from "@/components/ui/card"; import { ScrollArea } from "@/components/ui/scroll-area"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { Textarea } from "@/components/ui/textarea"; +import { useDesignAgent } from "@/hooks/use-design-agent"; import { cn } from "@/lib/utils"; interface AiSidebarProps { open: boolean; onClose: () => void; -} - -interface ChatMessage { - id: string; - role: "user" | "assistant"; - text: string; + /** Room the generated design is written into. Also the project access check. */ + projectId: string; } const STARTER_PROMPTS = [ @@ -27,27 +24,18 @@ const STARTER_PROMPTS = [ "Build a CI/CD pipeline", ] as const; -export function AiSidebar({ open, onClose }: AiSidebarProps) { +export function AiSidebar({ open, onClose, projectId }: AiSidebarProps) { const [inputValue, setInputValue] = useState(""); - const [messages, setMessages] = useState([]); - const messageCounterRef = useRef(0); + const { messages, statusText, isRunning, sendPrompt } = useDesignAgent(projectId); const showEmptyState = useMemo(() => messages.length === 0, [messages.length]); const sendMessage = (text: string) => { - const trimmed = text.trim(); - if (!trimmed) { + if (!text.trim() || isRunning) { return; } - const nextMessage: ChatMessage = { - id: `user-${messageCounterRef.current}`, - role: "user", - text: trimmed, - }; - messageCounterRef.current += 1; - - setMessages((prev) => [...prev, nextMessage]); + sendPrompt(text); setInputValue(""); }; @@ -108,7 +96,8 @@ export function AiSidebar({ open, onClose }: AiSidebarProps) { key={prompt} type="button" onClick={() => sendMessage(prompt)} - className="rounded-full bg-muted px-2.5 py-1 text-xs text-muted-foreground transition-colors hover:bg-muted/80 hover:text-foreground" + disabled={isRunning} + className="rounded-full bg-muted px-2.5 py-1 text-xs text-muted-foreground transition-colors hover:bg-muted/80 hover:text-foreground disabled:cursor-not-allowed disabled:opacity-50" > {prompt} @@ -117,19 +106,45 @@ export function AiSidebar({ open, onClose }: AiSidebarProps) { ) : null} - {messages.map((message) => ( -
+ {messages.map((message) => { + if (message.role === "user") { + return ( +
+
+ {message.text} +
+
+ ); + } + + return ( +
+
+ {message.text} +
+
+ ); + })} + + {statusText ? ( +
- {message.text} +
- ))} + ) : null}
@@ -139,12 +154,25 @@ export function AiSidebar({ open, onClose }: AiSidebarProps) { value={inputValue} onChange={(event) => setInputValue(event.target.value)} onKeyDown={handleKeyDown} + disabled={isRunning} className="max-h-36 min-h-24 resize-none bg-background pr-14 text-xs" />
-

Enter to send, Shift+Enter for a new line

-
diff --git a/components/editor/canvas-flow.tsx b/components/editor/canvas-flow.tsx index 097b43c..75d81b7 100644 --- a/components/editor/canvas-flow.tsx +++ b/components/editor/canvas-flow.tsx @@ -19,7 +19,11 @@ import { useUndo, useRedo, useUpdateMyPresence } from "@liveblocks/react"; import { CanvasNodeComponent } from "@/components/editor/canvas-node"; import { CanvasEdgeComponent, CanvasEdgeMarkerDefs } from "@/components/editor/canvas-edge"; -import { CanvasPresenceOverlay, LiveCursors } from "@/components/editor/canvas-presence"; +import { + CanvasPresenceOverlay, + CanvasThinkingIndicator, + LiveCursors, +} from "@/components/editor/canvas-presence"; import { ShapePanel, SHAPE_DRAG_MIME, type ShapeDragPayload } from "@/components/editor/shape-panel"; import { CanvasControlBar } from "@/components/editor/canvas-control-bar"; import { StarterTemplatesModal } from "@/components/editor/starter-templates-modal"; @@ -359,6 +363,7 @@ function CanvasFlowInner({ {/* Floating overlays */} + setIsTemplatesOpen(true)} isSidebarOpen={isSidebarOpen} diff --git a/components/editor/canvas-presence.tsx b/components/editor/canvas-presence.tsx index e1886cb..199febf 100644 --- a/components/editor/canvas-presence.tsx +++ b/components/editor/canvas-presence.tsx @@ -4,6 +4,7 @@ import { useMemo } from "react"; import { useUser, UserButton } from "@clerk/nextjs"; import { useOthers } from "@liveblocks/react/suspense"; import { ViewportPortal } from "@xyflow/react"; +import { Bot } from "lucide-react"; import { Avatar, AvatarFallback, AvatarImage } from "@/components/ui/avatar"; @@ -71,6 +72,37 @@ export function CanvasPresenceOverlay() { ); } +/** + * Shows that a collaborator has a design generation running in this room, + * driven by the `thinking` presence flag the AI sidebar sets. + * + * Rendered top-center so it stays visible when the AI sidebar is open. + */ +export function CanvasThinkingIndicator() { + const others = useOthers(); + + const thinkingCount = useMemo(() => { + return others.filter((other) => other.presence.thinking).length; + }, [others]); + + if (thinkingCount === 0) { + return null; + } + + return ( +
+
+
+
+ ); +} + export function LiveCursors() { const { user } = useUser(); const others = useOthers(); diff --git a/components/editor/canvas-wrapper.tsx b/components/editor/canvas-wrapper.tsx index 8190f19..cac24b7 100644 --- a/components/editor/canvas-wrapper.tsx +++ b/components/editor/canvas-wrapper.tsx @@ -48,6 +48,12 @@ interface CanvasWrapperProps { canAutosave: boolean; onSaveStatusChange?: (status: CanvasSaveStatus) => void; isSidebarOpen: boolean; + /** + * Rendered inside the room, next to the canvas. Overlays that need room + * presence (such as the AI sidebar) belong here rather than as a sibling of + * `CanvasWrapper`, which would place them outside the `RoomProvider`. + */ + children?: ReactNode; } /** @@ -58,12 +64,14 @@ interface CanvasWrapperProps { * - Initial presence includes cursor: null (no active cursor on join) * - ClientSideSuspense defers rendering until the room is ready * - CanvasErrorBoundary catches Liveblocks connection failures + * - `children` render inside the room so they can read and write presence */ export function CanvasWrapper({ roomId, canAutosave, onSaveStatusChange, isSidebarOpen, + children, }: CanvasWrapperProps) { return ( @@ -87,6 +95,7 @@ export function CanvasWrapper({ /> + {children} ); diff --git a/components/editor/editor-workspace-shell.tsx b/components/editor/editor-workspace-shell.tsx index 632a70d..5022b19 100644 --- a/components/editor/editor-workspace-shell.tsx +++ b/components/editor/editor-workspace-shell.tsx @@ -112,8 +112,13 @@ export function EditorWorkspaceShell({ canAutosave={isOwner} onSaveStatusChange={setSaveStatus} isSidebarOpen={isSidebarOpen} - /> - setIsAiSidebarOpen(false)} /> + > + setIsAiSidebarOpen(false)} + projectId={projectId} + /> + diff --git a/context/fix/01-AI-sidebar-fix.md b/context/fix/01-AI-sidebar-fix.md new file mode 100644 index 0000000..f12c111 --- /dev/null +++ b/context/fix/01-AI-sidebar-fix.md @@ -0,0 +1,156 @@ +Wire the AI sidebar to the design agent so a prompt typed in the sidebar produces nodes and +edges on the canvas. + +The sidebar is finished UI with no backend connection, and the design task is a stub. The +backend contract between them was built in `context/feature-specs/22-design-agent-api.md` and +is correct — this unit fills in the two ends and leaves that contract untouched. + +## Context + +What is broken: + +- `components/editor/ai-sidebar.tsx` is local state only. `sendMessage` pushes a `ChatMessage` + into `useState` and clears the input. Every message it can construct is `role: "user"`; there + is no code path that produces an assistant reply, and nothing calls `/api/ai/design`. +- `src/trigger/design-agent.ts` logs its payload and echoes it back. No AI call, no canvas write. +- `@trigger.dev/react-hooks` is installed and imported nowhere. +- `ai` and `@ai-sdk/google` are installed and imported nowhere. + +What already exists and must be reused rather than rebuilt: + +- `POST /api/ai/design` — Clerk auth, `getAccessibleProject` check, triggers the task, writes a + `TaskRun` row, returns `202 { runId }`. +- `POST /api/ai/design/token` — looks up `TaskRun` by `(runId, userId)` and mints a Trigger + public token scoped to `read: { runs: runId }`. This is the ownership gate; without the + `TaskRun` row any signed-in user could request a token for any run. +- `getLiveblocksClient()` in `lib/liveblocks.ts` — cached `@liveblocks/node` client. +- `CanvasNode`, `CanvasEdge`, `CANVAS_NODE_TYPE` (`"canvasNode"`), `CANVAS_EDGE_TYPE` + (`"canvasEdge"`), `SHAPE_DEFAULTS`, and `NODE_COLOR_PALETTE` in `types/canvas.ts`. +- `projectId` is already in scope in `components/editor/editor-workspace-shell.tsx`; it is + simply not passed to `AiSidebar` at the render site. +- The `thinking: boolean` field in the `Presence` type (`liveblocks.config.ts`) is initialised + in `canvas-wrapper.tsx` and never set anywhere. It is reserved for this feature. + +## Prerequisites + +1. `zod` is **not installed**. The `ai` package declares it as a peer dependency + (`^3.25.76 || ^4.1.8`) and `generateObject` needs it. Install it pinned to an exact version. + +2. `GOOGLE_GENERATIVE_AI_API_KEY` is present in `.env.example` but **not set in `.env.local`**. + Nothing will run without it. Add it, and document it in the same style as the existing + entries. It is server-only: it must never take a `NEXT_PUBLIC_` prefix. + +## Implementation + +1. Define the generation schema. + + Create a Zod schema describing what the model must return: a list of nodes and a list of + edges. Put it where both the task and any future caller can import it. + + Constrain it to what the canvas can actually render: + - node: `id`, `label`, `position: { x, y }`, and a `shape` drawn from `CanvasShape` + - edge: `id`, `source`, `target`, optional `label`, and an `arrowDirection` from + `CanvasEdgeData` + - do not let the model invent colors as free strings; if colors are generated at all, they + must come from `NODE_COLOR_PALETTE` ids + + Map the model output into `CanvasNode` / `CanvasEdge` in code, not in the prompt. The node + `type` must be `CANVAS_NODE_TYPE` and the edge `type` must be `CANVAS_EDGE_TYPE`, and sizes + should come from `SHAPE_DEFAULTS` rather than from the model. + +2. Implement the design task in `src/trigger/design-agent.ts`. + + Keep the existing task id (`design-agent`) and payload (`prompt`, `roomId`) — the route + triggers it by that id and the `TaskRun` row is already keyed to it. + + The task should: + - call `generateObject` with `@ai-sdk/google` and the schema from step 1 + - lay the nodes out with non-overlapping positions before writing them + - write them into the Liveblocks room with `mutateFlow` from `@liveblocks/react-flow/node`, + passing the client from `getLiveblocksClient()`: + + ```ts + await mutateFlow({ client, roomId }, (flow) => { + flow.addNodes(nodes); + flow.addEdges(edges); + }); + ``` + + - report progress with the `metadata` API from `@trigger.dev/sdk/v3` so the sidebar can show + stages rather than a spinner + + Do not override `storageKey`. The client calls `useLiveblocksFlow` in + `components/editor/canvas-flow.tsx` without one, so the server side must use the default or + the two will write to different stores. + +3. Pass `projectId` into the sidebar. + + `AiSidebar` currently takes only `open` and `onClose`. Add `projectId` and pass it from + `editor-workspace-shell.tsx`, where it is already available. + +4. Call the API from `sendMessage`. + + POST to `/api/ai/design` with `{ prompt, roomId: projectId, projectId }`. The route rejects + the request unless `roomId === projectId`, so send the same value for both. + + Handle the states the route actually returns: `202` with a `runId`, `400` for a malformed + body, `401` unauthenticated, `403` when the user cannot access the project. Keep the input + disabled while a run is in flight. + +5. Subscribe to the run. + + POST the `runId` to `/api/ai/design/token`, then pass the returned token and the run id to + `useRealtimeRun` from `@trigger.dev/react-hooks`. + + Render run status as an assistant-side message. The message list currently has no `assistant` + branch, so add one. Surface failures — a failed run must not leave the sidebar sitting on a + spinner. + + Set the `thinking` presence flag while a run is active so collaborators in the room see that + a generation is in progress, and clear it when the run settles. + +## Constraints That Must Not Be Broken + +1. **The generated canvas must arrive through Liveblocks, not through the HTTP response.** + The task writes into the room; every viewer sees the result, and the existing autosave in + `hooks/use-canvas-autosave.ts` persists it. Returning nodes from the route and applying them + locally would desync every other collaborator. + +2. **The token route is an authorization boundary.** Do not relax the `TaskRun` lookup, widen + the token scope beyond the single run, or mint tokens anywhere else. + +3. **Secrets stay on the server.** The Google API key and the Liveblocks secret are used only + inside the task and route handlers. The browser gets the run-scoped public token and nothing + else. + +4. **Do not write canvas JSON to blob storage from the task.** Persistence is the canvas + autosave path's job. Two writers to the same artifact will race. + +## Scope Limits + +- Do not implement the Specs tab. "Generate Spec" and its download stay inert; that is a + separate unit. +- Do not change `POST /api/ai/design` or `POST /api/ai/design/token` beyond what steps 4 and 5 + consume. They are already correct. +- Do not change the `TaskRun` model, `lib/project-access.ts`, or `lib/prisma.ts`. +- Do not change the Liveblocks auth route or the room/project id convention. +- Do not add streaming token-by-token chat. Structured output plus run status is the target. +- Do not add a second AI provider or an abstraction layer over providers. + +## Check When Done + +- Typing a prompt in the sidebar creates a run and the nodes appear on the canvas without a + page reload. +- A second browser signed in as a collaborator sees the same nodes appear, and sees the + `thinking` indicator while the run is active. +- The generated nodes are editable, movable, and survive a reload (the autosave path picked + them up). +- A prompt sent for a project the user cannot access returns 403 and is surfaced in the UI. +- A failed run shows an error in the sidebar rather than an indefinite loading state. +- `GOOGLE_GENERATIVE_AI_API_KEY` is set locally and carries no `NEXT_PUBLIC_` prefix. +- `pnpm lint`, `pnpm typecheck`, and `pnpm build` pass. + +## After Done + +Update `context/progress-tracker.md`. If the shape of the AI flow ends up differing from what +`context/architecture-context.md` describes, update that file too — otherwise leave it alone. diff --git a/context/progress-tracker.md b/context/progress-tracker.md index 692034f..db4ef30 100644 --- a/context/progress-tracker.md +++ b/context/progress-tracker.md @@ -8,10 +8,66 @@ Update this file whenever the current phase, active feature, or implementation s ## Current Goal -- Select and implement the next available feature spec unit after `22-design-agent-api.md`. +- Implement the Specs tab (Generate Spec + download), which remains inert. ## Completed +- Fix unit `context/fix/01-AI-sidebar-fix.md` completed (AI sidebar wired to the design agent): + - Installed `zod` pinned to an exact version (`4.6.2`) for `generateObject`. + - Added shared generation contract in `lib/design-generation.ts`: + - Zod schema for model output (`designGraphSchema`): nodes (`id`, `label`, + `position`, `shape` from `CANVAS_SHAPES`, optional `colorId` from + `NODE_COLOR_IDS`) and edges (`id`, `source`, `target`, optional `label`, + `arrowDirection` from `EDGE_ARROW_DIRECTIONS`) + - `buildCanvasGraph()` maps model output to `CanvasNode` / `CanvasEdge` in + code: applies `CANVAS_NODE_TYPE` / `CANVAS_EDGE_TYPE`, sizes from + `SHAPE_DEFAULTS`, palette colors from `NODE_COLOR_PALETTE`, re-keys ids + under a run-scoped prefix, drops dangling edges, and collapses model + position hints onto a non-overlapping grid + - Run progress contract shared by task and sidebar + (`DESIGN_AGENT_STAGE_KEY`, `DESIGN_AGENT_STAGES`, `parseDesignAgentStage`) + - Added value-list exports to `types/canvas.ts` so shapes, arrow directions, + and palette ids can be validated at runtime: `CANVAS_SHAPES`, + `EDGE_ARROW_DIRECTIONS`, `NODE_COLOR_IDS` / `NodeColorId`, + `CanvasArrowDirection` (existing types are now derived from them). + - Implemented `src/trigger/design-agent.ts` (task id and payload unchanged): + - Calls `generateObject` with `@ai-sdk/google` + (`GOOGLE_GENERATIVE_AI_MODEL`, default `gemini-3.5-flash`) + - Writes nodes and edges into the Liveblocks room with `mutateFlow` from + `@liveblocks/react-flow/node` using `getLiveblocksClient()`; default + `storageKey` (matches the client `useLiveblocksFlow`) + - Offsets each generation below existing room content + - Publishes `stage` / `nodeCount` / `edgeCount` on run metadata and returns + `{ nodeCount, edgeCount }` + - Added `hooks/use-design-agent.ts`: + - `POST /api/ai/design` with `{ prompt, roomId, projectId }`, handling + `202` / `400` / `401` / `403` and network failures + - Exchanges the run id for a run-scoped token via + `POST /api/ai/design/token`, then subscribes with `useRealtimeRun` + - Derives everything it shows (status line, closing summary, whether the + composer is locked) from the run's own reported state, so a settlement + step cannot leave the sidebar stuck mid-run + - Renders stage-based status, posts a summary naming the component and + connection counts on success, reports failures, and keeps the composer + disabled while a run is in flight + - Sets the Liveblocks `thinking` presence flag while a run is active + - Wired the sidebar and workspace: + - `AiSidebar` now takes `projectId` and renders assistant-side messages, + a live status line, and error messages + - `CanvasWrapper` accepts `children` rendered inside `RoomProvider`, and + `EditorWorkspaceShell` nests `AiSidebar` there so it can write presence + - Added `CanvasThinkingIndicator` in `components/editor/canvas-presence.tsx` + so collaborators see an in-progress generation + - Added `GOOGLE_GENERATIVE_AI_MODEL` to `.env.example` + (`GOOGLE_GENERATIVE_AI_API_KEY` is set locally, server-only, no + `NEXT_PUBLIC_` prefix). + - Out of scope and unchanged: the Specs tab, both AI routes, `TaskRun`, + `lib/project-access.ts`, `lib/prisma.ts`, the Liveblocks auth route, and + the canvas autosave persistence path. + - Validation checks: + - `pnpm lint`, `pnpm typecheck`, and `pnpm build` passed + - `trigger deploy --dry-run` built the task successfully + - Feature spec `22-design-agent-api.md` completed: - Added `TaskRun` Prisma model in `prisma/models/task-run.prisma` with: - `runId` unique @@ -597,3 +653,22 @@ Update this file whenever the current phase, active feature, or implementation s - Replaced SVG text labels with `foreignObject`-based wrapped label containers for diamond/hexagon/cylinder so long labels stack and clip within node bounds. - Updated CSS-shape label style to multiline wrapping with bounded height and hidden overflow instead of single-line ellipsis. - Reworked cylinder renderer into a stacked database-style cylinder (top, middle, and bottom elliptical bands). +- Rolled back the Prisma-to-Supabase data-layer migration on 2026-09-11: + - Two migration steps had been implemented and were reverted in full: the snake_case schema rename with database-side + defaults, and the addition of the Supabase client (`@supabase/supabase-js`, the `supabase` CLI, `lib/supabase.ts`, + `types/database.types.ts`, the Node 20 -> 22 bump, and the `SUPABASE_*` env vars). + - Decision: Prisma stays the data layer. The database remains Supabase Postgres — that predates this work and is + unchanged. Auth stays Clerk, real-time stays Liveblocks, blob storage stays Vercel Blob, so Supabase is used only + as a Postgres host. + - The database schema was rolled back to its original shape: `Project`, `ProjectCollaborator`, `TaskRun` with + camelCase columns, the `ProjectStatus` enum, no database-side `id` or `updated_at` defaults, and the + `moddatetime` trigger and extension removed. All three tables were empty, so no data was involved. + - The `20260911120000_snake_case_schema_and_db_defaults` migration was removed from `prisma/migrations` and its row + deleted from `_prisma_migrations`; the two original migrations are the full history again. + - Also fixed a pre-existing lockfile mismatch on `main`: `pnpm-lock.yaml` recorded `@trigger.dev/react-hooks` as + `^4.4.6` while `package.json` pins `4.4.6`, which made `pnpm install --frozen-lockfile` (what CI runs) fail. + - Validation checks: + - `prisma migrate status` reports 2 migrations and an up-to-date schema; `prisma migrate diff` reports no drift + - Prisma create/update/delete verified against the restored schema: slug ids, client-side `cuid()` generation and + `@updatedAt` all behave as before; test rows deleted + - `pnpm lint`, `pnpm typecheck`, and `pnpm build` passed diff --git a/hooks/use-design-agent.ts b/hooks/use-design-agent.ts new file mode 100644 index 0000000..48d49dc --- /dev/null +++ b/hooks/use-design-agent.ts @@ -0,0 +1,250 @@ +"use client"; + +import { useCallback, useEffect, useMemo, useState } from "react"; +import { useUpdateMyPresence } from "@liveblocks/react"; +import { useRealtimeRun } from "@trigger.dev/react-hooks"; + +import { + DESIGN_AGENT_STAGE_KEY, + parseDesignAgentStage, + type DesignAgentStage, +} from "@/lib/design-generation"; +import type { designAgentTask } from "@/src/trigger/design-agent"; + +export interface DesignChatMessage { + id: string; + role: "user" | "assistant"; + text: string; + /** Assistant messages reporting a failure render in the error style. */ + isError?: boolean; +} + +export interface UseDesignAgentResult { + messages: DesignChatMessage[]; + /** Live status line for the in-flight run, rendered assistant-side. */ + statusText: string | null; + /** True from submit until the run settles. The composer stays disabled meanwhile. */ + isRunning: boolean; + sendPrompt: (text: string) => void; +} + +type RealtimeDesignRun = NonNullable< + ReturnType>["run"] +>; + +interface RunOutcome { + text: string; + isError: boolean; +} + +const STAGE_LABELS: Record = { + generating: "Designing the architecture…", + writing: "Adding components to the canvas…", + done: "Finishing up…", +}; + +const GENERIC_ERROR = "Something went wrong generating the design. Try again."; + +function describeRequestFailure(status: number): string { + switch (status) { + case 400: + return "That prompt could not be sent. Try rephrasing it."; + case 401: + return "Your session expired. Sign in again to keep designing."; + case 403: + return "You do not have access to this project, so the design was not generated."; + default: + return GENERIC_ERROR; + } +} + +function summarize(run: RealtimeDesignRun): string { + const nodeCount = run.output?.nodeCount ?? 0; + const edgeCount = run.output?.edgeCount ?? 0; + + if (nodeCount === 0) { + return "The design run finished but produced no components. Try a more specific prompt."; + } + + return ( + `Added ${nodeCount} component${nodeCount === 1 ? "" : "s"} and ` + + `${edgeCount} connection${edgeCount === 1 ? "" : "s"} to the canvas.` + ); +} + +/** + * Drives one design run at a time: posts the prompt to `/api/ai/design`, + * exchanges the returned run id for a run-scoped public token, subscribes to + * the run over Trigger Realtime, and mirrors the run's liveness onto the + * Liveblocks `thinking` presence flag. + * + * The run's own reported state is the single source of truth. Everything the + * sidebar shows — the status line, the closing summary, whether the composer + * is locked — is derived from it, so the UI cannot be left mid-run by a + * settlement step that failed to fire. + * + * Must be used inside a `RoomProvider` — it writes presence for the room. + */ +export function useDesignAgent(projectId: string): UseDesignAgentResult { + const [history, setHistory] = useState([]); + const [isSubmitting, setIsSubmitting] = useState(false); + const [runId, setRunId] = useState(null); + const [accessToken, setAccessToken] = useState(null); + + const updateMyPresence = useUpdateMyPresence(); + + const { run, error } = useRealtimeRun(runId ?? undefined, { + accessToken: accessToken ?? undefined, + enabled: Boolean(runId && accessToken), + }); + + // A freshly started subscription can still be reporting the previous run, so + // only trust run data whose id matches the run currently being awaited. + const activeRun = run && run.id === runId ? run : undefined; + + /** The held run's result, or null while it is still in flight. */ + const outcome: RunOutcome | null = useMemo(() => { + if (!runId) { + return null; + } + + if (activeRun?.isSuccess) { + return { text: summarize(activeRun), isError: false }; + } + if (activeRun?.isCancelled) { + return { text: "The design run was cancelled.", isError: true }; + } + if (activeRun?.isFailed) { + return { text: activeRun.error?.message ?? GENERIC_ERROR, isError: true }; + } + if (error) { + return { + text: "Lost track of the design run. Reload to see whether it finished.", + isError: true, + }; + } + + return null; + }, [activeRun, error, runId]); + + const outcomeMessage: DesignChatMessage | null = useMemo(() => { + if (!runId || !outcome) { + return null; + } + return { + id: `assistant-${runId}`, + role: "assistant", + text: outcome.text, + isError: outcome.isError, + }; + }, [outcome, runId]); + + // A run is over the moment it reports an outcome; nothing has to be cleared + // for the composer to unlock. + const isRunning = isSubmitting || (runId !== null && outcome === null); + + const messages = useMemo( + () => (outcomeMessage ? [...history, outcomeMessage] : history), + [history, outcomeMessage], + ); + + const appendToHistory = useCallback((...entries: DesignChatMessage[]) => { + setHistory((previous) => [...previous, ...entries]); + }, []); + + const sendPrompt = useCallback( + (text: string) => { + const prompt = text.trim(); + if (!prompt || isRunning) { + return; + } + + // The finished run is about to be replaced, so commit its derived + // summary into the transcript before letting go of it. + const userMessage: DesignChatMessage = { + id: `user-${history.length}`, + role: "user", + text: prompt, + }; + appendToHistory(...(outcomeMessage ? [outcomeMessage, userMessage] : [userMessage])); + + setRunId(null); + setAccessToken(null); + setIsSubmitting(true); + + void (async () => { + try { + const response = await fetch("/api/ai/design", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ prompt, roomId: projectId, projectId }), + }); + + if (response.status !== 202) { + appendToHistory({ + id: `assistant-request-${Date.now()}`, + role: "assistant", + text: describeRequestFailure(response.status), + isError: true, + }); + setIsSubmitting(false); + return; + } + + const { runId: startedRunId } = (await response.json()) as { runId: string }; + + const tokenResponse = await fetch("/api/ai/design/token", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ runId: startedRunId }), + }); + + if (!tokenResponse.ok) { + appendToHistory({ + id: `assistant-token-${Date.now()}`, + role: "assistant", + text: "The design run started but its progress could not be tracked. Reload to see the result.", + isError: true, + }); + setIsSubmitting(false); + return; + } + + const { token } = (await tokenResponse.json()) as { token: string }; + + setRunId(startedRunId); + setAccessToken(token); + setIsSubmitting(false); + } catch { + appendToHistory({ + id: `assistant-error-${Date.now()}`, + role: "assistant", + text: GENERIC_ERROR, + isError: true, + }); + setIsSubmitting(false); + } + })(); + }, + [appendToHistory, history.length, isRunning, outcomeMessage, projectId], + ); + + // Let collaborators in the room see that a generation is in progress. + useEffect(() => { + updateMyPresence({ thinking: isRunning }); + }, [isRunning, updateMyPresence]); + + useEffect(() => { + return () => { + updateMyPresence({ thinking: false }); + }; + }, [updateMyPresence]); + + let statusText: string | null = null; + if (isRunning) { + const stage = activeRun ? parseDesignAgentStage(activeRun.metadata?.[DESIGN_AGENT_STAGE_KEY]) : null; + statusText = stage ? STAGE_LABELS[stage] : "Starting the design run…"; + } + + return { messages, statusText, isRunning, sendPrompt }; +} diff --git a/lib/design-generation.ts b/lib/design-generation.ts new file mode 100644 index 0000000..0ebcb0d --- /dev/null +++ b/lib/design-generation.ts @@ -0,0 +1,268 @@ +import { z } from "zod"; + +import { + CANVAS_EDGE_TYPE, + CANVAS_NODE_TYPE, + CANVAS_SHAPES, + EDGE_ARROW_DIRECTIONS, + NODE_COLOR_IDS, + NODE_COLOR_PALETTE, + 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. +// --------------------------------------------------------------------------- + +// --------------------------------------------------------------------------- +// Run progress contract +// +// The task publishes its current stage on run metadata; the AI sidebar reads +// it back over Realtime. Both ends import these constants so the key and the +// stage names cannot drift apart. +// --------------------------------------------------------------------------- + +/** Metadata key carrying the current {@link DesignAgentStage}. */ +export const DESIGN_AGENT_STAGE_KEY = "stage"; + +/** Progress stages a design run moves through, in order. */ +export const DESIGN_AGENT_STAGES = ["generating", "writing", "done"] as const; + +export type DesignAgentStage = (typeof DESIGN_AGENT_STAGES)[number]; + +/** Narrows an unknown metadata value to a known stage. */ +export function parseDesignAgentStage(value: unknown): DesignAgentStage | null { + return DESIGN_AGENT_STAGES.find((stage) => stage === value) ?? null; +} + +/** Upper bound on generated components, to keep a single run legible. */ +const MAX_NODES = 24; + +/** Upper bound on generated connections. */ +const MAX_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'."), + label: z + .string() + .min(1) + .describe("Human readable component name shown on the node, e.g. 'API Gateway'."), + position: z + .object({ + x: z.number().describe("Horizontal hint. Sibling components sit side by side."), + y: z.number().describe("Vertical hint. Requests flow from low y to high y."), + }) + .describe("Rough layout hint only. Exact pixel placement is resolved by the app."), + 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.", + ), + colorId: z + .enum(NODE_COLOR_IDS) + .optional() + .describe("Palette entry used to tint the node. Give related components the same entry."), +}); + +/** 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."), + 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."), +}); + +/** The full structured output requested from the model. */ +export const designGraphSchema = z.object({ + nodes: z.array(designNodeSchema).min(1).max(MAX_NODES), + edges: z.array(designEdgeSchema).max(MAX_EDGES), +}); + +export type DesignGraph = z.infer; + +// --------------------------------------------------------------------------- +// Layout +// --------------------------------------------------------------------------- + +/** Horizontal distance between two adjacent layout columns. */ +const COLUMN_WIDTH = 240; + +/** Vertical distance between two adjacent layout rows. */ +const ROW_HEIGHT = 170; + +/** + * Vertical distance within which two model-supplied `y` values are treated as + * belonging to the same row. + */ +const ROW_TOLERANCE = 60; + +interface LayoutOrigin { + x: number; + y: number; +} + +export interface BuildCanvasGraphOptions { + /** + * Prefix applied to every generated id, so a run can never collide with + * content already in the room. + */ + idPrefix: string; + /** Top-left corner the generated graph is laid out from. */ + origin?: LayoutOrigin; +} + +function resolveColor(colorId: NodeColorId | undefined) { + return NODE_COLOR_PALETTE.find((entry) => entry.id === colorId) ?? NODE_COLOR_PALETTE[0]; +} + +/** + * Groups nodes into rows by their model-supplied `y`, then orders each row by + * `x`. The model's coordinates are only a hint; collapsing them onto a fixed + * grid means nodes cannot overlap regardless of what the model returned. + */ +function assignGridCells(nodes: DesignGraph["nodes"]): Map { + const ordered = [...nodes].sort( + (a, b) => a.position.y - b.position.y || a.position.x - b.position.x, + ); + + const rows: DesignGraph["nodes"][] = []; + let currentRow: DesignGraph["nodes"] = []; + let rowAnchorY: number | null = null; + + for (const node of ordered) { + if (rowAnchorY === null) { + rowAnchorY = node.position.y; + currentRow.push(node); + continue; + } + + if (Math.abs(node.position.y - rowAnchorY) <= ROW_TOLERANCE) { + currentRow.push(node); + continue; + } + + rows.push(currentRow); + currentRow = [node]; + rowAnchorY = node.position.y; + } + + if (currentRow.length > 0) { + rows.push(currentRow); + } + + const cells = new Map(); + rows.forEach((row, rowIndex) => { + [...row] + .sort((a, b) => a.position.x - b.position.x) + .forEach((node, columnIndex) => { + cells.set(node.id, { column: columnIndex, row: rowIndex }); + }); + }); + + return cells; +} + +// --------------------------------------------------------------------------- +// Model output -> canvas graph +// --------------------------------------------------------------------------- + +/** + * Maps validated model output onto fully-formed canvas nodes and edges. + * + * Nodes and edges are re-keyed under `idPrefix`, laid out on a non-overlapping + * grid, and given the canvas node/edge types plus `SHAPE_DEFAULTS` sizing. + * Edges pointing at nodes the model did not define are dropped. + */ +export function buildCanvasGraph( + design: DesignGraph, + { idPrefix, origin = { x: 0, y: 0 } }: BuildCanvasGraphOptions, +): { nodes: CanvasNode[]; edges: CanvasEdge[] } { + const seenNodeIds = new Set(); + const uniqueNodes = design.nodes.filter((node) => { + if (seenNodeIds.has(node.id)) { + return false; + } + seenNodeIds.add(node.id); + return true; + }); + + const cells = assignGridCells(uniqueNodes); + const canvasNodeIds = new Map(); + + const nodes: CanvasNode[] = uniqueNodes.map((node) => { + const canvasNodeId = `${idPrefix}-${node.id}`; + 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 }; + + return { + id: canvasNodeId, + type: CANVAS_NODE_TYPE, + position: { + x: origin.x + cell.column * COLUMN_WIDTH + (COLUMN_WIDTH - dimensions.width) / 2, + y: origin.y + cell.row * ROW_HEIGHT + (ROW_HEIGHT - dimensions.height) / 2, + }, + style: { + width: dimensions.width, + height: dimensions.height, + }, + data: { + label: node.label, + shape: node.shape, + color: color.bg, + textColor: "var(--text-primary)", + strokeColor: "var(--text-primary)", + }, + }; + }); + + const seenEdgeIds = new Set(); + const edges: CanvasEdge[] = []; + + for (const edge of design.edges) { + const source = canvasNodeIds.get(edge.source); + const target = canvasNodeIds.get(edge.target); + if (!source || !target || source === target) { + continue; + } + + const canvasEdgeId = `${idPrefix}-${edge.id}`; + if (seenEdgeIds.has(canvasEdgeId)) { + continue; + } + seenEdgeIds.add(canvasEdgeId); + + const label = edge.label?.trim(); + + edges.push({ + id: canvasEdgeId, + type: CANVAS_EDGE_TYPE, + source, + target, + data: { + arrowDirection: edge.arrowDirection, + ...(label ? { label } : {}), + }, + }); + } + + return { nodes, edges }; +} diff --git a/prisma.config.ts b/prisma.config.ts index f7890be..74a9da8 100644 --- a/prisma.config.ts +++ b/prisma.config.ts @@ -1,8 +1,11 @@ // This file was generated by Prisma, and assumes you have installed the following: // npm install --save-dev prisma dotenv -import "dotenv/config"; +import { config } from "dotenv"; import { defineConfig } from "prisma/config"; +// Next.js loads .env.local automatically; the Prisma CLI does not, so load it here. +config({ path: [".env.local", ".env"] }); + export default defineConfig({ schema: "prisma/", migrations: { diff --git a/src/trigger/design-agent.ts b/src/trigger/design-agent.ts index 0f878ef..2c46116 100644 --- a/src/trigger/design-agent.ts +++ b/src/trigger/design-agent.ts @@ -1,17 +1,123 @@ -import { logger, task } from "@trigger.dev/sdk/v3"; +import { google } from "@ai-sdk/google"; +import { mutateFlow } from "@liveblocks/react-flow/node"; +import { logger, metadata, task } from "@trigger.dev/sdk/v3"; +import { generateObject } from "ai"; + +import { + buildCanvasGraph, + designGraphSchema, + DESIGN_AGENT_STAGE_KEY, + type DesignAgentStage, +} from "@/lib/design-generation"; +import { getLiveblocksClient } from "@/lib/liveblocks"; +import { SHAPE_DEFAULTS, type CanvasEdge, type CanvasNode, type CanvasShape } from "@/types/canvas"; export interface DesignAgentPayload { prompt: string; roomId: string; } +export interface DesignAgentResult { + nodeCount: number; + edgeCount: number; +} + +/** Vertical gap left between existing canvas content and a newly generated graph. */ +const EXISTING_CONTENT_GAP = 140; + +const DEFAULT_MODEL_ID = "gemini-3.5-flash"; + +const SYSTEM_PROMPT = [ + "You are a system design architect. Turn the user's description into a component diagram.", + "Return only components that belong on an architecture diagram: clients, gateways, services,", + "queues, caches, datastores, and external systems. Give every component a short, concrete", + "label. Connect components in the direction traffic actually flows, and 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.", + "Prefer a focused diagram of the components that matter over an exhaustive one.", +].join(" "); + +function resolveModelId(): string { + const configured = process.env.GOOGLE_GENERATIVE_AI_MODEL?.trim(); + return configured && configured.length > 0 ? configured : DEFAULT_MODEL_ID; +} + +function getNodeHeight(node: CanvasNode): number { + const styledHeight = node.style?.height; + if (typeof styledHeight === "number") { + return styledHeight; + } + + const shape = node.data.shape as CanvasShape | undefined; + return shape ? SHAPE_DEFAULTS[shape].height : SHAPE_DEFAULTS.rectangle.height; +} + +/** + * Places a generated graph below anything already on the canvas, so a second + * generation never lands on top of the first. + */ +function resolveOrigin(existingNodes: readonly CanvasNode[]): { x: number; y: number } { + if (existingNodes.length === 0) { + return { x: 0, y: 0 }; + } + + let left = Number.POSITIVE_INFINITY; + let bottom = Number.NEGATIVE_INFINITY; + + for (const node of existingNodes) { + left = Math.min(left, node.position.x); + bottom = Math.max(bottom, node.position.y + getNodeHeight(node)); + } + + return { x: left, y: bottom + EXISTING_CONTENT_GAP }; +} + export const designAgentTask = task({ id: "design-agent", - run: async (payload: DesignAgentPayload) => { - logger.log("Design agent task triggered", { payload }); + run: async (payload: DesignAgentPayload, { ctx }): Promise => { + logger.log("Design agent task triggered", { roomId: payload.roomId }); + + metadata.set(DESIGN_AGENT_STAGE_KEY, "generating" satisfies DesignAgentStage); + + const { object: design } = await generateObject({ + model: google(resolveModelId()), + schema: designGraphSchema, + system: SYSTEM_PROMPT, + prompt: payload.prompt, + }); + + logger.log("Design generated", { + nodeCount: design.nodes.length, + edgeCount: design.edges.length, + }); + + metadata.set(DESIGN_AGENT_STAGE_KEY, "writing" satisfies DesignAgentStage); + + const client = getLiveblocksClient(); + let result: DesignAgentResult = { nodeCount: 0, edgeCount: 0 }; + + await mutateFlow( + { client, roomId: payload.roomId }, + (flow) => { + const { nodes, edges } = buildCanvasGraph(design, { + idPrefix: ctx.run.id, + origin: resolveOrigin(flow.nodes), + }); + + flow.addNodes(nodes); + flow.addEdges(edges); + + result = { nodeCount: nodes.length, edgeCount: edges.length }; + }, + ); + + metadata.set(DESIGN_AGENT_STAGE_KEY, "done" satisfies DesignAgentStage); + metadata.set("nodeCount", result.nodeCount); + metadata.set("edgeCount", result.edgeCount); + + logger.log("Design written to room", { roomId: payload.roomId, ...result }); - return { - received: payload, - }; + return result; }, }); diff --git a/types/canvas.ts b/types/canvas.ts index bf329f2..989d2dd 100644 --- a/types/canvas.ts +++ b/types/canvas.ts @@ -4,14 +4,18 @@ import type { Node, Edge } from "@xyflow/react"; // Shape catalogue // --------------------------------------------------------------------------- +/** All supported draggable shapes, as a value list usable for runtime validation. */ +export const CANVAS_SHAPES = [ + "rectangle", + "circle", + "diamond", + "pill", + "cylinder", + "hexagon", +] as const; + /** All supported draggable shapes. */ -export type CanvasShape = - | "rectangle" - | "circle" - | "diamond" - | "pill" - | "cylinder" - | "hexagon"; +export type CanvasShape = (typeof CANVAS_SHAPES)[number]; /** Default width / height for each shape (pixels). */ export const SHAPE_DEFAULTS: Record = { @@ -27,9 +31,24 @@ export const SHAPE_DEFAULTS: Record { shape?: CanvasShape; } +/** Supported edge arrowhead directions, as a value list usable for runtime validation. */ +export const EDGE_ARROW_DIRECTIONS = ["none", "forward", "backward", "bidirectional"] as const; + +/** Direction of arrowheads rendered on a canvas edge. */ +export type CanvasArrowDirection = (typeof EDGE_ARROW_DIRECTIONS)[number]; + /** * Data payload carried by every canvas edge. */ @@ -93,7 +118,7 @@ export interface CanvasEdgeData extends Record { * 'backward' — arrowhead at source end only * 'bidirectional' — arrowheads at both ends */ - arrowDirection?: "none" | "forward" | "backward" | "bidirectional"; + arrowDirection?: CanvasArrowDirection; /** * Custom stroke color (hex). When defined uses pair.text from NODE_COLOR_PALETTE * for high visibility. Undefined = default zinc gray.