Merge branch 'main' into robertbrignull/JumpToUsageMessage

2023-10-10 10:50:21 +01:00
parent 47509a922a 319a54c32f
commit 6561bb0543
41 changed files with 585 additions and 233 deletions
--- a/extensions/ql-vscode/CHANGELOG.md
+++ b/extensions/ql-vscode/CHANGELOG.md
@@ -7,6 +7,7 @@
 - Increase the required version of VS Code to 1.82.0. [#2877](https://github.com/github/vscode-codeql/pull/2877)
 - Fix a bug where the query server was restarted twice after configuration changes. [#2884](https://github.com/github/vscode-codeql/pull/2884).
 - Add support for the `telemetry.telemetryLevel` setting. For more information, see the [telemetry documentation](https://codeql.github.com/docs/codeql-for-visual-studio-code/about-telemetry-in-codeql-for-visual-studio-code). [#2824](https://github.com/github/vscode-codeql/pull/2824).
+- Fix syntax highlighting directly after import statements with instantiation arguments. [#2792](https://github.com/github/vscode-codeql/pull/2792)

 ## 1.9.1 - 29 September 2023

--- a/extensions/ql-vscode/package.json
+++ b/extensions/ql-vscode/package.json
@@ -110,6 +110,10 @@
                  "string"
                ],
                "description": "Names of extension packs to include in the evaluation. These are resolved from the locations specified in `additionalPacks`."
+              },
+              "additionalRunQueryArgs": {
+                "type": "object",
+                "description": "**Internal use only**. Additional arguments to pass to the `runQuery` command of the query server, without validation."
              }
            }
          }
