Add creation of auto-model request V2

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

@@ -0,0 +1,43 @@
+import { gzip, gunzip, InputType as ZlibInputType, ZlibOptions } from "zlib";
+
+/**
+ * Promisified version of zlib.gzip
+ * @param buffer Buffer to compress
+ * @param options zlib options
+ */
+export function gzipEncode(
+  buffer: ZlibInputType,
+  options: ZlibOptions = {},
+): Promise<Buffer> {
+  return new Promise((resolve, reject) => {
+    gzip(buffer, options, (error, result) => {
+      if (error) {
+        reject(error);
+        return;
+      }
+
+      resolve(result);
+    });
+  });
+}
+
+/**
+ * Promisified version of zlib.gunzip
+ * @param buffer Buffer to decompress
+ * @param options zlib options
+ */
+export function gzipDecode(
+  buffer: ZlibInputType,
+  options: ZlibOptions = {},
+): Promise<Buffer> {
+  return new Promise((resolve, reject) => {
+    gunzip(buffer, options, (error, result) => {
+      if (error) {
+        reject(error);
+        return;
+      }
+
+      resolve(result);
+    });
+  });
+}

extensions/ql-vscode/src/config.ts

@@ -704,6 +704,7 @@ export function showQueriesPanel(): boolean {
 
 const DATA_EXTENSIONS = new Setting("dataExtensions", ROOT_SETTING);
 const LLM_GENERATION = new Setting("llmGeneration", DATA_EXTENSIONS);
+const LLM_GENERATION_V2 = new Setting("llmGenerationV2", DATA_EXTENSIONS);
 const FRAMEWORK_MODE = new Setting("frameworkMode", DATA_EXTENSIONS);
 const DISABLE_AUTO_NAME_EXTENSION_PACK = new Setting(
   "disableAutoNameExtensionPack",
@@ -718,6 +719,10 @@ export function showLlmGeneration(): boolean {
   return !!LLM_GENERATION.getValue<boolean>();
 }
 
+export function useLlmGenerationV2(): boolean {
+  return !!LLM_GENERATION_V2.getValue<boolean>();
+}
+
 export function enableFrameworkMode(): boolean {
   return !!FRAMEWORK_MODE.getValue<boolean>();
 }

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

@@ -0,0 +1,34 @@
+import { Credentials } from "../common/authentication";
+import { OctokitResponse } from "@octokit/types";
+
+export enum AutomodelMode {
+  Unspecified = "AUTOMODEL_MODE_UNSPECIFIED",
+  Framework = "AUTOMODEL_MODE_FRAMEWORK",
+  Application = "AUTOMODEL_MODE_APPLICATION",
+}
+
+export interface ModelRequest {
+  mode: AutomodelMode;
+  // Base64-encoded GZIP-compressed SARIF log
+  candidates: string;
+}
+
+export interface ModelResponse {
+  models: string;
+}
+
+export async function autoModelV2(
+  credentials: Credentials,
+  request: ModelRequest,
+): Promise<ModelResponse> {
+  const octokit = await credentials.getOctokit();
+
+  const response: OctokitResponse<ModelResponse> = await octokit.request(
+    "POST /repos/github/codeql/code-scanning/codeql/auto-model",
+    {
+      data: request,
+    },
+  );
+
+  return response.data;
+}

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

@@ -0,0 +1,223 @@
+import { CodeQLCliServer, SourceInfo } from "../codeql-cli/cli";
+import { QueryRunner } from "../query-server";
+import { DatabaseItem } from "../databases/local-databases";
+import { ProgressCallback } from "../common/vscode/progress";
+import * as Sarif from "sarif";
+import { qlpackOfDatabase, resolveQueries } from "../local-queries";
+import { extLogger } from "../common/logging/vscode";
+import { Mode } from "./shared/mode";
+import { QlPacksForLanguage } from "../databases/qlpack";
+import { createLockFileForStandardQuery } from "../local-queries/standard-queries";
+import { CancellationToken, CancellationTokenSource } from "vscode";
+import { getOnDiskWorkspaceFolders } from "../common/vscode/workspace-folders";
+import { showAndLogExceptionWithTelemetry, TeeLogger } from "../common/logging";
+import { QueryResultType } from "../query-server/new-messages";
+import { telemetryListener } from "../common/vscode/telemetry";
+import { redactableError } from "../common/errors";
+import { interpretResultsSarif } from "../query-results";
+import { join } from "path";
+import { assertNever } from "../common/helpers-pure";
+
+type Options = {
+  cliServer: CodeQLCliServer;
+  queryRunner: QueryRunner;
+  databaseItem: DatabaseItem;
+  queryStorageDir: string;
+
+  progress: ProgressCallback;
+};
+
+export type AutoModelQueriesResult = {
+  candidates: Sarif.Log;
+};
+
+type AutoModelQueryOptions = {
+  cliServer: CodeQLCliServer;
+  queryRunner: QueryRunner;
+  databaseItem: DatabaseItem;
+  qlpack: QlPacksForLanguage;
+  sourceInfo: SourceInfo | undefined;
+  extensionPacks: string[];
+  queryStorageDir: string;
+
+  progress: ProgressCallback;
+  token: CancellationToken;
+};
+
+function modeTag(mode: Mode): string {
+  switch (mode) {
+    case Mode.Application:
+      return "application-mode";
+    case Mode.Framework:
+      return "framework-mode";
+    default:
+      assertNever(mode);
+  }
+}
+
+async function runAutoModelQuery(
+  mode: Mode,
+  queryTag: string,
+  {
+    cliServer,
+    queryRunner,
+    databaseItem,
+    qlpack,
+    sourceInfo,
+    extensionPacks,
+    queryStorageDir,
+    progress,
+    token,
+  }: AutoModelQueryOptions,
+): Promise<Sarif.Log | undefined> {
+  // First, resolve the query that we want to run.
+  // All queries are tagged like this:
+  // internal extract automodel <mode> <queryTag>
+  // Example: internal extract automodel framework-mode candidates
+  const queries = await resolveQueries(
+    cliServer,
+    qlpack,
+    `Extract automodel ${queryTag}`,
+    {
+      kind: "problem",
+      "tags contain all": ["automodel", modeTag(mode), ...queryTag.split(" ")],
+    },
+  );
+  if (queries.length > 1) {
+    throw new Error(
+      `Found multiple auto model queries for ${mode} ${queryTag}. Can't continue`,
+    );
+  }
+  if (queries.length === 0) {
+    throw new Error(
+      `Did not found any auto model queries for ${mode} ${queryTag}. Can't continue`,
+    );
+  }
+
+  const queryPath = queries[0];
+  const { cleanup: cleanupLockFile } = await createLockFileForStandardQuery(
+    cliServer,
+    queryPath,
+  );
+
+  // Get metadata for the query. This is required to interpret the results. We already know the kind is problem
+  // (because of the constraint in resolveQueries), so we don't need any more checks on the metadata.
+  const metadata = await cliServer.resolveMetadata(queryPath);
+
+  const queryRun = queryRunner.createQueryRun(
+    databaseItem.databaseUri.fsPath,
+    {
+      queryPath,
+      quickEvalPosition: undefined,
+      quickEvalCountOnly: false,
+    },
+    false,
+    getOnDiskWorkspaceFolders(),
+    extensionPacks,
+    queryStorageDir,
+    undefined,
+    undefined,
+  );
+
+  const completedQuery = await queryRun.evaluate(
+    progress,
+    token,
+    new TeeLogger(queryRunner.logger, queryRun.outputDir.logPath),
+  );
+
+  await cleanupLockFile?.();
+
+  if (completedQuery.resultType !== QueryResultType.SUCCESS) {
+    void showAndLogExceptionWithTelemetry(
+      extLogger,
+      telemetryListener,
+      redactableError`Auto-model query ${queryTag} failed: ${
+        completedQuery.message ?? "No message"
+      }`,
+    );
+    return;
+  }
+
+  const interpretedResultsPath = join(
+    queryStorageDir,
+    `interpreted-results-${queryTag.replaceAll(" ", "-")}-${queryRun.id}.sarif`,
+  );
+
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars -- We only need the actual SARIF data, not the extra fields added by SarifInterpretationData
+  const { t, sortState, ...sarif } = await interpretResultsSarif(
+    cliServer,
+    metadata,
+    {
+      resultsPath: completedQuery.outputDir.bqrsPath,
+      interpretedResultsPath,
+    },
+    sourceInfo,
+    ["--sarif-add-snippets"],
+  );
+
+  return sarif;
+}
+
+export async function runAutoModelQueries(
+  mode: Mode,
+  { cliServer, queryRunner, databaseItem, queryStorageDir, progress }: Options,
+): 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
+  // snippets from it. The source location prefix is used to determine the
+  // base path of the database.
+  const sourceLocationPrefix = await databaseItem.getSourceLocationPrefix(
+    cliServer,
+  );
+  const sourceArchiveUri = databaseItem.sourceArchive;
+  const sourceInfo =
+    sourceArchiveUri === undefined
+      ? undefined
+      : {
+          sourceArchive: sourceArchiveUri.fsPath,
+          sourceLocationPrefix,
+        };
+
+  const additionalPacks = getOnDiskWorkspaceFolders();
+  const extensionPacks = Object.keys(
+    await cliServer.resolveQlpacks(additionalPacks, true),
+  );
+
+  progress({
+    step: 0,
+    maxStep,
+    message: "Finding candidates and examples",
+  });
+
+  const candidates = await runAutoModelQuery(mode, "candidates", {
+    cliServer,
+    queryRunner,
+    databaseItem,
+    qlpack,
+    sourceInfo,
+    extensionPacks,
+    queryStorageDir,
+    progress: (update) => {
+      progress({
+        step: update.step,
+        maxStep,
+        message: "Finding candidates and examples",
+      });
+    },
+    token: cancellationTokenSource.token,
+  });
+
+  if (!candidates) {
+    return undefined;
+  }
+
+  return {
+    candidates,
+  };
+}

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

@@ -0,0 +1,40 @@
+import { AutomodelMode, ModelRequest } from "./auto-model-api-v2";
+import { Mode } from "./shared/mode";
+import { AutoModelQueriesResult } from "./auto-model-codeml-queries";
+import { assertNever } from "../common/helpers-pure";
+import * as Sarif from "sarif";
+import { gzipEncode } from "../common/zlib";
+
+/**
+ * Encode a SARIF log to the format expected by the server: JSON, GZIP-compressed, base64-encoded
+ * @param log SARIF log to encode
+ * @returns base64-encoded GZIP-compressed SARIF log
+ */
+export async function encodeSarif(log: Sarif.Log): Promise<string> {
+  const json = JSON.stringify(log);
+  const buffer = Buffer.from(json, "utf-8");
+  const compressed = await gzipEncode(buffer);
+  return compressed.toString("base64");
+}
+
+export async function createAutoModelV2Request(
+  mode: Mode,
+  result: AutoModelQueriesResult,
+): Promise<ModelRequest> {
+  let requestMode: AutomodelMode;
+  switch (mode) {
+    case Mode.Application:
+      requestMode = AutomodelMode.Application;
+      break;
+    case Mode.Framework:
+      requestMode = AutomodelMode.Framework;
+      break;
+    default:
+      assertNever(mode);
+  }
+
+  return {
+    mode: requestMode,
+    candidates: await encodeSarif(result.candidates),
+  };
+}

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

@@ -39,17 +39,30 @@ 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,
 } from "./auto-model";
