Convert remaining extension host code to handle multiple models

This converts all remaining extension host code to handle multiple
models per method. The only place where we're using the legacy format
is in the webview and in the boundary between the webview and the
extension host.

Author: koesie10

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;
     }
 

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,
+          },
+        ];
       }
     }
 

extensions/ql-vscode/src/model-editor/method-modeling/method-modeling-view-provider.ts

@@ -13,6 +13,10 @@ import { AbstractWebviewViewProvider } from "../../common/vscode/abstract-webvie
 import { assertNever } from "../../common/helpers-pure";
 import { ModelEditorViewTracker } from "../model-editor-view-tracker";
 import { showMultipleModels } from "../../config";
+import {
+  convertFromLegacyModeledMethod,
+  convertToLegacyModeledMethod,
+} from "../modeled-methods-legacy";
 
 export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
   ToMethodModelingMessage,
@@ -62,7 +66,9 @@ export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
         await this.postMessage({
           t: "setSelectedMethod",
           method: selectedMethod.method,
-          modeledMethod: selectedMethod.modeledMethod,
+          modeledMethod: convertToLegacyModeledMethod(
+            selectedMethod.modeledMethods,
+          ),
           isModified: selectedMethod.isModified,
         });
       }
@@ -94,9 +100,10 @@ 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),
         );
         break;
       }
@@ -141,11 +148,11 @@ 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) {
+          const modeledMethods = e.modeledMethods[this.method?.signature ?? ""];
+          if (modeledMethods) {
             await this.postMessage({
               t: "setModeledMethod",
-              method: modeledMethod,
+              method: convertToLegacyModeledMethod(modeledMethods),
             });
           }
         }
@@ -171,7 +178,7 @@ export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
           await this.postMessage({
             t: "setSelectedMethod",
             method: e.method,
-            modeledMethod: e.modeledMethod,
+            modeledMethod: convertToLegacyModeledMethod(e.modeledMethods),
             isModified: e.isModified,
           });
         }

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

@@ -14,7 +14,7 @@ import { DatabaseItem } from "../../databases/local-databases";
 import { relative } from "path";
 import { CodeQLCliServer } from "../../codeql-cli/cli";
 import { INITIAL_HIDE_MODELED_METHODS_VALUE } from "../shared/hide-modeled-methods";
-import { getModelingStatus } from "../shared/modeling-status";
+import { getModelingStatusForModeledMethods } from "../shared/modeling-status";
 import { assertNever } from "../../common/helpers-pure";
 import { ModeledMethod } from "../modeled-method";
 
@@ -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 (
@@ -99,10 +99,13 @@ 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 = getModelingStatusForModeledMethods(
+      modeledMethods,
+      modifiedMethod,
+    );
     switch (status) {
       case "unmodeled":
         return new ThemeIcon("error", new ThemeColor("errorForeground"));

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(

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

@@ -48,6 +48,7 @@ import { telemetryListener } from "../common/vscode/telemetry";
 import { ModelingStore } from "./modeling-store";
 import { ModelEditorViewTracker } from "./model-editor-view-tracker";
 import {
+  convertFromLegacyModeledMethod,
   convertFromLegacyModeledMethods,
   convertToLegacyModeledMethods,
 } from "./modeled-methods-legacy";
@@ -259,7 +260,7 @@ export class ModelEditorView extends AbstractWebview<
         await this.generateModeledMethodsFromLlm(
           msg.packageName,
           msg.methods,
-          msg.modeledMethods,
+          convertFromLegacyModeledMethods(msg.modeledMethods),
         );
         void telemetryListener?.sendUIInteraction(
           "model-editor-generate-methods-from-llm",
@@ -303,7 +304,10 @@ export class ModelEditorView extends AbstractWebview<
         );
         break;
       case "setModeledMethod": {
-        this.setModeledMethod(msg.method);
+        this.setModeledMethods(
+          msg.method.signature,
+          convertFromLegacyModeledMethod(msg.method),
+        );
         break;
       }
       default:
@@ -363,10 +367,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,
@@ -438,10 +439,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);
@@ -466,7 +473,7 @@ export class ModelEditorView extends AbstractWebview<
   private async generateModeledMethodsFromLlm(
     packageName: string,
     methods: Method[],
-    modeledMethods: Record<string, ModeledMethod>,
+    modeledMethods: Record<string, ModeledMethod[]>,
   ): Promise<void> {
     await this.autoModeler.startModeling(
       packageName,
@@ -603,7 +610,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),
           });
         }
       }),
@@ -621,7 +628,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(
@@ -630,13 +637,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);
   }
 }