diff --git a/packages/ai/src/graph/code-review.ts b/packages/ai/src/graph/code-review.ts new file mode 100644 index 0000000..33133e5 --- /dev/null +++ b/packages/ai/src/graph/code-review.ts @@ -0,0 +1,346 @@ +import { generateObject, generateText, type LanguageModel } from "ai"; +import { z } from "zod"; + +import { + buildAnalyzeFilePrompt, + buildComposeReviewPrompt, + buildSummarizeChangesPrompt, + CODE_REVIEW_SYSTEM_PROMPT, + type FileChangeInput, +} from "../prompts/code-review"; +import { END, runGraph, type GraphDefinition, type RunResult } from "./runtime"; + +// --------------------------------------------------------------------------- +// Schemas +// --------------------------------------------------------------------------- + +const SeveritySchema = z.enum(["blocker", "major", "minor", "nit"]); +const CategorySchema = z.enum([ + "correctness", + "security", + "performance", + "readability", + "style", + "test-coverage", +]); + +const SummarySchema = z.object({ + summary: z.string().min(1), + affectedModules: z.array(z.string()).default([]), +}); + +const IssueSchema = z.object({ + title: z.string().min(1), + detail: z.string().min(1), + suggestion: z.string().optional(), + severity: SeveritySchema, + category: CategorySchema, + /** Optional 1-based line number in the new file. */ + line: z.number().int().positive().optional(), +}); + +const FileAnalysisSchema = z.object({ + issues: z.array(IssueSchema).default([]), +}); + +export type Severity = z.infer; +export type IssueCategory = z.infer; +export type Issue = z.infer & { file: string }; + +// --------------------------------------------------------------------------- +// Public input / output +// --------------------------------------------------------------------------- + +export type CodeReviewInput = + | { files: FileChangeInput[]; unifiedDiff?: never } + | { unifiedDiff: string; files?: never }; + +export type CodeReviewVerdict = + | "Approve" + | "Approve with comments" + | "Request changes" + | "Block"; + +export type CodeReviewOutput = { + summary: string; + affectedModules: string[]; + issues: Issue[]; + reviewMarkdown: string; + verdict: CodeReviewVerdict; +}; + +// --------------------------------------------------------------------------- +// Graph state +// --------------------------------------------------------------------------- + +type State = { + input: CodeReviewInput; + files: FileChangeInput[]; + summary: string; + affectedModules: string[]; + issues: Issue[]; + reviewMarkdown: string; + verdict: CodeReviewVerdict; + error?: string; +}; + +const NODES = [ + "parseInput", + "summarizeChanges", + "analyzeFiles", + "prioritize", + "composeReview", +] as const; +type NodeName = (typeof NODES)[number]; + +// --------------------------------------------------------------------------- +// Unified-diff parser (intentionally tolerant — best-effort path/status) +// --------------------------------------------------------------------------- + +export function parseUnifiedDiff(diff: string): FileChangeInput[] { + if (!diff.trim()) return []; + + // Split on the standard "diff --git a/ b/" header. The first chunk + // before any header is ignored. + const parts = diff.split(/^diff --git /m).slice(1); + + return parts.map((rawChunk) => { + const chunk = rawChunk.replace(/\r\n/g, "\n"); + const headerLineEnd = chunk.indexOf("\n"); + const headerLine = headerLineEnd >= 0 ? chunk.slice(0, headerLineEnd) : chunk; + const body = headerLineEnd >= 0 ? chunk.slice(headerLineEnd + 1) : ""; + + // Header: "a/ b/" + const headerMatch = headerLine.match(/^a\/(.+?)\s+b\/(.+)$/); + const oldPathFromHeader = headerMatch?.[1]; + const newPathFromHeader = headerMatch?.[2]; + + let status: FileChangeInput["status"] = "modified"; + if (/^new file mode/m.test(body)) status = "added"; + else if (/^deleted file mode/m.test(body)) status = "deleted"; + else if (/^rename from /m.test(body)) status = "renamed"; + + // Prefer the +++ b/ line when present (handles unusual quoting). + const plusMatch = body.match(/^\+\+\+ b\/(.+)$/m); + const minusMatch = body.match(/^--- a\/(.+)$/m); + const path = + (status === "deleted" ? minusMatch?.[1] : plusMatch?.[1]) ?? + newPathFromHeader ?? + oldPathFromHeader ?? + "unknown"; + + const oldPath = + status === "renamed" ? (minusMatch?.[1] ?? oldPathFromHeader) : undefined; + + return { + path, + status, + oldPath, + patch: body.trim(), + }; + }); +} + +// --------------------------------------------------------------------------- +// Prioritization (pure) +// --------------------------------------------------------------------------- + +const SEVERITY_RANK: Record = { + blocker: 0, + major: 1, + minor: 2, + nit: 3, +}; + +const CATEGORY_RANK: Record = { + correctness: 0, + security: 1, + performance: 2, + "test-coverage": 3, + readability: 4, + style: 5, +}; + +export function prioritizeIssues(issues: Issue[]): Issue[] { + // Stable sort by severity, then category. De-dupe by file+title (case-insensitive). + const seen = new Set(); + const deduped = issues.filter((i) => { + const key = `${i.file}::${i.title.toLowerCase().trim()}`; + if (seen.has(key)) return false; + seen.add(key); + return true; + }); + + return deduped.sort((a, b) => { + const s = SEVERITY_RANK[a.severity] - SEVERITY_RANK[b.severity]; + if (s !== 0) return s; + return CATEGORY_RANK[a.category] - CATEGORY_RANK[b.category]; + }); +} + +function deriveVerdict(issues: Issue[]): CodeReviewVerdict { + if (issues.some((i) => i.severity === "blocker")) return "Block"; + if (issues.some((i) => i.severity === "major")) return "Request changes"; + if (issues.length > 0) return "Approve with comments"; + return "Approve"; +} + +// --------------------------------------------------------------------------- +// Graph builder +// --------------------------------------------------------------------------- + +export type BuildCodeReviewGraphOptions = { + /** Language model to use for all nodes (e.g. selectOpenAIModel(provider)). */ + model: LanguageModel; + /** Cap how many files are analyzed per run. Default 25. */ + maxFiles?: number; + /** Per-file analysis concurrency. Default 4. */ + concurrency?: number; +}; + +export function buildCodeReviewGraph( + options: BuildCodeReviewGraphOptions, +): GraphDefinition { + const { model } = options; + const maxFiles = options.maxFiles ?? 25; + const concurrency = Math.max(1, options.concurrency ?? 4); + + return { + entry: "parseInput", + nodes: { + parseInput: (state) => { + const files = + state.input.files ?? + (state.input.unifiedDiff ? parseUnifiedDiff(state.input.unifiedDiff) : []); + const trimmed = files.slice(0, maxFiles); + if (trimmed.length === 0) { + return { + next: END, + patch: { + files: [], + error: "No file changes provided", + reviewMarkdown: + "# Code Review\n\nNo file changes were provided to review.", + verdict: "Approve", + }, + }; + } + return { next: "summarizeChanges", patch: { files: trimmed } }; + }, + + summarizeChanges: async (state, ctx) => { + const { object } = await generateObject({ + model, + schema: SummarySchema, + system: CODE_REVIEW_SYSTEM_PROMPT, + prompt: buildSummarizeChangesPrompt(state.files), + abortSignal: ctx.signal, + }); + return { + next: "analyzeFiles", + patch: { + summary: object.summary, + affectedModules: object.affectedModules, + }, + }; + }, + + analyzeFiles: async (state, ctx) => { + const issues: Issue[] = []; + const queue = [...state.files]; + + async function worker() { + while (queue.length > 0) { + const file = queue.shift(); + if (!file) return; + if (ctx.signal?.aborted) return; + const { object } = await generateObject({ + model, + schema: FileAnalysisSchema, + system: CODE_REVIEW_SYSTEM_PROMPT, + prompt: buildAnalyzeFilePrompt(file, state.summary), + abortSignal: ctx.signal, + }); + for (const issue of object.issues) { + issues.push({ ...issue, file: file.path }); + } + } + } + + const workers = Array.from({ length: Math.min(concurrency, state.files.length) }, () => + worker(), + ); + await Promise.all(workers); + + return { next: "prioritize", patch: { issues } }; + }, + + prioritize: (state) => { + const prioritized = prioritizeIssues(state.issues); + return { + next: "composeReview", + patch: { issues: prioritized, verdict: deriveVerdict(prioritized) }, + }; + }, + + composeReview: async (state, ctx) => { + const { text } = await generateText({ + model, + system: CODE_REVIEW_SYSTEM_PROMPT, + prompt: buildComposeReviewPrompt({ + summary: state.summary, + affectedModules: state.affectedModules, + prioritizedIssues: state.issues.map((i) => ({ + file: i.file, + severity: i.severity, + category: i.category, + title: i.title, + detail: i.detail, + suggestion: i.suggestion, + })), + }), + abortSignal: ctx.signal, + }); + return { next: END, patch: { reviewMarkdown: text } }; + }, + }, + }; +} + +// --------------------------------------------------------------------------- +// High-level convenience entry point +// --------------------------------------------------------------------------- + +export type RunCodeReviewOptions = BuildCodeReviewGraphOptions & { + signal?: AbortSignal; + onEvent?: Parameters[1]["onEvent"]; +}; + +export async function runCodeReview( + input: CodeReviewInput, + options: RunCodeReviewOptions, +): Promise, "steps" | "reason"> }> { + const graph = buildCodeReviewGraph(options); + const result = await runGraph(graph, { + initialState: { + input, + files: [], + summary: "", + affectedModules: [], + issues: [], + reviewMarkdown: "", + verdict: "Approve", + }, + signal: options.signal, + onEvent: options.onEvent, + }); + + return { + summary: result.state.summary, + affectedModules: result.state.affectedModules, + issues: result.state.issues, + reviewMarkdown: result.state.reviewMarkdown, + verdict: result.state.verdict, + run: { steps: result.steps, reason: result.reason }, + }; +} diff --git a/packages/ai/src/graph/runtime.ts b/packages/ai/src/graph/runtime.ts new file mode 100644 index 0000000..a73747e --- /dev/null +++ b/packages/ai/src/graph/runtime.ts @@ -0,0 +1,103 @@ +/** + * Tiny dependency-free graph runtime for AI workflows. + * + * A graph is a set of named nodes plus a start node. Each node receives the + * current state, performs work, and returns a `next` edge label plus an + * optional state patch that is shallow-merged into the running state. + * + * The special edge label `END` terminates the run. Nodes never throw to + * indicate flow control; failures should be returned as state and routed + * via an explicit edge. + */ + +export const END = "__end__" as const; +export type EndLabel = typeof END; + +export type NodeResult = { + next: E | EndLabel; + patch?: Partial; +}; + +export type GraphNode = ( + state: Readonly, + ctx: GraphContext, +) => Promise> | NodeResult; + +export type GraphContext = { + /** Current step number, starting at 0 for the entry node. */ + step: number; + /** Optional abort signal forwarded from the caller. */ + signal?: AbortSignal; + /** Per-run logger; defaults to a noop. */ + log: (event: GraphEvent) => void; +}; + +export type GraphEvent = + | { type: "node:start"; node: string; step: number } + | { type: "node:end"; node: string; step: number; next: string } + | { type: "run:end"; reason: "completed" | "max-steps" | "aborted" }; + +export type GraphDefinition = { + /** Map of node name -> node function. */ + nodes: Record>; + /** Entry node. Must exist in `nodes`. */ + entry: N; +}; + +export type RunOptions = { + /** Initial state passed to the entry node. */ + initialState: S; + /** Hard cap on node executions to prevent runaway loops. Default 32. */ + maxSteps?: number; + /** Optional abort signal forwarded into each node's context. */ + signal?: AbortSignal; + /** Optional event hook for tracing/inspection. */ + onEvent?: (event: GraphEvent) => void; +}; + +export type RunResult = { + state: S; + steps: number; + reason: "completed" | "max-steps" | "aborted"; +}; + +export async function runGraph( + graph: GraphDefinition, + options: RunOptions, +): Promise> { + const maxSteps = options.maxSteps ?? 32; + const log = options.onEvent ?? (() => {}); + + let state: S = { ...options.initialState }; + let current: N | EndLabel = graph.entry; + let step = 0; + + while (current !== END) { + if (options.signal?.aborted) { + log({ type: "run:end", reason: "aborted" }); + return { state, steps: step, reason: "aborted" }; + } + if (step >= maxSteps) { + log({ type: "run:end", reason: "max-steps" }); + return { state, steps: step, reason: "max-steps" }; + } + + const node = graph.nodes[current]; + if (!node) { + throw new Error(`Graph node "${String(current)}" is not defined`); + } + + log({ type: "node:start", node: current, step }); + const result = await node(state, { step, signal: options.signal, log }); + if (result.patch) { + state = { ...state, ...result.patch }; + } + log({ type: "node:end", node: current, step, next: String(result.next) }); + + current = result.next as N | EndLabel; + step += 1; + } + + log({ type: "run:end", reason: "completed" }); + return { state, steps: step, reason: "completed" }; +} diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index e7de3b4..ef3fdc6 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -19,3 +19,35 @@ export { type ActionPromptPair, type RewriteTone, } from "./actions"; +export { + CODE_REVIEW_SYSTEM_PROMPT, + buildAnalyzeFilePrompt, + buildComposeReviewPrompt, + buildSummarizeChangesPrompt, + type FileChangeInput, +} from "./prompts/code-review"; +export { + END, + runGraph, + type GraphContext, + type GraphDefinition, + type GraphEvent, + type GraphNode, + type NodeResult, + type RunOptions, + type RunResult, +} from "./graph/runtime"; +export { + buildCodeReviewGraph, + parseUnifiedDiff, + prioritizeIssues, + runCodeReview, + type BuildCodeReviewGraphOptions, + type CodeReviewInput, + type CodeReviewOutput, + type CodeReviewVerdict, + type Issue, + type IssueCategory, + type RunCodeReviewOptions, + type Severity, +} from "./graph/code-review"; diff --git a/packages/ai/src/prompts/code-review.ts b/packages/ai/src/prompts/code-review.ts new file mode 100644 index 0000000..dc671da --- /dev/null +++ b/packages/ai/src/prompts/code-review.ts @@ -0,0 +1,139 @@ +/** + * Prompts used by the code-review graph in `graph/code-review.ts`. + * + * The graph runs a small ensemble of focused steps; each step gets a tightly + * scoped user prompt while sharing one reviewer system prompt, so the model + * keeps a consistent voice and rubric across nodes. + */ + +export const CODE_REVIEW_SYSTEM_PROMPT = `You are a senior staff engineer performing a code review. + +Operating principles: +- Be precise. Cite the file and (when possible) the symbol or hunk you are referring to. +- Prefer correctness > security > performance > readability > style. +- Do not invent code that is not in the diff. If something is unclear, say so explicitly. +- When suggesting a change, show the smallest reasonable patch (a few lines) rather than rewriting whole files. +- Severity levels: "blocker" (must fix before merge), "major" (should fix), "minor" (nice to fix), "nit" (style/preference). +- Never fabricate test results, runtime behavior, or external API contracts.`; + +export type FileChangeInput = { + path: string; + status: "added" | "modified" | "deleted" | "renamed"; + oldPath?: string; + /** Unified diff for this file, or full file contents for added files. */ + patch: string; +}; + +function fenceForPath(path: string): string { + const ext = path.split(".").pop()?.toLowerCase() ?? ""; + switch (ext) { + case "ts": + case "tsx": + return "ts"; + case "js": + case "jsx": + return "js"; + case "py": + return "python"; + case "go": + return "go"; + case "rs": + return "rust"; + case "json": + return "json"; + case "md": + return "markdown"; + default: + return "diff"; + } +} + +export function buildSummarizeChangesPrompt(files: FileChangeInput[]): string { + const fileList = files + .map((f) => `- ${f.status.toUpperCase()} ${f.path}${f.oldPath ? ` (was ${f.oldPath})` : ""}`) + .join("\n"); + + const bodies = files + .map( + (f) => + `### ${f.path}\n\n\`\`\`${fenceForPath(f.path)}\n${f.patch.trim()}\n\`\`\``, + ) + .join("\n\n"); + + return `Summarize the intent of this change set in 2-4 sentences. Then list the most affected modules. + +Files: +${fileList} + +Diffs: +${bodies} + +Respond with: +1) "Summary:" — what the change does and why (inferred). +2) "Affected modules:" — bullet list of high-level areas (e.g., "auth/session", "tasks/api").`; +} + +export function buildAnalyzeFilePrompt(file: FileChangeInput, summary: string): string { + return `Review the following file change against the change-set summary below. Identify concrete issues only — do not list praise or generic advice. + +Change-set summary: +${summary} + +File: ${file.path} (${file.status}${file.oldPath ? `, was ${file.oldPath}` : ""}) + +\`\`\`${fenceForPath(file.path)} +${file.patch.trim()} +\`\`\` + +For each issue, choose exactly one category: "correctness", "security", "performance", "readability", "style", "test-coverage". +For each issue, choose exactly one severity: "blocker", "major", "minor", "nit". +If there are no issues, return an empty list.`; +} + +export function buildComposeReviewPrompt(input: { + summary: string; + affectedModules: string[]; + prioritizedIssues: Array<{ + file: string; + severity: string; + category: string; + title: string; + detail: string; + suggestion?: string; + }>; +}): string { + const issueLines = input.prioritizedIssues + .map( + (i, idx) => + `${idx + 1}. [${i.severity}/${i.category}] ${i.file} — ${i.title}\n ${i.detail}${i.suggestion ? `\n Suggested: ${i.suggestion}` : ""}`, + ) + .join("\n\n"); + + const modules = input.affectedModules.length + ? input.affectedModules.map((m) => `- ${m}`).join("\n") + : "- (none identified)"; + + return `Compose the final review as Markdown. Audience: the PR author. + +Required structure: +# Code Review + +## Summary +${input.summary} + +## Affected modules +${modules} + +## Findings +(For each finding below, render as a Markdown subsection with severity/category badges. +Group findings by severity in this order: blocker, major, minor, nit. Skip groups that are empty.) + +Findings to render: +${issueLines || "(none)"} + +## Verdict +End with one of: "Approve", "Approve with comments", "Request changes", or "Block". +Pick the strictest verdict justified by the findings (any blocker => "Block"; any major => "Request changes"). + +Do not add any sections beyond those listed.`; +}