pytorchBotHandler.ts

import { PullRequestReview } from "@octokit/webhooks-types";
import _ from "lodash";
import { updateDrciComments } from "pages/api/drci/drci";
import shlex from "shlex";
import { queryClickhouseSaved } from "../clickhouse";
import { fetchCrcrAllowlist } from "../crcrAllowlist";
// lib/reviewApproval owns the approval rules, so the merge command and the
// Dr.CI PR Status line cannot disagree about whether a PR is approved.
import {
  getApprovalStatusFromReviews,
  PR_APPROVED,
  PR_CHANGES_REQUESTED,
} from "../reviewApproval";
import { getHelp, getParser } from "./cliParser";
import { cherryPickClassifications } from "./Constants";
import { downstreamRepoFromCheckRunName } from "./crcrOncallBot";
import PytorchBotLogger from "./pytorchbotLogger";
import {
  hasWritePermissions as _hasWP,
  addLabels,
  CachedConfigTracker,
  hasApprovedPullRuns,
  isFirstTimeContributor,
  isPyTorchbotSupportedOrg,
  isPyTorchPyTorch,
  reactOnComment,
} from "./utils";

export const CIFLOW_TRUNK_LABEL = "ciflow/trunk";
export const CIFLOW_PULL_LABEL = "ciflow/pull";

export interface PytorchbotParams {
  owner: string;
  repo: string;
  prNum: number;
  ctx: any;
  url: string;
  login: string;
  commentId: number;
  commentBody: string;
  useReactions: boolean;
  cachedConfigTracker: CachedConfigTracker;
}

class PytorchBotHandler {
  ctx: any;
  useReactions: boolean;
  owner: string;
  repo: string;
  prNum: number;
  url: string;
  commentId: number;
  login: string;
  commentBody: string;
  headSha: string | undefined;
  cachedConfigTracker: CachedConfigTracker;

  forceMergeMessagePat = new RegExp("^\\s*\\S+\\s+\\S+.*");

  logger: PytorchBotLogger;

  constructor(params: PytorchbotParams) {
    this.owner = params.owner;
    this.repo = params.repo;
    this.prNum = params.prNum;
    this.ctx = params.ctx;
    this.url = params.url;
    this.login = params.login;
    this.commentId = params.commentId;
    this.commentBody = params.commentBody;
    this.useReactions = params.useReactions;
    this.cachedConfigTracker = params.cachedConfigTracker;

    this.logger = new PytorchBotLogger(params);
  }

  async ackComment() {
    if (this.useReactions) {
      await reactOnComment(this.ctx, "+1");
    } else {
      await this.addComment("+1");
    }
  }

  async dispatchEvent(event_type: string, payload: any) {
    const { owner, repo, url, ctx, prNum, commentId } = this;

    let filtered_payload = _.pickBy(payload, function (val) {
      return val !== "" && val !== false;
    });

    let client_payload = {
      pr_num: prNum,
      comment_id: commentId,
      ...filtered_payload,
    };

    ctx.log(
      `Creating dispatch event of type "${event_type}" for comment ${url}`
    );
    await this.ctx.octokit.repos.createDispatchEvent({
      owner: owner,
      repo: repo,
      event_type: event_type,
      client_payload: client_payload,
    });
  }

  async addComment(comment: string) {
    const { ctx, owner, repo, prNum, url } = this;
    ctx.log(`Commenting with "${comment}" for pull request ${url}`);
    await this.ctx.octokit.issues.createComment({
      issue_number: prNum,
      body: comment,
      owner: owner,
      repo: repo,
    });
  }

  async handleConfused(
    leaveMessage: boolean,
    message: string = "@pytorch bot did not understand your command. Please try `@pytorchbot --help` for other commands."
  ) {
    await this.logger.log("confused", { message });
    if (this.useReactions) {
      await reactOnComment(this.ctx, "confused");
    }
    if (leaveMessage) {
      await this.addComment(message);
    }
  }

  isValidForceMergeMessage(message: string): boolean {
    // We can enforce  the merge message format here, for example, rejecting
    // all messages not in the following format `[CATEGORY] description`.
    //
    // However, it seems too strict to enforce a fixed set of categories right
    // away without conducting a user study for all common use cases of force
    // merge first. So the message is just a free form text for now
    const matches = message?.match(this.forceMergeMessagePat);
    return matches != undefined && matches.length != 0;
  }

