Merge pull request #3377 from github/robertbrignull/disable-automodel-button

robertbrignull · web-flow · commit fee7eae7029e · 2024-02-21T16:11:21.000Z
Disable the automodel button if there are no methods that can be modeled
diff --git a/extensions/ql-vscode/src/model-editor/auto-model.ts b/extensions/ql-vscode/src/model-editor/auto-model.ts
@@ -5,59 +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[] {
-  // Filter out any methods already processed by auto-model
-  methods = methods.filter(
-    (m) => !processedByAutoModelMethods.has(m.signature),
-  );
-
-  // Sort the same way as the UI so we send the first ones listed in the UI first
-  const grouped = groupMethods(methods, mode);
-  const sortedGroupNames = sortGroupNames(grouped);
-  const sortedMethods = sortedGroupNames.flatMap((name) =>
-    sortMethods(grouped[name]),
-  );
-
-  const candidates: MethodSignature[] = [];
-
-  for (const method of sortedMethods) {
-    const modeledMethods: ModeledMethod[] = [
-      ...(modeledMethodsBySignature[method.signature] ?? []),
-    ];
-
-    // Anything that is modeled is not a candidate
-    if (modeledMethods.some((m) => m.type !== "none")) {
-      continue;
-    }
-
-    // A method that is supported is modeled outside of the model file, so it is not a candidate.
-    if (method.supported) {
-      continue;
-    }
-
-    // The rest are candidates
-    candidates.push(method);
-  }
-  return candidates;
-}
 
 /**
  * Encode a SARIF log to the format expected by the server: JSON, GZIP-compressed, base64-encoded
diff --git a/extensions/ql-vscode/src/model-editor/auto-modeler.ts b/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";
diff --git a/extensions/ql-vscode/src/model-editor/shared/auto-model-candidates.ts b/extensions/ql-vscode/src/model-editor/shared/auto-model-candidates.ts
@@ -0,0 +1,55 @@
+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[] {
+  // Filter out any methods already processed by auto-model
+  methods = methods.filter(
+    (m) => !processedByAutoModelMethods.has(m.signature),
+  );
+
+  // Sort the same way as the UI so we send the first ones listed in the UI first
+  const grouped = groupMethods(methods, mode);
+  const sortedGroupNames = sortGroupNames(grouped);
+  const sortedMethods = sortedGroupNames.flatMap((name) =>
+    sortMethods(grouped[name]),
+  );
+
+  const candidates: MethodSignature[] = [];
+
+  for (const method of sortedMethods) {
+    const modeledMethods: ModeledMethod[] = [
+      ...(modeledMethodsBySignature[method.signature] ?? []),
+    ];
+
+    // Anything that is modeled is not a candidate
+    if (modeledMethods.some((m) => m.type !== "none")) {
+      continue;
+    }
+
+    // A method that is supported is modeled outside of the model file, so it is not a candidate.
+    if (method.supported) {
+      continue;
+    }
+
+    // The rest are candidates
+    candidates.push(method);
+  }
+  return candidates;
+}
diff --git a/extensions/ql-vscode/src/view/model-editor/LibraryRow.tsx b/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>
diff --git a/extensions/ql-vscode/test/unit-tests/model-editor/auto-model.test.ts b/extensions/ql-vscode/test/unit-tests/model-editor/auto-model.test.ts
@@ -1,16 +1,12 @@
 import {
   createAutoModelRequest,
   encodeSarif,
-  getCandidates,
 } from "../../../src/model-editor/auto-model";
 import { Mode } from "../../../src/model-editor/shared/mode";
 import { AutomodelMode } from "../../../src/model-editor/auto-model-api";
 import type { AutoModelQueriesResult } from "../../../src/model-editor/auto-model-codeml-queries";
 import type { Log } from "sarif";
 import { gzipDecode } from "../../../src/common/zlib";
-import type { Method } from "../../../src/model-editor/method";
-import { EndpointType } from "../../../src/model-editor/method";
-import type { ModeledMethod } from "../../../src/model-editor/modeled-method";
 
 describe("createAutoModelRequest", () => {
   const createSarifLog = (queryId: string): Log => {
@@ -84,118 +80,3 @@ describe("createAutoModelRequest", () => {
     expect(parsed).toEqual(result.candidates);
   });
 });
-
-describe("getCandidates", () => {
-  it("doesn't return methods that are already modelled", () => {
-    const methods: Method[] = [
-      {
-        library: "my.jar",
-        signature: "org.my.A#x()",
-        endpointType: EndpointType.Method,
-        packageName: "org.my",
-        typeName: "A",
-        methodName: "x",
-        methodParameters: "()",
-        supported: false,
-        supportedType: "none",
-        usages: [],
-      },
-    ];
-    const modeledMethods: Record<string, ModeledMethod[]> = {
-      "org.my.A#x()": [
-        {
-          type: "neutral",
-          kind: "sink",
-          provenance: "manual",
-          signature: "org.my.A#x()",
-          endpointType: EndpointType.Method,
-          packageName: "org.my",
-          typeName: "A",
-          methodName: "x",
-          methodParameters: "()",
-        },
-      ],
-    };
-    const candidates = getCandidates(
-      Mode.Application,
-      methods,
-      modeledMethods,
-      new Set(),
-    );
-    expect(candidates.length).toEqual(0);
-  });
-
-  it("doesn't return methods that are supported from other sources", () => {
-    const methods: Method[] = [
-      {
-        library: "my.jar",
-        signature: "org.my.A#x()",
-        endpointType: EndpointType.Method,
-        packageName: "org.my",
-        typeName: "A",
-        methodName: "x",
-        methodParameters: "()",
-        supported: true,
-        supportedType: "none",
-        usages: [],
-      },
-    ];
-    const modeledMethods = {};
-    const candidates = getCandidates(
-      Mode.Application,
-      methods,
-      modeledMethods,
-      new Set(),
-    );
-    expect(candidates.length).toEqual(0);
-  });
-
-  it("doesn't return methods that are already processed by auto model", () => {
-    const methods: Method[] = [
-      {
-        library: "my.jar",
-        signature: "org.my.A#x()",
-        endpointType: EndpointType.Method,
-        packageName: "org.my",
-        typeName: "A",
-        methodName: "x",
-        methodParameters: "()",
-        supported: false,
-        supportedType: "none",
-        usages: [],
-      },
-    ];
-    const modeledMethods = {};
-    const candidates = getCandidates(
-      Mode.Application,
-      methods,
-      modeledMethods,
-      new Set(["org.my.A#x()"]),
-    );
-    expect(candidates.length).toEqual(0);
-  });
-
-  it("returns methods that are neither modeled nor supported from other sources", () => {
-    const methods: Method[] = [];
-    methods.push({
-      library: "my.jar",
-      signature: "org.my.A#x()",
-      endpointType: EndpointType.Method,
-      packageName: "org.my",
-      typeName: "A",
-      methodName: "x",
-      methodParameters: "()",
-      supported: false,
-      supportedType: "none",
-      usages: [],
-    });
-    const modeledMethods = {};
-    const candidates = getCandidates(
-      Mode.Application,
-      methods,
-      modeledMethods,
-      new Set(),
-    );
-    expect(candidates.length).toEqual(1);
-  });
-});
diff --git a/extensions/ql-vscode/test/unit-tests/model-editor/shared/auto-model-candidates.test.ts b/extensions/ql-vscode/test/unit-tests/model-editor/shared/auto-model-candidates.test.ts
