Merge pull request #3054 from github/v1.9.4

v1.9.4

.github/workflows/cli-test.yml

@@ -60,9 +60,9 @@ jobs:
       - name: Checkout
         uses: actions/checkout@v4
 
-      - uses: actions/setup-node@v3
+      - uses: actions/setup-node@v4
         with:
-          node-version: '16.17.1'
+          node-version-file: extensions/ql-vscode/.nvmrc
           cache: 'npm'
           cache-dependency-path: extensions/ql-vscode/package-lock.json
 

.github/workflows/main.yml

@@ -20,9 +20,9 @@ jobs:
         with:
           fetch-depth: 1
 
-      - uses: actions/setup-node@v3
+      - uses: actions/setup-node@v4
         with:
-          node-version: '16.17.1'
+          node-version-file: extensions/ql-vscode/.nvmrc
           cache: 'npm'
           cache-dependency-path: extensions/ql-vscode/package-lock.json
 
@@ -62,9 +62,9 @@ jobs:
         with:
           fetch-depth: 1
 
-      - uses: actions/setup-node@v3
+      - uses: actions/setup-node@v4
         with:
-          node-version: '16.17.1'
+          node-version-file: extensions/ql-vscode/.nvmrc
           cache: 'npm'
           cache-dependency-path: extensions/ql-vscode/package-lock.json
 
@@ -81,6 +81,8 @@ jobs:
 
       - name: Lint
         working-directory: extensions/ql-vscode
+        env:
+          NODE_OPTIONS: '--max-old-space-size=4096'
         run: |
           npm run lint
 
@@ -108,9 +110,9 @@ jobs:
         with:
           fetch-depth: 1
 
-      - uses: actions/setup-node@v3
+      - uses: actions/setup-node@v4
         with:
-          node-version: '16.17.1'
+          node-version-file: extensions/ql-vscode/.nvmrc
           cache: 'npm'
           cache-dependency-path: extensions/ql-vscode/package-lock.json
 
@@ -147,9 +149,9 @@ jobs:
         with:
           fetch-depth: 1
 
-      - uses: actions/setup-node@v3
+      - uses: actions/setup-node@v4
         with:
-          node-version: '16.17.1'
+          node-version-file: extensions/ql-vscode/.nvmrc
           cache: 'npm'
           cache-dependency-path: extensions/ql-vscode/package-lock.json
 
@@ -181,9 +183,9 @@ jobs:
         with:
           fetch-depth: 1
 
-      - uses: actions/setup-node@v3
+      - uses: actions/setup-node@v4
         with:
-          node-version: '16.17.1'
+          node-version-file: extensions/ql-vscode/.nvmrc
           cache: 'npm'
           cache-dependency-path: extensions/ql-vscode/package-lock.json
 
@@ -249,9 +251,9 @@ jobs:
       - name: Checkout
         uses: actions/checkout@v4
 
-      - uses: actions/setup-node@v3
+      - uses: actions/setup-node@v4
         with:
-          node-version: '16.17.1'
+          node-version-file: extensions/ql-vscode/.nvmrc
           cache: 'npm'
           cache-dependency-path: extensions/ql-vscode/package-lock.json
 

.github/workflows/release.yml

@@ -20,9 +20,9 @@ jobs:
       - name: Checkout
         uses: actions/checkout@v4
 
-      - uses: actions/setup-node@v3
+      - uses: actions/setup-node@v4
         with:
-          node-version: '16.17.1'
+          node-version-file: extensions/ql-vscode/.nvmrc
 
       - name: Install dependencies
         run: |

CODEOWNERS

@@ -1,5 +1,6 @@
 **/* @github/codeql-vscode-reviewers
 **/variant-analysis/ @github/code-scanning-secexp-reviewers
 **/databases/ @github/code-scanning-secexp-reviewers
+**/method-modeling/ @github/code-scanning-secexp-reviewers
 **/model-editor/ @github/code-scanning-secexp-reviewers
 **/queries-panel/ @github/code-scanning-secexp-reviewers

docs/node-version.md

@@ -15,12 +15,13 @@ You can find this info by seleting "About Visual Studio Code" from the top menu.
 
 The following files will need to be updated:
 
-- `.github/workflows/cli-test.yml` - the "node-version: '[VERSION]'" setting
-- `.github/workflows/main.yml` - all the "node-version: '[VERSION]'" settings
-- `.github/workflows/release.yml` - the "node-version: '[VERSION]'" setting
-- `extensions/ql-vscode/.nvmrc` - this will enable nvm to automatically switch to the correct node version when you're in the project folder
-- `extensions/ql-vscode/package-lock.json` - the "engines.node: '[VERSION]'" setting
+- `extensions/ql-vscode/.nvmrc` - this will enable nvm to automatically switch to the correct Node
+   version when you're in the project folder. It will also change the Node version the GitHub Actions
+   workflows use.
 - `extensions/ql-vscode/package.json` - the "engines.node: '[VERSION]'" setting
+- `extensions/ql-vscode/package.json` - the "@types/node: '[VERSION]'" dependency
+
+Then run `npm install` to update the `extensions/ql-vscode/package-lock.json` file.
 
 ## Node.js version used in tests
 

docs/test-plan.md

@@ -151,17 +151,20 @@ Run one of the above MRVAs, but cancel it from within VS Code:
 2. Open the Model Editor with the "CodeQL: Open CodeQL Model Editor" command from the command palette.
    - Check that the editor loads and shows methods to model.
    - Check that methods are grouped per library (e.g. `rocksdbjni@7.7.3` or `asm@6.0`)
-   - Check that the "Open database" link works.
+   - Check that the "Open source" link works.
+   - Check that the 'View' button works and the Method Usage panel highlight the correct method and usage
+   - Check that the Method Modeling panel shows the correct method and modeling state
 
 #### Test Case 2: Model methods
 
 1. Expand one of the libraries.
    - Change the model type and check that the other dropdowns change.
+   - Check that the method modeling panel updates accordingly
 2. Save the modeled methods.
 3. Click "Open extension pack"
    - Check that the file explorer opens a directory with a "models" directory
 4. Open the ".model.yml" file corresponding to the library that was changed.
-   - Check that the file contrains the entries that was modeled.
+   - Check that the file contains entries for the methods that were modeled.
 
 #### Test Case 3: Model with AI
 
@@ -173,6 +176,8 @@ Note that this test requires the feature flag: `codeQL.model.llmGeneration`
 
 #### Test Case 4: Model as dependency
 
+Note that this test requires the feature flag: `codeQL.model.flowGeneration`
+
 1. Click "Model as dependency"
    - Check that grouping are now per package (e.g. `com.alipay.sofa.rraft.option` or `com.google.protobuf`)
 2. Click "Generate".
@@ -187,9 +192,28 @@ Are there any components that are not showing up?
 
 ## Optional Test Cases
 
-These are mostly aimed at MRVA, but some of them are also applicable to non-MRVA queries.
+### Modeling Flow
 
-### Selecting repositories to run on
+1. Check that a method can have multiple models:
+   - Add a couple of new models for one method in the model editor
+   - Save and check that the modeling file (use the 'open extension pack' button to open it) shows multiple methods
+   - Check that the Method Modeling Panel shows the correct multiple models
+   - Check that you can browse through different models in the Method Modeling Panel
+   - Check that a 'duplicated classification' error appears in both model editor and modeling panel when a duplicate modeling occurs
+   - Check that a 'conflicting classification' error appears when a neutral model type is paired with a model of the same kind
+   - Check that clicking on the error highlights the correct modeling in both the editor and the modeling panel
+2. Check the Method Usage Panel
+   - Check that the Method Usage Panel opens and jumps to the correct usage when clicking on 'View' in the model editor
+   - Check that the first and following usages are opening when clicking on a usage
+   - Check that the usage icon color turns green when saving a newly modeled method
+   - Check that the usage icon color turns red when saving a newly unmodeld method
+3. Check the Method Modeling Panel
+   - Check that the 'Start modeling' button opens a new model editor
+   - Check that it refreshes the blank state when a model editor is opened/closed
+   - Check that when modeling in the editor the modeling panel updates accordingly
+   - Check that when modeling in the modeling panel the model editor updates accordingly
+
+### Selecting MRVA repositories to run on
 
 #### Test case 1: Running a query on a single repository
 
@@ -219,7 +243,7 @@ These are mostly aimed at MRVA, but some of them are also applicable to non-MRVA
    4. The org contains private repositories that are inaccessible
 2. The org does not exist
 
-### Using different types of controller repos
+### Using different types of controller repos for MRVA
 
 #### Test case 1: Running a query when the controller repository is public
 

docs/testing.md

@@ -27,7 +27,7 @@ Pre-requisites:
 Then, from the `extensions/ql-vscode` directory, use the appropriate command to run the tests:
 
 * Unit tests: `npm run test:unit`
-* View Tests: `npm test:view`
+* View Tests: `npm run test:view`
 * VSCode integration tests: `npm run test:vscode-integration`
 
 #### Running CLI integration tests from the terminal
@@ -48,8 +48,8 @@ Alternatively, you can run the tests inside of VSCode. There are several VSCode
 
 You will need to run tests using a task from inside of VS Code, under the "Run and Debug" view:
 
-* Unit tests: run the _Launch Unit Tests - React_ task
-* View Tests: run the _Launch Unit Tests_ task
+* Unit tests: run the _Launch Unit Tests_ task
+* View Tests: run the _Launch Unit Tests - React_ task
 * VSCode integration tests: run the _Launch Unit Tests - No Workspace_ and _Launch Unit Tests - Minimal Workspace_ tasks
 
 #### Running CLI integration tests from VSCode

extensions/ql-vscode/.eslintrc.js

@@ -71,6 +71,7 @@ const baseConfig = {
     "no-shadow": "off",
     "github/array-foreach": "off",
     "github/no-then": "off",
+    "react/jsx-key": ["error", { checkFragmentShorthand: true }],
   },
 };
 

extensions/ql-vscode/.nvmrc

@@ -1 +1 @@
-v16.17.1
+v18.15.0

extensions/ql-vscode/.storybook/main.ts

@@ -6,6 +6,7 @@ const config: StorybookConfig = {
     "@storybook/addon-links",
     "@storybook/addon-essentials",
     "@storybook/addon-interactions",
+    "@storybook/addon-a11y",
     "./vscode-theme-addon/preset.ts",
   ],
   framework: {

extensions/ql-vscode/CHANGELOG.md

@@ -1,5 +1,30 @@
 # CodeQL for Visual Studio Code: Changelog
 
+## 1.9.4 - 6 November 2023
+
+No user facing changes.
+
+## 1.9.3 - 26 October 2023
+
+- Sorted result set filenames now include a hash of the result set name instead of the full name. [#2955](https://github.com/github/vscode-codeql/pull/2955)
+- The "Install Pack Dependencies" will now only list CodeQL packs located in the workspace. [#2960](https://github.com/github/vscode-codeql/pull/2960)
+- Fix a bug where the "View Query Log" action for a query history item was not working. [#2984](https://github.com/github/vscode-codeql/pull/2984)
+- Add a command to sort items in the databases view by language. [#2993](https://github.com/github/vscode-codeql/pull/2993)
+- Fix not being able to open the results directory or evaluator log for a cancelled local query run. [#2996](https://github.com/github/vscode-codeql/pull/2996)
+- Fix empty row in alert path when the SARIF location was empty. [#3018](https://github.com/github/vscode-codeql/pull/3018)
+
+## 1.9.2 - 12 October 2023
+
+- Fix a bug where the query to Find Definitions in database source files would not be cancelled appropriately. [#2885](https://github.com/github/vscode-codeql/pull/2885)
+- It is now possible to show the language of query history items using the `%l` specifier in the `codeQL.queryHistory.format` setting. Note that this only works for queries run after this upgrade, and older items will show `unknown` as a language. [#2892](https://github.com/github/vscode-codeql/pull/2892)
+- 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).
+- Add a "CodeQL: Trim Cache" command that clears the evaluation cache of a database except for predicates annotated with the `cached` keyword. Its purpose is to get accurate performance measurements when tuning the final stage of a query, like a data-flow configuration. This is equivalent to the `codeql database cleanup --mode=normal` CLI command. In contrast, the existing "CodeQL: Clear Cache" command clears the entire cache. CodeQL CLI v2.15.1 or later is required. [#2928](https://github.com/github/vscode-codeql/pull/2928)
+- Fix syntax highlighting directly after import statements with instantiation arguments. [#2792](https://github.com/github/vscode-codeql/pull/2792)
+- The `debug.saveBeforeStart` setting is now respected when running variant analyses. [#2950](https://github.com/github/vscode-codeql/pull/2950)
+- The 'open database' button of the model editor was renamed to 'open source'. Also, it's now only available if the source archive is available as a workspace folder. [#2945](https://github.com/github/vscode-codeql/pull/2945)
+
 ## 1.9.1 - 29 September 2023
 
 - Add warning when using a VS Code version older than 1.82.0. [#2854](https://github.com/github/vscode-codeql/pull/2854)

extensions/ql-vscode/package.json

@@ -4,7 +4,7 @@
   "description": "CodeQL for Visual Studio Code",
   "author": "GitHub",
   "private": true,
-  "version": "1.9.1",
+  "version": "1.9.4",
   "publisher": "GitHub",
   "license": "MIT",
   "icon": "media/VS-marketplace-CodeQL-icon.png",
@@ -13,8 +13,8 @@
     "url": "https://github.com/github/vscode-codeql"
   },
   "engines": {
-    "vscode": "^1.67.0",
-    "node": "^16.17.1",
+    "vscode": "^1.82.0",
+    "node": "^18.15.0",
     "npm": ">=7.20.6"
   },
   "categories": [
@@ -34,27 +34,7 @@
     }
   },
   "activationEvents": [
-    "onLanguage:ql",
-    "onLanguage:ql-summary",
-    "onView:codeQLQueries",
-    "onView:codeQLDatabases",
-    "onView:codeQLVariantAnalysisRepositories",
-    "onView:codeQLQueryHistory",
-    "onView:codeQLAstViewer",
-    "onView:codeQLEvalLogViewer",
     "onView:test-explorer",
-    "onCommand:codeQL.checkForUpdatesToCLI",
-    "onCommand:codeQL.authenticateToGitHub",
-    "onCommand:codeQL.viewAst",
-    "onCommand:codeQL.viewCfg",
-    "onCommand:codeQL.openReferencedFile",
-    "onCommand:codeQL.previewQueryHelp",
-    "onCommand:codeQL.chooseDatabaseFolder",
-    "onCommand:codeQL.chooseDatabaseArchive",
-    "onCommand:codeQL.chooseDatabaseInternet",
-    "onCommand:codeQL.chooseDatabaseGithub",
-    "onCommand:codeQL.quickQuery",
-    "onCommand:codeQL.restartQueryServer",
     "onWebviewPanel:resultsView",
     "onWebviewPanel:codeQL.variantAnalysis",
     "onWebviewPanel:codeQL.dataFlowPaths",
@@ -110,6 +90,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."
               }
             }
           }
@@ -446,13 +430,20 @@
             "type": "boolean",
             "default": false,
             "scope": "application",
-            "markdownDescription": "Specifies whether to send CodeQL usage telemetry. This setting AND the global `#telemetry.enableTelemetry#` setting must be checked for telemetry to be sent to GitHub. 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)"
+            "markdownDescription": "Specifies whether to send CodeQL usage telemetry. This setting AND the one of the global telemetry settings (`#telemetry.enableTelemetry#` or `#telemetry.telemetryLevel#`) must be enabled for telemetry to be sent to GitHub. 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)",
+            "tags": [
+              "telemetry",
+              "usesOnlineServices"
+            ]
           },
           "codeQL.telemetry.logTelemetry": {
             "type": "boolean",
             "default": false,
             "scope": "application",
-            "description": "Specifies whether or not to write telemetry events to the extension log."
+            "description": "Specifies whether or not to write telemetry events to the extension log.",
+            "tags": [
+              "telemetry"
+            ]
           }
         }
       }
@@ -570,6 +561,10 @@
         "command": "codeQL.copyVersion",
         "title": "CodeQL: Copy Version Information"
       },
+      {
+        "command": "codeQLLanguageSelection.setSelectedItem",
+        "title": "Select"
+      },
       {
         "command": "codeQLQueries.runLocalQueryFromQueriesPanel",
         "title": "Run local query",
@@ -721,6 +716,10 @@
         "command": "codeQL.clearCache",
         "title": "CodeQL: Clear Cache"
       },
+      {
+        "command": "codeQL.trimCache",
+        "title": "CodeQL: Trim Cache"
+      },
       {
         "command": "codeQL.installPackDependencies",
         "title": "CodeQL: Install Pack Dependencies"
@@ -753,78 +752,6 @@
         "command": "codeQLDatabases.addDatabaseSource",
         "title": "Add Database Source to Workspace"
       },
-      {
-        "command": "codeQLDatabases.displayAllLanguages",
-        "title": "All languages"
-      },
-      {
-        "command": "codeQLDatabases.displayAllLanguagesSelected",
-        "title": "All languages (selected)"
-      },
-      {
-        "command": "codeQLDatabases.displayCpp",
-        "title": "C/C++"
-      },
-      {
-        "command": "codeQLDatabases.displayCppSelected",
-        "title": "C/C++ (selected)"
-      },
-      {
-        "command": "codeQLDatabases.displayCsharp",
-        "title": "C#"
-      },
-      {
-        "command": "codeQLDatabases.displayCsharpSelected",
-        "title": "C# (selected)"
-      },
-      {
-        "command": "codeQLDatabases.displayGo",
-        "title": "Go"
-      },
-      {
-        "command": "codeQLDatabases.displayGoSelected",
-        "title": "Go (selected)"
-      },
-      {
-        "command": "codeQLDatabases.displayJava",
-        "title": "Java/Kotlin"
-      },
-      {
-        "command": "codeQLDatabases.displayJavaSelected",
-        "title": "Java/Kotlin (selected)"
-      },
-      {
-        "command": "codeQLDatabases.displayJavascript",
-        "title": "JavaScript/TypeScript"
-      },
-      {
-        "command": "codeQLDatabases.displayJavascriptSelected",
-        "title": "JavaScript/TypeScript (selected)"
-      },
-      {
-        "command": "codeQLDatabases.displayPython",
-        "title": "Python"
-      },
-      {
-        "command": "codeQLDatabases.displayPythonSelected",
-        "title": "Python (selected)"
-      },
-      {
-        "command": "codeQLDatabases.displayRuby",
-        "title": "Ruby"
-      },
-      {
-        "command": "codeQLDatabases.displayRubySelected",
-        "title": "Ruby (selected)"
-      },
-      {
-        "command": "codeQLDatabases.displaySwift",
-        "title": "Swift"
-      },
-      {
-        "command": "codeQLDatabases.displaySwiftSelected",
-        "title": "Swift (selected)"
-      },
       {
         "command": "codeQL.chooseDatabaseFolder",
         "title": "CodeQL: Choose Database from Folder"
@@ -845,6 +772,10 @@
         "command": "codeQLDatabases.sortByName",
         "title": "Sort by Name"
       },
+      {
+        "command": "codeQLDatabases.sortByLanguage",
+        "title": "Sort by Language"
+      },
       {
         "command": "codeQLDatabases.sortByDateAdded",
         "title": "Sort by Date Added"
@@ -1071,14 +1002,14 @@
           "group": "1_databases@0"
         },
         {
-          "command": "codeQLDatabases.sortByDateAdded",
+          "command": "codeQLDatabases.sortByLanguage",
           "when": "view == codeQLDatabases",
           "group": "1_databases@1"
         },
         {
-          "submenu": "codeQLDatabases.languages",
-          "when": "view == codeQLDatabases && config.codeQL.canary && config.codeQL.showLanguageFilter",
-          "group": "2_databases@0"
+          "command": "codeQLDatabases.sortByDateAdded",
+          "when": "view == codeQLDatabases",
+          "group": "1_databases@2"
         },
         {
           "command": "codeQLQueries.createQuery",
@@ -1152,6 +1083,11 @@
           "when": "view == codeQLVariantAnalysisRepositories && viewItem =~ /canImportCodeSearch/",
           "group": "2_qlContextMenu@1"
         },
+        {
+          "command": "codeQLLanguageSelection.setSelectedItem",
+          "when": "view == codeQLLanguageSelection && viewItem =~ /canBeSelected/",
+          "group": "inline"
+        },
         {
           "command": "codeQLDatabases.setCurrentDatabase",
           "group": "inline",
@@ -1500,6 +1436,10 @@
         {
           "command": "codeQL.openModelEditor"
         },
+        {
+          "command": "codeQLLanguageSelection.setSelectedItem",
+          "when": "false"
+        },
         {
           "command": "codeQLQueries.runLocalQueryContextMenu",
           "when": "false"
@@ -1568,6 +1508,10 @@
           "command": "codeQLDatabases.sortByName",
           "when": "false"
         },
+        {
+          "command": "codeQLDatabases.sortByLanguage",
+          "when": "false"
+        },
         {
           "command": "codeQLDatabases.sortByDateAdded",
           "when": "false"
@@ -1600,78 +1544,6 @@
           "command": "codeQLDatabases.upgradeDatabase",
           "when": "false"
         },
-        {
-          "command": "codeQLDatabases.displayAllLanguages",
-          "when": "false"
-        },
-        {
-          "command": "codeQLDatabases.displayAllLanguagesSelected",
-          "when": "false"
-        },
-        {
-          "command": "codeQLDatabases.displayCpp",
-          "when": "false"
-        },
-        {
-          "command": "codeQLDatabases.displayCppSelected",
-          "when": "false"
-        },
-        {
-          "command": "codeQLDatabases.displayCsharp",
-          "when": "false"
-        },
-        {
-          "command": "codeQLDatabases.displayCsharpSelected",
-          "when": "false"
-        },
-        {
-          "command": "codeQLDatabases.displayGo",
-          "when": "false"
-        },
-        {
-          "command": "codeQLDatabases.displayGoSelected",
-          "when": "false"
-        },
-        {
-          "command": "codeQLDatabases.displayJava",
-          "when": "false"
-        },
-        {
-          "command": "codeQLDatabases.displayJavaSelected",
-          "when": "false"
-        },
-        {
-          "command": "codeQLDatabases.displayJavascript",
-          "when": "false"
-        },
-        {
-          "command": "codeQLDatabases.displayJavascriptSelected",
-          "when": "false"
-        },
-        {
-          "command": "codeQLDatabases.displayPython",
-          "when": "false"
-        },
-        {
-          "command": "codeQLDatabases.displayPythonSelected",
-          "when": "false"
-        },
-        {
-          "command": "codeQLDatabases.displayRuby",
-          "when": "false"
-        },
-        {
-          "command": "codeQLDatabases.displayRubySelected",
-          "when": "false"
-        },
-        {
-          "command": "codeQLDatabases.displaySwift",
-          "when": "false"
-        },
-        {
-          "command": "codeQLDatabases.displaySwiftSelected",
-          "when": "false"
-        },
         {
           "command": "codeQLQueryHistory.openQueryContextMenu",
           "when": "false"
@@ -1815,6 +1687,10 @@
         {
           "command": "codeQL.gotoQLContextEditor",
           "when": "false"
+        },
+        {
+          "command": "codeQL.trimCache",
+          "when": "codeql.supportsTrimCache"
         }
       ],
       "editor/context": [
@@ -1866,88 +1742,8 @@
           "command": "codeQL.gotoQLContextEditor",
           "when": "editorLangId == ql-summary && config.codeQL.canary"
         }
-      ],
-      "codeQLDatabases.languages": [
-        {
-          "command": "codeQLDatabases.displayAllLanguages",
-          "when": "codeQLDatabases.languageFilter"
-        },
-        {
-          "command": "codeQLDatabases.displayAllLanguagesSelected",
-          "when": "!codeQLDatabases.languageFilter"
-        },
-        {
-          "command": "codeQLDatabases.displayCpp",
-          "when": "codeQLDatabases.languageFilter != cpp"
-        },
-        {
-          "command": "codeQLDatabases.displayCppSelected",
-          "when": "codeQLDatabases.languageFilter == cpp"
-        },
-        {
-          "command": "codeQLDatabases.displayCsharp",
-          "when": "codeQLDatabases.languageFilter != csharp"
-        },
-        {
-          "command": "codeQLDatabases.displayCsharpSelected",
-          "when": "codeQLDatabases.languageFilter == csharp"
-        },
-        {
-          "command": "codeQLDatabases.displayGo",
-          "when": "codeQLDatabases.languageFilter != go"
-        },
-        {
-          "command": "codeQLDatabases.displayGoSelected",
-          "when": "codeQLDatabases.languageFilter == go"
-        },
-        {
-          "command": "codeQLDatabases.displayJava",
-          "when": "codeQLDatabases.languageFilter != java"
-        },
-        {
-          "command": "codeQLDatabases.displayJavaSelected",
-          "when": "codeQLDatabases.languageFilter == java"
-        },
-        {
-          "command": "codeQLDatabases.displayJavascript",
-          "when": "codeQLDatabases.languageFilter != javascript"
-        },
-        {
-          "command": "codeQLDatabases.displayJavascriptSelected",
-          "when": "codeQLDatabases.languageFilter == javascript"
-        },
-        {
-          "command": "codeQLDatabases.displayPython",
-          "when": "codeQLDatabases.languageFilter != python"
-        },
-        {
-          "command": "codeQLDatabases.displayPythonSelected",
-          "when": "codeQLDatabases.languageFilter == python"
-        },
-        {
-          "command": "codeQLDatabases.displayRuby",
-          "when": "codeQLDatabases.languageFilter != ruby"
-        },
-        {
-          "command": "codeQLDatabases.displayRubySelected",
-          "when": "codeQLDatabases.languageFilter == ruby"
-        },
-        {
-          "command": "codeQLDatabases.displaySwift",
-          "when": "codeQLDatabases.languageFilter != swift"
-        },
-        {
-          "command": "codeQLDatabases.displaySwiftSelected",
-          "when": "codeQLDatabases.languageFilter == swift"
-        }
       ]
     },
-    "submenus": [
-      {
-        "id": "codeQLDatabases.languages",
-        "label": "Languages"
-      }
-    ],
     "viewsContainers": {
       "activitybar": [
         {
@@ -1966,6 +1762,11 @@
     },
     "views": {
       "ql-container": [
+        {
+          "id": "codeQLLanguageSelection",
+          "name": "Language",
+          "when": "config.codeQL.canary && config.codeQL.showLanguageFilter"
+        },
         {
           "id": "codeQLDatabases",
           "name": "Databases"
@@ -1991,6 +1792,12 @@
           "id": "codeQLEvalLogViewer",
           "name": "Evaluator Log Viewer",
           "when": "config.codeQL.canary"
+        },
+        {
+          "id": "codeQLMethodModeling",
+          "type": "webview",
+          "name": "CodeQL Method Modeling",
+          "when": "config.codeQL.canary"
         }
       ],
       "codeql-methods-usage": [
@@ -1999,14 +1806,6 @@
           "name": "CodeQL Methods Usage",
           "when": "config.codeQL.canary && codeql.modelEditorOpen"
         }
-      ],
-      "explorer": [
-        {
-          "type": "webview",
-          "id": "codeQLMethodModeling",
-          "name": "CodeQL Method Modeling",
-          "when": "config.codeQL.canary && config.codeQL.model.methodModelingView && codeql.modelEditorOpen && !codeql.modelEditorActive"
-        }
       ]
     },
     "viewsWelcome": [
@@ -2070,8 +1869,8 @@
     "prepare": "cd ../.. && husky install"
   },
   "dependencies": {
-    "@octokit/plugin-retry": "^4.1.6",
-    "@octokit/rest": "^19.0.4",
+    "@octokit/plugin-retry": "^6.0.1",
+    "@octokit/rest": "^20.0.2",
     "@vscode/codicons": "^0.0.31",
     "@vscode/debugadapter": "^1.59.0",
     "@vscode/debugprotocol": "^1.59.0",
@@ -2085,10 +1884,10 @@
     "fs-extra": "^11.1.1",
     "immutable": "^4.0.0",
     "js-yaml": "^4.1.0",
-    "msw": "^1.2.0",
-    "nanoid": "^3.2.0",
+    "msw": "^2.0.0",
+    "nanoid": "^5.0.1",
     "node-fetch": "^2.6.7",
-    "p-queue": "^6.0.0",
+    "p-queue": "^7.4.1",
     "react": "^18.2.0",
     "react-dom": "^18.2.0",
     "semver": "^7.5.2",
@@ -2115,7 +1914,8 @@
     "@babel/preset-typescript": "^7.21.4",
     "@faker-js/faker": "^8.0.2",
     "@github/markdownlint-github": "^0.3.0",
-    "@octokit/plugin-throttling": "^5.0.1",
+    "@octokit/plugin-throttling": "^8.0.0",
+    "@storybook/addon-a11y": "^7.4.6",
     "@storybook/addon-actions": "^7.1.0",
     "@storybook/addon-essentials": "^7.1.0",
     "@storybook/addon-interactions": "^7.1.0",
@@ -2141,7 +1941,7 @@
     "@types/jest": "^29.0.2",
     "@types/js-yaml": "^4.0.6",
     "@types/nanoid": "^3.0.0",
-    "@types/node": "^16.11.25",
+    "@types/node": "18.15.0",
     "@types/node-fetch": "^2.5.2",
     "@types/react": "^18.0.28",
     "@types/react-dom": "^18.0.11",
@@ -2153,7 +1953,7 @@
     "@types/through2": "^2.0.36",
     "@types/tmp": "^0.1.0",
     "@types/unzipper": "^0.10.1",
-    "@types/vscode": "^1.67.0",
+    "@types/vscode": "^1.82.0",
     "@types/webpack": "^5.28.0",
     "@types/webpack-env": "^1.18.0",
     "@typescript-eslint/eslint-plugin": "^6.2.1",
@@ -2181,13 +1981,12 @@
     "gulp": "^4.0.2",
     "gulp-esbuild": "^0.10.5",
     "gulp-replace": "^1.1.3",
-    "gulp-sourcemaps": "^3.0.0",
     "gulp-typescript": "^5.0.1",
     "husky": "^8.0.0",
     "jest": "^29.0.3",
     "jest-environment-jsdom": "^29.0.3",
     "jest-runner-vscode": "^3.0.1",
-    "lint-staged": "^14.0.0",
+    "lint-staged": "^15.0.2",
     "markdownlint-cli2": "^0.6.0",
     "markdownlint-cli2-formatter-pretty": "^0.0.4",
     "mini-css-extract-plugin": "^2.6.1",

extensions/ql-vscode/scripts/add-fields-to-scenarios.ts

@@ -14,7 +14,8 @@
 import { pathExists, readJson, writeJson } from "fs-extra";
 import { resolve, relative } from "path";
 
-import { Octokit, type RestEndpointMethodTypes } from "@octokit/rest";
+import { Octokit } from "@octokit/core";
+import { type RestEndpointMethodTypes } from "@octokit/rest";
 import { throttling } from "@octokit/plugin-throttling";
 
 import { getFiles } from "./util/files";
@@ -22,6 +23,7 @@ import type { GitHubApiRequest } from "../src/common/mock-gh-api/gh-api-request"
 import { isGetVariantAnalysisRequest } from "../src/common/mock-gh-api/gh-api-request";
 import { VariantAnalysis } from "../src/variant-analysis/gh-api/variant-analysis";
 import { RepositoryWithMetadata } from "../src/variant-analysis/gh-api/repository";
+import { AppOctokit } from "../src/common/octokit";
 
 const extensionDirectory = resolve(__dirname, "..");
 const scenariosDirectory = resolve(
@@ -31,7 +33,7 @@ const scenariosDirectory = resolve(
 
 // Make sure we don't run into rate limits by automatically waiting until we can
 // make another request.
-const MyOctokit = Octokit.plugin(throttling);
+const MyOctokit = AppOctokit.plugin(throttling);
 
 const auth = process.env.GITHUB_TOKEN;
 

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/codeql-cli/cli.ts

@@ -10,7 +10,11 @@ import tk from "tree-kill";
 import { promisify } from "util";
 import { CancellationToken, Disposable, Uri } from "vscode";
 
-import { BQRSInfo, DecodedBqrsChunk } from "../common/bqrs-cli-types";
+import {
+  BQRSInfo,
+  DecodedBqrs,
+  DecodedBqrsChunk,
+} from "../common/bqrs-cli-types";
 import { allowCanaryQueryServer, CliConfig } from "../config";
 import {
   DistributionProvider,
@@ -1040,6 +1044,18 @@ export class CodeQLCliServer implements Disposable {
     );
   }
 
+  /**
+   * Gets all results from a bqrs.
+   * @param bqrsPath The path to the bqrs.
+   */
+  async bqrsDecodeAll(bqrsPath: string): Promise<DecodedBqrs> {
+    return await this.runJsonCodeQlCliCommand<DecodedBqrs>(
+      ["bqrs", "decode"],
+      [bqrsPath],
+      "Reading all bqrs data",
+    );
+  }
+
   async runInterpretCommand(
     format: string,
     additonalArgs: string[],
@@ -1244,11 +1260,13 @@ export class CodeQLCliServer implements Disposable {
    * @param additionalPacks A list of directories to search for qlpacks.
    * @param extensionPacksOnly Whether to only search for extension packs. If true, only extension packs will
    *    be returned. If false, all packs will be returned.
+   * @param kind Whether to only search for qlpacks with a certain kind.
    * @returns A dictionary mapping qlpack name to the directory it comes from
    */
   async resolveQlpacks(
     additionalPacks: string[],
     extensionPacksOnly = false,
+    kind?: "query" | "library" | "all",
   ): Promise<QlpacksInfo> {
     const args = this.getAdditionalPacksArg(additionalPacks);
     if (extensionPacksOnly) {
@@ -1259,6 +1277,8 @@ export class CodeQLCliServer implements Disposable {
         return {};
       }
       args.push("--kind", "extension", "--no-recursive");
+    } else if (kind) {
+      args.push("--kind", kind);
     }
 
     return this.runJsonCodeQlCliCommand<QlpacksInfo>(
@@ -1492,6 +1512,13 @@ export class CodeQLCliServer implements Disposable {
             CliVersionConstraint.CLI_VERSION_WITH_QUICK_EVAL_COUNT,
           ) >= 0,
         );
+        await this.app.commands.execute(
+          "setContext",
+          "codeql.supportsTrimCache",
+          newVersion.compare(
+            CliVersionConstraint.CLI_VERSION_WITH_TRIM_CACHE,
+          ) >= 0,
+        );
       } catch (e) {
         this._versionChangedListeners.forEach((listener) =>
           listener(undefined),
@@ -1755,6 +1782,12 @@ export class CliVersionConstraint {
     "2.14.0",
   );
 
+  /**
+   * CLI version where the query server supports the `evaluation/trimCache` method
+   * with `codeql database cleanup --mode=trim` semantics.
+   */
+  public static CLI_VERSION_WITH_TRIM_CACHE = new SemVer("2.15.1");
+
   constructor(private readonly cli: CodeQLCliServer) {
     /**/
   }

extensions/ql-vscode/src/codeql-cli/query-language.ts

@@ -1,6 +1,10 @@
 import { CodeQLCliServer } from "./cli";
 import { Uri, window } from "vscode";
-import { isQueryLanguage, QueryLanguage } from "../common/query-language";
+import {
+  getLanguageDisplayName,
+  isQueryLanguage,
+  QueryLanguage,
+} from "../common/query-language";
 import { getOnDiskWorkspaceFolders } from "../common/vscode/workspace-folders";
 import { extLogger } from "../common/logging/vscode";
 import { UserCancellationException } from "../common/vscode/progress";
@@ -46,14 +50,22 @@ export async function askForLanguage(
   cliServer: CodeQLCliServer,
   throwOnEmpty = true,
 ): Promise<QueryLanguage | undefined> {
-  const language = await window.showQuickPick(
-    await cliServer.getSupportedLanguages(),
-    {
-      placeHolder: "Select target language for your query",
-      ignoreFocusOut: true,
-    },
-  );
-  if (!language) {
+  const supportedLanguages = await cliServer.getSupportedLanguages();
+
+  const items = supportedLanguages
+    .filter((language) => isQueryLanguage(language))
+    .map((language) => ({
+      label: getLanguageDisplayName(language),
+      description: language,
+      language,
+    }))
+    .sort((a, b) => a.label.localeCompare(b.label));
+
+  const selectedItem = await window.showQuickPick(items, {
+    placeHolder: "Select target language for your query",
+    ignoreFocusOut: true,
+  });
+  if (!selectedItem) {
     // This only happens if the user cancels the quick pick.
     if (throwOnEmpty) {
       throw new UserCancellationException("Cancelled.");
@@ -66,6 +78,8 @@ export async function askForLanguage(
     return undefined;
   }
 
+  const language = selectedItem.language;
+
   if (!isQueryLanguage(language)) {
     void showAndLogErrorMessage(
       extLogger,

extensions/ql-vscode/src/common/bqrs-cli-types.ts

@@ -121,3 +121,5 @@ export interface DecodedBqrsChunk {
   next?: number;
   columns: BqrsColumn[];
 }
+
+export type DecodedBqrs = Record<string, DecodedBqrsChunk>;

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

@@ -12,6 +12,7 @@ import type {
 } from "../variant-analysis/shared/variant-analysis";
 import type { QLDebugConfiguration } from "../debugger/debug-configuration";
 import type { QueryTreeViewItem } from "../queries-panel/query-tree-view-item";
+import type { LanguageSelectionTreeViewItem } from "../language-selection-panel/language-selection-data-provider";
 import type { Method, Usage } from "../model-editor/method";
 
 // A command function matching the signature that VS Code calls when
@@ -199,6 +200,13 @@ export type QueryHistoryCommands = {
   "codeQL.exportSelectedVariantAnalysisResults": () => Promise<void>;
 };
 
+// Commands user for the language selector panel
+export type LanguageSelectionCommands = {
+  "codeQLLanguageSelection.setSelectedItem": (
+    item: LanguageSelectionTreeViewItem,
+  ) => Promise<void>;
+};
+
 // Commands used for the local databases panel
 export type LocalDatabasesCommands = {
   // Command palette commands
@@ -208,6 +216,7 @@ export type LocalDatabasesCommands = {
   "codeQL.chooseDatabaseGithub": () => Promise<void>;
   "codeQL.upgradeCurrentDatabase": () => Promise<void>;
   "codeQL.clearCache": () => Promise<void>;
+  "codeQL.trimCache": () => Promise<void>;
 
   // Explorer context menu
   "codeQL.setCurrentDatabase": (uri: Uri) => Promise<void>;
@@ -218,25 +227,8 @@ export type LocalDatabasesCommands = {
   "codeQLDatabases.chooseDatabaseInternet": () => Promise<void>;
   "codeQLDatabases.chooseDatabaseGithub": () => Promise<void>;
   "codeQLDatabases.sortByName": () => Promise<void>;
+  "codeQLDatabases.sortByLanguage": () => Promise<void>;
   "codeQLDatabases.sortByDateAdded": () => Promise<void>;
-  "codeQLDatabases.displayAllLanguages": () => Promise<void>;
-  "codeQLDatabases.displayCpp": () => Promise<void>;
-  "codeQLDatabases.displayCsharp": () => Promise<void>;
-  "codeQLDatabases.displayGo": () => Promise<void>;
-  "codeQLDatabases.displayJava": () => Promise<void>;
-  "codeQLDatabases.displayJavascript": () => Promise<void>;
-  "codeQLDatabases.displayPython": () => Promise<void>;
-  "codeQLDatabases.displayRuby": () => Promise<void>;
-  "codeQLDatabases.displaySwift": () => Promise<void>;
-  "codeQLDatabases.displayAllLanguagesSelected": () => Promise<void>;
-  "codeQLDatabases.displayCppSelected": () => Promise<void>;
-  "codeQLDatabases.displayCsharpSelected": () => Promise<void>;
-  "codeQLDatabases.displayGoSelected": () => Promise<void>;
-  "codeQLDatabases.displayJavaSelected": () => Promise<void>;
-  "codeQLDatabases.displayJavascriptSelected": () => Promise<void>;
-  "codeQLDatabases.displayPythonSelected": () => Promise<void>;
-  "codeQLDatabases.displayRubySelected": () => Promise<void>;
-  "codeQLDatabases.displaySwiftSelected": () => Promise<void>;
 
   // Database panel context menu
   "codeQLDatabases.setCurrentDatabase": (
@@ -323,7 +315,8 @@ export type PackagingCommands = {
 
 export type ModelEditorCommands = {
   "codeQL.openModelEditor": () => Promise<void>;
-  "codeQLModelEditor.jumpToUsageLocation": (
+  "codeQL.openModelEditorFromModelingPanel": () => Promise<void>;
+  "codeQLModelEditor.jumpToMethod": (
     method: Method,
     usage: Usage,
     databaseItem: DatabaseItem,
@@ -360,6 +353,7 @@ export type AllExtensionCommands = BaseCommands &
   QueryEditorCommands &
   ResultsViewCommands &
   QueryHistoryCommands &
+  LanguageSelectionCommands &
   LocalDatabasesCommands &
   DebuggerCommands &
   VariantAnalysisCommands &

extensions/ql-vscode/src/common/disposable-object.ts

@@ -9,10 +9,16 @@ export type DisposeHandler = (disposable: Disposable) => void;
 /**
  * Base class to make it easier to implement a `Disposable` that owns other disposable object.
  */
-export abstract class DisposableObject implements Disposable {
+export class DisposableObject implements Disposable {
   private disposables: Disposable[] = [];
   private tracked?: Set<Disposable> = undefined;
 
+  constructor(...dispoables: Disposable[]) {
+    for (const d of dispoables) {
+      this.push(d);
+    }
+  }
+
   /**
    * Adds `obj` to a list of objects to dispose when `this` is disposed. Objects added by `push` are
    * disposed in reverse order of being added.

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

@@ -22,6 +22,14 @@ export class RedactableError extends Error {
       .join("");
   }
 
+  public get fullMessageWithStack(): string {
+    if (!this.stack) {
+      return this.fullMessage;
+    }
+
+    return `${this.fullMessage}\n${this.stack}`;
+  }
+
   public get redactedMessage(): string {
     return this.strings
       .map((s, i) => s + (this.hasValue(i) ? this.getRedactedValue(i) : ""))

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

@@ -17,10 +17,14 @@ import {
 } from "../variant-analysis/shared/variant-analysis-filter-sort";
 import { ErrorLike } from "../common/errors";
 import { DataFlowPaths } from "../variant-analysis/shared/data-flow-paths";
-import { Method, Usage } from "../model-editor/method";
+import { Method } from "../model-editor/method";
 import { ModeledMethod } from "../model-editor/modeled-method";
-import { ModelEditorViewState } from "../model-editor/shared/view-state";
+import {
+  MethodModelingPanelViewState,
+  ModelEditorViewState,
+} from "../model-editor/shared/view-state";
 import { Mode } from "../model-editor/shared/mode";
+import { QueryLanguage } from "./query-language";
 
 /**
  * This module contains types and code that are shared between
@@ -51,6 +55,7 @@ export const RAW_RESULTS_LIMIT = 10000;
 export interface DatabaseInfo {
   name: string;
   databaseUri: string;
+  language?: QueryLanguage;
 }
 
 /** Arbitrary query metadata */
@@ -502,7 +507,7 @@ interface SetMethodsMessage {
 
 interface SetModeledMethodsMessage {
   t: "setModeledMethods";
-  methods: Record<string, ModeledMethod>;
+  methods: Record<string, ModeledMethod[]>;
 }
 
 interface SetModifiedMethodsMessage {
@@ -512,8 +517,7 @@ interface SetModifiedMethodsMessage {
 
 interface SetInProgressMethodsMessage {
   t: "setInProgressMethods";
-  packageName: string;
-  inProgressMethods: string[];
+  methods: string[];
 }
 
 interface SwitchModeMessage {
@@ -521,10 +525,9 @@ interface SwitchModeMessage {
   mode: Mode;
 }
 
-interface JumpToUsageMessage {
-  t: "jumpToUsage";
-  method: Method;
-  usage: Usage;
+interface JumpToMethodMessage {
+  t: "jumpToMethod";
+  methodSignature: string;
 }
 
 interface OpenDatabaseMessage {
@@ -541,8 +544,7 @@ interface RefreshMethods {
 
 interface SaveModeledMethods {
   t: "saveModeledMethods";
-  methods: Method[];
-  modeledMethods: Record<string, ModeledMethod>;
+  methodSignatures?: string[];
 }
 
 interface GenerateMethodMessage {
@@ -552,8 +554,7 @@ interface GenerateMethodMessage {
 interface GenerateMethodsFromLlmMessage {
   t: "generateMethodsFromLlm";
   packageName: string;
-  methods: Method[];
-  modeledMethods: Record<string, ModeledMethod>;
+  methodSignatures: string[];
 }
 
 interface StopGeneratingMethodsFromLlmMessage {
@@ -570,9 +571,25 @@ interface HideModeledMethodsMessage {
   hideModeledMethods: boolean;
 }
 
-interface SetModeledMethodMessage {
-  t: "setModeledMethod";
-  method: ModeledMethod;
+interface SetMultipleModeledMethodsMessage {
+  t: "setMultipleModeledMethods";
+  methodSignature: string;
+  modeledMethods: ModeledMethod[];
+}
+
+interface SetInModelingModeMessage {
+  t: "setInModelingMode";
+  inModelingMode: boolean;
+}
+
+interface SetInProgressMessage {
+  t: "setInProgress";
+  inProgress: boolean;
+}
+
+interface RevealMethodMessage {
+  t: "revealMethod";
+  methodSignature: string;
 }
 
 export type ToModelEditorMessage =
@@ -580,31 +597,47 @@ export type ToModelEditorMessage =
   | SetMethodsMessage
   | SetModeledMethodsMessage
   | SetModifiedMethodsMessage
-  | SetInProgressMethodsMessage;
+  | SetInProgressMethodsMessage
+  | RevealMethodMessage;
 
 export type FromModelEditorMessage =
-  | ViewLoadedMsg
+  | CommonFromViewMessages
   | SwitchModeMessage
   | RefreshMethods
   | OpenDatabaseMessage
   | OpenExtensionPackMessage
-  | JumpToUsageMessage
+  | JumpToMethodMessage
   | SaveModeledMethods
   | GenerateMethodMessage
   | GenerateMethodsFromLlmMessage
   | StopGeneratingMethodsFromLlmMessage
   | ModelDependencyMessage
   | HideModeledMethodsMessage
-  | SetModeledMethodMessage;
+  | SetMultipleModeledMethodsMessage;
+
+interface RevealInEditorMessage {
+  t: "revealInModelEditor";
+  method: Method;
+}
+
+interface StartModelingMessage {
+  t: "startModeling";
+}
 
 export type FromMethodModelingMessage =
-  | TelemetryMessage
-  | UnhandledErrorMessage
-  | SetModeledMethodMessage;
+  | CommonFromViewMessages
+  | SetMultipleModeledMethodsMessage
+  | RevealInEditorMessage
+  | StartModelingMessage;
+
+interface SetMethodModelingPanelViewStateMessage {
+  t: "setMethodModelingPanelViewState";
+  viewState: MethodModelingPanelViewState;
+}
 
 interface SetMethodMessage {
   t: "setMethod";
-  method: Method;
+  method: Method | undefined;
 }
 
 interface SetMethodModifiedMessage {
@@ -615,12 +648,16 @@ interface SetMethodModifiedMessage {
 interface SetSelectedMethodMessage {
   t: "setSelectedMethod";
   method: Method;
-  modeledMethod: ModeledMethod;
+  modeledMethods: ModeledMethod[];
   isModified: boolean;
+  isInProgress: boolean;
 }
 
 export type ToMethodModelingMessage =
+  | SetMethodModelingPanelViewStateMessage
   | SetMethodMessage
-  | SetModeledMethodMessage
+  | SetMultipleModeledMethodsMessage
   | SetMethodModifiedMessage
-  | SetSelectedMethodMessage;
+  | SetSelectedMethodMessage
+  | SetInModelingModeMessage
+  | SetInProgressMessage;

extensions/ql-vscode/src/common/logging/notifications.ts

@@ -112,5 +112,5 @@ export async function showAndLogExceptionWithTelemetry(
   options: ShowAndLogExceptionOptions = {},
 ): Promise<void> {
   telemetry?.sendError(error, options.extraTelemetryProperties);
-  return showAndLogErrorMessage(logger, error.fullMessage, options);
+  return showAndLogErrorMessage(logger, error.fullMessageWithStack, options);
 }

extensions/ql-vscode/src/common/mock-gh-api/gh-api-request.ts

@@ -17,7 +17,7 @@ export enum RequestKind {
   AutoModel = "autoModel",
 }
 
-interface BasicErorResponse {
+export interface BasicErrorResponse {
   message: string;
 }
 
@@ -27,7 +27,7 @@ interface GetRepoRequest {
   };
   response: {
     status: number;
-    body: Repository | BasicErorResponse | undefined;
+    body: Repository | BasicErrorResponse | undefined;
   };
 }
 
@@ -37,7 +37,7 @@ interface SubmitVariantAnalysisRequest {
   };
   response: {
     status: number;
-    body?: VariantAnalysis | BasicErorResponse;
+    body?: VariantAnalysis | BasicErrorResponse;
   };
 }
 
@@ -47,7 +47,7 @@ interface GetVariantAnalysisRequest {
   };
   response: {
     status: number;
-    body?: VariantAnalysis | BasicErorResponse;
+    body?: VariantAnalysis | BasicErrorResponse;
   };
 }
 
@@ -58,7 +58,7 @@ interface GetVariantAnalysisRepoRequest {
   };
   response: {
     status: number;
-    body?: VariantAnalysisRepoTask | BasicErorResponse;
+    body?: VariantAnalysisRepoTask | BasicErrorResponse;
   };
 }
 
@@ -74,6 +74,13 @@ export interface GetVariantAnalysisRepoResultRequest {
   };
 }
 
+export interface CodeSearchResponse {
+  total_count: number;
+  items: Array<{
+    repository: Repository;
+  }>;
+}
+
 interface CodeSearchRequest {
   request: {
     kind: RequestKind.CodeSearch;
@@ -81,16 +88,14 @@ interface CodeSearchRequest {
   };
   response: {
     status: number;
-    body?: {
-      total_count?: number;
-      items?: Array<{
-        repository: Repository;
-      }>;
-    };
-    message?: string;
+    body?: CodeSearchResponse | BasicErrorResponse;
   };
 }
 
+export interface AutoModelResponse {
+  models: string;
+}
+
 interface AutoModelRequest {
   request: {
     kind: RequestKind.AutoModel;
@@ -100,10 +105,7 @@ interface AutoModelRequest {
   };
   response: {
     status: number;
-    body?: {
-      models: string;
-    };
-    message?: string;
+    body?: AutoModelResponse | BasicErrorResponse;
   };
 }
 

extensions/ql-vscode/src/common/mock-gh-api/recorder.ts

@@ -1,30 +1,33 @@
 import { ensureDir, writeFile } from "fs-extra";
 import { join } from "path";
 
-import { MockedRequest } from "msw";
-import { SetupServer } from "msw/node";
-import { IsomorphicResponse } from "@mswjs/interceptors";
-
-import { Headers } from "headers-polyfill";
 import fetch from "node-fetch";
+import { SetupServer } from "msw/node";
 
 import { DisposableObject } from "../disposable-object";
+import { gzipDecode } from "../zlib";
 
 import {
+  AutoModelResponse,
+  BasicErrorResponse,
+  CodeSearchResponse,
   GetVariantAnalysisRepoResultRequest,
   GitHubApiRequest,
   RequestKind,
 } from "./gh-api-request";
+import {
+  VariantAnalysis,
+  VariantAnalysisRepoTask,
+} from "../../variant-analysis/gh-api/variant-analysis";
+import { Repository } from "../../variant-analysis/gh-api/repository";
 
 export class Recorder extends DisposableObject {
-  private readonly allRequests = new Map<string, MockedRequest>();
   private currentRecordedScenario: GitHubApiRequest[] = [];
 
   private _isRecording = false;
 
   constructor(private readonly server: SetupServer) {
     super();
-    this.onRequestStart = this.onRequestStart.bind(this);
     this.onResponseBypass = this.onResponseBypass.bind(this);
   }
 
@@ -45,7 +48,6 @@ export class Recorder extends DisposableObject {
 
     this.clear();
 
-    this.server.events.on("request:start", this.onRequestStart);
     this.server.events.on("response:bypass", this.onResponseBypass);
   }
 
@@ -56,13 +58,11 @@ export class Recorder extends DisposableObject {
 
     this._isRecording = false;
 
-    this.server.events.removeListener("request:start", this.onRequestStart);
     this.server.events.removeListener("response:bypass", this.onResponseBypass);
   }
 
   public clear() {
     this.currentRecordedScenario = [];
-    this.allRequests.clear();
   }
 
   public async save(scenariosPath: string, name: string): Promise<string> {
@@ -91,7 +91,7 @@ export class Recorder extends DisposableObject {
 
         let bodyFileLink = undefined;
         if (writtenRequest.response.body) {
-          await writeFile(bodyFilePath, writtenRequest.response.body || "");
+          await writeFile(bodyFilePath, writtenRequest.response.body);
           bodyFileLink = `file:${bodyFileName}`;
         }
 
@@ -112,33 +112,21 @@ export class Recorder extends DisposableObject {
     return scenarioDirectory;
   }
 
-  private onRequestStart(request: MockedRequest): void {
+  private async onResponseBypass({
+    response,
+    request,
+  }: {
+    response: Response;
+    request: Request;
+    requestId: string;
+  }): Promise<void> {
     if (request.headers.has("x-vscode-codeql-msw-bypass")) {
       return;
     }
 
-    this.allRequests.set(request.id, request);
-  }
-
-  private async onResponseBypass(
-    response: IsomorphicResponse,
-    requestId: string,
-  ): Promise<void> {
-    const request = this.allRequests.get(requestId);
-    this.allRequests.delete(requestId);
-    if (!request) {
-      return;
-    }
-
-    if (response.body === undefined) {
-      return;
-    }
-
     const gitHubApiRequest = await createGitHubApiRequest(
-      request.url.toString(),
-      response.status,
-      response.body,
-      response.headers,
+      request.url,
+      response,
     );
     if (!gitHubApiRequest) {
       return;
@@ -150,14 +138,14 @@ export class Recorder extends DisposableObject {
 
 async function createGitHubApiRequest(
   url: string,
-  status: number,
-  body: string,
-  headers: Headers,
+  response: Response,
 ): Promise<GitHubApiRequest | undefined> {
   if (!url) {
     return undefined;
   }
 
+  const status = response.status;
+
   if (url.match(/\/repos\/[a-zA-Z0-9-_.]+\/[a-zA-Z0-9-_.]+$/)) {
     return {
       request: {
@@ -165,7 +153,9 @@ async function createGitHubApiRequest(
       },
       response: {
         status,
-        body: JSON.parse(body),
+        body: await jsonResponseBody<
+          Repository | BasicErrorResponse | undefined
+        >(response),
       },
     };
   }
@@ -179,7 +169,9 @@ async function createGitHubApiRequest(
       },
       response: {
         status,
-        body: JSON.parse(body),
+        body: await jsonResponseBody<
+          VariantAnalysis | BasicErrorResponse | undefined
+        >(response),
       },
     };
   }
@@ -195,7 +187,9 @@ async function createGitHubApiRequest(
       },
       response: {
         status,
-        body: JSON.parse(body),
+        body: await jsonResponseBody<
+          VariantAnalysis | BasicErrorResponse | undefined
+        >(response),
       },
     };
   }
@@ -211,7 +205,9 @@ async function createGitHubApiRequest(
       },
       response: {
         status,
-        body: JSON.parse(body),
+        body: await jsonResponseBody<
+          VariantAnalysisRepoTask | BasicErrorResponse | undefined
+        >(response),
       },
     };
   }
@@ -238,9 +234,10 @@ async function createGitHubApiRequest(
         repositoryId: parseInt(repoDownloadMatch.groups.repositoryId, 10),
       },
       response: {
-        status,
+        status: response.status,
         body: responseBuffer,
-        contentType: headers.get("content-type") ?? "application/octet-stream",
+        contentType:
+          response.headers.get("content-type") ?? "application/octet-stream",
       },
     };
   }
@@ -254,7 +251,9 @@ async function createGitHubApiRequest(
       },
       response: {
         status,
-        body: JSON.parse(body),
+        body: await jsonResponseBody<
+          CodeSearchResponse | BasicErrorResponse | undefined
+        >(response),
       },
     };
   }
@@ -269,7 +268,9 @@ async function createGitHubApiRequest(
       },
       response: {
         status,
-        body: JSON.parse(body),
+        body: await jsonResponseBody<
+          BasicErrorResponse | AutoModelResponse | undefined
+        >(response),
       },
     };
   }
@@ -277,6 +278,26 @@ async function createGitHubApiRequest(
   return undefined;
 }
 
+async function responseBody(response: Response): Promise<Uint8Array> {
+  const body = await response.arrayBuffer();
+  const view = new Uint8Array(body);
+
+  if (view[0] === 0x1f && view[1] === 0x8b) {
+    // Response body is gzipped, so we need to un-gzip it.
+
+    return await gzipDecode(view);
+  } else {
+    return view;
+  }
+}
+
+async function jsonResponseBody<T>(response: Response): Promise<T> {
+  const body = await responseBody(response);
+  const text = new TextDecoder("utf-8").decode(body);
+
+  return JSON.parse(text);
+}
+
 function shouldWriteBodyToFile(
   request: GitHubApiRequest,
 ): request is GetVariantAnalysisRepoResultRequest {

extensions/ql-vscode/src/common/mock-gh-api/request-handlers.ts

@@ -1,6 +1,6 @@
 import { join } from "path";
 import { readdir, readJson, readFile } from "fs-extra";
-import { DefaultBodyType, MockedRequest, rest, RestHandler } from "msw";
+import { http, RequestHandler } from "msw";
 import {
   GitHubApiRequest,
   isAutoModelRequest,
@@ -14,7 +14,19 @@ import {
 
 const baseUrl = "https://api.github.com";
 
-type RequestHandler = RestHandler<MockedRequest<DefaultBodyType>>;
+const jsonResponse = <T>(
+  body: T,
+  init?: ResponseInit,
+  contentType = "application/json",
+): Response => {
+  return new Response(JSON.stringify(body), {
+    ...init,
+    headers: {
+      "Content-Type": contentType,
+      ...init?.headers,
+    },
+  });
+};
 
 export async function createRequestHandlers(
   scenarioDirPath: string,
@@ -82,11 +94,10 @@ function createGetRepoRequestHandler(
 
   const getRepoRequest = getRepoRequests[0];
 
-  return rest.get(`${baseUrl}/repos/:owner/:name`, (_req, res, ctx) => {
-    return res(
-      ctx.status(getRepoRequest.response.status),
-      ctx.json(getRepoRequest.response.body),
-    );
+  return http.get(`${baseUrl}/repos/:owner/:name`, () => {
+    return jsonResponse(getRepoRequest.response.body, {
+      status: getRepoRequest.response.status,
+    });
   });
 }
 
@@ -103,13 +114,12 @@ function createSubmitVariantAnalysisRequestHandler(
 
   const getRepoRequest = submitVariantAnalysisRequests[0];
 
-  return rest.post(
+  return http.post(
     `${baseUrl}/repositories/:controllerRepoId/code-scanning/codeql/variant-analyses`,
-    (_req, res, ctx) => {
-      return res(
-        ctx.status(getRepoRequest.response.status),
-        ctx.json(getRepoRequest.response.body),
-      );
+    () => {
+      return jsonResponse(getRepoRequest.response.body, {
+        status: getRepoRequest.response.status,
+      });
     },
   );
 }
@@ -125,9 +135,9 @@ function createGetVariantAnalysisRequestHandler(
   // During the lifetime of a variant analysis run, there are multiple requests
   // to get the variant analysis. We need to return different responses for each
   // request, so keep an index of the request and return the appropriate response.
-  return rest.get(
+  return http.get(
     `${baseUrl}/repositories/:controllerRepoId/code-scanning/codeql/variant-analyses/:variantAnalysisId`,
-    (_req, res, ctx) => {
+    () => {
       const request = getVariantAnalysisRequests[requestIndex];
 
       if (requestIndex < getVariantAnalysisRequests.length - 1) {
@@ -135,10 +145,9 @@ function createGetVariantAnalysisRequestHandler(
         requestIndex++;
       }
 
-      return res(
-        ctx.status(request.response.status),
-        ctx.json(request.response.body),
-      );
+      return jsonResponse(request.response.body, {
+        status: request.response.status,
+      });
     },
   );
 }
@@ -150,20 +159,19 @@ function createGetVariantAnalysisRepoRequestHandler(
     isGetVariantAnalysisRepoRequest,
   );
 
-  return rest.get(
+  return http.get(
     `${baseUrl}/repositories/:controllerRepoId/code-scanning/codeql/variant-analyses/:variantAnalysisId/repositories/:repoId`,
-    (req, res, ctx) => {
+    ({ request, params }) => {
       const scenarioRequest = getVariantAnalysisRepoRequests.find(
-        (r) => r.request.repositoryId.toString() === req.params.repoId,
+        (r) => r.request.repositoryId.toString() === params.repoId,
       );
       if (!scenarioRequest) {
-        throw Error(`No scenario request found for ${req.url}`);
+        throw Error(`No scenario request found for ${request.url}`);
       }
 
-      return res(
-        ctx.status(scenarioRequest.response.status),
-        ctx.json(scenarioRequest.response.body),
-      );
+      return jsonResponse(scenarioRequest.response.body, {
+        status: scenarioRequest.response.status,
+      });
     },
   );
 }
@@ -175,24 +183,25 @@ function createGetVariantAnalysisRepoResultRequestHandler(
     isGetVariantAnalysisRepoResultRequest,
   );
 
-  return rest.get(
+  return http.get(
     "https://objects-origin.githubusercontent.com/codeql-query-console/codeql-variant-analysis-repo-tasks/:variantAnalysisId/:repoId/*",
-    (req, res, ctx) => {
+    ({ request, params }) => {
       const scenarioRequest = getVariantAnalysisRepoResultRequests.find(
-        (r) => r.request.repositoryId.toString() === req.params.repoId,
+        (r) => r.request.repositoryId.toString() === params.repoId,
       );
       if (!scenarioRequest) {
-        throw Error(`No scenario request found for ${req.url}`);
+        throw Error(`No scenario request found for ${request.url}`);
       }
 
       if (scenarioRequest.response.body) {
-        return res(
-          ctx.status(scenarioRequest.response.status),
-          ctx.set("Content-Type", scenarioRequest.response.contentType),
-          ctx.body(scenarioRequest.response.body),
-        );
+        return new Response(scenarioRequest.response.body, {
+          status: scenarioRequest.response.status,
+          headers: {
+            "Content-Type": scenarioRequest.response.contentType,
+          },
+        });
       } else {
-        return res(ctx.status(scenarioRequest.response.status));
+        return new Response(null, { status: scenarioRequest.response.status });
       }
     },
   );
@@ -207,7 +216,7 @@ function createCodeSearchRequestHandler(
   // During a code search, there are multiple request to get pages of results. We
   // need to return different responses for each request, so keep an index of the
   // request and return the appropriate response.
-  return rest.get(`${baseUrl}/search/code?q=*`, (_req, res, ctx) => {
+  return http.get(`${baseUrl}/search/code`, () => {
     const request = codeSearchRequests[requestIndex];
 
     if (requestIndex < codeSearchRequests.length - 1) {
@@ -215,10 +224,9 @@ function createCodeSearchRequestHandler(
       requestIndex++;
     }
 
-    return res(
-      ctx.status(request.response.status),
-      ctx.json(request.response.body),
-    );
+    return jsonResponse(request.response.body, {
+      status: request.response.status,
+    });
   });
 }
 
@@ -231,9 +239,9 @@ function createAutoModelRequestHandler(
   // During automodeling there can be multiple API requests for each batch
   // of candidates we want to model. We need to return different responses for each request,
   // so keep an index of the request and return the appropriate response.
-  return rest.post(
+  return http.post(
     `${baseUrl}/repos/github/codeql/code-scanning/codeql/auto-model`,
-    (_req, res, ctx) => {
+    () => {
       const request = autoModelRequests[requestIndex];
 
       if (requestIndex < autoModelRequests.length - 1) {
@@ -241,10 +249,9 @@ function createAutoModelRequestHandler(
         requestIndex++;
       }
 
-      return res(
-        ctx.status(request.response.status),
-        ctx.json(request.response.body),
-      );
+      return jsonResponse(request.response.body, {
+        status: request.response.status,
+      });
     },
   );
 }

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

@@ -0,0 +1,6 @@
+/**
+ * Remove all readonly modifiers from a type.
+ */
+export type Mutable<T> = {
+  -readonly [P in keyof T]: T[P];
+};

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

@@ -0,0 +1,10 @@
+import * as Octokit from "@octokit/rest";
+import { retry } from "@octokit/plugin-retry";
+import fetch from "node-fetch";
+
+export const AppOctokit = Octokit.Octokit.defaults({
+  request: {
+    fetch,
+  },
+  retry,
+});

extensions/ql-vscode/src/common/query-language.ts

@@ -40,10 +40,7 @@ export const PACKS_BY_QUERY_LANGUAGE = {
   ],
   [QueryLanguage.Go]: ["codeql/go-queries"],
   [QueryLanguage.Java]: ["codeql/java-queries"],
-  [QueryLanguage.Javascript]: [
-    "codeql/javascript-queries",
-    "codeql/javascript-experimental-atm-queries",
-  ],
+  [QueryLanguage.Javascript]: ["codeql/javascript-queries"],
   [QueryLanguage.Python]: ["codeql/python-queries"],
   [QueryLanguage.Ruby]: ["codeql/ruby-queries"],
 };

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

@@ -0,0 +1,14 @@
+export type DeepReadonly<T> = T extends Array<infer R>
+  ? DeepReadonlyArray<R>
+  : // eslint-disable-next-line @typescript-eslint/ban-types
+  T extends Function
+  ? T
+  : T extends object
+  ? DeepReadonlyObject<T>
+  : T;
+
+interface DeepReadonlyArray<T> extends ReadonlyArray<DeepReadonly<T>> {}
+
+type DeepReadonlyObject<T> = {
+  readonly [P in keyof T]: DeepReadonly<T[P]>;
+};

extensions/ql-vscode/src/common/sarif-utils.ts

@@ -1,6 +1,7 @@
 import * as Sarif from "sarif";
 import type { HighlightedRegion } from "../variant-analysis/shared/analysis-result";
 import { ResolvableLocationValue } from "../common/bqrs-cli-types";
+import { isEmptyPath } from "./bqrs-utils";
 
 export interface SarifLink {
   dest: number;
@@ -111,6 +112,9 @@ export function parseSarifLocation(
     return { hint: "no artifact location" };
   if (physicalLocation.artifactLocation.uri === undefined)
     return { hint: "artifact location has no uri" };
+  if (isEmptyPath(physicalLocation.artifactLocation.uri)) {
+    return { hint: "artifact location has empty uri" };
+  }
 
   // This is not necessarily really an absolute uri; it could either be a
   // file uri or a relative uri.

extensions/ql-vscode/src/common/vscode/abstract-webview-view-provider.ts

@@ -0,0 +1,86 @@
+import * as vscode from "vscode";
+import { Uri, WebviewViewProvider } from "vscode";
+import { WebviewKind, WebviewMessage, getHtmlForWebview } from "./webview-html";
+import { Disposable } from "../disposable-object";
+import { App } from "../app";
+import { DeepReadonly } from "../readonly";
+
+export abstract class AbstractWebviewViewProvider<
+  ToMessage extends WebviewMessage,
+  FromMessage extends WebviewMessage,
+> implements WebviewViewProvider
+{
+  protected webviewView: vscode.WebviewView | undefined = undefined;
+  private disposables: Disposable[] = [];
+
+  constructor(
+    protected readonly app: App,
+    private readonly webviewKind: WebviewKind,
+  ) {}
+
+  /**
+   * This is called when a view first becomes visible. This may happen when the view is
+   * first loaded or when the user hides and then shows a view again.
+   */
+  public resolveWebviewView(
+    webviewView: vscode.WebviewView,
+    _context: vscode.WebviewViewResolveContext,
+    _token: vscode.CancellationToken,
+  ) {
+    webviewView.webview.options = {
+      enableScripts: true,
+      localResourceRoots: [Uri.file(this.app.extensionPath)],
+    };
+
+    const html = getHtmlForWebview(
+      this.app,
+      webviewView.webview,
+      this.webviewKind,
+      {
+        allowInlineStyles: true,
+        allowWasmEval: false,
+      },
+    );
+
+    webviewView.webview.html = html;
+
+    this.webviewView = webviewView;
+
+    webviewView.webview.onDidReceiveMessage(async (msg) => this.onMessage(msg));
+    webviewView.onDidDispose(() => this.dispose());
+  }
+
+  protected get isShowingView() {
+    return this.webviewView?.visible ?? false;
+  }
+
+  protected async postMessage(msg: DeepReadonly<ToMessage>): Promise<void> {
+    await this.webviewView?.webview.postMessage(msg);
+  }
+
+  protected dispose() {
+    while (this.disposables.length > 0) {
+      const disposable = this.disposables.pop()!;
+      disposable.dispose();
+    }
+
+    this.webviewView = undefined;
+  }
+
+  protected push<T extends Disposable>(obj: T): T {
+    if (obj !== undefined) {
+      this.disposables.push(obj);
+    }
+    return obj;
+  }
+
+  protected abstract onMessage(msg: FromMessage): Promise<void>;
+
+  /**
+   * This is called when a view first becomes visible. This may happen when the view is
+   * first loaded or when the user hides and then shows a view again.
+   */
+  protected onWebViewLoaded(): void {
+    // Do nothing by default.
+  }
+}

extensions/ql-vscode/src/common/vscode/abstract-webview.ts

@@ -12,6 +12,7 @@ import { App } from "../app";
 import { Disposable } from "../disposable-object";
 import { tmpDir } from "../../tmp-dir";
 import { getHtmlForWebview, WebviewMessage, WebviewKind } from "./webview-html";
+import { DeepReadonly } from "../readonly";
 
 export type WebviewPanelConfig = {
   viewId: string;
@@ -146,7 +147,7 @@ export abstract class AbstractWebview<
     this.panelLoadedCallBacks = [];
   }
 
-  protected async postMessage(msg: ToMessage): Promise<boolean> {
+  protected async postMessage(msg: DeepReadonly<ToMessage>): Promise<boolean> {
     const panel = await this.getPanel();
     return panel.webview.postMessage(msg);
   }

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

@@ -1,7 +1,7 @@
 import * as vscode from "vscode";
 import * as Octokit from "@octokit/rest";
-import { retry } from "@octokit/plugin-retry";
 import { Credentials } from "../authentication";
+import { AppOctokit } from "../octokit";
 
 export const GITHUB_AUTH_PROVIDER_ID = "github";
 
@@ -32,9 +32,8 @@ export class VSCodeCredentials implements Credentials {
 
     const accessToken = await this.getAccessToken();
 
-    return new Octokit.Octokit({
+    return new AppOctokit({
       auth: accessToken,
-      retry,
     });
   }
 

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

@@ -6,11 +6,7 @@ import {
   showAndLogExceptionWithTelemetry,
 } from "../logging";
 import { extLogger } from "../logging/vscode";
-import {
-  asError,
-  getErrorMessage,
-  getErrorStack,
-} from "../../common/helpers-pure";
+import { asError, getErrorMessage } from "../../common/helpers-pure";
 import { redactableError } from "../../common/errors";
 import { UserCancellationException } from "./progress";
 import { telemetryListener } from "./telemetry";
@@ -66,10 +62,7 @@ export function registerCommandWithErrorHandling(
         }
       } else {
         // Include the full stack in the error log only.
-        const errorStack = getErrorStack(e);
-        const fullMessage = errorStack
-          ? `${errorMessage.fullMessage}\n${errorStack}`
-          : errorMessage.fullMessage;
+        const fullMessage = errorMessage.fullMessageWithStack;
         void showAndLogExceptionWithTelemetry(logger, telemetry, errorMessage, {
           fullMessage,
           extraTelemetryProperties: {

extensions/ql-vscode/src/common/vscode/multi-cancellation-token.ts

@@ -0,0 +1,24 @@
+import { CancellationToken, Disposable } from "vscode";
+import { DisposableObject } from "../disposable-object";
+
+/**
+ * A cancellation token that cancels when any of its constituent
+ * cancellation tokens are cancelled.
+ */
+export class MultiCancellationToken implements CancellationToken {
+  private readonly tokens: CancellationToken[];
+
+  constructor(...tokens: CancellationToken[]) {
+    this.tokens = tokens;
+  }
+
+  get isCancellationRequested(): boolean {
+    return this.tokens.some((t) => t.isCancellationRequested);
+  }
+
+  onCancellationRequested<T>(listener: (e: T) => any): Disposable {
+    return new DisposableObject(
+      ...this.tokens.map((t) => t.onCancellationRequested(listener)),
+    );
+  }
+}

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

@@ -3,13 +3,13 @@ import {
   Extension,
   ExtensionContext,
   ConfigurationChangeEvent,
+  env,
 } from "vscode";
 import TelemetryReporter from "vscode-extension-telemetry";
 import {
   ConfigListener,
   CANARY_FEATURES,
   ENABLE_TELEMETRY,
-  GLOBAL_ENABLE_TELEMETRY,
   LOG_TELEMETRY,
   isIntegrationTestMode,
   isCanary,
@@ -59,8 +59,6 @@ export class ExtensionTelemetryListener
   extends ConfigListener
   implements AppTelemetry
 {
-  static relevantSettings = [ENABLE_TELEMETRY, CANARY_FEATURES];
-
   private reporter?: TelemetryReporter;
 
   private cliVersionStr = NOT_SET_CLI_VERSION;
@@ -72,6 +70,10 @@ export class ExtensionTelemetryListener
     private readonly ctx: ExtensionContext,
   ) {
     super();
+
+    env.onDidChangeTelemetryEnabled(async () => {
+      await this.initialize();
+    });
   }
 
   /**
@@ -91,10 +93,7 @@ export class ExtensionTelemetryListener
   async handleDidChangeConfiguration(
     e: ConfigurationChangeEvent,
   ): Promise<void> {
-    if (
-      e.affectsConfiguration("codeQL.telemetry.enableTelemetry") ||
-      e.affectsConfiguration("telemetry.enableTelemetry")
-    ) {
+    if (e.affectsConfiguration(ENABLE_TELEMETRY.qualifiedName)) {
       await this.initialize();
     }
 
@@ -102,7 +101,7 @@ export class ExtensionTelemetryListener
     // Re-request if codeQL.canary is being set to `true` and telemetry
     // is not currently enabled.
     if (
-      e.affectsConfiguration("codeQL.canary") &&
+      e.affectsConfiguration(CANARY_FEATURES.qualifiedName) &&
       CANARY_FEATURES.getValue() &&
       !ENABLE_TELEMETRY.getValue()
     ) {
@@ -212,7 +211,7 @@ export class ExtensionTelemetryListener
       properties.stack = error.stack;
     }
 
-    this.reporter.sendTelemetryEvent("error", properties, {});
+    this.reporter.sendTelemetryErrorEvent("error", properties, {});
   }
 
   /**
@@ -224,7 +223,7 @@ export class ExtensionTelemetryListener
       // if global telemetry is disabled, avoid showing the dialog or making any changes
       let result = undefined;
       if (
-        GLOBAL_ENABLE_TELEMETRY.getValue() &&
+        env.isTelemetryEnabled &&
         // Avoid showing the dialog if we are in integration test mode.
         !isIntegrationTestMode()
       ) {

extensions/ql-vscode/src/config.ts

@@ -72,15 +72,8 @@ export const VSCODE_SAVE_BEFORE_START_SETTING = new Setting(
 
 const ROOT_SETTING = new Setting("codeQL");
 
-// Global configuration
+// Telemetry configuration
 const TELEMETRY_SETTING = new Setting("telemetry", ROOT_SETTING);
-const AST_VIEWER_SETTING = new Setting("astViewer", ROOT_SETTING);
-const CONTEXTUAL_QUERIES_SETTINGS = new Setting(
-  "contextualQueries",
-  ROOT_SETTING,
-);
-const GLOBAL_TELEMETRY_SETTING = new Setting("telemetry");
-const LOG_INSIGHTS_SETTING = new Setting("logInsights", ROOT_SETTING);
 
 export const LOG_TELEMETRY = new Setting("logTelemetry", TELEMETRY_SETTING);
 export const ENABLE_TELEMETRY = new Setting(
@@ -88,11 +81,6 @@ export const ENABLE_TELEMETRY = new Setting(
   TELEMETRY_SETTING,
 );
 
-export const GLOBAL_ENABLE_TELEMETRY = new Setting(
-  "enableTelemetry",
-  GLOBAL_TELEMETRY_SETTING,
-);
-
 // Distribution configuration
 const DISTRIBUTION_SETTING = new Setting("cli", ROOT_SETTING);
 export const CUSTOM_CODEQL_PATH_SETTING = new Setting(
@@ -475,6 +463,7 @@ export function allowCanaryQueryServer() {
   return value === undefined ? true : !!value;
 }
 
+const LOG_INSIGHTS_SETTING = new Setting("logInsights", ROOT_SETTING);
 export const JOIN_ORDER_WARNING_THRESHOLD = new Setting(
   "joinOrderWarningThreshold",
   LOG_INSIGHTS_SETTING,
@@ -484,6 +473,7 @@ export function joinOrderWarningThreshold(): number {
   return JOIN_ORDER_WARNING_THRESHOLD.getValue<number>();
 }
 
+const AST_VIEWER_SETTING = new Setting("astViewer", ROOT_SETTING);
 /**
  * Hidden setting: Avoids caching in the AST viewer if the user is also a canary user.
  */
@@ -492,6 +482,10 @@ export const NO_CACHE_AST_VIEWER = new Setting(
   AST_VIEWER_SETTING,
 );
 
+const CONTEXTUAL_QUERIES_SETTINGS = new Setting(
+  "contextualQueries",
+  ROOT_SETTING,
+);
 /**
  * Hidden setting: Avoids caching in jump to def and find refs contextual queries if the user is also a canary user.
  */
@@ -708,18 +702,65 @@ 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 LLM_GENERATION_DEV_ENDPOINT = new Setting(
+  "llmGenerationDevEndpoint",
+  MODEL_SETTING,
+);
 const EXTENSIONS_DIRECTORY = new Setting("extensionsDirectory", MODEL_SETTING);
+const ENABLE_RUBY = new Setting("enableRuby", MODEL_SETTING);
 
-export function showFlowGeneration(): boolean {
-  return !!FLOW_GENERATION.getValue<boolean>();
+export interface ModelConfig {
+  flowGeneration: boolean;
+  llmGeneration: boolean;
+  getExtensionsDirectory(languageId: string): string | undefined;
+  showMultipleModels: boolean;
+  enableRuby: boolean;
 }
 
-export function showLlmGeneration(): boolean {
-  return !!LLM_GENERATION.getValue<boolean>();
-}
+export class ModelConfigListener extends ConfigListener implements ModelConfig {
+  protected handleDidChangeConfiguration(e: ConfigurationChangeEvent): void {
+    this.handleDidChangeConfigurationForRelevantSettings([MODEL_SETTING], e);
+  }
 
-export function getExtensionsDirectory(languageId: string): string | undefined {
-  return EXTENSIONS_DIRECTORY.getValue<string>({
-    languageId,
-  });
+  public get flowGeneration(): boolean {
+    return !!FLOW_GENERATION.getValue<boolean>();
+  }
+
+  public get llmGeneration(): boolean {
+    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>() || 5;
+  }
+
+  /**
+   * The URL of the endpoint to use for LLM generation. This should only be set
+   * if you want to test against a dev server.
+   */
+  public get llmGenerationDevEndpoint(): string | undefined {
+    return LLM_GENERATION_DEV_ENDPOINT.getValue<string | undefined>();
+  }
+
+  public getExtensionsDirectory(languageId: string): string | undefined {
+    return EXTENSIONS_DIRECTORY.getValue<string>({
+      languageId,
+    });
+  }
+
+  public get showMultipleModels(): boolean {
+    return isCanary();
+  }
+
+  public get enableRuby(): boolean {
+    return !!ENABLE_RUBY.getValue<boolean>();
+  }
 }

extensions/ql-vscode/src/databases/code-search-api.ts

@@ -1,22 +1,24 @@
-import { retry } from "@octokit/plugin-retry";
 import { throttling } from "@octokit/plugin-throttling";
 import { Octokit } from "@octokit/rest";
-import { Progress, CancellationToken } from "vscode";
+import { CancellationToken } from "vscode";
 import { Credentials } from "../common/authentication";
 import { BaseLogger } from "../common/logging";
+import { AppOctokit } from "../common/octokit";
+import {
+  ProgressCallback,
+  UserCancellationException,
+} from "../common/vscode/progress";
 
 export async function getCodeSearchRepositories(
   query: string,
-  progress: Progress<{
-    message?: string | undefined;
-    increment?: number | undefined;
-  }>,
+  progress: ProgressCallback,
   token: CancellationToken,
   credentials: Credentials,
   logger: BaseLogger,
 ): Promise<string[]> {
-  let nwos: string[] = [];
+  const nwos: string[] = [];
   const octokit = await provideOctokitWithThrottling(credentials, logger);
+  let i = 0;
 
   for await (const response of octokit.paginate.iterator(
     octokit.rest.search.code,
@@ -25,17 +27,19 @@ export async function getCodeSearchRepositories(
       per_page: 100,
     },
   )) {
+    i++;
     nwos.push(...response.data.map((item) => item.repository.full_name));
-    // calculate progress bar: 80% of the progress bar is used for the code search
-    const totalNumberOfRequests = Math.ceil(response.data.total_count / 100);
-    // Since we have a maximum of 1000 responses of the api, we can use a fixed increment whenever the totalNumberOfRequests would be greater than 10
-    const increment =
-      totalNumberOfRequests < 10 ? 80 / totalNumberOfRequests : 8;
-    progress.report({ increment });
+    const totalNumberOfResultPages = Math.ceil(response.data.total_count / 100);
+    const totalNumberOfRequests =
+      totalNumberOfResultPages > 10 ? 10 : totalNumberOfResultPages;
+    progress({
+      maxStep: totalNumberOfRequests,
+      step: i,
+      message: "Sending API requests to get Code Search results.",
+    });
 
     if (token.isCancellationRequested) {
-      nwos = [];
-      break;
+      throw new UserCancellationException("Code search cancelled.", true);
     }
   }
 
@@ -46,12 +50,11 @@ async function provideOctokitWithThrottling(
   credentials: Credentials,
   logger: BaseLogger,
 ): Promise<Octokit> {
-  const MyOctokit = Octokit.plugin(throttling);
+  const MyOctokit = AppOctokit.plugin(throttling);
   const auth = await credentials.getAccessToken();
 
   const octokit = new MyOctokit({
     auth,
-    retry,
     throttle: {
       onRateLimit: (retryAfter: number, options: any): boolean => {
         void logger.log(

extensions/ql-vscode/src/databases/database-fetcher.ts

@@ -14,7 +14,6 @@ import {
 } from "fs-extra";
 import { basename, join } from "path";
 import * as Octokit from "@octokit/rest";
-import { retry } from "@octokit/plugin-retry";
 
 import { DatabaseManager, DatabaseItem } from "./local-databases";
 import { tmpDir } from "../tmp-dir";
@@ -32,6 +31,8 @@ import { Credentials } from "../common/authentication";
 import { AppCommandManager } from "../common/commands";
 import { allowHttp } from "../config";
 import { showAndLogInformationMessage } from "../common/logging";
+import { AppOctokit } from "../common/octokit";
+import { getLanguageDisplayName } from "../common/query-language";
 
 /**
  * Prompts a user to fetch a database from a remote location. Database is assumed to be an archive file.
@@ -186,7 +187,7 @@ export async function downloadGitHubDatabase(
 
   const octokit = credentials
     ? await credentials.getOctokit()
-    : new Octokit.Octokit({ retry });
+    : new AppOctokit();
 
   const result = await convertGithubNwoToDatabaseUrl(
     nwo,
@@ -579,10 +580,23 @@ export async function promptForLanguage(
     return languages[0];
   }
 
-  return await window.showQuickPick(languages, {
+  const items = languages
+    .map((language) => ({
+      label: getLanguageDisplayName(language),
+      description: language,
+      language,
+    }))
+    .sort((a, b) => a.label.localeCompare(b.label));
+
+  const selectedItem = await window.showQuickPick(items, {
     placeHolder: "Select the database language to download:",
     ignoreFocusOut: true,
   });
+  if (!selectedItem) {
+    return undefined;
+  }
+
+  return selectedItem.language;
 }
 
 /**

extensions/ql-vscode/src/databases/local-databases-ui.ts

@@ -51,12 +51,14 @@ import {
   createMultiSelectionCommand,
   createSingleSelectionCommand,
 } from "../common/vscode/selection-commands";
-import { QueryLanguage, tryGetQueryLanguage } from "../common/query-language";
+import { tryGetQueryLanguage } from "../common/query-language";
 import { LanguageContextStore } from "../language-context-store";
 
 enum SortOrder {
   NameAsc = "NameAsc",
   NameDesc = "NameDesc",
+  LanguageAsc = "LanguageAsc",
+  LanguageDesc = "LanguageDesc",
   DateAddedAsc = "DateAddedAsc",
   DateAddedDesc = "DateAddedDesc",
 }
@@ -155,6 +157,10 @@ class DatabaseTreeDataProvider
             return db1.name.localeCompare(db2.name, env.language);
           case SortOrder.NameDesc:
             return db2.name.localeCompare(db1.name, env.language);
+          case SortOrder.LanguageAsc:
+            return db1.language.localeCompare(db2.language, env.language);
+          case SortOrder.LanguageDesc:
+            return db2.language.localeCompare(db1.language, env.language);
           case SortOrder.DateAddedAsc:
             return (db1.dateAdded || 0) - (db2.dateAdded || 0);
           case SortOrder.DateAddedDesc:
@@ -218,7 +224,7 @@ export class DatabaseUI extends DisposableObject {
   public constructor(
     private app: App,
     private databaseManager: DatabaseManager,
-    private languageContext: LanguageContextStore,
+    languageContext: LanguageContextStore,
     private readonly queryServer: QueryRunner | undefined,
     private readonly storagePath: string,
     readonly extensionPath: string,
@@ -252,6 +258,7 @@ export class DatabaseUI extends DisposableObject {
       "codeQL.upgradeCurrentDatabase":
         this.handleUpgradeCurrentDatabase.bind(this),
       "codeQL.clearCache": this.handleClearCache.bind(this),
+      "codeQL.trimCache": this.handleTrimCache.bind(this),
       "codeQLDatabases.chooseDatabaseFolder":
         this.handleChooseDatabaseFolder.bind(this),
       "codeQLDatabases.chooseDatabaseArchive":
@@ -263,61 +270,8 @@ export class DatabaseUI extends DisposableObject {
       "codeQLDatabases.setCurrentDatabase":
         this.handleMakeCurrentDatabase.bind(this),
       "codeQLDatabases.sortByName": this.handleSortByName.bind(this),
+      "codeQLDatabases.sortByLanguage": this.handleSortByLanguage.bind(this),
       "codeQLDatabases.sortByDateAdded": this.handleSortByDateAdded.bind(this),
-      "codeQLDatabases.displayAllLanguages":
-        this.handleClearLanguageFilter.bind(this),
-      "codeQLDatabases.displayCpp": this.handleChangeLanguageFilter.bind(
-        this,
-        QueryLanguage.Cpp,
-      ),
-      "codeQLDatabases.displayCsharp": this.handleChangeLanguageFilter.bind(
-        this,
-        QueryLanguage.CSharp,
-      ),
-      "codeQLDatabases.displayGo": this.handleChangeLanguageFilter.bind(
-        this,
-        QueryLanguage.Go,
-      ),
-      "codeQLDatabases.displayJava": this.handleChangeLanguageFilter.bind(
-        this,
-        QueryLanguage.Java,
-      ),
-      "codeQLDatabases.displayJavascript": this.handleChangeLanguageFilter.bind(
-        this,
-        QueryLanguage.Javascript,
-      ),
-      "codeQLDatabases.displayPython": this.handleChangeLanguageFilter.bind(
-        this,
-        QueryLanguage.Python,
-      ),
-      "codeQLDatabases.displayRuby": this.handleChangeLanguageFilter.bind(
-        this,
-        QueryLanguage.Ruby,
-      ),
-      "codeQLDatabases.displaySwift": this.handleChangeLanguageFilter.bind(
-        this,
-        QueryLanguage.Swift,
-      ),
-      "codeQLDatabases.displayAllLanguagesSelected":
-        this.handleClearLanguageFilter.bind(this),
-      "codeQLDatabases.displayCppSelected":
-        this.handleChangeLanguageFilter.bind(this, QueryLanguage.Cpp),
-      "codeQLDatabases.displayCsharpSelected":
-        this.handleChangeLanguageFilter.bind(this, QueryLanguage.CSharp),
-      "codeQLDatabases.displayGoSelected": this.handleChangeLanguageFilter.bind(
-        this,
-        QueryLanguage.Go,
-      ),
-      "codeQLDatabases.displayJavaSelected":
-        this.handleChangeLanguageFilter.bind(this, QueryLanguage.Java),
-      "codeQLDatabases.displayJavascriptSelected":
-        this.handleChangeLanguageFilter.bind(this, QueryLanguage.Javascript),
-      "codeQLDatabases.displayPythonSelected":
-        this.handleChangeLanguageFilter.bind(this, QueryLanguage.Python),
-      "codeQLDatabases.displayRubySelected":
-        this.handleChangeLanguageFilter.bind(this, QueryLanguage.Ruby),
-      "codeQLDatabases.displaySwiftSelected":
-        this.handleChangeLanguageFilter.bind(this, QueryLanguage.Swift),
       "codeQLDatabases.removeDatabase": createMultiSelectionCommand(
         this.handleRemoveDatabase.bind(this),
       ),
@@ -600,6 +554,14 @@ export class DatabaseUI extends DisposableObject {
     }
   }
 
+  private async handleSortByLanguage() {
+    if (this.treeDataProvider.sortOrder === SortOrder.LanguageAsc) {
+      this.treeDataProvider.sortOrder = SortOrder.LanguageDesc;
+    } else {
+      this.treeDataProvider.sortOrder = SortOrder.LanguageAsc;
+    }
+  }
+
   private async handleSortByDateAdded() {
     if (this.treeDataProvider.sortOrder === SortOrder.DateAddedAsc) {
       this.treeDataProvider.sortOrder = SortOrder.DateAddedDesc;
@@ -608,14 +570,6 @@ export class DatabaseUI extends DisposableObject {
     }
   }
 
-  private async handleClearLanguageFilter() {
-    await this.languageContext.clearLanguageContext();
-  }
-
-  private async handleChangeLanguageFilter(languageFilter: QueryLanguage) {
-    await this.languageContext.setLanguageContext(languageFilter);
-  }
-
   private async handleUpgradeCurrentDatabase(): Promise<void> {
     return withProgress(
       async (progress, token) => {
@@ -703,6 +657,25 @@ export class DatabaseUI extends DisposableObject {
     );
   }
 
+  private async handleTrimCache(): Promise<void> {
+    return withProgress(
+      async (_progress, token) => {
+        if (
+          this.queryServer !== undefined &&
+          this.databaseManager.currentDatabaseItem !== undefined
+        ) {
+          await this.queryServer.trimCacheInDatabase(
+            this.databaseManager.currentDatabaseItem,
+            token,
+          );
+        }
+      },
+      {
+        title: "Trimming cache",
+      },
+    );
+  }
+
   private async handleGetCurrentDatabase(): Promise<string | undefined> {
     const dbItem = await this.getDatabaseItemInternal(undefined);
     return dbItem?.databaseUri.fsPath;

extensions/ql-vscode/src/databases/local-databases/database-item-impl.ts

@@ -167,6 +167,15 @@ export class DatabaseItemImpl implements DatabaseItem {
     return encodeArchiveBasePath(sourceArchive.fsPath);
   }
 
+  /**
+   * Returns true if the database's source archive is in the workspace.
+   */
+  public hasSourceArchiveInExplorer(): boolean {
+    return (vscode.workspace.workspaceFolders || []).some((folder) =>
+      this.belongsToSourceArchiveExplorerUri(folder.uri),
+    );
+  }
+
   public verifyZippedSources(): string | undefined {
     const sourceArchive = this.sourceArchive;
     if (sourceArchive === undefined) {

extensions/ql-vscode/src/databases/local-databases/database-item.ts

@@ -56,6 +56,11 @@ export interface DatabaseItem {
    */
   getSourceArchiveExplorerUri(): vscode.Uri;
 
+  /**
+   * Returns true if the database's source archive is in the workspace.
+   */
+  hasSourceArchiveInExplorer(): boolean;
+
   /**
    * Holds if `uri` belongs to this database's source archive.
    */

extensions/ql-vscode/src/databases/local-databases/database-manager.ts

@@ -19,7 +19,10 @@ import {
   getFirstWorkspaceFolder,
   isFolderAlreadyInWorkspace,
 } from "../../common/vscode/workspace-folders";
-import { isQueryLanguage } from "../../common/query-language";
+import {
+  isQueryLanguage,
+  tryGetQueryLanguage,
+} from "../../common/query-language";
 import { existsSync } from "fs";
 import { QlPackGenerator } from "../../local-queries/qlpack-generator";
 import { asError, getErrorMessage } from "../../common/helpers-pure";
@@ -30,6 +33,7 @@ import { containsPath } from "../../common/files";
 import { DatabaseChangedEvent, DatabaseEventKind } from "./database-events";
 import { DatabaseResolver } from "./database-resolver";
 import { telemetryListener } from "../../common/vscode/telemetry";
+import { LanguageContextStore } from "../../language-context-store";
 
 /**
  * The name of the key in the workspaceState dictionary in which we
@@ -100,11 +104,25 @@ export class DatabaseManager extends DisposableObject {
     private readonly app: App,
     private readonly qs: QueryRunner,
     private readonly cli: cli.CodeQLCliServer,
+    private readonly languageContext: LanguageContextStore,
     public logger: Logger,
   ) {
     super();
 
     qs.onStart(this.reregisterDatabases.bind(this));
+
+    this.push(
+      this.languageContext.onLanguageContextChanged(async () => {
+        if (
+          this.currentDatabaseItem !== undefined &&
+          !this.languageContext.isSelectedLanguage(
+            tryGetQueryLanguage(this.currentDatabaseItem.language),
+          )
+        ) {
+          await this.setCurrentDatabaseItem(undefined);
+        }
+      }),
+    );
   }
 
   /**
@@ -230,8 +248,10 @@ export class DatabaseManager extends DisposableObject {
     const firstWorkspaceFolder = getFirstWorkspaceFolder();
     const folderName = `codeql-custom-queries-${databaseItem.language}`;
 
+    const qlpackStoragePath = join(firstWorkspaceFolder, folderName);
+
     if (
-      existsSync(join(firstWorkspaceFolder, folderName)) ||
+      existsSync(qlpackStoragePath) ||
       isFolderAlreadyInWorkspace(folderName)
     ) {
       return;
@@ -256,10 +276,10 @@ export class DatabaseManager extends DisposableObject {
 
     try {
       const qlPackGenerator = new QlPackGenerator(
-        folderName,
         databaseItem.language,
         this.cli,
-        firstWorkspaceFolder,
+        qlpackStoragePath,
+        qlpackStoragePath,
       );
       await qlPackGenerator.generate();
     } catch (e: unknown) {

extensions/ql-vscode/src/databases/ui/db-panel.ts

@@ -1,5 +1,4 @@
 import {
-  ProgressLocation,
   QuickPickItem,
   TreeView,
   TreeViewExpansionEvent,
@@ -7,7 +6,10 @@ import {
   window,
   workspace,
 } from "vscode";
-import { UserCancellationException } from "../../common/vscode/progress";
+import {
+  UserCancellationException,
+  withProgress,
+} from "../../common/vscode/progress";
 import {
   getNwoFromGitHubUrl,
   isValidGitHubNwo,
@@ -34,10 +36,7 @@ import { DatabasePanelCommands } from "../../common/commands";
 import { App } from "../../common/app";
 import { QueryLanguage } from "../../common/query-language";
 import { getCodeSearchRepositories } from "../code-search-api";
-import {
-  showAndLogErrorMessage,
-  showAndLogInformationMessage,
-} from "../../common/logging";
+import { showAndLogErrorMessage } from "../../common/logging";
 
 export interface RemoteDatabaseQuickPickItem extends QuickPickItem {
   remoteDatabaseKind: string;
@@ -409,15 +408,8 @@ export class DbPanel extends DisposableObject {
       return;
     }
 
-    await window.withProgress(
-      {
-        location: ProgressLocation.Notification,
-        title: "Searching for repositories... This might take a while",
-        cancellable: true,
-      },
+    await withProgress(
       async (progress, token) => {
-        progress.report({ increment: 10 });
-
         const repositories = await getCodeSearchRepositories(
           `${codeSearchQuery} ${languagePrompt}`,
           progress,
@@ -426,18 +418,22 @@ export class DbPanel extends DisposableObject {
           this.app.logger,
         );
 
-        token.onCancellationRequested(() => {
-          void showAndLogInformationMessage(
-            this.app.logger,
-            "Code search cancelled",
-          );
-          return;
+        if (token.isCancellationRequested) {
+          throw new UserCancellationException("Code search cancelled.", true);
+        }
+
+        progress({
+          maxStep: 12,
+          step: 12,
+          message: "Processing results...",
         });
 
-        progress.report({ increment: 10, message: "Processing results..." });
-
         await this.dbManager.addNewRemoteReposToList(repositories, listName);
       },
+      {
+        title: "Searching for repositories...",
+        cancellable: true,
+      },
     );
   }
 

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;

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 {

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,

extensions/ql-vscode/src/extension.ts

@@ -136,6 +136,7 @@ import { NewQueryRunner, QueryRunner, QueryServerClient } from "./query-server";
 import { QueriesModule } from "./queries-panel/queries-module";
 import { OpenReferencedFileCodeLensProvider } from "./local-queries/open-referenced-file-code-lens-provider";
 import { LanguageContextStore } from "./language-context-store";
+import { LanguageSelectionPanel } from "./language-selection-panel/language-selection-panel";
 
 /**
  * extension.ts
@@ -768,17 +769,28 @@ async function activateWithInstalledDistribution(
     fsWatcher.onDidDelete(clearPackCache);
   }
 
+  void extLogger.log("Initializing language context.");
+  const languageContext = new LanguageContextStore(app);
+
+  void extLogger.log("Initializing language selector.");
+  const languageSelectionPanel = new LanguageSelectionPanel(languageContext);
+  ctx.subscriptions.push(languageSelectionPanel);
+
   void extLogger.log("Initializing database manager.");
-  const dbm = new DatabaseManager(ctx, app, qs, cliServer, extLogger);
+  const dbm = new DatabaseManager(
+    ctx,
+    app,
+    qs,
+    cliServer,
+    languageContext,
+    extLogger,
+  );
 
   // Let this run async.
   void dbm.loadPersistedState();
 
   ctx.subscriptions.push(dbm);
 
-  void extLogger.log("Initializing language context.");
-  const languageContext = new LanguageContextStore(app);
-
   void extLogger.log("Initializing database panel.");
   const databaseUI = new DatabaseUI(
     app,
@@ -790,7 +802,11 @@ async function activateWithInstalledDistribution(
   );
   ctx.subscriptions.push(databaseUI);
 
-  QueriesModule.initialize(app, cliServer);
+  const queriesModule = QueriesModule.initialize(
+    app,
+    languageContext,
+    cliServer,
+  );
 
   void extLogger.log("Initializing evaluator log viewer.");
   const evalLogViewer = new EvalLogViewer();
@@ -871,6 +887,7 @@ async function activateWithInstalledDistribution(
     ctx,
     queryHistoryConfigurationListener,
     labelProvider,
+    languageContext,
     async (
       from: CompletedLocalQueryInfo,
       to: CompletedLocalQueryInfo,
@@ -924,9 +941,14 @@ async function activateWithInstalledDistribution(
     databaseUI,
     localQueryResultsView,
     queryStorageDir,
+    languageContext,
   );
   ctx.subscriptions.push(localQueries);
 
+  queriesModule.onDidChangeSelection((event) =>
+    localQueries.setSelectedQueryTreeViewItems(event.selection),
+  );
+
   void extLogger.log("Initializing debugger factory.");
   ctx.subscriptions.push(
     new QLDebugAdapterDescriptorFactory(queryStorageDir, qs, localQueries),
@@ -1015,6 +1037,7 @@ async function activateWithInstalledDistribution(
     ...getPackagingCommands({
       cliServer,
     }),
+    ...languageSelectionPanel.getCommands(),
     ...modelEditorModule.getCommands(),
     ...evalLogViewer.getCommands(),
     ...summaryLanguageSupport.getCommands(),
@@ -1163,13 +1186,13 @@ function addUnhandledRejectionListener() {
       const message = redactableError(
         asError(error),
       )`Unhandled error: ${getErrorMessage(error)}`;
+      const fullMessage = message.fullMessageWithStack;
+
       // Add a catch so that showAndLogExceptionWithTelemetry fails, we avoid
       // triggering "unhandledRejection" and avoid an infinite loop
-      showAndLogExceptionWithTelemetry(
-        extLogger,
-        telemetryListener,
-        message,
-      ).catch((telemetryError: unknown) => {
+      showAndLogExceptionWithTelemetry(extLogger, telemetryListener, message, {
+        fullMessage,
+      }).catch((telemetryError: unknown) => {
         void extLogger.log(
           `Failed to send error telemetry: ${getErrorMessage(telemetryError)}`,
         );

extensions/ql-vscode/src/language-context-store.ts

@@ -43,7 +43,36 @@ export class LanguageContextStore extends DisposableObject {
     );
   }
 
+  /**
+   * This returns true if the given language should be included.
+   *
+   * That means that either the given language is selected or the "All" option is selected.
+   *
+   * @param language a query language or undefined if the language is unknown.
+   */
   public shouldInclude(language: QueryLanguage | undefined): boolean {
     return this.languageFilter === "All" || this.languageFilter === language;
   }
+
+  /**
+   * This returns true if the given language is selected.
+   *
+   * If no language is given then it returns true if the "All" option is selected.
+   *
+   * @param language a query language or undefined.
+   */
+  public isSelectedLanguage(language: QueryLanguage | undefined): boolean {
+    return (
+      (this.languageFilter === "All" && language === undefined) ||
+      this.languageFilter === language
+    );
+  }
+
+  public get selectedLanguage(): QueryLanguage | undefined {
+    if (this.languageFilter === "All") {
+      return undefined;
+    }
+
+    return this.languageFilter;
+  }
 }

extensions/ql-vscode/src/language-selection-panel/language-selection-data-provider.ts

@@ -0,0 +1,91 @@
+import { DisposableObject } from "../common/disposable-object";
+import { LanguageContextStore } from "../language-context-store";
+import {
+  Event,
+  EventEmitter,
+  ThemeIcon,
+  TreeDataProvider,
+  TreeItem,
+} from "vscode";
+import {
+  QueryLanguage,
+  getLanguageDisplayName,
+} from "../common/query-language";
+
+const ALL_LANGUAGE_SELECTION_OPTIONS = [
+  undefined, // All languages
+  QueryLanguage.Cpp,
+  QueryLanguage.CSharp,
+  QueryLanguage.Go,
+  QueryLanguage.Java,
+  QueryLanguage.Javascript,
+  QueryLanguage.Python,
+  QueryLanguage.Ruby,
+  QueryLanguage.Swift,
+];
+
+// A tree view items consisting of of a language (or undefined for all languages)
+// and a boolean indicating whether it is selected or not.
+export class LanguageSelectionTreeViewItem extends TreeItem {
+  constructor(
+    public readonly language: QueryLanguage | undefined,
+    public readonly selected: boolean = false,
+  ) {
+    const label = language ? getLanguageDisplayName(language) : "All languages";
+    super(label);
+
+    this.iconPath = selected ? new ThemeIcon("check") : undefined;
+    this.contextValue = selected ? undefined : "canBeSelected";
+  }
+}
+
+export class LanguageSelectionTreeDataProvider
+  extends DisposableObject
+  implements TreeDataProvider<LanguageSelectionTreeViewItem>
+{
+  private treeItems: LanguageSelectionTreeViewItem[];
+  private readonly onDidChangeTreeDataEmitter = this.push(
+    new EventEmitter<void>(),
+  );
+
+  public constructor(private readonly languageContext: LanguageContextStore) {
+    super();
+
+    this.treeItems = this.createTree();
+
+    // If the language context changes, we need to update the tree.
+    this.push(
+      this.languageContext.onLanguageContextChanged(() => {
+        this.treeItems = this.createTree();
+        this.onDidChangeTreeDataEmitter.fire();
+      }),
+    );
+  }
+
+  public get onDidChangeTreeData(): Event<void> {
+    return this.onDidChangeTreeDataEmitter.event;
+  }
+
+  public getTreeItem(item: LanguageSelectionTreeViewItem): TreeItem {
+    return item;
+  }
+
+  public getChildren(
+    item?: LanguageSelectionTreeViewItem,
+  ): LanguageSelectionTreeViewItem[] {
+    if (!item) {
+      return this.treeItems;
+    } else {
+      return [];
+    }
+  }
+
+  private createTree(): LanguageSelectionTreeViewItem[] {
+    return ALL_LANGUAGE_SELECTION_OPTIONS.map((language) => {
+      return new LanguageSelectionTreeViewItem(
+        language,
+        this.languageContext.isSelectedLanguage(language),
+      );
+    });
+  }
+}

extensions/ql-vscode/src/language-selection-panel/language-selection-panel.ts

@@ -0,0 +1,41 @@
+import { DisposableObject } from "../common/disposable-object";
+import { window } from "vscode";
+import {
+  LanguageSelectionTreeDataProvider,
+  LanguageSelectionTreeViewItem,
+} from "./language-selection-data-provider";
+import { LanguageContextStore } from "../language-context-store";
+import { LanguageSelectionCommands } from "../common/commands";
+
+// This panel allows the selection of a single language, that will
+// then filter all other relevant views (e.g. db panel, query history).
+export class LanguageSelectionPanel extends DisposableObject {
+  constructor(private readonly languageContext: LanguageContextStore) {
+    super();
+
+    const dataProvider = new LanguageSelectionTreeDataProvider(languageContext);
+    this.push(dataProvider);
+
+    const treeView = window.createTreeView("codeQLLanguageSelection", {
+      treeDataProvider: dataProvider,
+    });
+    this.push(treeView);
+  }
+
+  public getCommands(): LanguageSelectionCommands {
+    return {
+      "codeQLLanguageSelection.setSelectedItem":
+        this.handleSetSelectedLanguage.bind(this),
+    };
+  }
+
+  private async handleSetSelectedLanguage(
+    item: LanguageSelectionTreeViewItem,
+  ): Promise<void> {
+    if (item.language) {
+      await this.languageContext.setLanguageContext(item.language);
+    } else {
+      await this.languageContext.clearLanguageContext();
+    }
+  }
+}

extensions/ql-vscode/src/language-support/contextual/query-resolver.ts

@@ -44,6 +44,7 @@ export async function runContextualQuery(
     false,
     getOnDiskWorkspaceFolders(),
     undefined,
+    {},
     queryStorageDir,
     undefined,
     templates,

extensions/ql-vscode/src/language-support/contextual/template-provider.ts

@@ -36,6 +36,7 @@ import {
 import { CoreCompletedQuery, QueryRunner } from "../../query-server";
 import { AstBuilder } from "../ast-viewer/ast-builder";
 import { qlpackOfDatabase } from "../../local-queries";
+import { MultiCancellationToken } from "../../common/vscode/multi-cancellation-token";
 
 /**
  * Runs templated CodeQL queries to find definitions in
@@ -43,6 +44,7 @@ import { qlpackOfDatabase } from "../../local-queries";
  * generalize this to other custom queries, e.g. showing dataflow to
  * or from a selected identifier.
  */
+
 export class TemplateQueryDefinitionProvider implements DefinitionProvider {
   private cache: CachedOperation<LocationLink[]>;
 
@@ -60,11 +62,11 @@ export class TemplateQueryDefinitionProvider implements DefinitionProvider {
   async provideDefinition(
     document: TextDocument,
     position: Position,
-    _token: CancellationToken,
+    token: CancellationToken,
   ): Promise<LocationLink[]> {
     const fileLinks = this.shouldUseCache()
-      ? await this.cache.get(document.uri.toString())
-      : await this.getDefinitions(document.uri.toString());
+      ? await this.cache.get(document.uri.toString(), token)
+      : await this.getDefinitions(document.uri.toString(), token);
 
     const locLinks: LocationLink[] = [];
     for (const link of fileLinks) {
@@ -79,9 +81,13 @@ export class TemplateQueryDefinitionProvider implements DefinitionProvider {
     return !(isCanary() && NO_CACHE_CONTEXTUAL_QUERIES.getValue<boolean>());
   }
 
-  private async getDefinitions(uriString: string): Promise<LocationLink[]> {
+  private async getDefinitions(
+    uriString: string,
+    token: CancellationToken,
+  ): Promise<LocationLink[]> {
     return withProgress(
-      async (progress, token) => {
+      async (progress, tokenInner) => {
+        const multiToken = new MultiCancellationToken(token, tokenInner);
         return getLocationsForUriString(
           this.cli,
           this.qs,
@@ -90,7 +96,7 @@ export class TemplateQueryDefinitionProvider implements DefinitionProvider {
           KeyType.DefinitionQuery,
           this.queryStorageDir,
           progress,
-          token,
+          multiToken,
           (src, _dest) => src === uriString,
         );
       },
@@ -126,11 +132,11 @@ export class TemplateQueryReferenceProvider implements ReferenceProvider {
     document: TextDocument,
     position: Position,
     _context: ReferenceContext,
-    _token: CancellationToken,
+    token: CancellationToken,
   ): Promise<Location[]> {
     const fileLinks = this.shouldUseCache()
-      ? await this.cache.get(document.uri.toString())
-      : await this.getReferences(document.uri.toString());
+      ? await this.cache.get(document.uri.toString(), token)
+      : await this.getReferences(document.uri.toString(), token);
 
     const locLinks: Location[] = [];
     for (const link of fileLinks) {
@@ -148,9 +154,14 @@ export class TemplateQueryReferenceProvider implements ReferenceProvider {
     return !(isCanary() && NO_CACHE_CONTEXTUAL_QUERIES.getValue<boolean>());
   }
 
-  private async getReferences(uriString: string): Promise<FullLocationLink[]> {
+  private async getReferences(
+    uriString: string,
+    token: CancellationToken,
+  ): Promise<FullLocationLink[]> {
     return withProgress(
-      async (progress, token) => {
+      async (progress, tokenInner) => {
+        const multiToken = new MultiCancellationToken(token, tokenInner);
+
         return getLocationsForUriString(
           this.cli,
           this.qs,
@@ -159,7 +170,7 @@ export class TemplateQueryReferenceProvider implements ReferenceProvider {
           KeyType.DefinitionQuery,
           this.queryStorageDir,
           progress,
-          token,
+          multiToken,
           (src, _dest) => src === uriString,
         );
       },

extensions/ql-vscode/src/local-queries/index.ts

@@ -1,5 +1,6 @@
 export * from "./local-queries";
 export * from "./local-query-run";
+export * from "./query-constraints";
 export * from "./query-resolver";
 export * from "./quick-eval-code-lens-provider";
 export * from "./quick-query";

extensions/ql-vscode/src/local-queries/local-queries.ts

@@ -49,6 +49,8 @@ import { LocalQueryRun } from "./local-query-run";
 import { createMultiSelectionCommand } from "../common/vscode/selection-commands";
 import { findLanguage } from "../codeql-cli/query-language";
 import type { QueryTreeViewItem } from "../queries-panel/query-tree-view-item";
+import { tryGetQueryLanguage } from "../common/query-language";
+import { LanguageContextStore } from "../language-context-store";
 
 interface DatabaseQuickPickItem extends QuickPickItem {
   databaseItem: DatabaseItem;
@@ -61,6 +63,8 @@ export enum QuickEvalType {
 }
 
 export class LocalQueries extends DisposableObject {
+  private selectedQueryTreeViewItems: readonly QueryTreeViewItem[] = [];
+
   public constructor(
     private readonly app: App,
     private readonly queryRunner: QueryRunner,
@@ -70,10 +74,17 @@ export class LocalQueries extends DisposableObject {
     private readonly databaseUI: DatabaseUI,
     private readonly localQueryResultsView: ResultsView,
     private readonly queryStorageDir: string,
+    private readonly languageContextStore: LanguageContextStore,
   ) {
     super();
   }
 
+  public setSelectedQueryTreeViewItems(
+    selection: readonly QueryTreeViewItem[],
+  ) {
+    this.selectedQueryTreeViewItems = selection;
+  }
+
   public getCommands(): LocalQueryCommands {
     return {
       "codeQL.runQuery": this.runQuery.bind(this),
@@ -322,13 +333,16 @@ export class LocalQueries extends DisposableObject {
         const credentials = isCanary() ? this.app.credentials : undefined;
         const contextStoragePath =
           this.app.workspaceStoragePath || this.app.globalStoragePath;
+        const language = this.languageContextStore.selectedLanguage;
         const skeletonQueryWizard = new SkeletonQueryWizard(
           this.cliServer,
           progress,
           credentials,
-          this.app.logger,
+          this.app,
           this.databaseManager,
           contextStoragePath,
+          this.selectedQueryTreeViewItems,
+          language,
         );
         await skeletonQueryWizard.execute();
       },
@@ -361,10 +375,15 @@ export class LocalQueries extends DisposableObject {
       );
     }
 
-    const initialInfo = await createInitialQueryInfo(selectedQuery, {
-      databaseUri: dbItem.databaseUri.toString(),
-      name: dbItem.name,
-    });
+    const initialInfo = await createInitialQueryInfo(
+      selectedQuery,
+      {
+        databaseUri: dbItem.databaseUri.toString(),
+        name: dbItem.name,
+        language: tryGetQueryLanguage(dbItem.language),
+      },
+      outputDir,
+    );
 
     // When cancellation is requested from the query history view, we just stop the debug session.
     const queryInfo = new LocalQueryInfo(initialInfo, tokenSource);
@@ -454,6 +473,7 @@ export class LocalQueries extends DisposableObject {
       true,
       additionalPacks,
       extensionPacks,
+      {},
       this.queryStorageDir,
       undefined,
       templates,

extensions/ql-vscode/src/local-queries/local-query-run.ts

@@ -97,6 +97,15 @@ export class LocalQueryRun {
    * Updates the UI in the case where query evaluation throws an exception.
    */
   public async fail(err: Error): Promise<void> {
+    const evalLogPaths = await this.summarizeEvalLog(
+      QueryResultType.OTHER_ERROR,
+      this.outputDir,
+      this.logger,
+    );
+    if (evalLogPaths !== undefined) {
+      this.queryInfo.setEvaluatorLogPaths(evalLogPaths);
+    }
+
     err.message = `Error running query: ${err.message}`;
     this.queryInfo.failureReason = err.message;
     await this.queryHistoryManager.refreshTreeView();

extensions/ql-vscode/src/local-queries/qlpack-generator.ts

@@ -1,35 +1,36 @@
 import { mkdir, writeFile } from "fs-extra";
 import { dump } from "js-yaml";
-import { join } from "path";
+import { dirname, join } from "path";
 import { Uri } from "vscode";
 import { CodeQLCliServer } from "../codeql-cli/cli";
 import { QueryLanguage } from "../common/query-language";
+import { getOnDiskWorkspaceFolders } from "../common/vscode/workspace-folders";
+import { basename } from "../common/path";
 
 export class QlPackGenerator {
-  private readonly qlpackName: string;
+  private qlpackName: string | undefined;
   private readonly qlpackVersion: string;
   private readonly header: string;
   private readonly qlpackFileName: string;
   private readonly folderUri: Uri;
 
   constructor(
-    private readonly folderName: string,
     private readonly queryLanguage: QueryLanguage,
     private readonly cliServer: CodeQLCliServer,
-    private readonly storagePath: string | undefined,
+    private readonly storagePath: string,
+    private readonly queryStoragePath: string,
+    private readonly includeFolderNameInQlpackName: boolean = false,
   ) {
-    if (this.storagePath === undefined) {
-      throw new Error("Workspace storage path is undefined");
-    }
-    this.qlpackName = `getting-started/codeql-extra-queries-${this.queryLanguage}`;
     this.qlpackVersion = "1.0.0";
     this.header = "# This is an automatically generated file.\n\n";
 
     this.qlpackFileName = "codeql-pack.yml";
-    this.folderUri = Uri.file(join(this.storagePath, this.folderName));
+    this.folderUri = Uri.file(this.storagePath);
   }
 
   public async generate() {
+    this.qlpackName = await this.determineQlpackName();
+
     // create QL pack folder and add to workspace
     await this.createWorkspaceFolder();
 
@@ -43,6 +44,37 @@ export class QlPackGenerator {
     await this.createCodeqlPackLockYaml();
   }
 
+  private async determineQlpackName(): Promise<string> {
+    let qlpackBaseName = `getting-started/codeql-extra-queries-${this.queryLanguage}`;
+    if (this.includeFolderNameInQlpackName) {
+      const folderBasename = basename(dirname(this.folderUri.fsPath));
+      if (
+        folderBasename.includes("codeql") ||
+        folderBasename.includes("queries")
+      ) {
+        // If the user has already included "codeql" or "queries" in the folder name, don't include it twice
+        qlpackBaseName = `getting-started/${folderBasename}-${this.queryLanguage}`;
+      } else {
+        qlpackBaseName = `getting-started/codeql-extra-queries-${folderBasename}-${this.queryLanguage}`;
+      }
+    }
+
+    const existingQlPacks = await this.cliServer.resolveQlpacks(
+      getOnDiskWorkspaceFolders(),
+    );
+    const existingQlPackNames = Object.keys(existingQlPacks);
+
+    let qlpackName = qlpackBaseName;
+    let i = 0;
+    while (existingQlPackNames.includes(qlpackName)) {
+      i++;
+
+      qlpackName = `${qlpackBaseName}-${i}`;
+    }
+
+    return qlpackName;
+  }
+
   private async createWorkspaceFolder() {
     await mkdir(this.folderUri.fsPath);
   }
@@ -60,7 +92,7 @@ export class QlPackGenerator {
   }
 
   public async createExampleQlFile(fileName = "example.ql") {
-    const exampleQlFilePath = join(this.folderUri.fsPath, fileName);
+    const exampleQlFilePath = join(this.queryStoragePath, fileName);
 
     const exampleQl = `
 /**

extensions/ql-vscode/src/local-queries/query-constraints.ts

@@ -0,0 +1,7 @@
+export interface QueryConstraints {
+  kind?: string;
+  "tags contain"?: string[];
+  "tags contain all"?: string[];
+  "query filename"?: string;
+  "query path"?: string;
+}

extensions/ql-vscode/src/local-queries/query-resolver.ts

@@ -14,6 +14,7 @@ import { showAndLogExceptionWithTelemetry } from "../common/logging";
 import { extLogger } from "../common/logging/vscode";
 import { telemetryListener } from "../common/vscode/telemetry";
 import { SuiteInstruction } from "../packaging/suite-instruction";
+import { QueryConstraints } from "./query-constraints";
 
 export async function qlpackOfDatabase(
   cli: Pick<CodeQLCliServer, "resolveQlpacks">,
@@ -27,12 +28,6 @@ export async function qlpackOfDatabase(
   return await getQlPackForDbscheme(cli, dbscheme);
 }
 
-export interface QueryConstraints {
-  kind?: string;
-  "tags contain"?: string[];
-  "tags contain all"?: string[];
-}
-
 /**
  * Finds the queries with the specified kind and tags in a list of CodeQL packs.
  *
@@ -132,6 +127,14 @@ export async function resolveQueries(
       `tagged all of "${constraints["tags contain all"].join(" ")}"`,
     );
   }
+  if (constraints["query filename"] !== undefined) {
+    humanConstraints.push(
+      `with query filename "${constraints["query filename"]}"`,
+    );
+  }
+  if (constraints["query path"] !== undefined) {
+    humanConstraints.push(`with query path "${constraints["query path"]}"`);
+  }
 
   const joinedPacksToSearch = packsToSearch.join(", ");
   const error = redactableError`No ${name} queries (${humanConstraints.join(

extensions/ql-vscode/src/local-queries/run-query.ts

@@ -41,6 +41,7 @@ export async function runQuery({
     false,
     additionalPacks,
     extensionPacks,
+    {},
     queryStorageDir,
     undefined,
     undefined,

extensions/ql-vscode/src/local-queries/skeleton-query-wizard.ts

@@ -1,19 +1,23 @@
-import { join } from "path";
-import { Uri, workspace, window as Window } from "vscode";
+import { dirname, join } from "path";
+import { Uri, window, window as Window, workspace } from "vscode";
 import { CodeQLCliServer } from "../codeql-cli/cli";
-import { BaseLogger } from "../common/logging";
+import { showAndLogExceptionWithTelemetry } from "../common/logging";
 import { Credentials } from "../common/authentication";
-import { QueryLanguage } from "../common/query-language";
+import {
+  getLanguageDisplayName,
+  QueryLanguage,
+} from "../common/query-language";
 import {
   getFirstWorkspaceFolder,
-  isFolderAlreadyInWorkspace,
+  getOnDiskWorkspaceFolders,
 } from "../common/vscode/workspace-folders";
-import { getErrorMessage } from "../common/helpers-pure";
+import { asError, getErrorMessage } from "../common/helpers-pure";
 import { QlPackGenerator } from "./qlpack-generator";
 import { DatabaseItem, DatabaseManager } from "../databases/local-databases";
 import {
   ProgressCallback,
   UserCancellationException,
+  withProgress,
 } from "../common/vscode/progress";
 import {
   askForGitHubRepo,
@@ -24,8 +28,16 @@ import {
   isCodespacesTemplate,
   setQlPackLocation,
 } from "../config";
-import { existsSync } from "fs-extra";
+import { lstat, pathExists, readFile } from "fs-extra";
 import { askForLanguage } from "../codeql-cli/query-language";
+import { showInformationMessageWithAction } from "../common/vscode/dialog";
+import { redactableError } from "../common/errors";
+import { App } from "../common/app";
+import { QueryTreeViewItem } from "../queries-panel/query-tree-view-item";
+import { containsPath, pathsEqual } from "../common/files";
+import { getQlPackPath } from "../common/ql";
+import { load } from "js-yaml";
+import { QlPackFile } from "../packaging/qlpack-file";
 
 type QueryLanguagesToDatabaseMap = Record<string, string>;
 
@@ -41,73 +53,139 @@ export const QUERY_LANGUAGE_TO_DATABASE_REPO: QueryLanguagesToDatabaseMap = {
 };
 
 export class SkeletonQueryWizard {
-  private language: QueryLanguage | undefined;
   private fileName = "example.ql";
   private qlPackStoragePath: string | undefined;
+  private queryStoragePath: string | undefined;
+  private downloadPromise: Promise<void> | undefined;
 
   constructor(
     private readonly cliServer: CodeQLCliServer,
     private readonly progress: ProgressCallback,
     private readonly credentials: Credentials | undefined,
-    private readonly logger: BaseLogger,
+    private readonly app: App,
     private readonly databaseManager: DatabaseManager,
     private readonly databaseStoragePath: string | undefined,
+    private readonly selectedItems: readonly QueryTreeViewItem[],
+    private language: QueryLanguage | undefined = undefined,
   ) {}
 
-  private get folderName() {
-    return `codeql-custom-queries-${this.language}`;
+  /**
+   * Wait for the download process to complete by waiting for the user to select
+   * either "Download database" or closing the dialog. This is used for testing.
+   */
+  public async waitForDownload() {
+    if (this.downloadPromise) {
+      await this.downloadPromise;
+    }
   }
 
   public async execute() {
-    // show quick pick to choose language
-    this.language = await this.chooseLanguage();
+    // First try detecting the language based on the existing qlpacks.
+    // This will override the selected language if there is an existing query pack.
+    const detectedLanguage = await this.detectLanguage();
+    if (detectedLanguage) {
+      this.language = detectedLanguage;
+    }
+
+    // If no existing qlpack was found, we need to ask the user for the language
+    if (!this.language) {
+      // show quick pick to choose language
+      this.language = await this.chooseLanguage();
+    }
+
     if (!this.language) {
       return;
     }
 
-    this.qlPackStoragePath = await this.determineStoragePath();
+    let createSkeletonQueryPack: boolean = false;
 
-    const skeletonPackAlreadyExists =
-      existsSync(join(this.qlPackStoragePath, this.folderName)) ||
-      isFolderAlreadyInWorkspace(this.folderName);
+    if (!this.qlPackStoragePath) {
+      // This means no existing qlpack was detected in the selected folder, so we need
+      // to find a new location to store the qlpack. This new location could potentially
+      // already exist.
+      const storagePath = await this.determineStoragePath();
+      this.qlPackStoragePath = join(
+        storagePath,
+        `codeql-custom-queries-${this.language}`,
+      );
 
-    if (skeletonPackAlreadyExists) {
-      // just create a new example query file in skeleton QL pack
-      await this.createExampleFile();
+      // Try to detect if there is already a qlpack in this location. We will assume that
+      // the user hasn't changed the language of the qlpack.
+      const qlPackPath = await getQlPackPath(this.qlPackStoragePath);
+
+      // If we are creating or using a qlpack in the user's selected folder, we will also
+      // create the query in that folder
+      this.queryStoragePath = this.qlPackStoragePath;
+
+      createSkeletonQueryPack = qlPackPath === undefined;
     } else {
+      // A query pack was detected in the selected folder or one of its ancestors, so we
+      // directly use the selected folder as the storage path for the query.
+      this.queryStoragePath = await this.determineStoragePathFromSelection();
+    }
+
+    if (createSkeletonQueryPack) {
       // generate a new skeleton QL pack with query file
       await this.createQlPack();
+    } else {
+      // just create a new example query file in skeleton QL pack
+      await this.createExampleFile();
+    }
+
+    // open the query file
+    try {
+      await this.openExampleFile();
+    } catch (e: unknown) {
+      void this.app.logger.log(
+        `Could not open example query file: ${getErrorMessage(e)}`,
+      );
     }
 
     // select existing database for language or download a new one
     await this.selectOrDownloadDatabase();
-
-    // open a query file
-
-    try {
-      await this.openExampleFile();
-    } catch (e: unknown) {
-      void this.logger.log(
-        `Could not open example query file: ${getErrorMessage(e)}`,
-      );
-    }
   }
 
   private async openExampleFile() {
-    if (this.folderName === undefined || this.qlPackStoragePath === undefined) {
+    if (this.queryStoragePath === undefined) {
       throw new Error("Path to folder is undefined");
     }
 
-    const queryFileUri = Uri.file(
-      join(this.qlPackStoragePath, this.folderName, this.fileName),
-    );
+    const queryFileUri = Uri.file(join(this.queryStoragePath, this.fileName));
 
     void workspace.openTextDocument(queryFileUri).then((doc) => {
-      void Window.showTextDocument(doc);
+      void Window.showTextDocument(doc, {
+        preview: false,
+      });
     });
   }
 
-  public async determineStoragePath() {
+  public async determineStoragePath(): Promise<string> {
+    if (this.selectedItems.length === 0) {
+      return this.determineRootStoragePath();
+    }
+
+    return this.determineStoragePathFromSelection();
+  }
+
+  private async determineStoragePathFromSelection(): Promise<string> {
+    // Just like VS Code's "New File" command, if the user has selected multiple files/folders in the queries panel,
+    // we will create the new file in the same folder as the first selected item.
+    // See https://github.com/microsoft/vscode/blob/a8b7239d0311d4915b57c837972baf4b01394491/src/vs/workbench/contrib/files/browser/fileActions.ts#L893-L900
+    const selectedItem = this.selectedItems[0];
+
+    const path = selectedItem.path;
+
+    // We use stat to protect against outdated query tree items
+    const fileStat = await lstat(path);
+
+    if (fileStat.isDirectory()) {
+      return path;
+    }
+
+    return dirname(path);
+  }
+
+  public async determineRootStoragePath() {
     const firstStorageFolder = getFirstWorkspaceFolder();
 
     if (isCodespacesTemplate()) {
@@ -116,7 +194,7 @@ export class SkeletonQueryWizard {
 
     let storageFolder = getQlPackLocation();
 
-    if (storageFolder === undefined || !existsSync(storageFolder)) {
+    if (storageFolder === undefined || !(await pathExists(storageFolder))) {
       storageFolder = await Window.showInputBox({
         title:
           "Please choose a folder in which to create your new query pack. You can change this in the extension settings.",
@@ -129,7 +207,7 @@ export class SkeletonQueryWizard {
       throw new UserCancellationException("No storage folder entered.");
     }
 
-    if (!existsSync(storageFolder)) {
+    if (!(await pathExists(storageFolder))) {
       throw new UserCancellationException(
         "Invalid folder. Must be a folder that already exists.",
       );
@@ -139,6 +217,62 @@ export class SkeletonQueryWizard {
     return storageFolder;
   }
 
+  private async detectLanguage(): Promise<QueryLanguage | undefined> {
+    if (this.selectedItems.length < 1) {
+      return undefined;
+    }
+
+    this.progress({
+      message: "Resolving existing query packs",
+      step: 1,
+      maxStep: 3,
+    });
+
+    const storagePath = await this.determineStoragePathFromSelection();
+
+    const queryPacks = await this.cliServer.resolveQlpacks(
+      getOnDiskWorkspaceFolders(),
+      false,
+      "query",
+    );
+
+    const matchingQueryPacks = Object.values(queryPacks)
+      .map((paths) => paths.find((path) => containsPath(path, storagePath)))
+      .filter((path): path is string => path !== undefined)
+      // Find the longest matching path
+      .sort((a, b) => b.length - a.length);
+
+    if (matchingQueryPacks.length === 0) {
+      return undefined;
+    }
+
+    const matchingQueryPackPath = matchingQueryPacks[0];
+
+    const qlPackPath = await getQlPackPath(matchingQueryPackPath);
+    if (!qlPackPath) {
+      return undefined;
+    }
+
+    const qlPack = load(await readFile(qlPackPath, "utf8")) as
+      | QlPackFile
+      | undefined;
+    const dependencies = qlPack?.dependencies;
+    if (!dependencies || typeof dependencies !== "object") {
+      return;
+    }
+
+    const matchingLanguages = Object.values(QueryLanguage).filter(
+      (language) => `codeql/${language}-all` in dependencies,
+    );
+    if (matchingLanguages.length !== 1) {
+      return undefined;
+    }
+
+    this.qlPackStoragePath = matchingQueryPackPath;
+
+    return matchingLanguages[0];
+  }
+
   private async chooseLanguage() {
     this.progress({
       message: "Choose language",
@@ -150,13 +284,6 @@ export class SkeletonQueryWizard {
   }
 
   private async createQlPack() {
-    if (this.folderName === undefined) {
-      throw new Error("Folder name is undefined");
-    }
-    if (this.language === undefined) {
-      throw new Error("Language is undefined");
-    }
-
     this.progress({
       message: "Creating skeleton QL pack around query",
       step: 2,
@@ -164,29 +291,17 @@ export class SkeletonQueryWizard {
     });
 
     try {
-      const qlPackGenerator = new QlPackGenerator(
-        this.folderName,
-        this.language,
-        this.cliServer,
-        this.qlPackStoragePath,
-      );
+      const qlPackGenerator = this.createQlPackGenerator();
 
       await qlPackGenerator.generate();
     } catch (e: unknown) {
-      void this.logger.log(
+      void this.app.logger.log(
         `Could not create skeleton QL pack: ${getErrorMessage(e)}`,
       );
     }
   }
 
   private async createExampleFile() {
-    if (this.folderName === undefined) {
-      throw new Error("Folder name is undefined");
-    }
-    if (this.language === undefined) {
-      throw new Error("Language is undefined");
-    }
-
     this.progress({
       message:
         "Skeleton query pack already exists. Creating additional query example file.",
@@ -195,29 +310,29 @@ export class SkeletonQueryWizard {
     });
 
     try {
-      const qlPackGenerator = new QlPackGenerator(
-        this.folderName,
-        this.language,
-        this.cliServer,
-        this.qlPackStoragePath,
-      );
+      const qlPackGenerator = this.createQlPackGenerator();
 
-      this.fileName = await this.determineNextFileName(this.folderName);
+      this.fileName = await this.determineNextFileName();
       await qlPackGenerator.createExampleQlFile(this.fileName);
     } catch (e: unknown) {
-      void this.logger.log(
-        `Could not create skeleton QL pack: ${getErrorMessage(e)}`,
+      void this.app.logger.log(
+        `Could not create query example file: ${getErrorMessage(e)}`,
       );
     }
   }
 
-  private async determineNextFileName(folderName: string): Promise<string> {
-    if (this.qlPackStoragePath === undefined) {
-      throw new Error("QL Pack storage path is undefined");
+  private async determineNextFileName(): Promise<string> {
+    if (this.queryStoragePath === undefined) {
+      throw new Error("Query storage path is undefined");
     }
 
-    const folderUri = Uri.file(join(this.qlPackStoragePath, folderName));
+    const folderUri = Uri.file(this.queryStoragePath);
     const files = await workspace.fs.readDirectory(folderUri);
+    // If the example.ql file doesn't exist yet, use that name
+    if (!files.some(([filename, _fileType]) => filename === this.fileName)) {
+      return this.fileName;
+    }
+
     const qlFiles = files.filter(([filename, _fileType]) =>
       filename.match(/^example[0-9]*\.ql$/),
     );
@@ -225,11 +340,43 @@ export class SkeletonQueryWizard {
     return `example${qlFiles.length + 1}.ql`;
   }
 
-  private async downloadDatabase() {
-    if (this.qlPackStoragePath === undefined) {
-      throw new Error("QL Pack storage path is undefined");
+  private async promptDownloadDatabase() {
+    if (this.language === undefined) {
+      throw new Error("Language is undefined");
     }
 
+    const openFileLink = this.openFileMarkdownLink;
+
+    const displayLanguage = getLanguageDisplayName(this.language);
+    const action = await showInformationMessageWithAction(
+      `New CodeQL query for ${displayLanguage} ${openFileLink} created, but no CodeQL databases for ${displayLanguage} were detected in your workspace. Would you like to download a CodeQL database for ${displayLanguage} to analyze with ${openFileLink}?`,
+      "Download database",
+    );
+
+    if (action) {
+      void withProgress(async (progress) => {
+        try {
+          await this.downloadDatabase(progress);
+        } catch (e: unknown) {
+          if (e instanceof UserCancellationException) {
+            return;
+          }
+
+          void showAndLogExceptionWithTelemetry(
+            this.app.logger,
+            this.app.telemetry,
+            redactableError(
+              asError(e),
+            )`An error occurred while downloading the GitHub repository: ${getErrorMessage(
+              e,
+            )}`,
+          );
+        }
+      });
+    }
+  }
+
+  private async downloadDatabase(progress: ProgressCallback) {
     if (this.databaseStoragePath === undefined) {
       throw new Error("Database storage path is undefined");
     }
@@ -238,10 +385,10 @@ export class SkeletonQueryWizard {
       throw new Error("Language is undefined");
     }
 
-    this.progress({
+    progress({
       message: "Downloading database",
-      step: 3,
-      maxStep: 3,
+      step: 1,
+      maxStep: 2,
     });
 
     const githubRepoNwo = QUERY_LANGUAGE_TO_DATABASE_REPO[this.language];
@@ -256,7 +403,7 @@ export class SkeletonQueryWizard {
       this.databaseManager,
       this.databaseStoragePath,
       this.credentials,
-      this.progress,
+      progress,
       this.cliServer,
       this.language,
     );
@@ -267,10 +414,6 @@ export class SkeletonQueryWizard {
       throw new Error("Language is undefined");
     }
 
-    if (this.qlPackStoragePath === undefined) {
-      throw new Error("QL Pack storage path is undefined");
-    }
-
     const existingDatabaseItem =
       await SkeletonQueryWizard.findExistingDatabaseItem(
         this.language,
@@ -278,14 +421,65 @@ export class SkeletonQueryWizard {
       );
 
     if (existingDatabaseItem) {
-      // select the found database
-      await this.databaseManager.setCurrentDatabaseItem(existingDatabaseItem);
+      const openFileLink = this.openFileMarkdownLink;
+
+      if (this.databaseManager.currentDatabaseItem !== existingDatabaseItem) {
+        // select the found database
+        await this.databaseManager.setCurrentDatabaseItem(existingDatabaseItem);
+
+        const displayLanguage = getLanguageDisplayName(this.language);
+        void window.showInformationMessage(
+          `New CodeQL query for ${displayLanguage} ${openFileLink} created. We have automatically selected your existing CodeQL ${displayLanguage} database ${existingDatabaseItem.name} for you to analyze with ${openFileLink}.`,
+        );
+      }
     } else {
       // download new database and select it
-      await this.downloadDatabase();
+      this.downloadPromise = this.promptDownloadDatabase().finally(() => {
+        this.downloadPromise = undefined;
+      });
     }
   }
 
+  private get openFileMarkdownLink() {
+    if (this.queryStoragePath === undefined) {
+      throw new Error("QL Pack storage path is undefined");
+    }
+
+    const queryPath = join(this.queryStoragePath, this.fileName);
+    const queryPathUri = Uri.file(queryPath);
+
+    const openFileArgs = [queryPathUri.toString(true)];
+    const queryString = encodeURI(JSON.stringify(openFileArgs));
+    return `[${this.fileName}](command:vscode.open?${queryString})`;
+  }
+
+  private createQlPackGenerator() {
+    if (this.qlPackStoragePath === undefined) {
+      throw new Error("QL pack storage path is undefined");
+    }
+    if (this.queryStoragePath === undefined) {
+      throw new Error("Query storage path is undefined");
+    }
+    if (this.language === undefined) {
+      throw new Error("Language is undefined");
+    }
+
+    const parentFolder = dirname(this.qlPackStoragePath);
+
+    // Only include the folder name in the qlpack name if the qlpack is not in the root of the workspace.
+    const includeFolderNameInQlpackName = !getOnDiskWorkspaceFolders().some(
+      (workspaceFolder) => pathsEqual(workspaceFolder, parentFolder),
+    );
+
+    return new QlPackGenerator(
+      this.language,
+      this.cliServer,
+      this.qlPackStoragePath,
+      this.queryStoragePath,
+      includeFolderNameInQlpackName,
+    );
+  }
+
   public static async findDatabaseItemByNwo(
     language: string,
     databaseNwo: string,

extensions/ql-vscode/src/log-insights/join-order.ts

@@ -40,17 +40,19 @@ function makeKey(
 const DEPENDENT_PREDICATES_REGEXP = (() => {
   const regexps = [
     // SCAN id
-    String.raw`SCAN\s+([0-9a-zA-Z:#_]+)\s`,
+    String.raw`SCAN\s+([0-9a-zA-Z:#_]+|\`[^\`\r\n]*\`)\s`,
     // JOIN id WITH id
-    String.raw`JOIN\s+([0-9a-zA-Z:#_]+)\s+WITH\s+([0-9a-zA-Z:#_]+)\s`,
+    String.raw`JOIN\s+([0-9a-zA-Z:#_]+|\`[^\`\r\n]*\`)\s+WITH\s+([0-9a-zA-Z:#_]+|\`[^\`\r\n]*\`)\s`,
     // AGGREGATE id, id
-    String.raw`AGGREGATE\s+([0-9a-zA-Z:#_]+)\s*,\s+([0-9a-zA-Z:#_]+)`,
+    String.raw`AGGREGATE\s+([0-9a-zA-Z:#_]+|\`[^\`\r\n]*\`)\s*,\s+([0-9a-zA-Z:#_]+|\`[^\`\r\n]*\`)`,
     // id AND NOT id
-    String.raw`([0-9a-zA-Z:#_]+)\s+AND\s+NOT\s+([0-9a-zA-Z:#_]+)`,
+    String.raw`([0-9a-zA-Z:#_]+|\`[^\`\r\n]*\`)\s+AND\s+NOT\s+([0-9a-zA-Z:#_]+|\`[^\`\r\n]*\`)`,
     // INVOKE HIGHER-ORDER RELATION rel ON <id, ..., id>
-    String.raw`INVOKE\s+HIGHER-ORDER\s+RELATION\s[^\s]+\sON\s+<([0-9a-zA-Z:#_<>]+)((?:,[0-9a-zA-Z:#_<>]+)*)>`,
+    String.raw`INVOKE\s+HIGHER-ORDER\s+RELATION\s[^\s]+\sON\s+<([0-9a-zA-Z:#_<>]+|\`[^\`\r\n]*\`)((?:,[0-9a-zA-Z:#_<>]+|,\`[^\`\r\n]*\`)*)>`,
     // SELECT id
-    String.raw`SELECT\s+([0-9a-zA-Z:#_]+)`,
+    String.raw`SELECT\s+([0-9a-zA-Z:#_]+|\`[^\`\r\n]*\`)`,
+    // REWRITE id WITH
+    String.raw`REWRITE\s+([0-9a-zA-Z:#_]+|\`[^\`\r\n]*\`)\s+WITH\s`,
   ];
   return new RegExp(
     `${String.raw`\{[0-9]+\}\s+[0-9a-zA-Z]+\s=\s(?:` + regexps.join("|")})`,
@@ -65,7 +67,12 @@ function getDependentPredicates(operations: string[]): I.List<string> {
         .rest() // Skip the first group as it's just the entire string
         .filter((x) => !!x && !x.match("r[0-9]+|PRIMITIVE")) // Only keep the references to predicates.
         .flatMap((x) => x.split(",")) // Group 2 in the INVOKE HIGHER_ORDER RELATION case is a comma-separated list of identifiers.
-        .filter((x) => !!x); // Remove empty strings
+        .filter((x) => !!x) // Remove empty strings
+        .map((x) =>
+          x.startsWith("`") && x.endsWith("`")
+            ? x.substring(1, x.length - 1)
+            : x,
+        ); // Remove quotes from quoted identifiers
     } else {
       return I.List();
     }

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

@@ -1,5 +1,7 @@
 import { Credentials } from "../common/authentication";
 import { OctokitResponse } from "@octokit/types";
+import fetch from "node-fetch";
+import { ModelConfigListener } from "../config";
 
 export enum AutomodelMode {
   Unspecified = "AUTOMODEL_MODE_UNSPECIFIED",
@@ -20,15 +22,44 @@ export interface ModelResponse {
 export async function autoModel(
   credentials: Credentials,
   request: ModelRequest,
+  modelingConfig: ModelConfigListener,
 ): Promise<ModelResponse> {
-  const octokit = await credentials.getOctokit();
+  const devEndpoint = modelingConfig.llmGenerationDevEndpoint;
+  if (devEndpoint) {
+    return callAutoModelDevEndpoint(devEndpoint, request);
+  } else {
+    const octokit = await credentials.getOctokit();
 
-  const response: OctokitResponse<ModelResponse> = await octokit.request(
-    "POST /repos/github/codeql/code-scanning/codeql/auto-model",
-    {
-      data: request,
-    },
-  );
+    const response: OctokitResponse<ModelResponse> = await octokit.request(
+      "POST /repos/github/codeql/code-scanning/codeql/auto-model",
+      {
+        data: request,
+      },
+    );
 
-  return response.data;
+    return response.data;
+  }
+}
+
+async function callAutoModelDevEndpoint(
+  endpoint: string,
+  request: ModelRequest,
+): Promise<ModelResponse> {
+  const json = JSON.stringify(request);
+  const response = await fetch(endpoint, {
+    method: "POST",
+    headers: {
+      "Content-Type": "application/json",
+    },
+    body: json,
+  });
+
+  if (!response.ok) {
+    throw new Error(
+      `Error calling auto-model API: ${response.status} ${response.statusText}`,
+    );
+  }
+
+  const data = await response.json();
+  return data as ModelResponse;
 }

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>,
+  methods: readonly Method[],
+  modeledMethodsBySignature: Record<string, readonly 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,12 @@ 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

@@ -1,6 +1,5 @@
 import { Method, MethodSignature } from "./method";
 import { ModeledMethod } from "./modeled-method";
-import { extLogger } from "../common/logging/vscode";
 import { load as loadYaml } from "js-yaml";
 import { ProgressCallback, withProgress } from "../common/vscode/progress";
 import { createAutoModelRequest, getCandidates } from "./auto-model";
@@ -16,11 +15,9 @@ import { QueryRunner } from "../query-server";
 import { DatabaseItem } from "../databases/local-databases";
 import { Mode } from "./shared/mode";
 import { CancellationTokenSource } from "vscode";
-
-// 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 { ModelingStore } from "./modeling-store";
+import { ModelConfigListener } from "../config";
+import { QueryLanguage } from "../common/query-language";
 
 /**
  * The auto-modeler holds state around auto-modeling jobs and allows
@@ -35,14 +32,13 @@ 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,
-    private readonly setInProgressMethods: (
-      packageName: string,
-      inProgressMethods: string[],
-    ) => Promise<void>,
+    private readonly language: QueryLanguage,
     private readonly addModeledMethods: (
-      modeledMethods: Record<string, ModeledMethod>,
+      modeledMethods: Record<string, ModeledMethod[]>,
     ) => Promise<void>,
   ) {
     this.jobs = new Map<string, CancellationTokenSource>();
@@ -58,8 +54,8 @@ export class AutoModeler {
    */
   public async startModeling(
     packageName: string,
-    methods: Method[],
-    modeledMethods: Record<string, ModeledMethod>,
+    methods: readonly Method[],
+    modeledMethods: Record<string, readonly ModeledMethod[]>,
     mode: Mode,
   ): Promise<void> {
     if (this.jobs.has(packageName)) {
@@ -87,7 +83,7 @@ export class AutoModeler {
    * @param packageName The name of the package to stop modeling.
    */
   public async stopModeling(packageName: string): Promise<void> {
-    void extLogger.log(`Stopping modeling for package ${packageName}`);
+    void this.app.logger.log(`Stopping modeling for package ${packageName}`);
     const cancellationTokenSource = this.jobs.get(packageName);
     if (cancellationTokenSource) {
       cancellationTokenSource.cancel();
@@ -105,19 +101,22 @@ export class AutoModeler {
 
   private async modelPackage(
     packageName: string,
-    methods: Method[],
-    modeledMethods: Record<string, ModeledMethod>,
+    methods: readonly Method[],
+    modeledMethods: Record<string, readonly ModeledMethod[]>,
     mode: Mode,
     cancellationTokenSource: CancellationTokenSource,
   ): Promise<void> {
-    void extLogger.log(`Modeling package ${packageName}`);
+    void this.app.logger.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);
 
       // If there are no candidates, there is nothing to model and we just return
       if (allCandidateMethods.length === 0) {
-        void extLogger.log("No candidates to model. Stopping.");
+        void this.app.logger.log("No candidates to model. Stopping.");
         return;
       }
 
@@ -135,11 +134,14 @@ export class AutoModeler {
           const start = i * candidateBatchSize;
           const end = start + candidateBatchSize;
           const candidatesToProcess = allCandidateMethods.slice(start, end);
+          const candidateSignatures = candidatesToProcess.map(
+            (c) => c.signature,
+          );
 
           // Let the UI know which candidates we are modeling
-          await this.setInProgressMethods(
-            packageName,
-            candidatesToProcess.map((c) => c.signature),
+          this.modelingStore.addInProgressMethods(
+            this.databaseItem,
+            candidateSignatures,
           );
 
           // Kick off the process to model the slice of candidates
@@ -149,10 +151,19 @@ export class AutoModeler {
             progress,
             cancellationTokenSource,
           );
+
+          // Let the UI know which candidates we are done modeling
+          this.modelingStore.removeInProgressMethods(
+            this.databaseItem,
+            candidateSignatures,
+          );
         }
       } finally {
-        // Clear out in progress methods
-        await this.setInProgressMethods(packageName, []);
+        // Clear out in progress methods in case anything went wrong
+        this.modelingStore.removeInProgressMethods(
+          this.databaseItem,
+          allCandidateMethods.map((c) => c.signature),
+        );
       }
     });
   }
@@ -163,7 +174,7 @@ export class AutoModeler {
     progress: ProgressCallback,
     cancellationTokenSource: CancellationTokenSource,
   ): Promise<void> {
-    void extLogger.log("Executing auto-model queries");
+    void this.app.logger.log("Executing auto-model queries");
 
     const usages = await runAutoModelQueries({
       mode,
@@ -181,7 +192,7 @@ export class AutoModeler {
 
     const request = await createAutoModelRequest(mode, usages);
 
-    void extLogger.log("Calling auto-model API");
+    void this.app.logger.log("Calling auto-model API");
 
     const response = await this.callAutoModelApi(request);
     if (!response) {
@@ -192,32 +203,11 @@ export class AutoModeler {
       filename: "auto-model.yml",
     });
 
-    const loadedMethods = loadDataExtensionYaml(models);
+    const loadedMethods = loadDataExtensionYaml(models, this.language);
     if (!loadedMethods) {
       return;
     }
 
-    // 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,
-        };
-      }
-    }
-
     await this.addModeledMethods(loadedMethods);
   }
 
@@ -225,7 +215,7 @@ export class AutoModeler {
     request: ModelRequest,
   ): Promise<ModelResponse | null> {
     try {
-      return await autoModel(this.app.credentials, request);
+      return await autoModel(this.app.credentials, request, this.modelConfig);
     } catch (e) {
       if (e instanceof RequestError && e.status === 429) {
         void showAndLogExceptionWithTelemetry(

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

@@ -4,13 +4,18 @@ import { ModeledMethodType } from "./modeled-method";
 import { parseLibraryFilename } from "./library";
 import { Mode } from "./shared/mode";
 import { ApplicationModeTuple, FrameworkModeTuple } from "./queries/query";
+import { QueryLanguage } from "../common/query-language";
+import { getModelsAsDataLanguage } from "./languages";
 
 export function decodeBqrsToMethods(
   chunk: DecodedBqrsChunk,
   mode: Mode,
+  language: QueryLanguage,
 ): Method[] {
   const methodsByApiName = new Map<string, Method>();
 
+  const definition = getModelsAsDataLanguage(language);
+
   chunk?.tuples.forEach((tuple) => {
     let usage: Call;
     let packageName: string;
@@ -51,7 +56,12 @@ export function decodeBqrsToMethods(
       classification = CallClassification.Unknown;
     }
 
-    const signature = `${packageName}.${typeName}#${methodName}${methodParameters}`;
+    const signature = definition.createMethodSignature({
+      packageName,
+      typeName,
+      methodName,
+      methodParameters,
+    });
 
     // For Java, we'll always get back a .jar file, and the library version may be bad because not all library authors
     // properly specify the version. Therefore, we'll always try to parse the name and version from the library filename
@@ -88,9 +98,16 @@ export function decodeBqrsToMethods(
     }
 
     const method = methodsByApiName.get(signature)!;
-    method.usages.push({
-      ...usage,
-      classification,
+    const usages = [
+      ...method.usages,
+      {
+        ...usage,
+        classification,
+      },
+    ];
+    methodsByApiName.set(signature, {
+      ...method,
+      usages,
     });
   });
 

extensions/ql-vscode/src/model-editor/extension-pack-picker.ts

@@ -11,7 +11,7 @@ import { getQlPackPath, QLPACK_FILENAMES } from "../common/ql";
 import { getErrorMessage } from "../common/helpers-pure";
 import { ExtensionPack } from "./shared/extension-pack";
 import { NotificationLogger, showAndLogErrorMessage } from "../common/logging";
-import { getExtensionsDirectory } from "../config";
+import { ModelConfig } from "../config";
 import {
   autoNameExtensionPack,
   ExtensionPackName,
@@ -28,6 +28,7 @@ const extensionPackValidate = ajv.compile(extensionPackMetadataSchemaJson);
 export async function pickExtensionPack(
   cliServer: Pick<CodeQLCliServer, "resolveQlpacks">,
   databaseItem: Pick<DatabaseItem, "name" | "language">,
+  modelConfig: ModelConfig,
   logger: NotificationLogger,
   progress: ProgressCallback,
   maxStep: number,
@@ -56,12 +57,14 @@ export async function pickExtensionPack(
   });
 
   // Get the `codeQL.model.extensionsDirectory` setting for the language
-  const userExtensionsDirectory = getExtensionsDirectory(databaseItem.language);
+  const userExtensionsDirectory = modelConfig.getExtensionsDirectory(
+    databaseItem.language,
+  );
 
   // If the setting is not set, automatically pick a suitable directory
   const extensionsDirectory = userExtensionsDirectory
     ? Uri.file(userExtensionsDirectory)
-    : await autoPickExtensionsDirectory();
+    : await autoPickExtensionsDirectory(logger);
 
   if (!extensionsDirectory) {
     return undefined;

extensions/ql-vscode/src/model-editor/extensions-workspace-folder.ts

@@ -1,7 +1,7 @@
 import { FileType, Uri, workspace, WorkspaceFolder } from "vscode";
 import { getOnDiskWorkspaceFoldersObjects } from "../common/vscode/workspace-folders";
-import { extLogger } from "../common/logging/vscode";
 import { tmpdir } from "../common/files";
+import { NotificationLogger, showAndLogErrorMessage } from "../common/logging";
 
 /**
  * Returns the ancestors of this path in order from furthest to closest (i.e. root of filesystem to parent directory)
@@ -143,9 +143,20 @@ async function findGitFolder(
  *    for which the .git directory is closest to a workspace folder
  * 6. If none of the above apply, return `undefined`
  */
-export async function autoPickExtensionsDirectory(): Promise<Uri | undefined> {
+export async function autoPickExtensionsDirectory(
+  logger: NotificationLogger,
+): Promise<Uri | undefined> {
   const workspaceFolders = getOnDiskWorkspaceFoldersObjects();
 
+  // If there are no on-disk workspace folders, we can't do anything
+  if (workspaceFolders.length === 0) {
+    void showAndLogErrorMessage(
+      logger,
+      `Could not find any on-disk workspace folders. Please ensure that you have opened a folder or workspace.`,
+    );
+    return undefined;
+  }
+
   // If there's only 1 workspace folder, use the `.github/codeql/extensions` directory in that folder
   if (workspaceFolders.length === 1) {
     return Uri.joinPath(
@@ -168,7 +179,7 @@ export async function autoPickExtensionsDirectory(): Promise<Uri | undefined> {
   // Get the root workspace directory, i.e. the common root directory of all workspace folders
   const rootDirectory = await getRootWorkspaceDirectory();
   if (!rootDirectory) {
-    void extLogger.log("Unable to determine root workspace directory");
+    void logger.log("Unable to determine root workspace directory");
 
     return undefined;
   }
@@ -192,7 +203,7 @@ export async function autoPickExtensionsDirectory(): Promise<Uri | undefined> {
       },
     )
   ) {
-    void extLogger.log(
+    void logger.log(
       `Failed to add workspace folder for extensions at ${extensionsUri.fsPath}`,
     );
     return undefined;

extensions/ql-vscode/src/model-editor/external-api-usage-queries.ts

@@ -1,189 +0,0 @@
-import { QueryRunner } from "../query-server";
-import { getOnDiskWorkspaceFolders } from "../common/vscode/workspace-folders";
-import { extLogger } from "../common/logging/vscode";
-import { showAndLogExceptionWithTelemetry } from "../common/logging";
-import { CancellationToken } from "vscode";
-import { CodeQLCliServer } from "../codeql-cli/cli";
-import { DatabaseItem } from "../databases/local-databases";
-import { ProgressCallback } from "../common/vscode/progress";
-import { redactableError } from "../common/errors";
-import { telemetryListener } from "../common/vscode/telemetry";
-import { join } from "path";
-import { Mode } from "./shared/mode";
-import { writeFile } from "fs-extra";
-import { QueryLanguage } from "../common/query-language";
-import { fetchExternalApiQueries } from "./queries";
-import { Method } from "./method";
-import { runQuery } from "../local-queries/run-query";
-import { decodeBqrsToMethods } from "./bqrs";
-import {
-  resolveEndpointsQuery,
-  syntheticQueryPackName,
-} from "./model-editor-queries";
-
-type RunQueryOptions = {
-  cliServer: CodeQLCliServer;
-  queryRunner: QueryRunner;
-  databaseItem: DatabaseItem;
-  queryStorageDir: string;
-  queryDir: string;
-
-  progress: ProgressCallback;
-  token: CancellationToken;
-};
-
-export async function prepareExternalApiQuery(
-  queryDir: string,
-  language: QueryLanguage,
-): Promise<boolean> {
-  // Resolve the query that we want to run.
-  const query = fetchExternalApiQueries[language];
-  if (!query) {
-    void showAndLogExceptionWithTelemetry(
-      extLogger,
-      telemetryListener,
-      redactableError`No external API usage query found for language ${language}`,
-    );
-    return false;
-  }
-  // Create the query file.
-  Object.values(Mode).map(async (mode) => {
-    const queryFile = join(queryDir, queryNameFromMode(mode));
-    await writeFile(queryFile, query[`${mode}ModeQuery`], "utf8");
-  });
-
-  // Create any dependencies
-  if (query.dependencies) {
-    for (const [filename, contents] of Object.entries(query.dependencies)) {
-      const dependencyFile = join(queryDir, filename);
-      await writeFile(dependencyFile, contents, "utf8");
-    }
-  }
-  return true;
-}
-
-export const externalApiQueriesProgressMaxStep = 2000;
-
-export async function runExternalApiQueries(
-  mode: Mode,
-  {
-    cliServer,
-    queryRunner,
-    databaseItem,
-    queryStorageDir,
-    queryDir,
-    progress,
-    token,
-  }: RunQueryOptions,
-): Promise<Method[] | undefined> {
-  // The below code is temporary to allow for rapid prototyping of the queries. Once the queries are stabilized, we will
-  // move these queries into the `github/codeql` repository and use them like any other contextual (e.g. AST) queries.
-  // This is intentionally not pretty code, as it will be removed soon.
-  // For a reference of what this should do in the future, see the previous implementation in
-  // https://github.com/github/vscode-codeql/blob/089d3566ef0bc67d9b7cc66e8fd6740b31c1c0b0/extensions/ql-vscode/src/data-extensions-editor/external-api-usage-query.ts#L33-L72
-
-  progress({
-    message: "Resolving QL packs",
-    step: 1,
-    maxStep: externalApiQueriesProgressMaxStep,
-  });
-  const additionalPacks = getOnDiskWorkspaceFolders();
-  const extensionPacks = Object.keys(
-    await cliServer.resolveQlpacks(additionalPacks, true),
-  );
-
-  progress({
-    message: "Resolving query",
-    step: 2,
-    maxStep: externalApiQueriesProgressMaxStep,
-  });
-
-  // Resolve the queries from either codeql/java-queries or from the temporary queryDir
-  const queryPath = await resolveEndpointsQuery(
-    cliServer,
-    databaseItem.language,
-    mode,
-    [syntheticQueryPackName],
-    [queryDir],
-  );
-  if (!queryPath) {
-    void showAndLogExceptionWithTelemetry(
-      extLogger,
-      telemetryListener,
-      redactableError`The ${mode} model editor query could not be found. Try re-opening the model editor. If that doesn't work, try upgrading the CodeQL libraries.`,
-    );
-    return;
-  }
-
-  // Run the actual query
-  const completedQuery = await runQuery({
-    queryRunner,
-    databaseItem,
-    queryPath,
-    queryStorageDir,
-    additionalPacks,
-    extensionPacks,
-    progress: (update) =>
-      progress({
-        step: update.step + 500,
-        maxStep: externalApiQueriesProgressMaxStep,
-        message: update.message,
-      }),
-    token,
-  });
-
-  if (!completedQuery) {
-    return;
-  }
-
-  // Read the results and covert to internal representation
-  progress({
-    message: "Decoding results",
-    step: 1600,
-    maxStep: externalApiQueriesProgressMaxStep,
-  });
-
-  const bqrsChunk = await readQueryResults({
-    cliServer,
-    bqrsPath: completedQuery.outputDir.bqrsPath,
-  });
-  if (!bqrsChunk) {
-    return;
-  }
-
-  progress({
-    message: "Finalizing results",
-    step: 1950,
-    maxStep: externalApiQueriesProgressMaxStep,
-  });
-
-  return decodeBqrsToMethods(bqrsChunk, mode);
-}
-
-type GetResultsOptions = {
-  cliServer: Pick<CodeQLCliServer, "bqrsInfo" | "bqrsDecode">;
-  bqrsPath: string;
-};
-
-export async function readQueryResults({
-  cliServer,
-  bqrsPath,
-}: GetResultsOptions) {
-  const bqrsInfo = await cliServer.bqrsInfo(bqrsPath);
-  if (bqrsInfo["result-sets"].length !== 1) {
-    void showAndLogExceptionWithTelemetry(
-      extLogger,
-      telemetryListener,
-      redactableError`Expected exactly one result set, got ${bqrsInfo["result-sets"].length}`,
-    );
-    return undefined;
-  }
-
-  const resultSet = bqrsInfo["result-sets"][0];
-
-  return cliServer.bqrsDecode(bqrsPath, resultSet.name);
-}
-
-function queryNameFromMode(mode: Mode): string {
-  return `${mode.charAt(0).toUpperCase() + mode.slice(1)}ModeEndpoints.ql`;
-}

extensions/ql-vscode/src/model-editor/flow-model-queries.ts

@@ -1,174 +0,0 @@
-import { CancellationToken } from "vscode";
-import { DatabaseItem } from "../databases/local-databases";
-import { basename } from "path";
-import { QueryRunner } from "../query-server";
-import { CodeQLCliServer } from "../codeql-cli/cli";
-import { showAndLogExceptionWithTelemetry } from "../common/logging";
-import { extLogger } from "../common/logging/vscode";
-import { extensiblePredicateDefinitions } from "./predicates";
-import { ProgressCallback } from "../common/vscode/progress";
-import { getOnDiskWorkspaceFolders } from "../common/vscode/workspace-folders";
-import { ModeledMethod, ModeledMethodType } from "./modeled-method";
-import { redactableError } from "../common/errors";
-import { telemetryListener } from "../common/vscode/telemetry";
-import { runQuery } from "../local-queries/run-query";
-import { resolveQueries } from "../local-queries";
-
-type FlowModelOptions = {
-  cliServer: CodeQLCliServer;
-  queryRunner: QueryRunner;
-  queryStorageDir: string;
-  databaseItem: DatabaseItem;
-  progress: ProgressCallback;
-  token: CancellationToken;
-  onResults: (results: ModeledMethod[]) => void | Promise<void>;
-};
-
-export async function runFlowModelQueries({
-  onResults,
-  ...options
-}: FlowModelOptions) {
-  const queries = await resolveFlowQueries(
-    options.cliServer,
-    options.databaseItem,
-  );
-
-  const queriesByBasename: Record<string, string> = {};
-  for (const query of queries) {
-    queriesByBasename[basename(query)] = query;
-  }
-
-  const summaryResults = await runSingleFlowQuery(
-    "summary",
-    queriesByBasename["CaptureSummaryModels.ql"],
-    0,
-    options,
-  );
-  if (summaryResults) {
-    await onResults(summaryResults);
-  }
-
-  const sinkResults = await runSingleFlowQuery(
-    "sink",
-    queriesByBasename["CaptureSinkModels.ql"],
-    1,
-    options,
-  );
-  if (sinkResults) {
-    await onResults(sinkResults);
-  }
-
-  const sourceResults = await runSingleFlowQuery(
-    "source",
-    queriesByBasename["CaptureSourceModels.ql"],
-    2,
-    options,
-  );
-  if (sourceResults) {
-    await onResults(sourceResults);
-  }
-
-  const neutralResults = await runSingleFlowQuery(
-    "neutral",
-    queriesByBasename["CaptureNeutralModels.ql"],
-    3,
-    options,
-  );
-  if (neutralResults) {
-    await onResults(neutralResults);
-  }
-}
-
-async function resolveFlowQueries(
-  cliServer: CodeQLCliServer,
-  databaseItem: DatabaseItem,
-): Promise<string[]> {
-  const packsToSearch = [`codeql/${databaseItem.language}-queries`];
-
-  return await resolveQueries(
-    cliServer,
-    packsToSearch,
-    "flow model generator",
-    {
-      "tags contain": ["modelgenerator"],
-    },
-  );
-}
-
-async function runSingleFlowQuery(
-  type: Exclude<ModeledMethodType, "none">,
-  queryPath: string | undefined,
-  queryStep: number,
-  {
-    cliServer,
-    queryRunner,
-    queryStorageDir,
-    databaseItem,
-    progress,
-    token,
-  }: Omit<FlowModelOptions, "onResults">,
-): Promise<ModeledMethod[]> {
-  // Check that the right query was found
-  if (queryPath === undefined) {
-    void showAndLogExceptionWithTelemetry(
-      extLogger,
-      telemetryListener,
-      redactableError`Failed to find ${type} query`,
-    );
-    return [];
-  }
-
-  // Run the query
-  const completedQuery = await runQuery({
-    queryRunner,
-    databaseItem,
-    queryPath,
-    queryStorageDir,
-    additionalPacks: getOnDiskWorkspaceFolders(),
-    extensionPacks: undefined,
-    progress: ({ step, message }) =>
-      progress({
-        message: `Generating ${type} model: ${message}`,
-        step: queryStep * 1000 + step,
-        maxStep: 4000,
-      }),
-    token,
-  });
-
-  if (!completedQuery) {
-    return [];
-  }
-
-  // Interpret the results
-  const definition = extensiblePredicateDefinitions[type];
-
-  const bqrsPath = completedQuery.outputDir.bqrsPath;
-
-  const bqrsInfo = await cliServer.bqrsInfo(bqrsPath);
-  if (bqrsInfo["result-sets"].length !== 1) {
-    void showAndLogExceptionWithTelemetry(
-      extLogger,
-      telemetryListener,
-      redactableError`Expected exactly one result set, got ${
-        bqrsInfo["result-sets"].length
-      } for ${basename(queryPath)}`,
-    );
-  }
-
-  const resultSet = bqrsInfo["result-sets"][0];
-
-  const decodedResults = await cliServer.bqrsDecode(bqrsPath, resultSet.name);
-
-  const results = decodedResults.tuples;
-
-  return (
-    results
-      // This is just a sanity check. The query should only return strings.
-      .filter((result) => typeof result[0] === "string")
-      .map((result) => {
-        const row = result[0] as string;
-
-        return definition.readModeledMethod(row.split(";"));
-      })
-  );
-}

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

@@ -0,0 +1,99 @@
+import { CancellationToken } from "vscode";
+import { DatabaseItem } from "../databases/local-databases";
+import { basename } from "path";
+import { QueryRunner } from "../query-server";
+import { CodeQLCliServer } from "../codeql-cli/cli";
+import { ProgressCallback } from "../common/vscode/progress";
+import { getOnDiskWorkspaceFolders } from "../common/vscode/workspace-folders";
+import { ModeledMethod } from "./modeled-method";
+import { runQuery } from "../local-queries/run-query";
+import { QueryConstraints, resolveQueries } from "../local-queries";
+import { DecodedBqrs } from "../common/bqrs-cli-types";
+type GenerateQueriesOptions = {
+  queryConstraints: QueryConstraints;
+  filterQueries?: (queryPath: string) => boolean;
+  parseResults: (
+    queryPath: string,
+    results: DecodedBqrs,
+  ) => ModeledMethod[] | Promise<ModeledMethod[]>;
+  onResults: (results: ModeledMethod[]) => void | Promise<void>;
+
+  cliServer: CodeQLCliServer;
+  queryRunner: QueryRunner;
+  queryStorageDir: string;
+  databaseItem: DatabaseItem;
+  progress: ProgressCallback;
+  token: CancellationToken;
+};
+
+export async function runGenerateQueries(options: GenerateQueriesOptions) {
+  const { queryConstraints, filterQueries, parseResults, onResults } = options;
+
+  options.progress({
+    message: "Resolving queries",
+    step: 1,
+    maxStep: 5000,
+  });
+
+  const packsToSearch = [`codeql/${options.databaseItem.language}-queries`];
+  const queryPaths = await resolveQueries(
+    options.cliServer,
+    packsToSearch,
+    "generate model",
+    queryConstraints,
+  );
+
+  const filteredQueryPaths = filterQueries
+    ? queryPaths.filter(filterQueries)
+    : queryPaths;
+
+  const maxStep = filteredQueryPaths.length * 1000;
+
+  for (let i = 0; i < filteredQueryPaths.length; i++) {
+    const queryPath = filteredQueryPaths[i];
+
+    const bqrs = await runSingleGenerateQuery(queryPath, i, maxStep, options);
+    if (bqrs) {
+      await onResults(await parseResults(queryPath, bqrs));
+    }
+  }
+}
+
+async function runSingleGenerateQuery(
+  queryPath: string,
+  queryStep: number,
+  maxStep: number,
+  {
+    cliServer,
+    queryRunner,
+    queryStorageDir,
+    databaseItem,
+    progress,
+    token,
+  }: GenerateQueriesOptions,
+): Promise<DecodedBqrs | undefined> {
+  const queryBasename = basename(queryPath);
+
+  // Run the query
+  const completedQuery = await runQuery({
+    queryRunner,
+    databaseItem,
+    queryPath,
+    queryStorageDir,
+    additionalPacks: getOnDiskWorkspaceFolders(),
+    extensionPacks: undefined,
+    progress: ({ step, message }) =>
+      progress({
+        message: `Generating model from ${queryBasename}: ${message}`,
+        step: queryStep * 1000 + step,
+        maxStep,
+      }),
+    token,
+  });
+
+  if (!completedQuery) {
+    return undefined;
+  }
+
+  return cliServer.bqrsDecodeAll(completedQuery.outputDir.bqrsPath);
+}

extensions/ql-vscode/src/model-editor/languages/index.ts

@@ -0,0 +1,2 @@
+export * from "./languages";
+export * from "./models-as-data";

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

@@ -0,0 +1,36 @@
+import { QueryLanguage } from "../../common/query-language";
+import {
+  ModelsAsDataLanguage,
+  ModelsAsDataLanguagePredicates,
+} from "./models-as-data";
+import { ruby } from "./ruby";
+import { staticLanguage } from "./static";
+
+const languages: Partial<Record<QueryLanguage, ModelsAsDataLanguage>> = {
+  [QueryLanguage.CSharp]: staticLanguage,
+  [QueryLanguage.Java]: staticLanguage,
+  [QueryLanguage.Ruby]: ruby,
+};
+
+export function getModelsAsDataLanguage(
+  language: QueryLanguage,
+): ModelsAsDataLanguage {
+  const definition = languages[language];
+  if (!definition) {
+    throw new Error(`No models-as-data definition for ${language}`);
+  }
+  return definition;
+}
+
+export function getModelsAsDataLanguageModel<
+  T extends keyof ModelsAsDataLanguagePredicates,
+>(
+  language: QueryLanguage,
+  model: T,
+): NonNullable<ModelsAsDataLanguagePredicates[T]> {
+  const definition = getModelsAsDataLanguage(language).predicates[model];
+  if (!definition) {
+    throw new Error(`No models-as-data predicate for ${model}`);
+  }
+  return definition;
+}

extensions/ql-vscode/src/model-editor/languages/models-as-data.ts

@@ -0,0 +1,67 @@
+import { MethodArgument, MethodDefinition } from "../method";
+import {
+  ModeledMethod,
+  NeutralModeledMethod,
+  SinkModeledMethod,
+  SourceModeledMethod,
+  SummaryModeledMethod,
+} from "../modeled-method";
+import { DataTuple } from "../model-extension-file";
+import { Mode } from "../shared/mode";
+import type { QueryConstraints } from "../../local-queries/query-constraints";
+import { DecodedBqrs } from "../../common/bqrs-cli-types";
+import { BaseLogger } from "../../common/logging";
+
+type GenerateMethodDefinition<T> = (method: T) => DataTuple[];
+type ReadModeledMethod = (row: DataTuple[]) => ModeledMethod;
+
+export type ModelsAsDataLanguagePredicate<T> = {
+  extensiblePredicate: string;
+  supportedKinds: string[];
+  generateMethodDefinition: GenerateMethodDefinition<T>;
+  readModeledMethod: ReadModeledMethod;
+};
+
+type ModelsAsDataLanguageModelGeneration = {
+  queryConstraints: QueryConstraints;
+  filterQueries?: (queryPath: string) => boolean;
+  parseResults: (
+    // The path to the query that generated the results.
+    queryPath: string,
+    // The results of the query.
+    bqrs: DecodedBqrs,
+    // The language-specific predicate that was used to generate the results. This is passed to allow
+    // sharing of code between different languages.
+    modelsAsDataLanguage: ModelsAsDataLanguage,
+    // The logger to use for logging.
+    logger: BaseLogger,
+  ) => ModeledMethod[];
+};
+
+export type ModelsAsDataLanguagePredicates = {
+  source?: ModelsAsDataLanguagePredicate<SourceModeledMethod>;
+  sink?: ModelsAsDataLanguagePredicate<SinkModeledMethod>;
+  summary?: ModelsAsDataLanguagePredicate<SummaryModeledMethod>;
+  neutral?: ModelsAsDataLanguagePredicate<NeutralModeledMethod>;
+};
+
+export type MethodArgumentOptions = {
+  options: MethodArgument[];
+  defaultArgumentPath: string;
+};
+
+export type ModelsAsDataLanguage = {
+  /**
+   * The modes that are available for this language. If not specified, all
+   * modes are available.
+   */
+  availableModes?: Mode[];
+  createMethodSignature: (method: MethodDefinition) => string;
+  predicates: ModelsAsDataLanguagePredicates;
+  modelGeneration?: ModelsAsDataLanguageModelGeneration;
+  /**
+   * Returns the list of valid arguments that can be selected for the given method.
+   * @param method The method to get the valid arguments for.
+   */
+  getArgumentOptions: (method: MethodDefinition) => MethodArgumentOptions;
+};

extensions/ql-vscode/src/model-editor/languages/ruby/generate.ts

@@ -0,0 +1,50 @@
+import { BaseLogger } from "../../../common/logging";
+import { DecodedBqrs } from "../../../common/bqrs-cli-types";
+import { ModelsAsDataLanguage } from "../models-as-data";
+import { ModeledMethod } from "../../modeled-method";
+import { DataTuple } from "../../model-extension-file";
+
+export function parseGenerateModelResults(
+  _queryPath: string,
+  bqrs: DecodedBqrs,
+  modelsAsDataLanguage: ModelsAsDataLanguage,
+  logger: BaseLogger,
+): ModeledMethod[] {
+  const modeledMethods: ModeledMethod[] = [];
+
+  for (const resultSetName in bqrs) {
+    const definition = Object.values(modelsAsDataLanguage.predicates).find(
+      (definition) => definition.extensiblePredicate === resultSetName,
+    );
+    if (definition === undefined) {
+      void logger.log(`No predicate found for ${resultSetName}`);
+
+      continue;
+    }
+
+    const resultSet = bqrs[resultSetName];
+
+    if (
+      resultSet.tuples.some((tuple) =>
+        tuple.some((value) => typeof value === "object"),
+      )
+    ) {
+      void logger.log(
+        `Skipping ${resultSetName} because it contains undefined values`,
+      );
+      continue;
+    }
+
+    modeledMethods.push(
+      ...resultSet.tuples.map((tuple) => {
+        const row = tuple.filter(
+          (value): value is DataTuple => typeof value !== "object",
+        );
+
+        return definition.readModeledMethod(row);
+      }),
+    );
+  }
+
+  return modeledMethods;
+}

extensions/ql-vscode/src/model-editor/languages/ruby/index.ts

@@ -0,0 +1,191 @@
+import { ModelsAsDataLanguage } from "../models-as-data";
+import { sharedExtensiblePredicates, sharedKinds } from "../shared";
+import { Mode } from "../../shared/mode";
+import { parseGenerateModelResults } from "./generate";
+import { getArgumentsList, MethodArgument } from "../../method";
+
+function parseRubyMethodFromPath(path: string): string {
+  const match = path.match(/Method\[([^\]]+)].*/);
+  if (match) {
+    return match[1];
+  } else {
+    return "";
+  }
+}
+
+function parseRubyAccessPath(path: string): {
+  methodName: string;
+  path: string;
+} {
+  const match = path.match(/Method\[([^\]]+)]\.(.*)/);
+  if (match) {
+    return { methodName: match[1], path: match[2] };
+  } else {
+    return { methodName: "", path: "" };
+  }
+}
+
+function rubyMethodSignature(typeName: string, methodName: string) {
+  return `${typeName}#${methodName}`;
+}
+
+export const ruby: ModelsAsDataLanguage = {
+  availableModes: [Mode.Framework],
+  createMethodSignature: ({ typeName, methodName }) =>
+    `${typeName}#${methodName}`,
+  predicates: {
+    source: {
+      extensiblePredicate: sharedExtensiblePredicates.source,
+      supportedKinds: sharedKinds.source,
+      // extensible predicate sourceModel(
+      //   string type, string path, string kind
+      // );
+      generateMethodDefinition: (method) => [
+        method.typeName,
+        `Method[${method.methodName}].${method.output}`,
+        method.kind,
+      ],
+      readModeledMethod: (row) => {
+        const typeName = row[0] as string;
+        const { methodName, path: output } = parseRubyAccessPath(
+          row[1] as string,
+        );
+        return {
+          type: "source",
+          input: "",
+          output,
+          kind: row[2] as string,
+          provenance: "manual",
+          signature: rubyMethodSignature(typeName, methodName),
+          packageName: "",
+          typeName,
+          methodName,
+          methodParameters: "",
+        };
+      },
+    },
+    sink: {
+      extensiblePredicate: sharedExtensiblePredicates.sink,
+      supportedKinds: sharedKinds.sink,
+      // extensible predicate sinkModel(
+      //   string type, string path, string kind
+      // );
+      generateMethodDefinition: (method) => {
+        const path = `Method[${method.methodName}].${method.input}`;
+        return [method.typeName, path, method.kind];
+      },
+      readModeledMethod: (row) => {
+        const typeName = row[0] as string;
+        const { methodName, path: input } = parseRubyAccessPath(
+          row[1] as string,
+        );
+        return {
+          type: "sink",
+          input,
+          output: "",
+          kind: row[2] as string,
+          provenance: "manual",
+          signature: rubyMethodSignature(typeName, methodName),
+          packageName: "",
+          typeName,
+          methodName,
+          methodParameters: "",
+        };
+      },
+    },
+    summary: {
+      extensiblePredicate: sharedExtensiblePredicates.summary,
+      supportedKinds: sharedKinds.summary,
+      // extensible predicate summaryModel(
+      //   string type, string path, string input, string output, string kind
+      // );
+      generateMethodDefinition: (method) => [
+        method.typeName,
+        `Method[${method.methodName}]`,
+        method.input,
+        method.output,
+        method.kind,
+      ],
+      readModeledMethod: (row) => {
+        const typeName = row[0] as string;
+        const methodName = parseRubyMethodFromPath(row[1] as string);
+        return {
+          type: "summary",
+          input: row[2] as string,
+          output: row[3] as string,
+          kind: row[4] as string,
+          provenance: "manual",
+          signature: rubyMethodSignature(typeName, methodName),
+          packageName: "",
+          typeName,
+          methodName,
+          methodParameters: "",
+        };
+      },
+    },
+    neutral: {
+      extensiblePredicate: sharedExtensiblePredicates.neutral,
+      supportedKinds: sharedKinds.neutral,
+      // extensible predicate neutralModel(
+      //   string type, string path, string kind
+      // );
+      generateMethodDefinition: (method) => [
+        method.typeName,
+        `Method[${method.methodName}]`,
+        method.kind,
+      ],
+      readModeledMethod: (row) => {
+        const typeName = row[0] as string;
+        const methodName = parseRubyMethodFromPath(row[1] as string);
+        return {
+          type: "neutral",
+          input: "",
+          output: "",
+          kind: row[2] as string,
+          provenance: "manual",
+          signature: rubyMethodSignature(typeName, methodName),
+          packageName: "",
+          typeName,
+          methodName,
+          methodParameters: "",
+        };
+      },
+    },
+  },
+  modelGeneration: {
+    queryConstraints: {
+      "query path": "queries/modeling/GenerateModel.ql",
+    },
+    parseResults: parseGenerateModelResults,
+  },
+  getArgumentOptions: (method) => {
+    const argumentsList = getArgumentsList(method.methodParameters).map(
+      (argument, index): MethodArgument => {
+        if (argument.endsWith(":")) {
+          return {
+            path: `Argument[${argument}]`,
+            label: `Argument[${argument}]`,
+          };
+        }
+
+        return {
+          path: `Argument[${index}]`,
+          label: `Argument[${index}]: ${argument}`,
+        };
+      },
+    );
+
+    return {
+      options: [
+        {
+          path: "Argument[self]",
+          label: "Argument[self]",
+        },
+        ...argumentsList,
+      ],
+      // If there are no arguments, we will default to "Argument[self]"
+      defaultArgumentPath:
+        argumentsList.length > 0 ? argumentsList[0].path : "Argument[self]",
+    };
+  },
+};

extensions/ql-vscode/src/model-editor/languages/shared.ts

@@ -0,0 +1,25 @@
+export const sharedExtensiblePredicates = {
+  source: "sourceModel",
+  sink: "sinkModel",
+  summary: "summaryModel",
+  neutral: "neutralModel",
+};
+
+export const sharedKinds = {
+  source: ["local", "remote"],
+  sink: [
+    "code-injection",
+    "command-injection",
+    "file-content-store",
+    "html-injection",
+    "js-injection",
+    "ldap-injection",
+    "log-injection",
+    "path-injection",
+    "request-forgery",
+    "sql-injection",
+    "url-redirection",
+  ],
+  summary: ["taint", "value"],
+  neutral: ["summary", "source", "sink"],
+};

extensions/ql-vscode/src/model-editor/languages/static/generate.ts

@@ -0,0 +1,60 @@
+import { BaseLogger } from "../../../common/logging";
+import {
+  ModelsAsDataLanguage,
+  ModelsAsDataLanguagePredicates,
+} from "../models-as-data";
+import { DecodedBqrs } from "../../../common/bqrs-cli-types";
+import { ModeledMethod } from "../../modeled-method";
+import { basename } from "../../../common/path";
+
+const queriesToModel: Record<string, keyof ModelsAsDataLanguagePredicates> = {
+  "CaptureSummaryModels.ql": "summary",
+  "CaptureSinkModels.ql": "sink",
+  "CaptureSourceModels.ql": "source",
+  "CaptureNeutralModels.ql": "neutral",
+};
+
+export function filterFlowModelQueries(queryPath: string): boolean {
+  return Object.keys(queriesToModel).includes(basename(queryPath));
+}
+
+export function parseFlowModelResults(
+  queryPath: string,
+  bqrs: DecodedBqrs,
+  modelsAsDataLanguage: ModelsAsDataLanguage,
+  logger: BaseLogger,
+): ModeledMethod[] {
+  if (Object.keys(bqrs).length !== 1) {
+    throw new Error(
+      `Expected exactly one result set from ${queryPath}, but got ${
+        Object.keys(bqrs).length
+      }`,
+    );
+  }
+
+  const modelType = queriesToModel[basename(queryPath)];
+  if (!modelType) {
+    void logger.log(`Unknown model type for ${queryPath}`);
+    return [];
+  }
+
+  const resultSet = bqrs[Object.keys(bqrs)[0]];
+
+  const results = resultSet.tuples;
+
+  const definition = modelsAsDataLanguage.predicates[modelType];
+  if (!definition) {
+    throw new Error(`No definition for ${modelType}`);
+  }
+
+  return (
+    results
+      // This is just a sanity check. The query should only return strings.
+      .filter((result) => typeof result[0] === "string")
+      .map((result) => {
+        const row = result[0] as string;
+
+        return definition.readModeledMethod(row.split(";"));
+      })
+  );
+}

extensions/ql-vscode/src/model-editor/languages/static/index.ts

@@ -0,0 +1,170 @@
+import { ModelsAsDataLanguage } from "../models-as-data";
+import { Provenance } from "../../modeled-method";
+import { DataTuple } from "../../model-extension-file";
+import { sharedExtensiblePredicates, sharedKinds } from "../shared";
+import { filterFlowModelQueries, parseFlowModelResults } from "./generate";
+import { getArgumentsList, MethodArgument } from "../../method";
+
+function readRowToMethod(row: DataTuple[]): string {
+  return `${row[0]}.${row[1]}#${row[3]}${row[4]}`;
+}
+
+export const staticLanguage: ModelsAsDataLanguage = {
+  createMethodSignature: ({
+    packageName,
+    typeName,
+    methodName,
+    methodParameters,
+  }) => `${packageName}.${typeName}#${methodName}${methodParameters}`,
+  predicates: {
+    source: {
+      extensiblePredicate: sharedExtensiblePredicates.source,
+      supportedKinds: sharedKinds.source,
+      // extensible predicate sourceModel(
+      //   string package, string type, boolean subtypes, string name, string signature, string ext,
+      //   string output, string kind, string provenance
+      // );
+      generateMethodDefinition: (method) => [
+        method.packageName,
+        method.typeName,
+        true,
+        method.methodName,
+        method.methodParameters,
+        "",
+        method.output,
+        method.kind,
+        method.provenance,
+      ],
+      readModeledMethod: (row) => ({
+        type: "source",
+        input: "",
+        output: row[6] as string,
+        kind: row[7] as string,
+        provenance: row[8] as Provenance,
+        signature: readRowToMethod(row),
+        packageName: row[0] as string,
+        typeName: row[1] as string,
+        methodName: row[3] as string,
+        methodParameters: row[4] as string,
+      }),
+    },
+    sink: {
+      extensiblePredicate: sharedExtensiblePredicates.sink,
+      supportedKinds: sharedKinds.sink,
+      // extensible predicate sinkModel(
+      //   string package, string type, boolean subtypes, string name, string signature, string ext,
+      //   string input, string kind, string provenance
+      // );
+      generateMethodDefinition: (method) => [
+        method.packageName,
+        method.typeName,
+        true,
+        method.methodName,
+        method.methodParameters,
+        "",
+        method.input,
+        method.kind,
+        method.provenance,
+      ],
+      readModeledMethod: (row) => ({
+        type: "sink",
+        input: row[6] as string,
+        output: "",
+        kind: row[7] as string,
+        provenance: row[8] as Provenance,
+        signature: readRowToMethod(row),
+        packageName: row[0] as string,
+        typeName: row[1] as string,
+        methodName: row[3] as string,
+        methodParameters: row[4] as string,
+      }),
+    },
+    summary: {
+      extensiblePredicate: sharedExtensiblePredicates.summary,
+      supportedKinds: sharedKinds.summary,
+      // extensible predicate summaryModel(
+      //   string package, string type, boolean subtypes, string name, string signature, string ext,
+      //   string input, string output, string kind, string provenance
+      // );
+      generateMethodDefinition: (method) => [
+        method.packageName,
+        method.typeName,
+        true,
+        method.methodName,
+        method.methodParameters,
+        "",
+        method.input,
+        method.output,
+        method.kind,
+        method.provenance,
+      ],
+      readModeledMethod: (row) => ({
+        type: "summary",
+        input: row[6] as string,
+        output: row[7] as string,
+        kind: row[8] as string,
+        provenance: row[9] as Provenance,
+        signature: readRowToMethod(row),
+        packageName: row[0] as string,
+        typeName: row[1] as string,
+        methodName: row[3] as string,
+        methodParameters: row[4] as string,
+      }),
+    },
+    neutral: {
+      extensiblePredicate: sharedExtensiblePredicates.neutral,
+      supportedKinds: sharedKinds.neutral,
+      // extensible predicate neutralModel(
+      //   string package, string type, string name, string signature, string kind, string provenance
+      // );
+      generateMethodDefinition: (method) => [
+        method.packageName,
+        method.typeName,
+        method.methodName,
+        method.methodParameters,
+        method.kind,
+        method.provenance,
+      ],
+      readModeledMethod: (row) => ({
+        type: "neutral",
+        input: "",
+        output: "",
+        kind: row[4] as string,
+        provenance: row[5] as Provenance,
+        signature: `${row[0]}.${row[1]}#${row[2]}${row[3]}`,
+        packageName: row[0] as string,
+        typeName: row[1] as string,
+        methodName: row[2] as string,
+        methodParameters: row[3] as string,
+      }),
+    },
+  },
+  modelGeneration: {
+    queryConstraints: {
+      "tags contain": ["modelgenerator"],
+    },
+    filterQueries: filterFlowModelQueries,
+    parseResults: parseFlowModelResults,
+  },
+  getArgumentOptions: (method) => {
+    const argumentsList = getArgumentsList(method.methodParameters).map(
+      (argument, index): MethodArgument => ({
+        path: `Argument[${index}]`,
+        label: `Argument[${index}]: ${argument}`,
+      }),
+    );
+
+    return {
+      options: [
+        {
+          path: "Argument[this]",
+          label: "Argument[this]",
+        },
+        ...argumentsList,
+      ],
+      // If there are no arguments, we will default to "Argument[this]"
+      defaultArgumentPath:
+        argumentsList.length > 0 ? argumentsList[0].path : "Argument[this]",
+    };
+  },
+};

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

@@ -4,15 +4,34 @@ import { DisposableObject } from "../../common/disposable-object";
 import { MethodModelingViewProvider } from "./method-modeling-view-provider";
 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";
+import { ModelingEvents } from "../modeling-events";
 
 export class MethodModelingPanel extends DisposableObject {
   private readonly provider: MethodModelingViewProvider;
 
-  constructor(app: App, modelingStore: ModelingStore) {
+  constructor(
+    app: App,
+    modelingStore: ModelingStore,
+    modelingEvents: ModelingEvents,
+    editorViewTracker: ModelEditorViewTracker,
+  ) {
     super();
 
-    this.provider = new MethodModelingViewProvider(app, modelingStore);
-    this.push(this.provider);
+    // This is here instead of in MethodModelingViewProvider because we need to
+    // dispose this when the extension gets disposed, not when the webview gets
+    // disposed.
+    const modelConfig = this.push(new ModelConfigListener());
+
+    this.provider = new MethodModelingViewProvider(
+      app,
+      modelingStore,
+      modelingEvents,
+      editorViewTracker,
+      modelConfig,
+    );
     this.push(
       window.registerWebviewViewProvider(
         MethodModelingViewProvider.viewType,
@@ -21,7 +40,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

@@ -1,153 +1,265 @@
-import * as vscode from "vscode";
-import { Uri, WebviewViewProvider } from "vscode";
-import { getHtmlForWebview } from "../../common/vscode/webview-html";
-import { FromMethodModelingMessage } from "../../common/interface-types";
+import {
+  FromMethodModelingMessage,
+  ToMethodModelingMessage,
+} from "../../common/interface-types";
 import { telemetryListener } from "../../common/vscode/telemetry";
 import { showAndLogExceptionWithTelemetry } from "../../common/logging/notifications";
-import { extLogger } from "../../common/logging/vscode/loggers";
 import { App } from "../../common/app";
 import { redactableError } from "../../common/errors";
 import { Method } from "../method";
-import { DisposableObject } from "../../common/disposable-object";
 import { ModelingStore } from "../modeling-store";
+import { AbstractWebviewViewProvider } from "../../common/vscode/abstract-webview-view-provider";
+import { assertNever } from "../../common/helpers-pure";
+import { ModelEditorViewTracker } from "../model-editor-view-tracker";
+import { ModelConfigListener } from "../../config";
+import { DatabaseItem } from "../../databases/local-databases";
+import { ModelingEvents } from "../modeling-events";
+import {
+  QueryLanguage,
+  tryGetQueryLanguage,
+} from "../../common/query-language";
 
-export class MethodModelingViewProvider
-  extends DisposableObject
-  implements WebviewViewProvider
-{
+export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
+  ToMethodModelingMessage,
+  FromMethodModelingMessage
+> {
   public static readonly viewType = "codeQLMethodModeling";
 
-  private webviewView: vscode.WebviewView | undefined = undefined;
-
   private method: Method | undefined = undefined;
+  private databaseItem: DatabaseItem | undefined = undefined;
+  private language: QueryLanguage | undefined = undefined;
 
   constructor(
-    private readonly app: App,
+    app: App,
     private readonly modelingStore: ModelingStore,
+    private readonly modelingEvents: ModelingEvents,
+    private readonly editorViewTracker: ModelEditorViewTracker,
+    private readonly modelConfig: ModelConfigListener,
   ) {
-    super();
+    super(app, "method-modeling");
   }
 
-  /**
-   * This is called when a view first becomes visible. This may happen when the view is
-   * first loaded or when the user hides and then shows a view again.
-   */
-  public resolveWebviewView(
-    webviewView: vscode.WebviewView,
-    _context: vscode.WebviewViewResolveContext,
-    _token: vscode.CancellationToken,
-  ) {
-    webviewView.webview.options = {
-      enableScripts: true,
-      localResourceRoots: [Uri.file(this.app.extensionPath)],
-    };
+  protected override async onWebViewLoaded(): Promise<void> {
+    await Promise.all([this.setViewState(), this.setInitialState()]);
+    this.registerToModelingEvents();
+    this.registerToModelConfigEvents();
+  }
 
-    const html = getHtmlForWebview(
-      this.app,
-      webviewView.webview,
-      "method-modeling",
-      {
-        allowInlineStyles: true,
-        allowWasmEval: false,
+  private async setViewState(): Promise<void> {
+    await this.postMessage({
+      t: "setMethodModelingPanelViewState",
+      viewState: {
+        language: this.language,
+        showMultipleModels: this.modelConfig.showMultipleModels,
       },
-    );
-
-    webviewView.webview.html = html;
-
-    webviewView.webview.onDidReceiveMessage(async (msg) => this.onMessage(msg));
-
-    this.webviewView = webviewView;
-
-    this.setInitialState(webviewView);
-    this.registerToModelingStoreEvents();
+    });
   }
 
-  public async setMethod(method: Method): Promise<void> {
+  public async setMethod(
+    databaseItem: DatabaseItem | undefined,
+    method: Method | undefined,
+  ): Promise<void> {
     this.method = method;
+    this.databaseItem = databaseItem;
+    this.language = databaseItem && tryGetQueryLanguage(databaseItem.language);
 
-    if (this.webviewView) {
-      await this.webviewView.webview.postMessage({
+    if (this.isShowingView) {
+      await this.postMessage({
         t: "setMethod",
         method,
       });
     }
   }
 
-  private setInitialState(webviewView: vscode.WebviewView): void {
-    const selectedMethod = this.modelingStore.getSelectedMethodDetails();
-    if (selectedMethod) {
-      void webviewView.webview.postMessage({
-        t: "setSelectedMethod",
-        method: selectedMethod.method,
-        modeledMethod: selectedMethod.modeledMethod,
-        isModified: selectedMethod.isModified,
+  private async setInitialState(): Promise<void> {
+    if (this.modelingStore.hasStateForActiveDb()) {
+      const selectedMethod = this.modelingStore.getSelectedMethodDetails();
+      if (selectedMethod) {
+        this.databaseItem = selectedMethod.databaseItem;
+        this.language = tryGetQueryLanguage(
+          selectedMethod.databaseItem.language,
+        );
+        this.method = selectedMethod.method;
+
+        await this.postMessage({
+          t: "setSelectedMethod",
+          method: selectedMethod.method,
+          modeledMethods: selectedMethod.modeledMethods,
+          isModified: selectedMethod.isModified,
+          isInProgress: selectedMethod.isInProgress,
+        });
+      }
+
+      await this.postMessage({
+        t: "setInModelingMode",
+        inModelingMode: true,
       });
     }
   }
 
-  private async onMessage(msg: FromMethodModelingMessage): Promise<void> {
+  protected override async onMessage(
+    msg: FromMethodModelingMessage,
+  ): Promise<void> {
     switch (msg.t) {
-      case "setModeledMethod": {
-        const activeState = this.modelingStore.getStateForActiveDb();
-        if (!activeState) {
-          throw new Error("No active state found in modeling store");
-        }
-        this.modelingStore.updateModeledMethod(
-          activeState.databaseItem,
-          msg.method,
-        );
+      case "viewLoaded":
+        await this.onWebViewLoaded();
         break;
-      }
 
-      case "telemetry": {
+      case "telemetry":
         telemetryListener?.sendUIInteraction(msg.action);
         break;
-      }
+
       case "unhandledError":
         void showAndLogExceptionWithTelemetry(
-          extLogger,
+          this.app.logger,
           telemetryListener,
           redactableError(
             msg.error,
           )`Unhandled error in method modeling view: ${msg.error.message}`,
         );
         break;
+
+      case "setMultipleModeledMethods": {
+        if (!this.databaseItem) {
+          return;
+        }
+
+        this.modelingStore.updateModeledMethods(
+          this.databaseItem,
+          msg.methodSignature,
+          msg.modeledMethods,
+        );
+        this.modelingStore.addModifiedMethod(
+          this.databaseItem,
+          msg.methodSignature,
+        );
+        break;
+      }
+      case "revealInModelEditor":
+        await this.revealInModelEditor(msg.method);
+        void telemetryListener?.sendUIInteraction(
+          "method-modeling-reveal-in-model-editor",
+        );
+
+        break;
+
+      case "startModeling":
+        await this.app.commands.execute(
+          "codeQL.openModelEditorFromModelingPanel",
+        );
+        break;
+      default:
+        assertNever(msg);
     }
   }
 
-  private registerToModelingStoreEvents(): void {
-    this.modelingStore.onModeledMethodsChanged(async (e) => {
-      if (this.webviewView && e.isActiveDb) {
-        const modeledMethod = e.modeledMethods[this.method?.signature ?? ""];
-        if (modeledMethod) {
-          await this.webviewView.webview.postMessage({
-            t: "setModeledMethod",
-            method: modeledMethod,
+  private async revealInModelEditor(method: Method): Promise<void> {
+    if (!this.databaseItem) {
+      return;
+    }
+
+    const view = this.editorViewTracker.getView(
+      this.databaseItem.databaseUri.toString(),
+    );
+    await view?.revealMethod(method);
+  }
+
+  private registerToModelingEvents(): void {
+    this.push(
+      this.modelingEvents.onModeledMethodsChanged(async (e) => {
+        if (this.webviewView && e.isActiveDb && this.method) {
+          const modeledMethods = e.modeledMethods[this.method.signature];
+          if (modeledMethods) {
+            await this.postMessage({
+              t: "setMultipleModeledMethods",
+              methodSignature: this.method.signature,
+              modeledMethods,
+            });
+          }
+        }
+      }),
+    );
+
+    this.push(
+      this.modelingEvents.onModifiedMethodsChanged(async (e) => {
+        if (this.webviewView && e.isActiveDb && this.method) {
+          const isModified = e.modifiedMethods.has(this.method.signature);
+          await this.postMessage({
+            t: "setMethodModified",
+            isModified,
           });
         }
-      }
-    });
+      }),
+    );
 
-    this.modelingStore.onModifiedMethodsChanged(async (e) => {
-      if (this.webviewView && e.isActiveDb && this.method) {
-        const isModified = e.modifiedMethods.has(this.method.signature);
-        await this.webviewView.webview.postMessage({
-          t: "setMethodModified",
-          isModified,
-        });
-      }
-    });
+    this.push(
+      this.modelingEvents.onSelectedMethodChanged(async (e) => {
+        if (this.webviewView) {
+          this.method = e.method;
+          this.databaseItem = e.databaseItem;
+          this.language = tryGetQueryLanguage(e.databaseItem.language);
 
-    this.modelingStore.onSelectedMethodChanged(async (e) => {
-      if (this.webviewView) {
-        this.method = e.method;
-        await this.webviewView.webview.postMessage({
-          t: "setSelectedMethod",
-          method: e.method,
-          modeledMethod: e.modeledMethod,
-          isModified: e.isModified,
+          await this.postMessage({
+            t: "setSelectedMethod",
+            method: e.method,
+            modeledMethods: e.modeledMethods,
+            isModified: e.isModified,
+            isInProgress: e.isInProgress,
+          });
+        }
+      }),
+    );
+
+    this.push(
+      this.modelingEvents.onDbOpened(async (databaseItem) => {
+        this.databaseItem = databaseItem;
+
+        await this.postMessage({
+          t: "setInModelingMode",
+          inModelingMode: true,
         });
-      }
-    });
+
+        this.language = tryGetQueryLanguage(databaseItem.language);
+        await this.setViewState();
+      }),
+    );
+
+    this.push(
+      this.modelingEvents.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);
+        }
+      }),
+    );
+
+    this.push(
+      this.modelingEvents.onInProgressMethodsChanged(async (e) => {
+        if (this.method && this.databaseItem) {
+          const dbUri = this.databaseItem.databaseUri.toString();
+          if (e.dbUri === dbUri) {
+            const inProgress = e.methods.has(this.method.signature);
+            await this.postMessage({
+              t: "setInProgress",
+              inProgress,
+            });
+          }
+        }
+      }),
+    );
+  }
+
+  private registerToModelConfigEvents(): void {
+    this.push(
+      this.modelConfig.onDidChangeConfiguration(() => {
+        void this.setViewState();
+      }),
+    );
   }
 }

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

@@ -1,9 +1,9 @@
 import { ResolvableLocationValue } from "../common/bqrs-cli-types";
-import { ModeledMethodType } from "./modeled-method";
+import { ModeledMethod, ModeledMethodType } from "./modeled-method";
 
 export type Call = {
-  label: string;
-  url: ResolvableLocationValue;
+  readonly label: string;
+  readonly url: Readonly<ResolvableLocationValue>;
 };
 
 export enum CallClassification {
@@ -14,14 +14,29 @@ export enum CallClassification {
 }
 
 export type Usage = Call & {
-  classification: CallClassification;
+  readonly classification: CallClassification;
 };
 
-export interface MethodSignature {
+export interface MethodDefinition {
+  /**
+   * The package name in Java, or the namespace in C#, e.g. `org.sql2o` or `System.Net.Http.Headers`.
+   *
+   * If the class is not in a package, the value should be an empty string.
+   */
+  readonly packageName: string;
+  readonly typeName: string;
+  readonly methodName: string;
+  /**
+   * The method parameters, including enclosing parentheses, e.g. `(String, String)`
+   */
+  readonly methodParameters: string;
+}
+
+export interface MethodSignature extends MethodDefinition {
   /**
    * Contains the version of the library if it can be determined by CodeQL, e.g. `4.2.2.2`
    */
-  libraryVersion?: string;
+  readonly libraryVersion?: string;
   /**
    * A unique signature that can be used to identify this external API usage.
    *
@@ -29,33 +44,26 @@ export interface MethodSignature {
    * in the form "packageName.typeName#methodName(methodParameters)".
    * e.g. `org.sql2o.Connection#createQuery(String)`
    */
-  signature: string;
-  /**
-   * The package name in Java, or the namespace in C#, e.g. `org.sql2o` or `System.Net.Http.Headers`.
-   *
-   * If the class is not in a package, the value should be an empty string.
-   */
-  packageName: string;
-  typeName: string;
-  methodName: string;
-  /**
-   * The method parameters, including enclosing parentheses, e.g. `(String, String)`
-   */
-  methodParameters: string;
+  readonly signature: string;
 }
 
 export interface Method extends MethodSignature {
   /**
    * Contains the name of the library containing the method declaration, e.g. `sql2o-1.6.0.jar` or `System.Runtime.dll`
    */
-  library: string;
+  readonly library: string;
   /**
    * Is this method already supported by CodeQL standard libraries.
    * If so, there is no need for the user to model it themselves.
    */
-  supported: boolean;
-  supportedType: ModeledMethodType;
-  usages: Usage[];
+  readonly supported: boolean;
+  readonly supportedType: ModeledMethodType;
+  readonly usages: readonly Usage[];
+}
+
+export interface MethodArgument {
+  path: string;
+  label: string;
 }
 
 export function getArgumentsList(methodParameters: string): string[] {
@@ -65,3 +73,21 @@ export function getArgumentsList(methodParameters: string): string[] {
 
   return methodParameters.substring(1, methodParameters.length - 1).split(",");
 }
+
+/**
+ * Should we present the user with the ability to edit to modelings for this method.
+ *
+ * A method may be unmodelable if it is already modeled by CodeQL or by an extension
+ * pack other than the one currently being edited.
+ */
+export function canMethodBeModeled(
+  method: Method,
+  modeledMethods: readonly ModeledMethod[],
+  methodIsUnsaved: boolean,
+): boolean {
+  return (
+    !method.supported ||
+    modeledMethods.some((modeledMethod) => modeledMethod.type !== "none") ||
+    methodIsUnsaved
+  );
+}

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

@@ -9,7 +9,7 @@ import {
   Uri,
 } from "vscode";
 import { DisposableObject } from "../../common/disposable-object";
-import { Method, Usage } from "../method";
+import { Method, Usage, canMethodBeModeled } from "../method";
 import { DatabaseItem } from "../../databases/local-databases";
 import { relative } from "path";
 import { CodeQLCliServer } from "../../codeql-cli/cli";
@@ -17,17 +17,24 @@ import { INITIAL_HIDE_MODELED_METHODS_VALUE } from "../shared/hide-modeled-metho
 import { getModelingStatus } from "../shared/modeling-status";
 import { assertNever } from "../../common/helpers-pure";
 import { ModeledMethod } from "../modeled-method";
+import { groupMethods, sortGroupNames, sortMethods } from "../shared/sorting";
+import { INITIAL_MODE, Mode } from "../shared/mode";
 
 export class MethodsUsageDataProvider
   extends DisposableObject
   implements TreeDataProvider<MethodsUsageTreeViewItem>
 {
-  private methods: Method[] = [];
+  private methods: readonly Method[] = [];
+  // sortedMethods is a separate field so we can check if the methods have changed
+  // by reference, which is faster than checking if the methods have changed by value.
+  private sortedTreeItems: readonly MethodTreeViewItem[] = [];
   private databaseItem: DatabaseItem | undefined = undefined;
   private sourceLocationPrefix: string | undefined = undefined;
   private hideModeledMethods: boolean = INITIAL_HIDE_MODELED_METHODS_VALUE;
-  private modeledMethods: Record<string, ModeledMethod> = {};
-  private modifiedMethodSignatures: Set<string> = new Set();
+  private mode: Mode = INITIAL_MODE;
+  private modeledMethods: Readonly<Record<string, readonly ModeledMethod[]>> =
+    {};
+  private modifiedMethodSignatures: ReadonlySet<string> = new Set();
 
   private readonly onDidChangeTreeDataEmitter = this.push(
     new EventEmitter<void>(),
@@ -49,24 +56,30 @@ export class MethodsUsageDataProvider
    * method and instead always pass new objects/arrays.
    */
   public async setState(
-    methods: Method[],
+    methods: readonly Method[],
     databaseItem: DatabaseItem,
     hideModeledMethods: boolean,
-    modeledMethods: Record<string, ModeledMethod>,
-    modifiedMethodSignatures: Set<string>,
+    mode: Mode,
+    modeledMethods: Readonly<Record<string, readonly ModeledMethod[]>>,
+    modifiedMethodSignatures: ReadonlySet<string>,
   ): Promise<void> {
     if (
       this.methods !== methods ||
       this.databaseItem !== databaseItem ||
       this.hideModeledMethods !== hideModeledMethods ||
+      this.mode !== mode ||
       this.modeledMethods !== modeledMethods ||
       this.modifiedMethodSignatures !== modifiedMethodSignatures
     ) {
       this.methods = methods;
+      this.sortedTreeItems = createTreeItems(
+        sortMethodsInGroups(methods, mode),
+      );
       this.databaseItem = databaseItem;
       this.sourceLocationPrefix =
         await this.databaseItem.getSourceLocationPrefix(this.cliServer);
       this.hideModeledMethods = hideModeledMethods;
+      this.mode = mode;
       this.modeledMethods = modeledMethods;
       this.modifiedMethodSignatures = modifiedMethodSignatures;
 
@@ -75,34 +88,37 @@ export class MethodsUsageDataProvider
   }
 
   getTreeItem(item: MethodsUsageTreeViewItem): TreeItem {
-    if (isExternalApiUsage(item)) {
+    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);
+      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.jumpToUsageLocation",
-          arguments: [method, item, this.databaseItem],
+          command: "codeQLModelEditor.jumpToMethod",
+          arguments: [method, usage, this.databaseItem],
         },
       };
     }
   }
 
   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"));
@@ -130,12 +146,18 @@ export class MethodsUsageDataProvider
   getChildren(item?: MethodsUsageTreeViewItem): MethodsUsageTreeViewItem[] {
     if (item === undefined) {
       if (this.hideModeledMethods) {
-        return this.methods.filter((api) => !api.supported);
+        return this.sortedTreeItems.filter((api) =>
+          canMethodBeModeled(
+            api.method,
+            this.modeledMethods[api.method.signature] ?? [],
+            this.modifiedMethodSignatures.has(api.method.signature),
+          ),
+        );
       } else {
-        return this.methods;
+        return [...this.sortedTreeItems];
       }
-    } else if (isExternalApiUsage(item)) {
-      return item.usages;
+    } else if (isMethodTreeViewItem(item)) {
+      return item.children;
     } else {
       return [];
     }
@@ -144,29 +166,45 @@ export class MethodsUsageDataProvider
   getParent(
     item: MethodsUsageTreeViewItem,
   ): MethodsUsageTreeViewItem | undefined {
-    if (isExternalApiUsage(item)) {
+    if (isMethodTreeViewItem(item)) {
       return undefined;
     } else {
-      return this.methods.find((e) => e.usages.includes(item));
+      return item.parent;
     }
   }
 
-  public resolveCanonicalUsage(usage: Usage): Usage | undefined {
-    for (const method of this.methods) {
-      for (const u of method.usages) {
-        if (usagesAreEqual(u, usage)) {
-          return u;
-        }
-      }
+  public resolveUsageTreeViewItem(
+    methodSignature: string,
+    usage: Usage,
+  ): UsageTreeViewItem | undefined {
+    const method = this.sortedTreeItems.find(
+      (m) => m.method.signature === methodSignature,
+    );
+    if (!method) {
+      return undefined;
     }
-    return undefined;
+
+    return method.children.find((u) => usagesAreEqual(u.usage, usage));
   }
 }
 
-export type MethodsUsageTreeViewItem = Method | Usage;
+type MethodTreeViewItem = {
+  method: Method;
+  children: UsageTreeViewItem[];
+};
 
-function isExternalApiUsage(item: MethodsUsageTreeViewItem): item is Method {
-  return (item as any).usages !== undefined;
+type UsageTreeViewItem = {
+  method: Method;
+  usage: Usage;
+  parent: MethodTreeViewItem;
+};
+
+export type MethodsUsageTreeViewItem = MethodTreeViewItem | UsageTreeViewItem;
+
+function isMethodTreeViewItem(
+  item: MethodsUsageTreeViewItem,
+): item is MethodTreeViewItem {
+  return "children" in item && "method" in item;
 }
 
 function usagesAreEqual(u1: Usage, u2: Usage): boolean {
@@ -180,3 +218,33 @@ function usagesAreEqual(u1: Usage, u2: Usage): boolean {
     u1.url.endColumn === u2.url.endColumn
   );
 }
+
+function sortMethodsInGroups(methods: readonly Method[], mode: Mode): Method[] {
+  const grouped = groupMethods(methods, mode);
+
+  const sortedGroupNames = sortGroupNames(grouped);
+
+  return sortedGroupNames.flatMap((groupName) => {
+    const group = grouped[groupName];
+
+    return sortMethods(group);
+  });
+}
+
+function createTreeItems(methods: readonly Method[]): MethodTreeViewItem[] {
+  return methods.map((method) => {
+    const newMethod: MethodTreeViewItem = {
+      method,
+      children: [],
+    };
+
+    newMethod.children = method.usages.map((usage) => ({
+      method,
+      usage,
+      // This needs to be a reference to the parent method, not a copy of it.
+      parent: newMethod,
+    }));
+
+    return newMethod;
+  });
+}

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

@@ -9,6 +9,8 @@ import { DatabaseItem } from "../../databases/local-databases";
 import { CodeQLCliServer } from "../../codeql-cli/cli";
 import { ModelingStore } from "../modeling-store";
 import { ModeledMethod } from "../modeled-method";
+import { Mode } from "../shared/mode";
+import { ModelingEvents } from "../modeling-events";
 
 export class MethodsUsagePanel extends DisposableObject {
   private readonly dataProvider: MethodsUsageDataProvider;
@@ -16,6 +18,7 @@ export class MethodsUsagePanel extends DisposableObject {
 
   public constructor(
     private readonly modelingStore: ModelingStore,
+    private readonly modelingEvents: ModelingEvents,
     cliServer: CodeQLCliServer,
   ) {
     super();
@@ -27,20 +30,22 @@ export class MethodsUsagePanel extends DisposableObject {
     });
     this.push(this.treeView);
 
-    this.registerToModelingStoreEvents();
+    this.registerToModelingEvents();
   }
 
   public async setState(
-    methods: Method[],
+    methods: readonly Method[],
     databaseItem: DatabaseItem,
     hideModeledMethods: boolean,
-    modeledMethods: Record<string, ModeledMethod>,
-    modifiedMethodSignatures: Set<string>,
+    mode: Mode,
+    modeledMethods: Readonly<Record<string, readonly ModeledMethod[]>>,
+    modifiedMethodSignatures: ReadonlySet<string>,
   ): Promise<void> {
     await this.dataProvider.setState(
       methods,
       databaseItem,
       hideModeledMethods,
+      mode,
       modeledMethods,
       modifiedMethodSignatures,
     );
@@ -53,22 +58,28 @@ export class MethodsUsagePanel extends DisposableObject {
     };
   }
 
-  public async revealItem(usage: Usage): Promise<void> {
-    const canonicalUsage = this.dataProvider.resolveCanonicalUsage(usage);
-    if (canonicalUsage !== undefined) {
-      await this.treeView.reveal(canonicalUsage);
+  public async revealItem(
+    methodSignature: string,
+    usage: Usage,
+  ): Promise<void> {
+    const usageTreeViewItem = this.dataProvider.resolveUsageTreeViewItem(
+      methodSignature,
+      usage,
+    );
+    if (usageTreeViewItem !== undefined) {
+      await this.treeView.reveal(usageTreeViewItem);
     }
   }
 
-  private registerToModelingStoreEvents(): void {
+  private registerToModelingEvents(): void {
     this.push(
-      this.modelingStore.onActiveDbChanged(async () => {
+      this.modelingEvents.onActiveDbChanged(async () => {
         await this.handleStateChangeEvent();
       }),
     );
 
     this.push(
-      this.modelingStore.onMethodsChanged(async (event) => {
+      this.modelingEvents.onMethodsChanged(async (event) => {
         if (event.isActiveDb) {
           await this.handleStateChangeEvent();
         }
@@ -76,7 +87,7 @@ export class MethodsUsagePanel extends DisposableObject {
     );
 
     this.push(
-      this.modelingStore.onHideModeledMethodsChanged(async (event) => {
+      this.modelingEvents.onHideModeledMethodsChanged(async (event) => {
         if (event.isActiveDb) {
           await this.handleStateChangeEvent();
         }
@@ -84,7 +95,15 @@ export class MethodsUsagePanel extends DisposableObject {
     );
 
     this.push(
-      this.modelingStore.onModifiedMethodsChanged(async (event) => {
+      this.modelingEvents.onModeChanged(async (event) => {
+        if (event.isActiveDb) {
+          await this.handleStateChangeEvent();
+        }
+      }),
+    );
+
+    this.push(
+      this.modelingEvents.onModifiedMethodsChanged(async (event) => {
         if (event.isActiveDb) {
           await this.handleStateChangeEvent();
         }
@@ -99,6 +118,7 @@ export class MethodsUsagePanel extends DisposableObject {
         activeState.methods,
         activeState.databaseItem,
         activeState.hideModeledMethods,
+        activeState.mode,
         activeState.modeledMethods,
         activeState.modifiedMethodSignatures,
       );

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

@@ -14,20 +14,26 @@ import { dir } from "tmp-promise";
 import { isQueryLanguage } from "../common/query-language";
 import { DisposableObject } from "../common/disposable-object";
 import { MethodsUsagePanel } from "./methods-usage/methods-usage-panel";
-import { Mode } from "./shared/mode";
 import { Method, Usage } from "./method";
-import { setUpPack } from "./model-editor-queries";
+import { setUpPack } from "./model-editor-queries-setup";
 import { MethodModelingPanel } from "./method-modeling/method-modeling-panel";
 import { ModelingStore } from "./modeling-store";
 import { showResolvableLocation } from "../databases/local-databases/locations";
-
-const SUPPORTED_LANGUAGES: string[] = ["java", "csharp"];
+import { ModelEditorViewTracker } from "./model-editor-view-tracker";
+import { ModelConfigListener } from "../config";
+import { ModelingEvents } from "./modeling-events";
+import { getModelsAsDataLanguage } from "./languages";
+import { INITIAL_MODE } from "./shared/mode";
+import { isSupportedLanguage } from "./supported-languages";
 
 export class ModelEditorModule extends DisposableObject {
   private readonly queryStorageDir: string;
   private readonly modelingStore: ModelingStore;
+  private readonly modelingEvents: ModelingEvents;
+  private readonly editorViewTracker: ModelEditorViewTracker<ModelEditorView>;
   private readonly methodsUsagePanel: MethodsUsagePanel;
   private readonly methodModelingPanel: MethodModelingPanel;
+  private readonly modelConfig: ModelConfigListener;
 
   private constructor(
     private readonly app: App,
@@ -38,15 +44,23 @@ export class ModelEditorModule extends DisposableObject {
   ) {
     super();
     this.queryStorageDir = join(baseQueryStorageDir, "model-editor-results");
-    this.modelingStore = new ModelingStore(app);
+    this.modelingEvents = new ModelingEvents(app);
+    this.modelingStore = new ModelingStore(this.modelingEvents);
+    this.editorViewTracker = new ModelEditorViewTracker();
     this.methodsUsagePanel = this.push(
-      new MethodsUsagePanel(this.modelingStore, cliServer),
+      new MethodsUsagePanel(this.modelingStore, this.modelingEvents, cliServer),
     );
     this.methodModelingPanel = this.push(
-      new MethodModelingPanel(app, this.modelingStore),
+      new MethodModelingPanel(
+        app,
+        this.modelingStore,
+        this.modelingEvents,
+        this.editorViewTracker,
+      ),
     );
+    this.modelConfig = this.push(new ModelConfigListener());
 
-    this.registerToModelingStoreEvents();
+    this.registerToModelingEvents();
   }
 
   public static async initialize(
@@ -70,115 +84,10 @@ export class ModelEditorModule extends DisposableObject {
 
   public getCommands(): ModelEditorCommands {
     return {
-      "codeQL.openModelEditor": async () => {
-        const db = this.databaseManager.currentDatabaseItem;
-        if (!db) {
-          void showAndLogErrorMessage(this.app.logger, "No database selected");
-          return;
-        }
-
-        const language = db.language;
-        if (
-          !SUPPORTED_LANGUAGES.includes(language) ||
-          !isQueryLanguage(language)
-        ) {
-          void showAndLogErrorMessage(
-            this.app.logger,
-            `The CodeQL Model Editor is not supported for ${language} databases.`,
-          );
-          return;
-        }
-
-        return withProgress(
-          async (progress) => {
-            const maxStep = 4;
-
-            if (!(await this.cliServer.cliConstraints.supportsQlpacksKind())) {
-              void showAndLogErrorMessage(
-                this.app.logger,
-                `This feature requires CodeQL CLI version ${CliVersionConstraint.CLI_VERSION_WITH_QLPACKS_KIND.format()} or later.`,
-              );
-              return;
-            }
-
-            if (
-              !(await this.cliServer.cliConstraints.supportsResolveExtensions())
-            ) {
-              void showAndLogErrorMessage(
-                this.app.logger,
-                `This feature requires CodeQL CLI version ${CliVersionConstraint.CLI_VERSION_WITH_RESOLVE_EXTENSIONS.format()} or later.`,
-              );
-              return;
-            }
-
-            const modelFile = await pickExtensionPack(
-              this.cliServer,
-              db,
-              this.app.logger,
-              progress,
-              maxStep,
-            );
-            if (!modelFile) {
-              return;
-            }
-
-            progress({
-              message: "Installing dependencies...",
-              step: 3,
-              maxStep,
-            });
-
-            // Create new temporary directory for query files and pack dependencies
-            const { path: queryDir, cleanup: cleanupQueryDir } = await dir({
-              unsafeCleanup: true,
-            });
-
-            const success = await setUpPack(this.cliServer, queryDir, language);
-            if (!success) {
-              await cleanupQueryDir();
-              return;
-            }
-
-            progress({
-              message: "Opening editor...",
-              step: 4,
-              maxStep,
-            });
-
-            const view = new ModelEditorView(
-              this.app,
-              this.modelingStore,
-              this.databaseManager,
-              this.cliServer,
-              this.queryRunner,
-              this.queryStorageDir,
-              queryDir,
-              db,
-              modelFile,
-              Mode.Application,
-            );
-
-            this.modelingStore.onDbClosed(async (dbUri) => {
-              if (dbUri === db.databaseUri.toString()) {
-                await cleanupQueryDir();
-              }
-            });
-
-            this.push(view);
-            this.push({
-              dispose(): void {
-                void cleanupQueryDir();
-              },
-            });
-
-            await view.openView();
-          },
-          {
-            title: "Opening CodeQL Model Editor",
-          },
-        );
-      },
-      "codeQLModelEditor.jumpToUsageLocation": async (
+      "codeQL.openModelEditor": this.openModelEditor.bind(this),
+      "codeQL.openModelEditorFromModelingPanel":
+        this.openModelEditor.bind(this),
+      "codeQLModelEditor.jumpToMethod": async (
         method: Method,
         usage: Usage,
         databaseItem: DatabaseItem,
@@ -192,9 +101,9 @@ export class ModelEditorModule extends DisposableObject {
     await ensureDir(this.queryStorageDir);
   }
 
-  private registerToModelingStoreEvents(): void {
+  private registerToModelingEvents(): void {
     this.push(
-      this.modelingStore.onSelectedMethodChanged(async (event) => {
+      this.modelingEvents.onSelectedMethodChanged(async (event) => {
         await this.showMethod(event.databaseItem, event.method, event.usage);
       }),
     );
@@ -205,8 +114,155 @@ export class ModelEditorModule extends DisposableObject {
     method: Method,
     usage: Usage,
   ): Promise<void> {
-    await this.methodsUsagePanel.revealItem(usage);
-    await this.methodModelingPanel.setMethod(method);
+    await this.methodsUsagePanel.revealItem(method.signature, usage);
+    await this.methodModelingPanel.setMethod(databaseItem, method);
     await showResolvableLocation(usage.url, databaseItem, this.app.logger);
   }
+
+  private async openModelEditor(): Promise<void> {
+    {
+      const db = this.databaseManager.currentDatabaseItem;
+      if (!db) {
+        void showAndLogErrorMessage(this.app.logger, "No database selected");
+        return;
+      }
+
+      const language = db.language;
+
+      if (
+        !isQueryLanguage(language) ||
+        !isSupportedLanguage(language, this.modelConfig)
+      ) {
+        void showAndLogErrorMessage(
+          this.app.logger,
+          `The CodeQL Model Editor is not supported for ${language} databases.`,
+        );
+        return;
+      }
+
+      const definition = getModelsAsDataLanguage(language);
+
+      const initialMode = definition.availableModes?.[0] ?? INITIAL_MODE;
+
+      const existingView = this.editorViewTracker.getView(
+        db.databaseUri.toString(),
+      );
+      if (existingView) {
+        await existingView.focusView();
+
+        return;
+      }
+
+      return withProgress(
+        async (progress) => {
+          const maxStep = 4;
+
+          if (!(await this.cliServer.cliConstraints.supportsQlpacksKind())) {
+            void showAndLogErrorMessage(
+              this.app.logger,
+              `This feature requires CodeQL CLI version ${CliVersionConstraint.CLI_VERSION_WITH_QLPACKS_KIND.format()} or later.`,
+            );
+            return;
+          }
+
+          if (
+            !(await this.cliServer.cliConstraints.supportsResolveExtensions())
+          ) {
+            void showAndLogErrorMessage(
+              this.app.logger,
+              `This feature requires CodeQL CLI version ${CliVersionConstraint.CLI_VERSION_WITH_RESOLVE_EXTENSIONS.format()} or later.`,
+            );
+            return;
+          }
+
+          const modelFile = await pickExtensionPack(
+            this.cliServer,
+            db,
+            this.modelConfig,
+            this.app.logger,
+            progress,
+            maxStep,
+          );
+          if (!modelFile) {
+            return;
+          }
+
+          progress({
+            message: "Installing dependencies...",
+            step: 3,
+            maxStep,
+          });
+
+          // Create new temporary directory for query files and pack dependencies
+          const { path: queryDir, cleanup: cleanupQueryDir } = await dir({
+            unsafeCleanup: true,
+          });
+
+          const success = await setUpPack(
+            this.cliServer,
+            this.app.logger,
+            queryDir,
+            language,
+            this.modelConfig,
+          );
+          if (!success) {
+            await cleanupQueryDir();
+            return;
+          }
+
+          progress({
+            message: "Opening editor...",
+            step: 4,
+            maxStep,
+          });
+
+          // Check again just before opening the editor to ensure no model editor has been opened between
+          // our first check and now.
+          const existingView = this.editorViewTracker.getView(
+            db.databaseUri.toString(),
+          );
+          if (existingView) {
+            await existingView.focusView();
+
+            return;
+          }
+
+          const view = new ModelEditorView(
+            this.app,
+            this.modelingStore,
+            this.modelingEvents,
+            this.editorViewTracker,
+            this.modelConfig,
+            this.databaseManager,
+            this.cliServer,
+            this.queryRunner,
+            this.queryStorageDir,
+            queryDir,
+            db,
+            modelFile,
+            language,
+            initialMode,
+          );
+
+          this.modelingEvents.onDbClosed(async (dbUri) => {
+            if (dbUri === db.databaseUri.toString()) {
+              await cleanupQueryDir();
+            }
+          });
+
+          this.push(view);
+          this.push({
+            dispose(): void {
+              void cleanupQueryDir();
+            },
+          });
+
+          await view.openView();
+        },
+        {
+          title: "Opening CodeQL Model Editor",
+        },
+      );
+    }
+  }
 }

extensions/ql-vscode/src/model-editor/model-editor-queries-setup.ts

@@ -0,0 +1,144 @@
+import { join } from "path";
+import { QueryLanguage } from "../common/query-language";
+import { writeFile } from "fs-extra";
+import { dump } from "js-yaml";
+import { prepareModelEditorQueries } from "./model-editor-queries";
+import { CodeQLCliServer } from "../codeql-cli/cli";
+import { ModelConfig } from "../config";
+import { Mode } from "./shared/mode";
+import { resolveQueriesFromPacks } from "../local-queries";
+import { modeTag } from "./mode-tag";
+import { NotificationLogger } from "../common/logging";
+
+export const syntheticQueryPackName = "codeql/model-editor-queries";
+
+/**
+ * setUpPack sets up a directory to use for the data extension editor queries if required.
+ *
+ * There are two cases (example language is Java):
+ * - In case the queries are present in the codeql/java-queries, we don't need to write our own queries
+ *   to disk. We still need to create a synthetic query pack so we can pass the queryDir to the query
+ *   resolver without caring about whether the queries are present in the pack or not.
+ * - In case the queries are not present in the codeql/java-queries, we need to write our own queries
+ *   to disk. We will create a synthetic query pack and install its dependencies so it is fully independent
+ *   and we can simply pass it through when resolving the queries.
+ *
+ * These steps together ensure that later steps of the process don't need to keep track of whether the queries
+ * are present in codeql/java-queries or in our own query pack. They just need to resolve the query.
+ *
+ * @param cliServer The CodeQL CLI server to use.
+ * @param logger The logger to use.
+ * @param queryDir The directory to set up.
+ * @param language The language to use for the queries.
+ * @param modelConfig The model config to use.
+ * @returns true if the setup was successful, false otherwise.
+ */
+export async function setUpPack(
+  cliServer: CodeQLCliServer,
+  logger: NotificationLogger,
+  queryDir: string,
+  language: QueryLanguage,
+  modelConfig: ModelConfig,
+): Promise<boolean> {
+  // Download the required query packs
+  await cliServer.packDownload([`codeql/${language}-queries`]);
+
+  // We'll only check if the application mode query exists in the pack and assume that if it does,
+  // the framework mode query will also exist.
+  const applicationModeQuery = await resolveEndpointsQuery(
+    cliServer,
+    language,
+    Mode.Application,
+    [],
+    [],
+  );
+
+  if (applicationModeQuery) {
+    // Set up a synthetic pack so CodeQL doesn't crash later when we try
+    // to resolve a query within this directory
+    const syntheticQueryPack = {
+      name: syntheticQueryPackName,
+      version: "0.0.0",
+      dependencies: {},
+    };
+
+    const qlpackFile = join(queryDir, "codeql-pack.yml");
+    await writeFile(qlpackFile, dump(syntheticQueryPack), "utf8");
+  } else {
+    // If we can't resolve the query, we need to write them to desk ourselves.
+    const externalApiQuerySuccess = await prepareModelEditorQueries(
+      logger,
+      queryDir,
+      language,
+    );
+    if (!externalApiQuerySuccess) {
+      return false;
+    }
+
+    // Set up a synthetic pack so that the query can be resolved later.
+    const syntheticQueryPack = {
+      name: syntheticQueryPackName,
+      version: "0.0.0",
+      dependencies: {
+        [`codeql/${language}-all`]: "*",
+      },
+    };
+
+    const qlpackFile = join(queryDir, "codeql-pack.yml");
+    await writeFile(qlpackFile, dump(syntheticQueryPack), "utf8");
+    await cliServer.packInstall(queryDir);
+  }
+
+  // Download any other required packs
+  if (language === "java" && modelConfig.llmGeneration) {
+    await cliServer.packDownload([`codeql/${language}-automodel-queries`]);
+  }
+
+  return true;
+}
+
+/**
+ * Resolve the query path to the model editor endpoints query. All queries are tagged like this:
+ * modeleditor endpoints <mode>
+ * Example: modeleditor endpoints framework-mode
+ *
+ * @param cliServer The CodeQL CLI server to use.
+ * @param language The language of the query pack to use.
+ * @param mode The mode to resolve the query for.
+ * @param additionalPackNames Additional pack names to search.
+ * @param additionalPackPaths Additional pack paths to search.
+ */
+export async function resolveEndpointsQuery(
+  cliServer: CodeQLCliServer,
+  language: string,
+  mode: Mode,
+  additionalPackNames: string[] = [],
+  additionalPackPaths: string[] = [],
+): Promise<string | undefined> {
+  const packsToSearch = [`codeql/${language}-queries`, ...additionalPackNames];
+
+  // First, resolve the query that we want to run.
+  // All queries are tagged like this:
+  // internal extract automodel <mode> <queryTag>
+  // Example: internal extract automodel framework-mode candidates
+  const queries = await resolveQueriesFromPacks(
+    cliServer,
+    packsToSearch,
+    {
+      kind: "table",
+      "tags contain all": ["modeleditor", "endpoints", modeTag(mode)],
+    },
+    additionalPackPaths,
+  );
+  if (queries.length > 1) {
+    throw new Error(
+      `Found multiple endpoints queries for ${mode}. Can't continue`,
+    );
+  }
+
+  if (queries.length === 0) {
+    return undefined;
+  }
+
+  return queries[0];
+}

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

@@ -1,138 +1,199 @@
-import { join } from "path";
-import { QueryLanguage } from "../common/query-language";
-import { writeFile } from "fs-extra";
-import { dump } from "js-yaml";
-import { prepareExternalApiQuery } from "./external-api-usage-queries";
+import { QueryRunner } from "../query-server";
+import { getOnDiskWorkspaceFolders } from "../common/vscode/workspace-folders";
+import {
+  NotificationLogger,
+  showAndLogExceptionWithTelemetry,
+} from "../common/logging";
+import { CancellationToken } from "vscode";
 import { CodeQLCliServer } from "../codeql-cli/cli";
-import { showLlmGeneration } from "../config";
+import { DatabaseItem } from "../databases/local-databases";
+import { ProgressCallback } from "../common/vscode/progress";
+import { redactableError } from "../common/errors";
+import { telemetryListener } from "../common/vscode/telemetry";
+import { join } from "path";
 import { Mode } from "./shared/mode";
-import { resolveQueriesFromPacks } from "../local-queries";
-import { modeTag } from "./mode-tag";
+import { outputFile, writeFile } from "fs-extra";
+import { QueryLanguage } from "../common/query-language";
+import { fetchExternalApiQueries } from "./queries";
+import { Method } from "./method";
+import { runQuery } from "../local-queries/run-query";
+import { decodeBqrsToMethods } from "./bqrs";
+import {
+  resolveEndpointsQuery,
+  syntheticQueryPackName,
+} from "./model-editor-queries-setup";
 
-export const syntheticQueryPackName = "codeql/external-api-usage";
+type RunQueryOptions = {
+  cliServer: CodeQLCliServer;
+  queryRunner: QueryRunner;
+  logger: NotificationLogger;
+  databaseItem: DatabaseItem;
+  language: QueryLanguage;
+  queryStorageDir: string;
+  queryDir: string;
 
-/**
- * setUpPack sets up a directory to use for the data extension editor queries if required.
- *
- * There are two cases (example language is Java):
- * - In case the queries are present in the codeql/java-queries, we don't need to write our own queries
- *   to disk. We still need to create a synthetic query pack so we can pass the queryDir to the query
- *   resolver without caring about whether the queries are present in the pack or not.
- * - In case the queries are not present in the codeql/java-queries, we need to write our own queries
- *   to disk. We will create a synthetic query pack and install its dependencies so it is fully independent
- *   and we can simply pass it through when resolving the queries.
- *
- * These steps together ensure that later steps of the process don't need to keep track of whether the queries
- * are present in codeql/java-queries or in our own query pack. They just need to resolve the query.
- *
- * @param cliServer The CodeQL CLI server to use.
- * @param queryDir The directory to set up.
- * @param language The language to use for the queries.
- * @returns true if the setup was successful, false otherwise.
- */
-export async function setUpPack(
-  cliServer: CodeQLCliServer,
+  progress: ProgressCallback;
+  token: CancellationToken;
+};
+
+export async function prepareModelEditorQueries(
+  logger: NotificationLogger,
   queryDir: string,
   language: QueryLanguage,
 ): Promise<boolean> {
-  // Download the required query packs
-  await cliServer.packDownload([`codeql/${language}-queries`]);
-
-  // We'll only check if the application mode query exists in the pack and assume that if it does,
-  // the framework mode query will also exist.
-  const applicationModeQuery = await resolveEndpointsQuery(
-    cliServer,
-    language,
-    Mode.Application,
-    [],
-    [],
-  );
-
-  if (applicationModeQuery) {
-    // Set up a synthetic pack so CodeQL doesn't crash later when we try
-    // to resolve a query within this directory
-    const syntheticQueryPack = {
-      name: syntheticQueryPackName,
-      version: "0.0.0",
-      dependencies: {},
-    };
-
-    const qlpackFile = join(queryDir, "codeql-pack.yml");
-    await writeFile(qlpackFile, dump(syntheticQueryPack), "utf8");
-  } else {
-    // If we can't resolve the query, we need to write them to desk ourselves.
-    const externalApiQuerySuccess = await prepareExternalApiQuery(
-      queryDir,
-      language,
+  // Resolve the query that we want to run.
+  const query = fetchExternalApiQueries[language];
+  if (!query) {
+    void showAndLogExceptionWithTelemetry(
+      logger,
+      telemetryListener,
+      redactableError`No bundled model editor query found for language ${language}`,
     );
-    if (!externalApiQuerySuccess) {
-      return false;
+    return false;
+  }
+  // Create the query file.
+  Object.values(Mode).map(async (mode) => {
+    const queryFile = join(queryDir, queryNameFromMode(mode));
+    await writeFile(queryFile, query[`${mode}ModeQuery`], "utf8");
+  });
+
+  // Create any dependencies
+  if (query.dependencies) {
+    for (const [filename, contents] of Object.entries(query.dependencies)) {
+      const dependencyFile = join(queryDir, filename);
+      await outputFile(dependencyFile, contents, "utf8");
     }
-
-    // Set up a synthetic pack so that the query can be resolved later.
-    const syntheticQueryPack = {
-      name: syntheticQueryPackName,
-      version: "0.0.0",
-      dependencies: {
-        [`codeql/${language}-all`]: "*",
-      },
-    };
-
-    const qlpackFile = join(queryDir, "codeql-pack.yml");
-    await writeFile(qlpackFile, dump(syntheticQueryPack), "utf8");
-    await cliServer.packInstall(queryDir);
   }
-
-  // Download any other required packs
-  if (language === "java" && showLlmGeneration()) {
-    await cliServer.packDownload([`codeql/${language}-automodel-queries`]);
-  }
-
   return true;
 }
 
-/**
- * Resolve the query path to the model editor endpoints query. All queries are tagged like this:
- * modeleditor endpoints <mode>
- * Example: modeleditor endpoints framework-mode
- *
- * @param cliServer The CodeQL CLI server to use.
- * @param language The language of the query pack to use.
- * @param mode The mode to resolve the query for.
- * @param additionalPackNames Additional pack names to search.
- * @param additionalPackPaths Additional pack paths to search.
- */
-export async function resolveEndpointsQuery(
-  cliServer: CodeQLCliServer,
-  language: string,
-  mode: Mode,
-  additionalPackNames: string[] = [],
-  additionalPackPaths: string[] = [],
-): Promise<string | undefined> {
-  const packsToSearch = [`codeql/${language}-queries`, ...additionalPackNames];
+export const externalApiQueriesProgressMaxStep = 2000;
 
-  // First, resolve the query that we want to run.
-  // All queries are tagged like this:
-  // internal extract automodel <mode> <queryTag>
-  // Example: internal extract automodel framework-mode candidates
-  const queries = await resolveQueriesFromPacks(
+export async function runModelEditorQueries(
+  mode: Mode,
+  {
     cliServer,
-    packsToSearch,
-    {
-      kind: "table",
-      "tags contain all": ["modeleditor", "endpoints", modeTag(mode)],
-    },
-    additionalPackPaths,
+    queryRunner,
+    logger,
+    databaseItem,
+    language,
+    queryStorageDir,
+    queryDir,
+    progress,
+    token,
+  }: RunQueryOptions,
+): Promise<Method[] | undefined> {
+  // The below code is temporary to allow for rapid prototyping of the queries. Once the queries are stabilized, we will
+  // move these queries into the `github/codeql` repository and use them like any other contextual (e.g. AST) queries.
+  // This is intentionally not pretty code, as it will be removed soon.
+  // For a reference of what this should do in the future, see the previous implementation in
+  // https://github.com/github/vscode-codeql/blob/089d3566ef0bc67d9b7cc66e8fd6740b31c1c0b0/extensions/ql-vscode/src/data-extensions-editor/external-api-usage-query.ts#L33-L72
+
+  progress({
+    message: "Resolving QL packs",
+    step: 1,
+    maxStep: externalApiQueriesProgressMaxStep,
+  });
+  const additionalPacks = getOnDiskWorkspaceFolders();
+  const extensionPacks = Object.keys(
+    await cliServer.resolveQlpacks(additionalPacks, true),
   );
-  if (queries.length > 1) {
-    throw new Error(
-      `Found multiple endpoints queries for ${mode}. Can't continue`,
+
+  progress({
+    message: "Resolving query",
+    step: 2,
+    maxStep: externalApiQueriesProgressMaxStep,
+  });
+
+  // Resolve the queries from either codeql/java-queries or from the temporary queryDir
+  const queryPath = await resolveEndpointsQuery(
+    cliServer,
+    databaseItem.language,
+    mode,
+    [syntheticQueryPackName],
+    [queryDir],
+  );
+  if (!queryPath) {
+    void showAndLogExceptionWithTelemetry(
+      logger,
+      telemetryListener,
+      redactableError`The ${mode} model editor query could not be found. Try re-opening the model editor. If that doesn't work, try upgrading the CodeQL libraries.`,
     );
+    return;
   }
 
-  if (queries.length === 0) {
+  // Run the actual query
+  const completedQuery = await runQuery({
+    queryRunner,
+    databaseItem,
+    queryPath,
+    queryStorageDir,
+    additionalPacks,
+    extensionPacks,
+    progress: (update) =>
+      progress({
+        step: update.step + 500,
+        maxStep: externalApiQueriesProgressMaxStep,
+        message: update.message,
+      }),
+    token,
+  });
+
+  if (!completedQuery) {
+    return;
+  }
+
+  // Read the results and covert to internal representation
+  progress({
+    message: "Decoding results",
+    step: 1600,
+    maxStep: externalApiQueriesProgressMaxStep,
+  });
+
+  const bqrsChunk = await readQueryResults({
+    cliServer,
+    logger,
+    bqrsPath: completedQuery.outputDir.bqrsPath,
+  });
+  if (!bqrsChunk) {
+    return;
+  }
+
+  progress({
+    message: "Finalizing results",
+    step: 1950,
+    maxStep: externalApiQueriesProgressMaxStep,
+  });
+
+  return decodeBqrsToMethods(bqrsChunk, mode, language);
+}
+
+type GetResultsOptions = {
+  cliServer: Pick<CodeQLCliServer, "bqrsInfo" | "bqrsDecode">;
+  logger: NotificationLogger;
+  bqrsPath: string;
+};
+
+export async function readQueryResults({
+  cliServer,
+  logger,
+  bqrsPath,
+}: GetResultsOptions) {
+  const bqrsInfo = await cliServer.bqrsInfo(bqrsPath);
+  if (bqrsInfo["result-sets"].length !== 1) {
+    void showAndLogExceptionWithTelemetry(
+      logger,
+      telemetryListener,
+      redactableError`Expected exactly one result set, got ${bqrsInfo["result-sets"].length}`,
+    );
     return undefined;
   }
 
-  return queries[0];
+  const resultSet = bqrsInfo["result-sets"][0];
+
+  return cliServer.bqrsDecode(bqrsPath, resultSet.name);
+}
+
+function queryNameFromMode(mode: Mode): string {
+  return `${mode.charAt(0).toUpperCase() + mode.slice(1)}ModeEndpoints.ql`;
 }

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

@@ -0,0 +1,33 @@
+import { Method } from "./method";
+
+interface ModelEditorViewInterface {
+  databaseUri: string;
+
+  revealMethod(method: Method): Promise<void>;
+}
+
+export class ModelEditorViewTracker<
+  T extends ModelEditorViewInterface = ModelEditorViewInterface,
+> {
+  private readonly views = new Map<string, T>();
+
+  constructor() {}
+
+  public registerView(view: T): void {
+    const databaseUri = view.databaseUri;
+
+    if (this.views.has(databaseUri)) {
+      throw new Error(`View for database ${databaseUri} already registered`);
+    }
+
+    this.views.set(databaseUri, view);
+  }
+
+  public unregisterView(view: T): void {
+    this.views.delete(view.databaseUri);
+  }
+
+  public getView(databaseUri: string): T | undefined {
+    return this.views.get(databaseUri);
+  }
+}

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

@@ -17,41 +17,51 @@ import {
 import { ProgressCallback, withProgress } from "../common/vscode/progress";
 import { QueryRunner } from "../query-server";
 import {
-  showAndLogExceptionWithTelemetry,
   showAndLogErrorMessage,
+  showAndLogExceptionWithTelemetry,
 } from "../common/logging";
 import { DatabaseItem, DatabaseManager } from "../databases/local-databases";
 import { CodeQLCliServer } from "../codeql-cli/cli";
 import { asError, assertNever, getErrorMessage } from "../common/helpers-pure";
-import { runFlowModelQueries } from "./flow-model-queries";
 import { promptImportGithubDatabase } from "../databases/database-fetcher";
 import { App } from "../common/app";
 import { redactableError } from "../common/errors";
 import {
   externalApiQueriesProgressMaxStep,
-  runExternalApiQueries,
-} from "./external-api-usage-queries";
-import { Method, Usage } from "./method";
+  runModelEditorQueries,
+} from "./model-editor-queries";
+import { Method } from "./method";
 import { ModeledMethod } from "./modeled-method";
 import { ExtensionPack } from "./shared/extension-pack";
-import { showFlowGeneration, showLlmGeneration } from "../config";
+import { ModelConfigListener } from "../config";
 import { Mode } from "./shared/mode";
 import { loadModeledMethods, saveModeledMethods } from "./modeled-method-fs";
 import { pickExtensionPack } from "./extension-pack-picker";
-import { getLanguageDisplayName } from "../common/query-language";
+import {
+  getLanguageDisplayName,
+  QueryLanguage,
+} from "../common/query-language";
 import { AutoModeler } from "./auto-modeler";
 import { telemetryListener } from "../common/vscode/telemetry";
 import { ModelingStore } from "./modeling-store";
+import { ModelEditorViewTracker } from "./model-editor-view-tracker";
+import { ModelingEvents } from "./modeling-events";
+import { getModelsAsDataLanguage, ModelsAsDataLanguage } from "./languages";
+import { runGenerateQueries } from "./generate";
 
 export class ModelEditorView extends AbstractWebview<
   ToModelEditorMessage,
   FromModelEditorMessage
 > {
   private readonly autoModeler: AutoModeler;
+  private readonly languageDefinition: ModelsAsDataLanguage;
 
   public constructor(
     protected readonly app: App,
     private readonly modelingStore: ModelingStore,
+    private readonly modelingEvents: ModelingEvents,
+    private readonly viewTracker: ModelEditorViewTracker<ModelEditorView>,
+    private readonly modelConfig: ModelConfigListener,
     private readonly databaseManager: DatabaseManager,
     private readonly cliServer: CodeQLCliServer,
     private readonly queryRunner: QueryRunner,
@@ -59,30 +69,32 @@ export class ModelEditorView extends AbstractWebview<
     private readonly queryDir: string,
     private readonly databaseItem: DatabaseItem,
     private readonly extensionPack: ExtensionPack,
-    private mode: Mode,
+    // The language is equal to databaseItem.language but is properly typed as QueryLanguage
+    private readonly language: QueryLanguage,
+    initialMode: Mode,
   ) {
     super(app);
 
-    this.modelingStore.initializeStateForDb(databaseItem);
-    this.registerToModelingStoreEvents();
+    this.modelingStore.initializeStateForDb(databaseItem, initialMode);
+    this.registerToModelingEvents();
+    this.registerToModelConfigEvents();
+
+    this.viewTracker.registerView(this);
 
     this.autoModeler = new AutoModeler(
       app,
       cliServer,
       queryRunner,
+      this.modelConfig,
+      modelingStore,
       queryStorageDir,
       databaseItem,
-      async (packageName, inProgressMethods) => {
-        await this.postMessage({
-          t: "setInProgressMethods",
-          packageName,
-          inProgressMethods,
-        });
-      },
+      language,
       async (modeledMethods) => {
         this.addModeledMethods(modeledMethods);
       },
     );
+    this.languageDefinition = getModelsAsDataLanguage(language);
   }
 
   public async openView() {
@@ -92,9 +104,6 @@ export class ModelEditorView extends AbstractWebview<
     panel.onDidChangeViewState(async () => {
       if (panel.active) {
         this.modelingStore.setActiveDb(this.databaseItem);
-        await this.markModelEditorAsActive();
-      } else {
-        await this.updateModelEditorActiveContext();
       }
     });
 
@@ -118,36 +127,12 @@ export class ModelEditorView extends AbstractWebview<
     );
   }
 
-  private async markModelEditorAsActive(): Promise<void> {
-    void this.app.commands.execute(
-      "setContext",
-      "codeql.modelEditorActive",
-      true,
-    );
-  }
-
-  private async updateModelEditorActiveContext(): Promise<void> {
-    await this.app.commands.execute(
-      "setContext",
-      "codeql.modelEditorActive",
-      this.isAModelEditorActive(),
-    );
-  }
-
   private isAModelEditorOpen(): boolean {
     return window.tabGroups.all.some((tabGroup) =>
       tabGroup.tabs.some((tab) => this.isTabModelEditorView(tab)),
     );
   }
 
-  private isAModelEditorActive(): boolean {
-    return window.tabGroups.all.some((tabGroup) =>
-      tabGroup.tabs.some(
-        (tab) => this.isTabModelEditorView(tab) && tab.isActive,
-      ),
-    );
-  }
-
   private isTabModelEditorView(tab: Tab): boolean {
     if (!(tab.input instanceof TabInputWebview)) {
       return false;
@@ -181,7 +166,7 @@ export class ModelEditorView extends AbstractWebview<
   }
 
   protected onPanelDispose(): void {
-    // Nothing to do here
+    this.viewTracker.unregisterView(this);
   }
 
   protected async onMessage(msg: FromModelEditorMessage): Promise<void> {
@@ -214,57 +199,72 @@ export class ModelEditorView extends AbstractWebview<
         );
 
         break;
-      case "jumpToUsage":
-        await this.handleJumpToUsage(msg.method, msg.usage);
-        void telemetryListener?.sendUIInteraction("model-editor-jump-to-usage");
+      case "jumpToMethod":
+        await this.handleJumpToMethod(msg.methodSignature);
+        void telemetryListener?.sendUIInteraction(
+          "model-editor-jump-to-method",
+        );
 
         break;
       case "saveModeledMethods":
-        await withProgress(
-          async (progress) => {
-            progress({
-              step: 1,
-              maxStep: 500 + externalApiQueriesProgressMaxStep,
-              message: "Writing model files",
-            });
-            await saveModeledMethods(
-              this.extensionPack,
-              this.databaseItem.language,
-              msg.methods,
-              msg.modeledMethods,
-              this.mode,
-              this.cliServer,
-              this.app.logger,
-            );
+        {
+          const methods = this.modelingStore.getMethods(
+            this.databaseItem,
+            msg.methodSignatures,
+          );
+          const modeledMethods = this.modelingStore.getModeledMethods(
+            this.databaseItem,
+            msg.methodSignatures,
+          );
+          const mode = this.modelingStore.getMode(this.databaseItem);
 
-            await Promise.all([
-              this.setViewState(),
-              this.loadMethods((update) =>
-                progress({
-                  ...update,
-                  step: update.step + 500,
-                  maxStep: 500 + externalApiQueriesProgressMaxStep,
-                }),
-              ),
-            ]);
-          },
-          {
-            cancellable: false,
-          },
-        );
+          await withProgress(
+            async (progress) => {
+              progress({
+                step: 1,
+                maxStep: 500 + externalApiQueriesProgressMaxStep,
+                message: "Writing model files",
+              });
+              await saveModeledMethods(
+                this.extensionPack,
+                this.language,
+                methods,
+                modeledMethods,
+                mode,
+                this.cliServer,
+                this.app.logger,
+              );
 
-        this.modelingStore.removeModifiedMethods(
-          this.databaseItem,
-          Object.keys(msg.modeledMethods),
-        );
+              await Promise.all([
+                this.setViewState(),
+                this.loadMethods((update) =>
+                  progress({
+                    ...update,
+                    step: update.step + 500,
+                    maxStep: 500 + externalApiQueriesProgressMaxStep,
+                  }),
+                ),
+              ]);
+            },
+            {
+              cancellable: false,
+            },
+          );
 
-        void telemetryListener?.sendUIInteraction(
-          "model-editor-save-modeled-methods",
-        );
+          this.modelingStore.removeModifiedMethods(
+            this.databaseItem,
+            Object.keys(modeledMethods),
+          );
+
+          void telemetryListener?.sendUIInteraction(
+            "model-editor-save-modeled-methods",
+          );
+        }
 
         break;
       case "generateMethod":
         await this.generateModeledMethods();
+
         void telemetryListener?.sendUIInteraction(
           "model-editor-generate-modeled-methods",
         );
@@ -273,8 +273,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",
@@ -293,7 +292,7 @@ export class ModelEditorView extends AbstractWebview<
         );
         break;
       case "switchMode":
-        this.mode = msg.mode;
+        this.modelingStore.setMode(this.databaseItem, msg.mode);
         this.modelingStore.setMethods(this.databaseItem, []);
         await Promise.all([
           this.postMessage({
@@ -317,10 +316,22 @@ export class ModelEditorView extends AbstractWebview<
           "model-editor-hide-modeled-methods",
         );
         break;
-      case "setModeledMethod": {
-        this.setModeledMethod(msg.method);
+      case "setMultipleModeledMethods": {
+        this.setModeledMethods(msg.methodSignature, msg.modeledMethods);
         break;
       }
+      case "telemetry":
+        telemetryListener?.sendUIInteraction(msg.action);
+        break;
+      case "unhandledError":
+        void showAndLogExceptionWithTelemetry(
+          this.app.logger,
+          telemetryListener,
+          redactableError(
+            msg.error,
+          )`Unhandled error in model editor view: ${msg.error.message}`,
+        );
+        break;
       default:
         assertNever(msg);
     }
@@ -338,29 +349,73 @@ export class ModelEditorView extends AbstractWebview<
     ]);
   }
 
+  public get databaseUri(): string {
+    return this.databaseItem.databaseUri.toString();
+  }
+
+  public async focusView(): Promise<void> {
+    this.panel?.reveal();
+  }
+
+  public async revealMethod(method: Method): Promise<void> {
+    this.panel?.reveal();
+
+    await this.postMessage({
+      t: "revealMethod",
+      methodSignature: method.signature,
+    });
+  }
+
   private async setViewState(): Promise<void> {
+    const modelsAsDataLanguage = getModelsAsDataLanguage(this.language);
+
+    const showGenerateButton =
+      this.modelConfig.flowGeneration && !!modelsAsDataLanguage.modelGeneration;
+
     const showLlmButton =
-      this.databaseItem.language === "java" && showLlmGeneration();
+      this.databaseItem.language === "java" && this.modelConfig.llmGeneration;
+
+    const sourceArchiveAvailable =
+      this.databaseItem.hasSourceArchiveInExplorer();
+
+    const showModeSwitchButton =
+      this.languageDefinition.availableModes === undefined ||
+      this.languageDefinition.availableModes.length > 1;
 
     await this.postMessage({
       t: "setModelEditorViewState",
       viewState: {
         extensionPack: this.extensionPack,
-        showFlowGeneration: showFlowGeneration(),
+        language: this.language,
+        showGenerateButton,
         showLlmButton,
-        mode: this.mode,
+        showMultipleModels: this.modelConfig.showMultipleModels,
+        mode: this.modelingStore.getMode(this.databaseItem),
+        showModeSwitchButton,
+        sourceArchiveAvailable,
       },
     });
   }
 
-  protected async handleJumpToUsage(method: Method, usage: Usage) {
-    this.modelingStore.setSelectedMethod(this.databaseItem, method, usage);
+  protected async handleJumpToMethod(methodSignature: string) {
+    const method = this.modelingStore.getMethod(
+      this.databaseItem,
+      methodSignature,
+    );
+    if (method) {
+      this.modelingStore.setSelectedMethod(
+        this.databaseItem,
+        method,
+        method.usages[0],
+      );
+    }
   }
 
   protected async loadExistingModeledMethods(): Promise<void> {
     try {
       const modeledMethods = await loadModeledMethods(
         this.extensionPack,
+        this.language,
         this.cliServer,
         this.app.logger,
       );
@@ -374,12 +429,16 @@ export class ModelEditorView extends AbstractWebview<
   }
 
   protected async loadMethods(progress: ProgressCallback): Promise<void> {
+    const mode = this.modelingStore.getMode(this.databaseItem);
+
     try {
       const cancellationTokenSource = new CancellationTokenSource();
-      const queryResult = await runExternalApiQueries(this.mode, {
+      const queryResult = await runModelEditorQueries(mode, {
         cliServer: this.cliServer,
         queryRunner: this.queryRunner,
+        logger: this.app.logger,
         databaseItem: this.databaseItem,
+        language: this.language,
         queryStorageDir: this.queryStorageDir,
         queryDir: this.queryDir,
         progress: (update) =>
@@ -398,9 +457,9 @@ export class ModelEditorView extends AbstractWebview<
       void showAndLogExceptionWithTelemetry(
         this.app.logger,
         this.app.telemetry,
-        redactableError(
-          asError(err),
-        )`Failed to load external API usages: ${getErrorMessage(err)}`,
+        redactableError(asError(err))`Failed to load results: ${getErrorMessage(
+          err,
+        )}`,
       );
     }
   }
@@ -410,17 +469,37 @@ export class ModelEditorView extends AbstractWebview<
       async (progress) => {
         const tokenSource = new CancellationTokenSource();
 
+        const mode = this.modelingStore.getMode(this.databaseItem);
+
+        const modelsAsDataLanguage = getModelsAsDataLanguage(this.language);
+        const modelGeneration = modelsAsDataLanguage.modelGeneration;
+        if (!modelGeneration) {
+          void showAndLogErrorMessage(
+            this.app.logger,
+            `Model generation is not supported for ${this.language}.`,
+          );
+          return;
+        }
+
         let addedDatabase: DatabaseItem | undefined;
 
         // In application mode, we need the database of a specific library to generate
         // the modeled methods. In framework mode, we'll use the current database.
-        if (this.mode === Mode.Application) {
+        if (mode === Mode.Application) {
           addedDatabase = await this.promptChooseNewOrExistingDatabase(
             progress,
           );
           if (!addedDatabase) {
             return;
           }
+
+          if (addedDatabase.language !== this.language) {
+            void showAndLogErrorMessage(
+              this.app.logger,
+              `The selected database is for ${addedDatabase.language}, but the current database is for ${this.language}.`,
+            );
+            return;
+          }
         }
 
         progress({
@@ -430,20 +509,23 @@ export class ModelEditorView extends AbstractWebview<
         });
 
         try {
-          await runFlowModelQueries({
+          await runGenerateQueries({
+            queryConstraints: modelGeneration.queryConstraints,
+            filterQueries: modelGeneration.filterQueries,
+            parseResults: (queryPath, results) =>
+              modelGeneration.parseResults(
+                queryPath,
+                results,
+                modelsAsDataLanguage,
+                this.app.logger,
+              ),
+            onResults: async (modeledMethods) => {
+              this.addModeledMethodsFromArray(modeledMethods);
+            },
             cliServer: this.cliServer,
             queryRunner: this.queryRunner,
             queryStorageDir: this.queryStorageDir,
             databaseItem: addedDatabase ?? this.databaseItem,
-            onResults: async (modeledMethods) => {
-              const modeledMethodsByName: Record<string, ModeledMethod> = {};
-
-              for (const modeledMethod of modeledMethods) {
-                modeledMethodsByName[modeledMethod.signature] = modeledMethod;
-              }
-
-              this.addModeledMethods(modeledMethodsByName);
-            },
             progress,
             token: tokenSource.token,
           });
@@ -453,7 +535,7 @@ export class ModelEditorView extends AbstractWebview<
             this.app.telemetry,
             redactableError(
               asError(e),
-            )`Failed to generate flow model: ${getErrorMessage(e)}`,
+            )`Failed to generate models: ${getErrorMessage(e)}`,
           );
         }
       },
@@ -463,14 +545,22 @@ 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,
+    );
+    const mode = this.modelingStore.getMode(this.databaseItem);
     await this.autoModeler.startModeling(
       packageName,
       methods,
       modeledMethods,
-      this.mode,
+      mode,
     );
   }
 
@@ -483,9 +573,19 @@ export class ModelEditorView extends AbstractWebview<
         return;
       }
 
+      let existingView = this.viewTracker.getView(
+        addedDatabase.databaseUri.toString(),
+      );
+      if (existingView) {
+        await existingView.focusView();
+
+        return;
+      }
+
       const modelFile = await pickExtensionPack(
         this.cliServer,
         addedDatabase,
+        this.modelConfig,
         this.app.logger,
         progress,
         3,
@@ -494,9 +594,23 @@ export class ModelEditorView extends AbstractWebview<
         return;
       }
 
+      // Check again just before opening the editor to ensure no model editor has been opened between
+      // our first check and now.
+      existingView = this.viewTracker.getView(
+        addedDatabase.databaseUri.toString(),
+      );
+      if (existingView) {
+        await existingView.focusView();
+
+        return;
+      }
+
       const view = new ModelEditorView(
         this.app,
         this.modelingStore,
+        this.modelingEvents,
+        this.viewTracker,
+        this.modelConfig,
         this.databaseManager,
         this.cliServer,
         this.queryRunner,
@@ -504,6 +618,7 @@ export class ModelEditorView extends AbstractWebview<
         this.queryDir,
         addedDatabase,
         modelFile,
+        this.language,
         Mode.Framework,
       );
       await view.openView();
@@ -583,9 +698,9 @@ export class ModelEditorView extends AbstractWebview<
     return addedDatabase;
   }
 
-  private registerToModelingStoreEvents() {
+  private registerToModelingEvents() {
     this.push(
-      this.modelingStore.onMethodsChanged(async (event) => {
+      this.modelingEvents.onMethodsChanged(async (event) => {
         if (event.dbUri === this.databaseItem.databaseUri.toString()) {
           await this.postMessage({
             t: "setMethods",
@@ -596,7 +711,7 @@ export class ModelEditorView extends AbstractWebview<
     );
 
     this.push(
-      this.modelingStore.onModeledMethodsChanged(async (event) => {
+      this.modelingEvents.onModeledMethodsChanged(async (event) => {
         if (event.dbUri === this.databaseItem.databaseUri.toString()) {
           await this.postMessage({
             t: "setModeledMethods",
@@ -607,7 +722,7 @@ export class ModelEditorView extends AbstractWebview<
     );
 
     this.push(
-      this.modelingStore.onModifiedMethodsChanged(async (event) => {
+      this.modelingEvents.onModifiedMethodsChanged(async (event) => {
         if (event.dbUri === this.databaseItem.databaseUri.toString()) {
           await this.postMessage({
             t: "setModifiedMethods",
@@ -616,9 +731,28 @@ export class ModelEditorView extends AbstractWebview<
         }
       }),
     );
+
+    this.push(
+      this.modelingEvents.onInProgressMethodsChanged(async (event) => {
+        if (event.dbUri === this.databaseItem.databaseUri.toString()) {
+          await this.postMessage({
+            t: "setInProgressMethods",
+            methods: Array.from(event.methods),
+          });
+        }
+      }),
+    );
   }
 
-  private addModeledMethods(modeledMethods: Record<string, ModeledMethod>) {
+  private registerToModelConfigEvents() {
+    this.push(
+      this.modelConfig.onDidChangeConfiguration(() => {
+        void this.setViewState();
+      }),
+    );
+  }
+
+  private addModeledMethods(modeledMethods: Record<string, ModeledMethod[]>) {
     this.modelingStore.addModeledMethods(this.databaseItem, modeledMethods);
 
     this.modelingStore.addModifiedMethods(
@@ -627,13 +761,26 @@ export class ModelEditorView extends AbstractWebview<
     );
   }
 
-  private setModeledMethod(method: ModeledMethod) {
-    const state = this.modelingStore.getStateForActiveDb();
-    if (!state) {
-      throw new Error("Attempting to set modeled method without active db");
+  private addModeledMethodsFromArray(modeledMethods: ModeledMethod[]) {
+    const modeledMethodsByName: Record<string, ModeledMethod[]> = {};
+
+    for (const modeledMethod of modeledMethods) {
+      if (!(modeledMethod.signature in modeledMethodsByName)) {
+        modeledMethodsByName[modeledMethod.signature] = [];
+      }
+
+      modeledMethodsByName[modeledMethod.signature].push(modeledMethod);
     }
 
-    this.modelingStore.updateModeledMethod(state.databaseItem, method);
-    this.modelingStore.addModifiedMethod(state.databaseItem, method.signature);
+    this.addModeledMethods(modeledMethodsByName);
+  }
+
+  private setModeledMethods(signature: string, methods: ModeledMethod[]) {
+    this.modelingStore.updateModeledMethods(
+      this.databaseItem,
+      signature,
+      methods,
+    );
+    this.modelingStore.addModifiedMethod(this.databaseItem, signature);
   }
 }

extensions/ql-vscode/src/model-editor/modeled-method-empty.ts

@@ -0,0 +1,88 @@
+import { ModeledMethod, SinkModeledMethod } from "./modeled-method";
+import { MethodSignature } from "./method";
+import { assertNever } from "../common/helpers-pure";
+
+export function createEmptyModeledMethod(
+  type: ModeledMethod["type"],
+  methodSignature: MethodSignature,
+) {
+  const canonicalMethodSignature: MethodSignature = {
+    packageName: methodSignature.packageName,
+    typeName: methodSignature.typeName,
+    methodName: methodSignature.methodName,
+    methodParameters: methodSignature.methodParameters,
+    signature: methodSignature.signature,
+  };
+
+  switch (type) {
+    case "none":
+      return createEmptyNoneModeledMethod(canonicalMethodSignature);
+    case "source":
+      return createEmptySourceModeledMethod(canonicalMethodSignature);
+    case "sink":
+      return createEmptySinkModeledMethod(canonicalMethodSignature);
+    case "summary":
+      return createEmptySummaryModeledMethod(canonicalMethodSignature);
+    case "neutral":
+      return createEmptyNeutralModeledMethod(canonicalMethodSignature);
+    default:
+      assertNever(type);
+  }
+}
+
+function createEmptyNoneModeledMethod(
+  methodSignature: MethodSignature,
+): ModeledMethod {
+  return {
+    ...methodSignature,
+    type: "none",
+  };
+}
+
+function createEmptySourceModeledMethod(
+  methodSignature: MethodSignature,
+): ModeledMethod {
+  return {
+    ...methodSignature,
+    type: "source",
+    output: "",
+    kind: "",
+    provenance: "manual",
+  };
+}
+
+function createEmptySinkModeledMethod(
+  methodSignature: MethodSignature,
+): SinkModeledMethod {
+  return {
+    ...methodSignature,
+    type: "sink",
+    input: "",
+    kind: "",
+    provenance: "manual",
+  };
+}
+
+function createEmptySummaryModeledMethod(
+  methodSignature: MethodSignature,
+): ModeledMethod {
+  return {
+    ...methodSignature,
+    type: "summary",
+    input: "",
+    output: "",
+    kind: "",
+    provenance: "manual",
+  };
+}
+
+function createEmptyNeutralModeledMethod(
+  methodSignature: MethodSignature,
+): ModeledMethod {
+  return {
+    ...methodSignature,
+    type: "neutral",
+    kind: "",
+    provenance: "manual",
+  };
+}

extensions/ql-vscode/src/model-editor/modeled-method-fs.ts

@@ -10,18 +10,20 @@ import { getOnDiskWorkspaceFolders } from "../common/vscode/workspace-folders";
 import { load as loadYaml } from "js-yaml";
 import { CodeQLCliServer } from "../codeql-cli/cli";
 import { pathsEqual } from "../common/files";
+import { QueryLanguage } from "../common/query-language";
 
 export async function saveModeledMethods(
   extensionPack: ExtensionPack,
-  language: string,
-  methods: Method[],
-  modeledMethods: Record<string, ModeledMethod>,
+  language: QueryLanguage,
+  methods: readonly Method[],
+  modeledMethods: Readonly<Record<string, readonly ModeledMethod[]>>,
   mode: Mode,
   cliServer: CodeQLCliServer,
   logger: NotificationLogger,
 ): Promise<void> {
   const existingModeledMethods = await loadModeledMethodFiles(
     extensionPack,
+    language,
     cliServer,
     logger,
   );
@@ -43,14 +45,15 @@ export async function saveModeledMethods(
 
 async function loadModeledMethodFiles(
   extensionPack: ExtensionPack,
+  language: QueryLanguage,
   cliServer: CodeQLCliServer,
   logger: NotificationLogger,
-): Promise<Record<string, Record<string, ModeledMethod>>> {
+): Promise<Record<string, Record<string, ModeledMethod[]>>> {
   const modelFiles = await listModelFiles(extensionPack.path, cliServer);
 
   const modeledMethodsByFile: Record<
     string,
-    Record<string, ModeledMethod>
+    Record<string, ModeledMethod[]>
   > = {};
 
   for (const modelFile of modelFiles) {
@@ -60,7 +63,7 @@ async function loadModeledMethodFiles(
       filename: modelFile,
     });
 
-    const modeledMethods = loadDataExtensionYaml(data);
+    const modeledMethods = loadDataExtensionYaml(data, language);
     if (!modeledMethods) {
       void showAndLogErrorMessage(
         logger,
@@ -76,19 +79,25 @@ async function loadModeledMethodFiles(
 
 export async function loadModeledMethods(
   extensionPack: ExtensionPack,
+  language: QueryLanguage,
   cliServer: CodeQLCliServer,
   logger: NotificationLogger,
-): Promise<Record<string, ModeledMethod>> {
-  const existingModeledMethods: Record<string, ModeledMethod> = {};
+): Promise<Record<string, ModeledMethod[]>> {
+  const existingModeledMethods: Record<string, ModeledMethod[]> = {};
 
   const modeledMethodsByFile = await loadModeledMethodFiles(
     extensionPack,
+    language,
     cliServer,
     logger,
   );
   for (const modeledMethods of Object.values(modeledMethodsByFile)) {
     for (const [key, value] of Object.entries(modeledMethods)) {
-      existingModeledMethods[key] = value;
+      if (!(key in existingModeledMethods)) {
+        existingModeledMethods[key] = [];
+      }
+
+      existingModeledMethods[key].push(...value);
     }
   }
 

extensions/ql-vscode/src/model-editor/modeled-method.ts

@@ -19,12 +19,85 @@ export type Provenance =
   // Entered by the user in the editor manually
   | "manual";
 
-export interface ModeledMethod extends MethodSignature {
-  type: ModeledMethodType;
-  input: string;
-  output: string;
-  kind: ModeledMethodKind;
-  provenance: Provenance;
+export interface NoneModeledMethod extends MethodSignature {
+  readonly type: "none";
 }
 
+export interface SourceModeledMethod extends MethodSignature {
+  readonly type: "source";
+  readonly output: string;
+  readonly kind: ModeledMethodKind;
+  readonly provenance: Provenance;
+}
+
+export interface SinkModeledMethod extends MethodSignature {
+  readonly type: "sink";
+  readonly input: string;
+  readonly kind: ModeledMethodKind;
+  readonly provenance: Provenance;
+}
+
+export interface SummaryModeledMethod extends MethodSignature {
+  readonly type: "summary";
+  readonly input: string;
+  readonly output: string;
+  readonly kind: ModeledMethodKind;
+  readonly provenance: Provenance;
+}
+
+export interface NeutralModeledMethod extends MethodSignature {
+  readonly type: "neutral";
+  readonly kind: ModeledMethodKind;
+  readonly provenance: Provenance;
+}
+
+export type ModeledMethod =
+  | NoneModeledMethod
+  | SourceModeledMethod
+  | SinkModeledMethod
+  | SummaryModeledMethod
+  | NeutralModeledMethod;
+
 export type ModeledMethodKind = string;
+
+export function modeledMethodSupportsKind(
+  modeledMethod: ModeledMethod,
+): modeledMethod is
+  | SourceModeledMethod
+  | SinkModeledMethod
+  | SummaryModeledMethod
+  | NeutralModeledMethod {
+  return (
+    modeledMethod.type === "source" ||
+    modeledMethod.type === "sink" ||
+    modeledMethod.type === "summary" ||
+    modeledMethod.type === "neutral"
+  );
+}
+
+export function modeledMethodSupportsInput(
+  modeledMethod: ModeledMethod,
+): modeledMethod is SinkModeledMethod | SummaryModeledMethod {
+  return modeledMethod.type === "sink" || modeledMethod.type === "summary";
+}
+
+export function modeledMethodSupportsOutput(
+  modeledMethod: ModeledMethod,
+): modeledMethod is SourceModeledMethod | SummaryModeledMethod {
+  return modeledMethod.type === "source" || modeledMethod.type === "summary";
+}
+
+export function modeledMethodSupportsProvenance(
+  modeledMethod: ModeledMethod,
+): modeledMethod is
+  | SourceModeledMethod
+  | SinkModeledMethod
+  | SummaryModeledMethod
+  | NeutralModeledMethod {
+  return (
+    modeledMethod.type === "source" ||
+    modeledMethod.type === "sink" ||
+    modeledMethod.type === "summary" ||
+    modeledMethod.type === "neutral"
+  );
+}

extensions/ql-vscode/src/model-editor/modeling-events.ts

@@ -0,0 +1,223 @@
+import { App } from "../common/app";
+import { DisposableObject } from "../common/disposable-object";
+import { AppEvent, AppEventEmitter } from "../common/events";
+import { DatabaseItem } from "../databases/local-databases";
+import { Method, Usage } from "./method";
+import { ModeledMethod } from "./modeled-method";
+import { Mode } from "./shared/mode";
+
+interface MethodsChangedEvent {
+  readonly methods: readonly Method[];
+  readonly dbUri: string;
+  readonly isActiveDb: boolean;
+}
+
+interface HideModeledMethodsChangedEvent {
+  readonly hideModeledMethods: boolean;
+  readonly isActiveDb: boolean;
+}
+
+interface ModeChangedEvent {
+  readonly mode: Mode;
+  readonly isActiveDb: boolean;
+}
+
+interface ModeledMethodsChangedEvent {
+  readonly modeledMethods: Readonly<Record<string, ModeledMethod[]>>;
+  readonly dbUri: string;
+  readonly isActiveDb: boolean;
+}
+
+interface ModifiedMethodsChangedEvent {
+  readonly modifiedMethods: ReadonlySet<string>;
+  readonly dbUri: string;
+  readonly isActiveDb: boolean;
+}
+
+interface SelectedMethodChangedEvent {
+  readonly databaseItem: DatabaseItem;
+  readonly method: Method;
+  readonly usage: Usage;
+  readonly modeledMethods: readonly ModeledMethod[];
+  readonly isModified: boolean;
+  readonly isInProgress: boolean;
+}
+
+interface InProgressMethodsChangedEvent {
+  readonly dbUri: string;
+  readonly methods: ReadonlySet<string>;
+}
+
+export class ModelingEvents extends DisposableObject {
+  public readonly onActiveDbChanged: AppEvent<void>;
+  public readonly onDbOpened: AppEvent<DatabaseItem>;
+  public readonly onDbClosed: AppEvent<string>;
+  public readonly onMethodsChanged: AppEvent<MethodsChangedEvent>;
+  public readonly onHideModeledMethodsChanged: AppEvent<HideModeledMethodsChangedEvent>;
+  public readonly onModeChanged: AppEvent<ModeChangedEvent>;
+  public readonly onModeledMethodsChanged: AppEvent<ModeledMethodsChangedEvent>;
+  public readonly onModifiedMethodsChanged: AppEvent<ModifiedMethodsChangedEvent>;
+  public readonly onSelectedMethodChanged: AppEvent<SelectedMethodChangedEvent>;
+  public readonly onInProgressMethodsChanged: AppEvent<InProgressMethodsChangedEvent>;
+
+  private readonly onActiveDbChangedEventEmitter: AppEventEmitter<void>;
+  private readonly onDbOpenedEventEmitter: AppEventEmitter<DatabaseItem>;
+  private readonly onDbClosedEventEmitter: AppEventEmitter<string>;
+  private readonly onMethodsChangedEventEmitter: AppEventEmitter<MethodsChangedEvent>;
+  private readonly onHideModeledMethodsChangedEventEmitter: AppEventEmitter<HideModeledMethodsChangedEvent>;
+  private readonly onModeChangedEventEmitter: AppEventEmitter<ModeChangedEvent>;
+  private readonly onModeledMethodsChangedEventEmitter: AppEventEmitter<ModeledMethodsChangedEvent>;
+  private readonly onModifiedMethodsChangedEventEmitter: AppEventEmitter<ModifiedMethodsChangedEvent>;
+  private readonly onSelectedMethodChangedEventEmitter: AppEventEmitter<SelectedMethodChangedEvent>;
+  private readonly onInProgressMethodsChangedEventEmitter: AppEventEmitter<InProgressMethodsChangedEvent>;
+
+  constructor(app: App) {
+    super();
+
+    this.onActiveDbChangedEventEmitter = this.push(
+      app.createEventEmitter<void>(),
+    );
+    this.onActiveDbChanged = this.onActiveDbChangedEventEmitter.event;
+
+    this.onDbOpenedEventEmitter = this.push(
+      app.createEventEmitter<DatabaseItem>(),
+    );
+    this.onDbOpened = this.onDbOpenedEventEmitter.event;
+
+    this.onDbClosedEventEmitter = this.push(app.createEventEmitter<string>());
+    this.onDbClosed = this.onDbClosedEventEmitter.event;
+
+    this.onMethodsChangedEventEmitter = this.push(
+      app.createEventEmitter<MethodsChangedEvent>(),
+    );
+    this.onMethodsChanged = this.onMethodsChangedEventEmitter.event;
+
+    this.onHideModeledMethodsChangedEventEmitter = this.push(
+      app.createEventEmitter<HideModeledMethodsChangedEvent>(),
+    );
+    this.onHideModeledMethodsChanged =
+      this.onHideModeledMethodsChangedEventEmitter.event;
+
+    this.onModeChangedEventEmitter = this.push(
+      app.createEventEmitter<ModeChangedEvent>(),
+    );
+    this.onModeChanged = this.onModeChangedEventEmitter.event;
+
+    this.onModeledMethodsChangedEventEmitter = this.push(
+      app.createEventEmitter<ModeledMethodsChangedEvent>(),
+    );
+    this.onModeledMethodsChanged =
+      this.onModeledMethodsChangedEventEmitter.event;
+
+    this.onModifiedMethodsChangedEventEmitter = this.push(
+      app.createEventEmitter<ModifiedMethodsChangedEvent>(),
+    );
+    this.onModifiedMethodsChanged =
+      this.onModifiedMethodsChangedEventEmitter.event;
+
+    this.onSelectedMethodChangedEventEmitter = this.push(
+      app.createEventEmitter<SelectedMethodChangedEvent>(),
+    );
+    this.onSelectedMethodChanged =
+      this.onSelectedMethodChangedEventEmitter.event;
+
+    this.onInProgressMethodsChangedEventEmitter = this.push(
+      app.createEventEmitter<InProgressMethodsChangedEvent>(),
+    );
+    this.onInProgressMethodsChanged =
+      this.onInProgressMethodsChangedEventEmitter.event;
+  }
+
+  public fireActiveDbChangedEvent() {
+    this.onActiveDbChangedEventEmitter.fire();
+  }
+
+  public fireDbOpenedEvent(databaseItem: DatabaseItem) {
+    this.onDbOpenedEventEmitter.fire(databaseItem);
+  }
+
+  public fireDbClosedEvent(dbUri: string) {
+    this.onDbClosedEventEmitter.fire(dbUri);
+  }
+
+  public fireMethodsChangedEvent(
+    methods: Method[],
+    dbUri: string,
+    isActiveDb: boolean,
+  ) {
+    this.onMethodsChangedEventEmitter.fire({
+      methods,
+      dbUri,
+      isActiveDb,
+    });
+  }
+
+  public fireHideModeledMethodsChangedEvent(
+    hideModeledMethods: boolean,
+    isActiveDb: boolean,
+  ) {
+    this.onHideModeledMethodsChangedEventEmitter.fire({
+      hideModeledMethods,
+      isActiveDb,
+    });
+  }
+
+  public fireModeChangedEvent(mode: Mode, isActiveDb: boolean) {
+    this.onModeChangedEventEmitter.fire({
+      mode,
+      isActiveDb,
+    });
+  }
+
+  public fireModeledMethodsChangedEvent(
+    modeledMethods: Record<string, ModeledMethod[]>,
+    dbUri: string,
+    isActiveDb: boolean,
+  ) {
+    this.onModeledMethodsChangedEventEmitter.fire({
+      modeledMethods,
+      dbUri,
+      isActiveDb,
+    });
+  }
+
+  public fireModifiedMethodsChangedEvent(
+    modifiedMethods: ReadonlySet<string>,
+    dbUri: string,
+    isActiveDb: boolean,
+  ) {
+    this.onModifiedMethodsChangedEventEmitter.fire({
+      modifiedMethods,
+      dbUri,
+      isActiveDb,
+    });
+  }
+
+  public fireSelectedMethodChangedEvent(
+    databaseItem: DatabaseItem,
+    method: Method,
+    usage: Usage,
+    modeledMethods: ModeledMethod[],
+    isModified: boolean,
+    isInProgress: boolean,
+  ) {
+    this.onSelectedMethodChangedEventEmitter.fire({
+      databaseItem,
+      method,
+      usage,
+      modeledMethods,
+      isModified,
+      isInProgress,
+    });
+  }
+
+  public fireInProgressMethodsChangedEvent(
+    dbUri: string,
+    methods: ReadonlySet<string>,
+  ) {
+    this.onInProgressMethodsChangedEventEmitter.fire({
+      dbUri,
+      methods,
+    });
+  }
+}

extensions/ql-vscode/src/model-editor/modeling-store.ts

@@ -1,133 +1,75 @@
-import { App } from "../common/app";
 import { DisposableObject } from "../common/disposable-object";
-import { AppEvent, AppEventEmitter } from "../common/events";
 import { DatabaseItem } from "../databases/local-databases";
 import { Method, Usage } from "./method";
 import { ModeledMethod } from "./modeled-method";
+import { ModelingEvents } from "./modeling-events";
 import { INITIAL_HIDE_MODELED_METHODS_VALUE } from "./shared/hide-modeled-methods";
+import { Mode } from "./shared/mode";
 
-interface DbModelingState {
+interface InternalDbModelingState {
   databaseItem: DatabaseItem;
   methods: Method[];
   hideModeledMethods: boolean;
-  modeledMethods: Record<string, ModeledMethod>;
+  mode: Mode;
+  modeledMethods: Record<string, ModeledMethod[]>;
   modifiedMethodSignatures: Set<string>;
+  inProgressMethods: Set<string>;
   selectedMethod: Method | undefined;
   selectedUsage: Usage | undefined;
 }
 
-interface MethodsChangedEvent {
-  methods: Method[];
-  dbUri: string;
-  isActiveDb: boolean;
+interface DbModelingState {
+  readonly databaseItem: DatabaseItem;
+  readonly methods: readonly Method[];
+  readonly hideModeledMethods: boolean;
+  readonly mode: Mode;
+  readonly modeledMethods: Readonly<Record<string, readonly ModeledMethod[]>>;
+  readonly modifiedMethodSignatures: ReadonlySet<string>;
+  readonly inProgressMethods: ReadonlySet<string>;
+  readonly selectedMethod: Method | undefined;
+  readonly selectedUsage: Usage | undefined;
 }
 
-interface HideModeledMethodsChangedEvent {
-  hideModeledMethods: boolean;
-  isActiveDb: boolean;
-}
-
-interface ModeledMethodsChangedEvent {
-  modeledMethods: Record<string, ModeledMethod>;
-  dbUri: string;
-  isActiveDb: boolean;
-}
-
-interface ModifiedMethodsChangedEvent {
-  modifiedMethods: Set<string>;
-  dbUri: string;
-  isActiveDb: boolean;
-}
-
-interface SelectedMethodChangedEvent {
-  databaseItem: DatabaseItem;
-  method: Method;
-  usage: Usage;
-  modeledMethod: ModeledMethod | undefined;
-  isModified: boolean;
+interface SelectedMethodDetails {
+  readonly databaseItem: DatabaseItem;
+  readonly method: Method;
+  readonly usage: Usage | undefined;
+  readonly modeledMethods: readonly ModeledMethod[];
+  readonly isModified: boolean;
+  readonly isInProgress: boolean;
 }
 
 export class ModelingStore extends DisposableObject {
-  public readonly onActiveDbChanged: AppEvent<void>;
-  public readonly onDbClosed: AppEvent<string>;
-  public readonly onMethodsChanged: AppEvent<MethodsChangedEvent>;
-  public readonly onHideModeledMethodsChanged: AppEvent<HideModeledMethodsChangedEvent>;
-  public readonly onModeledMethodsChanged: AppEvent<ModeledMethodsChangedEvent>;
-  public readonly onModifiedMethodsChanged: AppEvent<ModifiedMethodsChangedEvent>;
-  public readonly onSelectedMethodChanged: AppEvent<SelectedMethodChangedEvent>;
-
-  private readonly state: Map<string, DbModelingState>;
+  private readonly state: Map<string, InternalDbModelingState>;
   private activeDb: string | undefined;
 
-  private readonly onActiveDbChangedEventEmitter: AppEventEmitter<void>;
-  private readonly onDbClosedEventEmitter: AppEventEmitter<string>;
-  private readonly onMethodsChangedEventEmitter: AppEventEmitter<MethodsChangedEvent>;
-  private readonly onHideModeledMethodsChangedEventEmitter: AppEventEmitter<HideModeledMethodsChangedEvent>;
-  private readonly onModeledMethodsChangedEventEmitter: AppEventEmitter<ModeledMethodsChangedEvent>;
-  private readonly onModifiedMethodsChangedEventEmitter: AppEventEmitter<ModifiedMethodsChangedEvent>;
-  private readonly onSelectedMethodChangedEventEmitter: AppEventEmitter<SelectedMethodChangedEvent>;
-
-  constructor(app: App) {
+  constructor(private readonly modelingEvents: ModelingEvents) {
     super();
 
     // State initialization
-    this.state = new Map<string, DbModelingState>();
-
-    // Event initialization
-    this.onActiveDbChangedEventEmitter = this.push(
-      app.createEventEmitter<void>(),
-    );
-    this.onActiveDbChanged = this.onActiveDbChangedEventEmitter.event;
-
-    this.onDbClosedEventEmitter = this.push(app.createEventEmitter<string>());
-    this.onDbClosed = this.onDbClosedEventEmitter.event;
-
-    this.onMethodsChangedEventEmitter = this.push(
-      app.createEventEmitter<MethodsChangedEvent>(),
-    );
-    this.onMethodsChanged = this.onMethodsChangedEventEmitter.event;
-
-    this.onHideModeledMethodsChangedEventEmitter = this.push(
-      app.createEventEmitter<HideModeledMethodsChangedEvent>(),
-    );
-    this.onHideModeledMethodsChanged =
-      this.onHideModeledMethodsChangedEventEmitter.event;
-
-    this.onModeledMethodsChangedEventEmitter = this.push(
-      app.createEventEmitter<ModeledMethodsChangedEvent>(),
-    );
-    this.onModeledMethodsChanged =
-      this.onModeledMethodsChangedEventEmitter.event;
-
-    this.onModifiedMethodsChangedEventEmitter = this.push(
-      app.createEventEmitter<ModifiedMethodsChangedEvent>(),
-    );
-    this.onModifiedMethodsChanged =
-      this.onModifiedMethodsChangedEventEmitter.event;
-
-    this.onSelectedMethodChangedEventEmitter = this.push(
-      app.createEventEmitter<SelectedMethodChangedEvent>(),
-    );
-    this.onSelectedMethodChanged =
-      this.onSelectedMethodChangedEventEmitter.event;
+    this.state = new Map<string, InternalDbModelingState>();
   }
 
-  public initializeStateForDb(databaseItem: DatabaseItem) {
+  public initializeStateForDb(databaseItem: DatabaseItem, mode: Mode) {
     const dbUri = databaseItem.databaseUri.toString();
     this.state.set(dbUri, {
       databaseItem,
       methods: [],
       hideModeledMethods: INITIAL_HIDE_MODELED_METHODS_VALUE,
+      mode,
       modeledMethods: {},
       modifiedMethodSignatures: new Set(),
       selectedMethod: undefined,
       selectedUsage: undefined,
+      inProgressMethods: new Set(),
     });
+
+    this.modelingEvents.fireDbOpenedEvent(databaseItem);
   }
 
   public setActiveDb(databaseItem: DatabaseItem) {
     this.activeDb = databaseItem.databaseUri.toString();
-    this.onActiveDbChangedEventEmitter.fire();
+    this.modelingEvents.fireActiveDbChangedEvent();
   }
 
   public removeDb(databaseItem: DatabaseItem) {
@@ -139,11 +81,11 @@ export class ModelingStore extends DisposableObject {
 
     if (this.activeDb === dbUri) {
       this.activeDb = undefined;
-      this.onActiveDbChangedEventEmitter.fire();
+      this.modelingEvents.fireActiveDbChangedEvent();
     }
 
     this.state.delete(dbUri);
-    this.onDbClosedEventEmitter.fire(dbUri);
+    this.modelingEvents.fireDbClosedEvent(dbUri);
   }
 
   public getStateForActiveDb(): DbModelingState | undefined {
@@ -154,17 +96,63 @@ export class ModelingStore extends DisposableObject {
     return this.state.get(this.activeDb);
   }
 
+  private getInternalStateForActiveDb(): InternalDbModelingState | undefined {
+    if (!this.activeDb) {
+      return undefined;
+    }
+
+    return this.state.get(this.activeDb);
+  }
+
+  public hasStateForActiveDb(): boolean {
+    return !!this.getStateForActiveDb();
+  }
+
+  public anyDbsBeingModeled(): boolean {
+    return this.state.size > 0;
+  }
+
+  /**
+   * Returns the method for the given database item and method signature.
+   * Returns undefined if no method exists with that signature.
+   */
+  public getMethod(
+    dbItem: DatabaseItem,
+    methodSignature: string,
+  ): Method | undefined {
+    return this.getState(dbItem).methods.find(
+      (m) => m.signature === methodSignature,
+    );
+  }
+
+  /**
+   * Returns the methods for the given database item and method signatures.
+   * If the `methodSignatures` argument is not provided or is undefined, returns all methods.
+   */
+  public getMethods(
+    dbItem: DatabaseItem,
+    methodSignatures?: string[],
+  ): readonly Method[] {
+    const methods = this.getState(dbItem).methods;
+    if (!methodSignatures) {
+      return methods;
+    }
+    return methods.filter((method) =>
+      methodSignatures.includes(method.signature),
+    );
+  }
+
   public setMethods(dbItem: DatabaseItem, methods: Method[]) {
     const dbState = this.getState(dbItem);
     const dbUri = dbItem.databaseUri.toString();
 
     dbState.methods = [...methods];
 
-    this.onMethodsChangedEventEmitter.fire({
+    this.modelingEvents.fireMethodsChangedEvent(
       methods,
       dbUri,
-      isActiveDb: dbUri === this.activeDb,
-    });
+      dbUri === this.activeDb,
+    );
   }
 
   public setHideModeledMethods(
@@ -176,22 +164,55 @@ export class ModelingStore extends DisposableObject {
 
     dbState.hideModeledMethods = hideModeledMethods;
 
-    this.onHideModeledMethodsChangedEventEmitter.fire({
+    this.modelingEvents.fireHideModeledMethodsChangedEvent(
       hideModeledMethods,
-      isActiveDb: dbUri === this.activeDb,
-    });
+      dbUri === this.activeDb,
+    );
+  }
+
+  public setMode(dbItem: DatabaseItem, mode: Mode) {
+    const dbState = this.getState(dbItem);
+    const dbUri = dbItem.databaseUri.toString();
+
+    dbState.mode = mode;
+
+    this.modelingEvents.fireModeChangedEvent(mode, dbUri === this.activeDb);
+  }
+
+  public getMode(dbItem: DatabaseItem) {
+    return this.getState(dbItem).mode;
+  }
+
+  /**
+   * Returns the modeled methods for the given database item and method signatures.
+   * If the `methodSignatures` argument is not provided or is undefined, returns all modeled methods.
+   */
+  public getModeledMethods(
+    dbItem: DatabaseItem,
+    methodSignatures?: string[],
+  ): Readonly<Record<string, readonly ModeledMethod[]>> {
+    const modeledMethods = this.getState(dbItem).modeledMethods;
+    if (!methodSignatures) {
+      return modeledMethods;
+    }
+    return Object.fromEntries(
+      Object.entries(modeledMethods).filter(([key]) =>
+        methodSignatures.includes(key),
+      ),
+    );
   }
 
   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"),
           ),
         ),
       };
@@ -201,17 +222,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;
     });
   }
@@ -255,23 +280,59 @@ export class ModelingStore extends DisposableObject {
     });
   }
 
+  /**
+   * Sets which method is considered to be selected. This method will be shown in the method modeling panel.
+   *
+   * The `Method` and `Usage` objects must have been retrieved from the modeling store, and not from
+   * a webview. This is because we rely on object referential identity so it must be the same object
+   * that is held internally by the modeling store.
+   */
   public setSelectedMethod(dbItem: DatabaseItem, method: Method, usage: Usage) {
     const dbState = this.getState(dbItem);
 
     dbState.selectedMethod = method;
     dbState.selectedUsage = usage;
 
-    this.onSelectedMethodChangedEventEmitter.fire({
-      databaseItem: dbItem,
+    const modeledMethods = dbState.modeledMethods[method.signature] ?? [];
+    const isModified = dbState.modifiedMethodSignatures.has(method.signature);
+    const isInProgress = dbState.inProgressMethods.has(method.signature);
+    this.modelingEvents.fireSelectedMethodChangedEvent(
+      dbItem,
       method,
       usage,
-      modeledMethod: dbState.modeledMethods[method.signature],
-      isModified: dbState.modifiedMethodSignatures.has(method.signature),
+      modeledMethods,
+      isModified,
+      isInProgress,
+    );
+  }
+
+  public addInProgressMethods(
+    dbItem: DatabaseItem,
+    inProgressMethods: string[],
+  ) {
+    this.changeInProgressMethods(dbItem, (state) => {
+      state.inProgressMethods = new Set([
+        ...state.inProgressMethods,
+        ...inProgressMethods,
+      ]);
     });
   }
 
-  public getSelectedMethodDetails() {
-    const dbState = this.getStateForActiveDb();
+  public removeInProgressMethods(
+    dbItem: DatabaseItem,
+    methodSignatures: string[],
+  ) {
+    this.changeInProgressMethods(dbItem, (state) => {
+      state.inProgressMethods = new Set(
+        Array.from(state.inProgressMethods).filter(
+          (s) => !methodSignatures.includes(s),
+        ),
+      );
+    });
+  }
+
+  public getSelectedMethodDetails(): SelectedMethodDetails | undefined {
+    const dbState = this.getInternalStateForActiveDb();
     if (!dbState) {
       throw new Error("No active state found in modeling store");
     }
@@ -282,16 +343,18 @@ export class ModelingStore extends DisposableObject {
     }
 
     return {
+      databaseItem: dbState.databaseItem,
       method: selectedMethod,
       usage: dbState.selectedUsage,
-      modeledMethod: dbState.modeledMethods[selectedMethod.signature],
+      modeledMethods: dbState.modeledMethods[selectedMethod.signature] ?? [],
       isModified: dbState.modifiedMethodSignatures.has(
         selectedMethod.signature,
       ),
+      isInProgress: dbState.inProgressMethods.has(selectedMethod.signature),
     };
   }
 
-  private getState(databaseItem: DatabaseItem): DbModelingState {
+  private getState(databaseItem: DatabaseItem): InternalDbModelingState {
     if (!this.state.has(databaseItem.databaseUri.toString())) {
       throw Error(
         "Cannot get state for a database that has not been initialized",
@@ -303,31 +366,45 @@ export class ModelingStore extends DisposableObject {
 
   private changeModifiedMethods(
     dbItem: DatabaseItem,
-    updateState: (state: DbModelingState) => void,
+    updateState: (state: InternalDbModelingState) => void,
   ) {
     const state = this.getState(dbItem);
 
     updateState(state);
 
-    this.onModifiedMethodsChangedEventEmitter.fire({
-      modifiedMethods: state.modifiedMethodSignatures,
-      dbUri: dbItem.databaseUri.toString(),
-      isActiveDb: dbItem.databaseUri.toString() === this.activeDb,
-    });
+    this.modelingEvents.fireModifiedMethodsChangedEvent(
+      state.modifiedMethodSignatures,
+      dbItem.databaseUri.toString(),
+      dbItem.databaseUri.toString() === this.activeDb,
+    );
   }
 
   private changeModeledMethods(
     dbItem: DatabaseItem,
-    updateState: (state: DbModelingState) => void,
+    updateState: (state: InternalDbModelingState) => void,
   ) {
     const state = this.getState(dbItem);
 
     updateState(state);
 
-    this.onModeledMethodsChangedEventEmitter.fire({
-      modeledMethods: state.modeledMethods,
-      dbUri: dbItem.databaseUri.toString(),
-      isActiveDb: dbItem.databaseUri.toString() === this.activeDb,
-    });
+    this.modelingEvents.fireModeledMethodsChangedEvent(
+      state.modeledMethods,
+      dbItem.databaseUri.toString(),
+      dbItem.databaseUri.toString() === this.activeDb,
+    );
+  }
+
+  private changeInProgressMethods(
+    dbItem: DatabaseItem,
+    updateState: (state: InternalDbModelingState) => void,
+  ) {
+    const state = this.getState(dbItem);
+
+    updateState(state);
+
+    this.modelingEvents.fireInProgressMethodsChangedEvent(
+      dbItem.databaseUri.toString(),
+      state.inProgressMethods,
+    );
   }
 }

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

@@ -1,153 +0,0 @@
-import { ModeledMethod, ModeledMethodType, Provenance } from "./modeled-method";
-import { DataTuple } from "./model-extension-file";
-
-export type ExtensiblePredicateDefinition = {
-  extensiblePredicate: string;
-  generateMethodDefinition: (method: ModeledMethod) => DataTuple[];
-  readModeledMethod: (row: DataTuple[]) => ModeledMethod;
-
-  supportedKinds?: string[];
-};
-
-function readRowToMethod(row: DataTuple[]): string {
-  return `${row[0]}.${row[1]}#${row[3]}${row[4]}`;
-}
-
-export const extensiblePredicateDefinitions: Record<
-  Exclude<ModeledMethodType, "none">,
-  ExtensiblePredicateDefinition
-> = {
-  source: {
-    extensiblePredicate: "sourceModel",
-    // extensible predicate sourceModel(
-    //   string package, string type, boolean subtypes, string name, string signature, string ext,
-    //   string output, string kind, string provenance
-    // );
-    generateMethodDefinition: (method) => [
-      method.packageName,
-      method.typeName,
-      true,
-      method.methodName,
-      method.methodParameters,
-      "",
-      method.output,
-      method.kind,
-      method.provenance,
-    ],
-    readModeledMethod: (row) => ({
-      type: "source",
-      input: "",
-      output: row[6] as string,
-      kind: row[7] as string,
-      provenance: row[8] as Provenance,
-      signature: readRowToMethod(row),
-      packageName: row[0] as string,
-      typeName: row[1] as string,
-      methodName: row[3] as string,
-      methodParameters: row[4] as string,
-    }),
-    supportedKinds: ["local", "remote"],
-  },
-  sink: {
-    extensiblePredicate: "sinkModel",
-    // extensible predicate sinkModel(
-    //   string package, string type, boolean subtypes, string name, string signature, string ext,
-    //   string input, string kind, string provenance
-    // );
-    generateMethodDefinition: (method) => [
-      method.packageName,
-      method.typeName,
-      true,
-      method.methodName,
-      method.methodParameters,
-      "",
-      method.input,
-      method.kind,
-      method.provenance,
-    ],
-    readModeledMethod: (row) => ({
-      type: "sink",
-      input: row[6] as string,
-      output: "",
-      kind: row[7] as string,
-      provenance: row[8] as Provenance,
-      signature: readRowToMethod(row),
-      packageName: row[0] as string,
-      typeName: row[1] as string,
-      methodName: row[3] as string,
-      methodParameters: row[4] as string,
-    }),
-    supportedKinds: [
-      "code-injection",
-      "command-injection",
-      "file-content-store",
-      "html-injection",
-      "js-injection",
-      "ldap-injection",
-      "log-injection",
-      "path-injection",
-      "request-forgery",
-      "sql-injection",
-      "url-redirection",
-    ],
-  },
-  summary: {
-    extensiblePredicate: "summaryModel",
-    // extensible predicate summaryModel(
-    //   string package, string type, boolean subtypes, string name, string signature, string ext,
-    //   string input, string output, string kind, string provenance
-    // );
-    generateMethodDefinition: (method) => [
-      method.packageName,
-      method.typeName,
-      true,
-      method.methodName,
-      method.methodParameters,
-      "",
-      method.input,
-      method.output,
-      method.kind,
-      method.provenance,
-    ],
-    readModeledMethod: (row) => ({
-      type: "summary",
-      input: row[6] as string,
-      output: row[7] as string,
-      kind: row[8] as string,
-      provenance: row[9] as Provenance,
-      signature: readRowToMethod(row),
-      packageName: row[0] as string,
-      typeName: row[1] as string,
-      methodName: row[3] as string,
-      methodParameters: row[4] as string,
-    }),
-    supportedKinds: ["taint", "value"],
-  },
-  neutral: {
-    extensiblePredicate: "neutralModel",
-    // extensible predicate neutralModel(
-    //   string package, string type, string name, string signature, string kind, string provenance
-    // );
-    generateMethodDefinition: (method) => [
-      method.packageName,
-      method.typeName,
-      method.methodName,
-      method.methodParameters,
-      method.kind,
-      method.provenance,
-    ],
-    readModeledMethod: (row) => ({
-      type: "neutral",
-      input: "",
-      output: "",
-      kind: row[4] as string,
-      provenance: row[5] as Provenance,
-      signature: `${row[0]}.${row[1]}#${row[2]}${row[3]}`,
-      packageName: row[0] as string,
-      typeName: row[1] as string,
-      methodName: row[2] as string,
-      methodParameters: row[3] as string,
-    }),
-    supportedKinds: ["summary", "source", "sink"],
-  },
-};

extensions/ql-vscode/src/model-editor/queries/index.ts

@@ -1,9 +1,11 @@
 import { fetchExternalApisQuery as csharpFetchExternalApisQuery } from "./csharp";
 import { fetchExternalApisQuery as javaFetchExternalApisQuery } from "./java";
+import { fetchExternalApisQuery as rubyFetchExternalApisQuery } from "./ruby";
 import { Query } from "./query";
 import { QueryLanguage } from "../../common/query-language";
 
 export const fetchExternalApiQueries: Partial<Record<QueryLanguage, Query>> = {
   [QueryLanguage.CSharp]: csharpFetchExternalApisQuery,
   [QueryLanguage.Java]: javaFetchExternalApisQuery,
+  [QueryLanguage.Ruby]: rubyFetchExternalApisQuery,
 };

extensions/ql-vscode/src/model-editor/queries/ruby.ts

@@ -0,0 +1,404 @@
+import { Query } from "./query";
+
+export const fetchExternalApisQuery: Query = {
+  applicationModeQuery: `/**
+ * @name Fetch endpoints for use in the model editor (application mode)
+ * @description A list of 3rd party endpoints (methods and attributes) used in the codebase. Excludes test and generated code.
+ * @kind table
+ * @id rb/utils/modeleditor/application-mode-endpoints
+ * @tags modeleditor endpoints application-mode
+ */
+
+import ruby
+
+select "todo", "todo", "todo", "todo", "todo", false, "todo", "todo", "todo", "todo"
+`,
+  frameworkModeQuery: `/**
+ * @name Fetch endpoints for use in the model editor (framework mode)
+ * @description A list of endpoints accessible (methods and attributes) for consumers of the library. Excludes test and generated code.
+ * @kind table
+ * @id rb/utils/modeleditor/framework-mode-endpoints
+ * @tags modeleditor endpoints framework-mode
+ */
+
+import ruby
+import FrameworkModeEndpointsQuery
+import ModelEditor
+
+from PublicEndpointFromSource endpoint, boolean supported, string type
+where
+  supported = isSupported(endpoint) and
+  type = supportedType(endpoint)
+select endpoint, endpoint.getNamespace(), endpoint.getTypeName(), endpoint.getName(),
+  endpoint.getParameterTypes(), supported, endpoint.getFile().getBaseName(), type
+`,
+  dependencies: {
+    "FrameworkModeEndpointsQuery.qll": `private import ruby
+private import ModelEditor
+private import modeling.internal.Util as Util
+
+/**
+ * A class of effectively public callables from source code.
+ */
+class PublicEndpointFromSource extends Endpoint {
+  PublicEndpointFromSource() {
+    this.getFile() instanceof Util::RelevantFile
+  }
+
+  override predicate isSource() { this instanceof SourceCallable }
+
+  override predicate isSink() { this instanceof SinkCallable }
+}
+`,
+    "ModelEditor.qll": `/** Provides classes and predicates related to handling APIs for the VS Code extension. */
+
+private import ruby
+private import codeql.ruby.dataflow.FlowSummary
+private import codeql.ruby.dataflow.internal.DataFlowPrivate
+private import codeql.ruby.dataflow.internal.FlowSummaryImpl as FlowSummaryImpl
+private import codeql.ruby.dataflow.internal.FlowSummaryImplSpecific
+private import modeling.internal.Util as Util
+private import modeling.internal.Types
+private import codeql.ruby.frameworks.core.Gem
+
+/** Holds if the given callable is not worth supporting. */
+private predicate isUninteresting(DataFlow::MethodNode c) {
+  c.getLocation().getFile().getRelativePath().regexpMatch(".*(test|spec).*")
+}
+
+/**
+ * A callable method or accessor from either the Ruby Standard Library, a 3rd party library, or from the source.
+ */
+class Endpoint extends DataFlow::MethodNode {
+  Endpoint() {
+    this.isPublic() and not isUninteresting(this)
+  }
+
+  File getFile() { result = this.getLocation().getFile() }
+
+  string getName() { result = this.getMethodName() }
+
+  /**
+   * Gets the namespace of this endpoint.
+   */
+  bindingset[this]
+  string getNamespace() {
+    // Return the name of any gemspec file in the database.
+    // TODO: make this work for projects with multiple gems (and hence multiple gemspec files)
+    result = any(Gem::GemSpec g).getName()
+  }
+
+  /**
+   * Gets the unbound type name of this endpoint.
+   */
+  bindingset[this]
+  string getTypeName() {
+    // result = nestedName(this.getDeclaringType().getUnboundDeclaration())
+    // result = any(DataFlow::ClassNode c | Types::methodReturnsType(this, c) | c).getQualifiedName()
+    result = Util::getAnAccessPathPrefixWithoutSuffix(this)
+  }
+
+  /**
+   * Gets the parameter types of this endpoint.
+   */
+  bindingset[this]
+  string getParameterTypes() {
+    // For now, return the names of postional parameters. We don't always have type information, so we can't return type names.
+    // We don't yet handle keyword params, splat params or block params.
+    // result = "(" + parameterQualifiedTypeNamesToString(this) + ")"
+    result =
+      "(" +
+        concat(DataFlow::ParameterNode p, int i |
+          p = this.asCallable().getParameter(i)
+        |
+          p.getName(), "," order by i
+        ) + ")"
+  }
+
+  /** Holds if this API has a supported summary. */
+  pragma[nomagic]
+  predicate hasSummary() {
+    // this instanceof SummarizedCallable
+    none()
+  }
+
+  /** Holds if this API is a known source. */
+  pragma[nomagic]
+  abstract predicate isSource();
+
+  /** Holds if this API is a known sink. */
+  pragma[nomagic]
+  abstract predicate isSink();
+
+  /** Holds if this API is a known neutral. */
+  pragma[nomagic]
+  predicate isNeutral() {
+    // this instanceof FlowSummaryImpl::Public::NeutralCallable
+    none()
+  }
+
+  /**
+   * Holds if this API is supported by existing CodeQL libraries, that is, it is either a
+   * recognized source, sink or neutral or it has a flow summary.
+   */
+  predicate isSupported() {
+    this.hasSummary() or this.isSource() or this.isSink() or this.isNeutral()
+  }
+}
+
+boolean isSupported(Endpoint endpoint) {
+  if endpoint.isSupported() then result = true else result = false
+}
+
+string supportedType(Endpoint endpoint) {
+  endpoint.isSink() and result = "sink"
+  or
+  endpoint.isSource() and result = "source"
+  or
+  endpoint.hasSummary() and result = "summary"
+  or
+  endpoint.isNeutral() and result = "neutral"
+  or
+  not endpoint.isSupported() and result = ""
+}
+
+string methodClassification(Call method) {
+  result = "source"
+}
+
+/**
+ * A callable where there exists a MaD sink model that applies to it.
+ */
+class SinkCallable extends DataFlow::CallableNode {
+  SinkCallable() { sinkElement(this.asExpr().getExpr(), _, _, _) }
+}
+
+/**
+ * A callable where there exists a MaD source model that applies to it.
+ */
+class SourceCallable extends DataFlow::CallableNode {
+  SourceCallable() { sourceElement(this.asExpr().getExpr(), _, _, _) }
+}`,
+    "modeling/internal/Util.qll": `private import ruby
+
+// \`SomeClass#initialize\` methods are usually called indirectly via
+// \`SomeClass.new\`, so we need to account for this when generating access paths
+private string getNormalizedMethodName(DataFlow::MethodNode methodNode) {
+  exists(string actualMethodName | actualMethodName = methodNode.getMethodName() |
+    if actualMethodName = "initialize" then result = "new" else result = actualMethodName
+  )
+}
+
+private string getAccessPathSuffix(Ast::MethodBase method) {
+  if method instanceof Ast::SingletonMethod or method.getName() = "initialize"
+  then result = "!"
+  else result = ""
+}
+
+string getAnAccessPathPrefix(DataFlow::MethodNode methodNode) {
+  result =
+    getAnAccessPathPrefixWithoutSuffix(methodNode) +
+      getAccessPathSuffix(methodNode.asExpr().getExpr())
+}
+
+string getAnAccessPathPrefixWithoutSuffix(DataFlow::MethodNode methodNode) {
+  result =
+    methodNode
+        .asExpr()
+        .getExpr()
+        .getEnclosingModule()
+        .(Ast::ConstantWriteAccess)
+        .getAQualifiedName()
+}
+
+class RelevantFile extends File {
+  RelevantFile() { not this.getRelativePath().regexpMatch(".*/?test(case)?s?/.*") }
+}
+
+string getMethodPath(DataFlow::MethodNode methodNode) {
+  result = "Method[" + getNormalizedMethodName(methodNode) + "]"
+}
+
+private string getParameterPath(DataFlow::ParameterNode paramNode) {
+  exists(Ast::Parameter param, string paramSpec |
+    param = paramNode.asParameter() and
+    (
+      paramSpec = param.getPosition().toString()
+      or
+      paramSpec = param.(Ast::KeywordParameter).getName() + ":"
+      or
+      param instanceof Ast::BlockParameter and
+      paramSpec = "block"
+    )
+  |
+    result = "Parameter[" + paramSpec + "]"
+  )
+}
+
+string getMethodParameterPath(DataFlow::MethodNode methodNode, DataFlow::ParameterNode paramNode) {
+  result = getMethodPath(methodNode) + "." + getParameterPath(paramNode)
+}
+`,
+    "modeling/internal/Types.qll": `private import ruby
+private import codeql.ruby.ApiGraphs
+private import Util as Util
+
+module Types {
+  private module Config implements DataFlow::ConfigSig {
+    predicate isSource(DataFlow::Node source) {
+      // TODO: construction of type values not using a "new" call
+      source.(DataFlow::CallNode).getMethodName() = "new"
+    }
+
+    predicate isSink(DataFlow::Node sink) { sink = any(DataFlow::MethodNode m).getAReturnNode() }
+  }
+
+  private import DataFlow::Global<Config>
+
+  predicate methodReturnsType(DataFlow::MethodNode methodNode, DataFlow::ClassNode classNode) {
+    // ignore cases of initializing instance of self
+    not methodNode.getMethodName() = "initialize" and
+    exists(DataFlow::CallNode initCall |
+      flow(initCall, methodNode.getAReturnNode()) and
+      classNode.getAnImmediateReference().getAMethodCall() = initCall and
+      // constructed object does not have a type declared in test code
+      /*
+       * TODO: this may be too restrictive, e.g.
+       * - if a type is declared in both production and test code
+       * - if a built-in type is extended in test code
+       */
+
+      forall(Ast::ModuleBase classDecl | classDecl = classNode.getADeclaration() |
+        classDecl.getLocation().getFile() instanceof Util::RelevantFile
+      )
+    )
+  }
+
+  // \`exprNode\` is an instance of \`classNode\`
+  private predicate exprHasType(DataFlow::ExprNode exprNode, DataFlow::ClassNode classNode) {
+    exists(DataFlow::MethodNode methodNode, DataFlow::CallNode callNode |
+      methodReturnsType(methodNode, classNode) and
+      callNode.getATarget() = methodNode
+    |
+      exprNode.getALocalSource() = callNode
+    )
+    or
+    exists(DataFlow::MethodNode containingMethod |
+      classNode.getInstanceMethod(containingMethod.getMethodName()) = containingMethod
+    |
+      exprNode.getALocalSource() = containingMethod.getSelfParameter()
+    )
+  }
+
+  // extensible predicate typeModel(string type1, string type2, string path);
+  // the method node in type2 constructs an instance of classNode
+  private predicate typeModelReturns(string type1, string type2, string path) {
+    exists(DataFlow::MethodNode methodNode, DataFlow::ClassNode classNode |
+      methodNode.getLocation().getFile() instanceof Util::RelevantFile and
+      methodReturnsType(methodNode, classNode)
+    |
+      type1 = classNode.getQualifiedName() and
+      type2 = Util::getAnAccessPathPrefix(methodNode) and
+      path = Util::getMethodPath(methodNode) + ".ReturnValue"
+    )
+  }
+
+  predicate methodTakesParameterOfType(
+    DataFlow::MethodNode methodNode, DataFlow::ClassNode classNode,
+    DataFlow::ParameterNode parameterNode
+  ) {
+    exists(DataFlow::CallNode callToMethodNode, DataFlow::LocalSourceNode argumentNode |
+      callToMethodNode.getATarget() = methodNode and
+      // positional parameter
+      exists(int paramIndex |
+        argumentNode.flowsTo(callToMethodNode.getArgument(paramIndex)) and
+        parameterNode = methodNode.getParameter(paramIndex)
+      )
+      or
+      // keyword parameter
+      exists(string kwName |
+        argumentNode.flowsTo(callToMethodNode.getKeywordArgument(kwName)) and
+        parameterNode = methodNode.getKeywordParameter(kwName)
+      )
+      or
+      // block parameter
+      argumentNode.flowsTo(callToMethodNode.getBlock()) and
+      parameterNode = methodNode.getBlockParameter()
+    |
+      // parameter directly from new call
+      argumentNode.(DataFlow::CallNode).getMethodName() = "new" and
+      classNode.getAnImmediateReference().getAMethodCall() = argumentNode
+      or
+      // parameter from indirect new call
+      exists(DataFlow::ExprNode argExpr |
+        exprHasType(argExpr, classNode) and
+        argumentNode.(DataFlow::CallNode).getATarget() = argExpr
+      )
+    )
+  }
+
+  private predicate typeModelParameters(string type1, string type2, string path) {
+    exists(
+      DataFlow::MethodNode methodNode, DataFlow::ClassNode classNode,
+      DataFlow::ParameterNode parameterNode
+    |
+      methodNode.getLocation().getFile() instanceof Util::RelevantFile and
+      methodTakesParameterOfType(methodNode, classNode, parameterNode)
+    |
+      type1 = classNode.getQualifiedName() and
+      type2 = Util::getAnAccessPathPrefix(methodNode) and
+      path = Util::getMethodParameterPath(methodNode, parameterNode)
+    )
+  }
+
+  // TODO: non-positional params for block arg parameters
+  private predicate methodYieldsType(
+    DataFlow::CallableNode callableNode, int argIdx, DataFlow::ClassNode classNode
+  ) {
+    exprHasType(callableNode.getABlockCall().getArgument(argIdx), classNode)
+  }
+
+  /*
+   * e.g. for
+   * \`\`\`rb
+   * class Foo
+   *  def initialize
+   *    // do some stuff...
+   *    if block_given?
+   *      yield self
+   *    end
+   *  end
+   *
+   *  def do_something
+   *    // do something else
+   *  end
+   * end
+   *
+   * Foo.new do |foo| foo.do_something end
+   * \`\`\`
+   *
+   * the parameter foo to the block is an instance of Foo.
+   */
+
+  private predicate typeModelBlockArgumentParameters(string type1, string type2, string path) {
+    exists(DataFlow::MethodNode methodNode, DataFlow::ClassNode classNode, int argIdx |
+      methodNode.getLocation().getFile() instanceof Util::RelevantFile and
+      methodYieldsType(methodNode, argIdx, classNode)
+    |
+      type1 = classNode.getQualifiedName() and
+      type2 = Util::getAnAccessPathPrefix(methodNode) and
+      path = Util::getMethodPath(methodNode) + ".Argument[block].Parameter[" + argIdx + "]"
+    )
+  }
+
+  predicate typeModel(string type1, string type2, string path) {
+    typeModelReturns(type1, type2, path)
+    or
+    typeModelParameters(type1, type2, path)
+    or
+    typeModelBlockArgumentParameters(type1, type2, path)
+  }
+}
+`,
+  },
+};

extensions/ql-vscode/src/model-editor/shared/in-progress-methods.ts

@@ -1,34 +0,0 @@
-/**
- * A class that keeps track of which methods are in progress for each package.
- *
- * This class is immutable and therefore is safe to be used in a React useState hook.
- */
-export class InProgressMethods {
-  // A map of in-progress method signatures for each package.
-  private readonly methodMap: ReadonlyMap<string, Set<string>>;
-
-  constructor(methodMap?: ReadonlyMap<string, Set<string>>) {
-    this.methodMap = methodMap ?? new Map<string, Set<string>>();
-  }
-
-  /**
-   * Sets the in-progress methods for the given package.
-   * Returns a new InProgressMethods instance.
-   */
-  public setPackageMethods(
-    packageName: string,
-    methods: Set<string>,
-  ): InProgressMethods {
-    const newMethodMap = new Map<string, Set<string>>(this.methodMap);
-    newMethodMap.set(packageName, methods);
-    return new InProgressMethods(newMethodMap);
-  }
-
-  public hasMethod(packageName: string, method: string): boolean {
-    const methods = this.methodMap.get(packageName);
-    if (methods) {
-      return methods.has(method);
-    }
-    return false;
-  }
-}

extensions/ql-vscode/src/model-editor/shared/mode.ts

@@ -2,3 +2,5 @@ export enum Mode {
   Application = "application",
   Framework = "framework",
 }
+
+export const INITIAL_MODE = Mode.Application;