  async reasonToRejectForceRequest(
    forceMessage: string
  ): Promise<string | null> {
    const { ctx } = this;

    const hasWritePermission = await this.hasWritePermissions(
      ctx.payload?.comment?.user?.login
    );
    if (!hasWritePermission) {
      return "You are not authorized to force merges to this repository. Please use the regular `@pytorchmergebot merge` command instead";
    }

    const isValidMessage = this.isValidForceMergeMessage(forceMessage);
    if (!isValidMessage) {
      return `You need to provide a reason for using force merge, in the format @pytorchbot merge -f 'Explanation'.
The explanation needs to be clear on why this is needed. Here are some good examples:
* Bypass checks due to unrelated upstream failures from ...
* This is a minor fix to ..., which shouldn't break anything
* This is pre-tested in a previous CI run
* Bypass flaky ... check`;
    }

    return null;
  }

  async getApprovalStatus(): Promise<string> {
    const reviews: PullRequestReview[] = await this.ctx.octokit.paginate(
      this.ctx.octokit.pulls.listReviews,
      {
        owner: this.owner,
        repo: this.repo,
        pull_number: this.prNum,
        per_page: 100,
      }
    );

    if (!reviews.length) {
      this.ctx.log("Could not find any reviews for PR");
      return "no_reviews";
    }

    // The rules themselves live in lib/reviewApproval so the Dr.CI PR Status
    // line reaches the same verdict this does; see the note there.
    return getApprovalStatusFromReviews(
      reviews,
      isPyTorchPyTorch(this.owner, this.repo),
      (message: string) => this.ctx.log(message)
    );
  }

  async handleMerge(
    forceMessage: string,
    ignore_current: boolean,
    rebase: string | boolean,
    ic: boolean
  ) {
    const config: any = await this.cachedConfigTracker.loadConfig(this.ctx);
    if (config == null || !config["mergebot"]) {
      await this.handleConfused(
        true,
        "Mergebot is not configured for this repository. Please use the merge button provided by GitHub."
      );
      return;
    }

    const extra_data = {
      forceMessage,
      rebase,
    };
    const forceRequested = forceMessage != undefined;
    let rejection_reason = null;

    if (forceRequested) {
      rejection_reason = await this.reasonToRejectForceRequest(forceMessage);
    } else if (isPyTorchbotSupportedOrg(this.owner)) {
      // Ensure the PR has been signed off on
      let approval_status = await this.getApprovalStatus();
      if (approval_status == PR_CHANGES_REQUESTED) {
        rejection_reason =
          "This PR has pending changes requested. Please address the comments and update the PR before merging.";
      } else if (approval_status !== PR_APPROVED) {
        rejection_reason =
          "This PR needs to be approved by an authorized maintainer before merge.";
      }
    }

    if (ic) {
      rejection_reason =
        "`-ic` flag is deprecated, please use `-i` instead for the same effect.";
    }

    if (ignore_current) {
      if (
        !(await this.hasWritePermissions(
          this.ctx.payload?.comment?.user?.login
        ))
      ) {
        rejection_reason =
          "`-i` flag is only allowed for users with write permissions";
      }
    }

    if (rejection_reason) {
      await this.logger.log("merge-error", extra_data);
      await this.handleConfused(true, rejection_reason);
      return;
    }

    if (
      rebase &&
      !(await this.hasRebasePermissions(this.ctx.payload?.comment?.user?.login))
    ) {
      await this.addComment(
        "You don't have permissions to rebase this PR since you are a first time contributor.  If you think this is a mistake, please contact PyTorch Dev Infra."
      );
      rebase = false;
    }

    await this.logger.log("merge", extra_data);

    // Check for L4 CRCR blocking failures (L3 failures are non-blocking).
    // Force merge (-f) bypasses this check, consistent with the existing
    // force-merge semantics for in-repo CI failures.
    // Only applies to pytorch/pytorch — the CRCR allowlist and check runs
    // are specific to that repo.
    if (!forceRequested && isPyTorchPyTorch(this.owner, this.repo)) {
      const blockingRepos = await this.getCrcrBlockingFailures();
      if (blockingRepos.length > 0) {
        await this.addComment(
          `The following L4 downstream CI workflows are blocking this merge (failed or still running):\n\n` +
            blockingRepos.map((r) => `- \`${r}\``).join("\n") +
            `\n\nPlease investigate or use \`@pytorchbot merge -f\` to bypass.`
        );
        return;
      }

      let labels: string[] = this.ctx.payload?.issue?.labels.map(
        (e: any) => e["name"]
      );
      if (labels === undefined) {
        labels = this.ctx.payload?.pull_request?.labels.map(
          (e: any) => e["name"]
        );
      }
      if (
        labels !== undefined &&
        !labels.find((x) => x === CIFLOW_TRUNK_LABEL)
      ) {
        if (
          !(await this.hasWorkflowRunningPermissions(
            this.ctx.payload?.issue?.user?.login
          ))
        ) {
          await this.addComment(
            "Pull workflow has not been scheduled for the PR yet. It could be because author doesn't have permissions to run those or skip-checks keywords were added to PR/commits, aborting merge.  " +
              "Please get/give approval for the workflows and/or remove skip ci decorators before next merge attempt.  " +
              "If you think this is a mistake, please contact PyTorch Dev Infra."
          );
          return;
        }
        await addLabels(this.ctx, [CIFLOW_TRUNK_LABEL]);
      }
      if (!(await this.hasCiFlowPull())) {
        await addLabels(this.ctx, [CIFLOW_PULL_LABEL]);
      }
    }

    await this.dispatchEvent("try-merge", {
      force: forceRequested,
      ignore_current: ignore_current,
      rebase: rebase,
    });
    await this.ackComment();
  }

