codeflash-internal/js/cf-api/endpoints/code-context-hash.ts

import { prisma } from "@codeflash-ai/common"
import { Request, Response } from "express"
import { githubApp } from "../github/github-app.js"
import { getInstallationOctokitByOwner, isUserCollaborator } from "../github/github-utils.js"
import { userNickname } from "../auth0-mgmt.js"
import { posthog } from "../analytics.js"
import { logger } from "../utils/logger.js"
import {
  missingRequiredFields,
  unauthorized,
  githubNotCollaborator,
  validationFailure,
  internalServerError,
  conflict,
} from "../exceptions/index.js"

export async function is_code_being_optimized_again(req: Request, res: Response) {
  try {
    const { owner, repo, pr_number, code_contexts } = req.body
    const userId = (req as any).userId

    if (!repo || !owner || !pr_number || !code_contexts) {
      throw missingRequiredFields("repo, owner, pr_number, code_contexts")
    }
    const nickname: string | null = await userNickname(userId)
    if (nickname == null) {
      throw unauthorized("")
    }
    const octokit = await getInstallationOctokitByOwner(githubApp, owner, repo, userId)
    if (octokit instanceof Error) {
      throw internalServerError(octokit.message)
    }

    // Check collaborator status with error handling
    try {
      const isCollaborator = await isUserCollaborator(octokit, owner, repo, nickname)
      if (!isCollaborator) {
        logger.warn("User is not a collaborator on repository", req, {
          nickname,
          repo: `${owner}/${repo}`,
        })
        throw githubNotCollaborator(`${owner}/${repo}`)
      }
    } catch (error) {
      if (error && typeof error === "object" && "getHttpStatus" in error) {
        throw error
      }
      logger.error("Error checking collaborator status", req, {}, error as Error)
      throw internalServerError("Failed to verify collaborator status")
    }

    // Validate pr_number is an integer above 0
    const pr_number_int = parseInt(pr_number, 10)
    if (isNaN(pr_number_int) || pr_number_int <= 0) {
      throw validationFailure("pr_number must be a positive integer")
    }

    // Validate code_contexts is an array
    if (!Array.isArray(code_contexts)) {
      throw validationFailure("code_contexts must be an array of objects")
    }

    if (code_contexts.length === 0) {
      throw validationFailure("code_contexts cannot be empty")
    }

    // Validate each code context object has required fields
    for (let i = 0; i < code_contexts.length; i++) {
      const context = code_contexts[i]

      if (typeof context !== "object" || context === null) {
        throw validationFailure(`code_contexts[${i}] must be an object`)
      }

      const { file_path, function_name, code_hash } = context

      if (typeof file_path !== "string" || file_path.trim().length === 0) {
        throw validationFailure(`code_contexts[${i}].file_path must be a non-empty string`)
      }

      if (typeof function_name !== "string" || function_name.trim().length === 0) {
        throw validationFailure(`code_contexts[${i}].function_name must be a non-empty string`)
      }

      if (typeof code_hash !== "string" || !/^[a-f0-9]{64}$/.test(code_hash)) {
        throw validationFailure(`code_contexts[${i}].code_hash must be a valid SHA-256 hash`)
      }
    }

    // Build the compound keys for batch querying
    const compoundKeys = code_contexts.map(({ code_hash }) => ({
      repo,
      owner,
      pr_number: pr_number_int,
      context_hash: code_hash,
    }))

    const existingRecords = await prisma.pr_code_context_hash_cache.findMany({
      where: {
        OR: compoundKeys.map(key => ({
          repo: key.repo,
          owner: key.owner,
          pr_number: key.pr_number,
          context_hash: key.context_hash,
        })),
      },
      select: {
        context_hash: true,
      },
    })

    // Create a Set of existing hashes for O(1) lookup
    const existingHashes = new Set(existingRecords.map(record => record.context_hash))

    // Identify code contexts that are already optimized (return as tuples)
    const alreadyOptimizedTuples = code_contexts
      .filter(({ code_hash }) => existingHashes.has(code_hash))
      .map(({ file_path, function_name }) => [file_path, function_name])
    posthog?.capture({
      distinctId: userId,
      event: `cfapi-github-pr-optimization`,
      properties: {
        repo_owner: owner,
        repo_name: repo,
        pr_number,
      },
    })

    logger.info("Code context check completed", req, {
      repo: `${owner}/${repo}`,
      pr_number,
      total_contexts: code_contexts.length,
      already_optimized: alreadyOptimizedTuples.length,
      new_contexts: code_contexts.length - alreadyOptimizedTuples.length,
    })

    res.status(200).json({
      already_optimized_tuples: alreadyOptimizedTuples,
      total_contexts: code_contexts.length,
      new_contexts: code_contexts.length - alreadyOptimizedTuples.length,
    })
  } catch (error) {
    if (error && typeof error === "object" && "getHttpStatus" in error) {
      throw error
    }
    logger.errorWithSentry("Error in is_code_being_optimized_again", req, {}, error as Error)
    throw internalServerError("Error checking code optimization status")
  }
}

export async function add_optimized_code_context(req: Request, res: Response) {
  try {
    const { owner, repo, pr_number, code_hash } = req.body
    const userId = (req as any).userId

    if (!repo || !owner || !pr_number || !code_hash) {
      throw missingRequiredFields("repo, owner, pr_number, code_hash")
    }

    const nickname: string | null = await userNickname(userId)
    if (nickname == null) {
      throw unauthorized("")
    }

    const octokit = await getInstallationOctokitByOwner(githubApp, owner, repo, userId)
    if (octokit instanceof Error) {
      throw internalServerError(octokit.message)
    }

    // Check collaborator status with error handling
    try {
      const isCollaborator = await isUserCollaborator(octokit, owner, repo, nickname)
      if (!isCollaborator) {
        logger.warn("User is not a collaborator on repository", req, {
          nickname,
          repo: `${owner}/${repo}`,
        })
        throw githubNotCollaborator(`${owner}/${repo}`)
      }
    } catch (error) {
      if (error && typeof error === "object" && "getHttpStatus" in error) {
        throw error
      }
      logger.error("Error checking collaborator status", req, {}, error as Error)
      throw internalServerError("Failed to verify collaborator status")
    }

    // Validate pr_number is an integer above 0
    const pr_number_int = parseInt(pr_number, 10)
    if (isNaN(pr_number_int) || pr_number_int <= 0) {
      throw validationFailure("pr_number must be a positive integer")
    }

    // Validate code_hash is a valid SHA-256 hash
    if (typeof code_hash !== "string" || !/^[a-f0-9]{64}$/.test(code_hash)) {
      throw validationFailure("code_hash must be a valid SHA-256 hash")
    }

    // Create the new entry
    await prisma.pr_code_context_hash_cache.create({
      data: {
        repo,
        owner,
        pr_number: pr_number_int,
        context_hash: code_hash,
      },
    })

    logger.info("Code context successfully added", req, {
      repo: `${owner}/${repo}`,
      pr_number,
      code_hash,
    })

    res.status(201).json({
      message: "Code context successfully added",
    })
  } catch (error: any) {
    if (error && typeof error === "object" && "getHttpStatus" in error) {
      throw error
    }
    logger.errorWithSentry("Error in add_optimized_code_context", req, {}, error as Error)

    // Handle unique constraint violations gracefully
    if (error.code === "P2002") {
      logger.warn("Code context already exists for this PR", req, {})
      throw conflict("Code context already exists for this PR")
    }

    throw internalServerError("Error adding code context")
  }
}