Merge branch 'main' into robertbrignull/automodel-sort-order

Author: robertbrignull

extensions/ql-vscode/src/model-editor/auto-model.ts

@@ -5,57 +5,6 @@ import type { AutoModelQueriesResult } from "./auto-model-codeml-queries";
 import { assertNever } from "../common/helpers-pure";
 import type { Log } from "sarif";
 import { gzipEncode } from "../common/zlib";
-import type { Method, MethodSignature } from "./method";
-import type { ModeledMethod } from "./modeled-method";
-import { groupMethods, sortGroupNames, sortMethods } from "./shared/sorting";
-
-/**
- * Return the candidates that the model should be run on. This includes limiting the number of
- * candidates to the candidate limit and filtering out anything that is already modeled and respecting
- * the order in the UI.
- * @param mode Whether it is application or framework mode.
- * @param methods all methods.
- * @param modeledMethodsBySignature the currently modeled methods.
- * @returns list of modeled methods that are candidates for modeling.
- */
-export function getCandidates(
-  mode: Mode,
-  methods: readonly Method[],
-  modeledMethodsBySignature: Record<string, readonly ModeledMethod[]>,
-  processedByAutoModelMethods: Set<string>,
-): MethodSignature[] {
-  const candidateMethods = methods.filter((method) => {
-    // Filter out any methods already processed by auto-model
-    if (processedByAutoModelMethods.has(method.signature)) {
-      return false;
-    }
-
-    const modeledMethods: ModeledMethod[] = [
-      ...(modeledMethodsBySignature[method.signature] ?? []),
-    ];
-
-    // Anything that is modeled is not a candidate
-    if (modeledMethods.some((m) => m.type !== "none")) {
-      return false;
-    }
-
-    // A method that is supported is modeled outside of the model file, so it is not a candidate.
-    if (method.supported) {
-      return false;
-    }
-
-    return true;
-  });
-
-  // Sort the same way as the UI so we send the first ones listed in the UI first
-  const grouped = groupMethods(candidateMethods, mode);
-  const sortedGroupNames = sortGroupNames(grouped);
-  return sortedGroupNames.flatMap((name) =>
-    // We can safely pass empty sets for `modifiedSignatures` and `processedByAutoModelMethods`
-    // because we've filtered out all methods that are already modeled or have already been processed by auto-model.
-    sortMethods(grouped[name], modeledMethodsBySignature, new Set(), new Set()),
-  );
-}
 
 /**
  * Encode a SARIF log to the format expected by the server: JSON, GZIP-compressed, base64-encoded

extensions/ql-vscode/src/model-editor/auto-modeler.ts

@@ -3,7 +3,8 @@ import type { ModeledMethod } from "./modeled-method";
 import { load as loadYaml } from "js-yaml";
 import type { ProgressCallback } from "../common/vscode/progress";
 import { withProgress } from "../common/vscode/progress";
-import { createAutoModelRequest, getCandidates } from "./auto-model";
+import { createAutoModelRequest } from "./auto-model";
+import { getCandidates } from "./shared/auto-model-candidates";
 import { runAutoModelQueries } from "./auto-model-codeml-queries";
 import { loadDataExtensionYaml } from "./yaml";
 import type { ModelRequest, ModelResponse } from "./auto-model-api";

extensions/ql-vscode/src/model-editor/shared/auto-model-candidates.ts

@@ -0,0 +1,53 @@
+import type { Method, MethodSignature } from "../method";
+import type { ModeledMethod } from "../modeled-method";
+import type { Mode } from "./mode";
+import { groupMethods, sortGroupNames, sortMethods } from "./sorting";
+
+/**
+ * Return the candidates that the model should be run on. This includes limiting the number of
+ * candidates to the candidate limit and filtering out anything that is already modeled and respecting
+ * the order in the UI.
+ * @param mode Whether it is application or framework mode.
+ * @param methods all methods.
+ * @param modeledMethodsBySignature the currently modeled methods.
+ * @returns list of modeled methods that are candidates for modeling.
+ */
+
+export function getCandidates(
+  mode: Mode,
+  methods: readonly Method[],
+  modeledMethodsBySignature: Record<string, readonly ModeledMethod[]>,
+  processedByAutoModelMethods: Set<string>,
+): MethodSignature[] {
+  const candidateMethods = methods.filter((method) => {
+    // Filter out any methods already processed by auto-model
+    if (processedByAutoModelMethods.has(method.signature)) {
+      return false;
+    }
+
+    const modeledMethods: ModeledMethod[] = [
+      ...(modeledMethodsBySignature[method.signature] ?? []),
+    ];
+
+    // Anything that is modeled is not a candidate
+    if (modeledMethods.some((m) => m.type !== "none")) {
+      return false;
+    }
+
+    // A method that is supported is modeled outside of the model file, so it is not a candidate.
+    if (method.supported) {
+      return false;
+    }
+
+    return true;
+  });
+
+  // Sort the same way as the UI so we send the first ones listed in the UI first
+  const grouped = groupMethods(candidateMethods, mode);
+  const sortedGroupNames = sortGroupNames(grouped);
+  return sortedGroupNames.flatMap((name) =>
+    // We can safely pass empty sets for `modifiedSignatures` and `processedByAutoModelMethods`
+    // because we've filtered out all methods that are already modeled or have already been processed by auto-model.
+    sortMethods(grouped[name], modeledMethodsBySignature, new Set(), new Set()),
+  );
+}