  async handleRevert(reason: string) {
    await this.logger.log("revert", { reason });
    await this.dispatchEvent("try-revert", { reason: reason });
    await this.ackComment();
  }

  async handleRebase(branch: string) {
    await this.logger.log("rebase", { branch });
    const { ctx } = this;
    if (await this.hasRebasePermissions(ctx.payload?.comment?.user?.login)) {
      await this.dispatchEvent("try-rebase", { branch: branch });
      await this.ackComment();
    } else {
      await this.addComment(
        "You don't have permissions to rebase this PR since you are a first time contributor.  If you think this is a mistake, please contact PyTorch Dev Infra."
      );
    }
  }

  async existingRepoLabels(): Promise<string[]> {
    const { ctx, owner, repo } = this;
    const labels = await ctx.octokit.paginate(
      "GET /repos/{owner}/{repo}/labels",
      {
        owner: owner,
        repo: repo,
        per_page: 100,
      }
    );
    return labels.map((d: any) => d.name);
  }

  async hasWritePermissions(username: string): Promise<boolean> {
    return _hasWP(this.ctx, username);
  }

  async hasRebasePermissions(username: string): Promise<boolean> {
    return (
      (await _hasWP(this.ctx, username)) ||
      !(await isFirstTimeContributor(this.ctx, username))
    );
  }

  /** Lazy-load and cache the PR head SHA, then return it. */
  private async ensureHeadSha(): Promise<string> {
    if (this.headSha === undefined) {
      const pullRequest = await this.ctx.octokit.pulls.get({
        owner: this.owner,
        repo: this.repo,
        pull_number: this.prNum,
      });
      this.headSha = pullRequest.data.head.sha;
    }
    return this.headSha!;
  }

  async hasWorkflowRunningPermissions(username: string): Promise<boolean> {
    if (await _hasWP(this.ctx, username)) {
      return true;
    }

    return await hasApprovedPullRuns(
      this.ctx.octokit,
      this.ctx.payload.repository.owner.login,
      this.ctx.payload.repository.name,
      await this.ensureHeadSha()
    );
  }