-import { enableFrameworkMode, showLlmGeneration } from "../config";
+import {
+  enableFrameworkMode,
+  showLlmGeneration,
+  useLlmGenerationV2,
+} from "../config";
 import { getAutoModelUsages } from "./auto-model-usages-query";
 import { Mode } from "./shared/mode";
 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";
 
 export class DataExtensionsEditorView extends AbstractWebview<
   ToDataExtensionsEditorMessage,
@@ -367,49 +380,100 @@ export class DataExtensionsEditorView extends AbstractWebview<
       message: "Retrieving usages",
     });
 
-    const usages = await getAutoModelUsages({
+    let predictedModeledMethods: Record<string, ModeledMethod>;
-      cliServer: this.cliServer,
-      queryRunner: this.queryRunner,
-      queryStorageDir: this.queryStorageDir,
-      databaseItem: this.databaseItem,
-      progress: (update) => this.showProgress(update, maxStep),
-    });
 
-    await this.showProgress({
+    if (useLlmGenerationV2()) {
-      step: 1800,
+      const usages = await runAutoModelQueries(this.mode, {
-      maxStep,
+        cliServer: this.cliServer,
-      message: "Creating request",
+        queryRunner: this.queryRunner,
-    });
+        queryStorageDir: this.queryStorageDir,
+        databaseItem: this.databaseItem,
+        progress: (update) => this.showProgress(update, maxStep),
+      });
+      if (!usages) {
+        return;
+      }
 
-    const request = createAutoModelRequest(
+      await this.showProgress({
-      this.databaseItem.language,
+        step: 1800,
-      externalApiUsages,
+        maxStep,
-      modeledMethods,
+        message: "Creating request",
-      usages,
+      });
-      this.mode,
-    );
 
-    await this.showProgress({
+      const request = await createAutoModelV2Request(this.mode, usages);
-      step: 2000,
-      maxStep,
-      message: "Sending request",
-    });
 
-    const response = await this.callAutoModelApi(request);
+      await this.showProgress({
-    if (!response) {
+        step: 2000,
-      return;
+        maxStep,
+        message: "Sending request",
+      });
+
+      const response = await this.callAutoModelApiV2(request);
+      if (!response) {
+        return;
+      }
+
+      await this.showProgress({
+        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,
+        databaseItem: this.databaseItem,
+        progress: (update) => this.showProgress(update, maxStep),
+      });
+
+      await this.showProgress({
+        step: 1800,
+        maxStep,
+        message: "Creating request",
+      });
+
+      const request = createAutoModelRequest(
+        this.databaseItem.language,
+        externalApiUsages,
+        modeledMethods,
+        usages,
+        this.mode,
+      );
+
+      await this.showProgress({
+        step: 2000,
+        maxStep,
+        message: "Sending request",
+      });
+
+      const response = await this.callAutoModelApi(request);
+      if (!response) {
+        return;
+      }
+
+      await this.showProgress({
+        step: 2500,
+        maxStep,
+        message: "Parsing response",
+      });
+
+      predictedModeledMethods = parsePredictedClassifications(
+        response.predicted || [],
+      );
     }
 
-    await this.showProgress({
-      step: 2500,
-      maxStep,
-      message: "Parsing response",
-    });
-
-    const predictedModeledMethods = parsePredictedClassifications(
-      response.predicted || [],
-    );
-
     await this.showProgress({
       step: 2800,
       maxStep,
@@ -534,4 +598,25 @@ export class DataExtensionsEditorView extends AbstractWebview<
       }
     }
   }