extensions/ql-vscode/src/variant-analysis/code-scanning-pack.ts

@@ -0,0 +1,77 @@
+import { join } from "path";
+import type { BaseLogger } from "../common/logging";
+import type { QueryLanguage } from "../common/query-language";
+import type { CodeQLCliServer } from "../codeql-cli/cli";
+import type { QlPackDetails } from "./ql-pack-details";
+import { getQlPackFilePath } from "../common/ql";
+
+export async function resolveCodeScanningQueryPack(
+  logger: BaseLogger,
+  cliServer: CodeQLCliServer,
+  language: QueryLanguage,
+): Promise<QlPackDetails> {
+  // Get pack
+  void logger.log(`Downloading pack for language: ${language}`);
+  const packName = `codeql/${language}-queries`;
+  const packDownloadResult = await cliServer.packDownload([packName]);
+  const downloadedPack = packDownloadResult.packs[0];
+
+  const packDir = join(
+    packDownloadResult.packDir,
+    downloadedPack.name,
+    downloadedPack.version,
+  );
+
+  // Resolve queries
+  void logger.log(`Resolving queries for pack: ${packName}`);
+  const suitePath = join(
+    packDir,
+    "codeql-suites",
+    `${language}-code-scanning.qls`,
+  );
+  const resolvedQueries = await cliServer.resolveQueries(suitePath);
+
+  const problemQueries = await filterToOnlyProblemQueries(
+    logger,
+    cliServer,
+    resolvedQueries,
+  );
+
+  if (problemQueries.length === 0) {
+    throw Error(
+      `No problem queries found in published query pack: ${packName}.`,
+    );
+  }
+
+  // Return pack details
+  const qlPackFilePath = await getQlPackFilePath(packDir);
+
+  const qlPackDetails: QlPackDetails = {
+    queryFiles: problemQueries,
+    qlPackRootPath: packDir,
+    qlPackFilePath,
+    language,
+  };
+
+  return qlPackDetails;
+}
+
+async function filterToOnlyProblemQueries(
+  logger: BaseLogger,
+  cliServer: CodeQLCliServer,
+  queries: string[],
+): Promise<string[]> {
+  const problemQueries: string[] = [];
+  for (const query of queries) {
+    const queryMetadata = await cliServer.resolveMetadata(query);
+    if (
+      queryMetadata.kind === "problem" ||
+      queryMetadata.kind === "path-problem"
+    ) {
+      problemQueries.push(query);
+    } else {
+      void logger.log(`Skipping non-problem query ${query}`);
+    }
+  }
+  return problemQueries;
+}

extensions/ql-vscode/src/variant-analysis/variant-analysis-manager.ts

@@ -94,6 +94,7 @@ import { getQlPackFilePath } from "../common/ql";
 import { tryGetQueryMetadata } from "../codeql-cli/query-metadata";
 import { getOnDiskWorkspaceFolders } from "../common/vscode/workspace-folders";
 import { findVariantAnalysisQlPackRoot } from "./ql";
+import { resolveCodeScanningQueryPack } from "./code-scanning-pack";
 
 const maxRetryCount = 3;
 