  async handleLabel(labels: string[], is_pr_comment: boolean = true) {
    await this.logger.log("label", { labels });
    const { ctx } = this;
    /**
     * 1. Get all existing repo labels
     * 2. Parse labels from command
     * 3. Find valid and invalid labels
     * 4. Add valid labels to pr, report invalid labels
     */
    const repoLabels = new Set(await this.existingRepoLabels());
    // remove unnecessary spaces from labels
    const labelsToAdd = labels.map((s: string) => s.trim());

    const filteredLabels = labelsToAdd.filter((l: string) => repoLabels.has(l));
    const invalidLabels = labelsToAdd.filter((l: string) => !repoLabels.has(l));
    const ciflowLabels = labelsToAdd.filter((l: string) =>
      l.startsWith("ciflow/")
    );
    if (ciflowLabels.length > 0 && !is_pr_comment) {
      return await this.handleConfused(
        true,
        "Can't add ciflow labels to an Issue."
      );
    }

    if (
      labelsToAdd.includes("actionable") &&
      !(await this.hasWritePermissions(ctx.payload?.comment?.user?.login))
    ) {
      return await this.addComment(
        "Only regular contributors are expected to mark issues as actionable."
      );
    }

    // Labels only people with write access to the repo should be able to add
    const labels_requiring_write_access: string[] = ["skip-pr-sanity-check"];
    const write_required_labels = labelsToAdd.filter((l: string) =>
      labels_requiring_write_access.some(
        (write_required_label) => write_required_label === l
      )
    );

    if (
      write_required_labels.length > 0 &&
      !(await this.hasWritePermissions(ctx.payload?.comment?.user?.login))
    ) {
      return await this.addComment(
        "Only people with write access to the repo can add these labels: " +
          write_required_labels.join(", ") +
          ". Please ping one of the reviewers for help."
      );
    }

    if (
      ciflowLabels.length > 0 &&
      !(await this.hasWorkflowRunningPermissions(
        ctx.payload?.comment?.user?.login
      ))
    ) {
      // Still add the labels (they represent user intent), but inform
      // that CI won't be triggered until workflows are approved.
      // The ciflowPushTrigger will handle posting a detailed pending comment.
      await this.addComment(
        "The ciflow label(s) " +
          ciflowLabels.join(", ") +
          " will be added, but CI won't be triggered until " +
          "the workflows are approved (scroll to the bottom of this page).\n\n" +
          "Please ping one of the reviewers if you do not have access to approve and run workflows."
      );
      // Don't return -- let the labels be added below
    }
    if (invalidLabels.length > 0) {
      await this.addComment(
        "Didn't find following labels among repository labels: " +
          invalidLabels.join(",")
      );
    }
    if (filteredLabels.length > 0) {
      await addLabels(ctx, filteredLabels);
      await this.ackComment();
    }
  }

  async handleDrCI() {
    await this.logger.log("Dr. CI");
    const { ctx, prNum, repo, owner } = this;

    await this.ackComment();
    await updateDrciComments(ctx.octokit, owner, repo, [prNum]);
  }

  async handleLint(login: string) {
    await this.logger.log("lint");
    if (!(await this.hasWritePermissions(login))) {
      await this.addComment(
        "You don't have permissions to trigger lint fixes on this PR since it " +
          "requires write access to the repository. If you think this is a " +
          "mistake, please contact PyTorch Dev Infra."
      );
      return;
    }
    await this.dispatchEvent("apply-lint", {});
    await this.ackComment();
  }

  async handleCherryPick(
    branch: string,
    fixes: string,
    classification: string
  ) {
    await this.logger.log("cherry-pick", { branch, fixes, classification });

    await this.ackComment();
    const classificationData = cherryPickClassifications[classification];
    await this.dispatchEvent("try-cherry-pick", {
      branch: branch,
      fixes: fixes,
      classification: classification,
      requiresIssue: classificationData.requiresIssue,
      classificationHelp: classificationData.help,
    });
  }

