654 lines
19 KiB
TypeScript
654 lines
19 KiB
TypeScript
import * as Sentry from "@sentry/node"
|
|
import { fileDiffsToMap, isDiffContentsWellFormed } from "../diff_utils.js"
|
|
import { type FileDiffContent } from "@codeflash-ai/code-suggester/build/src/types.js"
|
|
import { userNickname } from "../auth0-mgmt.js"
|
|
import {
|
|
addLabelToPullRequest,
|
|
getInstallationOctokitByOwner,
|
|
isUserCollaborator,
|
|
} from "../github/github-utils.js"
|
|
import {
|
|
buildBenchmarkInfo,
|
|
buildPrTitle,
|
|
buildResultDetails,
|
|
buildResultFooter,
|
|
buildResultHeader,
|
|
buildResultTestReport,
|
|
} from "../github/pr-changes-utils.js"
|
|
import { githubApp } from "../github/github-app.js"
|
|
import {
|
|
assignReviewer,
|
|
createNewBranchFromDiffContents,
|
|
createNewPullRequest,
|
|
PrCommentFields,
|
|
} from "../github/create-pr-from-diffcontents.js"
|
|
import { posthog } from "../analytics.js"
|
|
import { PrismaClient } from "@prisma/client"
|
|
import { Request, Response } from "express"
|
|
import {
|
|
createAppInstallation,
|
|
createRepositoryMember,
|
|
getAppInstallationByInstalltionId,
|
|
prisma,
|
|
upsertRepository,
|
|
} from "@codeflash-ai/common"
|
|
import { AuthorizedUserReq } from "types.js"
|
|
import {
|
|
requestApproval,
|
|
requiresApproval,
|
|
isQualityMonitoringRepo,
|
|
sendQualityMonitoringNotification,
|
|
} from "../github/optimization_approval.js"
|
|
import { registerRepositoryAndMember } from "./utils/github-repo-setup.js"
|
|
|
|
// Define a comprehensive interface for PR title and body generation
|
|
export interface PrContentBuilder {
|
|
buildResultHeader: typeof buildResultHeader
|
|
buildBenchmarkInfo: typeof buildBenchmarkInfo
|
|
buildResultDetails: typeof buildResultDetails
|
|
buildResultTestReport: typeof buildResultTestReport
|
|
buildResultFooter: typeof buildResultFooter
|
|
buildPrTitle: typeof buildPrTitle
|
|
}
|
|
|
|
// Create a standalone PR title and body function with injectable dependencies
|
|
export function createStandalonePRTitleAndBody(
|
|
prCommentFields: PrCommentFields,
|
|
existingTests: string,
|
|
generatedTests: string,
|
|
coverage_message: string,
|
|
newBranchName: string,
|
|
builder: PrContentBuilder,
|
|
): { title: string; body: string } {
|
|
const prCommentHeader = builder.buildResultHeader(prCommentFields)
|
|
|
|
const benchmarkInfo =
|
|
prCommentFields.benchmark_details && prCommentFields.benchmark_details.length > 0
|
|
? builder.buildBenchmarkInfo(prCommentFields)
|
|
: ""
|
|
|
|
const prCommentBody = builder.buildResultDetails(prCommentFields)
|
|
const prCommentTestReport = builder.buildResultTestReport(
|
|
prCommentFields,
|
|
existingTests,
|
|
generatedTests,
|
|
coverage_message,
|
|
)
|
|
const prCommentFooter = builder.buildResultFooter(newBranchName)
|
|
const title: string = builder.buildPrTitle(
|
|
prCommentFields.function_name,
|
|
prCommentFields.speedup_pct,
|
|
prCommentFields.speedup_x,
|
|
)
|
|
|
|
const body: string = benchmarkInfo
|
|
? `${prCommentHeader}\n${benchmarkInfo}\n${prCommentBody}\n${prCommentTestReport}\n${prCommentFooter}`
|
|
: `${prCommentHeader}\n${prCommentBody}\n${prCommentTestReport}\n${prCommentFooter}`
|
|
|
|
return { title, body }
|
|
}
|
|
|
|
// Create a wrapper function for createNewPullRequest to match test expectations
|
|
export async function createStandalonePullRequest(
|
|
installationOctokit: any,
|
|
owner: string,
|
|
repo: string,
|
|
title: string,
|
|
body: string,
|
|
newBranchName: string,
|
|
baseBranch: string,
|
|
) {
|
|
return await createNewPullRequest(
|
|
installationOctokit,
|
|
owner,
|
|
repo,
|
|
title,
|
|
body,
|
|
newBranchName,
|
|
baseBranch,
|
|
)
|
|
}
|
|
|
|
// Dependencies interface for easier testing
|
|
export interface CreatePrDependencies {
|
|
prisma: PrismaClient
|
|
userNickname: typeof userNickname
|
|
getInstallationOctokitByOwner: typeof getInstallationOctokitByOwner
|
|
isUserCollaborator: typeof isUserCollaborator
|
|
requiresApproval: typeof requiresApproval
|
|
requestApproval: typeof requestApproval
|
|
triggerCreatePr: typeof triggerCreatePr
|
|
posthog: typeof posthog
|
|
isDiffContentsWellFormed: typeof isDiffContentsWellFormed
|
|
githubApp: typeof githubApp
|
|
}
|
|
|
|
// Dependencies for triggerCreatePr
|
|
export interface TriggerCreatePrDependencies {
|
|
prisma: PrismaClient
|
|
fileDiffsToMap: typeof fileDiffsToMap
|
|
buildPrTitle: typeof buildPrTitle
|
|
createNewBranchFromDiffContents: typeof createNewBranchFromDiffContents
|
|
createStandalonePullRequest: typeof createStandalonePullRequest
|
|
addLabelToPullRequest: typeof addLabelToPullRequest
|
|
assignReviewer: typeof assignReviewer
|
|
posthog: typeof posthog
|
|
createStandalonePRTitleAndBody: typeof createStandalonePRTitleAndBody
|
|
prContentBuilder: PrContentBuilder
|
|
}
|
|
|
|
// Default dependencies
|
|
const defaultPrContentBuilder: PrContentBuilder = {
|
|
buildResultHeader,
|
|
buildBenchmarkInfo,
|
|
buildResultDetails,
|
|
buildResultTestReport,
|
|
buildResultFooter,
|
|
buildPrTitle,
|
|
}
|
|
|
|
let dependencies: CreatePrDependencies = {
|
|
prisma: new PrismaClient(),
|
|
userNickname,
|
|
getInstallationOctokitByOwner,
|
|
isUserCollaborator,
|
|
requiresApproval,
|
|
requestApproval,
|
|
triggerCreatePr,
|
|
posthog,
|
|
isDiffContentsWellFormed,
|
|
githubApp,
|
|
}
|
|
|
|
let triggerCreatePrDeps: TriggerCreatePrDependencies = {
|
|
prisma: new PrismaClient(),
|
|
fileDiffsToMap,
|
|
buildPrTitle,
|
|
createNewBranchFromDiffContents,
|
|
createStandalonePullRequest,
|
|
addLabelToPullRequest,
|
|
assignReviewer,
|
|
posthog,
|
|
createStandalonePRTitleAndBody,
|
|
prContentBuilder: defaultPrContentBuilder,
|
|
}
|
|
|
|
// For testing - allow dependency injection
|
|
export function setCreatePrDependencies(deps: Partial<CreatePrDependencies>) {
|
|
dependencies = { ...dependencies, ...deps }
|
|
}
|
|
|
|
export function resetCreatePrDependencies() {
|
|
dependencies = {
|
|
prisma: new PrismaClient(),
|
|
userNickname,
|
|
getInstallationOctokitByOwner,
|
|
isUserCollaborator,
|
|
requiresApproval,
|
|
requestApproval,
|
|
triggerCreatePr,
|
|
posthog,
|
|
isDiffContentsWellFormed,
|
|
githubApp,
|
|
}
|
|
}
|
|
|
|
export function setTriggerCreatePrDependencies(deps: Partial<TriggerCreatePrDependencies>) {
|
|
triggerCreatePrDeps = { ...triggerCreatePrDeps, ...deps }
|
|
}
|
|
|
|
export function resetTriggerCreatePrDependencies() {
|
|
triggerCreatePrDeps = {
|
|
prisma: new PrismaClient(),
|
|
fileDiffsToMap,
|
|
buildPrTitle,
|
|
createNewBranchFromDiffContents,
|
|
createStandalonePullRequest,
|
|
addLabelToPullRequest,
|
|
assignReviewer,
|
|
posthog,
|
|
createStandalonePRTitleAndBody,
|
|
prContentBuilder: defaultPrContentBuilder,
|
|
}
|
|
}
|
|
|
|
export async function createPr(req: Request, res: Response) {
|
|
try {
|
|
const {
|
|
owner,
|
|
repo,
|
|
baseBranch,
|
|
diffContents,
|
|
prCommentFields,
|
|
existingTests,
|
|
generatedTests,
|
|
coverage_message,
|
|
} = req.body
|
|
|
|
const userId = (req as any).userId
|
|
const traceId = req.body.traceId || ""
|
|
|
|
if (!repo || !owner || !baseBranch || !dependencies.isDiffContentsWellFormed(diffContents)) {
|
|
res.status(400).send("Missing or malformed fields")
|
|
return
|
|
}
|
|
|
|
const nickname: string | null = await dependencies.userNickname(userId)
|
|
if (nickname == null) {
|
|
res.status(401).json({ error: "Unauthorized" })
|
|
return
|
|
}
|
|
|
|
const installationOctokit = await dependencies.getInstallationOctokitByOwner(
|
|
dependencies.githubApp,
|
|
owner,
|
|
repo,
|
|
)
|
|
if (installationOctokit instanceof Error) {
|
|
res.status(401).send(installationOctokit.message)
|
|
return
|
|
}
|
|
|
|
const isCollaborator = await dependencies.isUserCollaborator(
|
|
installationOctokit,
|
|
owner,
|
|
repo,
|
|
nickname,
|
|
)
|
|
if (!isCollaborator) {
|
|
console.log(`${nickname} is not a collaborator on ${owner}/${repo}`)
|
|
res.status(401).json({ error: "Unauthorized" })
|
|
return
|
|
}
|
|
// TODO: Remove this background upsert logic after ensuring all old repositories have been saved.
|
|
registerRepositoryAndMember(owner, repo, nickname, userId, installationOctokit)
|
|
.then(() => console.log(`Background repo and member upsert completed for ${owner}/${repo}`))
|
|
.catch(err => {
|
|
console.error(`Error in background upsertRepoAndCreateMember:`, err)
|
|
Sentry.captureException(err)
|
|
})
|
|
// Now check if approval is required
|
|
if (traceId && dependencies.requiresApproval(owner, repo)) {
|
|
const optimization = await dependencies.prisma.optimization_features.findUnique({
|
|
where: { trace_id: traceId },
|
|
select: {
|
|
approval_required: true,
|
|
approval_status: true,
|
|
},
|
|
})
|
|
|
|
if (optimization?.approval_status === "rejected") {
|
|
return res.status(403).json({
|
|
status: "rejected",
|
|
message: "This optimization request was rejected.",
|
|
})
|
|
}
|
|
|
|
if (optimization?.approval_status === "approved") {
|
|
console.log(`Request ${traceId} was previously approved, continuing with PR creation`)
|
|
const prNumber = await dependencies.triggerCreatePr(
|
|
owner,
|
|
repo,
|
|
baseBranch,
|
|
diffContents,
|
|
prCommentFields,
|
|
existingTests,
|
|
generatedTests,
|
|
coverage_message,
|
|
userId,
|
|
nickname,
|
|
installationOctokit,
|
|
traceId,
|
|
)
|
|
|
|
if (prNumber > 0) {
|
|
return res.json(prNumber)
|
|
} else {
|
|
return res.status(500).send("Error creating pull request")
|
|
}
|
|
} else {
|
|
const requestData = {
|
|
type: "create-pr",
|
|
owner,
|
|
repo,
|
|
baseBranch,
|
|
diffContents,
|
|
prCommentFields,
|
|
existingTests,
|
|
generatedTests,
|
|
coverage_message,
|
|
userId,
|
|
}
|
|
|
|
await dependencies.requestApproval(
|
|
traceId,
|
|
owner,
|
|
repo,
|
|
prCommentFields.function_name,
|
|
userId,
|
|
requestData,
|
|
)
|
|
|
|
return res.status(202).json({
|
|
status: "pending_approval",
|
|
message:
|
|
"This optimization requires approval. You will be notified when it is processed.",
|
|
})
|
|
}
|
|
}
|
|
|
|
const prNumber = await dependencies.triggerCreatePr(
|
|
owner,
|
|
repo,
|
|
baseBranch,
|
|
diffContents,
|
|
prCommentFields,
|
|
existingTests,
|
|
generatedTests,
|
|
coverage_message,
|
|
userId,
|
|
nickname,
|
|
installationOctokit,
|
|
traceId,
|
|
)
|
|
|
|
if (prNumber > 0) {
|
|
// Check if this is a quality monitoring repo and send notification
|
|
if (traceId && isQualityMonitoringRepo(owner, repo)) {
|
|
const requestData = {
|
|
type: "create-pr",
|
|
owner,
|
|
repo,
|
|
baseBranch,
|
|
diffContents,
|
|
prCommentFields,
|
|
existingTests,
|
|
generatedTests,
|
|
coverage_message,
|
|
userId,
|
|
}
|
|
|
|
const prUrl = `https://github.com/${owner}/${repo}/pull/${prNumber}`
|
|
|
|
// Send quality monitoring notification (non-blocking)
|
|
sendQualityMonitoringNotification(
|
|
traceId,
|
|
owner,
|
|
repo,
|
|
prCommentFields.function_name,
|
|
userId,
|
|
requestData,
|
|
prUrl,
|
|
).catch(error => {
|
|
console.error(`Failed to send quality monitoring notification: ${error}`)
|
|
})
|
|
|
|
console.log(
|
|
`Quality monitoring notification triggered for ${owner}/${repo} PR #${prNumber}`,
|
|
)
|
|
}
|
|
|
|
return res.json(prNumber)
|
|
} else {
|
|
return res.status(500).send("Error creating pull request")
|
|
}
|
|
} catch (error) {
|
|
console.log(`Error in /cfapi/create-pr: ${error}`)
|
|
if (error instanceof Error) {
|
|
console.log(`Error message: ${error.message}`)
|
|
console.log(`Error stack: ${error.stack}`)
|
|
dependencies.posthog.capture({
|
|
distinctId: (req as any).userId,
|
|
event: `cfapi-create-pr-failed-error-creating-standalone-pr`,
|
|
properties: {
|
|
error: error.message,
|
|
},
|
|
})
|
|
res.status(500).send(`Error creating pull request: ${error.message}`)
|
|
} else {
|
|
res.status(500).send(`Error creating pull request`)
|
|
}
|
|
}
|
|
}
|
|
|
|
export async function triggerCreatePr(
|
|
owner,
|
|
repo,
|
|
baseBranch,
|
|
diffContents,
|
|
prCommentFields,
|
|
existingTests,
|
|
generatedTests,
|
|
coverage_message,
|
|
userId,
|
|
nickname,
|
|
installationOctokit,
|
|
traceId = "",
|
|
): Promise<number> {
|
|
try {
|
|
const diffContentsMap: Map<string, FileDiffContent> =
|
|
triggerCreatePrDeps.fileDiffsToMap(diffContents)
|
|
|
|
const timestamp = Date.now()
|
|
const encodedTimestamp = timestamp.toString(36)
|
|
|
|
const newBranchName = `codeflash/optimize-${prCommentFields.function_name}-${encodedTimestamp}`
|
|
const commitMessage =
|
|
triggerCreatePrDeps.buildPrTitle(
|
|
prCommentFields.function_name,
|
|
prCommentFields.speedup_pct,
|
|
prCommentFields.speedup_x,
|
|
) + `\n${prCommentFields.optimization_explanation}`
|
|
|
|
const branchCreated = await triggerCreatePrDeps.createNewBranchFromDiffContents(
|
|
installationOctokit,
|
|
owner,
|
|
repo,
|
|
newBranchName,
|
|
baseBranch,
|
|
diffContentsMap,
|
|
commitMessage,
|
|
)
|
|
|
|
if (!branchCreated) {
|
|
throw new Error(`Failed to create branch ${newBranchName}`)
|
|
}
|
|
|
|
// Use the injectable function instead of the hardcoded one
|
|
const { title, body } = triggerCreatePrDeps.createStandalonePRTitleAndBody(
|
|
prCommentFields,
|
|
existingTests,
|
|
generatedTests,
|
|
coverage_message,
|
|
newBranchName,
|
|
triggerCreatePrDeps.prContentBuilder,
|
|
)
|
|
|
|
const newPrData = await triggerCreatePrDeps.createStandalonePullRequest(
|
|
installationOctokit,
|
|
owner,
|
|
repo,
|
|
title,
|
|
body,
|
|
newBranchName,
|
|
baseBranch,
|
|
)
|
|
|
|
try {
|
|
await triggerCreatePrDeps.prisma.optimization_events.update({
|
|
where: { trace_id: traceId },
|
|
data: {
|
|
pr_id: String(newPrData.data.id),
|
|
is_optimization_found: true,
|
|
event_type: "pr_created",
|
|
},
|
|
})
|
|
} catch (eventError) {
|
|
console.error("Failed to update optimization event:", eventError)
|
|
}
|
|
|
|
await triggerCreatePrDeps.addLabelToPullRequest(
|
|
installationOctokit,
|
|
owner,
|
|
repo,
|
|
newPrData.data.number,
|
|
)
|
|
await triggerCreatePrDeps.assignReviewer(
|
|
installationOctokit,
|
|
owner,
|
|
repo,
|
|
newPrData.data.number,
|
|
nickname,
|
|
)
|
|
|
|
console.log(`Created new PR #${newPrData.data.number} with branch ${newPrData.data.head.ref}`)
|
|
triggerCreatePrDeps.posthog.capture({
|
|
distinctId: userId,
|
|
event: `cfapi-create-pr-success-standalone-pr-created`,
|
|
properties: {
|
|
owner,
|
|
repo,
|
|
newPrNumber: newPrData.data.number,
|
|
newPrBranch: newPrData.data.head.ref,
|
|
PRURL: newPrData.data.html_url,
|
|
},
|
|
})
|
|
|
|
if (traceId !== "") {
|
|
let pull_request_db = await triggerCreatePrDeps.prisma.optimization_features.findUnique({
|
|
where: {
|
|
trace_id: traceId,
|
|
},
|
|
select: {
|
|
pull_request: true,
|
|
},
|
|
})
|
|
|
|
if (pull_request_db) {
|
|
if (pull_request_db.pull_request === null || pull_request_db.pull_request === undefined) {
|
|
pull_request_db.pull_request = {}
|
|
}
|
|
|
|
;(pull_request_db.pull_request as any).new_pr_url = newPrData.data.html_url
|
|
|
|
await triggerCreatePrDeps.prisma.optimization_features.update({
|
|
where: {
|
|
trace_id: traceId,
|
|
},
|
|
data: {
|
|
pull_request: pull_request_db.pull_request,
|
|
},
|
|
})
|
|
}
|
|
}
|
|
|
|
return newPrData.data.number
|
|
} catch (error) {
|
|
console.error(`Error in triggerCreatePr: ${error}`)
|
|
return -1
|
|
}
|
|
}
|
|
|
|
// Endpoint to manually add repositories, similar to webhook logic
|
|
export async function addRepositoryManually(req: AuthorizedUserReq, res: Response): Promise<void> {
|
|
try {
|
|
const { repositories_added, installation } = req.body
|
|
|
|
if (!repositories_added || !installation?.id) {
|
|
res.status(400).send("Missing required fields")
|
|
return
|
|
}
|
|
|
|
const installationExists = await getAppInstallationByInstalltionId(installation.id)
|
|
|
|
if (!installationExists) {
|
|
await createAppInstallation({
|
|
installation_id: installation.id,
|
|
account_id: installation.account.id,
|
|
account_login: installation.account.login,
|
|
account_type: installation.account.type,
|
|
})
|
|
console.log(`Installation created for ID: ${installation.id}`)
|
|
}
|
|
|
|
for (const repo of repositories_added) {
|
|
try {
|
|
await prisma.repositories.upsert({
|
|
where: { github_repo_id: String(repo.id) },
|
|
update: {
|
|
name: repo.name,
|
|
full_name: repo.full_name,
|
|
is_private: repo.private,
|
|
installation_id: installation.id,
|
|
},
|
|
create: {
|
|
github_repo_id: String(repo.id),
|
|
installation_id: installation.id,
|
|
name: repo.name,
|
|
full_name: repo.full_name,
|
|
is_private: repo.private,
|
|
is_active: true,
|
|
},
|
|
})
|
|
const savedRepo = await upsertRepository({
|
|
github_repo_id: String(repo.id),
|
|
installation_id: installation.id,
|
|
name: repo.name,
|
|
full_name: repo.full_name,
|
|
is_private: repo.private,
|
|
})
|
|
await createRepositoryMember({
|
|
repository_id: savedRepo.id,
|
|
user_id: req.userId,
|
|
role: "",
|
|
})
|
|
console.log(`Repository upserted: ${repo.full_name}`)
|
|
} catch (error) {
|
|
console.error(`Failed to add/reactivate repository ${repo.full_name}:`, error)
|
|
Sentry.captureException(error)
|
|
}
|
|
}
|
|
|
|
res.status(200).send("Repositories added successfully")
|
|
} catch (error) {
|
|
console.error(`Error adding repositories: ${error}`)
|
|
if (error instanceof Error) {
|
|
res.status(500).send(`Error adding repositories: ${error.message}`)
|
|
} else {
|
|
res.status(500).send("Error adding repositories")
|
|
}
|
|
}
|
|
}
|
|
|
|
function getStandalonePRTitleAndBody(
|
|
prCommentFields: PrCommentFields,
|
|
existingTests: string,
|
|
generatedTests: string,
|
|
coverage_message: string,
|
|
newBranchName: string,
|
|
) {
|
|
const prCommentHeader = buildResultHeader(prCommentFields)
|
|
// Build benchmark info if available
|
|
const benchmarkInfo =
|
|
prCommentFields.benchmark_details && prCommentFields.benchmark_details.length > 0
|
|
? buildBenchmarkInfo(prCommentFields)
|
|
: ""
|
|
// Open a new PR from the new branch onto the original PR's head branch
|
|
const prCommentBody = buildResultDetails(prCommentFields)
|
|
const prCommentTestReport = buildResultTestReport(
|
|
prCommentFields,
|
|
existingTests,
|
|
generatedTests,
|
|
coverage_message,
|
|
)
|
|
const prCommentFooter = buildResultFooter(newBranchName)
|
|
const title: string = buildPrTitle(
|
|
prCommentFields.function_name,
|
|
prCommentFields.speedup_pct,
|
|
prCommentFields.speedup_x,
|
|
)
|
|
const body: string = benchmarkInfo
|
|
? `${prCommentHeader}\n${benchmarkInfo}\n${prCommentBody}\n${prCommentTestReport}\n${prCommentFooter}`
|
|
: `${prCommentHeader}\n${prCommentBody}\n${prCommentTestReport}\n${prCommentFooter}`
|
|
return { title: title, body: body }
|
|
}
|