From cd7af48bf7eb849bdd991ddcec9b87cf49ebb3a6 Mon Sep 17 00:00:00 2001 From: AdarshVMore Date: Sat, 27 Jun 2026 10:12:04 +0530 Subject: [PATCH] API key addition and usage feature implemented --- .../controllers/aiUsage.controller.ts | 140 +++++++++++++++++ Server/app/api-service/index.ts | 2 + .../app/api-service/routes/aiUsage.routes.ts | 11 ++ Server/app/webhook-service/queue/producer.ts | 0 .../webhook-service/routes/github.webhook.ts | 16 +- .../app/worker-service/services/ai.service.ts | 2 +- .../worker-service/services/review.service.ts | 141 ++++++++++++++---- Server/package/ai/gif.ts | 13 +- Server/package/ai/review.ts | 100 +++++++++---- Server/package/lib/apiKey.service.ts | 68 +++++++++ Server/package/lib/encryption.ts | 31 ++++ .../migration.sql | 43 ++++++ Server/prisma/schema.prisma | 55 +++++-- client/src/app/(dashboard)/ai-usage/page.tsx | 25 +++- client/src/lib/api/aiUsage.ts | 9 ++ package-lock.json | 6 + 16 files changed, 577 insertions(+), 85 deletions(-) create mode 100644 Server/app/api-service/controllers/aiUsage.controller.ts create mode 100644 Server/app/api-service/routes/aiUsage.routes.ts delete mode 100644 Server/app/webhook-service/queue/producer.ts create mode 100644 Server/package/lib/apiKey.service.ts create mode 100644 Server/package/lib/encryption.ts create mode 100644 Server/prisma/migrations/20260627000000_add_api_keys_and_token_usage/migration.sql create mode 100644 package-lock.json diff --git a/Server/app/api-service/controllers/aiUsage.controller.ts b/Server/app/api-service/controllers/aiUsage.controller.ts new file mode 100644 index 0000000..42b5336 --- /dev/null +++ b/Server/app/api-service/controllers/aiUsage.controller.ts @@ -0,0 +1,140 @@ +import { Request, Response } from "express"; +import { db } from "../../../package/db/prisma.js"; +import { encrypt, maskApiKey } from "../../../package/lib/encryption.js"; +import { PLATFORM_FREE_REVIEWS } from "../../../package/lib/apiKey.service.js"; + +async function getDbUserId(req: Request): Promise { + const githubId = (req as any).githubUser?.id; + if (!githubId) return null; + const user = await db.user.findUnique({ where: { githubId } }); + return user?.id ?? null; +} + +function startOfMonth(): Date { + const now = new Date(); + return new Date(now.getFullYear(), now.getMonth(), 1); +} + +async function buildUsageSummary(userId: string) { + const monthStart = startOfMonth(); + + const [monthlyUsage, reviewCount, dailyUsage, user] = await Promise.all([ + db.tokenUsage.aggregate({ + where: { userId, createdAt: { gte: monthStart } }, + _sum: { totalTokens: true }, + }), + db.tokenUsage.count({ + where: { userId, reviewId: { not: null }, createdAt: { gte: monthStart } }, + }), + db.tokenUsage.findMany({ + where: { userId, createdAt: { gte: monthStart } }, + select: { totalTokens: true, createdAt: true }, + orderBy: { createdAt: "asc" }, + }), + db.user.findUnique({ + where: { id: userId }, + select: { platformReviewsUsed: true }, + }), + ]); + + const totalTokensThisMonth = monthlyUsage._sum.totalTokens ?? 0; + const avgTokensPerPR = + reviewCount > 0 ? Math.round(totalTokensThisMonth / reviewCount) : 0; + + const byDay = new Map(); + for (const row of dailyUsage) { + const date = row.createdAt.toISOString().slice(0, 10); + byDay.set(date, (byDay.get(date) ?? 0) + row.totalTokens); + } + + const usageOverTime = Array.from(byDay.entries()).map(([date, tokens]) => ({ + date, + tokens, + })); + + const platformReviewsRemaining = Math.max( + 0, + PLATFORM_FREE_REVIEWS - (user?.platformReviewsUsed ?? 0), + ); + + return { + totalTokensThisMonth, + avgTokensPerPR, + usageOverTime, + platformReviewsUsed: user?.platformReviewsUsed ?? 0, + platformReviewsRemaining, + platformFreeLimit: PLATFORM_FREE_REVIEWS, + }; +} + +export async function getAIUsage(req: Request, res: Response) { + const userId = await getDbUserId(req); + if (!userId) { + res.status(401).json({ error: "unauthorized" }); + return; + } + + const [keys, usage] = await Promise.all([ + db.userApiKey.findMany({ + where: { userId }, + orderBy: { createdAt: "desc" }, + select: { id: true, name: true, maskedKey: true, createdAt: true }, + }), + buildUsageSummary(userId), + ]); + + res.json({ keys, usage }); +} + +export async function addApiKey(req: Request, res: Response) { + const userId = await getDbUserId(req); + if (!userId) { + res.status(401).json({ error: "unauthorized" }); + return; + } + + const { name, key } = req.body as { name?: string; key?: string }; + if (!name?.trim() || !key?.trim()) { + res.status(400).json({ error: "name and key are required" }); + return; + } + + if (!key.startsWith("sk-ant-")) { + res.status(400).json({ error: "invalid Anthropic API key format" }); + return; + } + + const saved = await db.userApiKey.create({ + data: { + userId, + name: name.trim(), + encryptedKey: encrypt(key.trim()), + maskedKey: maskApiKey(key.trim()), + }, + select: { id: true, name: true, maskedKey: true, createdAt: true }, + }); + + res.status(201).json(saved); +} + +export async function deleteApiKey(req: Request, res: Response) { + const userId = await getDbUserId(req); + if (!userId) { + res.status(401).json({ error: "unauthorized" }); + return; + } + + const { id } = req.params; + if (!id || Array.isArray(id)) { + res.status(400).json({ error: "key id required" }); + return; + } + const existing = await db.userApiKey.findFirst({ where: { id, userId } }); + if (!existing) { + res.status(404).json({ error: "key not found" }); + return; + } + + await db.userApiKey.delete({ where: { id } }); + res.status(204).send(); +} diff --git a/Server/app/api-service/index.ts b/Server/app/api-service/index.ts index 22e3dbf..39a3c60 100644 --- a/Server/app/api-service/index.ts +++ b/Server/app/api-service/index.ts @@ -7,6 +7,7 @@ import repoRoutes from "./routes/repos.route.js" import prRoutes from "./routes/pr.routes.js" import userRoutes from "./routes/user.routes.js" import deploymentRoutes from "./routes/deployment.routes.js" +import aiUsageRoutes from "./routes/aiUsage.routes.js" const app = express() @@ -25,5 +26,6 @@ app.use("/api/repo/", repoRoutes) app.use("/api/pr/", prRoutes) app.use("/api/user/", userRoutes) app.use("/api/deployment/", deploymentRoutes) +app.use("/api/ai-usage/", aiUsageRoutes) app.listen(3001, () => console.log("API service listening on port 3001")) \ No newline at end of file diff --git a/Server/app/api-service/routes/aiUsage.routes.ts b/Server/app/api-service/routes/aiUsage.routes.ts new file mode 100644 index 0000000..ec2f1b6 --- /dev/null +++ b/Server/app/api-service/routes/aiUsage.routes.ts @@ -0,0 +1,11 @@ +import { Router } from "express"; +import { requireAuth } from "../auth/auth.js"; +import { addApiKey, deleteApiKey, getAIUsage } from "../controllers/aiUsage.controller.js"; + +const router = Router(); + +router.get("/", requireAuth, getAIUsage); +router.post("/keys", requireAuth, addApiKey); +router.delete("/keys/:id", requireAuth, deleteApiKey); + +export default router; diff --git a/Server/app/webhook-service/queue/producer.ts b/Server/app/webhook-service/queue/producer.ts deleted file mode 100644 index e69de29..0000000 diff --git a/Server/app/webhook-service/routes/github.webhook.ts b/Server/app/webhook-service/routes/github.webhook.ts index e74b8b2..83ef44c 100644 --- a/Server/app/webhook-service/routes/github.webhook.ts +++ b/Server/app/webhook-service/routes/github.webhook.ts @@ -5,7 +5,6 @@ import { db } from "../../../package/db/prisma.js"; import { getRedisConnection } from "../../../package/lib/redis.client.js"; - const router = Router(); const GITHUB_SECRET = process.env.GITHUB_SECRET as string; @@ -100,9 +99,22 @@ if (!client.isOpen){ create: { owner, name: repo, installationId, userId: user.id }, }); + const repository = await db.repository.findUnique({ + where: { owner_name: { owner, name: repo } }, + select: { userId: true }, + }); + const pushed = await client.lPush( "reviewQueue", - JSON.stringify({ installationId, owner, repo, prNumber, prTitle: payload.pull_request.title, prAuthor: payload.pull_request.user.login }) + JSON.stringify({ + installationId, + owner, + repo, + prNumber, + prTitle: payload.pull_request.title, + prAuthor: payload.pull_request.user.login, + userId: repository?.userId ?? user.id, + }), ); console.log("pushed to queue, queue length:", pushed); } diff --git a/Server/app/worker-service/services/ai.service.ts b/Server/app/worker-service/services/ai.service.ts index ec5cd3c..69d41bc 100644 --- a/Server/app/worker-service/services/ai.service.ts +++ b/Server/app/worker-service/services/ai.service.ts @@ -1 +1 @@ -export { getAIReview, reviewPrompt, parseAIResponse, getRevieType, getCodeDiff, generateRelevantSearchQuery } from "../../../package/ai/review.js"; +export { getAIReview, reviewPrompt, parseAIResponse, getRevieType, getCodeDiff, generateRelevantSearchQuery, createTokenAccumulator } from "../../../package/ai/review.js"; diff --git a/Server/app/worker-service/services/review.service.ts b/Server/app/worker-service/services/review.service.ts index abbe728..68cb965 100644 --- a/Server/app/worker-service/services/review.service.ts +++ b/Server/app/worker-service/services/review.service.ts @@ -3,11 +3,31 @@ import { getDifferenceData, getReviewRules, } from "./github.service.js"; -import { getAIReview, reviewPrompt, parseAIResponse, getRevieType, getCodeDiff, generateRelevantSearchQuery } from "./ai.service.js"; +import { + getAIReview, + reviewPrompt, + parseAIResponse, + getRevieType, + getCodeDiff, + generateRelevantSearchQuery, + createTokenAccumulator, +} from "./ai.service.js"; import { db } from "../../../package/db/prisma.js"; -import {vectorDBExists, createEmbeddings, saveToVectorDB, searchRelevantEmbeddings, toIndexName} from "./rag.service.js" +import { + vectorDBExists, + createEmbeddings, + saveToVectorDB, + searchRelevantEmbeddings, + toIndexName, +} from "./rag.service.js"; import { findGIF } from "./gif.service.js"; import { getGifName } from "../../../package/ai/gif.js"; +import { + resolveReviewApiKey, + incrementPlatformReviewCount, + recordTokenUsage, + PLATFORM_FREE_REVIEWS, +} from "../../../package/lib/apiKey.service.js"; export type PRReviewJobData = { installationId: number; @@ -15,7 +35,8 @@ export type PRReviewJobData = { repo: string; prNumber: number; prTitle?: string; - prAuthor?:string; + prAuthor?: string; + userId?: string; }; export function formatReviewComment(parsed: any, gifURL?: string | null): string { @@ -30,8 +51,8 @@ export function formatReviewComment(parsed: any, gifURL?: string | null): string let md = `## Code Review\n\n`; - if(gifURL){ - md+= `![review-gif](${gifURL})\n\n` + if (gifURL) { + md += `![review-gif](${gifURL})\n\n`; } md += `${parsed.summary}\n\n`; @@ -62,7 +83,7 @@ ${arr - Category: ${issue.category} - Issue: ${issue.problem} - Suggestion: ${issue.fix} -` +`, ) .join("\n")} `; @@ -78,9 +99,21 @@ ${arr return md; } -async function getReviewGifUrl(summary?: string | null): Promise { +function formatQuotaExceededComment(): string { + return `## Code Review Unavailable + +You've used all **${PLATFORM_FREE_REVIEWS} free platform reviews**. To continue receiving AI code reviews, add your own Claude API key in the [CodeRefyn dashboard](https://custom-ai-code-reviewer.vercel.app/ai-usage). + +Once configured, all reviews will use your key and token usage will appear on the AI Usage page.`; +} + +async function getReviewGifUrl( + summary: string | null | undefined, + apiKey: string, + usage: ReturnType, +): Promise { try { - const gifQuery = await getGifName(summary); + const gifQuery = await getGifName(summary, apiKey, usage); return await findGIF(gifQuery); } catch (error) { console.error("Skipping review GIF", error); @@ -88,46 +121,80 @@ async function getReviewGifUrl(summary?: string | null): Promise } } +async function resolveUserId( + owner: string, + repo: string, + userId?: string, +): Promise { + if (userId) return userId; + + const repository = await db.repository.findUnique({ + where: { owner_name: { owner, name: repo } }, + select: { userId: true }, + }); + return repository?.userId ?? null; +} + export async function runPRReview(data: PRReviewJobData) { const { installationId, owner, repo, prNumber, prTitle, prAuthor } = data; - let dbName = "" + const dbName = ""; - // const isGIFEnabled = + const createVectorDB = (await vectorDBExists(dbName)) as boolean; - const createVectorDB = await vectorDBExists(dbName) as boolean + if (createVectorDB) { + const values = {}; - if(createVectorDB){ - let values = {} - - const embeddings = await createEmbeddings(values) as any - const finalSaved = await saveToVectorDB(embeddings) - console.log(finalSaved) + const embeddings = (await createEmbeddings(values)) as any; + const finalSaved = await saveToVectorDB(embeddings); + console.log(finalSaved); } const octokit = await getOctokit(installationId); + const userId = await resolveUserId(owner, repo, data.userId); + + if (!userId) { + console.log(`No user found for ${owner}/${repo} — skipping review`); + return; + } + + const resolvedKey = await resolveReviewApiKey(userId); + if (!resolvedKey) { + console.log(`User ${userId} has no API key and free reviews exhausted`); + await octokit.issues.createComment({ + owner, + repo, + issue_number: prNumber, + body: formatQuotaExceededComment(), + }); + return; + } + + const { apiKey, source } = resolvedKey; + const tokenUsage = createTokenAccumulator(); + const difference = await getDifferenceData(octokit, owner, repo, prNumber); - const { data: collaborators } = await octokit.rest.repos.listCollaborators({ + await octokit.rest.repos.listCollaborators({ owner, repo, }); const rules = await getReviewRules(octokit, owner, repo, prNumber); - const reviewType = await getRevieType(difference) as any - const indexName = toIndexName(owner, repo) - const DBExsists = await vectorDBExists(indexName) - let relevantCode:any - const processedDiff = await getCodeDiff(difference) as any - if(reviewType === "feature" && DBExsists) { - const searchQuery = await generateRelevantSearchQuery(processedDiff) + const reviewType = (await getRevieType(difference, apiKey, tokenUsage)) as any; + const indexName = toIndexName(owner, repo); + const DBExsists = await vectorDBExists(indexName); + let relevantCode: any; + const processedDiff = (await getCodeDiff(difference)) as any; + if (reviewType === "feature" && DBExsists) { + const searchQuery = await generateRelevantSearchQuery(processedDiff, apiKey, tokenUsage); const query = { text: searchQuery, indexName: indexName, topK: 5, - } - relevantCode = await searchRelevantEmbeddings(query) + }; + relevantCode = await searchRelevantEmbeddings(query); } const prompt = reviewPrompt(difference, rules, relevantCode); - const aiResponse = await getAIReview(prompt); + const aiResponse = await getAIReview(prompt, apiKey, tokenUsage); const cleanText = aiResponse .replace(/```json/g, "") .replace(/```/g, "") @@ -138,7 +205,7 @@ export async function runPRReview(data: PRReviewJobData) { console.log("Failed to parse AI response"); return; } - const gifUrl = await getReviewGifUrl(parsedPrompt.summary); + const gifUrl = await getReviewGifUrl(parsedPrompt.summary, apiKey, tokenUsage); const commentBody = formatReviewComment(parsedPrompt, gifUrl); await octokit.issues.createComment({ @@ -186,6 +253,20 @@ export async function runPRReview(data: PRReviewJobData) { }); } - console.log(`review saved to DB for ${owner}/${repo}#${prNumber}`); + await recordTokenUsage({ + userId, + reviewId: savedReview.id, + inputTokens: tokenUsage.inputTokens, + outputTokens: tokenUsage.outputTokens, + source, + }); + + if (source === "platform") { + await incrementPlatformReviewCount(userId); + } + + console.log( + `review saved to DB for ${owner}/${repo}#${prNumber} (${tokenUsage.inputTokens + tokenUsage.outputTokens} tokens, ${source})`, + ); return aiResponse; } diff --git a/Server/package/ai/gif.ts b/Server/package/ai/gif.ts index 661eadd..5e1e5b6 100644 --- a/Server/package/ai/gif.ts +++ b/Server/package/ai/gif.ts @@ -1,15 +1,17 @@ import Anthropic from "@anthropic-ai/sdk"; +import { addUsage, createTokenAccumulator, type TokenAccumulator } from "./review.js"; -const anthropic = new Anthropic({ - apiKey: process.env.ANTHROPIC_API_KEY!, -}); - -export async function getGifName(summary?: string | null): Promise { +export async function getGifName( + summary: string | null | undefined, + apiKey: string, + usage?: TokenAccumulator, +): Promise { const fallback = "code review"; const cleanSummary = summary?.trim(); if (!cleanSummary) return fallback; + const anthropic = new Anthropic({ apiKey }); const res = await anthropic.messages.create({ model: "claude-sonnet-4-5", max_tokens: 40, @@ -23,6 +25,7 @@ export async function getGifName(summary?: string | null): Promise { }, ], }); + addUsage(usage ?? createTokenAccumulator(), res.usage); const block = res.content.find((b) => b.type === "text"); const text = block && block.type === "text" ? block.text.trim() : ""; diff --git a/Server/package/ai/review.ts b/Server/package/ai/review.ts index 40876b8..156462e 100644 --- a/Server/package/ai/review.ts +++ b/Server/package/ai/review.ts @@ -1,10 +1,26 @@ import Anthropic from "@anthropic-ai/sdk"; -const anthropic = new Anthropic({ - apiKey: process.env.ANTHROPIC_API_KEY!, -}); +export type TokenAccumulator = { + inputTokens: number; + outputTokens: number; +}; -export async function getRevieType(diff: string) { +export function createTokenAccumulator(): TokenAccumulator { + return { inputTokens: 0, outputTokens: 0 }; +} + +export function addUsage(acc: TokenAccumulator, usage?: Anthropic.Messages.Usage | null) { + if (!usage) return; + acc.inputTokens += usage.input_tokens ?? 0; + acc.outputTokens += usage.output_tokens ?? 0; +} + +function createClient(apiKey: string) { + return new Anthropic({ apiKey }); +} + +export async function getRevieType(diff: string, apiKey: string, usage?: TokenAccumulator) { + const anthropic = createClient(apiKey); const prompt = `You are supposed to see the Code Diff ${diff} and identify the PR Type if it is a Bug Fix / New Feature / Code Update , etc and only return the name of the type in response . "Return ONLY: feature | bugfix | refactor"`; const res = await anthropic.messages.create({ model: "claude-sonnet-4-5", @@ -17,8 +33,9 @@ export async function getRevieType(diff: string) { }, ], }); - const block = res.content.find((b) => b.type === "text") - return block && block.type === "text" ? block.text.trim().toLowerCase() : "" + addUsage(usage ?? createTokenAccumulator(), res.usage); + const block = res.content.find((b) => b.type === "text"); + return block && block.type === "text" ? block.text.trim().toLowerCase() : ""; } export async function getCodeDiff(diff: string) { @@ -29,31 +46,42 @@ export async function getCodeDiff(diff: string) { .slice(0, 500); } -export async function generateRelevantSearchQuery(processedDiff: string) { +export async function generateRelevantSearchQuery( + processedDiff: string, + apiKey: string, + usage?: TokenAccumulator, +) { + const anthropic = createClient(apiKey); const res = await anthropic.messages.create({ - model: "claude-sonnet-4-5", - max_tokens: 1500, - temperature: 0, - system: "Convert the given git diff into a short semantic search query to find relevant code in a repository Return ONLY a short phrase.", - messages: [ - { - role: "user", - content: processedDiff, - }, - ], + model: "claude-sonnet-4-5", + max_tokens: 1500, + temperature: 0, + system: + "Convert the given git diff into a short semantic search query to find relevant code in a repository Return ONLY a short phrase.", + messages: [ + { + role: "user", + content: processedDiff, + }, + ], }); - - const block = res.content.find((b) => b.type === "text") - return block && block.type === "text" ? block.text.trim() : "" + addUsage(usage ?? createTokenAccumulator(), res.usage); + const block = res.content.find((b) => b.type === "text"); + return block && block.type === "text" ? block.text.trim() : ""; } -export async function getAIReview(prompt: string | null, retries = 4) { +export async function getAIReview( + prompt: string | null, + apiKey: string, + usage?: TokenAccumulator, + retries = 4, +) { for (let attempt = 0; attempt <= retries; attempt++) { try { - return await _callAI(prompt); + return await _callAI(prompt, apiKey, usage); } catch (err: any) { if (err?.status === 529 && attempt < retries) { - const delay = Math.pow(2, attempt) * 1000; // 1s, 2s, 4s, 8s + const delay = Math.pow(2, attempt) * 1000; console.log( `Anthropic overloaded (529), retrying in ${delay / 1000}s... (attempt ${attempt + 1}/${retries})`, ); @@ -66,7 +94,8 @@ export async function getAIReview(prompt: string | null, retries = 4) { throw new Error("Max retries exceeded for AI review"); } -async function _callAI(prompt: string | null) { +async function _callAI(prompt: string | null, apiKey: string, usage?: TokenAccumulator) { + const anthropic = createClient(apiKey); const res = await anthropic.messages.create({ model: "claude-sonnet-4-5", max_tokens: 1500, @@ -136,6 +165,7 @@ Do NOT include markdown, explanations, or text outside JSON.`, }, ], }); + addUsage(usage ?? createTokenAccumulator(), res.usage); const block = res.content.find((b) => b.type === "text"); if (!block || block.type !== "text") @@ -144,24 +174,28 @@ Do NOT include markdown, explanations, or text outside JSON.`, } type RelevantCodeMatch = { - id: string - score: number - content: string - filePath: string -} + id: string; + score: number; + content: string; + filePath: string; +}; function formatRelevantCode(matches: RelevantCodeMatch[]): string { - if (!matches.length) return "" + if (!matches.length) return ""; return matches .map((m) => `// ${m.filePath}\n${m.content}`) - .join("\n\n---\n\n") + .join("\n\n---\n\n"); } -export function reviewPrompt(diff: string, rules: any, relevantCode?: RelevantCodeMatch[] | string) { +export function reviewPrompt( + diff: string, + rules: any, + relevantCode?: RelevantCodeMatch[] | string, +) { const relevantSection = Array.isArray(relevantCode) && relevantCode.length > 0 ? `\nEXISTING CODEBASE CONTEXT (retrieved for this feature):\nUse this to check for consistency, duplication, or conflicts with existing patterns.\n\n${formatRelevantCode(relevantCode)}\n` - : "" + : ""; return ` REPOSITORY RULES (HIGHEST PRIORITY): diff --git a/Server/package/lib/apiKey.service.ts b/Server/package/lib/apiKey.service.ts new file mode 100644 index 0000000..645992c --- /dev/null +++ b/Server/package/lib/apiKey.service.ts @@ -0,0 +1,68 @@ +import { db } from "../db/prisma.js"; +import { decrypt } from "./encryption.js"; + +export const PLATFORM_FREE_REVIEWS = 5; + +export type ResolvedApiKey = { + apiKey: string; + source: "user_key" | "platform"; +}; + +export async function getActiveUserApiKey(userId: string): Promise { + const record = await db.userApiKey.findFirst({ + where: { userId }, + orderBy: { createdAt: "desc" }, + }); + if (!record) return null; + return decrypt(record.encryptedKey); +} + +export async function resolveReviewApiKey(userId: string): Promise { + const userKey = await getActiveUserApiKey(userId); + if (userKey) { + return { apiKey: userKey, source: "user_key" }; + } + + const user = await db.user.findUnique({ + where: { id: userId }, + select: { platformReviewsUsed: true }, + }); + if (!user || user.platformReviewsUsed >= PLATFORM_FREE_REVIEWS) { + return null; + } + + const platformKey = process.env.ANTHROPIC_API_KEY; + if (!platformKey) { + throw new Error("Platform ANTHROPIC_API_KEY is not configured"); + } + + return { apiKey: platformKey, source: "platform" }; +} + +export async function incrementPlatformReviewCount(userId: string): Promise { + await db.user.update({ + where: { id: userId }, + data: { platformReviewsUsed: { increment: 1 } }, + }); +} + +export type TokenUsageInput = { + userId: string; + reviewId?: string; + inputTokens: number; + outputTokens: number; + source: "user_key" | "platform"; +}; + +export async function recordTokenUsage(data: TokenUsageInput): Promise { + await db.tokenUsage.create({ + data: { + userId: data.userId, + reviewId: data.reviewId ?? null, + inputTokens: data.inputTokens, + outputTokens: data.outputTokens, + totalTokens: data.inputTokens + data.outputTokens, + source: data.source, + }, + }); +} diff --git a/Server/package/lib/encryption.ts b/Server/package/lib/encryption.ts new file mode 100644 index 0000000..efc8a03 --- /dev/null +++ b/Server/package/lib/encryption.ts @@ -0,0 +1,31 @@ +import crypto from "crypto"; + +const ALGORITHM = "aes-256-gcm"; + +function getKey(): Buffer { + const secret = process.env.ENCRYPTION_SECRET ?? "dev-encryption-secret-change-me"; + return crypto.scryptSync(secret, "coderefiyn-salt", 32); +} + +export function encrypt(text: string): string { + const iv = crypto.randomBytes(16); + const cipher = crypto.createCipheriv(ALGORITHM, getKey(), iv); + const encrypted = Buffer.concat([cipher.update(text, "utf8"), cipher.final()]); + const tag = cipher.getAuthTag(); + return Buffer.concat([iv, tag, encrypted]).toString("base64"); +} + +export function decrypt(data: string): string { + const buf = Buffer.from(data, "base64"); + const iv = buf.subarray(0, 16); + const tag = buf.subarray(16, 32); + const encrypted = buf.subarray(32); + const decipher = crypto.createDecipheriv(ALGORITHM, getKey(), iv); + decipher.setAuthTag(tag); + return Buffer.concat([decipher.update(encrypted), decipher.final()]).toString("utf8"); +} + +export function maskApiKey(key: string): string { + if (key.length <= 8) return "••••••••"; + return `${key.slice(0, 7)}••••••••${key.slice(-4)}`; +} diff --git a/Server/prisma/migrations/20260627000000_add_api_keys_and_token_usage/migration.sql b/Server/prisma/migrations/20260627000000_add_api_keys_and_token_usage/migration.sql new file mode 100644 index 0000000..72eaec7 --- /dev/null +++ b/Server/prisma/migrations/20260627000000_add_api_keys_and_token_usage/migration.sql @@ -0,0 +1,43 @@ +-- AlterTable +ALTER TABLE "User" ADD COLUMN "platformReviewsUsed" INTEGER NOT NULL DEFAULT 0; + +-- CreateTable +CREATE TABLE "UserApiKey" ( + "id" TEXT NOT NULL, + "userId" TEXT NOT NULL, + "name" TEXT NOT NULL, + "encryptedKey" TEXT NOT NULL, + "maskedKey" TEXT NOT NULL, + "createdAt" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + + CONSTRAINT "UserApiKey_pkey" PRIMARY KEY ("id") +); + +-- CreateTable +CREATE TABLE "TokenUsage" ( + "id" TEXT NOT NULL, + "userId" TEXT NOT NULL, + "reviewId" TEXT, + "inputTokens" INTEGER NOT NULL DEFAULT 0, + "outputTokens" INTEGER NOT NULL DEFAULT 0, + "totalTokens" INTEGER NOT NULL DEFAULT 0, + "source" TEXT NOT NULL DEFAULT 'user_key', + "createdAt" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + + CONSTRAINT "TokenUsage_pkey" PRIMARY KEY ("id") +); + +-- CreateIndex +CREATE INDEX "UserApiKey_userId_idx" ON "UserApiKey"("userId"); + +-- CreateIndex +CREATE INDEX "TokenUsage_userId_createdAt_idx" ON "TokenUsage"("userId", "createdAt"); + +-- AddForeignKey +ALTER TABLE "UserApiKey" ADD CONSTRAINT "UserApiKey_userId_fkey" FOREIGN KEY ("userId") REFERENCES "User"("id") ON DELETE CASCADE ON UPDATE CASCADE; + +-- AddForeignKey +ALTER TABLE "TokenUsage" ADD CONSTRAINT "TokenUsage_userId_fkey" FOREIGN KEY ("userId") REFERENCES "User"("id") ON DELETE CASCADE ON UPDATE CASCADE; + +-- AddForeignKey +ALTER TABLE "TokenUsage" ADD CONSTRAINT "TokenUsage_reviewId_fkey" FOREIGN KEY ("reviewId") REFERENCES "PRReview"("id") ON DELETE SET NULL ON UPDATE CASCADE; diff --git a/Server/prisma/schema.prisma b/Server/prisma/schema.prisma index cfda725..8999ee4 100644 --- a/Server/prisma/schema.prisma +++ b/Server/prisma/schema.prisma @@ -8,14 +8,44 @@ datasource db { } model User { - id String @id @default(cuid()) - githubId Int @unique - githubUsername String - githubAvatar String? - email String? - plan String @default("free") - createdAt DateTime @default(now()) - repos Repository[] + id String @id @default(cuid()) + githubId Int @unique + githubUsername String + githubAvatar String? + email String? + plan String @default("free") + platformReviewsUsed Int @default(0) + createdAt DateTime @default(now()) + repos Repository[] + apiKeys UserApiKey[] + tokenUsage TokenUsage[] +} + +model UserApiKey { + id String @id @default(cuid()) + userId String + name String + encryptedKey String + maskedKey String + createdAt DateTime @default(now()) + user User @relation(fields: [userId], references: [id], onDelete: Cascade) + + @@index([userId]) +} + +model TokenUsage { + id String @id @default(cuid()) + userId String + reviewId String? + inputTokens Int @default(0) + outputTokens Int @default(0) + totalTokens Int @default(0) + source String @default("user_key") + createdAt DateTime @default(now()) + user User @relation(fields: [userId], references: [id], onDelete: Cascade) + review PRReview? @relation(fields: [reviewId], references: [id], onDelete: SetNull) + + @@index([userId, createdAt]) } model Repository { @@ -33,18 +63,19 @@ model Repository { } model PRReview { - id String @id @default(cuid()) + id String @id @default(cuid()) prNumber Int prTitle String? summary String? repositoryId String - createdAt DateTime @default(now()) + createdAt DateTime @default(now()) score Int? author String commitSHA String - prAuthor String @default("") + prAuthor String @default("") issues Issue[] - repository Repository @relation(fields: [repositoryId], references: [id]) + tokenUsage TokenUsage[] + repository Repository @relation(fields: [repositoryId], references: [id]) } model Issue { diff --git a/client/src/app/(dashboard)/ai-usage/page.tsx b/client/src/app/(dashboard)/ai-usage/page.tsx index d68524c..5d82902 100644 --- a/client/src/app/(dashboard)/ai-usage/page.tsx +++ b/client/src/app/(dashboard)/ai-usage/page.tsx @@ -48,7 +48,7 @@ export default function AIUsagePage() { if (loading) return
if (error) return
{error}
- + async function onAdd(e: React.FormEvent) { e.preventDefault() if (!keyValue.trim() || !keyName.trim()) return @@ -82,12 +82,33 @@ export default function AIUsagePage() {
{/* Stats */} -
+
+ 0 ? 'Your key' : `${usage?.platformReviewsRemaining ?? 0} / ${usage?.platformFreeLimit ?? 5}`} + label={keys.length > 0 ? 'Billing source' : 'Free reviews left'} + trend={keys.length > 0 ? 'Reviews use your Claude API key' : 'Platform key until limit reached'} + trendDirection="neutral" + />
+ {keys.length === 0 && ( + +

+ No API key configured. You have{' '} + {usage?.platformReviewsRemaining ?? 0}{' '} + of{' '} + {usage?.platformFreeLimit ?? 5}{' '} + free platform reviews remaining. +

+

+ Add your Claude API key below for unlimited reviews billed to your account. +

+ + )} + {/* Usage chart */} Token usage over time diff --git a/client/src/lib/api/aiUsage.ts b/client/src/lib/api/aiUsage.ts index 3c3695d..c477dec 100644 --- a/client/src/lib/api/aiUsage.ts +++ b/client/src/lib/api/aiUsage.ts @@ -16,6 +16,9 @@ export type AIUsageSummary = { totalTokensThisMonth: number avgTokensPerPR: number usageOverTime: AIUsagePoint[] + platformReviewsUsed: number + platformReviewsRemaining: number + platformFreeLimit: number } export type AIUsageData = { @@ -29,6 +32,9 @@ export const emptyAIUsage: AIUsageData = { totalTokensThisMonth: 0, avgTokensPerPR: 0, usageOverTime: [], + platformReviewsUsed: 0, + platformReviewsRemaining: 5, + platformFreeLimit: 5, }, } @@ -40,6 +46,9 @@ export async function fetchAIUsage(): Promise { totalTokensThisMonth: data.usage?.totalTokensThisMonth ?? 0, avgTokensPerPR: data.usage?.avgTokensPerPR ?? 0, usageOverTime: data.usage?.usageOverTime ?? [], + platformReviewsUsed: data.usage?.platformReviewsUsed ?? 0, + platformReviewsRemaining: data.usage?.platformReviewsRemaining ?? 5, + platformFreeLimit: data.usage?.platformFreeLimit ?? 5, }, } } diff --git a/package-lock.json b/package-lock.json new file mode 100644 index 0000000..2ea8fd5 --- /dev/null +++ b/package-lock.json @@ -0,0 +1,6 @@ +{ + "name": "AI-Code-Reviewer", + "lockfileVersion": 3, + "requires": true, + "packages": {} +}