codeflash-internal/js/cf-api/endpoints/create-pr.ts

import { fileDiffsToMap, isDiffContentsWellFormed } from "../diff_utils.js"
import { type FileDiffContent } from "code-suggester/build/src/types.js"
import { userNickname } from "../auth0-mgmt.js"
import {
  addLabelToPullRequest,
  getInstallationOctokitByOwner,
  isUserCollaborator,
} from "../github/github-utils.js"
import { buildPrTitle } from "../github/pr-changes-utils.js"
import { githubApp } from "../github/github-app.js"
import {
  assignReviewer,
  createNewBranchFromDiffContents,
  createStandalonePullRequest,
} from "../github/create-pr-from-diffcontents.js"
import { posthog } from "../analytics.js"
import { PrismaClient } from "@prisma/client"
import { Request, Response } from "express"

const prisma = new PrismaClient()

export async function createPr(req: Request, res: Response): Promise<void> {
  try {
    const {
      owner,
      repo,
      baseBranch,
      diffContents,
      prCommentFields,
      existingTests,
      generatedTests,
      coverage,
    } = req.body

    const userId = (req as any).userId
    //traceId is optional to allow for backwards compatibility, can make this required in the future
    const traceId = req.body.traceId || ""
    //coveragePct is optional to allow for backwards compatibility, can make this required in the future

    if (!repo || !owner || !baseBranch || !isDiffContentsWellFormed(diffContents)) {
      res.status(400).send("Missing or malformed fields")
      return
    }

    const nickname: string | null = await userNickname(userId)
    if (nickname == null) {
      res.status(401).send("Unauthorized") // Error getting user nickname
      return
    }

    const installationOctokit = await getInstallationOctokitByOwner(githubApp, owner, repo)
    if (installationOctokit instanceof Error) {
      res.status(401).send(installationOctokit.message)
      return
    }
    const isCollaborator = await isUserCollaborator(installationOctokit, owner, repo, nickname)
    if (!isCollaborator) {
      console.log(`${nickname} is not a collaborator on ${owner}/${repo}`)
      res.status(401).send("Unauthorized") // User is not a collaborator
      return
    }

    console.log(`${nickname} is a collaborator on ${owner}/${repo}`)

    const diffContentsMap: Map<string, FileDiffContent> = fileDiffsToMap(diffContents)

    // Generate a Base36-encoded timestamp
    const timestamp = Date.now() // Current timestamp in milliseconds
    const encodedTimestamp = timestamp.toString(36) // e.g., 'kf12oi0'

    // If you change this, please also change the regex in github-app.ts in the pull_request.closed event
    const newBranchName = `codeflash/optimize-${prCommentFields.function_name}-${encodedTimestamp}`
    const commitMessage =
      buildPrTitle(prCommentFields.function_name, prCommentFields.speedup_pct) +
      `\n${prCommentFields.optimization_explanation}`
    const branchCreated = await createNewBranchFromDiffContents(
      installationOctokit,
      owner,
      repo,
      newBranchName,
      baseBranch,
      diffContentsMap,
      commitMessage,
    )
    if (!branchCreated) {
      throw new Error(`Failed to create branch ${newBranchName}`)
    }

    const newPrData = await createStandalonePullRequest(
      installationOctokit,
      owner,
      repo,
      newBranchName,
      baseBranch,
      prCommentFields,
      existingTests,
      generatedTests,
      coverage,
    )

    await addLabelToPullRequest(installationOctokit, owner, repo, newPrData.data.number)

    // Assign the user who initiated the create-pr as the reviewer
    await assignReviewer(installationOctokit, owner, repo, newPrData.data.number, nickname)

    // Respond with the new PR details
    console.log(`Created new PR #${newPrData.data.number} with branch ${newPrData.data.head.ref}`)
    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 prisma.optimization_features.findUnique({
        where: {
          trace_id: traceId,
        },
        select: {
          pull_request: true,
        },
      })
      if (pull_request_db) {
        // the trace_id is not in the database then ignore it, because it should already exist by this stage.
        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 prisma.optimization_features.update({
          where: {
            trace_id: traceId,
          },
          data: {
            pull_request: pull_request_db.pull_request,
          },
        })
      }
    }
    res.json(newPrData.data.number)
  } 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}`)
      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`)
    }
  }
}