codeflash-internal/js/cf-api/endpoints/suggest-pr-changes.ts

import { determineValidHunks, fileDiffsToMap, isDiffContentsWellFormed } from "../diff_utils"
import { userNickname } from "../auth0-mgmt"
import {
  buildDependentPrTitle,
  getInstallationOctokitByOwner,
  isUserCollaborator,
} from "../github/github-utils"
import { githubApp } from "../github/github-app"
import {
  buildResultDetails,
  buildResultHeader,
  buildResultTestReport,
  createDependentPullRequest,
  createNewBranchFromDiffContents,
} from "../github/create-pr-from-diffcontents"
import { posthog } from "../analytics"
import suggester from "code-suggester"
import { type FileDiffContent } from "code-suggester/build/src/types"
import { PrismaClient } from "@prisma/client"
import { sendSlackMessage } from "github/slack_util"

const prisma = new PrismaClient()

export async function suggestPrChanges(req, res): Promise<void> {
  try {
    const {
      owner,
      repo,
      pullNumber,
      diffContents,
      prCommentFields,
      existingTests,
      generatedTests,
      coverage_message,
    } = req.body
    const userId = req.userId
    //traceId is optional to allow for backwards compatibility, can make this required in the future
    const traceId = req.body.traceId || ""
    console.log(`traceId: ${traceId}`)

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

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

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

    // Suggest changes

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

    const { validHunks, invalidHunks } = await determineValidHunks(
      installationOctokit.rest,
      { owner, repo },
      pullNumber,
      100,
      diffContentsMap,
    )

    if (invalidHunks.size > 0) {
      // we can't suggest all of the hunks for this PR, because some of them are invalid (out of scope for this PR).
      // so instead, let's make a new branch, commit the changes,
      // and make a new PR from that branch onto the PR's branch

      console.log(
        `Can't suggest all of the hunks for this PR, because some of them are invalid (out of scope for this PR).`,
      )
      console.log(`Making a new dependent PR...`)

      // a timestamp format like 2024-01-31-12.59.48
      const timestamp = new Date()
        .toISOString()
        .replace(/:/g, ".")
        .replace(/\.\d+Z$/, "")

      // If you change this, please also change the regex in github-app.ts in the pull_request.closed event
      const newBranchName = `codeflash/optimize-pr${pullNumber}-${timestamp}`

      // Get the head branch of the original pull request
      const originalPrData = await installationOctokit.rest.pulls.get({
        owner,
        repo,
        pull_number: pullNumber,
      })
      const baseBranch = originalPrData.data.head.ref

      const commitMessage =
        buildDependentPrTitle(
          prCommentFields.function_name,
          prCommentFields.speedup_pct,
          pullNumber,
          baseBranch,
        ) + `\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 createDependentPullRequest(
        installationOctokit,
        owner,
        repo,
        pullNumber,
        newBranchName,
        baseBranch,
        prCommentFields,
        existingTests,
        generatedTests,
      )

      // Respond with the new PR details
      console.log(
        `Created new dependent PR #${newPrData.data.number} from branch ${newPrData.data.head.ref}`,
      )

      if (owner === "Future-House" && (repo === "aviary" || repo === "paper-qa")) {
        await sendSlackMessage(
          `new dependent PR created: ${newPrData.data.html_url} for ${owner}/${repo}`,
        )
      }

      posthog.capture({
        distinctId: userId,
        event: `cfapi-suggest-pr-changes-success-dependent-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) {
          if (pull_request_db.pull_request === null || pull_request_db.pull_request === undefined) {
            pull_request_db.pull_request = {}
          }
          pull_request_db.pull_request.dependent_pr_url = newPrData.data.html_url

          await prisma.optimization_features.update({
            where: {
              trace_id: traceId,
            },
            data: pull_request_db,
          })
        }
      }
      res.json(newPrData.data.number)
    } else {
      // all suggestions are valid and can be made on this PR.
      // let's make the suggestions

      console.log(`All suggestions are valid and can be made on this PR.`)
      console.log(`Making suggestions...`)

      const reviewNumber = await suggester.reviewPullRequest(
        installationOctokit.rest,
        diffContentsMap,
        {
          repo,
          owner,
          pullNumber,
          pageSize: 100,
        },
      )
      if (reviewNumber == null) {
        console.log(`Failed to create a suggester review for PR #${pullNumber}.`)
        console.log(`This may be because there are no changes to be made.`)
        console.log(`diffContentsMap: ${JSON.stringify(diffContentsMap)}`)
        posthog.capture({
          distinctId: userId,
          event: `cfapi-suggest-pr-changes-failed-no-changes`,
        })
        return res.status(500).send("Failed creating suggestions for review.")
      }
      console.log(`Created suggestions for review #${reviewNumber} for PR #${pullNumber}`)

      const prCommentBody =
        `## ⚡️ Codeflash found optimizations for this PR\n` +
        `${buildResultHeader(prCommentFields)}\n` +
        `${buildResultDetails(prCommentFields)}\n` +
        `${buildResultTestReport(prCommentFields, existingTests, generatedTests, coverage_message)}\n`

      // Add explanation comment
      const result = await installationOctokit.rest.issues.createComment({
        owner,
        repo,
        issue_number: pullNumber,
        body: prCommentBody,
      })
      console.log(`Added review comment to PR #${pullNumber}: ${result.data.html_url}`)

      if (owner === "Future-House" && (repo === "aviary" || repo === "paper-qa")) {
        await sendSlackMessage(`Suggestions made for PR #${pullNumber} in ${owner}/${repo}`)
      }
      posthog.capture({
        distinctId: userId,
        event: `cfapi-suggest-pr-changes-success-suggestions-made`,
        properties: {
          owner,
          repo,
          reviewNumber,
          PRURL: result.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.review_suggestion_pr_url = result.data.html_url

          await prisma.optimization_features.update({
            where: {
              trace_id: traceId,
            },
            data: pull_request_db,
          })
        }
      }
      res.json(reviewNumber)
    }
  } catch (error: any) {
    console.log(`Error in /cfapi/suggest-pr-changes: ${error}`)
    console.log(`Error message: ${error.message}`)
    console.log(`Error stack: ${error.stack}`)
    posthog.capture({
      distinctId: req.userId,
      event: `cfapi-suggest-pr-changes-failed-error`,
      properties: {
        error: error.message,
        stack: error.stack,
      },
    })
    res.status(500).send(`Error creating pull request: ${error.message}`)
  }
}