Merge remote-tracking branch 'origin/main' into koesie10/convert-remaining-multiple-models

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

@@ -556,8 +556,7 @@ interface GenerateMethodMessage {
 interface GenerateMethodsFromLlmMessage {
   t: "generateMethodsFromLlm";
   packageName: string;
-  methods: Method[];
-  modeledMethods: Record<string, ModeledMethod>;
+  methodSignatures: string[];
 }
 
 interface StopGeneratingMethodsFromLlmMessage {
@@ -634,7 +633,7 @@ interface SetMethodModelingPanelViewStateMessage {
 
 interface SetMethodMessage {
   t: "setMethod";
-  method: Method;
+  method: Method | undefined;
 }
 
 interface SetMethodModifiedMessage {

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

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

@@ -13,6 +13,7 @@ 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,
@@ -25,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,
@@ -50,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({
@@ -74,6 +80,11 @@ export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
           isModified: selectedMethod.isModified,
         });
       }
+
+      await this.postMessage({
+        t: "setInModelingMode",
+        inModelingMode: true,
+      });
     }
   }
 
@@ -107,6 +118,10 @@ export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
           msg.method.signature,
           convertFromLegacyModeledMethod(msg.method),
         );
+        this.modelingStore.addModifiedMethod(
+          activeState.databaseItem,
+          msg.method.signature,
+        );
         break;
       }
       case "revealInModelEditor":
@@ -180,6 +195,8 @@ 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,
@@ -200,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);
+        }
       }),
     );
   }

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

@@ -105,7 +105,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);
   }
 

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

@@ -45,7 +45,6 @@ import { ModelingStore } from "./modeling-store";
 import { ModelEditorViewTracker } from "./model-editor-view-tracker";
 import {
   convertFromLegacyModeledMethod,
-  convertFromLegacyModeledMethods,
   convertToLegacyModeledMethods,
 } from "./modeled-methods-legacy";
 
@@ -268,8 +267,7 @@ export class ModelEditorView extends AbstractWebview<
       case "generateMethodsFromLlm":
         await this.generateModeledMethodsFromLlm(
           msg.packageName,
-          msg.methods,
-          convertFromLegacyModeledMethods(msg.modeledMethods),
+          msg.methodSignatures,
         );
         void telemetryListener?.sendUIInteraction(
           "model-editor-generate-methods-from-llm",
@@ -481,9 +479,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,

extensions/ql-vscode/src/model-editor/modeled-methods-legacy.ts

@@ -1,27 +1,5 @@
 import { ModeledMethod } from "./modeled-method";
 
-/**
- * Converts a record of ModeledMethod[] indexed by signature to a record of a single 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 `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 modeledMethods The record of ModeledMethod[] indexed by signature
- */
-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, convertFromLegacyModeledMethod(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

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(

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

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;