Merge branch 'main' into robertbrignull/ResultTables-Header

extensions/ql-vscode/CHANGELOG.md

@@ -7,6 +7,7 @@
 - Remove "last updated" information and sorting from variant analysis results view. [#2637](https://github.com/github/vscode-codeql/pull/2637)
 - Links to code on GitHub now include column numbers as well as line numbers. [#2406](https://github.com/github/vscode-codeql/pull/2406)
 - No longer highlight trailing commas for jump to definition.  [#2615](https://github.com/github/vscode-codeql/pull/2615)
+- Fix a bug where the QHelp preview page was not being refreshed after changes to the underlying `.qhelp` file. [#2660](https://github.com/github/vscode-codeql/pull/2660)
 
 ## 1.8.8 - 17 July 2023
 

extensions/ql-vscode/package.json

@@ -1655,6 +1655,13 @@
           "title": "CodeQL",
           "icon": "media/logo.svg"
         }
+      ],
+      "panel": [
+        {
+          "id": "codeql-model-details",
+          "title": "CodeQL Model Details",
+          "icon": "media/logo.svg"
+        }
       ]
     },
     "views": {
@@ -1685,9 +1692,20 @@
           "name": "Evaluator Log Viewer",
           "when": "config.codeQL.canary"
         }
+      ],
+      "codeql-model-details": [
+        {
+          "id": "codeQLModelDetails",
+          "name": "CodeQL Model Details",
+          "when": "config.codeQL.canary && config.codeQL.dataExtensions.modelDetailsView"
+        }
       ]
     },
     "viewsWelcome": [
+      {
+        "view": "codeQLModelDetails",
+        "contents": "Loading..."
+      },
       {
         "view": "codeQLAstViewer",
         "contents": "Run the 'CodeQL: View AST' command on an open source file from a CodeQL database.\n[View AST](command:codeQL.viewAst)"
@@ -1756,8 +1774,6 @@
     "fs-extra": "^11.1.1",
     "immutable": "^4.0.0",
     "js-yaml": "^4.1.0",
-    "minimatch": "^9.0.0",
-    "minimist": "^1.2.6",
     "msw": "^1.2.0",
     "nanoid": "^3.2.0",
     "node-fetch": "^2.6.7",
@@ -1767,10 +1783,8 @@
     "semver": "^7.5.2",
     "source-map": "^0.7.4",
     "source-map-support": "^0.5.21",
-    "stream": "^0.0.2",
-    "stream-chain": "^2.2.4",
     "stream-json": "^1.7.3",
-    "styled-components": "^5.3.3",
+    "styled-components": "^6.0.2",
     "tmp": "^0.1.0",
     "tmp-promise": "^3.0.2",
     "tree-kill": "^1.2.2",
@@ -1811,13 +1825,10 @@
     "@types/d3-graphviz": "^2.6.6",
     "@types/del": "^4.0.0",
     "@types/fs-extra": "^11.0.1",
-    "@types/google-protobuf": "^3.2.7",
     "@types/gulp": "^4.0.9",
     "@types/gulp-replace": "^1.1.0",
-    "@types/gulp-sourcemaps": "0.0.32",
     "@types/jest": "^29.0.2",
     "@types/js-yaml": "^3.12.5",
-    "@types/jszip": "^3.1.6",
     "@types/nanoid": "^3.0.0",
     "@types/node": "^16.11.25",
     "@types/node-fetch": "^2.5.2",
@@ -1825,7 +1836,6 @@
     "@types/react-dom": "^18.0.11",
     "@types/sarif": "^2.1.2",
     "@types/semver": "^7.2.0",
-    "@types/stream-chain": "^2.0.1",
     "@types/stream-json": "^1.7.1",
     "@types/styled-components": "^5.1.11",
     "@types/tar-stream": "^2.2.2",
@@ -1835,19 +1845,19 @@
     "@types/vscode": "^1.67.0",
     "@types/webpack": "^5.28.0",
     "@types/webpack-env": "^1.18.0",
-    "@typescript-eslint/eslint-plugin": "^5.38.0",
-    "@typescript-eslint/parser": "^5.38.0",
+    "@typescript-eslint/eslint-plugin": "^6.2.1",
+    "@typescript-eslint/parser": "^6.2.1",
     "@vscode/test-electron": "^2.2.0",
     "@vscode/vsce": "^2.19.0",
     "ansi-colors": "^4.1.1",
     "applicationinsights": "^2.3.5",
-    "cosmiconfig": "^7.1.0",
+    "cosmiconfig": "^8.2.0",
     "cross-env": "^7.0.3",
     "css-loader": "^6.8.1",
     "del": "^6.0.0",
     "esbuild": "^0.15.15",
     "eslint": "^8.23.1",
-    "eslint-config-prettier": "^8.5.0",
+    "eslint-config-prettier": "^9.0.0",
     "eslint-plugin-etc": "^2.0.2",
     "eslint-plugin-github": "^4.4.1",
     "eslint-plugin-jest-dom": "^5.0.1",
@@ -1871,7 +1881,7 @@
     "markdownlint-cli2-formatter-pretty": "^0.0.4",
     "mini-css-extract-plugin": "^2.6.1",
     "npm-run-all": "^4.1.5",
-    "patch-package": "^7.0.0",
+    "patch-package": "^8.0.0",
     "prettier": "^3.0.0",
     "storybook": "^7.1.0",
     "tar-stream": "^3.0.0",
@@ -1880,8 +1890,7 @@
     "ts-json-schema-generator": "^1.1.2",
     "ts-loader": "^9.4.2",
     "ts-node": "^10.7.0",
-    "ts-protoc-gen": "^0.9.0",
-    "ts-unused-exports": "^9.0.5",
+    "ts-unused-exports": "^10.0.0",
     "typescript": "^5.0.2",
     "webpack": "^5.76.0",
     "webpack-cli": "^5.0.1"

extensions/ql-vscode/src/common/commands.ts

@@ -12,6 +12,7 @@ import type {
 } from "../variant-analysis/shared/variant-analysis";
 import type { QLDebugConfiguration } from "../debugger/debug-configuration";
 import type { QueryTreeViewItem } from "../queries-panel/query-tree-view-item";
+import type { Usage } from "../data-extensions-editor/external-api-usage";
 
 // A command function matching the signature that VS Code calls when
 // a command is invoked from a context menu on a TreeView with
@@ -58,7 +59,9 @@ export type ExplorerSelectionCommandFunction<Item> = (
 type BuiltInVsCodeCommands = {
   // The codeQLDatabases.focus command is provided by VS Code because we've registered the custom view
   "codeQLDatabases.focus": () => Promise<void>;
+  "codeQLModelDetails.focus": () => Promise<void>;
   "markdown.showPreviewToSide": (uri: Uri) => Promise<void>;
+  "workbench.action.closeActiveEditor": () => Promise<void>;
   revealFileInOS: (uri: Uri) => Promise<void>;
   setContext: (
     key: `${"codeql" | "codeQL"}${string}`,
@@ -302,6 +305,10 @@ export type PackagingCommands = {
 
 export type DataExtensionsEditorCommands = {
   "codeQL.openDataExtensionsEditor": () => Promise<void>;
+  "codeQLDataExtensionsEditor.jumpToUsageLocation": (
+    usage: Usage,
+    databaseItem: DatabaseItem,
+  ) => Promise<void>;
 };
 
 export type EvalLogViewerCommands = {

extensions/ql-vscode/src/common/interface-types.ts

@@ -510,6 +510,12 @@ interface AddModeledMethodsMessage {
   modeledMethods: Record<string, ModeledMethod>;
 }
 
+interface SetInProgressMethodsMessage {
+  t: "setInProgressMethods";
+  packageName: string;
+  inProgressMethods: string[];
+}
+
 interface SwitchModeMessage {
   t: "switchMode";
   mode: Mode;
@@ -544,10 +550,16 @@ interface GenerateExternalApiMessage {
 
 interface GenerateExternalApiFromLlmMessage {
   t: "generateExternalApiFromLlm";
+  packageName: string;
   externalApiUsages: ExternalApiUsage[];
   modeledMethods: Record<string, ModeledMethod>;
 }
 
+interface StopGeneratingExternalApiFromLlmMessage {
+  t: "stopGeneratingExternalApiFromLlm";
+  packageName: string;
+}
+
 interface ModelDependencyMessage {
   t: "modelDependency";
 }
@@ -556,7 +568,8 @@ export type ToDataExtensionsEditorMessage =
   | SetExtensionPackStateMessage
   | SetExternalApiUsagesMessage
   | LoadModeledMethodsMessage
-  | AddModeledMethodsMessage;
+  | AddModeledMethodsMessage
+  | SetInProgressMethodsMessage;
 
 export type FromDataExtensionsEditorMessage =
   | ViewLoadedMsg
@@ -568,4 +581,5 @@ export type FromDataExtensionsEditorMessage =
   | SaveModeledMethods
   | GenerateExternalApiMessage
   | GenerateExternalApiFromLlmMessage
+  | StopGeneratingExternalApiFromLlmMessage
   | ModelDependencyMessage;

extensions/ql-vscode/src/common/sarif-utils.ts

@@ -2,7 +2,7 @@ import * as Sarif from "sarif";
 import type { HighlightedRegion } from "../variant-analysis/shared/analysis-result";
 import { ResolvableLocationValue } from "../common/bqrs-cli-types";
 
-interface SarifLink {
+export interface SarifLink {
   dest: number;
   text: string;
 }

extensions/ql-vscode/src/config.ts

@@ -714,6 +714,11 @@ const EXTENSIONS_DIRECTORY = new Setting(
   "extensionsDirectory",
   DATA_EXTENSIONS,
 );
+const MODEL_DETAILS_VIEW = new Setting("modelDetailsView", DATA_EXTENSIONS);
+
+export function showModelDetailsView(): boolean {
+  return !!MODEL_DETAILS_VIEW.getValue<boolean>();
+}
 
 export function showLlmGeneration(): boolean {
   return !!LLM_GENERATION.getValue<boolean>();

extensions/ql-vscode/src/data-extensions-editor/auto-model-codeml-queries.ts

@@ -17,6 +17,10 @@ import { redactableError } from "../common/errors";
 import { interpretResultsSarif } from "../query-results";
 import { join } from "path";
 import { assertNever } from "../common/helpers-pure";
+import { dir } from "tmp-promise";
+import { writeFile, outputFile } from "fs-extra";
+import { dump as dumpYaml } from "js-yaml";
+import { MethodSignature } from "./external-api-usage";
 
 type AutoModelQueryOptions = {
   queryTag: string;
@@ -26,6 +30,7 @@ type AutoModelQueryOptions = {
   databaseItem: DatabaseItem;
   qlpack: QlPacksForLanguage;
   sourceInfo: SourceInfo | undefined;
+  additionalPacks: string[];
   extensionPacks: string[];
   queryStorageDir: string;
 
@@ -52,6 +57,7 @@ async function runAutoModelQuery({
   databaseItem,
   qlpack,
   sourceInfo,
+  additionalPacks,
   extensionPacks,
   queryStorageDir,
   progress,
@@ -99,7 +105,7 @@ async function runAutoModelQuery({
       quickEvalCountOnly: false,
     },
     false,
-    getOnDiskWorkspaceFolders(),
+    additionalPacks,
     extensionPacks,
     queryStorageDir,
     undefined,
@@ -147,12 +153,14 @@ async function runAutoModelQuery({
 
 type AutoModelQueriesOptions = {
   mode: Mode;
+  candidateMethods: MethodSignature[];
   cliServer: CodeQLCliServer;
   queryRunner: QueryRunner;
   databaseItem: DatabaseItem;
   queryStorageDir: string;
 
   progress: ProgressCallback;
+  cancellationTokenSource: CancellationTokenSource;
 };
 
 export type AutoModelQueriesResult = {
@@ -161,17 +169,14 @@ export type AutoModelQueriesResult = {
 
 export async function runAutoModelQueries({
   mode,
+  candidateMethods,
   cliServer,
   queryRunner,
   databaseItem,
   queryStorageDir,
   progress,
+  cancellationTokenSource,
 }: AutoModelQueriesOptions): Promise<AutoModelQueriesResult | undefined> {
-  // maxStep for this part is 1500
-  const maxStep = 1500;
-
-  const cancellationTokenSource = new CancellationTokenSource();
-
   const qlpack = await qlpackOfDatabase(cliServer, databaseItem);
 
   // CodeQL needs to have access to the database to be able to retrieve the
@@ -189,17 +194,17 @@ export async function runAutoModelQueries({
           sourceLocationPrefix,
         };
 
-  const additionalPacks = getOnDiskWorkspaceFolders();
+  // Generate a pack containing the candidate filters
+  const filterPackDir = await generateCandidateFilterPack(
+    databaseItem.language,
+    candidateMethods,
+  );
+
+  const additionalPacks = [...getOnDiskWorkspaceFolders(), filterPackDir];
   const extensionPacks = Object.keys(
     await cliServer.resolveQlpacks(additionalPacks, true),
   );
 
-  progress({
-    step: 0,
-    maxStep,
-    message: "Finding candidates and examples",
-  });
-
   const candidates = await runAutoModelQuery({
     mode,
     queryTag: "candidates",
@@ -208,15 +213,10 @@ export async function runAutoModelQueries({
     databaseItem,
     qlpack,
     sourceInfo,
+    additionalPacks,
     extensionPacks,
     queryStorageDir,
-    progress: (update) => {
-      progress({
-        step: update.step,
-        maxStep,
-        message: "Finding candidates and examples",
-      });
-    },
+    progress,
     token: cancellationTokenSource.token,
   });
 
@@ -228,3 +228,59 @@ export async function runAutoModelQueries({
     candidates,
   };
 }
+
+/**
+ * generateCandidateFilterPack will create a temporary extension pack.
+ * This pack will contain a filter that will restrict the automodel queries
+ * to the specified candidate methods only.
+ * This is done using the `extensible` predicate "automodelCandidateFilter".
+ * @param language
+ * @param candidateMethods
+ * @returns
+ */
+export async function generateCandidateFilterPack(
+  language: string,
+  candidateMethods: MethodSignature[],
+): Promise<string> {
+  // Pack resides in a temporary directory, to not pollute the workspace.
+  const packDir = (await dir({ unsafeCleanup: true })).path;
+
+  const syntheticConfigPack = {
+    name: "codeql/automodel-filter",
+    version: "0.0.0",
+    library: true,
+    extensionTargets: {
+      [`codeql/${language}-queries`]: "*",
+    },
+    dataExtensions: ["filter.yml"],
+  };
+
+  const qlpackFile = join(packDir, "codeql-pack.yml");
+  await outputFile(qlpackFile, dumpYaml(syntheticConfigPack), "utf8");
+
+  // The predicate has the following defintion:
+  // extensible predicate automodelCandidateFilter(string package, string type, string name, string signature)
+  const dataRows = candidateMethods.map((method) => [
+    method.packageName,
+    method.typeName,
+    method.methodName,
+    method.methodParameters,
+  ]);
+
+  const filter = {
+    extensions: [
+      {
+        addsTo: {
+          pack: `codeql/${language}-queries`,
+          extensible: "automodelCandidateFilter",
+        },
+        data: dataRows,
+      },
+    ],
+  };
+
+  const filterFile = join(packDir, "filter.yml");
+  await writeFile(filterFile, dumpYaml(filter), "utf8");
+
+  return packDir;
+}

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

@@ -4,6 +4,55 @@ import { AutoModelQueriesResult } from "./auto-model-codeml-queries";
 import { assertNever } from "../common/helpers-pure";
 import * as Sarif from "sarif";
 import { gzipEncode } from "../common/zlib";
+import { ExternalApiUsage, MethodSignature } from "./external-api-usage";
+import { 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 externalApiUsages all external API usages.
+ * @param modeledMethods the currently modeled methods.
+ * @returns list of modeled methods that are candidates for modeling.
+ */
+export function getCandidates(
+  mode: Mode,
+  externalApiUsages: ExternalApiUsage[],
+  modeledMethods: Record<string, ModeledMethod>,
+): MethodSignature[] {
+  // Sort the same way as the UI so we send the first ones listed in the UI first
+  const grouped = groupMethods(externalApiUsages, mode);
+  const sortedGroupNames = sortGroupNames(grouped);
+  const sortedExternalApiUsages = sortedGroupNames.flatMap((name) =>
+    sortMethods(grouped[name]),
+  );
+
+  const candidates: MethodSignature[] = [];
+
+  for (const externalApiUsage of sortedExternalApiUsages) {
+    const modeledMethod: ModeledMethod = modeledMethods[
+      externalApiUsage.signature
+    ] ?? {
+      type: "none",
+    };
+
+    // Anything that is modeled is not a candidate
+    if (modeledMethod.type !== "none") {
+      continue;
+    }
+
+    // A method that is supported is modeled outside of the model file, so it is not a candidate.
+    if (externalApiUsage.supported) {
+      continue;
+    }
+
+    // The rest are candidates
+    candidates.push(externalApiUsage);
+  }
+  return candidates;
+}
 
 /**
  * Encode a SARIF log to the format expected by the server: JSON, GZIP-compressed, base64-encoded

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

@@ -0,0 +1,246 @@
+import { ExternalApiUsage, MethodSignature } from "./external-api-usage";
+import { ModeledMethod } from "./modeled-method";
+import { extLogger } from "../common/logging/vscode";
+import { load as loadYaml } from "js-yaml";
+import { ProgressCallback, withProgress } from "../common/vscode/progress";
+import { createAutoModelV2Request, getCandidates } from "./auto-model-v2";
+import { runAutoModelQueries } from "./auto-model-codeml-queries";
+import { loadDataExtensionYaml } from "./yaml";
+import { ModelRequest, ModelResponse, autoModelV2 } from "./auto-model-api-v2";
+import { RequestError } from "@octokit/request-error";
+import { showAndLogExceptionWithTelemetry } from "../common/logging";
+import { redactableError } from "../common/errors";
+import { App } from "../common/app";
+import { CodeQLCliServer } from "../codeql-cli/cli";
+import { QueryRunner } from "../query-server";
+import { DatabaseItem } from "../databases/local-databases";
+import { Mode } from "./shared/mode";
+import { CancellationTokenSource } from "vscode";
+
+// Limit the number of candidates we send to the model in each request
+// to avoid long requests.
+// Note that the model may return fewer than this number of candidates.
+const candidateBatchSize = 20;
+
+/**
+ * The auto-modeler holds state around auto-modeling jobs and allows
+ * starting and stopping them.
+ */
+export class AutoModeler {
+  // Keep track of auto-modeling jobs that are in progress
+  // so that we can stop them.
+  private readonly jobs: Map<string, CancellationTokenSource>;
+
+  constructor(
+    private readonly app: App,
+    private readonly cliServer: CodeQLCliServer,
+    private readonly queryRunner: QueryRunner,
+    private readonly queryStorageDir: string,
+    private readonly databaseItem: DatabaseItem,
+    private readonly setInProgressMethods: (
+      packageName: string,
+      inProgressMethods: string[],
+    ) => Promise<void>,
+    private readonly addModeledMethods: (
+      modeledMethods: Record<string, ModeledMethod>,
+    ) => Promise<void>,
+  ) {
+    this.jobs = new Map<string, CancellationTokenSource>();
+  }
+
+  /**
+   * Models the given package's external API usages, except
+   * the ones that are already modeled.
+   * @param packageName The name of the package to model.
+   * @param externalApiUsages The external API usages.
+   * @param modeledMethods The currently modeled methods.
+   * @param mode The mode we are modeling in.
+   */
+  public async startModeling(
+    packageName: string,
+    externalApiUsages: ExternalApiUsage[],
+    modeledMethods: Record<string, ModeledMethod>,
+    mode: Mode,
+  ): Promise<void> {
+    if (this.jobs.has(packageName)) {
+      return;
+    }
+
+    const cancellationTokenSource = new CancellationTokenSource();
+    this.jobs.set(packageName, cancellationTokenSource);
+
+    try {
+      await this.modelPackage(
+        packageName,
+        externalApiUsages,
+        modeledMethods,
+        mode,
+        cancellationTokenSource,
+      );
+    } finally {
+      this.jobs.delete(packageName);
+    }
+  }
+
+  /**
+   * Stops modeling the given package.
+   * @param packageName The name of the package to stop modeling.
+   */
+  public async stopModeling(packageName: string): Promise<void> {
+    void extLogger.log(`Stopping modeling for package ${packageName}`);
+    const cancellationTokenSource = this.jobs.get(packageName);
+    if (cancellationTokenSource) {
+      cancellationTokenSource.cancel();
+    }
+  }
+
+  /**
+   * Stops all in-progress modeling jobs.
+   */
+  public async stopAllModeling(): Promise<void> {
+    for (const cancellationTokenSource of this.jobs.values()) {
+      cancellationTokenSource.cancel();
+    }
+  }
+
+  private async modelPackage(
+    packageName: string,
+    externalApiUsages: ExternalApiUsage[],
+    modeledMethods: Record<string, ModeledMethod>,
+    mode: Mode,
+    cancellationTokenSource: CancellationTokenSource,
+  ): Promise<void> {
+    void extLogger.log(`Modeling package ${packageName}`);
+    await withProgress(async (progress) => {
+      // Fetch the candidates to send to the model
+      const allCandidateMethods = getCandidates(
+        mode,
+        externalApiUsages,
+        modeledMethods,
+      );
+
+      // If there are no candidates, there is nothing to model and we just return
+      if (allCandidateMethods.length === 0) {
+        void extLogger.log("No candidates to model. Stopping.");
+        return;
+      }
+
+      // Find number of slices to make
+      const batchNumber = Math.ceil(
+        allCandidateMethods.length / candidateBatchSize,
+      );
+      try {
+        for (let i = 0; i < batchNumber; i++) {
+          // Check if we should stop
+          if (cancellationTokenSource.token.isCancellationRequested) {
+            break;
+          }
+
+          const start = i * candidateBatchSize;
+          const end = start + candidateBatchSize;
+          const candidatesToProcess = allCandidateMethods.slice(start, end);
+
+          // Let the UI know which candidates we are modeling
+          await this.setInProgressMethods(
+            packageName,
+            candidatesToProcess.map((c) => c.signature),
+          );
+
+          // Kick off the process to model the slice of candidates
+          await this.modelCandidates(
+            candidatesToProcess,
+            mode,
+            progress,
+            cancellationTokenSource,
+          );
+        }
+      } finally {
+        // Clear out in progress methods
+        await this.setInProgressMethods(packageName, []);
+      }
+    });
+  }
+
+  private async modelCandidates(
+    candidateMethods: MethodSignature[],
+    mode: Mode,
+    progress: ProgressCallback,
+    cancellationTokenSource: CancellationTokenSource,
+  ): Promise<void> {
+    void extLogger.log("Executing auto-model queries");
+
+    const usages = await runAutoModelQueries({
+      mode,
+      candidateMethods,
+      cliServer: this.cliServer,
+      queryRunner: this.queryRunner,
+      queryStorageDir: this.queryStorageDir,
+      databaseItem: this.databaseItem,
+      progress: (update) => progress({ ...update }),
+      cancellationTokenSource,
+    });
+    if (!usages) {
+      return;
+    }
+
+    const request = await createAutoModelV2Request(mode, usages);
+
+    void extLogger.log("Calling auto-model API");
+
+    const response = await this.callAutoModelApi(request);
+    if (!response) {
+      return;
+    }
+
+    const models = loadYaml(response.models, {
+      filename: "auto-model.yml",
+    });
+
+    const loadedMethods = loadDataExtensionYaml(models);
+    if (!loadedMethods) {
+      return;
+    }
+
+    // Any candidate that was part of the response is a negative result
+    // meaning that the canidate is not a sink for the kinds that the LLM is checking for.
+    // For now we model this as a sink neutral method, however this is subject
+    // to discussion.
+    for (const candidate of candidateMethods) {
+      if (!(candidate.signature in loadedMethods)) {
+        loadedMethods[candidate.signature] = {
+          type: "neutral",
+          kind: "sink",
+          input: "",
+          output: "",
+          provenance: "ai-generated",
+          signature: candidate.signature,
+          packageName: candidate.packageName,
+          typeName: candidate.typeName,
+          methodName: candidate.methodName,
+          methodParameters: candidate.methodParameters,
+        };
+      }
+    }
+
+    await this.addModeledMethods(loadedMethods);
+  }
+
+  private async callAutoModelApi(
+    request: ModelRequest,
+  ): Promise<ModelResponse | null> {
+    try {
+      return await autoModelV2(this.app.credentials, request);
+    } catch (e) {
+      if (e instanceof RequestError && e.status === 429) {
+        void showAndLogExceptionWithTelemetry(
+          this.app.logger,
+          this.app.telemetry,
+          redactableError(e)`Rate limit hit, please try again soon.`,
+        );
+        return null;
+      } else {
+        throw e;
+      }
+    }
+  }
+}

extensions/ql-vscode/src/data-extensions-editor/data-extensions-editor-module.ts

@@ -3,7 +3,7 @@ import { DataExtensionsEditorView } from "./data-extensions-editor-view";
 import { DataExtensionsEditorCommands } from "../common/commands";
 import { CliVersionConstraint, CodeQLCliServer } from "../codeql-cli/cli";
 import { QueryRunner } from "../query-server";
-import { DatabaseManager } from "../databases/local-databases";
+import { DatabaseItem, DatabaseManager } from "../databases/local-databases";
 import { ensureDir } from "fs-extra";
 import { join } from "path";
 import { App } from "../common/app";
@@ -20,11 +20,17 @@ import { redactableError } from "../common/errors";
 import { extLogger } from "../common/logging/vscode";
 import { isQueryLanguage } from "../common/query-language";
 import { setUpPack } from "./external-api-usage-query";
+import { DisposableObject } from "../common/disposable-object";
+import { ModelDetailsPanel } from "./model-details/model-details-panel";
+import { Mode } from "./shared/mode";
+import { showResolvableLocation } from "../databases/local-databases/locations";
+import { Usage } from "./external-api-usage";
 
 const SUPPORTED_LANGUAGES: string[] = ["java", "csharp"];
 
-export class DataExtensionsEditorModule {
+export class DataExtensionsEditorModule extends DisposableObject {
   private readonly queryStorageDir: string;
+  private readonly modelDetailsPanel: ModelDetailsPanel;
 
   private constructor(
     private readonly ctx: ExtensionContext,
@@ -34,10 +40,12 @@ export class DataExtensionsEditorModule {
     private readonly queryRunner: QueryRunner,
     baseQueryStorageDir: string,
   ) {
+    super();
     this.queryStorageDir = join(
       baseQueryStorageDir,
       "data-extensions-editor-results",
     );
+    this.modelDetailsPanel = this.push(new ModelDetailsPanel(cliServer));
   }
 
   public static async initialize(
@@ -138,6 +146,8 @@ export class DataExtensionsEditorModule {
               queryDir,
               db,
               modelFile,
+              Mode.Application,
+              this.modelDetailsPanel.setState.bind(this.modelDetailsPanel),
             );
             await view.openView();
           },
@@ -146,6 +156,12 @@ export class DataExtensionsEditorModule {
           },
         );
       },
+      "codeQLDataExtensionsEditor.jumpToUsageLocation": async (
+        usage: Usage,
+        databaseItem: DatabaseItem,
+      ) => {
+        await showResolvableLocation(usage.url, databaseItem, this.app.logger);
+      },
     };
   }
 

extensions/ql-vscode/src/data-extensions-editor/data-extensions-editor-view.ts

@@ -35,11 +35,6 @@ import { ExternalApiUsage } from "./external-api-usage";
 import { ModeledMethod } from "./modeled-method";
 import { ExtensionPack } from "./shared/extension-pack";
 import { autoModel, ModelRequest, ModelResponse } from "./auto-model-api";
-import {
-  autoModelV2,
-  ModelRequest as ModelRequestV2,
-  ModelResponse as ModelResponseV2,
-} from "./auto-model-api-v2";
 import {
   createAutoModelRequest,
   parsePredictedClassifications,
@@ -47,6 +42,7 @@ import {
 import {
   enableFrameworkMode,
   showLlmGeneration,
+  showModelDetailsView,
   useLlmGenerationV2,
 } from "../config";
 import { getAutoModelUsages } from "./auto-model-usages-query";
@@ -55,15 +51,14 @@ import { loadModeledMethods, saveModeledMethods } from "./modeled-method-fs";
 import { join } from "path";
 import { pickExtensionPack } from "./extension-pack-picker";
 import { getLanguageDisplayName } from "../common/query-language";
-import { runAutoModelQueries } from "./auto-model-codeml-queries";
-import { createAutoModelV2Request } from "./auto-model-v2";
-import { load as loadYaml } from "js-yaml";
-import { loadDataExtensionYaml } from "./yaml";
+import { AutoModeler } from "./auto-modeler";
 
 export class DataExtensionsEditorView extends AbstractWebview<
   ToDataExtensionsEditorMessage,
   FromDataExtensionsEditorMessage
 > {
+  private readonly autoModeler: AutoModeler;
+
   public constructor(
     ctx: ExtensionContext,
     private readonly app: App,
@@ -74,9 +69,31 @@ export class DataExtensionsEditorView extends AbstractWebview<
     private readonly queryDir: string,
     private readonly databaseItem: DatabaseItem,
     private readonly extensionPack: ExtensionPack,
-    private mode: Mode = Mode.Application,
+    private mode: Mode,
+    private readonly updateModelDetailsPanelState: (
+      externalApiUsages: ExternalApiUsage[],
+      databaseItem: DatabaseItem,
+    ) => Promise<void>,
   ) {
     super(ctx);
+
+    this.autoModeler = new AutoModeler(
+      app,
+      cliServer,
+      queryRunner,
+      queryStorageDir,
+      databaseItem,
+      async (packageName, inProgressMethods) => {
+        await this.postMessage({
+          t: "setInProgressMethods",
+          packageName,
+          inProgressMethods,
+        });
+      },
+      async (modeledMethods) => {
+        await this.postMessage({ t: "addModeledMethods", modeledMethods });
+      },
+    );
   }
 
   public async openView() {
@@ -137,7 +154,7 @@ export class DataExtensionsEditorView extends AbstractWebview<
 
         break;
       case "jumpToUsage":
-        await this.jumpToUsage(msg.location);
+        await this.handleJumpToUsage(msg.location);
 
         break;
       case "saveModeledMethods":
@@ -159,15 +176,24 @@ export class DataExtensionsEditorView extends AbstractWebview<
 
         break;
       case "generateExternalApiFromLlm":
-        await this.generateModeledMethodsFromLlm(
-          msg.externalApiUsages,
-          msg.modeledMethods,
-        );
-
+        if (useLlmGenerationV2()) {
+          await this.generateModeledMethodsFromLlmV2(
+            msg.packageName,
+            msg.externalApiUsages,
+            msg.modeledMethods,
+          );
+        } else {
+          await this.generateModeledMethodsFromLlmV1(
+            msg.externalApiUsages,
+            msg.modeledMethods,
+          );
+        }
+        break;
+      case "stopGeneratingExternalApiFromLlm":
+        await this.autoModeler.stopModeling(msg.packageName);
         break;
       case "modelDependency":
         await this.modelDependency();
-
         break;
       case "switchMode":
         this.mode = msg.mode;
@@ -205,24 +231,22 @@ export class DataExtensionsEditorView extends AbstractWebview<
     });
   }
 
+  protected async handleJumpToUsage(location: ResolvableLocationValue) {
+    if (showModelDetailsView()) {
+      await this.openModelDetailsView();
+    } else {
+      await this.jumpToUsage(location);
+    }
+  }
+
+  protected async openModelDetailsView() {
+    await this.app.commands.execute("codeQLModelDetails.focus");
+  }
+
   protected async jumpToUsage(
     location: ResolvableLocationValue,
   ): Promise<void> {
-    try {
-      await showResolvableLocation(location, this.databaseItem);
-    } catch (e) {
-      if (e instanceof Error) {
-        if (e.message.match(/File not found/)) {
-          void window.showErrorMessage(
-            "Original file of this result is not in the database's source archive.",
-          );
-        } else {
-          void this.app.logger.log(`Unable to handleMsgFromView: ${e.message}`);
-        }
-      } else {
-        void this.app.logger.log(`Unable to handleMsgFromView: ${e}`);
-      }
-    }
+    await showResolvableLocation(location, this.databaseItem, this.app.logger);
   }
 
   protected async loadExistingModeledMethods(): Promise<void> {
@@ -288,6 +312,10 @@ export class DataExtensionsEditorView extends AbstractWebview<
             t: "setExternalApiUsages",
             externalApiUsages,
           });
+          await this.updateModelDetailsPanelState(
+            externalApiUsages,
+            this.databaseItem,
+          );
         } catch (err) {
           void showAndLogExceptionWithTelemetry(
             this.app.logger,
@@ -361,7 +389,7 @@ export class DataExtensionsEditorView extends AbstractWebview<
     );
   }
 
-  private async generateModeledMethodsFromLlm(
+  private async generateModeledMethodsFromLlmV1(
     externalApiUsages: ExternalApiUsage[],
     modeledMethods: Record<string, ModeledMethod>,
   ): Promise<void> {
@@ -374,102 +402,50 @@ export class DataExtensionsEditorView extends AbstractWebview<
         message: "Retrieving usages",
       });
 
-      let predictedModeledMethods: Record<string, ModeledMethod>;
+      const usages = await getAutoModelUsages({
+        cliServer: this.cliServer,
+        queryRunner: this.queryRunner,
+        queryStorageDir: this.queryStorageDir,
+        queryDir: this.queryDir,
+        databaseItem: this.databaseItem,
+        progress: (update) => progress({ ...update, maxStep }),
+      });
 
-      if (useLlmGenerationV2()) {
-        const usages = await runAutoModelQueries({
-          mode: this.mode,
-          cliServer: this.cliServer,
-          queryRunner: this.queryRunner,
-          queryStorageDir: this.queryStorageDir,
-          databaseItem: this.databaseItem,
-          progress: (update) => progress({ ...update, maxStep }),
-        });
-        if (!usages) {
-          return;
-        }
+      progress({
+        step: 1800,
+        maxStep,
+        message: "Creating request",
+      });
 
-        progress({
-          step: 1800,
-          maxStep,
-          message: "Creating request",
-        });
+      const request = createAutoModelRequest(
+        this.databaseItem.language,
+        externalApiUsages,
+        modeledMethods,
+        usages,
+        this.mode,
+      );
 
-        const request = await createAutoModelV2Request(this.mode, usages);
+      progress({
+        step: 2000,
+        maxStep,
+        message: "Sending request",
+      });
 
-        progress({
-          step: 2000,
-          maxStep,
-          message: "Sending request",
-        });
-
-        const response = await this.callAutoModelApiV2(request);
-        if (!response) {
-          return;
-        }
-
-        progress({
-          step: 2500,
-          maxStep,
-          message: "Parsing response",
-        });
-
-        const models = loadYaml(response.models, {
-          filename: "auto-model.yml",
-        });
-
-        const modeledMethods = loadDataExtensionYaml(models);
-        if (!modeledMethods) {
-          return;
-        }
-
-        predictedModeledMethods = modeledMethods;
-      } else {
-        const usages = await getAutoModelUsages({
-          cliServer: this.cliServer,
-          queryRunner: this.queryRunner,
-          queryStorageDir: this.queryStorageDir,
-          queryDir: this.queryDir,
-          databaseItem: this.databaseItem,
-          progress: (update) => progress({ ...update, maxStep }),
-        });
-
-        progress({
-          step: 1800,
-          maxStep,
-          message: "Creating request",
-        });
-
-        const request = createAutoModelRequest(
-          this.databaseItem.language,
-          externalApiUsages,
-          modeledMethods,
-          usages,
-          this.mode,
-        );
-
-        progress({
-          step: 2000,
-          maxStep,
-          message: "Sending request",
-        });
-
-        const response = await this.callAutoModelApi(request);
-        if (!response) {
-          return;
-        }
-
-        progress({
-          step: 2500,
-          maxStep,
-          message: "Parsing response",
-        });
-
-        predictedModeledMethods = parsePredictedClassifications(
-          response.predicted || [],
-        );
+      const response = await this.callAutoModelApi(request);
+      if (!response) {
+        return;
       }
 
+      progress({
+        step: 2500,
+        maxStep,
+        message: "Parsing response",
+      });
+
+      const predictedModeledMethods = parsePredictedClassifications(
+        response.predicted || [],
+      );
+
       progress({
         step: 2800,
         maxStep,
@@ -483,6 +459,19 @@ export class DataExtensionsEditorView extends AbstractWebview<
     });
   }
 
+  private async generateModeledMethodsFromLlmV2(
+    packageName: string,
+    externalApiUsages: ExternalApiUsage[],
+    modeledMethods: Record<string, ModeledMethod>,
+  ): Promise<void> {
+    await this.autoModeler.startModeling(
+      packageName,
+      externalApiUsages,
+      modeledMethods,
+      this.mode,
+    );
+  }
+
   private async modelDependency(): Promise<void> {
     return withProgress(async (progress, token) => {
       const addedDatabase = await this.promptChooseNewOrExistingDatabase(
@@ -514,6 +503,7 @@ export class DataExtensionsEditorView extends AbstractWebview<
         addedDatabase,
         modelFile,
         Mode.Framework,
+        this.updateModelDetailsPanelState,
       );
       await view.openView();
     });
@@ -609,23 +599,4 @@ export class DataExtensionsEditorView extends AbstractWebview<
       }
     }
   }
-
-  private async callAutoModelApiV2(
-    request: ModelRequestV2,
-  ): Promise<ModelResponseV2 | null> {
-    try {
-      return await autoModelV2(this.app.credentials, request);
-    } catch (e) {
-      if (e instanceof RequestError && e.status === 429) {
-        void showAndLogExceptionWithTelemetry(
-          this.app.logger,
-          this.app.telemetry,
-          redactableError(e)`Rate limit hit, please try again soon.`,
-        );
-        return null;
-      } else {
-        throw e;
-      }
-    }
-  }
 }

extensions/ql-vscode/src/data-extensions-editor/external-api-usage.ts

@@ -13,7 +13,7 @@ export enum CallClassification {
   Generated = "generated",
 }
 
-type Usage = Call & {
+export type Usage = Call & {
   classification: CallClassification;
 };
 

extensions/ql-vscode/src/data-extensions-editor/model-details/model-details-data-provider.ts

@@ -0,0 +1,99 @@
+import {
+  Event,
+  EventEmitter,
+  ThemeColor,
+  ThemeIcon,
+  TreeDataProvider,
+  TreeItem,
+  TreeItemCollapsibleState,
+  Uri,
+} from "vscode";
+import { DisposableObject } from "../../common/disposable-object";
+import { ExternalApiUsage, Usage } from "../external-api-usage";
+import { DatabaseItem } from "../../databases/local-databases";
+import { relative } from "path";
+import { CodeQLCliServer } from "../../codeql-cli/cli";
+
+export class ModelDetailsDataProvider
+  extends DisposableObject
+  implements TreeDataProvider<ModelDetailsTreeViewItem>
+{
+  private externalApiUsages: ExternalApiUsage[] = [];
+  private databaseItem: DatabaseItem | undefined = undefined;
+  private sourceLocationPrefix: string | undefined = undefined;
+
+  private readonly onDidChangeTreeDataEmitter = this.push(
+    new EventEmitter<void>(),
+  );
+
+  public constructor(private readonly cliServer: CodeQLCliServer) {
+    super();
+  }
+
+  public get onDidChangeTreeData(): Event<void> {
+    return this.onDidChangeTreeDataEmitter.event;
+  }
+
+  public async setState(
+    externalApiUsages: ExternalApiUsage[],
+    databaseItem: DatabaseItem,
+  ): Promise<void> {
+    this.externalApiUsages = externalApiUsages;
+    this.databaseItem = databaseItem;
+    this.sourceLocationPrefix = await this.databaseItem.getSourceLocationPrefix(
+      this.cliServer,
+    );
+    this.onDidChangeTreeDataEmitter.fire();
+  }
+
+  getTreeItem(item: ModelDetailsTreeViewItem): TreeItem {
+    if (isExternalApiUsage(item)) {
+      return {
+        label: `${item.packageName}.${item.typeName}.${item.methodName}${item.methodParameters}`,
+        collapsibleState: TreeItemCollapsibleState.Collapsed,
+        iconPath: new ThemeIcon("symbol-method"),
+      };
+    } else {
+      return {
+        label: item.label,
+        description: `${this.relativePathWithinDatabase(item.url.uri)} [${
+          item.url.startLine
+        }, ${item.url.endLine}]`,
+        collapsibleState: TreeItemCollapsibleState.None,
+        command: {
+          title: "Show usage",
+          command: "codeQLDataExtensionsEditor.jumpToUsageLocation",
+          arguments: [item, this.databaseItem],
+        },
+        iconPath: new ThemeIcon("error", new ThemeColor("errorForeground")),
+      };
+    }
+  }
+
+  private relativePathWithinDatabase(uri: string): string {
+    const parsedUri = Uri.parse(uri);
+    if (this.sourceLocationPrefix) {
+      return relative(this.sourceLocationPrefix, parsedUri.fsPath);
+    } else {
+      return parsedUri.fsPath;
+    }
+  }
+
+  getChildren(item?: ModelDetailsTreeViewItem): ModelDetailsTreeViewItem[] {
+    if (item === undefined) {
+      return this.externalApiUsages;
+    } else if (isExternalApiUsage(item)) {
+      return item.usages;
+    } else {
+      return [];
+    }
+  }
+}
+
+type ModelDetailsTreeViewItem = ExternalApiUsage | Usage;
+
+function isExternalApiUsage(
+  item: ModelDetailsTreeViewItem,
+): item is ExternalApiUsage {
+  return (item as any).usages !== undefined;
+}

extensions/ql-vscode/src/data-extensions-editor/model-details/model-details-panel.ts

@@ -0,0 +1,33 @@
+import { TreeView, window } from "vscode";
+import { DisposableObject } from "../../common/disposable-object";
+import { ModelDetailsDataProvider } from "./model-details-data-provider";
+import { DatabaseItem } from "../../databases/local-databases";
+import { ExternalApiUsage, Usage } from "../external-api-usage";
+import { CodeQLCliServer } from "../../codeql-cli/cli";
+
+export class ModelDetailsPanel extends DisposableObject {
+  private readonly dataProvider: ModelDetailsDataProvider;
+  private readonly treeView: TreeView<ExternalApiUsage | Usage>;
+
+  public constructor(cliServer: CodeQLCliServer) {
+    super();
+
+    this.dataProvider = new ModelDetailsDataProvider(cliServer);
+
+    this.treeView = window.createTreeView("codeQLModelDetails", {
+      treeDataProvider: this.dataProvider,
+    });
+    this.push(this.treeView);
+  }
+
+  public async setState(
+    externalApiUsages: ExternalApiUsage[],
+    databaseItem: DatabaseItem,
+  ): Promise<void> {
+    await this.dataProvider.setState(externalApiUsages, databaseItem);
+    this.treeView.badge = {
+      value: externalApiUsages.length,
+      tooltip: "Number of external APIs",
+    };
+  }
+}

extensions/ql-vscode/src/data-extensions-editor/predicates.ts

@@ -47,7 +47,7 @@ export const extensiblePredicateDefinitions: Record<
       methodName: row[3] as string,
       methodParameters: row[4] as string,
     }),
-    supportedKinds: ["remote"],
+    supportedKinds: ["local", "remote"],
   },
   sink: {
     extensiblePredicate: "sinkModel",
@@ -78,7 +78,19 @@ export const extensiblePredicateDefinitions: Record<
       methodName: row[3] as string,
       methodParameters: row[4] as string,
     }),
-    supportedKinds: ["sql", "xss", "logging"],
+    supportedKinds: [
+      "code-injection",
+      "command-injection",
+      "file-content-store",
+      "html-injection",
+      "js-injection",
+      "ldap-injection",
+      "log-injection",
+      "path-injection",
+      "request-forgery",
+      "sql-injection",
+      "url-redirection",
+    ],
   },
   summary: {
     extensiblePredicate: "summaryModel",

extensions/ql-vscode/src/data-extensions-editor/shared/in-progress-methods.ts

@@ -0,0 +1,31 @@
+/**
+ * A class that keeps track of which methods are in progress for each package.
+ */
+export class InProgressMethods {
+  // A map of in-progress method signatures for each package.
+  private readonly methodMap: Map<string, Set<string>>;
+
+  constructor() {
+    this.methodMap = new Map<string, Set<string>>();
+  }
+
+  public setPackageMethods(packageName: string, methods: Set<string>): void {
+    this.methodMap.set(packageName, methods);
+  }
+
+  public hasMethod(packageName: string, method: string): boolean {
+    const methods = this.methodMap.get(packageName);
+    if (methods) {
+      return methods.has(method);
+    }
+    return false;
+  }
+
+  public static fromExisting(methods: InProgressMethods): InProgressMethods {
+    const newInProgressMethods = new InProgressMethods();
+    methods.methodMap.forEach((value, key) => {
+      newInProgressMethods.methodMap.set(key, new Set<string>(value));
+    });
+    return newInProgressMethods;
+  }
+}

extensions/ql-vscode/src/databases/local-databases/locations.ts

@@ -97,8 +97,19 @@ export function tryResolveLocation(
 export async function showResolvableLocation(
   loc: ResolvableLocationValue,
   databaseItem: DatabaseItem,
+  logger: Logger,
 ): Promise<void> {
-  await showLocation(tryResolveLocation(loc, databaseItem));
+  try {
+    await showLocation(tryResolveLocation(loc, databaseItem));
+  } catch (e) {
+    if (e instanceof Error && e.message.match(/File not found/)) {
+      void Window.showErrorMessage(
+        "Original file of this result is not in the database's source archive.",
+      );
+    } else {
+      void logger.log(`Unable to jump to location: ${getErrorMessage(e)}`);
+    }
+  }
 }
 
 export async function showLocation(location?: Location) {
@@ -146,16 +157,6 @@ export async function jumpToLocation(
 ) {
   const databaseItem = databaseManager.findDatabaseItem(Uri.parse(databaseUri));
   if (databaseItem !== undefined) {
-    try {
-      await showResolvableLocation(loc, databaseItem);
-    } catch (e) {
-      if (e instanceof Error && e.message.match(/File not found/)) {
-        void Window.showErrorMessage(
-          "Original file of this result is not in the database's source archive.",
-        );
-      } else {
-        void logger.log(`Unable to jump to location: ${getErrorMessage(e)}`);
-      }
-    }
+    await showResolvableLocation(loc, databaseItem, logger);
   }
 }

extensions/ql-vscode/src/debugger/debug-protocol.ts

@@ -94,7 +94,6 @@ export type Response = DebugProtocol.Response & { type: "response" };
 export type InitializeResponse = DebugProtocol.InitializeResponse &
   Response & { command: "initialize" };
 
-// eslint-disable-next-line @typescript-eslint/no-empty-interface
 export interface QuickEvalResponse extends Response {}
 
 export type AnyResponse = InitializeResponse | QuickEvalResponse;

extensions/ql-vscode/src/language-support/query-editor.ts

@@ -1,4 +1,4 @@
-import { Uri, window } from "vscode";
+import { Uri, ViewColumn, window } from "vscode";
 import { CodeQLCliServer } from "../codeql-cli/cli";
 import { QueryRunner } from "../query-server";
 import { basename, join } from "path";
@@ -74,6 +74,16 @@ async function previewQueryHelp(
     const uri = Uri.file(absolutePathToMd);
     try {
       await cliServer.generateQueryHelp(pathToQhelp, absolutePathToMd);
+      // Open and then close the raw markdown file first. This ensures that the preview
+      // is refreshed when we open it in the next step.
+      // This will mean that the users will see a the raw markdown file for a brief moment,
+      // but this is the best we can do for now to ensure that the preview is refreshed.
+      await window.showTextDocument(uri, {
+        viewColumn: ViewColumn.Active,
+      });
+      await commandManager.execute("workbench.action.closeActiveEditor");
+
+      // Now open the preview
       await commandManager.execute("markdown.showPreviewToSide", uri);
     } catch (e) {
       const errorMessage = getErrorMessage(e).includes(

extensions/ql-vscode/src/query-history/query-history-scrubber.ts

@@ -10,10 +10,6 @@ import { getErrorMessage } from "../common/helpers-pure";
 
 const LAST_SCRUB_TIME_KEY = "lastScrubTime";
 
-type Counter = {
-  increment: () => void;
-};
-
 /**
  * Registers an interval timer that will periodically check for queries old enought
  * to be deleted.
@@ -37,8 +33,8 @@ export function registerQueryHistoryScrubber(
   qhm: QueryHistoryManager,
   ctx: ExtensionContext,
 
-  // optional counter to keep track of how many times the scrubber has run
-  counter?: Counter,
+  // optional callback to keep track of how many times the scrubber has run
+  onScrubberRun?: () => void,
 ): Disposable {
   const deregister = setInterval(
     scrubQueries,
@@ -48,7 +44,7 @@ export function registerQueryHistoryScrubber(
     queryHistoryDirs,
     qhm,
     ctx,
-    counter,
+    onScrubberRun,
   );
 
   return {
@@ -64,7 +60,7 @@ async function scrubQueries(
   queryHistoryDirs: QueryHistoryDirs,
   qhm: QueryHistoryManager,
   ctx: ExtensionContext,
-  counter?: Counter,
+  onScrubberRun?: () => void,
 ) {
   const lastScrubTime = ctx.globalState.get<number>(LAST_SCRUB_TIME_KEY);
   const now = Date.now();
@@ -76,7 +72,7 @@ async function scrubQueries(
 
     let scrubCount = 0; // total number of directories deleted
     try {
-      counter?.increment();
+      onScrubberRun?.();
       void extLogger.log(
         "Cleaning up query history directories. Removing old entries.",
       );

extensions/ql-vscode/src/variant-analysis/gh-api/mocks/gh-api-request.ts

@@ -11,6 +11,7 @@ export enum RequestKind {
   GetVariantAnalysisRepo = "getVariantAnalysisRepo",
   GetVariantAnalysisRepoResult = "getVariantAnalysisRepoResult",
   CodeSearch = "codeSearch",
+  AutoModel = "autoModel",
 }
 
 interface BasicErorResponse {
@@ -87,13 +88,30 @@ interface CodeSearchRequest {
   };
 }
 
+interface AutoModelRequest {
+  request: {
+    kind: RequestKind.AutoModel;
+    body?: {
+      candidates: string;
+    };
+  };
+  response: {
+    status: number;
+    body?: {
+      models: string;
+    };
+    message?: string;
+  };
+}
+
 export type GitHubApiRequest =
   | GetRepoRequest
   | SubmitVariantAnalysisRequest
   | GetVariantAnalysisRequest
   | GetVariantAnalysisRepoRequest
   | GetVariantAnalysisRepoResultRequest
-  | CodeSearchRequest;
+  | CodeSearchRequest
+  | AutoModelRequest;
 
 export const isGetRepoRequest = (
   request: GitHubApiRequest,
@@ -123,3 +141,8 @@ export const isCodeSearchRequest = (
   request: GitHubApiRequest,
 ): request is CodeSearchRequest =>
   request.request.kind === RequestKind.CodeSearch;
+
+export const isAutoModelRequest = (
+  request: GitHubApiRequest,
+): request is AutoModelRequest =>
+  request.request.kind === RequestKind.AutoModel;

extensions/ql-vscode/src/variant-analysis/gh-api/mocks/recorder.ts

@@ -259,6 +259,21 @@ async function createGitHubApiRequest(
     };
   }
 
+  const autoModelMatch = url.match(
+    /\/repos\/github\/codeql\/code-scanning\/codeql\/auto-model/,
+  );
+  if (autoModelMatch) {
+    return {
+      request: {
+        kind: RequestKind.AutoModel,
+      },
+      response: {
+        status,
+        body: JSON.parse(body),
+      },
+    };
+  }
+
   return undefined;
 }
 

extensions/ql-vscode/src/variant-analysis/gh-api/mocks/request-handlers.ts

@@ -3,6 +3,7 @@ import { readdir, readJson, readFile } from "fs-extra";
 import { DefaultBodyType, MockedRequest, rest, RestHandler } from "msw";
 import {
   GitHubApiRequest,
+  isAutoModelRequest,
   isCodeSearchRequest,
   isGetRepoRequest,
   isGetVariantAnalysisRepoRequest,
@@ -27,6 +28,7 @@ export async function createRequestHandlers(
     createGetVariantAnalysisRepoRequestHandler(requests),
     createGetVariantAnalysisRepoResultRequestHandler(requests),
     createCodeSearchRequestHandler(requests),
+    createAutoModelRequestHandler(requests),
   ];
 
   return handlers;
@@ -219,3 +221,30 @@ function createCodeSearchRequestHandler(
     );
   });
 }
+
+function createAutoModelRequestHandler(
+  requests: GitHubApiRequest[],
+): RequestHandler {
+  const autoModelRequests = requests.filter(isAutoModelRequest);
+  let requestIndex = 0;
+
+  // During automodeling there can be multiple API requests for each batch
+  // of candidates we want to model. We need to return different responses for each request,
+  // so keep an index of the request and return the appropriate response.
+  return rest.post(
+    `${baseUrl}/repos/github/codeql/code-scanning/codeql/auto-model`,
+    (_req, res, ctx) => {
+      const request = autoModelRequests[requestIndex];
+
+      if (requestIndex < autoModelRequests.length - 1) {
+        // If there are more requests to come, increment the index.
+        requestIndex++;
+      }
+
+      return res(
+        ctx.status(request.response.status),
+        ctx.json(request.response.body),
+      );
+    },
+  );
+}

extensions/ql-vscode/src/variant-analysis/gh-api/mocks/scenarios/auto-model-success/0-autoModel.json

@@ -0,0 +1,11 @@
+{
+  "request": {
+    "kind": "autoModel"
+  },
+  "response": {
+    "status": 200,
+    "body": {
+      "models": "extensions:\n- addsTo: {extensible: sinkModel, pack: codeql/java-all}\n  data:\n  - [javax.servlet.http, HttpServletResponse, true, sendRedirect, (String), '', 'Argument[this]',\n    request-forgery, ai-generated]\n  - [javax.servlet.http, HttpServletResponse, true, sendRedirect, (String), '', 'Argument[0]',\n    request-forgery, ai-generated]\n"
+    }
+  }
+}

extensions/ql-vscode/src/variant-analysis/gh-api/mocks/scenarios/auto-model-success/1-autoModel.json

@@ -0,0 +1,11 @@
+{
+  "request": {
+    "kind": "autoModel"
+  },
+  "response": {
+    "status": 200,
+    "body": {
+      "models": "extensions:\n- addsTo: {extensible: sinkModel, pack: codeql/java-all}\n  data:\n  - [javax.servlet, MultipartConfigElement, true, MultipartConfigElement, (String),\n    '', 'Argument[0]', request-forgery, ai-generated]\n"
+    }
+  }
+}

extensions/ql-vscode/src/variant-analysis/gh-api/mocks/scenarios/auto-model-success/README.md

@@ -0,0 +1 @@
+This scenario is best when modeling the `javax.servlet-api` package.

extensions/ql-vscode/src/variant-analysis/sarif-processing.ts

@@ -1,5 +1,6 @@
 import * as sarif from "sarif";
 import {
+  SarifLink,
   parseHighlightedLine,
   parseSarifPlainTextMessage,
   parseSarifRegion,
@@ -14,6 +15,7 @@ import {
   ThreadFlow,
   CodeSnippet,
   HighlightedRegion,
+  AnalysisMessageLocationTokenLocation,
 } from "./shared/analysis-result";
 
 // A line of more than 8k characters is probably generated.
@@ -303,24 +305,47 @@ function getMessage(
     if (typeof messagePart === "string") {
       tokens.push({ t: "text", text: messagePart });
     } else {
-      const relatedLocation = result.relatedLocations!.find(
-        (rl) => rl.id === messagePart.dest,
-      );
-      tokens.push({
-        t: "location",
-        text: messagePart.text,
-        location: {
-          fileLink: {
-            fileLinkPrefix,
-            filePath: relatedLocation!.physicalLocation!.artifactLocation!.uri!,
-          },
-          highlightedRegion: getHighlightedRegion(
-            relatedLocation!.physicalLocation!.region!,
-          ),
-        },
-      });
+      const location = getRelatedLocation(messagePart, result, fileLinkPrefix);
+      if (location === undefined) {
+        tokens.push({ t: "text", text: messagePart.text });
+      } else {
+        tokens.push({
+          t: "location",
+          text: messagePart.text,
+          location,
+        });
+      }
     }
   }
 
   return { tokens };
 }
+
+function getRelatedLocation(
+  messagePart: SarifLink,
+  result: sarif.Result,
+  fileLinkPrefix: string,
+): AnalysisMessageLocationTokenLocation | undefined {
+  const relatedLocation = result.relatedLocations!.find(
+    (rl) => rl.id === messagePart.dest,
+  );
+  if (
+    relatedLocation === undefined ||
+    relatedLocation.physicalLocation?.artifactLocation?.uri === undefined ||
+    relatedLocation.physicalLocation?.artifactLocation?.uri?.startsWith(
+      "file:",
+    ) ||
+    relatedLocation.physicalLocation?.region === undefined
+  ) {
+    return undefined;
+  }
+  return {
+    fileLink: {
+      fileLinkPrefix,
+      filePath: relatedLocation.physicalLocation.artifactLocation.uri,
+    },
+    highlightedRegion: getHighlightedRegion(
+      relatedLocation.physicalLocation.region,
+    ),
+  };
+}

extensions/ql-vscode/src/variant-analysis/shared/analysis-result.ts

@@ -63,10 +63,12 @@ interface AnalysisMessageTextToken {
 export interface AnalysisMessageLocationToken {
   t: "location";
   text: string;
-  location: {
-    fileLink: FileLink;
-    highlightedRegion?: HighlightedRegion;
-  };
+  location: AnalysisMessageLocationTokenLocation;
+}
+
+export interface AnalysisMessageLocationTokenLocation {
+  fileLink: FileLink;
+  highlightedRegion?: HighlightedRegion;
 }
 
 export type ResultSeverity = "Recommendation" | "Warning" | "Error";

extensions/ql-vscode/src/view/common/Alert.tsx

@@ -1,6 +1,6 @@
 import * as React from "react";
 import { ReactNode } from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 
 type ContainerProps = {
   type: "warning" | "error";

extensions/ql-vscode/src/view/common/CodePaths/CodePaths.tsx

@@ -1,5 +1,5 @@
 import * as React from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 import { VSCodeLink } from "@vscode/webview-ui-toolkit/react";
 
 import {

extensions/ql-vscode/src/view/common/CodePaths/ThreadPath.tsx

@@ -1,5 +1,5 @@
 import * as React from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 import { VSCodeTag } from "@vscode/webview-ui-toolkit/react";
 
 import {

extensions/ql-vscode/src/view/common/DeterminateProgressRing.tsx

@@ -1,5 +1,5 @@
 import * as React from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 
 type Props = {
   percent: number;

extensions/ql-vscode/src/view/common/Dropdown.tsx

@@ -1,6 +1,6 @@
 import * as React from "react";
 import { ChangeEvent } from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 
 const StyledDropdown = styled.select`
   width: 100%;

extensions/ql-vscode/src/view/common/FileCodeSnippet/CodeSnippetCode.tsx

@@ -1,5 +1,5 @@
 import * as React from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 
 import { HighlightedRegion } from "../../../variant-analysis/shared/analysis-result";
 import {

extensions/ql-vscode/src/view/common/FileCodeSnippet/CodeSnippetLine.tsx

@@ -1,5 +1,5 @@
 import * as React from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 
 import {
   AnalysisMessage,

extensions/ql-vscode/src/view/common/FileCodeSnippet/CodeSnippetMessage.tsx

@@ -1,5 +1,5 @@
 import * as React from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 import { VSCodeLink } from "@vscode/webview-ui-toolkit/react";
 
 import {

extensions/ql-vscode/src/view/common/FileCodeSnippet/FileCodeSnippet.tsx

@@ -1,5 +1,5 @@
 import * as React from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 import { VSCodeLink } from "@vscode/webview-ui-toolkit/react";
 
 import {

extensions/ql-vscode/src/view/common/HorizontalSpace.tsx

@@ -1,4 +1,4 @@
-import styled from "styled-components";
+import { styled } from "styled-components";
 
 export const HorizontalSpace = styled.div<{ size: 1 | 2 | 3 }>`
   flex: 0 0 auto;

extensions/ql-vscode/src/view/common/SectionTitle.tsx

@@ -1,4 +1,4 @@
-import styled from "styled-components";
+import { styled } from "styled-components";
 
 export const SectionTitle = styled.h2`
   font-size: medium;

extensions/ql-vscode/src/view/common/StarCount.tsx

@@ -1,5 +1,5 @@
 import * as React from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 import { Codicon } from "./icon";
 
 const Star = styled.span`

extensions/ql-vscode/src/view/common/TextButton.tsx

@@ -1,5 +1,5 @@
 import * as React from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 
 type Size = "x-small" | "small" | "medium" | "large" | "x-large";
 

extensions/ql-vscode/src/view/common/VerticalSpace.tsx

@@ -1,4 +1,4 @@
-import styled from "styled-components";
+import { styled } from "styled-components";
 
 export const VerticalSpace = styled.div<{ size: 1 | 2 | 3 }>`
   flex: 0 0 auto;

extensions/ql-vscode/src/view/common/ViewTitle.tsx

@@ -1,4 +1,4 @@
-import styled from "styled-components";
+import { styled } from "styled-components";
 
 export const ViewTitle = styled.h1`
   font-size: 2em;

extensions/ql-vscode/src/view/common/icon/Codicon.tsx

@@ -1,5 +1,5 @@
 import * as React from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 import classNames from "classnames";
 
 type Props = {

extensions/ql-vscode/src/view/common/icon/ErrorIcon.tsx

@@ -1,5 +1,5 @@
 import * as React from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 import { Codicon } from "./Codicon";
 
 const Icon = styled(Codicon)`

extensions/ql-vscode/src/view/common/icon/SuccessIcon.tsx

@@ -1,5 +1,5 @@
 import * as React from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 import { Codicon } from "./Codicon";
 
 const Icon = styled(Codicon)`

extensions/ql-vscode/src/view/common/icon/WarningIcon.tsx

@@ -1,5 +1,5 @@
 import * as React from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 import { Codicon } from "./Codicon";
 
 const Icon = styled(Codicon)`

extensions/ql-vscode/src/view/data-extensions-editor/DataExtensionsEditor.tsx

@@ -6,7 +6,7 @@ import {
   VSCodeCheckbox,
   VSCodeTag,
 } from "@vscode/webview-ui-toolkit/react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 import { ExternalApiUsage } from "../../data-extensions-editor/external-api-usage";
 import { ModeledMethod } from "../../data-extensions-editor/modeled-method";
 import { assertNever } from "../../common/helpers-pure";
@@ -17,6 +17,7 @@ import { DataExtensionEditorViewState } from "../../data-extensions-editor/share
 import { ModeledMethodsList } from "./ModeledMethodsList";
 import { percentFormatter } from "./formatters";
 import { Mode } from "../../data-extensions-editor/shared/mode";
+import { InProgressMethods } from "../../data-extensions-editor/shared/in-progress-methods";
 import { getLanguageDisplayName } from "../../common/query-language";
 
 const LoadingContainer = styled.div`
@@ -91,6 +92,10 @@ export function DataExtensionsEditor({
     new Set(),
   );
 
+  const [inProgressMethods, setInProgressMethods] = useState<InProgressMethods>(
+    new InProgressMethods(),
+  );
+
   const [hideModeledApis, setHideModeledApis] = useState(true);
 
   const [modeledMethods, setModeledMethods] = useState<
@@ -135,6 +140,17 @@ export function DataExtensionsEditor({
                 ]),
             );
             break;
+          case "setInProgressMethods":
+            setInProgressMethods((oldInProgressMethods) => {
+              const methods =
+                InProgressMethods.fromExisting(oldInProgressMethods);
+              methods.setPackageMethods(
+                msg.packageName,
+                new Set(msg.inProgressMethods),
+              );
+              return methods;
+            });
+            break;
           default:
             assertNever(msg);
         }
@@ -220,11 +236,13 @@ export function DataExtensionsEditor({
 
   const onGenerateFromLlmClick = useCallback(
     (
+      packageName: string,
       externalApiUsages: ExternalApiUsage[],
       modeledMethods: Record<string, ModeledMethod>,
     ) => {
       vscode.postMessage({
         t: "generateExternalApiFromLlm",
+        packageName,
         externalApiUsages,
         modeledMethods,
       });
@@ -232,6 +250,13 @@ export function DataExtensionsEditor({
     [],
   );
 
+  const onStopGenerateFromLlmClick = useCallback((packageName: string) => {
+    vscode.postMessage({
+      t: "stopGeneratingExternalApiFromLlm",
+      packageName,
+    });
+  }, []);
+
   const onOpenDatabaseClick = useCallback(() => {
     vscode.postMessage({
       t: "openDatabase",
@@ -330,11 +355,13 @@ export function DataExtensionsEditor({
           externalApiUsages={externalApiUsages}
           modeledMethods={modeledMethods}
           modifiedSignatures={modifiedSignatures}
+          inProgressMethods={inProgressMethods}
           viewState={viewState}
           hideModeledApis={hideModeledApis}
           onChange={onChange}
           onSaveModelClick={onSaveModelClick}
           onGenerateFromLlmClick={onGenerateFromLlmClick}
+          onStopGenerateFromLlmClick={onStopGenerateFromLlmClick}
           onGenerateFromSourceClick={onGenerateFromSourceClick}
           onModelDependencyClick={onModelDependencyClick}
         />

extensions/ql-vscode/src/view/data-extensions-editor/InProgressDropdown.tsx

@@ -0,0 +1,23 @@
+import * as React from "react";
+import { Dropdown } from "../common/Dropdown";
+
+export const InProgressDropdown = () => {
+  const options: Array<{ label: string; value: string }> = [
+    {
+      label: "Thinking...",
+      value: "Thinking...",
+    },
+  ];
+  const noop = () => {
+    // Do nothing
+  };
+
+  return (
+    <Dropdown
+      value="Thinking..."
+      options={options}
+      disabled={false}
+      onChange={noop}
+    />
+  );
+};

extensions/ql-vscode/src/view/data-extensions-editor/KindInput.tsx

@@ -1,6 +1,5 @@
 import * as React from "react";
 import { ChangeEvent, useCallback, useEffect, useMemo } from "react";
-
 import type { ModeledMethod } from "../../data-extensions-editor/modeled-method";
 import { Dropdown } from "../common/Dropdown";
 

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

@@ -1,6 +1,6 @@
 import * as React from "react";
 import { useCallback, useMemo, useState } from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 import { ExternalApiUsage } from "../../data-extensions-editor/external-api-usage";
 import { ModeledMethod } from "../../data-extensions-editor/modeled-method";
 import { ModeledMethodDataGrid } from "./ModeledMethodDataGrid";
@@ -14,6 +14,7 @@ import {
   VSCodeTag,
 } from "@vscode/webview-ui-toolkit/react";
 import { DataExtensionEditorViewState } from "../../data-extensions-editor/shared/view-state";
+import { InProgressMethods } from "../../data-extensions-editor/shared/in-progress-methods";
 
 const LibraryContainer = styled.div`
   background-color: var(--vscode-peekViewResult-background);
@@ -72,6 +73,7 @@ type Props = {
   externalApiUsages: ExternalApiUsage[];
   modeledMethods: Record<string, ModeledMethod>;
   modifiedSignatures: Set<string>;
+  inProgressMethods: InProgressMethods;
   viewState: DataExtensionEditorViewState;
   hideModeledApis: boolean;
   onChange: (
@@ -84,9 +86,11 @@ type Props = {
     modeledMethods: Record<string, ModeledMethod>,
   ) => void;
   onGenerateFromLlmClick: (
+    dependencyName: string,
     externalApiUsages: ExternalApiUsage[],
     modeledMethods: Record<string, ModeledMethod>,
   ) => void;
+  onStopGenerateFromLlmClick: (dependencyName: string) => void;
   onGenerateFromSourceClick: () => void;
   onModelDependencyClick: () => void;
 };
@@ -97,11 +101,13 @@ export const LibraryRow = ({
   externalApiUsages,
   modeledMethods,
   modifiedSignatures,
+  inProgressMethods,
   viewState,
   hideModeledApis,
   onChange,
   onSaveModelClick,
   onGenerateFromLlmClick,
+  onStopGenerateFromLlmClick,
   onGenerateFromSourceClick,
   onModelDependencyClick,
 }: Props) => {
@@ -117,11 +123,20 @@ export const LibraryRow = ({
 
   const handleModelWithAI = useCallback(
     async (e: React.MouseEvent) => {
-      onGenerateFromLlmClick(externalApiUsages, modeledMethods);
+      onGenerateFromLlmClick(title, externalApiUsages, modeledMethods);
       e.stopPropagation();
       e.preventDefault();
     },
-    [externalApiUsages, modeledMethods, onGenerateFromLlmClick],
+    [title, externalApiUsages, modeledMethods, onGenerateFromLlmClick],
+  );
+
+  const handleStopModelWithAI = useCallback(
+    async (e: React.MouseEvent) => {
+      onStopGenerateFromLlmClick(title);
+      e.stopPropagation();
+      e.preventDefault();
+    },
+    [title, onStopGenerateFromLlmClick],
   );
 
   const handleModelFromSource = useCallback(
@@ -164,6 +179,12 @@ export const LibraryRow = ({
     );
   }, [externalApiUsages, modifiedSignatures]);
 
+  const canStopAutoModeling = useMemo(() => {
+    return externalApiUsages.some((externalApiUsage) =>
+      inProgressMethods.hasMethod(title, externalApiUsage.signature),
+    );
+  }, [externalApiUsages, title, inProgressMethods]);
+
   return (
     <LibraryContainer>
       <TitleContainer onClick={toggleExpanded} aria-expanded={isExpanded}>
@@ -182,12 +203,18 @@ export const LibraryRow = ({
           </ModeledPercentage>
           {hasUnsavedChanges ? <VSCodeTag>UNSAVED</VSCodeTag> : null}
         </NameContainer>
-        {viewState.showLlmButton && (
+        {viewState.showLlmButton && !canStopAutoModeling && (
           <VSCodeButton appearance="icon" onClick={handleModelWithAI}>
             <Codicon name="lightbulb-autofix" label="Model with AI" />
             &nbsp;Model with AI
           </VSCodeButton>
         )}
+        {viewState.showLlmButton && canStopAutoModeling && (
+          <VSCodeButton appearance="icon" onClick={handleStopModelWithAI}>
+            <Codicon name="debug-stop" label="Stop model with AI" />
+            &nbsp;Stop
+          </VSCodeButton>
+        )}
         {viewState.mode === Mode.Application && (
           <VSCodeButton appearance="icon" onClick={handleModelFromSource}>
             <Codicon name="code" label="Model from source" />
@@ -206,9 +233,11 @@ export const LibraryRow = ({
         <>
           <SectionDivider />
           <ModeledMethodDataGrid
+            packageName={title}
             externalApiUsages={externalApiUsages}
             modeledMethods={modeledMethods}
             modifiedSignatures={modifiedSignatures}
+            inProgressMethods={inProgressMethods}
             mode={viewState.mode}
             hideModeledApis={hideModeledApis}
             onChange={onChangeWithModelName}

extensions/ql-vscode/src/view/data-extensions-editor/MethodClassifications.tsx

@@ -5,7 +5,7 @@ import {
   ExternalApiUsage,
 } from "../../data-extensions-editor/external-api-usage";
 import { VSCodeTag } from "@vscode/webview-ui-toolkit/react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 
 const ClassificationsContainer = styled.div`
   display: inline-flex;

extensions/ql-vscode/src/view/data-extensions-editor/MethodRow.tsx

@@ -2,10 +2,11 @@ import {
   VSCodeDataGridCell,
   VSCodeDataGridRow,
   VSCodeLink,
+  VSCodeProgressRing,
 } from "@vscode/webview-ui-toolkit/react";
 import * as React from "react";
 import { ChangeEvent, useCallback, useMemo } from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 import { vscode } from "../vscode-api";
 
 import { ExternalApiUsage } from "../../data-extensions-editor/external-api-usage";
@@ -23,6 +24,7 @@ import {
   ModelingStatus,
   ModelingStatusIndicator,
 } from "./ModelingStatusIndicator";
+import { InProgressDropdown } from "./InProgressDropdown";
 
 const ApiOrMethodCell = styled(VSCodeDataGridCell)`
   display: flex;
@@ -43,6 +45,12 @@ const ViewLink = styled(VSCodeLink)`
   white-space: nowrap;
 `;
 
+const ProgressRing = styled(VSCodeProgressRing)`
+  width: 16px;
+  height: 16px;
+  margin-left: auto;
+`;
+
 const modelTypeOptions: Array<{ value: ModeledMethodType; label: string }> = [
   { value: "none", label: "Unmodeled" },
   { value: "source", label: "Source" },
@@ -55,6 +63,7 @@ type Props = {
   externalApiUsage: ExternalApiUsage;
   modeledMethod: ModeledMethod | undefined;
   methodIsUnsaved: boolean;
+  modelingInProgress: boolean;
   mode: Mode;
   hideModeledApis: boolean;
   onChange: (
@@ -216,38 +225,59 @@ function ModelableMethodRow(props: Props) {
           </UsagesButton>
         )}
         <ViewLink onClick={jumpToUsage}>View</ViewLink>
+        {props.modelingInProgress && <ProgressRing />}
       </ApiOrMethodCell>
-      <VSCodeDataGridCell gridColumn={2}>
-        <Dropdown
-          value={modeledMethod?.type ?? "none"}
-          options={modelTypeOptions}
-          onChange={handleTypeInput}
-        />
-      </VSCodeDataGridCell>
-      <VSCodeDataGridCell gridColumn={3}>
-        <Dropdown
-          value={modeledMethod?.input}
-          options={inputOptions}
-          disabled={!showInputCell}
-          onChange={handleInputInput}
-        />
-      </VSCodeDataGridCell>
-      <VSCodeDataGridCell gridColumn={4}>
-        <Dropdown
-          value={modeledMethod?.output}
-          options={outputOptions}
-          disabled={!showOutputCell}
-          onChange={handleOutputInput}
-        />
-      </VSCodeDataGridCell>
-      <VSCodeDataGridCell gridColumn={5}>
-        <KindInput
-          kinds={predicate?.supportedKinds || []}
-          value={modeledMethod?.kind}
-          disabled={!showKindCell}
-          onChange={handleKindChange}
-        />
-      </VSCodeDataGridCell>
+      {props.modelingInProgress && (
+        <>
+          <VSCodeDataGridCell gridColumn={2}>
+            <InProgressDropdown />
+          </VSCodeDataGridCell>
+          <VSCodeDataGridCell gridColumn={3}>
+            <InProgressDropdown />
+          </VSCodeDataGridCell>
+          <VSCodeDataGridCell gridColumn={4}>
+            <InProgressDropdown />
+          </VSCodeDataGridCell>
+          <VSCodeDataGridCell gridColumn={5}>
+            <InProgressDropdown />
+          </VSCodeDataGridCell>
+        </>
+      )}
+      {!props.modelingInProgress && (
+        <>
+          <VSCodeDataGridCell gridColumn={2}>
+            <Dropdown
+              value={modeledMethod?.type ?? "none"}
+              options={modelTypeOptions}
+              onChange={handleTypeInput}
+            />
+          </VSCodeDataGridCell>
+          <VSCodeDataGridCell gridColumn={3}>
+            <Dropdown
+              value={modeledMethod?.input}
+              options={inputOptions}
+              disabled={!showInputCell}
+              onChange={handleInputInput}
+            />
+          </VSCodeDataGridCell>
+          <VSCodeDataGridCell gridColumn={4}>
+            <Dropdown
+              value={modeledMethod?.output}
+              options={outputOptions}
+              disabled={!showOutputCell}
+              onChange={handleOutputInput}
+            />
+          </VSCodeDataGridCell>
+          <VSCodeDataGridCell gridColumn={5}>
+            <KindInput
+              kinds={predicate?.supportedKinds || []}
+              value={modeledMethod?.kind}
+              disabled={!showKindCell}
+              onChange={handleKindChange}
+            />
+          </VSCodeDataGridCell>
+        </>
+      )}
     </VSCodeDataGridRow>
   );
 }

extensions/ql-vscode/src/view/data-extensions-editor/ModeledMethodDataGrid.tsx

@@ -10,11 +10,14 @@ import { ModeledMethod } from "../../data-extensions-editor/modeled-method";
 import { useMemo } from "react";
 import { Mode } from "../../data-extensions-editor/shared/mode";
 import { sortMethods } from "../../data-extensions-editor/shared/sorting";
+import { InProgressMethods } from "../../data-extensions-editor/shared/in-progress-methods";
 
 type Props = {
+  packageName: string;
   externalApiUsages: ExternalApiUsage[];
   modeledMethods: Record<string, ModeledMethod>;
   modifiedSignatures: Set<string>;
+  inProgressMethods: InProgressMethods;
   mode: Mode;
   hideModeledApis: boolean;
   onChange: (
@@ -24,9 +27,11 @@ type Props = {
 };
 
 export const ModeledMethodDataGrid = ({
+  packageName,
   externalApiUsages,
   modeledMethods,
   modifiedSignatures,
+  inProgressMethods,
   mode,
   hideModeledApis,
   onChange,
@@ -37,7 +42,7 @@ export const ModeledMethodDataGrid = ({
   );
 
   return (
-    <VSCodeDataGrid gridTemplateColumns="0.5fr 0.125fr 0.125fr 0.125fr 0.125fr">
+    <VSCodeDataGrid gridTemplateColumns="0.5fr 0.125fr 0.125fr 0.125fr 0.125fr 0.125fr">
       <VSCodeDataGridRow rowType="header">
         <VSCodeDataGridCell cellType="columnheader" gridColumn={1}>
           API or method
@@ -61,6 +66,10 @@ export const ModeledMethodDataGrid = ({
           externalApiUsage={externalApiUsage}
           modeledMethod={modeledMethods[externalApiUsage.signature]}
           methodIsUnsaved={modifiedSignatures.has(externalApiUsage.signature)}
+          modelingInProgress={inProgressMethods.hasMethod(
+            packageName,
+            externalApiUsage.signature,
+          )}
           mode={mode}
           hideModeledApis={hideModeledApis}
           onChange={onChange}

extensions/ql-vscode/src/view/data-extensions-editor/ModeledMethodsList.tsx

@@ -9,11 +9,13 @@ import {
   sortGroupNames,
 } from "../../data-extensions-editor/shared/sorting";
 import { DataExtensionEditorViewState } from "../../data-extensions-editor/shared/view-state";
+import { InProgressMethods } from "../../data-extensions-editor/shared/in-progress-methods";
 
 type Props = {
   externalApiUsages: ExternalApiUsage[];
   modeledMethods: Record<string, ModeledMethod>;
   modifiedSignatures: Set<string>;
+  inProgressMethods: InProgressMethods;
   viewState: DataExtensionEditorViewState;
   hideModeledApis: boolean;
   onChange: (
@@ -26,9 +28,11 @@ type Props = {
     modeledMethods: Record<string, ModeledMethod>,
   ) => void;
   onGenerateFromLlmClick: (
+    packageName: string,
     externalApiUsages: ExternalApiUsage[],
     modeledMethods: Record<string, ModeledMethod>,
   ) => void;
+  onStopGenerateFromLlmClick: (packageName: string) => void;
   onGenerateFromSourceClick: () => void;
   onModelDependencyClick: () => void;
 };
@@ -41,11 +45,13 @@ export const ModeledMethodsList = ({
   externalApiUsages,
   modeledMethods,
   modifiedSignatures,
+  inProgressMethods,
   viewState,
   hideModeledApis,
   onChange,
   onSaveModelClick,
   onGenerateFromLlmClick,
+  onStopGenerateFromLlmClick,
   onGenerateFromSourceClick,
   onModelDependencyClick,
 }: Props) => {
@@ -84,11 +90,13 @@ export const ModeledMethodsList = ({
           externalApiUsages={grouped[libraryName]}
           modeledMethods={modeledMethods}
           modifiedSignatures={modifiedSignatures}
+          inProgressMethods={inProgressMethods}
           viewState={viewState}
           hideModeledApis={hideModeledApis}
           onChange={onChange}
           onSaveModelClick={onSaveModelClick}
           onGenerateFromLlmClick={onGenerateFromLlmClick}
+          onStopGenerateFromLlmClick={onStopGenerateFromLlmClick}
           onGenerateFromSourceClick={onGenerateFromSourceClick}
           onModelDependencyClick={onModelDependencyClick}
         />

extensions/ql-vscode/src/view/data-flow-paths/DataFlowPaths.tsx

@@ -1,5 +1,5 @@
 import * as React from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 import { useState } from "react";
 
 import { useTelemetryOnChange } from "../common/telemetry";

extensions/ql-vscode/src/view/results/__tests__/results.spec.tsx

@@ -12,10 +12,7 @@ import { ColumnKindCode } from "../../../common/bqrs-cli-types";
 import { postMessage } from "../../common/post-message";
 
 const exampleSarif = fs.readJSONSync(
-  resolve(
-    __dirname,
-    "../../../../test/vscode-tests/no-workspace/data/sarif/validSarif.sarif",
-  ),
+  resolve(__dirname, "../../../../test/data/sarif/validSarif.sarif"),
 );
 
 describe(ResultsApp.name, () => {

extensions/ql-vscode/src/view/variant-analysis/AnalyzedRepoItemContent.tsx

@@ -1,5 +1,5 @@
 import * as React from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 import {
   AnalysisAlert,
   AnalysisRawResults,

extensions/ql-vscode/src/view/variant-analysis/LinkIconButton.tsx

@@ -1,5 +1,5 @@
 import { VSCodeLink } from "@vscode/webview-ui-toolkit/react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 
 export const LinkIconButton = styled(VSCodeLink)`
   .codicon {

extensions/ql-vscode/src/view/variant-analysis/QueryDetails.tsx

@@ -1,5 +1,5 @@
 import * as React from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 import { ViewTitle } from "../common";
 import { LinkIconButton } from "./LinkIconButton";
 

extensions/ql-vscode/src/view/variant-analysis/RawResultsTable.tsx

@@ -1,6 +1,6 @@
 import * as React from "react";
 import { useState } from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 import { VSCodeLink } from "@vscode/webview-ui-toolkit/react";
 import {
   CellValue,

extensions/ql-vscode/src/view/variant-analysis/RepoRow.tsx

@@ -1,6 +1,6 @@
 import * as React from "react";
 import { ChangeEvent, useCallback, useEffect, useState } from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 import { VSCodeBadge, VSCodeCheckbox } from "@vscode/webview-ui-toolkit/react";
 import {
   isCompletedAnalysisRepoStatus,

extensions/ql-vscode/src/view/variant-analysis/RepositoriesFilter.tsx

@@ -1,6 +1,6 @@
 import * as React from "react";
 import { useCallback } from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 import { VSCodeDropdown, VSCodeOption } from "@vscode/webview-ui-toolkit/react";
 import { Codicon } from "../common";
 import { FilterKey } from "../../variant-analysis/shared/variant-analysis-filter-sort";

extensions/ql-vscode/src/view/variant-analysis/RepositoriesSearch.tsx

@@ -1,6 +1,6 @@
 import * as React from "react";
 import { useCallback } from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 import { VSCodeTextField } from "@vscode/webview-ui-toolkit/react";
 import { Codicon } from "../common";
 

extensions/ql-vscode/src/view/variant-analysis/RepositoriesSearchSortRow.tsx

@@ -1,6 +1,6 @@
 import * as React from "react";
 import { Dispatch, SetStateAction, useCallback } from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 import {
   FilterKey,
   RepositoriesFilterSortState,

extensions/ql-vscode/src/view/variant-analysis/RepositoriesSort.tsx

@@ -1,6 +1,6 @@
 import * as React from "react";
 import { useCallback } from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 import { VSCodeDropdown, VSCodeOption } from "@vscode/webview-ui-toolkit/react";
 import { SortKey } from "../../variant-analysis/shared/variant-analysis-filter-sort";
 import { Codicon } from "../common";

extensions/ql-vscode/src/view/variant-analysis/StatItem.tsx

@@ -1,6 +1,6 @@
 import * as React from "react";
 import type { ReactNode } from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 
 type Props = {
   title: ReactNode;

extensions/ql-vscode/src/view/variant-analysis/VariantAnalysisActions.tsx

@@ -1,5 +1,5 @@
 import * as React from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 import { VSCodeButton } from "@vscode/webview-ui-toolkit/react";
 import { VariantAnalysisStatus } from "../../variant-analysis/shared/variant-analysis";
 

extensions/ql-vscode/src/view/variant-analysis/VariantAnalysisAnalyzedRepos.tsx

@@ -1,6 +1,6 @@
 import * as React from "react";
 import { Dispatch, SetStateAction, useCallback, useMemo } from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 import { RepoRow } from "./RepoRow";
 import {
   VariantAnalysis,

extensions/ql-vscode/src/view/variant-analysis/VariantAnalysisContainer.tsx

@@ -1,4 +1,4 @@
-import styled from "styled-components";
+import { styled } from "styled-components";
 
 export const VariantAnalysisContainer = styled.div`
   max-width: 55em;

extensions/ql-vscode/src/view/variant-analysis/VariantAnalysisHeader.tsx

@@ -1,6 +1,6 @@
 import * as React from "react";
 import { useMemo } from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 import {
   getSkippedRepoCount,
   getTotalResultCount,

extensions/ql-vscode/src/view/variant-analysis/VariantAnalysisLoading.tsx

@@ -1,5 +1,5 @@
 import * as React from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 
 const Container = styled.div`
   display: flex;

extensions/ql-vscode/src/view/variant-analysis/VariantAnalysisOutcomePanels.tsx

@@ -1,6 +1,6 @@
 import * as React from "react";
 import { Dispatch, SetStateAction } from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 import {
   VSCodeBadge,
   VSCodePanels,

extensions/ql-vscode/src/view/variant-analysis/VariantAnalysisSkippedRepositoriesTab.tsx

@@ -1,6 +1,6 @@
 import * as React from "react";
 import { useMemo } from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 import { VariantAnalysisSkippedRepositoryGroup } from "../../variant-analysis/shared/variant-analysis";
 import { Alert } from "../common";
 import { RepoRow } from "./RepoRow";

extensions/ql-vscode/src/view/variant-analysis/VariantAnalysisStats.tsx

@@ -1,6 +1,6 @@
 import * as React from "react";
 import { useMemo } from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 import { VariantAnalysisStatus } from "../../variant-analysis/shared/variant-analysis";
 import { StatItem } from "./StatItem";
 import { formatDecimal } from "../../common/number";

extensions/ql-vscode/src/view/variant-analysis/VariantAnalysisStatusStats.tsx

@@ -1,5 +1,5 @@
 import * as React from "react";
-import styled from "styled-components";
+import { styled } from "styled-components";
 import { VSCodeLink } from "@vscode/webview-ui-toolkit/react";
 import { formatDate } from "../../common/date";
 import { VariantAnalysisStatus } from "../../variant-analysis/shared/variant-analysis";

extensions/ql-vscode/test/unit-tests/command-lint.test.ts

@@ -44,7 +44,8 @@ describe("commands declared in package.json", () => {
       command.match(/^codeQLQueryHistory\./) ||
       command.match(/^codeQLAstViewer\./) ||
       command.match(/^codeQLEvalLogViewer\./) ||
-      command.match(/^codeQLTests\./)
+      command.match(/^codeQLTests\./) ||
+      command.match(/^codeQLDataExtensionsEditor\./)
     ) {
       scopedCmds.add(command);
       expect(title).toBeDefined();

extensions/ql-vscode/test/unit-tests/common/sarif-parser.test.ts

@@ -1,9 +1,9 @@
 import { join } from "path";
 
-import { sarifParser } from "../../../../src/common/sarif-parser";
+import { sarifParser } from "../../../src/common/sarif-parser";
 
 describe("sarif parser", () => {
-  const sarifDir = join(__dirname, "../data/sarif");
+  const sarifDir = join(__dirname, "../../data/sarif");
   it("should parse a valid SARIF file", async () => {
     const result = await sarifParser(join(sarifDir, "validSarif.sarif"));
     expect(result.version).toBeDefined();

extensions/ql-vscode/test/unit-tests/data-extensions-editor/auto-model-v2.test.ts

@@ -1,12 +1,15 @@
 import {
   createAutoModelV2Request,
   encodeSarif,
+  getCandidates,
 } from "../../../src/data-extensions-editor/auto-model-v2";
 import { Mode } from "../../../src/data-extensions-editor/shared/mode";
 import { AutomodelMode } from "../../../src/data-extensions-editor/auto-model-api-v2";
 import { AutoModelQueriesResult } from "../../../src/data-extensions-editor/auto-model-codeml-queries";
 import * as sarif from "sarif";
 import { gzipDecode } from "../../../src/common/zlib";
+import { ExternalApiUsage } from "../../../src/data-extensions-editor/external-api-usage";
+import { ModeledMethod } from "../../../src/data-extensions-editor/modeled-method";
 
 describe("createAutoModelV2Request", () => {
   const createSarifLog = (queryId: string): sarif.Log => {
@@ -80,3 +83,86 @@ describe("createAutoModelV2Request", () => {
     expect(parsed).toEqual(result.candidates);
   });
 });
+
+describe("getCandidates", () => {
+  it("doesn't return methods that are already modelled", () => {
+    const externalApiUsages: ExternalApiUsage[] = [
+      {
+        library: "my.jar",
+        signature: "org.my.A#x()",
+        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: "",
+        input: "",
+        output: "",
+        provenance: "manual",
+        signature: "org.my.A#x()",
+        packageName: "org.my",
+        typeName: "A",
+        methodName: "x",
+        methodParameters: "()",
+      },
+    };
+    const candidates = getCandidates(
+      Mode.Application,
+      externalApiUsages,
+      modeledMethods,
+    );
+    expect(candidates.length).toEqual(0);
+  });
+
+  it("doesn't return methods that are supported from other sources", () => {
+    const externalApiUsages: ExternalApiUsage[] = [
+      {
+        library: "my.jar",
+        signature: "org.my.A#x()",
+        packageName: "org.my",
+        typeName: "A",
+        methodName: "x",
+        methodParameters: "()",
+        supported: true,
+        supportedType: "none",
+        usages: [],
+      },
+    ];
+    const modeledMethods = {};
+    const candidates = getCandidates(
+      Mode.Application,
+      externalApiUsages,
+      modeledMethods,
+    );
+    expect(candidates.length).toEqual(0);
+  });
+
+  it("returns methods that are neither modeled nor supported from other sources", () => {
+    const externalApiUsages: ExternalApiUsage[] = [];
+    externalApiUsages.push({
+      library: "my.jar",
+      signature: "org.my.A#x()",
+      packageName: "org.my",
+      typeName: "A",
+      methodName: "x",
+      methodParameters: "()",
+      supported: false,
+      supportedType: "none",
+      usages: [],
+    });
+    const modeledMethods = {};
+    const candidates = getCandidates(
+      Mode.Application,
+      externalApiUsages,
+      modeledMethods,
+    );
+    expect(candidates.length).toEqual(1);
+  });
+});

extensions/ql-vscode/test/unit-tests/sarif-processing.test.ts

@@ -720,6 +720,45 @@ describe("SARIF processing", () => {
       expectNoParsingError(result);
       expect(actualCodeSnippet).not.toBeUndefined();
     });
+
+    it("should be able to handle when a location has no uri", () => {
+      const sarif = buildValidSarifLog();
+      sarif.runs![0].results![0].message.text = "message [String](1)";
+      sarif.runs![0].results![0].relatedLocations = [
+        {
+          id: 1,
+          physicalLocation: {
+            artifactLocation: {
+              uri: "file:/modules/java.base/java/lang/String.class",
+              index: 1,
+            },
+          },
+          message: {
+            text: "String",
+          },
+        },
+      ];
+
+      const result = extractAnalysisAlerts(sarif, fakefileLinkPrefix);
+
+      expect(result).toBeTruthy();
+      expectNoParsingError(result);
+      expect(result.alerts[0].codeSnippet).not.toBeUndefined();
+      expect(result.alerts[0].message.tokens).toStrictEqual([
+        {
+          t: "text",
+          text: "message ",
+        },
+        {
+          t: "text",
+          text: "String",
+        },
+        {
+          t: "text",
+          text: "",
+        },
+      ]);
+    });
   });
 
   function expectResultParsingError(msg: string) {

extensions/ql-vscode/test/vscode-tests/activated-extension/variant-analysis/variant-analysis-results-manager.test.ts

@@ -279,10 +279,7 @@ describe(VariantAnalysisResultsManager.name, () => {
           await fs.outputJson(
             join(repoTaskStorageDirectory, "results/results.sarif"),
             await fs.readJson(
-              resolve(
-                __dirname,
-                "../../no-workspace/data/sarif/validSarif.sarif",
-              ),
+              resolve(__dirname, "../../../data/sarif/validSarif.sarif"),
             ),
           );
         });

extensions/ql-vscode/test/vscode-tests/no-workspace/data-extensions-editor/auto-model-codeml-queries.test.ts

@@ -5,13 +5,21 @@ import {
 } from "../../../../src/databases/local-databases";
 import { file } from "tmp-promise";
 import { QueryResultType } from "../../../../src/query-server/new-messages";
-import { runAutoModelQueries } from "../../../../src/data-extensions-editor/auto-model-codeml-queries";
+import {
+  generateCandidateFilterPack,
+  runAutoModelQueries,
+} from "../../../../src/data-extensions-editor/auto-model-codeml-queries";
 import { Mode } from "../../../../src/data-extensions-editor/shared/mode";
 import { mockedObject, mockedUri } from "../../utils/mocking.helpers";
 import { CodeQLCliServer } from "../../../../src/codeql-cli/cli";
 import { QueryRunner } from "../../../../src/query-server";
 import * as queryResolver from "../../../../src/local-queries/query-resolver";
 import * as standardQueries from "../../../../src/local-queries/standard-queries";
+import { MethodSignature } from "../../../../src/data-extensions-editor/external-api-usage";
+import { join } from "path";
+import { exists, readFile } from "fs-extra";
+import { load as loadYaml } from "js-yaml";
+import { CancellationTokenSource } from "vscode-jsonrpc";
 
 describe("runAutoModelQueries", () => {
   const qlpack = {
@@ -60,6 +68,7 @@ describe("runAutoModelQueries", () => {
 
     const options = {
       mode: Mode.Application,
+      candidateMethods: [],
       cliServer: mockedObject<CodeQLCliServer>({
         resolveQlpacks: jest.fn().mockResolvedValue({
           "/a/b/c/my-extension-pack": {},
@@ -134,13 +143,17 @@ describe("runAutoModelQueries", () => {
       }),
       queryStorageDir: "/tmp/queries",
       progress: jest.fn(),
+      cancellationTokenSource: new CancellationTokenSource(),
     };
 
     const result = await runAutoModelQueries(options);
     expect(result).not.toBeUndefined();
 
     expect(options.cliServer.resolveQlpacks).toHaveBeenCalledTimes(1);
-    expect(options.cliServer.resolveQlpacks).toHaveBeenCalledWith([], true);
+    expect(options.cliServer.resolveQlpacks).toHaveBeenCalledWith(
+      expect.arrayContaining([expect.stringContaining("tmp")]),
+      true,
+    );
     expect(resolveQueriesSpy).toHaveBeenCalledTimes(1);
     expect(resolveQueriesSpy).toHaveBeenCalledWith(
       options.cliServer,
@@ -165,7 +178,7 @@ describe("runAutoModelQueries", () => {
         quickEvalCountOnly: false,
       },
       false,
-      [],
+      expect.arrayContaining([expect.stringContaining("tmp")]),
       ["/a/b/c/my-extension-pack"],
       "/tmp/queries",
       undefined,
@@ -173,3 +186,34 @@ describe("runAutoModelQueries", () => {
     );
   });
 });
+
+describe("generateCandidateFilterPack", () => {
+  it("should create a temp pack containing the candidate filters", async () => {
+    const candidateMethods: MethodSignature[] = [
+      {
+        signature: "org.my.A#x()",
+        packageName: "org.my",
+        typeName: "A",
+        methodName: "x",
+        methodParameters: "()",
+      },
+    ];
+    const packDir = await generateCandidateFilterPack("java", candidateMethods);
+    expect(packDir).not.toBeUndefined();
+    const qlpackFile = join(packDir, "codeql-pack.yml");
+    expect(await exists(qlpackFile)).toBe(true);
+    const filterFile = join(packDir, "filter.yml");
+    expect(await exists(filterFile)).toBe(true);
+    // Read the contents of filterFile and parse as yaml
+    const yaml = await loadYaml(await readFile(filterFile, "utf8"));
+    const extensions = yaml.extensions;
+    expect(extensions).toBeInstanceOf(Array);
+    expect(extensions).toHaveLength(1);
+    const extension = extensions[0];
+    expect(extension.addsTo.pack).toEqual("codeql/java-queries");
+    expect(extension.addsTo.extensible).toEqual("automodelCandidateFilter");
+    expect(extension.data).toBeInstanceOf(Array);
+    expect(extension.data).toHaveLength(1);
+    expect(extension.data[0]).toEqual(["org.my", "A", "x", "()"]);
+  });
+});

extensions/ql-vscode/test/vscode-tests/no-workspace/query-history/query-history-scrubber.test.ts

@@ -13,46 +13,28 @@ import {
   TWO_HOURS_IN_MS,
 } from "../../../../src/common/time";
 import { mockedObject } from "../../utils/mocking.helpers";
+import { DirResult } from "tmp";
+
+const now = Date.now();
+// We don't want our times to align exactly with the hour,
+// so we can better mimic real life
+const LESS_THAN_ONE_DAY = ONE_DAY_IN_MS - 1000;
 
 describe("query history scrubber", () => {
-  const now = Date.now();
-
   let deregister: vscode.Disposable | undefined;
-  let mockCtx: vscode.ExtensionContext;
-  let runCount = 0;
-
-  // We don't want our times to align exactly with the hour,
-  // so we can better mimic real life
-  const LESS_THAN_ONE_DAY = ONE_DAY_IN_MS - 1000;
-  const tmpDir = dirSync({
-    unsafeCleanup: true,
-  });
+  let tmpDir: DirResult;
 
   beforeEach(() => {
+    tmpDir = dirSync({
+      unsafeCleanup: true,
+    });
+
     jest.spyOn(extLogger, "log").mockResolvedValue(undefined);
 
     jest.useFakeTimers({
       doNotFake: ["setTimeout"],
       now,
     });
-
-    mockCtx = {
-      globalState: {
-        lastScrubTime: now,
-        get(key: string) {
-          if (key !== "lastScrubTime") {
-            throw new Error(`Unexpected key: ${key}`);
-          }
-          return this.lastScrubTime;
-        },
-        async update(key: string, value: any) {
-          if (key !== "lastScrubTime") {
-            throw new Error(`Unexpected key: ${key}`);
-          }
-          this.lastScrubTime = value;
-        },
-      },
-    } as any as vscode.ExtensionContext;
   });
 
   afterEach(() => {
@@ -60,30 +42,32 @@ describe("query history scrubber", () => {
       deregister.dispose();
       deregister = undefined;
     }
+    tmpDir.removeCallback();
   });
 
   it("should not throw an error when the query directory does not exist", async () => {
-    registerScrubber("idontexist");
+    const mockCtx = createMockContext();
+    const runCounter = registerScrubber("idontexist", mockCtx);
 
     jest.advanceTimersByTime(ONE_HOUR_IN_MS);
     await wait();
     // "Should not have called the scrubber"
-    expect(runCount).toBe(0);
+    expect(runCounter).toHaveBeenCalledTimes(0);
 
     jest.advanceTimersByTime(ONE_HOUR_IN_MS - 1);
     await wait();
     // "Should not have called the scrubber"
-    expect(runCount).toBe(0);
+    expect(runCounter).toHaveBeenCalledTimes(0);
 
     jest.advanceTimersByTime(1);
     await wait();
     // "Should have called the scrubber once"
-    expect(runCount).toBe(1);
+    expect(runCounter).toHaveBeenCalledTimes(1);
 
     jest.advanceTimersByTime(TWO_HOURS_IN_MS);
     await wait();
     // "Should have called the scrubber a second time"
-    expect(runCount).toBe(2);
+    expect(runCounter).toHaveBeenCalledTimes(2);
 
     expect((mockCtx.globalState as any).lastScrubTime).toBe(
       now + TWO_HOURS_IN_MS * 2,
@@ -97,7 +81,7 @@ describe("query history scrubber", () => {
       TWO_HOURS_IN_MS,
       THREE_HOURS_IN_MS,
     );
-    registerScrubber(queryDir);
+    registerScrubber(queryDir, createMockContext());
 
     jest.advanceTimersByTime(TWO_HOURS_IN_MS);
     await wait();
@@ -176,7 +160,31 @@ describe("query history scrubber", () => {
     return `query-${timestamp}`;
   }
 
-  function registerScrubber(dir: string) {
+  function createMockContext(): vscode.ExtensionContext {
+    return {
+      globalState: {
+        lastScrubTime: now,
+        get(key: string) {
+          if (key !== "lastScrubTime") {
+            throw new Error(`Unexpected key: ${key}`);
+          }
+          return this.lastScrubTime;
+        },
+        async update(key: string, value: any) {
+          if (key !== "lastScrubTime") {
+            throw new Error(`Unexpected key: ${key}`);
+          }
+          this.lastScrubTime = value;
+        },
+      },
+    } as any as vscode.ExtensionContext;
+  }
+
+  function registerScrubber(
+    dir: string,
+    ctx: vscode.ExtensionContext,
+  ): jest.Mock {
+    const onScrubberRun = jest.fn();
     deregister = registerQueryHistoryScrubber(
       ONE_HOUR_IN_MS,
       TWO_HOURS_IN_MS,
@@ -187,11 +195,10 @@ describe("query history scrubber", () => {
           return Promise.resolve();
         },
       }),
-      mockCtx,
-      {
-        increment: () => runCount++,
-      },
+      ctx,
+      onScrubberRun,
     );
+    return onScrubberRun;
   }
 
   async function wait(ms = 500) {