diff --git a/app/api/projects/[projectId]/ai-sessions/[sessionId]/turns/route.ts b/app/api/projects/[projectId]/ai-sessions/[sessionId]/turns/route.ts index 218a8b2..2afc175 100644 --- a/app/api/projects/[projectId]/ai-sessions/[sessionId]/turns/route.ts +++ b/app/api/projects/[projectId]/ai-sessions/[sessionId]/turns/route.ts @@ -1,35 +1,68 @@ import { auth } from "@clerk/nextjs/server"; +import { MAX_QUESTIONS_PER_ROUND } from "@/lib/ai/agent-schema"; import { MAX_MESSAGE_LENGTH } from "@/lib/ai/session-limits"; -import { startTurn } from "@/lib/ai/session-turns"; +import { startTurn, type TurnInput } from "@/lib/ai/session-turns"; import { getAccessibleProject, getCurrentIdentity } from "@/lib/project-access"; type TurnRouteContext = { params: Promise<{ projectId: string; sessionId: string }> }; -/** Only `{ type: "message", text }` for now; answers/generate/skip arrive with the turn engine. */ -function parseMessageText(value: unknown): string | null { - if (typeof value !== "object" || value === null) { - return null; - } +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null; +} - const candidate = value as Record; - if (candidate.type !== "message" || typeof candidate.text !== "string") { +/** + * Accepts one of: + * - `{ type: "message", text }` — free text, 1 to MAX_MESSAGE_LENGTH characters + * - `{ type: "answers", answers: [{ questionId, answer }] }` — replies to the open questions + * - `{ type: "generate" }` — draw the latest plan + * - `{ type: "skip" }` — stop asking questions and plan with assumptions + */ +function parseTurnInput(value: unknown): TurnInput | null { + if (!isRecord(value)) { return null; } - const text = candidate.text.trim(); - if (text.length === 0 || text.length > MAX_MESSAGE_LENGTH) { - return null; - } + switch (value.type) { + case "message": { + if (typeof value.text !== "string") { + return null; + } + const text = value.text.trim(); + return text.length > 0 && text.length <= MAX_MESSAGE_LENGTH ? { type: "message", text } : null; + } + + case "answers": { + if (!Array.isArray(value.answers) || value.answers.length === 0) { + return null; + } + if (value.answers.length > MAX_QUESTIONS_PER_ROUND) { + return null; + } + const answers = value.answers.flatMap((entry) => + isRecord(entry) && typeof entry.questionId === "string" && typeof entry.answer === "string" + ? [{ questionId: entry.questionId, answer: entry.answer }] + : [], + ); + return answers.length === value.answers.length ? { type: "answers", answers } : null; + } + + case "generate": + return { type: "generate" }; - return text; + case "skip": + return { type: "skip" }; + + default: + return null; + } } /** - * Sends a message in one of the current user's sessions and starts the design - * run that answers it. Responds 202 with the updated session (the reply is a - * PENDING message carrying the run id), or 502 with the session when the run - * could not be started. + * Sends a turn in one of the current user's sessions and starts the design + * agent run that answers it. Responds 202 with the updated session (the reply + * is a PENDING message carrying the run id), or 502 with the session when the + * run could not be started. */ export async function POST(request: Request, context: TurnRouteContext) { const { userId } = await auth(); @@ -44,12 +77,9 @@ export async function POST(request: Request, context: TurnRouteContext) { return Response.json({ error: "Invalid request body" }, { status: 400 }); } - const text = parseMessageText(body); - if (!text) { - return Response.json( - { error: `Message text is required and must be at most ${MAX_MESSAGE_LENGTH} characters` }, - { status: 400 }, - ); + const input = parseTurnInput(body); + if (!input) { + return Response.json({ error: "Invalid turn" }, { status: 400 }); } const { projectId, sessionId } = await context.params; @@ -59,7 +89,7 @@ export async function POST(request: Request, context: TurnRouteContext) { return Response.json({ error: "Forbidden" }, { status: 403 }); } - const result = await startTurn({ projectId: project.id, userId }, sessionId, text); + const result = await startTurn({ projectId: project.id, userId }, sessionId, input); if (result.ok) { return Response.json({ session: result.session }, { status: 202 }); } diff --git a/context/progress-tracker.md b/context/progress-tracker.md index 36aa4ce..8149439 100644 --- a/context/progress-tracker.md +++ b/context/progress-tracker.md @@ -9,8 +9,8 @@ Update this file whenever the current phase, active feature, or implementation s ## Current Goal - Persistent, multi-turn AI design sessions (plan: storage → wire sessions → turn engine → UI cards → generation - quality → docs). Steps 1 (storage, PR #21) and 2 (sessions wired into the sidebar) are done; next is step 3, - the clarify → plan → generate turn engine in the design-agent task. + quality → docs). Steps 1 (storage, PR #21), 2 (sessions in the sidebar, PR #22) and 3 (clarify → plan → generate + turn engine) are done; next is step 4, question/plan/result cards in the sidebar. - After that: the Specs tab (Generate Spec + automatic Markdown download), which remains inert. ## Completed @@ -1244,3 +1244,96 @@ Update this file whenever the current phase, active feature, or implementation s room, sessions and `TaskRun` removed. - Signed-out requests to the turns and session routes are stopped by Clerk. - Not checked in a browser (sidebar UI, reload resume, history switching and deleting, Realtime stage line). +- AI sessions, step 3: multi-turn design agent — clarify → plan → generate (2026-09-15, branch + `feat/ai-design-turn-engine`): + - Each turn is still one `design-agent` run; the task now decides what the turn does. + - `lib/ai/agent-schema.ts` (zod, client-safe): `designBriefSchema` (goal, scale, core features, non-functional, + constraints, assumptions, open questions), `clarifyQuestionSchema` (id, question, why, 2–4 options), + `turnAnalysisSchema` (updated brief + decision `ask|plan|generate` + reply + ≤3 questions), `designPlanSchema` + (summary, components with role/responsibility, flows, decisions with choice/rationale/alternatives, + assumptions), and `designAgentResultSchema` — the run output (`ask` / `plan` / `generated`), validated by the + server before storing. Also the task payload (`intent`, `input`, last 12 history entries, brief, latest plan, + `clarifyRounds`), limits (`MAX_CLARIFY_ROUNDS` 3, `MAX_QUESTIONS_PER_ROUND` 3), and text renderings of + questions/plans/answers used as message content (what the model reads back and the sidebar shows until step 4). + - `lib/ai/prompts.ts`: analyze, plan, and generate system prompts plus context builders. + - `lib/ai/design-agent-engine.ts` (no canvas/DB access): `analyzeTurn`, `draftPlan`, `generateDesignGraph` on + `generateText` + `Output.object` (replacing the deprecated `generateObject`), and `enforceDecision`, which + overrides the model: no questions after a skip or past 3 rounds, no `generate` without a plan, no `ask` with + zero questions. Gemini `thinkingLevel` per step: analyze `low`, plan `medium`, generate `low` — default thinking + took 34–82s per call; with these levels analyze calls took 14–24s in the eval. + - `src/trigger/design-agent.ts`: `generate` intent with a plan → draw it; otherwise analyze → questions, or draw + the existing plan when the user approved it in text, or draft/revise a plan. Stages `analyzing` → `planning` → + `generating` → `writing` → `done`. `retry: { maxAttempts: 1 }` — model calls already retry with backoff, and + turn-level retries multiplied quota use and could draw a diagram twice. + - `lib/ai/session-turns.ts`: + - `TurnInput`: `message` (text), `answers` (`{ questionId, answer }[]`, only while the latest reply is + questions, paired with the question text), `generate` (409 without a plan), `skip`. + - Settling stores `QUESTIONS` (payload `questions`, phase `CLARIFYING`, `clarifyRounds + 1`), `PLAN` (payload + `plan`, phase `PLANNED`), or `RESULT` (payload counts + the plan's `decisions`, phase `COMPLETE`), and saves + the brief on the session. Invalid run output is stored as an error. A failed turn leaves the phase unchanged. + - Run failures are stored as friendly messages (usage limit / busy / timeout / generic); the raw provider error + is only logged. + - Phase is no longer set to `GENERATING` when a turn starts. + - Turns route accepts all four inputs (`{ type: "message" | "answers" | "generate" | "skip" }`); 400 for anything else. + - `hooks/use-ai-session.ts`: stage labels for the new stages. The composer still sends only `message` turns, so for + now users answer questions and approve plans in free text ("looks good, draw it"); buttons come in step 4. + - Validation checks: + - `pnpm typecheck` and `pnpm lint` passed + - DB turn script (15 checks) passed with the new input types, including failed turns leaving the phase unchanged + - Engine eval against Gemini (default thinking): vague prompt → 3 relevant questions with options; answers → plan + with a goal and no repeated questions; detailed URL-shortener prompt → straight to a plan (10 components, 4 + decisions with alternatives, validates against the result schema); approval → `generate`; drawn graph used + exactly the 8 plan components with sensible edges; "add a Redis cache" → revised plan that kept every + component and added "Link Cache". All six enforced rules passed. + - With the new thinking levels: vague prompt and answers passed again (24s and 14s, from 75s and 39s). + - Full flow via the local Trigger worker confirmed the 409 guards (generate before a plan, answers without open + questions or after a plan), that a failed run is stored as the friendly usage-limit message, and cleanup. + - The Gemini key hit its free-tier quota on `gemini-3.5-flash` (20 requests), so testing moved to + `gemini-2.5-flash` (`GOOGLE_GENERATIVE_AI_MODEL` in `.env.local`; free tier 5 requests/minute): + - `design-agent-engine.ts` now sends `thinkingBudget` (low 1024, medium 4096 tokens) to `gemini-2.x` models, + which don't take `thinkingLevel`; Gemini 3+ keeps `thinkingLevel`. + - Eval on 2.5-flash, 4–20s per call: vague prompt → 3 questions with options; answers → plan; detailed prompt → + straight to a plan (13 components, 5 decisions with alternatives, validates); drawn graph used exactly the 8 + plan components; "add a Redis cache" → revised plan keeping every component plus "Link Cache"; skip → plan + with 6 recorded assumptions; "Looks good, go ahead and draw it" → `generate`; "Looks good, but use Postgres + instead" → `plan`. All passed. + - Full flow via Trigger: 409 guards passed; the single task attempt failed in 16s and stored the friendly + usage-limit message. + - Full flow via the restarted Trigger worker on 2.5-flash: first message → 3 questions (12.6s); answers stored as + ANSWERS paired with the questions → a second round of 3 questions (12.4s); `clarifyRounds` and phase correct + each round; brief stored. The skip turn then failed after 39.6s with "No object generated: response did not + match schema". + - Fix for that failure — list limits are no longer hard schema constraints: + - Cause: the model regularly fills lists up to their `maxItems` (a replay produced exactly 8 plan assumptions and + 8 brief features against caps of 8 and 12), and Gemini doesn't reliably enforce `maxItems`, so one extra item + failed the whole turn. + - `agent-schema.ts`: `.max()` removed from every model-filled list (brief lists, question options, plan + components/flows/decisions/alternatives/assumptions, questions per round) and the `options` minimum dropped; + limits live in `LIST_LIMITS` and are stated in the schema descriptions. New `clampBrief`, `clampQuestions`, + `clampPlan` trim to the limits, keeping the first items. Required minimums that matter stay (a plan needs a + component; an `ask` result needs a question). + - `design-agent-engine.ts`: every result is clamped, and each structured call retries once on + `NoObjectGeneratedError` (transport errors are already retried by the SDK). + - `design-generation.ts`: `designGraphSchema` no longer enforces 24 nodes / 48 edges; `buildCanvasGraph` trims to + those limits, and edges pointing at a trimmed node are still dropped. + - Offline check (no model calls, 15 checks): over-long briefs, questions, plans and graphs parse and are trimmed + to their limits, a single-option question parses, an empty plan or graph is still rejected, clamped plans + validate as stored run output, and no edge references a trimmed node. `pnpm typecheck` and `pnpm lint` pass. + - Moved testing to `gemini-3-flash-preview` (`gemini-3-flash` is not an id on this key; free tier 20 requests/day, + 5/minute). The first full-flow run hung: the run sat in `analyzing` for over 5 minutes with one attempt and no + error, while probes showed the model under load ("high demand" on a bare call) but answering `analyzeTurn` in + 30.7s. Nothing bounded a model call or the task (config `maxDuration` 3600s), so a held request would keep the + reply pending and the composer locked. Fix: + - `design-agent-engine.ts`: `timeout: { totalMs }` on every structured call — analyze 90s, plan 150s, generate + 120s. A timeout fails the turn with the existing "took too long" message. + - `src/trigger/design-agent.ts`: `maxDuration: 300` as the backstop. + - The stuck test run was cancelled. + - Re-run after the fix: the first turn again got no model response inside the worker and failed cleanly at the 90s + analyze timeout (settled after 97.1s as "The design agent took too long to respond. Try again."), so the + timeout path is verified. Two worker runs on `gemini-3-flash-preview` got no response while a direct call took + 30.7s and `gemini-2.5-flash` runs through the same worker succeeded, which points to preview-model overload + rather than code. + - Not yet verified: the full flow getting past the skip turn to a plan and a generated diagram through the Trigger + worker, with the fix. Both free-tier Gemini quotas on this key are used up (`gemini-3.5-flash` and + `gemini-2.5-flash`, 20 requests each); the last run failed on the first turn with the usage-limit message. + Re-run `verify-ai-turn-flow-e2e.ts` once quota resets or with a billed key. diff --git a/context/project-overview.md b/context/project-overview.md index 7a934eb..dbf1d82 100644 --- a/context/project-overview.md +++ b/context/project-overview.md @@ -19,8 +19,9 @@ Ghost AI is a real-time collaborative system design workspace. Users describe a 2. User creates or selects a project. 3. User enters the project workspace. 4. User optionally imports a starter system design template into the canvas. -5. User prompts the AI to generate or extend the system design. -6. AI generates nodes and edges in the shared canvas. +5. User describes the system to the AI; the AI asks clarifying questions when important requirements are missing, + then proposes a plan (components, flows, key decisions, assumptions) for the user to approve or revise. +6. On approval, AI generates nodes and edges in the shared canvas. 7. Collaborators edit and refine the design. 8. User triggers spec generation. 9. App persists the generated Markdown spec. @@ -49,7 +50,14 @@ Ghost AI is a real-time collaborative system design workspace. Users describe a ### AI Architecture Generation -- AI generates a system design from a user-supplied prompt. +- AI generates a system design through a short conversation instead of a single prompt: + - It keeps a running requirements brief across turns. + - It asks up to 3 rounds of clarifying questions (at most 3 per round, each with suggested answers) when an + unknown would change the architecture; users can skip to planning with stated assumptions. + - It proposes a plan — components, key flows, major decisions with rationale and alternatives, and assumptions — + that the user approves or asks to revise before anything is drawn. +- AI chat sessions are saved per user per project for 7 days after last use, so a conversation and an in-flight + generation survive a reload. ### Spec Generation diff --git a/hooks/use-ai-session.ts b/hooks/use-ai-session.ts index d23c102..de18aba 100644 --- a/hooks/use-ai-session.ts +++ b/hooks/use-ai-session.ts @@ -44,12 +44,14 @@ export interface UseAiSessionResult { } const STAGE_LABELS: Record = { + analyzing: "Reading your requirements…", + planning: "Drafting a plan…", generating: "Designing the architecture…", writing: "Adding components to the canvas…", done: "Finishing up…", }; -const STARTING_STATUS = "Starting the design run…"; +const STARTING_STATUS = "Thinking…"; const GENERIC_ERROR = "Something went wrong sending that message. Try again."; diff --git a/lib/ai/agent-schema.ts b/lib/ai/agent-schema.ts new file mode 100644 index 0000000..6ae7744 --- /dev/null +++ b/lib/ai/agent-schema.ts @@ -0,0 +1,282 @@ +import { z } from "zod"; + +// --------------------------------------------------------------------------- +// Design agent contract +// +// A design session moves through turns: the agent keeps a running brief of +// the requirements, asks clarifying questions while something important is +// unknown, proposes a plan, and generates the diagram once the plan is +// approved. Each turn is one Trigger.dev run. These schemas are shared by the +// task (model output), the server (validating run output before storing it), +// and the client (rendering questions and plans). +// --------------------------------------------------------------------------- + +/** Clarifying rounds allowed before the agent must plan with assumptions. */ +export const MAX_CLARIFY_ROUNDS = 3; + +export const MAX_QUESTIONS_PER_ROUND = 3; + +/** Most recent transcript messages sent to the model with each turn. */ +export const AGENT_HISTORY_LIMIT = 12; + +/** + * Size limits for lists the model fills in. They are stated in the schema + * descriptions but not enforced as `maxItems`: the model does not reliably + * honour them, and one item over a hard limit failed the whole turn. Lists + * are trimmed to these sizes by the clamp functions below instead. + */ +export const LIST_LIMITS = { + coreFeatures: 12, + briefItems: 8, + openQuestions: 6, + questionOptions: 4, + components: 24, + flows: 6, + decisions: 6, + alternatives: 3, + assumptions: 8, +} as const; + +function atMost(limit: number, description: string): string { + return `${description} At most ${limit} items.`; +} + +export const designBriefSchema = z.object({ + goal: z + .string() + .describe("One sentence: what the system does and for whom. Empty string if still unknown."), + scale: z + .string() + .describe("Expected users, traffic, or data volume, as stated or assumed. Empty string if unknown."), + coreFeatures: z + .array(z.string()) + .describe(atMost(LIST_LIMITS.coreFeatures, "Capabilities the system must provide.")), + nonFunctional: z + .array(z.string()) + .describe(atMost(LIST_LIMITS.briefItems, "Latency, availability, consistency, security, or compliance requirements.")), + constraints: z + .array(z.string()) + .describe( + atMost(LIST_LIMITS.briefItems, "Technology preferences, cloud, budget, team, or existing systems to integrate with."), + ), + assumptions: z + .array(z.string()) + .describe(atMost(LIST_LIMITS.briefItems, "Things decided without the user stating them.")), + openQuestions: z + .array(z.string()) + .describe(atMost(LIST_LIMITS.openQuestions, "Unknowns that would still change the architecture.")), +}); + +export type DesignBrief = z.infer; + +export const clarifyQuestionSchema = z.object({ + id: z.string().min(1).describe("Short slug, unique within this round, e.g. 'expected-scale'."), + question: z.string().min(1).describe("The question, phrased for a product owner, not a DBA."), + why: z.string().describe("One short clause on how the answer changes the design."), + options: z + .array(z.string()) + .describe(atMost(LIST_LIMITS.questionOptions, "2 to 4 concrete likely answers the user can pick from.")), +}); + +export type ClarifyQuestion = z.infer; + +export const TURN_DECISIONS = ["ask", "plan", "generate"] as const; + +export type TurnDecision = (typeof TURN_DECISIONS)[number]; + +/** What the model returns when it reads a user turn. */ +export const turnAnalysisSchema = z.object({ + brief: designBriefSchema.describe("The full updated brief, merging everything said so far."), + decision: z + .enum(TURN_DECISIONS) + .describe( + "'ask' to gather missing requirements, 'plan' to propose or revise a plan, " + + "'generate' only when a plan exists and the user approved it.", + ), + reply: z + .string() + .describe("One or two plain sentences to the user. No markdown, and do not restate the questions."), + questions: z + .array(clarifyQuestionSchema) + .describe( + atMost(MAX_QUESTIONS_PER_ROUND, "Questions for this round, most important first. Empty unless decision is 'ask'."), + ), +}); + +export type TurnAnalysis = z.infer; + +export const planDecisionSchema = z.object({ + title: z.string().min(1).describe("The question the decision answers, e.g. 'Primary datastore'."), + choice: z.string().min(1).describe("What was chosen."), + rationale: z.string().min(1).describe("Why, tied to the brief."), + alternatives: z + .array(z.string()) + .describe(atMost(LIST_LIMITS.alternatives, "Realistic options that were not chosen.")), +}); + +export type PlanDecision = z.infer; + +export const designPlanSchema = z.object({ + summary: z.string().min(1).describe("Two or three sentences describing the architecture."), + components: z + .array( + z.object({ + name: z.string().min(1).describe("Short unique component name, used verbatim on the diagram."), + role: z + .string() + .min(1) + .describe("One or two word kind: Client, Gateway, Service, Worker, Queue, Cache, Database, External."), + responsibility: z.string().min(1).describe("What it does, in one sentence."), + }), + ) + .min(1) + .describe(atMost(LIST_LIMITS.components, "The components to draw.")), + flows: z + .array(z.string()) + .describe(atMost(LIST_LIMITS.flows, "Key request or data flows, each written as 'A → B → C: purpose'.")), + decisions: z + .array(planDecisionSchema) + .describe(atMost(LIST_LIMITS.decisions, "The major architectural choices, most consequential first.")), + assumptions: z.array(z.string()).describe(atMost(LIST_LIMITS.assumptions, "Assumptions the plan depends on.")), +}); + +export type DesignPlan = z.infer; + +/** Output of one design-agent run. The server validates it before storing it. */ +export const designAgentResultSchema = z.discriminatedUnion("action", [ + z.object({ + action: z.literal("ask"), + reply: z.string(), + questions: z.array(clarifyQuestionSchema).min(1), + brief: designBriefSchema, + }), + z.object({ + action: z.literal("plan"), + reply: z.string(), + plan: designPlanSchema, + brief: designBriefSchema, + }), + z.object({ + action: z.literal("generated"), + reply: z.string(), + nodeCount: z.number().int().nonnegative(), + edgeCount: z.number().int().nonnegative(), + decisions: z.array(planDecisionSchema), + brief: designBriefSchema, + }), +]); + +export type DesignAgentResult = z.infer; + +// --------------------------------------------------------------------------- +// Clamping +// --------------------------------------------------------------------------- + +export function clampBrief(brief: DesignBrief): DesignBrief { + return { + ...brief, + coreFeatures: brief.coreFeatures.slice(0, LIST_LIMITS.coreFeatures), + nonFunctional: brief.nonFunctional.slice(0, LIST_LIMITS.briefItems), + constraints: brief.constraints.slice(0, LIST_LIMITS.briefItems), + assumptions: brief.assumptions.slice(0, LIST_LIMITS.briefItems), + openQuestions: brief.openQuestions.slice(0, LIST_LIMITS.openQuestions), + }; +} + +export function clampQuestions(questions: readonly ClarifyQuestion[]): ClarifyQuestion[] { + return questions.slice(0, MAX_QUESTIONS_PER_ROUND).map((question) => ({ + ...question, + options: question.options.slice(0, LIST_LIMITS.questionOptions), + })); +} + +export function clampPlan(plan: DesignPlan): DesignPlan { + return { + ...plan, + components: plan.components.slice(0, LIST_LIMITS.components), + flows: plan.flows.slice(0, LIST_LIMITS.flows), + decisions: plan.decisions.slice(0, LIST_LIMITS.decisions).map((decision) => ({ + ...decision, + alternatives: decision.alternatives.slice(0, LIST_LIMITS.alternatives), + })), + assumptions: plan.assumptions.slice(0, LIST_LIMITS.assumptions), + }; +} + +// --------------------------------------------------------------------------- +// Turn input +// --------------------------------------------------------------------------- + +/** + * - `message`: free text from the composer + * - `answers`: replies to the latest clarifying questions + * - `generate`: approve the latest plan and draw it + * - `skip`: stop asking questions and plan with assumptions + */ +export const TURN_INTENTS = ["message", "answers", "generate", "skip"] as const; + +export type TurnIntent = (typeof TURN_INTENTS)[number]; + +export interface ClarifyAnswer { + questionId: string; + answer: string; +} + +export interface AgentHistoryEntry { + role: "user" | "assistant"; + content: string; +} + +export interface DesignAgentPayload { + roomId: string; + intent: TurnIntent; + /** The user's turn rendered as text (answers are paired with their questions). */ + input: string; + /** Prior transcript, oldest first, at most {@link AGENT_HISTORY_LIMIT} entries. */ + history: AgentHistoryEntry[]; + brief: DesignBrief | null; + /** The latest plan in the session, if one was proposed. */ + plan: DesignPlan | null; + clarifyRounds: number; +} + +// --------------------------------------------------------------------------- +// Text renderings +// +// Stored as message content: they are what the model reads back as history, +// and what the sidebar shows for messages it has no dedicated card for. +// --------------------------------------------------------------------------- + +export function formatQuestionsText(reply: string, questions: readonly ClarifyQuestion[]): string { + const lines = questions.map((question, index) => { + const options = question.options.length > 0 ? ` (e.g. ${question.options.join(" / ")})` : ""; + return `${index + 1}. ${question.question}${options}`; + }); + return [reply.trim(), ...lines].filter(Boolean).join("\n"); +} + +export function formatPlanText(reply: string, plan: DesignPlan): string { + const sections = [ + reply.trim(), + plan.summary, + `Components:\n${plan.components.map((c) => `- ${c.name} (${c.role}): ${c.responsibility}`).join("\n")}`, + plan.flows.length > 0 ? `Flows:\n${plan.flows.map((flow) => `- ${flow}`).join("\n")}` : "", + plan.decisions.length > 0 + ? `Key decisions:\n${plan.decisions.map((d) => `- ${d.title}: ${d.choice} — ${d.rationale}`).join("\n")}` + : "", + plan.assumptions.length > 0 ? `Assumptions:\n${plan.assumptions.map((a) => `- ${a}`).join("\n")}` : "", + ]; + return sections.filter(Boolean).join("\n\n"); +} + +export function formatAnswersText( + questions: readonly ClarifyQuestion[], + answers: readonly ClarifyAnswer[], +): string { + return answers + .map((entry) => { + const question = questions.find((candidate) => candidate.id === entry.questionId); + return question ? `${question.question}\n→ ${entry.answer}` : entry.answer; + }) + .join("\n\n"); +} diff --git a/lib/ai/design-agent-engine.ts b/lib/ai/design-agent-engine.ts new file mode 100644 index 0000000..bb7d7b3 --- /dev/null +++ b/lib/ai/design-agent-engine.ts @@ -0,0 +1,204 @@ +import { google, type GoogleLanguageModelOptions } from "@ai-sdk/google"; +import { generateText, NoObjectGeneratedError, Output, type LanguageModel, type ModelMessage } from "ai"; + +import { + clampBrief, + clampPlan, + clampQuestions, + designBriefSchema, + designPlanSchema, + MAX_CLARIFY_ROUNDS, + turnAnalysisSchema, + type DesignAgentPayload, + type DesignBrief, + type DesignPlan, + type TurnAnalysis, +} from "@/lib/ai/agent-schema"; +import { + ANALYZE_SYSTEM_PROMPT, + buildAnalyzeContext, + buildGeneratePrompt, + buildPlanPrompt, + GENERATE_SYSTEM_PROMPT, + PLAN_SYSTEM_PROMPT, +} from "@/lib/ai/prompts"; +import { designGraphSchema, type DesignGraph } from "@/lib/design-generation"; + +// Model calls for one design turn. No canvas or database access, so the steps +// can be exercised directly from a script. + +const DEFAULT_MODEL_ID = "gemini-3.5-flash"; + +/** + * Upper bound for each structured call, retries included. Under provider load a + * request was held open for over five minutes with no response, which left the + * turn pending and the composer locked; a timeout fails the turn instead. + */ +const CALL_TIMEOUT_MS = { + analyze: 90_000, + plan: 150_000, + generate: 120_000, +} as const; + +type ThinkingLevel = "low" | "medium"; + +/** Gemini 2.x models take a token budget instead of a thinking level. */ +const THINKING_BUDGETS: Record = { + low: 1024, + medium: 4096, +}; + +export function resolveDesignModel(): LanguageModel { + const configured = process.env.GOOGLE_GENERATIVE_AI_MODEL?.trim(); + return google(configured && configured.length > 0 ? configured : DEFAULT_MODEL_ID); +} + +function modelIdOf(model: LanguageModel): string { + return typeof model === "string" ? model : model.modelId; +} + +/** + * Model reasoning depth per step. Default thinking made a single turn take up + * to ~80s; reading a turn and drawing an agreed plan need little reasoning, + * while drafting the plan (where the architectural decisions are made) keeps + * more. Gemini 3+ takes `thinkingLevel`; Gemini 2.x only understands + * `thinkingBudget`. + */ +function thinking(model: LanguageModel, level: ThinkingLevel) { + const thinkingConfig = /^gemini-2\./.test(modelIdOf(model)) + ? { thinkingBudget: THINKING_BUDGETS[level] } + : { thinkingLevel: level }; + return { google: { thinkingConfig } satisfies GoogleLanguageModelOptions }; +} + +/** + * Runs a structured-output call, retrying once when the response does not + * match the schema. Transport errors are already retried by the SDK; a schema + * mismatch is usually a one-off bad sample, so a single fresh attempt is + * cheaper than failing the whole turn. + */ +async function withSchemaRetry(call: () => Promise): Promise { + try { + return await call(); + } catch (error) { + if (!NoObjectGeneratedError.isInstance(error)) { + throw error; + } + console.warn("[design-agent] model output did not match the schema; retrying once", error.cause); + return call(); + } +} + +/** + * Builds the model conversation from stored history plus the new turn. Leading + * assistant entries are dropped, since a trimmed history can start mid-reply + * and providers expect a conversation to open with the user. + */ +function toMessages(payload: DesignAgentPayload): ModelMessage[] { + const firstUser = payload.history.findIndex((entry) => entry.role === "user"); + const history = firstUser === -1 ? [] : payload.history.slice(firstUser); + + return [ + ...history.map((entry): ModelMessage => ({ role: entry.role, content: entry.content })), + { role: "user", content: payload.input }, + ]; +} + +/** + * Applies the rules the model is asked to follow but cannot be trusted to: + * no questions past the round limit or after a skip, no generating without a + * plan, no "ask" without questions. + */ +export function enforceDecision(analysis: TurnAnalysis, payload: DesignAgentPayload): TurnAnalysis { + let { decision } = analysis; + + if (decision === "ask") { + const mayAsk = + payload.intent !== "skip" && + payload.clarifyRounds < MAX_CLARIFY_ROUNDS && + analysis.questions.length > 0; + if (!mayAsk) { + decision = "plan"; + } + } + + if (decision === "generate" && !payload.plan) { + decision = "plan"; + } + + return { + ...analysis, + decision, + questions: decision === "ask" ? analysis.questions : [], + }; +} + +export async function analyzeTurn(model: LanguageModel, payload: DesignAgentPayload): Promise { + const { output } = await withSchemaRetry(() => + generateText({ + model, + system: `${ANALYZE_SYSTEM_PROMPT}\n\n${buildAnalyzeContext(payload)}`, + messages: toMessages(payload), + output: Output.object({ schema: turnAnalysisSchema, name: "turn_analysis" }), + providerOptions: thinking(model, "low"), + timeout: { totalMs: CALL_TIMEOUT_MS.analyze }, + }), + ); + + return enforceDecision( + { ...output, brief: clampBrief(output.brief), questions: clampQuestions(output.questions) }, + payload, + ); +} + +export async function draftPlan( + model: LanguageModel, + brief: DesignBrief, + previousPlan: DesignPlan | null, + latestInput: string, +): Promise { + const { output } = await withSchemaRetry(() => + generateText({ + model, + system: PLAN_SYSTEM_PROMPT, + prompt: buildPlanPrompt(brief, previousPlan, latestInput), + output: Output.object({ schema: designPlanSchema, name: "design_plan" }), + providerOptions: thinking(model, "medium"), + timeout: { totalMs: CALL_TIMEOUT_MS.plan }, + }), + ); + + return clampPlan(output); +} + +export async function generateDesignGraph( + model: LanguageModel, + brief: DesignBrief, + plan: DesignPlan, +): Promise { + const { output } = await withSchemaRetry(() => + generateText({ + model, + system: GENERATE_SYSTEM_PROMPT, + prompt: buildGeneratePrompt(brief, plan), + output: Output.object({ schema: designGraphSchema, name: "design_graph" }), + providerOptions: thinking(model, "low"), + timeout: { totalMs: CALL_TIMEOUT_MS.generate }, + }), + ); + + return output; +} + +/** An empty brief, for a generate turn that somehow has none stored. */ +export function emptyBrief(): DesignBrief { + return designBriefSchema.parse({ + goal: "", + scale: "", + coreFeatures: [], + nonFunctional: [], + constraints: [], + assumptions: [], + openQuestions: [], + }); +} diff --git a/lib/ai/prompts.ts b/lib/ai/prompts.ts new file mode 100644 index 0000000..b57b763 --- /dev/null +++ b/lib/ai/prompts.ts @@ -0,0 +1,108 @@ +import { + MAX_CLARIFY_ROUNDS, + MAX_QUESTIONS_PER_ROUND, + type DesignAgentPayload, + type DesignBrief, + type DesignPlan, +} from "@/lib/ai/agent-schema"; + +// --------------------------------------------------------------------------- +// Analyze: read the user's turn, update the brief, decide what happens next. +// --------------------------------------------------------------------------- + +export const ANALYZE_SYSTEM_PROMPT = [ + "You are Draftly's system design architect. You gather requirements in a short conversation,", + "then propose an architecture plan, then draw it once the user approves.", + "", + "On every turn:", + "1. Update the brief from the whole conversation. Keep earlier facts unless the user contradicts", + " them. Record anything you decide on the user's behalf under assumptions.", + "2. Choose a decision:", + ` - "ask": an unknown would materially change the architecture — scale, real-time needs,`, + " consistency, data sensitivity, integrations, or deployment constraints. Ask at most", + ` ${MAX_QUESTIONS_PER_ROUND} questions, most important first, each with 2 to 4 concrete options.`, + " Never ask what can reasonably be assumed, and never repeat a question already answered.", + ` - "plan": the brief is good enough to design, or the user asked to change an existing plan.`, + ` - "generate": only when a plan already exists and the user's latest message approves it`, + " 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.", +].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 { + const notes = [ + `Current brief:\n${describeBrief(payload.brief)}`, + payload.plan + ? `A plan has been proposed:\n${describePlan(payload.plan)}` + : "No plan has been proposed yet, so 'generate' is not available.", + `Clarifying rounds used: ${payload.clarifyRounds} of ${MAX_CLARIFY_ROUNDS}.`, + ]; + + if (payload.intent === "answers") { + notes.push("The latest message answers your clarifying questions."); + } + if (payload.intent === "skip") { + notes.push("The user chose to skip further questions: decide 'plan' and record assumptions for the gaps."); + } + if (payload.clarifyRounds >= MAX_CLARIFY_ROUNDS) { + notes.push("The question limit is reached: do not decide 'ask'."); + } + + return notes.join("\n\n"); +} + +// --------------------------------------------------------------------------- +// Plan: turn the brief into components, flows, and decisions. +// --------------------------------------------------------------------------- + +export const PLAN_SYSTEM_PROMPT = [ + "You are Draftly's system design architect. Turn the requirements brief into an architecture plan", + "the user will review before it is drawn.", + "- Components: a focused set (usually 4 to 16) of clients, gateways, services, workers, queues,", + " caches, datastores, and external systems. Names are short, concrete, and unique.", + "- Flows: the few request or data paths that explain how the system works.", + "- Decisions: the major architectural choices — datastore, sync vs async messaging, caching,", + " authentication, scaling approach, and anything the brief makes contentious. Tie each rationale", + " 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.", +].join("\n"); + +function describePlan(plan: DesignPlan): string { + return JSON.stringify(plan, null, 2); +} + +export function buildPlanPrompt(brief: DesignBrief, previousPlan: DesignPlan | null, latestInput: string): string { + return [ + `Requirements brief:\n${JSON.stringify(brief, null, 2)}`, + previousPlan ? `Previous plan:\n${describePlan(previousPlan)}` : "", + `The user's latest message:\n${latestInput}`, + ] + .filter(Boolean) + .join("\n\n"); +} + +// --------------------------------------------------------------------------- +// Generate: draw the approved plan as a diagram. +// --------------------------------------------------------------------------- + +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.", + "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.", +].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", + ); +} diff --git a/lib/ai/session-turns.ts b/lib/ai/session-turns.ts index abb5e4f..805ab75 100644 --- a/lib/ai/session-turns.ts +++ b/lib/ai/session-turns.ts @@ -1,15 +1,42 @@ import { ApiError, runs, tasks } from "@trigger.dev/sdk/v3"; +import { z } from "zod"; +import type { Prisma } from "@/app/generated/prisma/client"; +import { + AGENT_HISTORY_LIMIT, + clarifyQuestionSchema, + designAgentResultSchema, + designBriefSchema, + designPlanSchema, + formatAnswersText, + formatPlanText, + formatQuestionsText, + type AgentHistoryEntry, + type ClarifyAnswer, + type ClarifyQuestion, + type DesignAgentPayload, + type DesignAgentResult, + type DesignBrief, + type DesignPlan, + type TurnIntent, +} from "@/lib/ai/agent-schema"; import { computeExpiresAt, DEFAULT_SESSION_TITLE, getSession, + MAX_MESSAGE_LENGTH, MAX_MESSAGES_PER_SESSION, toSessionTitle, } from "@/lib/ai/session-store"; import { prisma } from "@/lib/prisma"; import type { designAgentTask } from "@/src/trigger/design-agent"; -import type { AiMessageKind, AiMessageStatus, AiSessionDetail, AiSessionPhase } from "@/types/ai-session"; +import type { + AiMessageDto, + AiMessageKind, + AiMessageStatus, + AiSessionDetail, + AiSessionPhase, +} from "@/types/ai-session"; // --------------------------------------------------------------------------- // Turns @@ -40,24 +67,132 @@ export function summarizeDesignResult(nodeCount: number, edgeCount: number): str ); } +// --------------------------------------------------------------------------- +// Reading the session back for the agent +// --------------------------------------------------------------------------- + +function latestAssistantMessage(session: AiSessionDetail): AiMessageDto | undefined { + return session.messages.findLast((message) => message.role === "ASSISTANT" && message.status === "COMPLETE"); +} + +function readPayloadField(message: AiMessageDto | undefined, field: string): unknown { + const payload = message?.payload; + return typeof payload === "object" && payload !== null ? (payload as Record)[field] : undefined; +} + +/** The most recent plan proposed in the session, if any. */ +function latestPlan(session: AiSessionDetail): DesignPlan | null { + const message = session.messages.findLast((entry) => entry.kind === "PLAN" && entry.status === "COMPLETE"); + const parsed = designPlanSchema.safeParse(readPayloadField(message, "plan")); + return parsed.success ? parsed.data : null; +} + +/** Questions still awaiting answers: only while they are the latest reply. */ +function openQuestions(session: AiSessionDetail): ClarifyQuestion[] { + const message = latestAssistantMessage(session); + if (message?.kind !== "QUESTIONS") { + return []; + } + const parsed = z.array(clarifyQuestionSchema).safeParse(readPayloadField(message, "questions")); + return parsed.success ? parsed.data : []; +} + +function storedBrief(session: AiSessionDetail): DesignBrief | null { + const parsed = designBriefSchema.safeParse(session.brief); + return parsed.success ? parsed.data : null; +} + +/** Completed transcript the model sees, oldest first. Failed replies are left out. */ +function agentHistory(session: AiSessionDetail): AgentHistoryEntry[] { + return session.messages + .filter((message) => message.status === "COMPLETE" && message.content.length > 0) + .map((message): AgentHistoryEntry => ({ + role: message.role === "USER" ? "user" : "assistant", + content: message.content, + })) + .slice(-AGENT_HISTORY_LIMIT); +} + +// --------------------------------------------------------------------------- +// Settling +// --------------------------------------------------------------------------- + interface TurnOutcome { kind: AiMessageKind; status: AiMessageStatus; content: string; - payload?: { nodeCount: number; edgeCount: number }; - phase: AiSessionPhase; + payload?: Prisma.InputJsonObject; + /** Omitted for failures, so a failed turn leaves the session where it was. */ + phase?: AiSessionPhase; + brief?: DesignBrief; + countsClarifyRound?: boolean; } function failedOutcome(content: string): TurnOutcome { - return { kind: "ERROR", status: "FAILED", content, phase: "CLARIFYING" }; + return { kind: "ERROR", status: "FAILED", content }; +} + +/** + * Turns a failed run's error into something a user can act on. Raw provider + * errors (quota details, retry counts, URLs) are logged, never stored. + */ +function describeRunFailure(runId: string, message: string | undefined): string { + console.error("[ai-sessions] design-agent run failed", runId, message); + + if (!message) { + return GENERIC_RUN_ERROR; + } + if (/quota|rate.?limit|too many requests|\b429\b/i.test(message)) { + return "Draftly AI has hit its usage limit for the moment. Wait a minute and try again."; + } + if (/high demand|overloaded|unavailable|\b503\b/i.test(message)) { + return "Draftly AI is busy right now. Try again in a moment."; + } + if (/timed? ?out|timeout|deadline/i.test(message)) { + return "The design agent took too long to respond. Try again."; + } + return GENERIC_RUN_ERROR; +} + +function outcomeFromResult(result: DesignAgentResult): TurnOutcome { + switch (result.action) { + case "ask": + return { + kind: "QUESTIONS", + status: "COMPLETE", + content: formatQuestionsText(result.reply, result.questions), + payload: { questions: result.questions }, + phase: "CLARIFYING", + brief: result.brief, + countsClarifyRound: true, + }; + case "plan": + return { + kind: "PLAN", + status: "COMPLETE", + content: formatPlanText(result.reply, result.plan), + payload: { plan: result.plan }, + phase: "PLANNED", + brief: result.brief, + }; + case "generated": + return { + kind: "RESULT", + status: "COMPLETE", + content: summarizeDesignResult(result.nodeCount, result.edgeCount), + payload: { nodeCount: result.nodeCount, edgeCount: result.edgeCount, decisions: result.decisions }, + phase: "COMPLETE", + brief: result.brief, + }; + } } /** * Writes a run's outcome onto its pending message. The status guard makes this * safe to call concurrently: only the first caller changes anything. */ -async function applyOutcome(sessionId: string, messageId: string, outcome: TurnOutcome): Promise { - return prisma.$transaction(async (tx) => { +async function applyOutcome(sessionId: string, messageId: string, outcome: TurnOutcome): Promise { + await prisma.$transaction(async (tx) => { const { count } = await tx.aiMessage.updateMany({ where: { id: messageId, sessionId, status: "PENDING" }, data: { @@ -68,12 +203,18 @@ async function applyOutcome(sessionId: string, messageId: string, outcome: TurnO }, }); - if (count === 0) { - return false; + if (count === 0 || (!outcome.phase && !outcome.brief)) { + return; } - await tx.aiSession.update({ where: { id: sessionId }, data: { phase: outcome.phase } }); - return true; + await tx.aiSession.update({ + where: { id: sessionId }, + data: { + ...(outcome.phase ? { phase: outcome.phase } : {}), + ...(outcome.brief ? { brief: outcome.brief } : {}), + ...(outcome.countsClarifyRound ? { clarifyRounds: { increment: 1 } } : {}), + }, + }); }); } @@ -100,19 +241,18 @@ async function settleTurn(sessionId: string, messageId: string, runId: string): let outcome: TurnOutcome; if (run.isSuccess) { - const nodeCount = run.output?.nodeCount ?? 0; - const edgeCount = run.output?.edgeCount ?? 0; - outcome = { - kind: "RESULT", - status: "COMPLETE", - content: summarizeDesignResult(nodeCount, edgeCount), - payload: { nodeCount, edgeCount }, - phase: "COMPLETE", - }; + // Run output crosses a service boundary; don't store it unchecked. + const parsed = designAgentResultSchema.safeParse(run.output); + if (parsed.success) { + outcome = outcomeFromResult(parsed.data); + } else { + console.error("[ai-sessions] unexpected design-agent output", runId, parsed.error.issues); + outcome = failedOutcome("The design agent returned an unexpected response. Try again."); + } } else if (run.isCancelled) { outcome = failedOutcome("The design run was cancelled."); } else if (run.isFailed) { - outcome = failedOutcome(run.error?.message ?? GENERIC_RUN_ERROR); + outcome = failedOutcome(describeRunFailure(runId, run.error?.message)); } else { return false; } @@ -140,18 +280,79 @@ export async function getSettledSession(scope: SessionScope, sessionId: string): return changed.some(Boolean) ? getSession(scope, sessionId) : session; } +// --------------------------------------------------------------------------- +// Starting a turn +// --------------------------------------------------------------------------- + +export type TurnInput = + | { type: "message"; text: string } + | { type: "answers"; answers: ClarifyAnswer[] } + | { type: "generate" } + | { type: "skip" }; + +interface ResolvedTurn { + intent: TurnIntent; + /** What is stored and shown as the user's message, and what the model reads. */ + content: string; + kind: AiMessageKind; + payload?: Prisma.InputJsonObject; +} + +type ResolveTurnResult = { ok: true; turn: ResolvedTurn } | { ok: false; status: 400 | 409; error: string }; + +/** Checks a turn against the session state and renders it as a user message. */ +function resolveTurn(session: AiSessionDetail, input: TurnInput): ResolveTurnResult { + switch (input.type) { + case "message": + return { ok: true, turn: { intent: "message", content: input.text, kind: "TEXT" } }; + + case "answers": { + const questions = openQuestions(session); + if (questions.length === 0) { + return { ok: false, status: 409, error: "There are no open questions to answer" }; + } + + const answers = input.answers + .map((entry) => ({ questionId: entry.questionId, answer: entry.answer.trim() })) + .filter((entry) => entry.answer.length > 0 && questions.some((question) => question.id === entry.questionId)); + if (answers.length === 0) { + return { ok: false, status: 400, error: "Answer at least one question" }; + } + + const content = formatAnswersText(questions, answers); + if (content.length > MAX_MESSAGE_LENGTH) { + return { ok: false, status: 400, error: `Answers must be at most ${MAX_MESSAGE_LENGTH} characters` }; + } + + return { ok: true, turn: { intent: "answers", content, kind: "ANSWERS", payload: { answers } } }; + } + + case "generate": + if (!latestPlan(session)) { + return { ok: false, status: 409, error: "There is no plan to generate yet" }; + } + return { ok: true, turn: { intent: "generate", content: "Generate this plan.", kind: "TEXT" } }; + + case "skip": + return { + ok: true, + turn: { intent: "skip", content: "Skip the questions and plan with sensible assumptions.", kind: "TEXT" }, + }; + } +} + export type StartTurnResult = | { ok: true; session: AiSessionDetail } - | { ok: false; status: 404 | 409 | 502; error: string; session?: AiSessionDetail }; + | { ok: false; status: 400 | 404 | 409 | 502; error: string; session?: AiSessionDetail }; /** - * Records a user message and starts the design run that answers it. + * Records a user turn and starts the design-agent run that answers it. * * The run is triggered before anything is written, so a stored PENDING message * always has a run id to settle from. If the trigger fails, the user message is * still stored with an error reply, so the transcript explains what happened. */ -export async function startTurn(scope: SessionScope, sessionId: string, text: string): Promise { +export async function startTurn(scope: SessionScope, sessionId: string, input: TurnInput): Promise { const session = await getSettledSession(scope, sessionId); if (!session) { return { ok: false, status: 404, error: "Session not found" }; @@ -165,12 +366,25 @@ export async function startTurn(scope: SessionScope, sessionId: string, text: st return { ok: false, status: 409, error: "This chat is full. Start a new chat to keep designing." }; } + const resolved = resolveTurn(session, input); + if (!resolved.ok) { + return resolved; + } + const { turn } = resolved; + + const payload: DesignAgentPayload = { + roomId: scope.projectId, + intent: turn.intent, + input: turn.content, + history: agentHistory(session), + brief: storedBrief(session), + plan: latestPlan(session), + clarifyRounds: session.clarifyRounds, + }; + let runId: string | null = null; try { - const handle = await tasks.trigger("design-agent", { - prompt: text, - roomId: scope.projectId, - }); + const handle = await tasks.trigger("design-agent", payload); runId = handle.id; } catch (error) { // wrong-environment TRIGGER_SECRET_KEY, a branch env that does not exist, or @@ -185,7 +399,14 @@ export async function startTurn(scope: SessionScope, sessionId: string, text: st await prisma.$transaction(async (tx) => { await tx.aiMessage.create({ - data: { sessionId, role: "USER", kind: "TEXT", content: text, createdAt: userAt }, + data: { + sessionId, + role: "USER", + kind: turn.kind, + content: turn.content, + ...(turn.payload ? { payload: turn.payload } : {}), + createdAt: userAt, + }, }); await tx.aiMessage.create({ @@ -210,8 +431,9 @@ export async function startTurn(scope: SessionScope, sessionId: string, text: st data: { lastActivityAt: userAt, expiresAt: computeExpiresAt(userAt), - ...(runId ? { phase: "GENERATING" } : {}), - ...(session.title === DEFAULT_SESSION_TITLE ? { title: toSessionTitle(text) } : {}), + ...(input.type === "message" && session.title === DEFAULT_SESSION_TITLE + ? { title: toSessionTitle(input.text) } + : {}), }, }); }); diff --git a/lib/design-generation.ts b/lib/design-generation.ts index e04b622..a80bac5 100644 --- a/lib/design-generation.ts +++ b/lib/design-generation.ts @@ -36,8 +36,12 @@ import { /** 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; +/** + * Progress stages a design run can report. A turn starts at `analyzing`, then + * either finishes (questions), goes through `planning`, or — once a plan is + * approved — through `generating` and `writing`. + */ +export const DESIGN_AGENT_STAGES = ["analyzing", "planning", "generating", "writing", "done"] as const; export type DesignAgentStage = (typeof DESIGN_AGENT_STAGES)[number]; @@ -90,9 +94,11 @@ export const designEdgeSchema = z.object({ }); /** 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).max(MAX_NODES), - edges: z.array(designEdgeSchema).max(MAX_EDGES), + nodes: z.array(designNodeSchema).min(1).describe(`At most ${MAX_NODES} components.`), + edges: z.array(designEdgeSchema).describe(`At most ${MAX_EDGES} connections.`), }); export type DesignGraph = z.infer; @@ -196,13 +202,15 @@ export function buildCanvasGraph( { 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 uniqueNodes = design.nodes + .filter((node) => { + if (seenNodeIds.has(node.id)) { + return false; + } + seenNodeIds.add(node.id); + return true; + }) + .slice(0, MAX_NODES); const cells = assignGridCells(uniqueNodes); const canvasNodeIds = new Map(); @@ -240,6 +248,9 @@ export function buildCanvasGraph( const edges: CanvasEdge[] = []; for (const edge of design.edges) { + if (edges.length >= MAX_EDGES) { + break; + } const source = canvasNodeIds.get(edge.source); const target = canvasNodeIds.get(edge.target); if (!source || !target || source === target) { diff --git a/src/trigger/design-agent.ts b/src/trigger/design-agent.ts index 018872d..e74541d 100644 --- a/src/trigger/design-agent.ts +++ b/src/trigger/design-agent.ts @@ -1,48 +1,26 @@ -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 type { LanguageModel } from "ai"; +import type { DesignAgentPayload, DesignAgentResult, DesignBrief, DesignPlan } from "@/lib/ai/agent-schema"; import { - buildCanvasGraph, - designGraphSchema, - DESIGN_AGENT_STAGE_KEY, - type DesignAgentStage, -} from "@/lib/design-generation"; + analyzeTurn, + draftPlan, + emptyBrief, + generateDesignGraph, + resolveDesignModel, +} from "@/lib/ai/design-agent-engine"; +import { buildCanvasGraph, 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; -} +export type { DesignAgentPayload, DesignAgentResult }; /** 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.", - "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.", - "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 setStage(stage: DesignAgentStage) { + metadata.set(DESIGN_AGENT_STAGE_KEY, stage); } function getNodeHeight(node: CanvasNode): number { @@ -75,51 +53,89 @@ function resolveOrigin(existingNodes: readonly CanvasNode[]): { x: number; y: nu return { x: left, y: bottom + EXISTING_CONTENT_GAP }; } -export const designAgentTask = task({ - id: "design-agent", - 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, +/** Generates the diagram for an approved plan and writes it into the room. */ +async function drawPlan( + model: LanguageModel, + roomId: string, + runId: string, + brief: DesignBrief, + plan: DesignPlan, +): 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 }); + + setStage("writing"); + let counts = { nodeCount: 0, edgeCount: 0 }; + + await mutateFlow({ client: getLiveblocksClient(), roomId }, (flow) => { + const { nodes, edges } = buildCanvasGraph(design, { + idPrefix: runId, + origin: resolveOrigin(flow.nodes), }); - metadata.set(DESIGN_AGENT_STAGE_KEY, "writing" satisfies DesignAgentStage); + flow.addNodes(nodes); + flow.addEdges(edges); - const client = getLiveblocksClient(); - let result: DesignAgentResult = { nodeCount: 0, edgeCount: 0 }; + counts = { nodeCount: nodes.length, edgeCount: edges.length }; + }); - 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 }); + logger.log("Design written to room", { roomId, ...counts }); + return counts; +} - return result; +/** + * 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. + */ +export const designAgentTask = task({ + id: "design-agent", + // Each model call already retries transient errors with backoff. Retrying the + // 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. + maxDuration: 300, + run: async (payload: DesignAgentPayload, { ctx }): Promise => { + logger.log("Design agent turn", { roomId: payload.roomId, intent: payload.intent }); + const model = resolveDesignModel(); + + // 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); + setStage("done"); + return { action: "generated", reply: "", ...counts, decisions: payload.plan.decisions, brief }; + } + + setStage("analyzing"); + const analysis = await analyzeTurn(model, payload); + logger.log("Turn analyzed", { decision: analysis.decision, questions: analysis.questions.length }); + + if (analysis.decision === "ask") { + setStage("done"); + return { action: "ask", reply: analysis.reply, questions: analysis.questions, brief: analysis.brief }; + } + + if (analysis.decision === "generate" && payload.plan) { + const counts = await drawPlan(model, payload.roomId, ctx.run.id, analysis.brief, payload.plan); + setStage("done"); + return { + action: "generated", + reply: analysis.reply, + ...counts, + decisions: payload.plan.decisions, + brief: analysis.brief, + }; + } + + setStage("planning"); + const plan = await draftPlan(model, analysis.brief, payload.plan, payload.input); + logger.log("Plan drafted", { components: plan.components.length, decisions: plan.decisions.length }); + + setStage("done"); + return { action: "plan", reply: analysis.reply, plan, brief: analysis.brief }; }, });