  async handlePytorchCommands(
    inputArgs: string,
    is_pr_comment: boolean = true
  ) {
    let args;
    let split_args: string[] = [];

    try {
      const parser = getParser();

      split_args = shlex.split(inputArgs);
      args = parser.parse_args(split_args);
    } catch (err: any) {
      // If the args are invalid, comment with the error + some help.
      await this.addComment(
        "❌ 🤖 pytorchbot command failed: \n```\n" +
          err.message +
          "```\n" +
          "Try `@pytorchbot --help` for more info."
      );
      return;
    }

    // if help is present as an option on the main command, or -h or --help is in any location in the args (parseargs fails to get -h at the start of the args)
    if (
      args.help ||
      split_args.includes("-h") ||
      split_args.includes("--help")
    ) {
      return await this.addComment(getHelp());
    }

    // commands which only make sense in the context of a PR
    if (is_pr_comment) {
      switch (args.command) {
        case "revert":
          return await this.handleRevert(args.message);
        case "merge":
          return await this.handleMerge(
            args.force,
            args.ignore_current,
            args.rebase,
            args.ic
          );
        case "rebase": {
          if (!args.branch) {
            args.branch = "viable/strict";
          }
          return await this.handleRebase(args.branch);
        }
        case "drci": {
          return await this.handleDrCI();
        }
        case "lint":
        case "fix-lint":
        case "apply-lint": {
          return await this.handleLint(this.ctx.payload?.comment?.user?.login);
        }
      }
    }
    switch (args.command) {
      case "label": {
        return await this.handleLabel(args.labels, is_pr_comment);
      }
      case "cherry-pick": {
        return await this.handleCherryPick(
          args.onto,
          args.fixes,
          args.classification
        );
      }
      default:
        return await this.handleConfused(false);
    }
  }

  async hasCiFlowPull(): Promise<boolean> {
    try {
      const workflowNames = await this.getWorkflowsLatest();
      return (
        workflowNames?.some(
          (workflow: any) => workflow.workflow_name === "pull"
        ) ?? false
      );
    } catch (error: any) {
      // Return true if we cannot read workflow data so that we don't unneccisarily tag the PR
      await this.logger.log("workflow-pull-error", error);
      return true;
    }
  }

  // Returns the workflows attached to the PR only for the latest commit
  async getWorkflowsLatest(): Promise<any> {
    return await queryClickhouseSaved("get_workflows_for_commit", {
      prNumber: this.prNum,
      headSha: this.headSha,
    });
  }

  /**
   * Return the list of L4 downstream repos whose CRCR check runs are failing
   * or still pending on this PR's head commit.  L3 workflows are intentionally
   * omitted — they are non-blocking.
   *
   * A pending (not-yet-completed) L4 check run also blocks merge: a still-
   * running check could fail, so merging before it completes would bypass
   * the downstream gating.  Use ``@pytorchbot merge -f`` to override.
   */
  async getCrcrBlockingFailures(): Promise<string[]> {
    // Only "failure" blocks merge — "cancelled" and "timed_out" are often
    // superseded / infra-related and should not gate merge.
    const BLOCKING_CONCLUSIONS = new Set(["failure"]);

    // Query GitHub Check Runs API for all check runs on this commit.
    // Use paginate to handle PRs with more than 100 check runs.
    let checkRuns: any[] = [];
    try {
      const headSha = await this.ensureHeadSha();
      checkRuns = await this.ctx.octokit.paginate(
        this.ctx.octokit.checks.listForRef,
        {
          owner: this.owner,
          repo: this.repo,
          ref: headSha,
          filter: "latest",
          per_page: 100,
        }
      );
    } catch {
      // If we can't fetch check runs, fail open (don't block merge on
      // an infrastructure error, including transient ensureHeadSha failure)
      this.ctx.log("getCrcrBlockingFailures: failed to list check runs");
      return [];
    }

    // Find all CRCR check runs (not just failures — we also need to
    // gate on pending L4 checks so a still-running check isn't bypassed).
    const crcrCheckRuns = checkRuns.filter((cr: any) =>
      cr.name.startsWith("crcr/")
    );

    if (crcrCheckRuns.length === 0) {
      return [];
    }

    // Load the allowlist to classify each repo as L3 or L4
    let allowlist;
    try {
      allowlist = await fetchCrcrAllowlist(this.ctx.octokit);
    } catch {
      this.ctx.log("getCrcrBlockingFailures: failed to load allowlist");
      return [];
    }

    const blocking = new Set<string>();
    for (const cr of crcrCheckRuns) {
      const downstreamRepo = downstreamRepoFromCheckRunName(cr.name);
      if (!downstreamRepo || !allowlist.isBlocking(downstreamRepo)) {
        continue; // Not L4
      }

      // Block when the L4 check has failed OR is still pending.
      // Successful/completed L4 checks don't block.
      const isFailure = BLOCKING_CONCLUSIONS.has(cr.conclusion ?? "");
      const isPending = cr.status !== "completed";

      if (isFailure || isPending) {
        blocking.add(downstreamRepo);
      }
    }

    return [...blocking];
  }
}

export default PytorchBotHandler;