--- a/extensions/ql-vscode/src/common/interface-types.ts
+++ b/extensions/ql-vscode/src/common/interface-types.ts
@@ -555,8 +555,7 @@ interface GenerateMethodMessage {
 interface GenerateMethodsFromLlmMessage {
  t: "generateMethodsFromLlm";
  packageName: string;
-  methods: Method[];
-  modeledMethods: Record<string, ModeledMethod>;
+  methodSignatures: string[];
 }

 interface StopGeneratingMethodsFromLlmMessage {
@@ -633,7 +632,7 @@ interface SetMethodModelingPanelViewStateMessage {

 interface SetMethodMessage {
  t: "setMethod";
-  method: Method;
+  method: Method | undefined;
 }

 interface SetMethodModifiedMessage {
--- a/extensions/ql-vscode/src/debugger/debug-configuration.ts
+++ b/extensions/ql-vscode/src/debugger/debug-configuration.ts
@@ -22,6 +22,7 @@ export interface QLDebugArgs {
  extensionPacks?: string[] | string;
  quickEval?: boolean;
  noDebug?: boolean;
+  additionalRunQueryArgs?: Record<string, any>;
 }

 /**
@@ -120,6 +121,7 @@ export class QLDebugConfigurationProvider
        extensionPacks,
        quickEvalContext,
        noDebug: qlConfiguration.noDebug ?? false,
+        additionalRunQueryArgs: qlConfiguration.additionalRunQueryArgs ?? {},
      };

      return resultConfiguration;
--- a/extensions/ql-vscode/src/debugger/debug-protocol.ts
+++ b/extensions/ql-vscode/src/debugger/debug-protocol.ts
@@ -70,6 +70,8 @@ export interface LaunchConfig {
  quickEvalContext: QuickEvalContext | undefined;
  /** Run the query without debugging it. */
  noDebug: boolean;
+  /** Undocumented: Additional arguments to be passed to the `runQuery` API on the query server. */
+  additionalRunQueryArgs: Record<string, any>;
 }

 export interface LaunchRequest extends Request, DebugProtocol.LaunchRequest {
--- a/extensions/ql-vscode/src/debugger/debug-session.ts
+++ b/extensions/ql-vscode/src/debugger/debug-session.ts
@@ -161,6 +161,7 @@ class RunningQuery extends DisposableObject {
      true,
      config.additionalPacks,
      config.extensionPacks,
+      config.additionalRunQueryArgs,
      queryStorageDir,
      undefined,
      undefined,
--- a/extensions/ql-vscode/src/language-support/contextual/query-resolver.ts
+++ b/extensions/ql-vscode/src/language-support/contextual/query-resolver.ts
@@ -44,6 +44,7 @@ export async function runContextualQuery(
    false,
    getOnDiskWorkspaceFolders(),
    undefined,
+    {},
    queryStorageDir,
    undefined,
    templates,
--- a/extensions/ql-vscode/src/local-queries/local-queries.ts
+++ b/extensions/ql-vscode/src/local-queries/local-queries.ts
@@ -456,6 +456,7 @@ export class LocalQueries extends DisposableObject {
      true,
      additionalPacks,
      extensionPacks,
+      {},
      this.queryStorageDir,
      undefined,
      templates,
--- a/extensions/ql-vscode/src/local-queries/run-query.ts
+++ b/extensions/ql-vscode/src/local-queries/run-query.ts
@@ -41,6 +41,7 @@ export async function runQuery({
    false,
    additionalPacks,
    extensionPacks,
+    {},
    queryStorageDir,
    undefined,
    undefined,
--- a/extensions/ql-vscode/src/model-editor/auto-model.ts
+++ b/extensions/ql-vscode/src/model-editor/auto-model.ts
@@ -14,13 +14,13 @@ import { groupMethods, sortGroupNames, sortMethods } from "./shared/sorting";
 * the order in the UI.
 * @param mode Whether it is application or framework mode.
 * @param methods all methods.
- * @param modeledMethods the currently modeled methods.
+ * @param modeledMethodsBySignature the currently modeled methods.
 * @returns list of modeled methods that are candidates for modeling.
 */
 export function getCandidates(
  mode: Mode,
  methods: Method[],
-  modeledMethods: Record<string, ModeledMethod>,
+  modeledMethodsBySignature: 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(methods, mode);
@@ -32,12 +32,11 @@ export function getCandidates(
  const candidates: MethodSignature[] = [];

  for (const method of sortedMethods) {
-    const modeledMethod: ModeledMethod = modeledMethods[method.signature] ?? {
-      type: "none",
-    };
+    const modeledMethods: ModeledMethod[] =
+      modeledMethodsBySignature[method.signature] ?? [];

    // Anything that is modeled is not a candidate
-    if (modeledMethod.type !== "none") {
+    if (modeledMethods.some((m) => m.type !== "none")) {
      continue;
    }

--- a/extensions/ql-vscode/src/model-editor/auto-modeler.ts
+++ b/extensions/ql-vscode/src/model-editor/auto-modeler.ts
@@ -16,7 +16,6 @@ import { QueryRunner } from "../query-server";
 import { DatabaseItem } from "../databases/local-databases";
 import { Mode } from "./shared/mode";
 import { CancellationTokenSource } from "vscode";
-import { convertToLegacyModeledMethods } from "./modeled-methods-legacy";

 // Limit the number of candidates we send to the model in each request
 // to avoid long requests.
@@ -43,7 +42,7 @@ export class AutoModeler {
      inProgressMethods: string[],
    ) => Promise<void>,
    private readonly addModeledMethods: (
-      modeledMethods: Record<string, ModeledMethod>,
+      modeledMethods: Record<string, ModeledMethod[]>,
    ) => Promise<void>,
  ) {
    this.jobs = new Map<string, CancellationTokenSource>();
@@ -60,7 +59,7 @@ export class AutoModeler {
  public async startModeling(
    packageName: string,
    methods: Method[],
-    modeledMethods: Record<string, ModeledMethod>,
+    modeledMethods: Record<string, ModeledMethod[]>,
    mode: Mode,
  ): Promise<void> {
    if (this.jobs.has(packageName)) {
@@ -107,7 +106,7 @@ export class AutoModeler {
  private async modelPackage(
    packageName: string,
    methods: Method[],
-    modeledMethods: Record<string, ModeledMethod>,
+    modeledMethods: Record<string, ModeledMethod[]>,
    mode: Mode,
    cancellationTokenSource: CancellationTokenSource,
  ): Promise<void> {
@@ -193,31 +192,31 @@ export class AutoModeler {
      filename: "auto-model.yml",
    });

-    const rawLoadedMethods = loadDataExtensionYaml(models);
-    if (!rawLoadedMethods) {
+    const loadedMethods = loadDataExtensionYaml(models);
+    if (!loadedMethods) {
      return;
    }

-    const loadedMethods = convertToLegacyModeledMethods(rawLoadedMethods);
-
    // 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,
-        };
+        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,
+          },
+        ];
      }
    }

--- a/extensions/ql-vscode/src/model-editor/method-modeling/method-modeling-panel.ts
+++ b/extensions/ql-vscode/src/model-editor/method-modeling/method-modeling-panel.ts
@@ -6,6 +6,7 @@ import { Method } from "../method";
 import { ModelingStore } from "../modeling-store";
 import { ModelEditorViewTracker } from "../model-editor-view-tracker";
 import { ModelConfigListener } from "../../config";
+import { DatabaseItem } from "../../databases/local-databases";

 export class MethodModelingPanel extends DisposableObject {
  private readonly provider: MethodModelingViewProvider;
@@ -36,7 +37,10 @@ export class MethodModelingPanel extends DisposableObject {
    );
  }

-  public async setMethod(method: Method): Promise<void> {
-    await this.provider.setMethod(method);
+  public async setMethod(
+    databaseItem: DatabaseItem,
+    method: Method,
+  ): Promise<void> {
+    await this.provider.setMethod(databaseItem, method);
  }
 }
--- a/extensions/ql-vscode/src/model-editor/method-modeling/method-modeling-view-provider.ts
+++ b/extensions/ql-vscode/src/model-editor/method-modeling/method-modeling-view-provider.ts
@@ -13,6 +13,11 @@ import { AbstractWebviewViewProvider } from "../../common/vscode/abstract-webvie
 import { assertNever } from "../../common/helpers-pure";
 import { ModelEditorViewTracker } from "../model-editor-view-tracker";
 import { ModelConfigListener } from "../../config";
+import { DatabaseItem } from "../../databases/local-databases";
+import {
+  convertFromLegacyModeledMethod,
+  convertToLegacyModeledMethod,
+} from "../modeled-methods-legacy";

 export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
  ToMethodModelingMessage,
@@ -21,6 +26,7 @@ export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
  public static readonly viewType = "codeQLMethodModeling";

  private method: Method | undefined = undefined;
+  private databaseItem: DatabaseItem | undefined = undefined;

  constructor(
    app: App,
@@ -46,8 +52,12 @@ export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
    });
  }

-  public async setMethod(method: Method): Promise<void> {
+  public async setMethod(
+    databaseItem: DatabaseItem | undefined,
+    method: Method | undefined,
+  ): Promise<void> {
    this.method = method;
+    this.databaseItem = databaseItem;

    if (this.isShowingView) {
      await this.postMessage({
@@ -64,10 +74,17 @@ export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
        await this.postMessage({
          t: "setSelectedMethod",
          method: selectedMethod.method,
-          modeledMethod: selectedMethod.modeledMethod,
+          modeledMethod: convertToLegacyModeledMethod(
+            selectedMethod.modeledMethods,
+          ),
          isModified: selectedMethod.isModified,
        });
      }
+
+      await this.postMessage({
+        t: "setInModelingMode",
+        inModelingMode: true,
+      });
    }
  }

@@ -96,9 +113,14 @@ export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
      case "setModeledMethod": {
        const activeState = this.ensureActiveState();

-        this.modelingStore.updateModeledMethod(
+        this.modelingStore.updateModeledMethods(
          activeState.databaseItem,
-          msg.method,
+          msg.method.signature,
+          convertFromLegacyModeledMethod(msg.method),
+        );
+        this.modelingStore.addModifiedMethod(
+          activeState.databaseItem,
+          msg.method.signature,
        );
        break;
      }
@@ -143,12 +165,15 @@ export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
    this.push(
      this.modelingStore.onModeledMethodsChanged(async (e) => {
        if (this.webviewView && e.isActiveDb) {
-          const modeledMethod = e.modeledMethods[this.method?.signature ?? ""];
-          if (modeledMethod) {
-            await this.postMessage({
-              t: "setModeledMethod",
-              method: modeledMethod,
-            });
+          const modeledMethods = e.modeledMethods[this.method?.signature ?? ""];
+          if (modeledMethods) {
+            const modeledMethod = convertToLegacyModeledMethod(modeledMethods);
+            if (modeledMethod) {
+              await this.postMessage({
+                t: "setModeledMethod",
+                method: modeledMethod,
+              });
+            }
          }
        }
      }),
@@ -170,10 +195,12 @@ export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
      this.modelingStore.onSelectedMethodChanged(async (e) => {
        if (this.webviewView) {
          this.method = e.method;
+          this.databaseItem = e.databaseItem;
+
          await this.postMessage({
            t: "setSelectedMethod",
            method: e.method,
-            modeledMethod: e.modeledMethod,
+            modeledMethod: convertToLegacyModeledMethod(e.modeledMethods),
            isModified: e.isModified,
          });
        }
@@ -190,13 +217,17 @@ export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
    );

    this.push(
-      this.modelingStore.onDbClosed(async () => {
+      this.modelingStore.onDbClosed(async (dbUri) => {
        if (!this.modelingStore.anyDbsBeingModeled()) {
          await this.postMessage({
            t: "setInModelingMode",
            inModelingMode: false,
          });
        }
+
+        if (dbUri === this.databaseItem?.databaseUri.toString()) {
+          await this.setMethod(undefined, undefined);
+        }
      }),
    );
  }
--- a/extensions/ql-vscode/src/model-editor/methods-usage/methods-usage-data-provider.ts
+++ b/extensions/ql-vscode/src/model-editor/methods-usage/methods-usage-data-provider.ts
@@ -26,7 +26,7 @@ export class MethodsUsageDataProvider
  private databaseItem: DatabaseItem | undefined = undefined;
  private sourceLocationPrefix: string | undefined = undefined;
  private hideModeledMethods: boolean = INITIAL_HIDE_MODELED_METHODS_VALUE;
-  private modeledMethods: Record<string, ModeledMethod> = {};
+  private modeledMethods: Record<string, ModeledMethod[]> = {};
  private modifiedMethodSignatures: Set<string> = new Set();

  private readonly onDidChangeTreeDataEmitter = this.push(
@@ -52,7 +52,7 @@ export class MethodsUsageDataProvider
    methods: Method[],
    databaseItem: DatabaseItem,
    hideModeledMethods: boolean,
-    modeledMethods: Record<string, ModeledMethod>,
+    modeledMethods: Record<string, ModeledMethod[]>,
    modifiedMethodSignatures: Set<string>,
  ): Promise<void> {
    if (
@@ -102,10 +102,10 @@ export class MethodsUsageDataProvider
  }

  private getModelingStatusIcon(method: Method): ThemeIcon {
-    const modeledMethod = this.modeledMethods[method.signature];
+    const modeledMethods = this.modeledMethods[method.signature];
    const modifiedMethod = this.modifiedMethodSignatures.has(method.signature);

-    const status = getModelingStatus(modeledMethod, modifiedMethod);
+    const status = getModelingStatus(modeledMethods, modifiedMethod);
    switch (status) {
      case "unmodeled":
        return new ThemeIcon("error", new ThemeColor("errorForeground"));
--- a/extensions/ql-vscode/src/model-editor/methods-usage/methods-usage-panel.ts
+++ b/extensions/ql-vscode/src/model-editor/methods-usage/methods-usage-panel.ts
@@ -34,7 +34,7 @@ export class MethodsUsagePanel extends DisposableObject {
    methods: Method[],
    databaseItem: DatabaseItem,
    hideModeledMethods: boolean,
-    modeledMethods: Record<string, ModeledMethod>,
+    modeledMethods: Record<string, ModeledMethod[]>,
    modifiedMethodSignatures: Set<string>,
  ): Promise<void> {
    await this.dataProvider.setState(
--- a/extensions/ql-vscode/src/model-editor/model-editor-module.ts
+++ b/extensions/ql-vscode/src/model-editor/model-editor-module.ts
@@ -104,7 +104,7 @@ export class ModelEditorModule extends DisposableObject {
    usage: Usage,
  ): Promise<void> {
    await this.methodsUsagePanel.revealItem(usage);
-    await this.methodModelingPanel.setMethod(method);
+    await this.methodModelingPanel.setMethod(databaseItem, method);
    await showResolvableLocation(usage.url, databaseItem, this.app.logger);
  }

--- a/extensions/ql-vscode/src/model-editor/model-editor-view.ts
+++ b/extensions/ql-vscode/src/model-editor/model-editor-view.ts
@@ -44,7 +44,7 @@ import { telemetryListener } from "../common/vscode/telemetry";
 import { ModelingStore } from "./modeling-store";
 import { ModelEditorViewTracker } from "./model-editor-view-tracker";
 import {
-  convertFromLegacyModeledMethods,
+  convertFromLegacyModeledMethod,
  convertToLegacyModeledMethods,
 } from "./modeled-methods-legacy";

@@ -226,7 +226,7 @@ export class ModelEditorView extends AbstractWebview<
                this.extensionPack,
                this.databaseItem.language,
                methods,
-                convertFromLegacyModeledMethods(modeledMethods),
+                modeledMethods,
                this.mode,
                this.cliServer,
                this.app.logger,
@@ -269,8 +269,7 @@ export class ModelEditorView extends AbstractWebview<
      case "generateMethodsFromLlm":
        await this.generateModeledMethodsFromLlm(
          msg.packageName,
-          msg.methods,
-          msg.modeledMethods,
+          msg.methodSignatures,
        );
        void telemetryListener?.sendUIInteraction(
          "model-editor-generate-methods-from-llm",
@@ -314,7 +313,10 @@ export class ModelEditorView extends AbstractWebview<
        );
        break;
      case "setModeledMethod": {
-        this.setModeledMethod(msg.method);
+        this.setModeledMethods(
+          msg.method.signature,
+          convertFromLegacyModeledMethod(msg.method),
+        );
        break;
      }
      default:
@@ -374,10 +376,7 @@ export class ModelEditorView extends AbstractWebview<
        this.cliServer,
        this.app.logger,
      );
-      this.modelingStore.setModeledMethods(
-        this.databaseItem,
-        convertToLegacyModeledMethods(modeledMethods),
-      );
+      this.modelingStore.setModeledMethods(this.databaseItem, modeledMethods);
    } catch (e: unknown) {
      void showAndLogErrorMessage(
        this.app.logger,
@@ -449,10 +448,16 @@ export class ModelEditorView extends AbstractWebview<
            queryStorageDir: this.queryStorageDir,
            databaseItem: addedDatabase ?? this.databaseItem,
            onResults: async (modeledMethods) => {
-              const modeledMethodsByName: Record<string, ModeledMethod> = {};
+              const modeledMethodsByName: Record<string, ModeledMethod[]> = {};

              for (const modeledMethod of modeledMethods) {
-                modeledMethodsByName[modeledMethod.signature] = modeledMethod;
+                if (!(modeledMethod.signature in modeledMethodsByName)) {
+                  modeledMethodsByName[modeledMethod.signature] = [];
+                }
+
+                modeledMethodsByName[modeledMethod.signature].push(
+                  modeledMethod,
+                );
              }

              this.addModeledMethods(modeledMethodsByName);
@@ -476,9 +481,16 @@ export class ModelEditorView extends AbstractWebview<

  private async generateModeledMethodsFromLlm(
    packageName: string,
-    methods: Method[],
-    modeledMethods: Record<string, ModeledMethod>,
+    methodSignatures: string[],
  ): Promise<void> {
+    const methods = this.modelingStore.getMethods(
+      this.databaseItem,
+      methodSignatures,
+    );
+    const modeledMethods = this.modelingStore.getModeledMethods(
+      this.databaseItem,
+      methodSignatures,
+    );
    await this.autoModeler.startModeling(
      packageName,
      methods,
@@ -616,7 +628,7 @@ export class ModelEditorView extends AbstractWebview<
        if (event.dbUri === this.databaseItem.databaseUri.toString()) {
          await this.postMessage({
            t: "setModeledMethods",
-            methods: event.modeledMethods,
+            methods: convertToLegacyModeledMethods(event.modeledMethods),
          });
        }
      }),
@@ -642,7 +654,7 @@ export class ModelEditorView extends AbstractWebview<
    );
  }

-  private addModeledMethods(modeledMethods: Record<string, ModeledMethod>) {
+  private addModeledMethods(modeledMethods: Record<string, ModeledMethod[]>) {
    this.modelingStore.addModeledMethods(this.databaseItem, modeledMethods);

    this.modelingStore.addModifiedMethods(
@@ -651,13 +663,17 @@ export class ModelEditorView extends AbstractWebview<
    );
  }

-  private setModeledMethod(method: ModeledMethod) {
+  private setModeledMethods(signature: string, methods: ModeledMethod[]) {
    const state = this.modelingStore.getStateForActiveDb();
    if (!state) {
      throw new Error("Attempting to set modeled method without active db");
    }

-    this.modelingStore.updateModeledMethod(state.databaseItem, method);
-    this.modelingStore.addModifiedMethod(state.databaseItem, method.signature);
+    this.modelingStore.updateModeledMethods(
+      state.databaseItem,
+      signature,
+      methods,
+    );
+    this.modelingStore.addModifiedMethod(state.databaseItem, signature);
  }
 }
--- a/extensions/ql-vscode/src/model-editor/modeled-methods-legacy.ts
+++ b/extensions/ql-vscode/src/model-editor/modeled-methods-legacy.ts
@@ -1,23 +1,57 @@
 import { ModeledMethod } from "./modeled-method";

-export function convertFromLegacyModeledMethods(
-  modeledMethods: Record<string, ModeledMethod>,
-): Record<string, ModeledMethod[]> {
-  // Convert a single ModeledMethod to an array of ModeledMethods
-  return Object.fromEntries(
-    Object.entries(modeledMethods).map(([signature, modeledMethod]) => {
-      return [signature, [modeledMethod]];
-    }),
-  );
-}
-
+/**
+ * Converts a record of a single ModeledMethod indexed by signature to a record of ModeledMethod[] indexed by signature
+ * for legacy usage. This function should always be used instead of the trivial conversion to track usages of this
+ * conversion.
+ *
+ * This method should only be called inside a `postMessage` call. If it's used anywhere else, consider whether the
+ * boundary is correct: the boundary should as close as possible to the extension host -> webview boundary.
+ *
+ * @param modeledMethods The record of a single ModeledMethod indexed by signature
+ */
 export function convertToLegacyModeledMethods(
  modeledMethods: Record<string, ModeledMethod[]>,
 ): Record<string, ModeledMethod> {
  // Always take the first modeled method in the array
  return Object.fromEntries(
-    Object.entries(modeledMethods).map(([signature, modeledMethods]) => {
-      return [signature, modeledMethods[0]];
-    }),
+    Object.entries(modeledMethods)
+      .map(([signature, modeledMethods]) => {
+        const modeledMethod = convertToLegacyModeledMethod(modeledMethods);
+        if (!modeledMethod) {
+          return null;
+        }
+        return [signature, modeledMethod];
+      })
+      .filter((entry): entry is [string, ModeledMethod] => entry !== null),
  );
 }
+
+/**
+ * Converts a single ModeledMethod to a ModeledMethod[] for legacy usage. This function should always be used instead
+ * of the trivial conversion to track usages of this conversion.
+ *
+ * This method should only be called inside a `onMessage` function (or its equivalent). If it's used anywhere else,
+ * consider whether the boundary is correct: the boundary should as close as possible to the webview -> extension host
+ * boundary.
+ *
+ * @param modeledMethod The single ModeledMethod
+ */
+export function convertFromLegacyModeledMethod(modeledMethod: ModeledMethod) {
+  return [modeledMethod];
+}
+
+/**
+ * Converts a ModeledMethod[] to a single ModeledMethod for legacy usage. This function should always be used instead
+ * of the trivial conversion to track usages of this conversion.
+ *
+ * This method should only be called inside a `postMessage` call. If it's used anywhere else, consider whether the
+ * boundary is correct: the boundary should as close as possible to the extension host -> webview boundary.
+ *
+ * @param modeledMethods The ModeledMethod[]
+ */
+export function convertToLegacyModeledMethod(
+  modeledMethods: ModeledMethod[],
+): ModeledMethod | undefined {
+  return modeledMethods[0];
+}
--- a/extensions/ql-vscode/src/model-editor/modeling-store.ts
+++ b/extensions/ql-vscode/src/model-editor/modeling-store.ts
@@ -10,7 +10,7 @@ export interface DbModelingState {
  databaseItem: DatabaseItem;
  methods: Method[];
  hideModeledMethods: boolean;
-  modeledMethods: Record<string, ModeledMethod>;
+  modeledMethods: Record<string, ModeledMethod[]>;
  modifiedMethodSignatures: Set<string>;
  selectedMethod: Method | undefined;
  selectedUsage: Usage | undefined;
@@ -28,7 +28,7 @@ interface HideModeledMethodsChangedEvent {
 }

 interface ModeledMethodsChangedEvent {
-  modeledMethods: Record<string, ModeledMethod>;
+  modeledMethods: Record<string, ModeledMethod[]>;
  dbUri: string;
  isActiveDb: boolean;
 }
@@ -43,7 +43,7 @@ interface SelectedMethodChangedEvent {
  databaseItem: DatabaseItem;
  method: Method;
  usage: Usage;
-  modeledMethod: ModeledMethod | undefined;
+  modeledMethods: ModeledMethod[];
  isModified: boolean;
 }

@@ -221,7 +221,7 @@ export class ModelingStore extends DisposableObject {
  public getModeledMethods(
    dbItem: DatabaseItem,
    methodSignatures?: string[],
-  ): Record<string, ModeledMethod> {
+  ): Record<string, ModeledMethod[]> {
    const modeledMethods = this.getState(dbItem).modeledMethods;
    if (!methodSignatures) {
      return modeledMethods;
@@ -235,14 +235,15 @@ export class ModelingStore extends DisposableObject {

  public addModeledMethods(
    dbItem: DatabaseItem,
-    methods: Record<string, ModeledMethod>,
+    methods: Record<string, ModeledMethod[]>,
  ) {
    this.changeModeledMethods(dbItem, (state) => {
      const newModeledMethods = {
        ...methods,
+        // Keep all methods that are already modeled in some form in the state
        ...Object.fromEntries(
-          Object.entries(state.modeledMethods).filter(
-            ([_, value]) => value.type !== "none",
+          Object.entries(state.modeledMethods).filter(([_, value]) =>
+            value.some((m) => m.type !== "none"),
          ),
        ),
      };
@@ -252,17 +253,21 @@ export class ModelingStore extends DisposableObject {

  public setModeledMethods(
    dbItem: DatabaseItem,
-    methods: Record<string, ModeledMethod>,
+    methods: Record<string, ModeledMethod[]>,
  ) {
    this.changeModeledMethods(dbItem, (state) => {
      state.modeledMethods = { ...methods };
    });
  }

-  public updateModeledMethod(dbItem: DatabaseItem, method: ModeledMethod) {
+  public updateModeledMethods(
+    dbItem: DatabaseItem,
+    signature: string,
+    modeledMethods: ModeledMethod[],
+  ) {
    this.changeModeledMethods(dbItem, (state) => {
      const newModeledMethods = { ...state.modeledMethods };
-      newModeledMethods[method.signature] = method;
+      newModeledMethods[signature] = modeledMethods;
      state.modeledMethods = newModeledMethods;
    });
  }
@@ -325,7 +330,7 @@ export class ModelingStore extends DisposableObject {
      databaseItem: dbItem,
      method,
      usage,
-      modeledMethod: dbState.modeledMethods[method.signature],
+      modeledMethods: dbState.modeledMethods[method.signature],
      isModified: dbState.modifiedMethodSignatures.has(method.signature),
    });
  }
@@ -344,7 +349,7 @@ export class ModelingStore extends DisposableObject {
    return {
      method: selectedMethod,
      usage: dbState.selectedUsage,
-      modeledMethod: dbState.modeledMethods[selectedMethod.signature],
+      modeledMethods: dbState.modeledMethods[selectedMethod.signature],
      isModified: dbState.modifiedMethodSignatures.has(
        selectedMethod.signature,
      ),
--- a/extensions/ql-vscode/src/model-editor/shared/modeling-status.ts
+++ b/extensions/ql-vscode/src/model-editor/shared/modeling-status.ts
@@ -3,13 +3,13 @@ import { ModeledMethod } from "../modeled-method";
 export type ModelingStatus = "unmodeled" | "unsaved" | "saved";

 export function getModelingStatus(
-  modeledMethod: ModeledMethod | undefined,
+  modeledMethods: ModeledMethod[],
  methodIsUnsaved: boolean,
 ): ModelingStatus {
-  if (modeledMethod) {
+  if (modeledMethods.length > 0) {
    if (methodIsUnsaved) {
      return "unsaved";
-    } else if (modeledMethod.type !== "none") {
+    } else if (modeledMethods.some((m) => m.type !== "none")) {
      return "saved";
    }
  }
--- a/extensions/ql-vscode/src/query-server/legacy/legacy-query-runner.ts
+++ b/extensions/ql-vscode/src/query-server/legacy/legacy-query-runner.ts
@@ -65,6 +65,7 @@ export class LegacyQueryRunner extends QueryRunner {
    query: CoreQueryTarget,
    additionalPacks: string[],
    extensionPacks: string[] | undefined,
+    _additionalRunQueryArgs: Record<string, any>, // Ignored in legacy query server
    generateEvalLog: boolean,
    outputDir: QueryOutputDir,
    progress: ProgressCallback,
--- a/extensions/ql-vscode/src/query-server/new-query-runner.ts
+++ b/extensions/ql-vscode/src/query-server/new-query-runner.ts
@@ -75,6 +75,7 @@ export class NewQueryRunner extends QueryRunner {
    query: CoreQueryTarget,
    additionalPacks: string[],
    extensionPacks: string[] | undefined,
+    additionalRunQueryArgs: Record<string, any>,
    generateEvalLog: boolean,
    outputDir: QueryOutputDir,
    progress: ProgressCallback,
@@ -89,6 +90,7 @@ export class NewQueryRunner extends QueryRunner {
      generateEvalLog,
      additionalPacks,
      extensionPacks,
+      additionalRunQueryArgs,
      outputDir,
      progress,
      token,
--- a/extensions/ql-vscode/src/query-server/query-runner.ts
+++ b/extensions/ql-vscode/src/query-server/query-runner.ts
@@ -75,6 +75,7 @@ export abstract class QueryRunner {
    query: CoreQueryTarget,
    additionalPacks: string[],
    extensionPacks: string[] | undefined,
+    additionalRunQueryArgs: Record<string, any>,
    generateEvalLog: boolean,
    outputDir: QueryOutputDir,
    progress: ProgressCallback,
@@ -107,6 +108,7 @@ export abstract class QueryRunner {
    generateEvalLog: boolean,
    additionalPacks: string[],
    extensionPacks: string[] | undefined,
+    additionalRunQueryArgs: Record<string, any>,
    queryStorageDir: string,
    id = `${basename(query.queryPath)}-${nanoid()}`,
    templates: Record<string, string> | undefined,
@@ -133,6 +135,7 @@ export abstract class QueryRunner {
            query,
            additionalPacks,
            extensionPacks,
+            additionalRunQueryArgs,
            generateEvalLog,
            outputDir,
            progress,
--- a/extensions/ql-vscode/src/query-server/run-queries.ts
+++ b/extensions/ql-vscode/src/query-server/run-queries.ts
@@ -27,6 +27,7 @@ export async function compileAndRunQueryAgainstDatabaseCore(
  generateEvalLog: boolean,
  additionalPacks: string[],
  extensionPacks: string[] | undefined,
+  additionalRunQueryArgs: Record<string, any>,
  outputDir: QueryOutputDir,
  progress: ProgressCallback,
  token: CancellationToken,
@@ -55,6 +56,8 @@ export async function compileAndRunQueryAgainstDatabaseCore(
    logPath: evalLogPath,
    target,
    extensionPacks,
+    // Add any additional arguments without interpretation.
+    ...additionalRunQueryArgs,
  };

  // Update the active query logger every time there is a new request to compile.
--- a/extensions/ql-vscode/src/stories/model-editor/MethodRow.stories.tsx
+++ b/extensions/ql-vscode/src/stories/model-editor/MethodRow.stories.tsx
@@ -7,6 +7,9 @@ import { CallClassification, Method } from "../../model-editor/method";
 import { ModeledMethod } from "../../model-editor/modeled-method";
 import { VSCodeDataGrid } from "@vscode/webview-ui-toolkit/react";
 import { GRID_TEMPLATE_COLUMNS } from "../../view/model-editor/ModeledMethodDataGrid";
+import { ModelEditorViewState } from "../../model-editor/shared/view-state";
+import { createMockExtensionPack } from "../../../test/factories/model-editor/extension-pack";
+import { Mode } from "../../model-editor/shared/mode";

 export default {
  title: "CodeQL Model Editor/Method Row",
@@ -66,51 +69,78 @@ const modeledMethod: ModeledMethod = {
  methodParameters: "()",
 };

+const viewState: ModelEditorViewState = {
+  extensionPack: createMockExtensionPack(),
+  showFlowGeneration: true,
+  showLlmButton: true,
+  showMultipleModels: true,
+  mode: Mode.Application,
+};
+
 export const Unmodeled = Template.bind({});
 Unmodeled.args = {
  method,
-  modeledMethod: undefined,
+  modeledMethods: [],
  methodCanBeModeled: true,
+  viewState,
 };

 export const Source = Template.bind({});
 Source.args = {
  method,
-  modeledMethod: { ...modeledMethod, type: "source" },
+  modeledMethods: [{ ...modeledMethod, type: "source" }],
  methodCanBeModeled: true,
+  viewState,
 };

 export const Sink = Template.bind({});
 Sink.args = {
  method,
-  modeledMethod: { ...modeledMethod, type: "sink" },
+  modeledMethods: [{ ...modeledMethod, type: "sink" }],
  methodCanBeModeled: true,
+  viewState,
 };

 export const Summary = Template.bind({});
 Summary.args = {
  method,
-  modeledMethod: { ...modeledMethod, type: "summary" },
+  modeledMethods: [{ ...modeledMethod, type: "summary" }],
  methodCanBeModeled: true,
+  viewState,
 };

 export const Neutral = Template.bind({});
 Neutral.args = {
  method,
-  modeledMethod: { ...modeledMethod, type: "neutral" },
+  modeledMethods: [{ ...modeledMethod, type: "neutral" }],
  methodCanBeModeled: true,
+  viewState,
 };

 export const AlreadyModeled = Template.bind({});
 AlreadyModeled.args = {
  method: { ...method, supported: true },
-  modeledMethod: undefined,
+  modeledMethods: [],
+  viewState,
 };

 export const ModelingInProgress = Template.bind({});
 ModelingInProgress.args = {
  method,
-  modeledMethod,
+  modeledMethods: [modeledMethod],
  modelingInProgress: true,
  methodCanBeModeled: true,
+  viewState,
+};
+
+export const MultipleModelings = Template.bind({});
+MultipleModelings.args = {
+  method,
+  modeledMethods: [
+    { ...modeledMethod, type: "source" },
+    { ...modeledMethod, type: "sink" },
+    { ...modeledMethod },
+  ],
+  methodCanBeModeled: true,
+  viewState,
 };
--- a/extensions/ql-vscode/src/view/method-modeling/MethodModelingView.tsx
+++ b/extensions/ql-vscode/src/view/method-modeling/MethodModelingView.tsx
@@ -30,7 +30,8 @@ export function MethodModelingView({ initialViewState }: Props): JSX.Element {
  const [isMethodModified, setIsMethodModified] = useState<boolean>(false);

  const modelingStatus = useMemo(
-    () => getModelingStatus(modeledMethod, isMethodModified),
+    () =>
+      getModelingStatus(modeledMethod ? [modeledMethod] : [], isMethodModified),
    [modeledMethod, isMethodModified],
  );

--- a/extensions/ql-vscode/src/view/model-editor/LibraryRow.tsx
+++ b/extensions/ql-vscode/src/view/model-editor/LibraryRow.tsx
@@ -81,8 +81,7 @@ export type LibraryRowProps = {
  onSaveModelClick: (methodSignatures: string[]) => void;
  onGenerateFromLlmClick: (
    dependencyName: string,
-    methods: Method[],
-    modeledMethods: Record<string, ModeledMethod>,
+    methodSignatures: string[],
  ) => void;
  onStopGenerateFromLlmClick: (dependencyName: string) => void;
  onGenerateFromSourceClick: () => void;
@@ -126,11 +125,14 @@ export const LibraryRow = ({

  const handleModelWithAI = useCallback(
    async (e: React.MouseEvent) => {
-      onGenerateFromLlmClick(title, methods, modeledMethods);
+      onGenerateFromLlmClick(
+        title,
+        methods.map((m) => m.signature),
+      );
      e.stopPropagation();
      e.preventDefault();
    },
-    [title, methods, modeledMethods, onGenerateFromLlmClick],
+    [title, methods, onGenerateFromLlmClick],
  );

  const handleStopModelWithAI = useCallback(
@@ -232,7 +234,7 @@ export const LibraryRow = ({
            modeledMethods={modeledMethods}
            modifiedSignatures={modifiedSignatures}
            inProgressMethods={inProgressMethods}
-            mode={viewState.mode}
+            viewState={viewState}
            hideModeledMethods={hideModeledMethods}
            revealedMethodSignature={revealedMethodSignature}
            onChange={onChange}
--- a/extensions/ql-vscode/src/view/model-editor/MethodRow.tsx
+++ b/extensions/ql-vscode/src/view/model-editor/MethodRow.tsx
@@ -21,8 +21,16 @@ import { MethodName } from "./MethodName";
 import { ModelTypeDropdown } from "./ModelTypeDropdown";
 import { ModelInputDropdown } from "./ModelInputDropdown";
 import { ModelOutputDropdown } from "./ModelOutputDropdown";
+import { ModelEditorViewState } from "../../model-editor/shared/view-state";

-const ApiOrMethodCell = styled(VSCodeDataGridCell)`
+const MultiModelColumn = styled(VSCodeDataGridCell)`
+  display: flex;
+  flex-direction: column;
+  gap: 0.5em;
+`;
+
+const ApiOrMethodRow = styled.div`
+  min-height: calc(var(--input-height) * 1px);
  display: flex;
  flex-direction: row;
  align-items: center;
@@ -55,10 +63,10 @@ const DataGridRow = styled(VSCodeDataGridRow)<{ focused?: boolean }>`
 export type MethodRowProps = {
  method: Method;
  methodCanBeModeled: boolean;
-  modeledMethod: ModeledMethod | undefined;
+  modeledMethods: ModeledMethod[];
  methodIsUnsaved: boolean;
  modelingInProgress: boolean;
-  mode: Mode;
+  viewState: ModelEditorViewState;
  revealedMethodSignature: string | null;
  onChange: (modeledMethod: ModeledMethod) => void;
 };
@@ -88,19 +96,23 @@ const ModelableMethodRow = forwardRef<HTMLElement | undefined, MethodRowProps>(
  (props, ref) => {
    const {
      method,
-      modeledMethod,
+      modeledMethods: modeledMethodsProp,
      methodIsUnsaved,
-      mode,
+      viewState,
      revealedMethodSignature,
      onChange,
    } = props;

+    const modeledMethods = viewState.showMultipleModels
+      ? modeledMethodsProp
+      : modeledMethodsProp.slice(0, 1);
+
    const jumpToMethod = useCallback(
      () => sendJumpToMethodMessage(method),
      [method],
    );

-    const modelingStatus = getModelingStatus(modeledMethod, methodIsUnsaved);
+    const modelingStatus = getModelingStatus(modeledMethods, methodIsUnsaved);

    return (
      <DataGridRow
@@ -108,18 +120,20 @@ const ModelableMethodRow = forwardRef<HTMLElement | undefined, MethodRowProps>(
        ref={ref}
        focused={revealedMethodSignature === method.signature}
      >
-        <ApiOrMethodCell gridColumn={1}>
-          <ModelingStatusIndicator status={modelingStatus} />
-          <MethodClassifications method={method} />
-          <MethodName {...props.method} />
-          {mode === Mode.Application && (
-            <UsagesButton onClick={jumpToMethod}>
-              {method.usages.length}
-            </UsagesButton>
-          )}
-          <ViewLink onClick={jumpToMethod}>View</ViewLink>
-          {props.modelingInProgress && <ProgressRing />}
-        </ApiOrMethodCell>
+        <VSCodeDataGridCell gridColumn={1}>
+          <ApiOrMethodRow>
+            <ModelingStatusIndicator status={modelingStatus} />
+            <MethodClassifications method={method} />
+            <MethodName {...props.method} />
+            {viewState.mode === Mode.Application && (
+              <UsagesButton onClick={jumpToMethod}>
+                {method.usages.length}
+              </UsagesButton>
+            )}
+            <ViewLink onClick={jumpToMethod}>View</ViewLink>
+            {props.modelingInProgress && <ProgressRing />}
+          </ApiOrMethodRow>
+        </VSCodeDataGridCell>
        {props.modelingInProgress && (
          <>
            <VSCodeDataGridCell gridColumn={2}>
@@ -138,34 +152,46 @@ const ModelableMethodRow = forwardRef<HTMLElement | undefined, MethodRowProps>(
        )}
        {!props.modelingInProgress && (
          <>
-            <VSCodeDataGridCell gridColumn={2}>
-              <ModelTypeDropdown
-                method={method}
-                modeledMethod={modeledMethod}
-                onChange={onChange}
-              />
-            </VSCodeDataGridCell>
-            <VSCodeDataGridCell gridColumn={3}>
-              <ModelInputDropdown
-                method={method}
-                modeledMethod={modeledMethod}
-                onChange={onChange}
-              />
-            </VSCodeDataGridCell>
-            <VSCodeDataGridCell gridColumn={4}>
-              <ModelOutputDropdown
-                method={method}
-                modeledMethod={modeledMethod}
-                onChange={onChange}
-              />
-            </VSCodeDataGridCell>
-            <VSCodeDataGridCell gridColumn={5}>
-              <ModelKindDropdown
-                method={method}
-                modeledMethod={modeledMethod}
-                onChange={onChange}
-              />
-            </VSCodeDataGridCell>
+            <MultiModelColumn gridColumn={2}>
+              {forEachModeledMethod(modeledMethods, (modeledMethod, index) => (
+                <ModelTypeDropdown
+                  key={index}
+                  method={method}
+                  modeledMethod={modeledMethod}
+                  onChange={onChange}
+                />
+              ))}
+            </MultiModelColumn>
+            <MultiModelColumn gridColumn={3}>
+              {forEachModeledMethod(modeledMethods, (modeledMethod, index) => (
+                <ModelInputDropdown
+                  key={index}
+                  method={method}
+                  modeledMethod={modeledMethod}
+                  onChange={onChange}
+                />
+              ))}
+            </MultiModelColumn>
+            <MultiModelColumn gridColumn={4}>
+              {forEachModeledMethod(modeledMethods, (modeledMethod, index) => (
+                <ModelOutputDropdown
+                  key={index}
+                  method={method}
+                  modeledMethod={modeledMethod}
+                  onChange={onChange}
+                />
+              ))}
+            </MultiModelColumn>
+            <MultiModelColumn gridColumn={5}>
+              {forEachModeledMethod(modeledMethods, (modeledMethod, index) => (
+                <ModelKindDropdown
+                  key={index}
+                  method={method}
+                  modeledMethod={modeledMethod}
+                  onChange={onChange}
+                />
+              ))}
+            </MultiModelColumn>
          </>
        )}
      </DataGridRow>
@@ -178,7 +204,7 @@ const UnmodelableMethodRow = forwardRef<
  HTMLElement | undefined,
  MethodRowProps
 >((props, ref) => {
-  const { method, mode, revealedMethodSignature } = props;
+  const { method, viewState, revealedMethodSignature } = props;

  const jumpToMethod = useCallback(
    () => sendJumpToMethodMessage(method),
@@ -191,17 +217,19 @@ const UnmodelableMethodRow = forwardRef<
      ref={ref}
      focused={revealedMethodSignature === method.signature}
    >
-      <ApiOrMethodCell gridColumn={1}>
-        <ModelingStatusIndicator status="saved" />
-        <MethodName {...props.method} />
-        {mode === Mode.Application && (
-          <UsagesButton onClick={jumpToMethod}>
-            {method.usages.length}
-          </UsagesButton>
-        )}
-        <ViewLink onClick={jumpToMethod}>View</ViewLink>
-        <MethodClassifications method={method} />
-      </ApiOrMethodCell>
+      <VSCodeDataGridCell gridColumn={1}>
+        <ApiOrMethodRow>
+          <ModelingStatusIndicator status="saved" />
+          <MethodName {...props.method} />
+          {viewState.mode === Mode.Application && (
+            <UsagesButton onClick={jumpToMethod}>
+              {method.usages.length}
+            </UsagesButton>
+          )}
+          <ViewLink onClick={jumpToMethod}>View</ViewLink>
+          <MethodClassifications method={method} />
+        </ApiOrMethodRow>
+      </VSCodeDataGridCell>
      <VSCodeDataGridCell gridColumn="span 4">
        Method already modeled
      </VSCodeDataGridCell>
@@ -216,3 +244,17 @@ function sendJumpToMethodMessage(method: Method) {
    methodSignature: method.signature,
  });
 }
+
+function forEachModeledMethod(
+  modeledMethods: ModeledMethod[],
+  renderer: (
+    modeledMethod: ModeledMethod | undefined,
+    index: number,
+  ) => JSX.Element,
+): JSX.Element | JSX.Element[] {
+  if (modeledMethods.length === 0) {
+    return renderer(undefined, 0);
+  } else {
+    return modeledMethods.map(renderer);
+  }
+}
--- a/extensions/ql-vscode/src/view/model-editor/ModelEditor.tsx
+++ b/extensions/ql-vscode/src/view/model-editor/ModelEditor.tsx
@@ -219,16 +219,11 @@ export function ModelEditor({
  }, []);

  const onGenerateFromLlmClick = useCallback(
-    (
-      packageName: string,
-      methods: Method[],
-      modeledMethods: Record<string, ModeledMethod>,
-    ) => {
+    (packageName: string, methodSignatures: string[]) => {
      vscode.postMessage({
        t: "generateMethodsFromLlm",
        packageName,
-        methods,
-        modeledMethods,
+        methodSignatures,
      });
    },
    [],
--- a/extensions/ql-vscode/src/view/model-editor/ModeledMethodDataGrid.tsx
+++ b/extensions/ql-vscode/src/view/model-editor/ModeledMethodDataGrid.tsx
@@ -8,10 +8,10 @@ import { MethodRow } from "./MethodRow";
 import { Method } from "../../model-editor/method";
 import { ModeledMethod } from "../../model-editor/modeled-method";
 import { useMemo } from "react";
-import { Mode } from "../../model-editor/shared/mode";
 import { sortMethods } from "../../model-editor/shared/sorting";
 import { InProgressMethods } from "../../model-editor/shared/in-progress-methods";
 import { HiddenMethodsRow } from "./HiddenMethodsRow";
+import { ModelEditorViewState } from "../../model-editor/shared/view-state";

 export const GRID_TEMPLATE_COLUMNS = "0.5fr 0.125fr 0.125fr 0.125fr 0.125fr";

@@ -21,7 +21,7 @@ export type ModeledMethodDataGridProps = {
  modeledMethods: Record<string, ModeledMethod>;
  modifiedSignatures: Set<string>;
  inProgressMethods: InProgressMethods;
-  mode: Mode;
+  viewState: ModelEditorViewState;
  hideModeledMethods: boolean;
  revealedMethodSignature: string | null;
  onChange: (modeledMethod: ModeledMethod) => void;
@@ -33,7 +33,7 @@ export const ModeledMethodDataGrid = ({
  modeledMethods,
  modifiedSignatures,
  inProgressMethods,
-  mode,
+  viewState,
  hideModeledMethods,
  revealedMethodSignature,
  onChange,
@@ -84,22 +84,25 @@ export const ModeledMethodDataGrid = ({
              Kind
            </VSCodeDataGridCell>
          </VSCodeDataGridRow>
-          {methodsWithModelability.map(({ method, methodCanBeModeled }) => (
-            <MethodRow
-              key={method.signature}
-              method={method}
-              methodCanBeModeled={methodCanBeModeled}
-              modeledMethod={modeledMethods[method.signature]}
-              methodIsUnsaved={modifiedSignatures.has(method.signature)}
-              modelingInProgress={inProgressMethods.hasMethod(
-                packageName,
-                method.signature,
-              )}
-              mode={mode}
-              revealedMethodSignature={revealedMethodSignature}
-              onChange={onChange}
-            />
-          ))}
+          {methodsWithModelability.map(({ method, methodCanBeModeled }) => {
+            const modeledMethod = modeledMethods[method.signature];
+            return (
+              <MethodRow
+                key={method.signature}
+                method={method}
+                methodCanBeModeled={methodCanBeModeled}
+                modeledMethods={modeledMethod ? [modeledMethod] : []}
+                methodIsUnsaved={modifiedSignatures.has(method.signature)}
+                modelingInProgress={inProgressMethods.hasMethod(
+                  packageName,
+                  method.signature,
+                )}
+                viewState={viewState}
+                revealedMethodSignature={revealedMethodSignature}
+                onChange={onChange}
+              />
+            );
+          })}
        </>
      )}
      <HiddenMethodsRow
--- a/extensions/ql-vscode/src/view/model-editor/ModeledMethodsList.tsx
+++ b/extensions/ql-vscode/src/view/model-editor/ModeledMethodsList.tsx
@@ -23,8 +23,7 @@ export type ModeledMethodsListProps = {
  onSaveModelClick: (methodSignatures: string[]) => void;
  onGenerateFromLlmClick: (
    packageName: string,
-    methods: Method[],
-    modeledMethods: Record<string, ModeledMethod>,
+    methodSignatures: string[],
  ) => void;
  onStopGenerateFromLlmClick: (packageName: string) => void;
  onGenerateFromSourceClick: () => void;
--- a/extensions/ql-vscode/src/view/model-editor/tests/MethodRow.spec.tsx
+++ b/extensions/ql-vscode/src/view/model-editor/tests/MethodRow.spec.tsx
@@ -9,6 +9,8 @@ import { Mode } from "../../../model-editor/shared/mode";
 import { MethodRow, MethodRowProps } from "../MethodRow";
 import { ModeledMethod } from "../../../model-editor/modeled-method";
 import userEvent from "@testing-library/user-event";
+import { ModelEditorViewState } from "../../../model-editor/shared/view-state";
+import { createMockExtensionPack } from "../../../../test/factories/model-editor/extension-pack";

 describe(MethodRow.name, () => {
  const method = createMethod({
@@ -31,16 +33,24 @@ describe(MethodRow.name, () => {
  };
  const onChange = jest.fn();

+  const viewState: ModelEditorViewState = {
+    mode: Mode.Application,
+    showFlowGeneration: false,
+    showLlmButton: false,
+    showMultipleModels: false,
+    extensionPack: createMockExtensionPack(),
+  };
+
  const render = (props: Partial<MethodRowProps> = {}) =>
    reactRender(
      <MethodRow
        method={method}
        methodCanBeModeled={true}
-        modeledMethod={modeledMethod}
+        modeledMethods={[modeledMethod]}
        methodIsUnsaved={false}
        modelingInProgress={false}
        revealedMethodSignature={null}
-        mode={Mode.Application}
+        viewState={viewState}
        onChange={onChange}
        {...props}
      />,
@@ -54,6 +64,14 @@ describe(MethodRow.name, () => {
    expect(screen.queryByLabelText("Loading")).not.toBeInTheDocument();
  });

+  it("renders when there is no modeled method", () => {
+    render({ modeledMethods: [] });
+
+    expect(screen.queryAllByRole("combobox")).toHaveLength(4);
+    expect(screen.getByLabelText("Method not modeled")).toBeInTheDocument();
+    expect(screen.queryByLabelText("Loading")).not.toBeInTheDocument();
+  });
+
  it("can change the kind", async () => {
    render();

@@ -110,7 +128,7 @@ describe(MethodRow.name, () => {

  it("shows the modeling status indicator when unmodeled", () => {
    render({
-      modeledMethod: undefined,
+      modeledMethods: [],
    });

    expect(screen.getByLabelText("Method not modeled")).toBeInTheDocument();
@@ -124,10 +142,48 @@ describe(MethodRow.name, () => {
    expect(screen.getByLabelText("Loading")).toBeInTheDocument();
  });

+  it("can render multiple models", () => {
+    render({
+      modeledMethods: [
+        { ...modeledMethod, type: "source" },
+        { ...modeledMethod, type: "sink" },
+        { ...modeledMethod, type: "summary" },
+      ],
+      viewState: {
+        ...viewState,
+        showMultipleModels: true,
+      },
+    });
+
+    const kindInputs = screen.getAllByRole("combobox", { name: "Model type" });
+    expect(kindInputs).toHaveLength(3);
+    expect(kindInputs[0]).toHaveValue("source");
+    expect(kindInputs[1]).toHaveValue("sink");
+    expect(kindInputs[2]).toHaveValue("summary");
+  });
+
+  it("renders only first model when showMultipleModels feature flag is disabled", () => {
+    render({
+      modeledMethods: [
+        { ...modeledMethod, type: "source" },
+        { ...modeledMethod, type: "sink" },
+        { ...modeledMethod, type: "summary" },
+      ],
+      viewState: {
+        ...viewState,
+        showMultipleModels: false,
+      },
+    });
+
+    const kindInputs = screen.getAllByRole("combobox", { name: "Model type" });
+    expect(kindInputs.length).toBe(1);
+    expect(kindInputs[0]).toHaveValue("source");
+  });
+
  it("renders an unmodelable method", () => {
    render({
      methodCanBeModeled: false,
-      modeledMethod: undefined,
+      modeledMethods: [],
    });

    expect(screen.queryByRole("combobox")).not.toBeInTheDocument();
--- a/extensions/ql-vscode/src/view/model-editor/tests/ModeledMethodDataGrid.spec.tsx
+++ b/extensions/ql-vscode/src/view/model-editor/tests/ModeledMethodDataGrid.spec.tsx
@@ -7,6 +7,8 @@ import {
  ModeledMethodDataGrid,
  ModeledMethodDataGridProps,
 } from "../ModeledMethodDataGrid";
+import { ModelEditorViewState } from "../../../model-editor/shared/view-state";
+import { createMockExtensionPack } from "../../../../test/factories/model-editor/extension-pack";

 describe(ModeledMethodDataGrid.name, () => {
  const method1 = createMethod({
@@ -41,6 +43,14 @@ describe(ModeledMethodDataGrid.name, () => {
  });
  const onChange = jest.fn();

+  const viewState: ModelEditorViewState = {
+    mode: Mode.Application,
+    showFlowGeneration: false,
+    showLlmButton: false,
+    showMultipleModels: false,
+    extensionPack: createMockExtensionPack(),
+  };
+
  const render = (props: Partial<ModeledMethodDataGridProps> = {}) =>
    reactRender(
      <ModeledMethodDataGrid
@@ -58,7 +68,7 @@ describe(ModeledMethodDataGrid.name, () => {
        }}
        modifiedSignatures={new Set([method1.signature])}
        inProgressMethods={new InProgressMethods()}
-        mode={Mode.Application}
+        viewState={viewState}
        hideModeledMethods={false}
        revealedMethodSignature={null}
        onChange={onChange}
--- a/extensions/ql-vscode/syntaxes/ql.tmLanguage.yml
+++ b/extensions/ql-vscode/syntaxes/ql.tmLanguage.yml
@@ -170,6 +170,14 @@ repository:
    match: '\]'
    name: punctuation.squarebracket.close.ql

+  open-angle:
+    match: '<'
+    name: punctuation.anglebracket.open.ql
+
+  close-angle:
+    match: '>'
+    name: punctuation.anglebracket.close.ql
+
  operator-or-punctuation:
    patterns:
    - include: '#relational-operator'
@@ -186,6 +194,8 @@ repository:
    - include: '#close-brace'
    - include: '#open-bracket'
    - include: '#close-bracket'
+    - include: '#open-angle'
+    - include: '#close-angle'

 # Keywords
  dont-care:
@@ -651,18 +661,36 @@ repository:
    - include: '#non-context-sensitive'
    - include: '#annotation'

+  # The argument list of an instantiation, enclosed in angle brackets.
+  instantiation-args:
+    beginPattern: '#open-angle'
+    endPattern: '#close-angle'
+    name: meta.type.parameters.ql
+    patterns:
+    # Include `#instantiation-args` first so that `#open-angle` and `#close-angle` take precedence
+    # over `#relational-operator`.
+    - include: '#instantiation-args'
+    - include: '#non-context-sensitive'
+    - match: '(?#simple-id)'
+      name: entity.name.type.namespace.ql
+
  # An `import` directive. Note that we parse the optional `as` clause as a separate top-level
  # directive, because otherwise it's too hard to figure out where the `import` directive ends.
  import-directive:
    beginPattern: '#import'
-    # Ends with a simple-id that is not followed by a `.` or a `::`. This does not handle comments or
-    # line breaks between the simple-id and the `.` or `::`. 
-    end: '(?#simple-id) (?!\s*(\.|\:\:))'
-    endCaptures:
-      '0':
-        name: entity.name.type.namespace.ql
+    # TextMate makes it tricky to tell whether an identifier that we encounter is part of the
+    # `import` directive or whether it's the first token of the next module-level declaration.
+    # To find the end of the import directive, we'll look for a zero-width match where the previous
+    # token is either an identifier (other than `import`) or a `>`, and the next token is not a `.`,
+    # `<`, `,`, or `::`. This works for nearly all real-world `import` directives, but it will end the
+    # `import` directive too early if there is a comment or line break between two components of the
+    # module expression.
+    end: '(?<!\bimport)(?<=(?:\>)|[A-Za-z0-9_]) (?!\s*(\.|\:\:|\,|(?#open-angle)))'
    name: meta.block.import-directive.ql
    patterns:
+    # Include `#instantiation-args` first so that `#open-angle` and `#close-angle` take precedence
+    # over `#relational-operator`.
+    - include: '#instantiation-args'
    - include: '#non-context-sensitive'
    - match: '(?#simple-id)'
      name: entity.name.type.namespace.ql
@@ -703,7 +731,6 @@ repository:
    - match: '(?#simple-id)|(?#at-lower-id)'
      name: entity.name.type.ql

-
  # A `module` declaration, whether a module definition or an alias declaration.
  module-declaration:
    # Starts with the `module` keyword.
--- a/extensions/ql-vscode/test/unit-tests/model-editor/auto-model.test.ts
+++ b/extensions/ql-vscode/test/unit-tests/model-editor/auto-model.test.ts
@@ -99,19 +99,21 @@ describe("getCandidates", () => {
        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 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, methods, modeledMethods);
    expect(candidates.length).toEqual(0);
--- a/extensions/ql-vscode/test/vscode-tests/cli-integration/debugger/debugger.test.ts
+++ b/extensions/ql-vscode/test/vscode-tests/cli-integration/debugger/debugger.test.ts
@@ -177,4 +177,31 @@ describeWithCodeQL()("Debugger", () => {
      expect(editor.document.isDirty).toBe(false);
    });
  });
+
+  it("should pass additionalArgs through to query server", async () => {
+    if (!(await cli.cliConstraints.supportsNewQueryServerForTests())) {
+      // Only works with the new query server.
+      return;
+    }
+    await withDebugController(appCommands, async (controller) => {
+      await controller.startDebugging(
+        {
+          query: quickEvalQueryPath,
+          additionalRunQueryArgs: {
+            // Overrides the value passed to the query server
+            queryPath: simpleQueryPath,
+          },
+        },
+        true,
+      );
+      await controller.expectLaunched();
+      const result = await controller.expectSucceeded();
+      await controller.expectExited();
+      await controller.expectTerminated();
+      await controller.expectSessionClosed();
+
+      // Expect the number of results to be the same as if we had run the simple query, not the quick eval query.
+      expect(await getResultCount(result.results.outputDir, cli)).toBe(2);
+    });
+  });
 });
--- a/extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/auto-model-codeml-queries.test.ts
+++ b/extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/auto-model-codeml-queries.test.ts
@@ -159,6 +159,7 @@ describe("runAutoModelQueries", () => {
      false,
      expect.arrayContaining([expect.stringContaining("tmp")]),
      ["/a/b/c/my-extension-pack"],
+      {},
      "/tmp/queries",
      undefined,
      undefined,
--- a/extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/external-api-usage-query.test.ts
+++ b/extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/external-api-usage-query.test.ts
@@ -165,6 +165,7 @@ describe("external api usage query", () => {
        false,
        [],
        ["my/extensions"],
+        {},
        "/tmp/queries",
        undefined,
        undefined,
--- a/extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/methods-usage/methods-usage-data-provider.test.ts
+++ b/extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/methods-usage/methods-usage-data-provider.test.ts
@@ -20,7 +20,7 @@ describe("MethodsUsageDataProvider", () => {
  describe("setState", () => {
    const hideModeledMethods = false;
    const methods: Method[] = [];
-    const modeledMethods: Record<string, ModeledMethod> = {};
+    const modeledMethods: Record<string, ModeledMethod[]> = {};
    const modifiedMethodSignatures: Set<string> = new Set();
    const dbItem = mockedObject<DatabaseItem>({
      getSourceLocationPrefix: () => "test",
@@ -125,7 +125,7 @@ describe("MethodsUsageDataProvider", () => {
    });

    it("should emit onDidChangeTreeData event when modeled methods has changed", async () => {
-      const modeledMethods2: Record<string, ModeledMethod> = {};
+      const modeledMethods2: Record<string, ModeledMethod[]> = {};

      await dataProvider.setState(
        methods,
@@ -213,7 +213,7 @@ describe("MethodsUsageDataProvider", () => {
    });

    const methods: Method[] = [supportedMethod, unsupportedMethod];
-    const modeledMethods: Record<string, ModeledMethod> = {};
+    const modeledMethods: Record<string, ModeledMethod[]> = {};
    const modifiedMethodSignatures: Set<string> = new Set();

    const dbItem = mockedObject<DatabaseItem>({
--- a/extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/methods-usage/methods-usage-panel.test.ts
+++ b/extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/methods-usage/methods-usage-panel.test.ts
@@ -21,7 +21,7 @@ describe("MethodsUsagePanel", () => {
  describe("setState", () => {
    const hideModeledMethods = false;
    const methods: Method[] = [createMethod()];
-    const modeledMethods: Record<string, ModeledMethod> = {};
+    const modeledMethods: Record<string, ModeledMethod[]> = {};
    const modifiedMethodSignatures: Set<string> = new Set();

    it("should update the tree view with the correct batch number", async () => {
@@ -50,7 +50,7 @@ describe("MethodsUsagePanel", () => {
    let modelingStore: ModelingStore;

    const hideModeledMethods: boolean = false;
-    const modeledMethods: Record<string, ModeledMethod> = {};
+    const modeledMethods: Record<string, ModeledMethod[]> = {};
    const modifiedMethodSignatures: Set<string> = new Set();
    const usage = createUsage();

--- a/syntaxes/ql.tmLanguage.json
+++ b/syntaxes/ql.tmLanguage.json
@@ -108,6 +108,14 @@
      "match": "(?x)\\]",
      "name": "punctuation.squarebracket.close.ql"
    },
+    "open-angle": {
+      "match": "(?x)<",
+      "name": "punctuation.anglebracket.open.ql"
+    },
+    "close-angle": {
+      "match": "(?x)>",
+      "name": "punctuation.anglebracket.close.ql"
+    },
    "operator-or-punctuation": {
      "patterns": [
        {
@@ -151,6 +159,12 @@
        },
        {
          "include": "#close-bracket"
+        },
+        {
+          "include": "#open-angle"
+        },
+        {
+          "include": "#close-angle"
        }
      ]
    },
@@ -723,15 +737,48 @@
        }
      ]
    },
-    "import-directive": {
-      "end": "(?x)(?:\\b [A-Za-z][0-9A-Za-z_]* (?:(?!(?:[0-9A-Za-z_])))) (?!\\s*(\\.|\\:\\:))",
-      "endCaptures": {
-        "0": {
+    "instantiation-args": {
+      "name": "meta.type.parameters.ql",
+      "patterns": [
+        {
+          "include": "#instantiation-args"
+        },
+        {
+          "include": "#non-context-sensitive"
+        },
+        {
+          "match": "(?x)(?:\\b [A-Za-z][0-9A-Za-z_]* (?:(?!(?:[0-9A-Za-z_]))))",
          "name": "entity.name.type.namespace.ql"
        }
+      ],
+      "begin": "(?x)((?:<))",
+      "beginCaptures": {
+        "1": {
+          "patterns": [
+            {
+              "include": "#open-angle"
+            }
+          ]
+        }
      },
+      "end": "(?x)((?:>))",
+      "endCaptures": {
+        "1": {
+          "patterns": [
+            {
+              "include": "#close-angle"
+            }
+          ]
+        }
+      }
+    },
+    "import-directive": {
+      "end": "(?x)(?<!\\bimport)(?<=(?:\\>)|[A-Za-z0-9_]) (?!\\s*(\\.|\\:\\:|\\,|(?:<)))",
      "name": "meta.block.import-directive.ql",
      "patterns": [
+        {
+          "include": "#instantiation-args"
+        },
        {
          "include": "#non-context-sensitive"
        },