@@ -219,7 +220,7 @@ export class VariantAnalysisManager
   public async runVariantAnalysisFromPublishedPack(): Promise<void> {
     return withProgress(async (progress, token) => {
       progress({
-        maxStep: 8,
+        maxStep: 7,
         step: 0,
         message: "Determining query language",
       });
@@ -230,53 +231,17 @@ export class VariantAnalysisManager
       }
 
       progress({
-        maxStep: 8,
-        step: 1,
-        message: "Downloading query pack",
-      });
-
-      const packName = `codeql/${language}-queries`;
-      const packDownloadResult = await this.cliServer.packDownload([packName]);
-      const downloadedPack = packDownloadResult.packs[0];
-
-      const packDir = join(
-        packDownloadResult.packDir,
-        downloadedPack.name,
-        downloadedPack.version,
-      );
-
-      progress({
-        maxStep: 8,
+        maxStep: 7,
         step: 2,
-        message: "Resolving queries in pack",
+        message: "Downloading query pack and resolving queries",
       });
 
-      const suitePath = join(
-        packDir,
-        "codeql-suites",
-        `${language}-code-scanning.qls`,
-      );
-      const resolvedQueries = await this.cliServer.resolveQueries(suitePath);
-
-      const problemQueries =
-        await this.filterToOnlyProblemQueries(resolvedQueries);
-
-      if (problemQueries.length === 0) {
-        void this.app.logger.showErrorMessage(
-          `Unable to trigger variant analysis. No problem queries found in published query pack: ${packName}.`,
-        );
-        return;
-      }
-
-      const qlPackFilePath = await getQlPackFilePath(packDir);
-
       // Build up details to pass to the functions that run the variant analysis.
-      const qlPackDetails: QlPackDetails = {
-        queryFiles: problemQueries,
-        qlPackRootPath: packDir,
-        qlPackFilePath,
+      const qlPackDetails = await resolveCodeScanningQueryPack(
+        this.app.logger,
+        this.cliServer,
         language,
-      };
+      );
 
       await this.runVariantAnalysis(
         qlPackDetails,
@@ -291,24 +256,6 @@ export class VariantAnalysisManager
     });
   }
 
-  private async filterToOnlyProblemQueries(
-    queries: string[],
-  ): Promise<string[]> {
-    const problemQueries: string[] = [];
-    for (const query of queries) {
-      const queryMetadata = await this.cliServer.resolveMetadata(query);
-      if (
-        queryMetadata.kind === "problem" ||
-        queryMetadata.kind === "path-problem"
-      ) {
-        problemQueries.push(query);
-      } else {
-        void this.app.logger.log(`Skipping non-problem query ${query}`);
-      }
-    }
-    return problemQueries;
-  }
-
   private async runVariantAnalysisCommand(queryFiles: Uri[]): Promise<void> {
     if (queryFiles.length === 0) {
       throw new Error("Please select a .ql file to run as a variant analysis");

extensions/ql-vscode/src/view/model-editor/LibraryRow.tsx

@@ -14,6 +14,7 @@ import {
 } from "@vscode/webview-ui-toolkit/react";
 import type { ModelEditorViewState } from "../../model-editor/shared/view-state";
 import type { AccessPathSuggestionOptions } from "../../model-editor/suggestions";
+import { getCandidates } from "../../model-editor/shared/auto-model-candidates";
 
 const LibraryContainer = styled.div`
   background-color: var(--vscode-peekViewResult-background);
@@ -186,6 +187,17 @@ export const LibraryRow = ({
     return methods.some((method) => inProgressMethods.has(method.signature));
   }, [methods, inProgressMethods]);
 
+  const modelWithAIDisabled = useMemo(() => {
+    return (
+      getCandidates(
+        viewState.mode,
+        methods,
+        modeledMethodsMap,
+        processedByAutoModelMethods,
+      ).length === 0
+    );
+  }, [methods, modeledMethodsMap, processedByAutoModelMethods, viewState.mode]);
+
   return (
     <LibraryContainer>
       <TitleContainer onClick={toggleExpanded} aria-expanded={isExpanded}>
@@ -205,7 +217,11 @@ export const LibraryRow = ({
           {hasUnsavedChanges ? <VSCodeTag>UNSAVED</VSCodeTag> : null}
         </NameContainer>
         {viewState.showLlmButton && !canStopAutoModeling && (
-          <VSCodeButton appearance="icon" onClick={handleModelWithAI}>
+          <VSCodeButton
+            appearance="icon"
+            disabled={modelWithAIDisabled}
+            onClick={handleModelWithAI}
+          >
             <Codicon name="lightbulb-autofix" label="Model with AI" />
             &nbsp;Model with AI
           </VSCodeButton>