Merge branch 'main' into robertbrignull/remove_selected_model

Author: robertbrignull

extensions/ql-vscode/scripts/source-map.ts

@@ -115,21 +115,35 @@ async function extractSourceMap() {
   }
 
   if (stacktrace.includes("at")) {
-    const rawSourceMaps = new Map<string, RawSourceMap>();
+    const rawSourceMaps = new Map<string, RawSourceMap | null>();
 
     const mappedStacktrace = await replaceAsync(
       stacktrace,
       stackLineRegex,
       async (match, name, file, line, column) => {
         if (!rawSourceMaps.has(file)) {
-          const rawSourceMap: RawSourceMap = await readJSON(
-            resolve(sourceMapsDirectory, `${basename(file)}.map`),
-          );
-          rawSourceMaps.set(file, rawSourceMap);
+          try {
+            const rawSourceMap: RawSourceMap = await readJSON(
+              resolve(sourceMapsDirectory, `${basename(file)}.map`),
+            );
+            rawSourceMaps.set(file, rawSourceMap);
+          } catch (e: unknown) {
+            // If the file is not found, we will not decode it and not try reading this source map again
+            if (e instanceof Error && "code" in e && e.code === "ENOENT") {
+              rawSourceMaps.set(file, null);
+            } else {
+              throw e;
+            }
+          }
+        }
+
+        const sourceMap = rawSourceMaps.get(file);
+        if (!sourceMap) {
+          return match;
         }
 
         const originalPosition = await SourceMapConsumer.with(
-          rawSourceMaps.get(file) as RawSourceMap,
+          sourceMap,
           null,
           async function (consumer) {
             return consumer.originalPositionFor({

extensions/ql-vscode/src/config.ts

@@ -702,6 +702,10 @@ export function showQueriesPanel(): boolean {
 const MODEL_SETTING = new Setting("model", ROOT_SETTING);
 const FLOW_GENERATION = new Setting("flowGeneration", MODEL_SETTING);
 const LLM_GENERATION = new Setting("llmGeneration", MODEL_SETTING);
+const LLM_GENERATION_BATCH_SIZE = new Setting(
+  "llmGenerationBatchSize",
+  MODEL_SETTING,
+);
 const EXTENSIONS_DIRECTORY = new Setting("extensionsDirectory", MODEL_SETTING);
 const SHOW_MULTIPLE_MODELS = new Setting("showMultipleModels", MODEL_SETTING);
 
@@ -725,6 +729,14 @@ export class ModelConfigListener extends ConfigListener implements ModelConfig {
     return !!LLM_GENERATION.getValue<boolean>();
   }
 
+  /**
+   * Limits the number of candidates we send to the model in each request to avoid long requests.
+   * Note that the model may return fewer than this number of candidates.
+   */
+  public get llmGenerationBatchSize(): number {
+    return LLM_GENERATION_BATCH_SIZE.getValue<number | null>() || 10;
+  }
+
   public getExtensionsDirectory(languageId: string): string | undefined {
     return EXTENSIONS_DIRECTORY.getValue<string>({
       languageId,

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

@@ -17,11 +17,7 @@ import { DatabaseItem } from "../databases/local-databases";
 import { Mode } from "./shared/mode";
 import { CancellationTokenSource } from "vscode";
 import { ModelingStore } from "./modeling-store";
-
-// Limit the number of candidates we send to the model in each request
-// to avoid long requests.
-// Note that the model may return fewer than this number of candidates.
-const candidateBatchSize = 20;
+import { ModelConfigListener } from "../config";
 
 /**
  * The auto-modeler holds state around auto-modeling jobs and allows
@@ -36,6 +32,7 @@ export class AutoModeler {
     private readonly app: App,
     private readonly cliServer: CodeQLCliServer,
     private readonly queryRunner: QueryRunner,
+    private readonly modelConfig: ModelConfigListener,
     private readonly modelingStore: ModelingStore,
     private readonly queryStorageDir: string,
     private readonly databaseItem: DatabaseItem,
@@ -109,6 +106,9 @@ export class AutoModeler {
     cancellationTokenSource: CancellationTokenSource,
   ): Promise<void> {
     void extLogger.log(`Modeling package ${packageName}`);
+
+    const candidateBatchSize = this.modelConfig.llmGenerationBatchSize;
+
     await withProgress(async (progress) => {
       // Fetch the candidates to send to the model
       const allCandidateMethods = getCandidates(mode, methods, modeledMethods);

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

@@ -89,26 +89,26 @@ export class MethodsUsageDataProvider
 
   getTreeItem(item: MethodsUsageTreeViewItem): TreeItem {
     if (isMethodTreeViewItem(item)) {
+      const { method } = item;
+
       return {
-        label: `${item.packageName}.${item.typeName}.${item.methodName}${item.methodParameters}`,
+        label: `${method.packageName}.${method.typeName}.${method.methodName}${method.methodParameters}`,
         collapsibleState: TreeItemCollapsibleState.Collapsed,
-        iconPath: this.getModelingStatusIcon(item),
+        iconPath: this.getModelingStatusIcon(method),
       };
     } else {
-      const method = this.getParent(item);
-      if (!method || !isMethodTreeViewItem(method)) {
-        throw new Error("Parent not found for tree item");
-      }
+      const { method, usage } = item;
+
       return {
-        label: item.label,
-        description: `${this.relativePathWithinDatabase(item.url.uri)} [${
-          item.url.startLine
-        }, ${item.url.endLine}]`,
+        label: usage.label,
+        description: `${this.relativePathWithinDatabase(usage.url.uri)} [${
+          usage.url.startLine
+        }, ${usage.url.endLine}]`,
         collapsibleState: TreeItemCollapsibleState.None,
         command: {
           title: "Show usage",
           command: "codeQLModelEditor.jumpToMethod",
-          arguments: [method, item, this.databaseItem],
+          arguments: [method, usage, this.databaseItem],
         },
       };
     }
@@ -146,7 +146,7 @@ export class MethodsUsageDataProvider
   getChildren(item?: MethodsUsageTreeViewItem): MethodsUsageTreeViewItem[] {
     if (item === undefined) {
       if (this.hideModeledMethods) {
-        return this.sortedTreeItems.filter((api) => !api.supported);
+        return this.sortedTreeItems.filter((api) => !api.method.supported);
       } else {
         return [...this.sortedTreeItems];
       }
@@ -172,21 +172,24 @@ export class MethodsUsageDataProvider
     usage: Usage,
   ): UsageTreeViewItem | undefined {
     const method = this.sortedTreeItems.find(
-      (m) => m.signature === methodSignature,
+      (m) => m.method.signature === methodSignature,
     );
     if (!method) {
       return undefined;
     }
 
-    return method.children.find((u) => usagesAreEqual(u, usage));
+    return method.children.find((u) => usagesAreEqual(u.usage, usage));
   }
 }
 
-type MethodTreeViewItem = Method & {
+type MethodTreeViewItem = {
+  method: Method;
   children: UsageTreeViewItem[];
 };
 
-type UsageTreeViewItem = Usage & {
+type UsageTreeViewItem = {
+  method: Method;
+  usage: Usage;
   parent: MethodTreeViewItem;
 };
 
@@ -195,7 +198,7 @@ export type MethodsUsageTreeViewItem = MethodTreeViewItem | UsageTreeViewItem;
 function isMethodTreeViewItem(
   item: MethodsUsageTreeViewItem,
 ): item is MethodTreeViewItem {
-  return "children" in item && "usages" in item;
+  return "children" in item && "method" in item;
 }
 
 function usagesAreEqual(u1: Usage, u2: Usage): boolean {
@@ -225,12 +228,13 @@ function sortMethodsInGroups(methods: readonly Method[], mode: Mode): Method[] {
 function createTreeItems(methods: readonly Method[]): MethodTreeViewItem[] {
   return methods.map((method) => {
     const newMethod: MethodTreeViewItem = {
-      ...method,
+      method,
       children: [],
     };
 
     newMethod.children = method.usages.map((usage) => ({
-      ...usage,
+      method,
+      usage,
       // This needs to be a reference to the parent method, not a copy of it.
       parent: newMethod,
     }));

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

@@ -76,6 +76,7 @@ export class ModelEditorView extends AbstractWebview<
       app,
       cliServer,
       queryRunner,
+      this.modelConfig,
       modelingStore,
       queryStorageDir,
       databaseItem,

extensions/ql-vscode/src/stories/model-editor/InProgressDropdown.stories.tsx

@@ -0,0 +1,16 @@
+import * as React from "react";
+
+import { Meta, StoryFn } from "@storybook/react";
+
+import { InProgressDropdown as InProgressDropdownComponent } from "../../view/model-editor/InProgressDropdown";
+
+export default {
+  title: "CodeQL Model Editor/In Progress Dropdown",
+  component: InProgressDropdownComponent,
+} as Meta<typeof InProgressDropdownComponent>;
+
+const Template: StoryFn<typeof InProgressDropdownComponent> = (args) => (
+  <InProgressDropdownComponent />
+);
+
+export const InProgressDropdown = Template.bind({});

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

@@ -19,6 +19,7 @@ type Props = {
   value: string | undefined;
   options: Array<{ value: string; label: string }>;
   disabled?: boolean;
+  className?: string;
   disabledPlaceholder?: string;
   onChange?: (event: ChangeEvent<HTMLSelectElement>) => void;
 
@@ -40,6 +41,7 @@ export function Dropdown({
   options,
   disabled,
   disabledPlaceholder,
+  className,
   onChange,
   ...props
 }: Props) {
@@ -49,6 +51,7 @@ export function Dropdown({
       value={disabled ? disabledValue : value}
       disabled={disabled}
       onChange={onChange}
+      className={className}
       {...props}
     >
       {disabled ? (

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

@@ -1,9 +1,14 @@
 import * as React from "react";
+import { styled } from "styled-components";
 import { Dropdown } from "../common/Dropdown";
 
+const StyledDropdown = styled(Dropdown)`
+  font-style: italic;
+`;
+
 export const InProgressDropdown = () => {
   return (
-    <Dropdown
+    <StyledDropdown
       value="Thinking..."
       options={[]}
       disabled={true}

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

@@ -126,6 +126,33 @@ const ModelableMethodRow = forwardRef<HTMLElement | undefined, MethodRowProps>(
       [method, modeledMethods, onChange],
     );
 
+    const removeModelClickedHandlers = useMemo(
+      () =>
+        modeledMethods.map((_, index) => () => {
+          const newModeledMethods = [...modeledMethods];
+          newModeledMethods.splice(index, 1);
+          onChange(method.signature, newModeledMethods);
+        }),
+      [method, modeledMethods, onChange],
+    );
+
+    const handleAddModelClick = useCallback(() => {
+      const newModeledMethod: ModeledMethod = {
+        type: "none",
+        input: "",
+        output: "",
+        kind: "",
+        provenance: "manual",
+        signature: method.signature,
+        packageName: method.packageName,
+        typeName: method.typeName,
+        methodName: method.methodName,
+        methodParameters: method.methodParameters,
+      };
+      const newModeledMethods = [...modeledMethods, newModeledMethod];
+      onChange(method.signature, newModeledMethods);
+    }, [method, modeledMethods, onChange]);
+
     const jumpToMethod = useCallback(
       () => sendJumpToMethodMessage(method),
       [method],
@@ -228,6 +255,7 @@ const ModelableMethodRow = forwardRef<HTMLElement | undefined, MethodRowProps>(
                       key={index}
                       appearance="icon"
                       aria-label="Add new model"
+                      onClick={handleAddModelClick}
                       disabled={addModelButtonDisabled}
                     >
                       <Codicon name="add" />
@@ -237,6 +265,7 @@ const ModelableMethodRow = forwardRef<HTMLElement | undefined, MethodRowProps>(
                       key={index}
                       appearance="icon"
                       aria-label="Remove model"
+                      onClick={removeModelClickedHandlers[index]}
                     >
                       <Codicon name="trash" />
                     </CodiconRow>

extensions/ql-vscode/src/view/model-editor/__tests__/MethodRow.spec.tsx

@@ -358,4 +358,84 @@ describe(MethodRow.name, () => {
       expect(removeButton?.getElementsByTagName("input")[0]).toBeEnabled();
     }
   });
+
+  it("can add a new model", async () => {
+    render({
+      modeledMethods: [modeledMethod],
+      viewState: {
+        ...viewState,
+        showMultipleModels: true,
+      },
+    });
+
+    onChange.mockReset();
+    await userEvent.click(screen.getByLabelText("Add new model"));
+
+    expect(onChange).toHaveBeenCalledTimes(1);
+    expect(onChange).toHaveBeenCalledWith(method.signature, [
+      modeledMethod,
+      {
+        type: "none",
+        input: "",
+        output: "",
+        kind: "",
+        provenance: "manual",
+        signature: method.signature,
+        packageName: method.packageName,
+        typeName: method.typeName,
+        methodName: method.methodName,
+        methodParameters: method.methodParameters,
+      },
+    ]);
+  });
+
+  it("can delete the first modeled method", async () => {
+    render({
+      modeledMethods: [
+        { ...modeledMethod, type: "source" },
+        { ...modeledMethod, type: "sink" },
+        { ...modeledMethod, type: "none" },
+        { ...modeledMethod, type: "summary" },
+      ],
+      viewState: {
+        ...viewState,
+        showMultipleModels: true,
+      },
+    });
+
+    onChange.mockReset();
+    await userEvent.click(screen.getAllByLabelText("Remove model")[0]);
+
+    expect(onChange).toHaveBeenCalledTimes(1);
+    expect(onChange).toHaveBeenCalledWith(method.signature, [
+      { ...modeledMethod, type: "sink" },
+      { ...modeledMethod, type: "none" },
+      { ...modeledMethod, type: "summary" },
+    ]);
+  });
+
+  it("can delete a modeled method in the middle", async () => {
+    render({
+      modeledMethods: [
+        { ...modeledMethod, type: "source" },
+        { ...modeledMethod, type: "sink" },
+        { ...modeledMethod, type: "none" },
+        { ...modeledMethod, type: "summary" },
+      ],
+      viewState: {
+        ...viewState,
+        showMultipleModels: true,
+      },
+    });
+
+    onChange.mockReset();
+    await userEvent.click(screen.getAllByLabelText("Remove model")[2]);
+
+    expect(onChange).toHaveBeenCalledTimes(1);
+    expect(onChange).toHaveBeenCalledWith(method.signature, [
+      { ...modeledMethod, type: "source" },
+      { ...modeledMethod, type: "sink" },
+      { ...modeledMethod, type: "summary" },
+    ]);
+  });
 });