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

import { determineValidHunks, fileDiffsToMap, isDiffContentsWellFormed } from "../diff_utils.js"
import { userNickname } from "../auth0-mgmt.js"
import { getInstallationOctokitByOwner, isUserCollaborator } from "../github/github-utils.js"
import { githubApp } from "../github/github-app.js"
import { buildDependentPrTitle, buildPrCommentBody } from "../github/pr-changes-utils.js"
import {
  createDependentPullRequest,
  createNewBranchFromDiffContents,
} from "../github/create-pr-from-diffcontents.js"
import { posthog } from "../analytics.js"
import suggester from "@codeflash-ai/code-suggester"
import { type FileDiffContent } from "@codeflash-ai/code-suggester/build/src/types.js"
import { PrismaClient } from "@prisma/client"
import { sendSlackMessage } from "../github/slack_util.js"
import { Request, Response } from "express"
import { AnyOctokit, AuthorizedUserReq, PullRequestDB } from "../types.js"

interface CustomRequest extends Request {
  userId?: string
}
const prisma = new PrismaClient()

const slackNotificationConfig = {
  "Future-House": ["aviary", "paper-qa"],
  "langflow-ai": ["langflow"],
  "albumentations-team": ["albumentations"],
  "Skyvern-AI": ["skyvern"],
  roboflow: ["inference"],
}

export async function suggestPrChanges(
  req: AuthorizedUserReq,
  res: Response,
): Promise<Response | undefined> {
  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}`)

    // Check if the owner is roboflow
    if (owner === "roboflow") {
      console.log(`Rejecting request for roboflow repository`)
      return res.status(401).send("Unauthorized for roboflow repositories")
    }

    // Suggest changes

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

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

    // 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
    console.log(`Attempting to access ref for: ${owner}/${repo}, branch: ${baseBranch}`)
    const commitMessage =
      buildDependentPrTitle(
        prCommentFields.function_name,
        prCommentFields.speedup_pct,
        prCommentFields.speedup_x,
        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}`)
    }

    let hasMultipleHunksInSameFile = false
    let hasMultipleFiles = validHunks.size > 1
    for (const [filePath, hunks] of validHunks.entries()) {
      if (hunks.length > 1) {
        console.log(
          `File ${filePath} has ${hunks.length} hunks, using dependent PR instead of review comments`,
        )
        hasMultipleHunksInSameFile = true
        break
      }
    }
    if (hasMultipleFiles) {
      console.log(
        `Found ${validHunks.size} files with changes, using dependent PR instead of review comments`,
      )
    }
    // Modified condition: Create a dependent PR if there are invalid hunks OR multiple valid hunks
    // This addresses both the invalid hunks case and the code suggester library issues
    if (
      invalidHunks.size > 0 ||
      validHunks.size > 1 ||
      hasMultipleFiles ||
      hasMultipleHunksInSameFile
    ) {
      // 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(
        `Creating a dependent PR because there are ${invalidHunks.size > 0 ? "invalid hunks" : "multiple valid hunks"}.`,
      )
      console.log(`Making a new dependent PR...`)

      const newPrData = await createDependentPullRequest(
        installationOctokit as AnyOctokit,
        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 (slackNotificationConfig[owner as keyof typeof slackNotificationConfig]?.includes(repo)) {
        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 as PullRequestDB).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 {
      // Only use the code suggester library when there is exactly one valid hunk and no invalid hunks
      console.log(`Creating unified review for a single valid hunk.`)
      // Include all the information in a single PR comment
      const prCommentBody = buildPrCommentBody(
        prCommentFields,
        existingTests,
        generatedTests,
        coverage_message,
        newBranchName,
        { isUnifiedReview: true, includeHeader: false },
      )
      let reviewComments = []
      let foundInvalidHunk = false

      // Process each valid hunk entry in the Map
      for (const [filePath, hunks] of validHunks.entries()) {
        // For each hunk in this file
        for (const hunk of hunks) {
          if (hunk.oldStart < hunk.oldEnd) {
            const newContent = hunk.newContent.join("\n")
            const isLongDiff = newContent.length > 500
            let commentBody
            if (isLongDiff) {
              // Create a collapsible section for long diffs
              commentBody =
                prCommentBody +
                "\n\n" +
                "<details>\n" +
                "<summary>Click to see suggested changes</summary>\n\n" +
                "```suggestion\n" +
                newContent +
                "\n```\n" +
                "</details>"
            } else {
              // For smaller diffs, show them directly
              commentBody = prCommentBody + "\n\n" + "```suggestion\n" + newContent + "\n```"
            }
            reviewComments.push({
              path: filePath,
              line: hunk.oldEnd,
              start_line: hunk.oldStart,
              side: "RIGHT",
              body: commentBody, // Format as a suggestion
            })
          } else {
            console.log(
              `Found invalid review range for ${filePath}: start_line (${hunk.oldStart}) must be less than line (${hunk.oldEnd}) with content ${hunk.newContent.join("\n")}`,
            )
            foundInvalidHunk = true
            break
          }
        }

        if (foundInvalidHunk) {
          break
        }
      }

      // Add explanation comment
      if (!foundInvalidHunk && reviewComments.length > 0) {
        const review = await installationOctokit.rest.pulls.createReview({
          owner,
          repo,
          pull_number: pullNumber,
          commit_id: originalPrData.data.head.sha,
          event: "COMMENT",
          comments: reviewComments,
        })
        console.log(`Added review comment to PR #${pullNumber}: ${review.data.html_url}`)

        if (
          slackNotificationConfig[owner as keyof typeof slackNotificationConfig]?.includes(repo)
        ) {
          await sendSlackMessage(`Suggestions made for ${review.data.html_url} in ${owner}/${repo}`)
        }
        posthog.capture({
          distinctId: userId,
          event: `cfapi-suggest-pr-changes-success-suggestions-made`,
          properties: {
            owner,
            repo,
            reviewId: review.data.id,
            PRURL: review.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 as PullRequestDB).pull_request.review_suggestion_pr_url =
              review.data.html_url

            await prisma.optimization_features.update({
              where: {
                trace_id: traceId,
              },
              data: pull_request_db,
            })
          }
        }
        res.json(review.data.id)
      } else {
        // Don't create a dependent PR, just return an error response
        const reason = foundInvalidHunk
          ? "invalid line ordering in hunks"
          : "no valid review comments could be created"

        console.log(`Cannot create review due to ${reason}`)

        // Return a meaningful error response
        return res.status(422).json({
          error: `Cannot create review comments due to ${reason}`,
          message: "Please consider creating a dependent PR instead",
        })
      }
    }
  } 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}`)
  }
}