+
+  private async callAutoModelApiV2(
+    request: ModelRequestV2,
+  ): Promise<ModelResponseV2 | null> {
+    try {
+      return await autoModelV2(this.app.credentials, request);
+    } catch (e) {
+      await this.clearProgress();
+
+      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/test/unit-tests/data-extensions-editor/auto-model-v2.test.ts

@@ -0,0 +1,82 @@
+import {
+  createAutoModelV2Request,
+  encodeSarif,
+} 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";
+
+describe("createAutoModelV2Request", () => {
+  const createSarifLog = (queryId: string): sarif.Log => {
+    return {
+      version: "2.1.0",
+      $schema: "http://json.schemastore.org/sarif-2.1.0-rtm.4",
+      runs: [
+        {
+          tool: {
+            driver: {
+              name: "CodeQL",
+              rules: [
+                {
+                  id: queryId,
+                },
+              ],
+            },
+          },
+          results: [
+            {
+              message: {
+                text: "msg",
+              },
+              locations: [
+                {
+                  physicalLocation: {
+                    contextRegion: {
+                      startLine: 10,
+                      endLine: 12,
+                      snippet: {
+                        text: "Foo",
+                      },
+                    },
+                    region: {
+                      startLine: 10,
+                      startColumn: 1,
+                      endColumn: 3,
+                    },
+                    artifactLocation: {
+                      uri: "foo.js",
+                    },
+                  },
+                },
+              ],
+            },
+          ],
+        },
+      ],
+    };
+  };
+
+  const result: AutoModelQueriesResult = {
+    candidates: createSarifLog(
+      "java/ml/extract-automodel-application-candidates",
+    ),
+  };
+
+  it("creates a matching request", async () => {
+    expect(await createAutoModelV2Request(Mode.Application, result)).toEqual({
+      mode: AutomodelMode.Application,
+      candidates: await encodeSarif(result.candidates),
+    });
+  });
+
+  it("can decode the SARIF", async () => {
+    const request = await createAutoModelV2Request(Mode.Application, result);
+    const decoded = Buffer.from(request.candidates, "base64");
+    const decompressed = await gzipDecode(decoded);
+    const json = decompressed.toString("utf-8");
+    const parsed = JSON.parse(json);
+    expect(parsed).toEqual(result.candidates);
+  });
+});

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

@@ -0,0 +1,188 @@
+import { createMockLogger } from "../../../__mocks__/loggerMock";
+import type { Uri } from "vscode";
+import {
+  DatabaseItem,
+  DatabaseKind,
+} 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 { Mode } from "../../../../src/data-extensions-editor/shared/mode";
+import { mockedObject } 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";
+
+function createMockUri(path = "/a/b/c/foo"): Uri {
+  return {
+    scheme: "file",
+    authority: "",
+    path,
+    query: "",
+    fragment: "",
+    fsPath: path,
+    with: jest.fn(),
+    toJSON: jest.fn(),
+  };
+}
+
+describe("runAutoModelQueries", () => {
+  const qlpack = {
+    dbschemePack: "dbschemePack",
+    dbschemePackIsLibraryPack: false,
+  };
+
+  let resolveQueriesSpy: jest.SpiedFunction<
+    typeof queryResolver.resolveQueries
+  >;
+  let createLockFileForStandardQuerySpy: jest.SpiedFunction<
+    typeof standardQueries.createLockFileForStandardQuery
+  >;
+
+  beforeEach(() => {
+    jest.spyOn(queryResolver, "qlpackOfDatabase").mockResolvedValue(qlpack);
+
+    resolveQueriesSpy = jest
+      .spyOn(queryResolver, "resolveQueries")
+      .mockImplementation(async (_cliServer, _qlPack, _name, constraints) => {
+        if (constraints["tags contain all"]?.includes("candidates")) {
+          return ["/a/b/c/ql/candidates.ql"];
+        }
+        if (constraints["tags contain all"]?.includes("positive")) {
+          return ["/a/b/c/ql/positive-examples.ql"];
+        }
+        if (constraints["tags contain all"]?.includes("negative")) {
+          return ["/a/b/c/ql/negative-examples.ql"];
+        }
+
+        return [];
+      });
+
+    createLockFileForStandardQuerySpy = jest
+      .spyOn(standardQueries, "createLockFileForStandardQuery")
+      .mockResolvedValue({});
+  });
+
+  it("should run the query and return the results", async () => {
+    const logPath = (await file()).path;
+    const bqrsPath = (await file()).path;
+    const outputDir = {
+      logPath,
+      bqrsPath,
+    };
+
+    const options = {
+      cliServer: mockedObject<CodeQLCliServer>({
+        resolveQlpacks: jest.fn().mockResolvedValue({
+          "/a/b/c/my-extension-pack": {},
+        }),
+        resolveMetadata: jest.fn().mockResolvedValue({
+          kind: "problem",
+        }),
+        interpretBqrsSarif: jest.fn().mockResolvedValue({
+          version: "2.1.0",
+          $schema: "http://json.schemastore.org/sarif-2.1.0-rtm.4",
+          runs: [
+            {
+              tool: {
+                driver: {
+                  name: "CodeQL",
+                },
+              },
+              results: [
+                {
+                  message: {
+                    text: "msg",
+                  },
+                  locations: [
+                    {
+                      physicalLocation: {
+                        contextRegion: {
+                          startLine: 10,
+                          endLine: 12,
+                          snippet: {
+                            text: "Foo",
+                          },
+                        },
+                        region: {
+                          startLine: 10,
+                          startColumn: 1,
+                          endColumn: 3,
+                        },
+                        artifactLocation: {
+                          uri: "foo.js",
+                        },
+                      },
+                    },
+                  ],
+                },
+              ],
+            },
+          ],
+        }),
+      }),
+      queryRunner: mockedObject<QueryRunner>({
+        createQueryRun: jest.fn().mockReturnValue({
+          evaluate: jest.fn().mockResolvedValue({
+            resultType: QueryResultType.SUCCESS,
+            outputDir,
+          }),
+          outputDir,
+        }),
+        logger: createMockLogger(),
+      }),
+      databaseItem: mockedObject<DatabaseItem>({
+        databaseUri: createMockUri("/a/b/c/src.zip"),
+        contents: {
+          kind: DatabaseKind.Database,
+          name: "foo",
+          datasetUri: createMockUri(),
+        },
+        language: "java",
+        getSourceLocationPrefix: jest
+          .fn()
+          .mockResolvedValue("/home/runner/work/my-repo/my-repo"),
+        sourceArchive: createMockUri("/a/b/c/src.zip"),
+      }),
+      queryStorageDir: "/tmp/queries",
+      progress: jest.fn(),
+    };
+
+    const result = await runAutoModelQueries(Mode.Application, options);
+    expect(result).not.toBeUndefined();
+
+    expect(options.cliServer.resolveQlpacks).toHaveBeenCalledTimes(1);
+    expect(options.cliServer.resolveQlpacks).toHaveBeenCalledWith([], true);
+    expect(resolveQueriesSpy).toHaveBeenCalledTimes(1);
+    expect(resolveQueriesSpy).toHaveBeenCalledWith(
+      options.cliServer,
+      qlpack,
+      "Extract automodel candidates",
+      {
+        kind: "problem",
+        "tags contain all": ["automodel", "application-mode", "candidates"],
+      },
+    );
+    expect(createLockFileForStandardQuerySpy).toHaveBeenCalledTimes(1);
+    expect(createLockFileForStandardQuerySpy).toHaveBeenCalledWith(
+      options.cliServer,
+      "/a/b/c/ql/candidates.ql",
+    );
+    expect(options.queryRunner.createQueryRun).toHaveBeenCalledTimes(1);
+    expect(options.queryRunner.createQueryRun).toHaveBeenCalledWith(
+      "/a/b/c/src.zip",
+      {
+        queryPath: "/a/b/c/ql/candidates.ql",
+        quickEvalPosition: undefined,
+        quickEvalCountOnly: false,
+      },
+      false,
+      [],
+      ["/a/b/c/my-extension-pack"],
+      "/tmp/queries",
+      undefined,
+      undefined,
+    );
+  });
+});

extensions/ql-vscode/test/vscode-tests/utils/mocking.helpers.ts

@@ -40,6 +40,18 @@ export function mockedObject<T extends object>(
         return undefined;
       }
 
+      // The `asymmetricMatch` is accessed by jest to check if the object is a matcher.
+      // We don't want to throw an error when this happens.
+      if (prop === "asymmetricMatch") {
+        return undefined;
+      }
+
+      // The `Symbol.iterator` is accessed by jest to check if the object is iterable.
+      // We don't want to throw an error when this happens.
+      if (prop === Symbol.iterator) {
+        return undefined;
+      }
+
       throw new Error(`Method ${String(prop)} not mocked`);
     },
   });