Merge pull request #2891 from github/koesie10/dispose-webview

Dispose tracked objects when panel is disposed

.github/workflows/cli-test.yml

@@ -62,7 +62,7 @@ jobs:
 
       - uses: actions/setup-node@v3
         with:
-          node-version: '18.15.0'
+          node-version: '16.17.1'
           cache: 'npm'
           cache-dependency-path: extensions/ql-vscode/package-lock.json
 

.github/workflows/main.yml

@@ -22,7 +22,7 @@ jobs:
 
       - uses: actions/setup-node@v3
         with:
-          node-version: '18.15.0'
+          node-version: '16.17.1'
           cache: 'npm'
           cache-dependency-path: extensions/ql-vscode/package-lock.json
 
@@ -64,7 +64,7 @@ jobs:
 
       - uses: actions/setup-node@v3
         with:
-          node-version: '18.15.0'
+          node-version: '16.17.1'
           cache: 'npm'
           cache-dependency-path: extensions/ql-vscode/package-lock.json
 
@@ -99,6 +99,42 @@ jobs:
         run: |
           npm run find-deadcode
 
+  generated:
+    name: Check generated code
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+        with:
+          fetch-depth: 1
+
+      - uses: actions/setup-node@v3
+        with:
+          node-version: '16.17.1'
+          cache: 'npm'
+          cache-dependency-path: extensions/ql-vscode/package-lock.json
+
+      - name: Install dependencies
+        working-directory: extensions/ql-vscode
+        run: |
+          npm ci
+        shell: bash
+
+      - name: Check that repo is clean
+        run: |
+          git diff --exit-code
+          git diff --exit-code --cached
+
+      - name: Generate code
+        working-directory: extensions/ql-vscode
+        run: |
+          npm run generate
+
+      - name: Check for changes
+        run: |
+          git diff --exit-code
+          git diff --exit-code --cached
+
   unit-test:
     name: Unit Test
     runs-on: ${{ matrix.os }}
@@ -113,7 +149,7 @@ jobs:
 
       - uses: actions/setup-node@v3
         with:
-          node-version: '18.15.0'
+          node-version: '16.17.1'
           cache: 'npm'
           cache-dependency-path: extensions/ql-vscode/package-lock.json
 
@@ -147,7 +183,7 @@ jobs:
 
       - uses: actions/setup-node@v3
         with:
-          node-version: '18.15.0'
+          node-version: '16.17.1'
           cache: 'npm'
           cache-dependency-path: extensions/ql-vscode/package-lock.json
 
@@ -215,7 +251,7 @@ jobs:
 
       - uses: actions/setup-node@v3
         with:
-          node-version: '18.15.0'
+          node-version: '16.17.1'
           cache: 'npm'
           cache-dependency-path: extensions/ql-vscode/package-lock.json
 

.github/workflows/release.yml

@@ -22,7 +22,7 @@ jobs:
 
       - uses: actions/setup-node@v3
         with:
-          node-version: '18.15.0'
+          node-version: '16.17.1'
 
       - name: Install dependencies
         run: |

docs/node-version.md

@@ -0,0 +1,29 @@
+# Node version
+
+The CodeQL for VS Code extension defines the version of Node.js that it is intended to run with. This Node.js version is used when running most CI and unit tests.
+
+When running in production (i.e. as an extension for a VS Code application) it will use the Node.js version provided by VS Code. This can mean a different Node.js version is used by different users with different versions of VS Code.
+We should make sure the CodeQL for VS Code extension works with the Node.js version supplied by all versions of VS Code that we support.
+
+## Checking the version of Node.js supplied by VS Code
+
+You can find this info by seleting "About Visual Studio Code" from the top menu.
+
+![about-vscode](images/about-vscode.png)
+
+## Updating the Node.js version
+
+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/package.json` - the "engines.node: '[VERSION]'" setting
+
+## Node.js version used in tests
+
+Unit tests will use whatever version of Node.js is installed locally. In CI this will be the version specified in the workflow.
+
+Integration tests download a copy of VS Code and then will use whatever version of Node.js is provided by VS Code. Our integration tests are currently pinned to an older version of VS Code. See [VS Code version used in tests](./vscode-version.md#vs-code-version-used-in-tests) for more information.

docs/releasing.md

@@ -11,10 +11,7 @@
       * New telemetry events are added.
       * Deprecation or removal of commands.
       * Accumulation of many changes, none of which are individually big enough to warrant a minor bump, but which together are. This does not include changes which are purely internal to the extension, such as refactoring, or which are only available behind a feature flag.
-1. Double-check that the node version we're using matches the one used for VS Code. You can find this info by seleting "About Visual Studio Code" from the top menu. If it doesn't match, you will then need to update the node version in the following files:
-    * `.nvmrc` - this will enable `nvm` to automatically switch to the correct node version when you're in the project folder
-    * `.github/workflows/main.yml` - all the "node-version: '[VERSION]'" settings
-    * `.github/workflows/release.yml` - the "node-version: '[VERSION]'" setting
+1. Double-check that the node version we're using matches the one used for VS Code. See the [Node.js version instructions](./node-version.md) for more information.
 1. Double-check that the extension `package.json` and `package-lock.json` have the version you intend to release. If you are doing a patch release (as opposed to minor or major version) this should already be correct.
 1. Create a PR for this release:
     * This PR will contain any missing bits from steps 1, 2 and 3. Most of the time, this will just be updating `CHANGELOG.md` with today's date.

docs/vscode-version.md

@@ -0,0 +1,33 @@
+# VS Code version
+
+The CodeQL for VS Code extension specifies the versions of VS Code that it is compatible with. VS Code will only offer to install and upgrade the extension when this version range is satisfied.
+
+## Where is the VS Code version specified
+
+1. Hard limit in [`package.json`](https://github.com/github/vscode-codeql/blob/606bfd7f877d9fffe4ff83b78015ab15f8840b12/extensions/ql-vscode/package.json#L16)
+
+    This is the value that VS Code understands and respects. If a user does not meet this version requirement then VS Code will not offer to install the CodeQL for VS Code extension, and if the extension is already installed then it will silently refuse to upgrade the extension.
+
+1. Soft limit in [`extension.ts`](https://github.com/github/vscode-codeql/blob/606bfd7f877d9fffe4ff83b78015ab15f8840b12/extensions/ql-vscode/src/extension.ts#L307)
+
+    This value is used internally by the CodeQL for VS Code extension and is used to provide a warning to users without blocking them from installing or upgrading. If the extension detects that this version range is not met it will output a warning message to the user prompting them to upgrade their VS Code version to ge the latest features of CodeQL.
+
+## When to update the VS Code version
+
+Generally we should aim to support as wide a range of VS Code versions as we can, so unless there is a reason to do so we do not update the minimum VS Code version requirement.
+Reasons for updating the minimum VS Code version include:
+
+- A new feature is included in VS Code. We may want to ensure that it is available to use so we do not have to provide an alternative code path.
+- A breaking change has happened in VS Code, and it is not possible to support both new and old versions.
+
+Also consider what percentage of our users are using each VS Code version. This information is available in our telemetry.
+
+## How to update the VS Code version
+
+To provide a good experience to users, it is recommented to update the `MIN_VERSION` in `extension.ts` first and release, and then update the `vscode` version in `package.json` and release again. By stagging this update across two releases it gives users on older VS Code versions a chance to upgrade before it silently refuses to upgrade them.
+
+## VS Code version used in tests
+
+Our integration tests are currently pinned to use an older version of VS Code due to <https://github.com/github/vscode-codeql/issues/2402>.
+This version is specified in [`jest-runner-vscode.config.base.js`](https://github.com/github/vscode-codeql/blob/d93f2b67c84e79737b0ce4bb74e31558b5f5166e/extensions/ql-vscode/test/vscode-tests/jest-runner-vscode.config.base.js#L17).
+Until this is resolved this will limit us updating our minimum supported version of VS Code.

extensions/ql-vscode/.nvmrc

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

extensions/ql-vscode/CHANGELOG.md

@@ -1,5 +1,12 @@
 # CodeQL for Visual Studio Code: Changelog
 
+## 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)
+- Fix a bug when parsing large evaluation log summaries. [#2858](https://github.com/github/vscode-codeql/pull/2858)
+- Right-align and format numbers in raw result tables. [#2864](https://github.com/github/vscode-codeql/pull/2864)
+- Remove rate limit warning notifications when using Code Search to add repositories to a variant analysis list. [#2812](https://github.com/github/vscode-codeql/pull/2812)
+
 ## 1.9.0 - 19 September 2023
 
 - Release the [CodeQL model editor](https://codeql.github.com/docs/codeql/codeql-for-visual-studio-code/using-the-codeql-model-editor) to create CodeQL model packs for Java frameworks. Open the editor using the "CodeQL: Open CodeQL Model Editor (Beta)" command. [#2823](https://github.com/github/vscode-codeql/pull/2823)

extensions/ql-vscode/media/dark/sort-alpha.svg

@@ -1,17 +0,0 @@
-<?xml version="1.0" encoding="iso-8859-1"?>
-<!-- Generator: Adobe Illustrator 19.0.0, SVG Export Plug-In . SVG Version: 6.00 Build 0)  -->
-<svg version="1.1" id="Capa_1" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" x="0px" y="0px" fill="none"
-	 viewBox="0 0 432 432" style="enable-background:new 0 0 432 432;" xml:space="preserve">
-<g>
-	<g>
-		<g>
-			<polygon points="234.24,9.067 183.893,59.413 284.587,59.413" fill="#C5C5C5"/>
-			<polygon points="301.44,304.32 427.947,120.853 427.947,93.973 250.88,93.973 250.88,128.107 376.32,128.107 250.027,310.72
-				250.027,338.24 432,338.24 432,304.32" fill="#C5C5C5"/>
-			<polygon points="234.24,422.933 283.947,373.227 184.533,373.227" fill="#C5C5C5"/>
-			<path d="M226.773,338.24L130.987,93.76H96L0,338.24h39.253l19.627-52.267h109.013l19.627,52.267H226.773z M71.893,250.987
-				L113.28,140.48l41.387,110.507H71.893z" fill="#C5C5C5"/>
-		</g>
-	</g>
-</g>
-</svg>

extensions/ql-vscode/media/dark/sort-date.svg

@@ -1,3 +0,0 @@
-<svg width="16" height="16" viewBox="0 0 16 16" fill="none" xmlns="http://www.w3.org/2000/svg">
-<path fill-rule="evenodd" clip-rule="evenodd" d="M7 2L6 3V6H7V3H14V5.45306L14.2071 5.29286L15 6.08576V3L14 2H7ZM8 4H10V6H8V4ZM5 9H3V11H5V9ZM2 7L1 8V13L2 14H9L10 13V8L9 7H2ZM2 13V8H9V13H2ZM8 10H6V12H8V10ZM13 4H12V7.86388L10.818 6.68192L10.1109 7.38903L12.1465 9.42454L12.8536 9.42454L14.889 7.38908L14.1819 6.68197L13 7.86388V4Z" fill="#C5C5C5"/>
-</svg>

extensions/ql-vscode/media/light/sort-alpha.svg

@@ -1,17 +0,0 @@
-<?xml version="1.0" encoding="iso-8859-1"?>
-<!-- Generator: Adobe Illustrator 19.0.0, SVG Export Plug-In . SVG Version: 6.00 Build 0)  -->
-<svg version="1.1" id="Capa_1" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" x="0px" y="0px"
-	 viewBox="0 0 432 432" style="enable-background:new 0 0 432 432;" xml:space="preserve">
-<g>
-	<g>
-		<g>
-			<polygon points="234.24,9.067 183.893,59.413 284.587,59.413 			"/>
-			<polygon points="301.44,304.32 427.947,120.853 427.947,93.973 250.88,93.973 250.88,128.107 376.32,128.107 250.027,310.72 
-				250.027,338.24 432,338.24 432,304.32 			"/>
-			<polygon points="234.24,422.933 283.947,373.227 184.533,373.227 			"/>
-			<path d="M226.773,338.24L130.987,93.76H96L0,338.24h39.253l19.627-52.267h109.013l19.627,52.267H226.773z M71.893,250.987
-				L113.28,140.48l41.387,110.507H71.893z"/>
-		</g>
-	</g>
-</g>
-</svg>

extensions/ql-vscode/media/light/sort-date.svg

@@ -1,3 +0,0 @@
-<svg width="16" height="16" viewBox="0 0 16 16" fill="none" xmlns="http://www.w3.org/2000/svg">
-<path fill-rule="evenodd" clip-rule="evenodd" d="M7 2L6 3V6H7V3H14V5.45306L14.2071 5.29286L15 6.08576V3L14 2H7ZM8 4H10V6H8V4ZM5 9H3V11H5V9ZM2 7L1 8V13L2 14H9L10 13V8L9 7H2ZM2 13V8H9V13H2ZM8 10H6V12H8V10ZM13 4H12V7.86388L10.818 6.68192L10.1109 7.38903L12.1465 9.42454L12.8536 9.42454L14.889 7.38908L14.1819 6.68197L13 7.86388V4Z" fill="#424242"/>
-</svg>

extensions/ql-vscode/package.json

@@ -4,7 +4,7 @@
   "description": "CodeQL for Visual Studio Code",
   "author": "GitHub",
   "private": true,
-  "version": "1.9.0",
+  "version": "1.9.1",
   "publisher": "GitHub",
   "license": "MIT",
   "icon": "media/VS-marketplace-CodeQL-icon.png",
@@ -14,7 +14,7 @@
   },
   "engines": {
     "vscode": "^1.67.0",
-    "node": "^18.15.0",
+    "node": "^16.17.1",
     "npm": ">=7.20.6"
   },
   "categories": [
@@ -753,6 +753,78 @@
         "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"
@@ -771,19 +843,11 @@
       },
       {
         "command": "codeQLDatabases.sortByName",
-        "title": "Sort by Name",
-        "icon": {
-          "light": "media/light/sort-alpha.svg",
-          "dark": "media/dark/sort-alpha.svg"
-        }
+        "title": "Sort by Name"
       },
       {
         "command": "codeQLDatabases.sortByDateAdded",
-        "title": "Sort by Date Added",
-        "icon": {
-          "light": "media/light/sort-date.svg",
-          "dark": "media/dark/sort-date.svg"
-        }
+        "title": "Sort by Date Added"
       },
       {
         "command": "codeQL.checkForUpdatesToCLI",
@@ -981,16 +1045,6 @@
         }
       ],
       "view/title": [
-        {
-          "command": "codeQLDatabases.sortByName",
-          "when": "view == codeQLDatabases",
-          "group": "navigation"
-        },
-        {
-          "command": "codeQLDatabases.sortByDateAdded",
-          "when": "view == codeQLDatabases",
-          "group": "navigation"
-        },
         {
           "command": "codeQLDatabases.chooseDatabaseFolder",
           "when": "view == codeQLDatabases",
@@ -1011,6 +1065,21 @@
           "when": "view == codeQLDatabases",
           "group": "navigation"
         },
+        {
+          "command": "codeQLDatabases.sortByName",
+          "when": "view == codeQLDatabases",
+          "group": "1_databases@0"
+        },
+        {
+          "command": "codeQLDatabases.sortByDateAdded",
+          "when": "view == codeQLDatabases",
+          "group": "1_databases@1"
+        },
+        {
+          "submenu": "codeQLDatabases.languages",
+          "when": "view == codeQLDatabases && config.codeQL.canary && config.codeQL.showLanguageFilter",
+          "group": "2_databases@0"
+        },
         {
           "command": "codeQLQueries.createQuery",
           "when": "view == codeQLQueries",
@@ -1531,6 +1600,78 @@
           "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"
@@ -1725,8 +1866,88 @@
           "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": [
         {
@@ -1784,7 +2005,7 @@
           "type": "webview",
           "id": "codeQLMethodModeling",
           "name": "CodeQL Method Modeling",
-          "when": "config.codeQL.canary && config.codeQL.model.methodModelingView && codeql.modelEditorOpen"
+          "when": "config.codeQL.canary && config.codeQL.model.methodModelingView && codeql.modelEditorOpen && !codeql.modelEditorActive"
         }
       ]
     },
@@ -1842,6 +2063,8 @@
     "storybook": "storybook dev -p 6006",
     "build-storybook": "storybook build",
     "lint:scenarios": "ts-node scripts/lint-scenarios.ts",
+    "generate": "npm-run-all -p generate:*",
+    "generate:schemas": "ts-node scripts/generate-schemas.ts",
     "check-types": "find . -type f -name \"tsconfig.json\" -not -path \"./node_modules/*\" | sed -r 's|/[^/]+$||' | sort | uniq | xargs -I {} sh -c \"echo Checking types in {} && cd {} && npx tsc --noEmit\"",
     "postinstall": "patch-package",
     "prepare": "cd ../.. && husky install"
@@ -1882,7 +2105,7 @@
     "vscode-languageclient": "^8.0.2",
     "vscode-test-adapter-api": "^1.7.0",
     "vscode-test-adapter-util": "^0.7.0",
-    "zip-a-folder": "^2.0.0"
+    "zip-a-folder": "^3.1.3"
   },
   "devDependencies": {
     "@babel/core": "^7.18.13",
@@ -1916,7 +2139,7 @@
     "@types/gulp": "^4.0.9",
     "@types/gulp-replace": "^1.1.0",
     "@types/jest": "^29.0.2",
-    "@types/js-yaml": "^3.12.5",
+    "@types/js-yaml": "^4.0.6",
     "@types/nanoid": "^3.0.0",
     "@types/node": "^16.11.25",
     "@types/node-fetch": "^2.5.2",

extensions/ql-vscode/scripts/generate-schemas.ts

@@ -0,0 +1,72 @@
+import { createGenerator } from "ts-json-schema-generator";
+import { join, resolve } from "path";
+import { outputFile } from "fs-extra";
+import { format, resolveConfig } from "prettier";
+
+const extensionDirectory = resolve(__dirname, "..");
+
+const schemas = [
+  {
+    path: join(
+      extensionDirectory,
+      "src",
+      "model-editor",
+      "extension-pack-metadata.ts",
+    ),
+    type: "ExtensionPackMetadata",
+    schemaPath: join(
+      extensionDirectory,
+      "src",
+      "model-editor",
+      "extension-pack-metadata.schema.json",
+    ),
+  },
+  {
+    path: join(
+      extensionDirectory,
+      "src",
+      "model-editor",
+      "model-extension-file.ts",
+    ),
+    type: "ModelExtensionFile",
+    schemaPath: join(
+      extensionDirectory,
+      "src",
+      "model-editor",
+      "model-extension-file.schema.json",
+    ),
+  },
+];
+
+async function generateSchema(
+  schemaDefinition: (typeof schemas)[number],
+): Promise<void> {
+  const schema = createGenerator({
+    path: schemaDefinition.path,
+    tsconfig: resolve(extensionDirectory, "tsconfig.json"),
+    type: schemaDefinition.type,
+    skipTypeCheck: true,
+    topRef: true,
+    additionalProperties: true,
+  }).createSchema(schemaDefinition.type);
+
+  const schemaJson = JSON.stringify(schema, null, 2);
+
+  const prettierOptions = await resolveConfig(schemaDefinition.schemaPath);
+
+  const formattedSchemaJson = await format(schemaJson, {
+    ...prettierOptions,
+    filepath: schemaDefinition.schemaPath,
+  });
+
+  await outputFile(schemaDefinition.schemaPath, formattedSchemaJson);
+}
+
+async function generateSchemas() {
+  await Promise.all(schemas.map(generateSchema));
+}
+
+generateSchemas().catch((e: unknown) => {
+  console.error(e);
+  process.exit(2);
+});

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

@@ -6,7 +6,6 @@ import { dirname, join, delimiter } from "path";
 import * as sarif from "sarif";
 import { SemVer } from "semver";
 import { Readable } from "stream";
-import { StringDecoder } from "string_decoder";
 import tk from "tree-kill";
 import { promisify } from "util";
 import { CancellationToken, Disposable, Uri } from "vscode";
@@ -31,6 +30,7 @@ import { CompilationMessage } from "../query-server/legacy-messages";
 import { sarifParser } from "../common/sarif-parser";
 import { App } from "../common/app";
 import { QueryLanguage } from "../common/query-language";
+import { LINE_ENDINGS, splitStreamAtSeparators } from "../common/split-stream";
 
 /**
  * The version of the SARIF format that we are using.
@@ -1649,120 +1649,13 @@ export async function runCodeQlCliCommand(
   }
 }
 
-/**
- * Buffer to hold state used when splitting a text stream into lines.
- */
-class SplitBuffer {
-  private readonly decoder = new StringDecoder("utf8");
-  private readonly maxSeparatorLength: number;
-  private buffer = "";
-  private searchIndex = 0;
-
-  constructor(private readonly separators: readonly string[]) {
-    this.maxSeparatorLength = separators
-      .map((s) => s.length)
-      .reduce((a, b) => Math.max(a, b), 0);
-  }
-
-  /**
-   * Append new text data to the buffer.
-   * @param chunk The chunk of data to append.
-   */
-  public addChunk(chunk: Buffer): void {
-    this.buffer += this.decoder.write(chunk);
-  }
-
-  /**
-   * Signal that the end of the input stream has been reached.
-   */
-  public end(): void {
-    this.buffer += this.decoder.end();
-    this.buffer += this.separators[0]; // Append a separator to the end to ensure the last line is returned.
-  }
-
-  /**
-   * A version of startsWith that isn't overriden by a broken version of ms-python.
-   *
-   * The definition comes from
-   * https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/String/startsWith
-   * which is CC0/public domain
-   *
-   * See https://github.com/github/vscode-codeql/issues/802 for more context as to why we need it.
-   */
-  private static startsWith(
-    s: string,
-    searchString: string,
-    position: number,
-  ): boolean {
-    const pos = position > 0 ? position | 0 : 0;
-    return s.substring(pos, pos + searchString.length) === searchString;
-  }
-
-  /**
-   * Extract the next full line from the buffer, if one is available.
-   * @returns The text of the next available full line (without the separator), or `undefined` if no
-   * line is available.
-   */
-  public getNextLine(): string | undefined {
-    while (this.searchIndex <= this.buffer.length - this.maxSeparatorLength) {
-      for (const separator of this.separators) {
-        if (SplitBuffer.startsWith(this.buffer, separator, this.searchIndex)) {
-          const line = this.buffer.slice(0, this.searchIndex);
-          this.buffer = this.buffer.slice(this.searchIndex + separator.length);
-          this.searchIndex = 0;
-          return line;
-        }
-      }
-      this.searchIndex++;
-    }
-
-    return undefined;
-  }
-}
-
-/**
- * Splits a text stream into lines based on a list of valid line separators.
- * @param stream The text stream to split. This stream will be fully consumed.
- * @param separators The list of strings that act as line separators.
- * @returns A sequence of lines (not including separators).
- */
-async function* splitStreamAtSeparators(
-  stream: Readable,
-  separators: string[],
-): AsyncGenerator<string, void, unknown> {
-  const buffer = new SplitBuffer(separators);
-  for await (const chunk of stream) {
-    buffer.addChunk(chunk);
-    let line: string | undefined;
-    do {
-      line = buffer.getNextLine();
-      if (line !== undefined) {
-        yield line;
-      }
-    } while (line !== undefined);
-  }
-  buffer.end();
-  let line: string | undefined;
-  do {
-    line = buffer.getNextLine();
-    if (line !== undefined) {
-      yield line;
-    }
-  } while (line !== undefined);
-}
-
-/**
- *  Standard line endings for splitting human-readable text.
- */
-const lineEndings = ["\r\n", "\r", "\n"];
-
 /**
  * Log a text stream to a `Logger` interface.
  * @param stream The stream to log.
  * @param logger The logger that will consume the stream output.
  */
 async function logStream(stream: Readable, logger: BaseLogger): Promise<void> {
-  for await (const line of splitStreamAtSeparators(stream, lineEndings)) {
+  for await (const line of splitStreamAtSeparators(stream, LINE_ENDINGS)) {
     // Await the result of log here in order to ensure the logs are written in the correct order.
     await logger.log(line);
   }

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

@@ -219,6 +219,24 @@ export type LocalDatabasesCommands = {
   "codeQLDatabases.chooseDatabaseGithub": () => Promise<void>;
   "codeQLDatabases.sortByName": () => 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": (

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

@@ -500,14 +500,14 @@ interface SetMethodsMessage {
   methods: Method[];
 }
 
-interface LoadModeledMethodsMessage {
-  t: "loadModeledMethods";
-  modeledMethods: Record<string, ModeledMethod>;
+interface SetModeledMethodsMessage {
+  t: "setModeledMethods";
+  methods: Record<string, ModeledMethod>;
 }
 
-interface AddModeledMethodsMessage {
-  t: "addModeledMethods";
-  modeledMethods: Record<string, ModeledMethod>;
+interface SetModifiedMethodsMessage {
+  t: "setModifiedMethods";
+  methodSignatures: string[];
 }
 
 interface SetInProgressMethodsMessage {
@@ -570,11 +570,16 @@ interface HideModeledMethodsMessage {
   hideModeledMethods: boolean;
 }
 
+interface SetModeledMethodMessage {
+  t: "setModeledMethod";
+  method: ModeledMethod;
+}
+
 export type ToModelEditorMessage =
   | SetExtensionPackStateMessage
   | SetMethodsMessage
-  | LoadModeledMethodsMessage
-  | AddModeledMethodsMessage
+  | SetModeledMethodsMessage
+  | SetModifiedMethodsMessage
   | SetInProgressMethodsMessage;
 
 export type FromModelEditorMessage =
@@ -589,15 +594,33 @@ export type FromModelEditorMessage =
   | GenerateMethodsFromLlmMessage
   | StopGeneratingMethodsFromLlmMessage
   | ModelDependencyMessage
-  | HideModeledMethodsMessage;
+  | HideModeledMethodsMessage
+  | SetModeledMethodMessage;
 
 export type FromMethodModelingMessage =
   | TelemetryMessage
-  | UnhandledErrorMessage;
+  | UnhandledErrorMessage
+  | SetModeledMethodMessage;
 
 interface SetMethodMessage {
   t: "setMethod";
   method: Method;
 }
 
-export type ToMethodModelingMessage = SetMethodMessage;
+interface SetMethodModifiedMessage {
+  t: "setMethodModified";
+  isModified: boolean;
+}
+
+interface SetSelectedMethodMessage {
+  t: "setSelectedMethod";
+  method: Method;
+  modeledMethod: ModeledMethod;
+  isModified: boolean;
+}
+
+export type ToMethodModelingMessage =
+  | SetMethodMessage
+  | SetModeledMethodMessage
+  | SetMethodModifiedMessage
+  | SetSelectedMethodMessage;

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

@@ -62,3 +62,9 @@ export const dbSchemeToLanguage: Record<string, QueryLanguage> = {
 export function isQueryLanguage(language: string): language is QueryLanguage {
   return Object.values(QueryLanguage).includes(language as QueryLanguage);
 }
+
+export function tryGetQueryLanguage(
+  language: string,
+): QueryLanguage | undefined {
+  return isQueryLanguage(language) ? language : undefined;
+}

extensions/ql-vscode/src/common/split-stream.ts

@@ -0,0 +1,125 @@
+import { Readable } from "stream";
+import { StringDecoder } from "string_decoder";
+
+/**
+ * Buffer to hold state used when splitting a text stream into lines.
+ */
+export class SplitBuffer {
+  private readonly decoder = new StringDecoder("utf8");
+  private readonly maxSeparatorLength: number;
+  private buffer = "";
+  private searchIndex = 0;
+  private ended = false;
+
+  constructor(private readonly separators: readonly string[]) {
+    this.maxSeparatorLength = separators
+      .map((s) => s.length)
+      .reduce((a, b) => Math.max(a, b), 0);
+  }
+
+  /**
+   * Append new text data to the buffer.
+   * @param chunk The chunk of data to append.
+   */
+  public addChunk(chunk: Buffer): void {
+    this.buffer += this.decoder.write(chunk);
+  }
+
+  /**
+   * Signal that the end of the input stream has been reached.
+   */
+  public end(): void {
+    this.buffer += this.decoder.end();
+    this.ended = true;
+  }
+
+  /**
+   * A version of startsWith that isn't overriden by a broken version of ms-python.
+   *
+   * The definition comes from
+   * https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/String/startsWith
+   * which is CC0/public domain
+   *
+   * See https://github.com/github/vscode-codeql/issues/802 for more context as to why we need it.
+   */
+  private static startsWith(
+    s: string,
+    searchString: string,
+    position: number,
+  ): boolean {
+    const pos = position > 0 ? position | 0 : 0;
+    return s.substring(pos, pos + searchString.length) === searchString;
+  }
+
+  /**
+   * Extract the next full line from the buffer, if one is available.
+   * @returns The text of the next available full line (without the separator), or `undefined` if no
+   * line is available.
+   */
+  public getNextLine(): string | undefined {
+    // If we haven't received all of the input yet, don't search too close to the end of the buffer,
+    // or we could match a separator that's split across two chunks. For example, we could see "\r"
+    // at the end of the buffer and match that, even though we were about to receive a "\n" right
+    // after it.
+    const maxSearchIndex = this.ended
+      ? this.buffer.length - 1
+      : this.buffer.length - this.maxSeparatorLength;
+    while (this.searchIndex <= maxSearchIndex) {
+      for (const separator of this.separators) {
+        if (SplitBuffer.startsWith(this.buffer, separator, this.searchIndex)) {
+          const line = this.buffer.slice(0, this.searchIndex);
+          this.buffer = this.buffer.slice(this.searchIndex + separator.length);
+          this.searchIndex = 0;
+          return line;
+        }
+      }
+      this.searchIndex++;
+    }
+
+    if (this.ended && this.buffer.length > 0) {
+      // If we still have some text left in the buffer, return it as the last line.
+      const line = this.buffer;
+      this.buffer = "";
+      this.searchIndex = 0;
+      return line;
+    } else {
+      return undefined;
+    }
+  }
+}
+
+/**
+ * Splits a text stream into lines based on a list of valid line separators.
+ * @param stream The text stream to split. This stream will be fully consumed.
+ * @param separators The list of strings that act as line separators.
+ * @returns A sequence of lines (not including separators).
+ */
+export async function* splitStreamAtSeparators(
+  stream: Readable,
+  separators: string[],
+): AsyncGenerator<string, void, unknown> {
+  const buffer = new SplitBuffer(separators);
+  for await (const chunk of stream) {
+    buffer.addChunk(chunk);
+    let line: string | undefined;
+    do {
+      line = buffer.getNextLine();
+      if (line !== undefined) {
+        yield line;
+      }
+    } while (line !== undefined);
+  }
+  buffer.end();
+  let line: string | undefined;
+  do {
+    line = buffer.getNextLine();
+    if (line !== undefined) {
+      yield line;
+    }
+  } while (line !== undefined);
+}
+
+/**
+ *  Standard line endings for splitting human-readable text.
+ */
+export const LINE_ENDINGS = ["\r\n", "\r", "\n"];

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

@@ -9,7 +9,7 @@ import {
 import { join } from "path";
 
 import { App } from "../app";
-import { DisposableObject, DisposeHandler } from "../disposable-object";
+import { Disposable } from "../disposable-object";
 import { tmpDir } from "../../tmp-dir";
 import { getHtmlForWebview, WebviewMessage, WebviewKind } from "./webview-html";
 
@@ -27,16 +27,16 @@ export type WebviewPanelConfig = {
 export abstract class AbstractWebview<
   ToMessage extends WebviewMessage,
   FromMessage extends WebviewMessage,
-> extends DisposableObject {
+> {
   protected panel: WebviewPanel | undefined;
   protected panelLoaded = false;
   protected panelLoadedCallBacks: Array<() => void> = [];
 
   private panelResolves?: Array<(panel: WebviewPanel) => void>;
 
-  constructor(protected readonly app: App) {
-    super();
-  }
+  private disposables: Disposable[] = [];
+
+  constructor(protected readonly app: App) {}
 
   public async restoreView(panel: WebviewPanel): Promise<void> {
     this.panel = panel;
@@ -101,6 +101,7 @@ export abstract class AbstractWebview<
         this.panel = undefined;
         this.panelLoaded = false;
         this.onPanelDispose();
+        this.disposeAll();
       }, null),
     );
 
@@ -150,8 +151,27 @@ export abstract class AbstractWebview<
     return panel.webview.postMessage(msg);
   }
 
-  public dispose(disposeHandler?: DisposeHandler) {
+  public dispose() {
     this.panel?.dispose();
-    super.dispose(disposeHandler);
+    this.disposeAll();
+  }
+
+  private disposeAll() {
+    while (this.disposables.length > 0) {
+      const disposable = this.disposables.pop()!;
+      disposable.dispose();
+    }
+  }
+
+  /**
+   * Adds `obj` to a list of objects to dispose when the panel is disposed. Objects added by `push` are
+   * disposed in reverse order of being added.
+   * @param obj The object to take ownership of.
+   */
+  protected push<T extends Disposable>(obj: T): T {
+    if (obj !== undefined) {
+      this.disposables.push(obj);
+    }
+    return obj;
   }
 }

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

@@ -3,10 +3,7 @@ import { throttling } from "@octokit/plugin-throttling";
 import { Octokit } from "@octokit/rest";
 import { Progress, CancellationToken } from "vscode";
 import { Credentials } from "../common/authentication";
-import {
-  NotificationLogger,
-  showAndLogWarningMessage,
-} from "../common/logging";
+import { BaseLogger } from "../common/logging";
 
 export async function getCodeSearchRepositories(
   query: string,
@@ -16,7 +13,7 @@ export async function getCodeSearchRepositories(
   }>,
   token: CancellationToken,
   credentials: Credentials,
-  logger: NotificationLogger,
+  logger: BaseLogger,
 ): Promise<string[]> {
   let nwos: string[] = [];
   const octokit = await provideOctokitWithThrottling(credentials, logger);
@@ -47,7 +44,7 @@ export async function getCodeSearchRepositories(
 
 async function provideOctokitWithThrottling(
   credentials: Credentials,
-  logger: NotificationLogger,
+  logger: BaseLogger,
 ): Promise<Octokit> {
   const MyOctokit = Octokit.plugin(throttling);
   const auth = await credentials.getAccessToken();
@@ -57,16 +54,14 @@ async function provideOctokitWithThrottling(
     retry,
     throttle: {
       onRateLimit: (retryAfter: number, options: any): boolean => {
-        void showAndLogWarningMessage(
-          logger,
+        void logger.log(
           `Rate Limit detected for request ${options.method} ${options.url}. Retrying after ${retryAfter} seconds!`,
         );
 
         return true;
       },
       onSecondaryRateLimit: (_retryAfter: number, options: any): void => {
-        void showAndLogWarningMessage(
-          logger,
+        void logger.log(
           `Secondary Rate Limit detected for request ${options.method} ${options.url}`,
         );
       },

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

@@ -51,6 +51,8 @@ import {
   createMultiSelectionCommand,
   createSingleSelectionCommand,
 } from "../common/vscode/selection-commands";
+import { QueryLanguage, tryGetQueryLanguage } from "../common/query-language";
+import { LanguageContextStore } from "../language-context-store";
 
 enum SortOrder {
   NameAsc = "NameAsc",
@@ -73,7 +75,10 @@ class DatabaseTreeDataProvider
   );
   private currentDatabaseItem: DatabaseItem | undefined;
 
-  constructor(private databaseManager: DatabaseManager) {
+  constructor(
+    private databaseManager: DatabaseManager,
+    private languageContext: LanguageContextStore,
+  ) {
     super();
 
     this.currentDatabaseItem = databaseManager.currentDatabaseItem;
@@ -88,6 +93,11 @@ class DatabaseTreeDataProvider
         this.handleDidChangeCurrentDatabaseItem.bind(this),
       ),
     );
+    this.push(
+      this.languageContext.onLanguageContextChanged(async () => {
+        this._onDidChangeTreeData.fire(undefined);
+      }),
+    );
   }
 
   public get onDidChangeTreeData(): Event<DatabaseItem | undefined> {
@@ -131,7 +141,15 @@ class DatabaseTreeDataProvider
 
   public getChildren(element?: DatabaseItem): ProviderResult<DatabaseItem[]> {
     if (element === undefined) {
-      return this.databaseManager.databaseItems.slice(0).sort((db1, db2) => {
+      // Filter items by language
+      const displayItems = this.databaseManager.databaseItems.filter((item) => {
+        return this.languageContext.shouldInclude(
+          tryGetQueryLanguage(item.language),
+        );
+      });
+
+      // Sort items
+      return displayItems.slice(0).sort((db1, db2) => {
         switch (this.sortOrder) {
           case SortOrder.NameAsc:
             return db1.name.localeCompare(db2.name, env.language);
@@ -200,6 +218,7 @@ export class DatabaseUI extends DisposableObject {
   public constructor(
     private app: App,
     private databaseManager: DatabaseManager,
+    private languageContext: LanguageContextStore,
     private readonly queryServer: QueryRunner | undefined,
     private readonly storagePath: string,
     readonly extensionPath: string,
@@ -207,7 +226,7 @@ export class DatabaseUI extends DisposableObject {
     super();
 
     this.treeDataProvider = this.push(
-      new DatabaseTreeDataProvider(databaseManager),
+      new DatabaseTreeDataProvider(databaseManager, languageContext),
     );
     this.push(
       window.createTreeView("codeQLDatabases", {
@@ -245,6 +264,60 @@ export class DatabaseUI extends DisposableObject {
         this.handleMakeCurrentDatabase.bind(this),
       "codeQLDatabases.sortByName": this.handleSortByName.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),
       ),
@@ -535,6 +608,14 @@ 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) => {

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

@@ -409,7 +409,7 @@ export class DbPanel extends DisposableObject {
       return;
     }
 
-    void window.withProgress(
+    await window.withProgress(
       {
         location: ProgressLocation.Notification,
         title: "Searching for repositories... This might take a while",

extensions/ql-vscode/src/extension.ts

@@ -135,6 +135,7 @@ import { TestManagerBase } from "./query-testing/test-manager-base";
 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";
 
 /**
  * extension.ts
@@ -299,12 +300,12 @@ const shouldUpdateOnNextActivationKey = "shouldUpdateOnNextActivation";
 
 const codeQlVersionRange = DEFAULT_DISTRIBUTION_VERSION_RANGE;
 
-// This is the minimum version of vscode that we _want_ to support. We want to update the language server library, but that
-// requires 1.67 or later. If we change the minimum version in the package.json, then anyone on an older version of vscode
+// This is the minimum version of vscode that we _want_ to support. We want to update to Node 18, but that
+// requires 1.82 or later. If we change the minimum version in the package.json, then anyone on an older version of vscode will
 // silently be unable to upgrade. So, the solution is to first bump the minimum version here and release. Then
 // bump the version in the package.json and release again. This way, anyone on an older version of vscode will get a warning
 // before silently being refused to upgrade.
-const MIN_VERSION = "1.67.0";
+const MIN_VERSION = "1.82.0";
 
 /**
  * Returns the CodeQLExtensionInterface, or an empty object if the interface is not
@@ -774,10 +775,15 @@ async function activateWithInstalledDistribution(
   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,
     dbm,
+    languageContext,
     qs,
     getContextStoragePath(ctx),
     ctx.extensionPath,

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

@@ -0,0 +1,49 @@
+import { App } from "./common/app";
+import { DisposableObject } from "./common/disposable-object";
+import { AppEvent, AppEventEmitter } from "./common/events";
+import { QueryLanguage } from "./common/query-language";
+
+type LanguageFilter = QueryLanguage | "All";
+
+export class LanguageContextStore extends DisposableObject {
+  public readonly onLanguageContextChanged: AppEvent<void>;
+  private readonly onLanguageContextChangedEmitter: AppEventEmitter<void>;
+
+  private languageFilter: LanguageFilter;
+
+  constructor(private readonly app: App) {
+    super();
+    // State initialization
+    this.languageFilter = "All";
+
+    // Set up event emitters
+    this.onLanguageContextChangedEmitter = this.push(
+      app.createEventEmitter<void>(),
+    );
+    this.onLanguageContextChanged = this.onLanguageContextChangedEmitter.event;
+  }
+
+  public async clearLanguageContext() {
+    this.languageFilter = "All";
+    this.onLanguageContextChangedEmitter.fire();
+    await this.app.commands.execute(
+      "setContext",
+      "codeQLDatabases.languageFilter",
+      "",
+    );
+  }
+
+  public async setLanguageContext(language: QueryLanguage) {
+    this.languageFilter = language;
+    this.onLanguageContextChangedEmitter.fire();
+    await this.app.commands.execute(
+      "setContext",
+      "codeQLDatabases.languageFilter",
+      language,
+    );
+  }
+
+  public shouldInclude(language: QueryLanguage | undefined): boolean {
+    return this.languageFilter === "All" || this.languageFilter === language;
+  }
+}

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

@@ -13,6 +13,7 @@ import { redactableError } from "../common/errors";
 import { showAndLogExceptionWithTelemetry } from "../common/logging";
 import { extLogger } from "../common/logging/vscode";
 import { telemetryListener } from "../common/vscode/telemetry";
+import { SuiteInstruction } from "../packaging/suite-instruction";
 
 export async function qlpackOfDatabase(
   cli: Pick<CodeQLCliServer, "resolveQlpacks">,
@@ -38,24 +39,26 @@ export interface QueryConstraints {
  * @param cli The CLI instance to use.
  * @param qlpacks The list of packs to search.
  * @param constraints Constraints on the queries to search for.
+ * @param additionalPacks Additional pack paths to search.
  * @returns The found queries from the first pack in which any matching queries were found.
  */
-async function resolveQueriesFromPacks(
+export async function resolveQueriesFromPacks(
   cli: CodeQLCliServer,
   qlpacks: string[],
   constraints: QueryConstraints,
+  additionalPacks: string[] = [],
 ): Promise<string[]> {
   const suiteFile = (
     await file({
       postfix: ".qls",
     })
   ).path;
-  const suiteYaml = [];
+  const suiteYaml: SuiteInstruction[] = [];
   for (const qlpack of qlpacks) {
     suiteYaml.push({
       from: qlpack,
       queries: ".",
-      include: constraints,
+      include: constraints as Record<string, string[]>,
     });
   }
   await writeFile(
@@ -66,10 +69,10 @@ async function resolveQueriesFromPacks(
     "utf8",
   );
 
-  return await cli.resolveQueriesInSuite(
-    suiteFile,
-    getOnDiskWorkspaceFolders(),
-  );
+  return await cli.resolveQueriesInSuite(suiteFile, [
+    ...getOnDiskWorkspaceFolders(),
+    ...additionalPacks,
+  ]);
 }
 
 export async function resolveQueriesByLanguagePack(
@@ -96,6 +99,7 @@ export async function resolveQueriesByLanguagePack(
  * @param packsToSearch The list of packs to search.
  * @param name The name of the query to use in error messages.
  * @param constraints Constraints on the queries to search for.
+ * @param additionalPacks Additional pack paths to search.
  * @returns The found queries from the first pack in which any matching queries were found.
  */
 export async function resolveQueries(
@@ -103,11 +107,13 @@ export async function resolveQueries(
   packsToSearch: string[],
   name: string,
   constraints: QueryConstraints,
+  additionalPacks: string[] = [],
 ): Promise<string[]> {
   const queries = await resolveQueriesFromPacks(
     cli,
     packsToSearch,
     constraints,
+    additionalPacks,
   );
   if (queries.length > 0) {
     return queries;

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

@@ -75,6 +75,7 @@ import { telemetryListener } from "../common/vscode/telemetry";
 import { redactableError } from "../common/errors";
 import { ResultsViewCommands } from "../common/commands";
 import { App } from "../common/app";
+import { Disposable } from "../common/disposable-object";
 
 /**
  * results-view.ts
@@ -157,6 +158,12 @@ function numInterpretedPages(
   return Math.ceil(n / pageSize);
 }
 
+/**
+ * The results view is used for displaying the results of a local query. It is a singleton; only 1 results view exists
+ * in the extension. It is created when the extension is activated and disposed of when the extension is deactivated.
+ * There can be multiple panels linked to this view over the lifetime of the extension, but there is only ever 1 panel
+ * active at a time.
+ */
 export class ResultsView extends AbstractWebview<
   IntoResultsViewMsg,
   FromResultsViewMsg
@@ -168,6 +175,9 @@ export class ResultsView extends AbstractWebview<
     "codeql-query-results",
   );
 
+  // Event listeners that should be disposed of when the view is disposed.
+  private disposableEventListeners: Disposable[] = [];
+
   constructor(
     app: App,
     private databaseManager: DatabaseManager,
@@ -176,14 +186,16 @@ export class ResultsView extends AbstractWebview<
     private labelProvider: HistoryItemLabelProvider,
   ) {
     super(app);
-    this.push(this._diagnosticCollection);
-    this.push(
+
+    // We can't use this.push for these two event listeners because they need to be disposed of when the view is
+    // disposed, not when the panel is disposed. The results view is a singleton, so we shouldn't be calling this.push.
+    this.disposableEventListeners.push(
       vscode.window.onDidChangeTextEditorSelection(
         this.handleSelectionChange.bind(this),
       ),
     );
 
-    this.push(
+    this.disposableEventListeners.push(
       this.databaseManager.onDidChangeDatabaseItem(({ kind }) => {
         if (kind === DatabaseEventKind.Remove) {
           this._diagnosticCollection.clear();
@@ -981,4 +993,12 @@ export class ResultsView extends AbstractWebview<
       editor.setDecorations(shownLocationLineDecoration, []);
     }
   }
+
+  dispose() {
+    super.dispose();
+
+    this._diagnosticCollection.dispose();
+    this.disposableEventListeners.forEach((d) => d.dispose());
+    this.disposableEventListeners = [];
+  }
 }

extensions/ql-vscode/src/log-insights/summary-parser.ts

@@ -1,4 +1,5 @@
-import { writeFile, promises } from "fs-extra";
+import { createReadStream, writeFile } from "fs-extra";
+import { LINE_ENDINGS, splitStreamAtSeparators } from "../common/split-stream";
 
 /**
  * Location information for a single pipeline invocation in the RA.
@@ -64,59 +65,64 @@ export async function generateSummarySymbolsFile(
 async function generateSummarySymbols(
   summaryPath: string,
 ): Promise<SummarySymbols> {
-  const summary = await promises.readFile(summaryPath, {
+  const stream = createReadStream(summaryPath, {
     encoding: "utf-8",
   });
-  const symbols: SummarySymbols = {
-    predicates: {},
-  };
+  try {
+    const lines = splitStreamAtSeparators(stream, LINE_ENDINGS);
 
-  const lines = summary.split(/\r?\n/);
-  let lineNumber = 0;
-  while (lineNumber < lines.length) {
-    const startLineNumber = lineNumber;
-    lineNumber++;
-    const startLine = lines[startLineNumber];
-    const nonRecursiveMatch = startLine.match(NON_RECURSIVE_TUPLE_COUNT_REGEXP);
-    let predicateName: string | undefined = undefined;
+    const symbols: SummarySymbols = {
+      predicates: {},
+    };
+
+    let lineNumber = 0;
+    let raStartLine = 0;
     let iteration = 0;
-    if (nonRecursiveMatch) {
-      predicateName = nonRecursiveMatch.groups!.predicateName;
-    } else {
-      const recursiveMatch = startLine.match(RECURSIVE_TUPLE_COUNT_REGEXP);
-      if (recursiveMatch?.groups) {
-        predicateName = recursiveMatch.groups.predicateName;
-        iteration = parseInt(recursiveMatch.groups.iteration);
-      }
-    }
-
-    if (predicateName !== undefined) {
-      const raStartLine = lineNumber;
-      let raEndLine: number | undefined = undefined;
-      while (lineNumber < lines.length && raEndLine === undefined) {
-        const raLine = lines[lineNumber];
-        const returnMatch = raLine.match(RETURN_REGEXP);
+    let predicateName: string | undefined = undefined;
+    let startLine = 0;
+    for await (const line of lines) {
+      if (predicateName === undefined) {
+        // Looking for the start of the predicate.
+        const nonRecursiveMatch = line.match(NON_RECURSIVE_TUPLE_COUNT_REGEXP);
+        if (nonRecursiveMatch) {
+          iteration = 0;
+          predicateName = nonRecursiveMatch.groups!.predicateName;
+        } else {
+          const recursiveMatch = line.match(RECURSIVE_TUPLE_COUNT_REGEXP);
+          if (recursiveMatch?.groups) {
+            predicateName = recursiveMatch.groups.predicateName;
+            iteration = parseInt(recursiveMatch.groups.iteration);
+          }
+        }
+        if (predicateName !== undefined) {
+          startLine = lineNumber;
+          raStartLine = lineNumber + 1;
+        }
+      } else {
+        const returnMatch = line.match(RETURN_REGEXP);
         if (returnMatch) {
-          raEndLine = lineNumber;
-        }
-        lineNumber++;
-      }
-      if (raEndLine !== undefined) {
-        let symbol = symbols.predicates[predicateName];
-        if (symbol === undefined) {
-          symbol = {
-            iterations: {},
+          let symbol = symbols.predicates[predicateName];
+          if (symbol === undefined) {
+            symbol = {
+              iterations: {},
+            };
+            symbols.predicates[predicateName] = symbol;
+          }
+          symbol.iterations[iteration] = {
+            startLine,
+            raStartLine,
+            raEndLine: lineNumber,
           };
-          symbols.predicates[predicateName] = symbol;
-        }
-        symbol.iterations[iteration] = {
-          startLine: lineNumber,
-          raStartLine,
-          raEndLine,
-        };
-      }
-    }
-  }
 
-  return symbols;
+          predicateName = undefined;
+        }
+      }
+
+      lineNumber++;
+    }
+
+    return symbols;
+  } finally {
+    stream.close();
+  }
 }

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

@@ -7,7 +7,6 @@ import { Mode } from "./shared/mode";
 import { getOnDiskWorkspaceFolders } from "../common/vscode/workspace-folders";
 import { interpretResultsSarif } from "../query-results";
 import { join } from "path";
-import { assertNever } from "../common/helpers-pure";
 import { dir } from "tmp-promise";
 import { writeFile, outputFile } from "fs-extra";
 import { dump as dumpYaml } from "js-yaml";
@@ -16,17 +15,7 @@ import { runQuery } from "../local-queries/run-query";
 import { QueryMetadata } from "../common/interface-types";
 import { CancellationTokenSource } from "vscode";
 import { resolveQueries } from "../local-queries";
-
-function modeTag(mode: Mode): string {
-  switch (mode) {
-    case Mode.Application:
-      return "application-mode";
-    case Mode.Framework:
-      return "framework-mode";
-    default:
-      assertNever(mode);
-  }
-}
+import { modeTag } from "./mode-tag";
 
 type AutoModelQueriesOptions = {
   mode: Mode;

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

@@ -2,41 +2,65 @@ import { DecodedBqrsChunk } from "../common/bqrs-cli-types";
 import { Call, CallClassification, Method } from "./method";
 import { ModeledMethodType } from "./modeled-method";
 import { parseLibraryFilename } from "./library";
+import { Mode } from "./shared/mode";
+import { ApplicationModeTuple, FrameworkModeTuple } from "./queries/query";
 
-export function decodeBqrsToMethods(chunk: DecodedBqrsChunk): Method[] {
+export function decodeBqrsToMethods(
+  chunk: DecodedBqrsChunk,
+  mode: Mode,
+): Method[] {
   const methodsByApiName = new Map<string, Method>();
 
   chunk?.tuples.forEach((tuple) => {
-    const usage = tuple[0] as Call;
-    const signature = tuple[1] as string;
-    const supported = (tuple[2] as string) === "true";
-    let library = tuple[4] as string;
-    let libraryVersion: string | undefined = tuple[5] as string;
-    const type = tuple[6] as ModeledMethodType;
-    const classification = tuple[8] as CallClassification;
+    let usage: Call;
+    let packageName: string;
+    let typeName: string;
+    let methodName: string;
+    let methodParameters: string;
+    let supported: boolean;
+    let library: string;
+    let libraryVersion: string | undefined;
+    let type: ModeledMethodType;
+    let classification: CallClassification;
 
-    const [packageWithType, methodDeclaration] = signature.split("#");
+    if (mode === Mode.Application) {
+      [
+        usage,
+        packageName,
+        typeName,
+        methodName,
+        methodParameters,
+        supported,
+        library,
+        libraryVersion,
+        type,
+        classification,
+      ] = tuple as ApplicationModeTuple;
+    } else {
+      [
+        usage,
+        packageName,
+        typeName,
+        methodName,
+        methodParameters,
+        supported,
+        library,
+        type,
+      ] = tuple as FrameworkModeTuple;
 
-    const packageName = packageWithType.substring(
-      0,
-      packageWithType.lastIndexOf("."),
-    );
-    const typeName = packageWithType.substring(
-      packageWithType.lastIndexOf(".") + 1,
-    );
+      classification = CallClassification.Unknown;
+    }
 
-    const methodName = methodDeclaration.substring(
-      0,
-      methodDeclaration.indexOf("("),
-    );
-    const methodParameters = methodDeclaration.substring(
-      methodDeclaration.indexOf("("),
-    );
+    const signature = `${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
     // for Java.
-    if (library.endsWith(".jar") || libraryVersion === "") {
+    if (
+      library.endsWith(".jar") ||
+      libraryVersion === "" ||
+      libraryVersion === undefined
+    ) {
       const { name, version } = parseLibraryFilename(library);
       library = name;
       if (version) {

extensions/ql-vscode/src/model-editor/data-schema.json

@@ -1,45 +0,0 @@
-{
-  "type": "object",
-  "properties": {
-    "extensions": {
-      "type": "array",
-      "items": {
-        "type": "object",
-        "required": ["addsTo", "data"],
-        "properties": {
-          "addsTo": {
-            "type": "object",
-            "required": ["pack", "extensible"],
-            "properties": {
-              "pack": {
-                "type": "string"
-              },
-              "extensible": {
-                "type": "string"
-              }
-            }
-          },
-          "data": {
-            "type": "array",
-            "items": {
-              "type": "array",
-              "items": {
-                "oneOf": [
-                  {
-                    "type": "string"
-                  },
-                  {
-                    "type": "boolean"
-                  },
-                  {
-                    "type": "number"
-                  }
-                ]
-              }
-            }
-          }
-        }
-      }
-    }
-  }
-}

extensions/ql-vscode/src/model-editor/extension-pack-metadata.schema.json

@@ -0,0 +1,97 @@
+{
+  "$schema": "http://json-schema.org/draft-07/schema#",
+  "$ref": "#/definitions/ExtensionPackMetadata",
+  "definitions": {
+    "ExtensionPackMetadata": {
+      "type": "object",
+      "properties": {
+        "extensionTargets": {
+          "type": "object",
+          "additionalProperties": {
+            "type": "string"
+          }
+        },
+        "dataExtensions": {
+          "anyOf": [
+            {
+              "type": "array",
+              "items": {
+                "type": "string"
+              }
+            },
+            {
+              "type": "string"
+            }
+          ]
+        },
+        "name": {
+          "type": "string"
+        },
+        "version": {
+          "type": "string"
+        },
+        "dependencies": {
+          "type": "object",
+          "additionalProperties": {
+            "type": "string"
+          }
+        },
+        "dbscheme": {
+          "type": "string"
+        },
+        "library": {
+          "type": "boolean"
+        },
+        "defaultSuite": {
+          "type": "array",
+          "items": {
+            "$ref": "#/definitions/SuiteInstruction"
+          }
+        },
+        "defaultSuiteFile": {
+          "type": "string"
+        }
+      },
+      "required": ["dataExtensions", "extensionTargets", "name", "version"]
+    },
+    "SuiteInstruction": {
+      "type": "object",
+      "properties": {
+        "qlpack": {
+          "type": "string"
+        },
+        "query": {
+          "type": "string"
+        },
+        "queries": {
+          "type": "string"
+        },
+        "include": {
+          "type": "object",
+          "additionalProperties": {
+            "type": "array",
+            "items": {
+              "type": "string"
+            }
+          }
+        },
+        "exclude": {
+          "type": "object",
+          "additionalProperties": {
+            "type": "array",
+            "items": {
+              "type": "string"
+            }
+          }
+        },
+        "description": {
+          "type": "string"
+        },
+        "from": {
+          "type": "string"
+        }
+      },
+      "description": "A single entry in a .qls file."
+    }
+  }
+}

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

@@ -0,0 +1,7 @@
+import { QlPackFile } from "../packaging/qlpack-file";
+
+export type ExtensionPackMetadata = QlPackFile & {
+  // Make both extensionTargets and dataExtensions required
+  extensionTargets: Record<string, string>;
+  dataExtensions: string[] | string;
+};

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

@@ -2,6 +2,7 @@ import { join } from "path";
 import { outputFile, pathExists, readFile } from "fs-extra";
 import { dump as dumpYaml, load as loadYaml } from "js-yaml";
 import { Uri } from "vscode";
+import Ajv from "ajv";
 import { CodeQLCliServer } from "../codeql-cli/cli";
 import { getOnDiskWorkspaceFolders } from "../common/vscode/workspace-folders";
 import { ProgressCallback } from "../common/vscode/progress";
@@ -18,6 +19,12 @@ import {
 } from "./extension-pack-name";
 import { autoPickExtensionsDirectory } from "./extensions-workspace-folder";
 
+import { ExtensionPackMetadata } from "./extension-pack-metadata";
+import * as extensionPackMetadataSchemaJson from "./extension-pack-metadata.schema.json";
+
+const ajv = new Ajv({ allErrors: true });
+const extensionPackValidate = ajv.compile(extensionPackMetadataSchemaJson);
+
 export async function pickExtensionPack(
   cliServer: Pick<CodeQLCliServer, "resolveQlpacks">,
   databaseItem: Pick<DatabaseItem, "name" | "language">,
@@ -170,6 +177,22 @@ async function writeExtensionPack(
   return extensionPack;
 }
 
+function validateExtensionPack(
+  extensionPack: unknown,
+): extensionPack is ExtensionPackMetadata {
+  extensionPackValidate(extensionPack);
+
+  if (extensionPackValidate.errors) {
+    throw new Error(
+      `Invalid extension pack YAML: ${extensionPackValidate.errors
+        .map((error) => `${error.instancePath} ${error.message}`)
+        .join(", ")}`,
+    );
+  }
+
+  return true;
+}
+
 async function readExtensionPack(
   path: string,
   language: string,
@@ -188,6 +211,10 @@ async function readExtensionPack(
     throw new Error(`Could not parse ${qlpackPath}`);
   }
 
+  if (!validateExtensionPack(qlpack)) {
+    throw new Error(`Could not validate ${qlpackPath}`);
+  }
+
   const dataExtensionValue = qlpack.dataExtensions;
   if (
     !(

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

@@ -16,6 +16,10 @@ 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;
@@ -88,7 +92,28 @@ export async function runExternalApiQueries(
     await cliServer.resolveQlpacks(additionalPacks, true),
   );
 
-  const queryPath = join(queryDir, queryNameFromMode(mode));
+  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({
@@ -132,7 +157,7 @@ export async function runExternalApiQueries(
     maxStep: externalApiQueriesProgressMaxStep,
   });
 
-  return decodeBqrsToMethods(bqrsChunk);
+  return decodeBqrsToMethods(bqrsChunk, mode);
 }
 
 type GetResultsOptions = {
@@ -160,7 +185,5 @@ export async function readQueryResults({
 }
 
 function queryNameFromMode(mode: Mode): string {
-  return `FetchExternalApis${
-    mode.charAt(0).toUpperCase() + mode.slice(1)
-  }Mode.ql`;
+  return `${mode.charAt(0).toUpperCase() + mode.slice(1)}ModeEndpoints.ql`;
 }

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

@@ -3,14 +3,16 @@ import { App } from "../../common/app";
 import { DisposableObject } from "../../common/disposable-object";
 import { MethodModelingViewProvider } from "./method-modeling-view-provider";
 import { Method } from "../method";
+import { ModelingStore } from "../modeling-store";
 
 export class MethodModelingPanel extends DisposableObject {
   private readonly provider: MethodModelingViewProvider;
 
-  constructor(app: App) {
+  constructor(app: App, modelingStore: ModelingStore) {
     super();
 
-    this.provider = new MethodModelingViewProvider(app);
+    this.provider = new MethodModelingViewProvider(app, modelingStore);
+    this.push(this.provider);
     this.push(
       window.registerWebviewViewProvider(
         MethodModelingViewProvider.viewType,

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

@@ -8,13 +8,25 @@ 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";
 
-export class MethodModelingViewProvider implements WebviewViewProvider {
+export class MethodModelingViewProvider
+  extends DisposableObject
+  implements WebviewViewProvider
+{
   public static readonly viewType = "codeQLMethodModeling";
 
   private webviewView: vscode.WebviewView | undefined = undefined;
 
-  constructor(private readonly app: App) {}
+  private method: Method | undefined = undefined;
+
+  constructor(
+    private readonly app: App,
+    private readonly modelingStore: ModelingStore,
+  ) {
+    super();
+  }
 
   /**
    * This is called when a view first becomes visible. This may happen when the view is
@@ -45,9 +57,14 @@ export class MethodModelingViewProvider implements WebviewViewProvider {
     webviewView.webview.onDidReceiveMessage(async (msg) => this.onMessage(msg));
 
     this.webviewView = webviewView;
+
+    this.setInitialState(webviewView);
+    this.registerToModelingStoreEvents();
   }
 
   public async setMethod(method: Method): Promise<void> {
+    this.method = method;
+
     if (this.webviewView) {
       await this.webviewView.webview.postMessage({
         t: "setMethod",
@@ -56,8 +73,32 @@ export class MethodModelingViewProvider implements WebviewViewProvider {
     }
   }
 
+  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 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,
+        );
+        break;
+      }
+
       case "telemetry": {
         telemetryListener?.sendUIInteraction(msg.action);
         break;
@@ -73,4 +114,40 @@ export class MethodModelingViewProvider implements WebviewViewProvider {
         break;
     }
   }
+
+  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,
+          });
+        }
+      }
+    });
+
+    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.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,
+        });
+      }
+    });
+  }
 }

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

@@ -57,3 +57,11 @@ export interface Method extends MethodSignature {
   supportedType: ModeledMethodType;
   usages: Usage[];
 }
+
+export function getArgumentsList(methodParameters: string): string[] {
+  if (methodParameters === "()") {
+    return [];
+  }
+
+  return methodParameters.substring(1, methodParameters.length - 1).split(",");
+}

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

@@ -14,6 +14,9 @@ import { DatabaseItem } from "../../databases/local-databases";
 import { relative } from "path";
 import { CodeQLCliServer } from "../../codeql-cli/cli";
 import { INITIAL_HIDE_MODELED_METHODS_VALUE } from "../shared/hide-modeled-methods";
+import { getModelingStatus } from "../shared/modeling-status";
+import { assertNever } from "../../common/helpers-pure";
+import { ModeledMethod } from "../modeled-method";
 
 export class MethodsUsageDataProvider
   extends DisposableObject
@@ -23,6 +26,8 @@ export class MethodsUsageDataProvider
   private databaseItem: DatabaseItem | undefined = undefined;
   private sourceLocationPrefix: string | undefined = undefined;
   private hideModeledMethods: boolean = INITIAL_HIDE_MODELED_METHODS_VALUE;
+  private modeledMethods: Record<string, ModeledMethod> = {};
+  private modifiedMethodSignatures: Set<string> = new Set();
 
   private readonly onDidChangeTreeDataEmitter = this.push(
     new EventEmitter<void>(),
@@ -47,17 +52,23 @@ export class MethodsUsageDataProvider
     methods: Method[],
     databaseItem: DatabaseItem,
     hideModeledMethods: boolean,
+    modeledMethods: Record<string, ModeledMethod>,
+    modifiedMethodSignatures: Set<string>,
   ): Promise<void> {
     if (
       this.methods !== methods ||
       this.databaseItem !== databaseItem ||
-      this.hideModeledMethods !== hideModeledMethods
+      this.hideModeledMethods !== hideModeledMethods ||
+      this.modeledMethods !== modeledMethods ||
+      this.modifiedMethodSignatures !== modifiedMethodSignatures
     ) {
       this.methods = methods;
       this.databaseItem = databaseItem;
       this.sourceLocationPrefix =
         await this.databaseItem.getSourceLocationPrefix(this.cliServer);
       this.hideModeledMethods = hideModeledMethods;
+      this.modeledMethods = modeledMethods;
+      this.modifiedMethodSignatures = modifiedMethodSignatures;
 
       this.onDidChangeTreeDataEmitter.fire();
     }
@@ -68,7 +79,7 @@ export class MethodsUsageDataProvider
       return {
         label: `${item.packageName}.${item.typeName}.${item.methodName}${item.methodParameters}`,
         collapsibleState: TreeItemCollapsibleState.Collapsed,
-        iconPath: new ThemeIcon("symbol-method"),
+        iconPath: this.getModelingStatusIcon(item),
       };
     } else {
       const method = this.getParent(item);
@@ -83,11 +94,30 @@ export class MethodsUsageDataProvider
           command: "codeQLModelEditor.jumpToUsageLocation",
           arguments: [method, item, this.databaseItem],
         },
-        iconPath: new ThemeIcon("error", new ThemeColor("errorForeground")),
       };
     }
   }
 
+  private getModelingStatusIcon(method: Method): ThemeIcon {
+    const modeledMethod = this.modeledMethods[method.signature];
+    const modifiedMethod = this.modifiedMethodSignatures.has(method.signature);
+
+    const status = getModelingStatus(modeledMethod, modifiedMethod);
+    switch (status) {
+      case "unmodeled":
+        return new ThemeIcon("error", new ThemeColor("errorForeground"));
+      case "unsaved":
+        return new ThemeIcon("pass", new ThemeColor("testing.iconPassed"));
+      case "saved":
+        return new ThemeIcon(
+          "pass-filled",
+          new ThemeColor("testing.iconPassed"),
+        );
+      default:
+        assertNever(status);
+    }
+  }
+
   private relativePathWithinDatabase(uri: string): string {
     const parsedUri = Uri.parse(uri);
     if (this.sourceLocationPrefix) {

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

@@ -7,12 +7,17 @@ import {
 import { Method, Usage } from "../method";
 import { DatabaseItem } from "../../databases/local-databases";
 import { CodeQLCliServer } from "../../codeql-cli/cli";
+import { ModelingStore } from "../modeling-store";
+import { ModeledMethod } from "../modeled-method";
 
 export class MethodsUsagePanel extends DisposableObject {
   private readonly dataProvider: MethodsUsageDataProvider;
   private readonly treeView: TreeView<MethodsUsageTreeViewItem>;
 
-  public constructor(cliServer: CodeQLCliServer) {
+  public constructor(
+    private readonly modelingStore: ModelingStore,
+    cliServer: CodeQLCliServer,
+  ) {
     super();
 
     this.dataProvider = new MethodsUsageDataProvider(cliServer);
@@ -21,14 +26,24 @@ export class MethodsUsagePanel extends DisposableObject {
       treeDataProvider: this.dataProvider,
     });
     this.push(this.treeView);
+
+    this.registerToModelingStoreEvents();
   }
 
   public async setState(
     methods: Method[],
     databaseItem: DatabaseItem,
     hideModeledMethods: boolean,
+    modeledMethods: Record<string, ModeledMethod>,
+    modifiedMethodSignatures: Set<string>,
   ): Promise<void> {
-    await this.dataProvider.setState(methods, databaseItem, hideModeledMethods);
+    await this.dataProvider.setState(
+      methods,
+      databaseItem,
+      hideModeledMethods,
+      modeledMethods,
+      modifiedMethodSignatures,
+    );
     const numOfApis = hideModeledMethods
       ? methods.filter((api) => !api.supported).length
       : methods.length;
@@ -44,4 +59,49 @@ export class MethodsUsagePanel extends DisposableObject {
       await this.treeView.reveal(canonicalUsage);
     }
   }
+
+  private registerToModelingStoreEvents(): void {
+    this.push(
+      this.modelingStore.onActiveDbChanged(async () => {
+        await this.handleStateChangeEvent();
+      }),
+    );
+
+    this.push(
+      this.modelingStore.onMethodsChanged(async (event) => {
+        if (event.isActiveDb) {
+          await this.handleStateChangeEvent();
+        }
+      }),
+    );
+
+    this.push(
+      this.modelingStore.onHideModeledMethodsChanged(async (event) => {
+        if (event.isActiveDb) {
+          await this.handleStateChangeEvent();
+        }
+      }),
+    );
+
+    this.push(
+      this.modelingStore.onModifiedMethodsChanged(async (event) => {
+        if (event.isActiveDb) {
+          await this.handleStateChangeEvent();
+        }
+      }),
+    );
+  }
+
+  private async handleStateChangeEvent(): Promise<void> {
+    const activeState = this.modelingStore.getStateForActiveDb();
+    if (activeState !== undefined) {
+      await this.setState(
+        activeState.methods,
+        activeState.databaseItem,
+        activeState.hideModeledMethods,
+        activeState.modeledMethods,
+        activeState.modifiedMethodSignatures,
+      );
+    }
+  }
 }

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

@@ -0,0 +1,13 @@
+import { Mode } from "./shared/mode";
+import { assertNever } from "../common/helpers-pure";
+
+export function modeTag(mode: Mode): string {
+  switch (mode) {
+    case Mode.Application:
+      return "application-mode";
+    case Mode.Framework:
+      return "framework-mode";
+    default:
+      assertNever(mode);
+  }
+}

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

@@ -15,20 +15,20 @@ 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 { showResolvableLocation } from "../databases/local-databases/locations";
 import { Method, Usage } from "./method";
 import { setUpPack } from "./model-editor-queries";
 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"];
 
 export class ModelEditorModule extends DisposableObject {
   private readonly queryStorageDir: string;
+  private readonly modelingStore: ModelingStore;
   private readonly methodsUsagePanel: MethodsUsagePanel;
   private readonly methodModelingPanel: MethodModelingPanel;
 
-  private mostRecentlyActiveView: ModelEditorView | undefined = undefined;
-
   private constructor(
     private readonly app: App,
     private readonly databaseManager: DatabaseManager,
@@ -38,22 +38,15 @@ export class ModelEditorModule extends DisposableObject {
   ) {
     super();
     this.queryStorageDir = join(baseQueryStorageDir, "model-editor-results");
-    this.methodsUsagePanel = this.push(new MethodsUsagePanel(cliServer));
-    this.methodModelingPanel = this.push(new MethodModelingPanel(app));
-  }
+    this.modelingStore = new ModelingStore(app);
+    this.methodsUsagePanel = this.push(
+      new MethodsUsagePanel(this.modelingStore, cliServer),
+    );
+    this.methodModelingPanel = this.push(
+      new MethodModelingPanel(app, this.modelingStore),
+    );
 
-  private handleViewBecameActive(view: ModelEditorView): void {
-    this.mostRecentlyActiveView = view;
-  }
-
-  private handleViewWasDisposed(view: ModelEditorView): void {
-    if (this.mostRecentlyActiveView === view) {
-      this.mostRecentlyActiveView = undefined;
-    }
-  }
-
-  private isMostRecentlyActiveView(view: ModelEditorView): boolean {
-    return this.mostRecentlyActiveView === view;
+    this.registerToModelingStoreEvents();
   }
 
   public static async initialize(
@@ -139,6 +132,7 @@ export class ModelEditorModule extends DisposableObject {
             const { path: queryDir, cleanup: cleanupQueryDir } = await dir({
               unsafeCleanup: true,
             });
+
             const success = await setUpPack(this.cliServer, queryDir, language);
             if (!success) {
               await cleanupQueryDir();
@@ -153,6 +147,7 @@ export class ModelEditorModule extends DisposableObject {
 
             const view = new ModelEditorView(
               this.app,
+              this.modelingStore,
               this.databaseManager,
               this.cliServer,
               this.queryRunner,
@@ -161,16 +156,14 @@ export class ModelEditorModule extends DisposableObject {
               db,
               modelFile,
               Mode.Application,
-              this.methodsUsagePanel.setState.bind(this.methodsUsagePanel),
-              this.showMethod.bind(this),
-              this.handleViewBecameActive.bind(this),
-              (view) => {
-                this.handleViewWasDisposed(view);
-                void cleanupQueryDir();
-              },
-              this.isMostRecentlyActiveView.bind(this),
             );
 
+            this.modelingStore.onDbClosed(async (dbUri) => {
+              if (dbUri === db.databaseUri.toString()) {
+                await cleanupQueryDir();
+              }
+            });
+
             this.push(view);
             this.push({
               dispose(): void {
@@ -190,8 +183,7 @@ export class ModelEditorModule extends DisposableObject {
         usage: Usage,
         databaseItem: DatabaseItem,
       ) => {
-        await this.methodModelingPanel.setMethod(method);
-        await showResolvableLocation(usage.url, databaseItem, this.app.logger);
+        this.modelingStore.setSelectedMethod(databaseItem, method, usage);
       },
     };
   }
@@ -200,8 +192,21 @@ export class ModelEditorModule extends DisposableObject {
     await ensureDir(this.queryStorageDir);
   }
 
-  private async showMethod(method: Method, usage: Usage): Promise<void> {
+  private registerToModelingStoreEvents(): void {
+    this.push(
+      this.modelingStore.onSelectedMethodChanged(async (event) => {
+        await this.showMethod(event.databaseItem, event.method, event.usage);
+      }),
+    );
+  }
+
+  private async showMethod(
+    databaseItem: DatabaseItem,
+    method: Method,
+    usage: Usage,
+  ): Promise<void> {
     await this.methodsUsagePanel.revealItem(usage);
     await this.methodModelingPanel.setMethod(method);
+    await showResolvableLocation(usage.url, databaseItem, this.app.logger);
   }
 }

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

@@ -4,9 +4,28 @@ import { writeFile } from "fs-extra";
 import { dump } from "js-yaml";
 import { prepareExternalApiQuery } from "./external-api-usage-queries";
 import { CodeQLCliServer } from "../codeql-cli/cli";
+import { showLlmGeneration } from "../config";
+import { Mode } from "./shared/mode";
+import { resolveQueriesFromPacks } from "../local-queries";
+import { modeTag } from "./mode-tag";
+
+export const syntheticQueryPackName = "codeql/external-api-usage";
 
 /**
- * setUpPack sets up a directory to use for the data extension 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 queryDir The directory to set up.
  * @param language The language to use for the queries.
  * @returns true if the setup was successful, false otherwise.
@@ -16,31 +35,104 @@ export async function setUpPack(
   queryDir: string,
   language: QueryLanguage,
 ): Promise<boolean> {
-  // Create the external API query
-  const externalApiQuerySuccess = await prepareExternalApiQuery(
-    queryDir,
+  // 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 (!externalApiQuerySuccess) {
-    return false;
+
+  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,
+    );
+    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);
   }
 
-  // Set up a synthetic pack so that the query can be resolved later.
-  const syntheticQueryPack = {
-    name: "codeql/external-api-usage",
-    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);
-
-  // Install the other needed query packs
-  await cliServer.packDownload([`codeql/${language}-queries`]);
-  await cliServer.packDownload([`codeql/${language}-automodel-queries`]);
+  // 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];
+
+  // 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-view.ts

@@ -1,4 +1,11 @@
-import { CancellationTokenSource, Uri, ViewColumn, window } from "vscode";
+import {
+  CancellationTokenSource,
+  Tab,
+  TabInputWebview,
+  Uri,
+  ViewColumn,
+  window,
+} from "vscode";
 import {
   AbstractWebview,
   WebviewPanelConfig,
@@ -19,7 +26,6 @@ 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 { showResolvableLocation } from "../databases/local-databases/locations";
 import { redactableError } from "../common/errors";
 import {
   externalApiQueriesProgressMaxStep,
@@ -34,8 +40,8 @@ import { loadModeledMethods, saveModeledMethods } from "./modeled-method-fs";
 import { pickExtensionPack } from "./extension-pack-picker";
 import { getLanguageDisplayName } from "../common/query-language";
 import { AutoModeler } from "./auto-modeler";
-import { INITIAL_HIDE_MODELED_METHODS_VALUE } from "./shared/hide-modeled-methods";
 import { telemetryListener } from "../common/vscode/telemetry";
+import { ModelingStore } from "./modeling-store";
 
 export class ModelEditorView extends AbstractWebview<
   ToModelEditorMessage,
@@ -43,11 +49,9 @@ export class ModelEditorView extends AbstractWebview<
 > {
   private readonly autoModeler: AutoModeler;
 
-  private methods: Method[];
-  private hideModeledMethods: boolean;
-
   public constructor(
     protected readonly app: App,
+    private readonly modelingStore: ModelingStore,
     private readonly databaseManager: DatabaseManager,
     private readonly cliServer: CodeQLCliServer,
     private readonly queryRunner: QueryRunner,
@@ -56,23 +60,12 @@ export class ModelEditorView extends AbstractWebview<
     private readonly databaseItem: DatabaseItem,
     private readonly extensionPack: ExtensionPack,
     private mode: Mode,
-    private readonly updateMethodsUsagePanelState: (
-      methods: Method[],
-      databaseItem: DatabaseItem,
-      hideModeledMethods: boolean,
-    ) => Promise<void>,
-    private readonly showMethod: (
-      method: Method,
-      usage: Usage,
-    ) => Promise<void>,
-    private readonly handleViewBecameActive: (view: ModelEditorView) => void,
-    private readonly handleViewWasDisposed: (view: ModelEditorView) => void,
-    private readonly isMostRecentlyActiveView: (
-      view: ModelEditorView,
-    ) => boolean,
   ) {
     super(app);
 
+    this.modelingStore.initializeStateForDb(databaseItem);
+    this.registerToModelingStoreEvents();
+
     this.autoModeler = new AutoModeler(
       app,
       cliServer,
@@ -87,11 +80,9 @@ export class ModelEditorView extends AbstractWebview<
         });
       },
       async (modeledMethods) => {
-        await this.postMessage({ t: "addModeledMethods", modeledMethods });
+        this.addModeledMethods(modeledMethods);
       },
     );
-    this.methods = [];
-    this.hideModeledMethods = INITIAL_HIDE_MODELED_METHODS_VALUE;
   }
 
   public async openView() {
@@ -100,17 +91,15 @@ export class ModelEditorView extends AbstractWebview<
 
     panel.onDidChangeViewState(async () => {
       if (panel.active) {
-        this.handleViewBecameActive(this);
-        await this.updateMethodsUsagePanelState(
-          this.methods,
-          this.databaseItem,
-          this.hideModeledMethods,
-        );
+        this.modelingStore.setActiveDb(this.databaseItem);
+        await this.markModelEditorAsActive();
+      } else {
+        await this.updateModelEditorActiveContext();
       }
     });
 
     panel.onDidDispose(() => {
-      this.handleViewWasDisposed(this);
+      this.modelingStore.removeDb(this.databaseItem);
       // onDidDispose is called after the tab has been closed,
       // so we want to check if there are any others still open.
       void this.app.commands.execute(
@@ -129,17 +118,46 @@ 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) => {
-        const viewType: string | undefined = (tab.input as any)?.viewType;
-        // The viewType has a prefix, such as "mainThreadWebview-", but if the
-        // suffix matches that should be enough to identify the view.
-        return viewType && viewType.endsWith("model-editor");
-      }),
+      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;
+    }
+
+    // The viewType has a prefix, such as "mainThreadWebview-", but if the
+    // suffix matches that should be enough to identify the view.
+    return tab.input.viewType.endsWith("model-editor");
+  }
+
   protected async getPanelConfig(): Promise<WebviewPanelConfig> {
     return {
       viewId: "model-editor",
@@ -234,6 +252,12 @@ export class ModelEditorView extends AbstractWebview<
             cancellable: false,
           },
         );
+
+        this.modelingStore.removeModifiedMethods(
+          this.databaseItem,
+          Object.keys(msg.modeledMethods),
+        );
+
         void telemetryListener?.sendUIInteraction(
           "model-editor-save-modeled-methods",
         );
@@ -270,11 +294,11 @@ export class ModelEditorView extends AbstractWebview<
         break;
       case "switchMode":
         this.mode = msg.mode;
-        this.methods = [];
+        this.modelingStore.setMethods(this.databaseItem, []);
         await Promise.all([
           this.postMessage({
             t: "setMethods",
-            methods: this.methods,
+            methods: [],
           }),
           this.setViewState(),
           withProgress((progress) => this.loadMethods(progress), {
@@ -285,16 +309,18 @@ export class ModelEditorView extends AbstractWebview<
 
         break;
       case "hideModeledMethods":
-        this.hideModeledMethods = msg.hideModeledMethods;
-        await this.updateMethodsUsagePanelState(
-          this.methods,
+        this.modelingStore.setHideModeledMethods(
           this.databaseItem,
-          this.hideModeledMethods,
+          msg.hideModeledMethods,
         );
         void telemetryListener?.sendUIInteraction(
           "model-editor-hide-modeled-methods",
         );
         break;
+      case "setModeledMethod": {
+        this.setModeledMethod(msg.method);
+        break;
+      }
       default:
         assertNever(msg);
     }
@@ -328,8 +354,7 @@ export class ModelEditorView extends AbstractWebview<
   }
 
   protected async handleJumpToUsage(method: Method, usage: Usage) {
-    await this.showMethod(method, usage);
-    await showResolvableLocation(usage.url, this.databaseItem, this.app.logger);
+    this.modelingStore.setSelectedMethod(this.databaseItem, method, usage);
   }
 
   protected async loadExistingModeledMethods(): Promise<void> {
@@ -339,10 +364,7 @@ export class ModelEditorView extends AbstractWebview<
         this.cliServer,
         this.app.logger,
       );
-      await this.postMessage({
-        t: "loadModeledMethods",
-        modeledMethods,
-      });
+      this.modelingStore.setModeledMethods(this.databaseItem, modeledMethods);
     } catch (e: unknown) {
       void showAndLogErrorMessage(
         this.app.logger,
@@ -370,19 +392,8 @@ export class ModelEditorView extends AbstractWebview<
       if (!queryResult) {
         return;
       }
-      this.methods = queryResult;
 
-      await this.postMessage({
-        t: "setMethods",
-        methods: this.methods,
-      });
-      if (this.isMostRecentlyActiveView(this)) {
-        await this.updateMethodsUsagePanelState(
-          this.methods,
-          this.databaseItem,
-          this.hideModeledMethods,
-        );
-      }
+      this.modelingStore.setMethods(this.databaseItem, queryResult);
     } catch (err) {
       void showAndLogExceptionWithTelemetry(
         this.app.logger,
@@ -431,10 +442,7 @@ export class ModelEditorView extends AbstractWebview<
                 modeledMethodsByName[modeledMethod.signature] = modeledMethod;
               }
 
-              await this.postMessage({
-                t: "addModeledMethods",
-                modeledMethods: modeledMethodsByName,
-              });
+              this.addModeledMethods(modeledMethodsByName);
             },
             progress,
             token: tokenSource.token,
@@ -488,6 +496,7 @@ export class ModelEditorView extends AbstractWebview<
 
       const view = new ModelEditorView(
         this.app,
+        this.modelingStore,
         this.databaseManager,
         this.cliServer,
         this.queryRunner,
@@ -496,11 +505,6 @@ export class ModelEditorView extends AbstractWebview<
         addedDatabase,
         modelFile,
         Mode.Framework,
-        this.updateMethodsUsagePanelState,
-        this.showMethod,
-        this.handleViewBecameActive,
-        this.handleViewWasDisposed,
-        this.isMostRecentlyActiveView,
       );
       await view.openView();
     });
@@ -578,4 +582,58 @@ export class ModelEditorView extends AbstractWebview<
 
     return addedDatabase;
   }
+
+  private registerToModelingStoreEvents() {
+    this.push(
+      this.modelingStore.onMethodsChanged(async (event) => {
+        if (event.dbUri === this.databaseItem.databaseUri.toString()) {
+          await this.postMessage({
+            t: "setMethods",
+            methods: event.methods,
+          });
+        }
+      }),
+    );
+
+    this.push(
+      this.modelingStore.onModeledMethodsChanged(async (event) => {
+        if (event.dbUri === this.databaseItem.databaseUri.toString()) {
+          await this.postMessage({
+            t: "setModeledMethods",
+            methods: event.modeledMethods,
+          });
+        }
+      }),
+    );
+
+    this.push(
+      this.modelingStore.onModifiedMethodsChanged(async (event) => {
+        if (event.dbUri === this.databaseItem.databaseUri.toString()) {
+          await this.postMessage({
+            t: "setModifiedMethods",
+            methodSignatures: [...event.modifiedMethods],
+          });
+        }
+      }),
+    );
+  }
+
+  private addModeledMethods(modeledMethods: Record<string, ModeledMethod>) {
+    this.modelingStore.addModeledMethods(this.databaseItem, modeledMethods);
+
+    this.modelingStore.addModifiedMethods(
+      this.databaseItem,
+      new Set(Object.keys(modeledMethods)),
+    );
+  }
+
+  private setModeledMethod(method: ModeledMethod) {
+    const state = this.modelingStore.getStateForActiveDb();
+    if (!state) {
+      throw new Error("Attempting to set modeled method without active db");
+    }
+
+    this.modelingStore.updateModeledMethod(state.databaseItem, method);
+    this.modelingStore.addModifiedMethod(state.databaseItem, method.signature);
+  }
 }

extensions/ql-vscode/src/model-editor/model-extension-file.schema.json

@@ -0,0 +1,45 @@
+{
+  "$schema": "http://json-schema.org/draft-07/schema#",
+  "$ref": "#/definitions/ModelExtensionFile",
+  "definitions": {
+    "ModelExtensionFile": {
+      "type": "object",
+      "properties": {
+        "extensions": {
+          "type": "array",
+          "items": {
+            "type": "object",
+            "properties": {
+              "addsTo": {
+                "type": "object",
+                "properties": {
+                  "pack": {
+                    "type": "string"
+                  },
+                  "extensible": {
+                    "type": "string"
+                  }
+                },
+                "required": ["pack", "extensible"]
+              },
+              "data": {
+                "type": "array",
+                "items": {
+                  "type": "array",
+                  "items": {
+                    "$ref": "#/definitions/DataTuple"
+                  }
+                }
+              }
+            },
+            "required": ["addsTo", "data"]
+          }
+        }
+      },
+      "required": ["extensions"]
+    },
+    "DataTuple": {
+      "type": ["boolean", "number", "string"]
+    }
+  }
+}

extensions/ql-vscode/src/model-editor/model-extension-file.ts

@@ -0,0 +1,17 @@
+type ExtensibleReference = {
+  pack: string;
+  extensible: string;
+};
+
+export type DataTuple = boolean | number | string;
+
+type DataRow = DataTuple[];
+
+type ModelExtension = {
+  addsTo: ExtensibleReference;
+  data: DataRow[];
+};
+
+export type ModelExtensionFile = {
+  extensions: ModelExtension[];
+};

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

@@ -23,6 +23,8 @@ export interface ModeledMethod extends MethodSignature {
   type: ModeledMethodType;
   input: string;
   output: string;
-  kind: string;
+  kind: ModeledMethodKind;
   provenance: Provenance;
 }
+
+export type ModeledMethodKind = string;

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

@@ -0,0 +1,333 @@
+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 { INITIAL_HIDE_MODELED_METHODS_VALUE } from "./shared/hide-modeled-methods";
+
+interface DbModelingState {
+  databaseItem: DatabaseItem;
+  methods: Method[];
+  hideModeledMethods: boolean;
+  modeledMethods: Record<string, ModeledMethod>;
+  modifiedMethodSignatures: Set<string>;
+  selectedMethod: Method | undefined;
+  selectedUsage: Usage | undefined;
+}
+
+interface MethodsChangedEvent {
+  methods: Method[];
+  dbUri: string;
+  isActiveDb: boolean;
+}
+
+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;
+}
+
+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 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) {
+    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;
+  }
+
+  public initializeStateForDb(databaseItem: DatabaseItem) {
+    const dbUri = databaseItem.databaseUri.toString();
+    this.state.set(dbUri, {
+      databaseItem,
+      methods: [],
+      hideModeledMethods: INITIAL_HIDE_MODELED_METHODS_VALUE,
+      modeledMethods: {},
+      modifiedMethodSignatures: new Set(),
+      selectedMethod: undefined,
+      selectedUsage: undefined,
+    });
+  }
+
+  public setActiveDb(databaseItem: DatabaseItem) {
+    this.activeDb = databaseItem.databaseUri.toString();
+    this.onActiveDbChangedEventEmitter.fire();
+  }
+
+  public removeDb(databaseItem: DatabaseItem) {
+    const dbUri = databaseItem.databaseUri.toString();
+
+    if (!this.state.has(dbUri)) {
+      throw Error("Cannot remove a database that has not been initialized");
+    }
+
+    if (this.activeDb === dbUri) {
+      this.activeDb = undefined;
+      this.onActiveDbChangedEventEmitter.fire();
+    }
+
+    this.state.delete(dbUri);
+    this.onDbClosedEventEmitter.fire(dbUri);
+  }
+
+  public getStateForActiveDb(): DbModelingState | undefined {
+    if (!this.activeDb) {
+      return undefined;
+    }
+
+    return this.state.get(this.activeDb);
+  }
+
+  public setMethods(dbItem: DatabaseItem, methods: Method[]) {
+    const dbState = this.getState(dbItem);
+    const dbUri = dbItem.databaseUri.toString();
+
+    dbState.methods = [...methods];
+
+    this.onMethodsChangedEventEmitter.fire({
+      methods,
+      dbUri,
+      isActiveDb: dbUri === this.activeDb,
+    });
+  }
+
+  public setHideModeledMethods(
+    dbItem: DatabaseItem,
+    hideModeledMethods: boolean,
+  ) {
+    const dbState = this.getState(dbItem);
+    const dbUri = dbItem.databaseUri.toString();
+
+    dbState.hideModeledMethods = hideModeledMethods;
+
+    this.onHideModeledMethodsChangedEventEmitter.fire({
+      hideModeledMethods,
+      isActiveDb: dbUri === this.activeDb,
+    });
+  }
+
+  public addModeledMethods(
+    dbItem: DatabaseItem,
+    methods: Record<string, ModeledMethod>,
+  ) {
+    this.changeModeledMethods(dbItem, (state) => {
+      const newModeledMethods = {
+        ...methods,
+        ...Object.fromEntries(
+          Object.entries(state.modeledMethods).filter(
+            ([_, value]) => value.type !== "none",
+          ),
+        ),
+      };
+      state.modeledMethods = newModeledMethods;
+    });
+  }
+
+  public setModeledMethods(
+    dbItem: DatabaseItem,
+    methods: Record<string, ModeledMethod>,
+  ) {
+    this.changeModeledMethods(dbItem, (state) => {
+      state.modeledMethods = { ...methods };
+    });
+  }
+
+  public updateModeledMethod(dbItem: DatabaseItem, method: ModeledMethod) {
+    this.changeModeledMethods(dbItem, (state) => {
+      const newModeledMethods = { ...state.modeledMethods };
+      newModeledMethods[method.signature] = method;
+      state.modeledMethods = newModeledMethods;
+    });
+  }
+
+  public setModifiedMethods(
+    dbItem: DatabaseItem,
+    methodSignatures: Set<string>,
+  ) {
+    this.changeModifiedMethods(dbItem, (state) => {
+      state.modifiedMethodSignatures = new Set(methodSignatures);
+    });
+  }
+
+  public addModifiedMethods(
+    dbItem: DatabaseItem,
+    methodSignatures: Iterable<string>,
+  ) {
+    this.changeModifiedMethods(dbItem, (state) => {
+      const newModifiedMethods = new Set([
+        ...state.modifiedMethodSignatures,
+        ...methodSignatures,
+      ]);
+      state.modifiedMethodSignatures = newModifiedMethods;
+    });
+  }
+
+  public addModifiedMethod(dbItem: DatabaseItem, methodSignature: string) {
+    this.addModifiedMethods(dbItem, [methodSignature]);
+  }
+
+  public removeModifiedMethods(
+    dbItem: DatabaseItem,
+    methodSignatures: string[],
+  ) {
+    this.changeModifiedMethods(dbItem, (state) => {
+      const newModifiedMethods = Array.from(
+        state.modifiedMethodSignatures,
+      ).filter((s) => !methodSignatures.includes(s));
+
+      state.modifiedMethodSignatures = new Set(newModifiedMethods);
+    });
+  }
+
+  public setSelectedMethod(dbItem: DatabaseItem, method: Method, usage: Usage) {
+    const dbState = this.getState(dbItem);
+
+    dbState.selectedMethod = method;
+    dbState.selectedUsage = usage;
+
+    this.onSelectedMethodChangedEventEmitter.fire({
+      databaseItem: dbItem,
+      method,
+      usage,
+      modeledMethod: dbState.modeledMethods[method.signature],
+      isModified: dbState.modifiedMethodSignatures.has(method.signature),
+    });
+  }
+
+  public getSelectedMethodDetails() {
+    const dbState = this.getStateForActiveDb();
+    if (!dbState) {
+      throw new Error("No active state found in modeling store");
+    }
+
+    const selectedMethod = dbState.selectedMethod;
+    if (!selectedMethod) {
+      return undefined;
+    }
+
+    return {
+      method: selectedMethod,
+      usage: dbState.selectedUsage,
+      modeledMethod: dbState.modeledMethods[selectedMethod.signature],
+      isModified: dbState.modifiedMethodSignatures.has(
+        selectedMethod.signature,
+      ),
+    };
+  }
+
+  private getState(databaseItem: DatabaseItem): DbModelingState {
+    if (!this.state.has(databaseItem.databaseUri.toString())) {
+      throw Error(
+        "Cannot get state for a database that has not been initialized",
+      );
+    }
+
+    return this.state.get(databaseItem.databaseUri.toString())!;
+  }
+
+  private changeModifiedMethods(
+    dbItem: DatabaseItem,
+    updateState: (state: DbModelingState) => void,
+  ) {
+    const state = this.getState(dbItem);
+
+    updateState(state);
+
+    this.onModifiedMethodsChangedEventEmitter.fire({
+      modifiedMethods: state.modifiedMethodSignatures,
+      dbUri: dbItem.databaseUri.toString(),
+      isActiveDb: dbItem.databaseUri.toString() === this.activeDb,
+    });
+  }
+
+  private changeModeledMethods(
+    dbItem: DatabaseItem,
+    updateState: (state: DbModelingState) => void,
+  ) {
+    const state = this.getState(dbItem);
+
+    updateState(state);
+
+    this.onModeledMethodsChangedEventEmitter.fire({
+      modeledMethods: state.modeledMethods,
+      dbUri: dbItem.databaseUri.toString(),
+      isActiveDb: dbItem.databaseUri.toString() === this.activeDb,
+    });
+  }
+}

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

@@ -1,16 +1,15 @@
 import { ModeledMethod, ModeledMethodType, Provenance } from "./modeled-method";
+import { DataTuple } from "./model-extension-file";
 
 export type ExtensiblePredicateDefinition = {
   extensiblePredicate: string;
-  generateMethodDefinition: (method: ModeledMethod) => Tuple[];
-  readModeledMethod: (row: Tuple[]) => ModeledMethod;
+  generateMethodDefinition: (method: ModeledMethod) => DataTuple[];
+  readModeledMethod: (row: DataTuple[]) => ModeledMethod;
 
   supportedKinds?: string[];
 };
 
-type Tuple = boolean | number | string;
-
-function readRowToMethod(row: Tuple[]): string {
+function readRowToMethod(row: DataTuple[]): string {
   return `${row[0]}.${row[1]}#${row[3]}${row[4]}`;
 }
 

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

@@ -2,130 +2,151 @@ import { Query } from "./query";
 
 export const fetchExternalApisQuery: Query = {
   applicationModeQuery: `/**
- * @name Usage of APIs coming from external libraries
- * @description A list of 3rd party APIs used in the codebase.
- * @tags telemetry
- * @kind problem
- * @id cs/telemetry/fetch-external-apis
+ * @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 csharp/utils/modeleditor/application-mode-endpoints
+ * @tags modeleditor endpoints application-mode
  */
 
-private import csharp
-private import AutomodelVsCode
+import csharp
+import ApplicationModeEndpointsQuery
+import ModelEditor
 
-class ExternalApi extends CallableMethod {
-  ExternalApi() {
-    this.isUnboundDeclaration() and
-    this.fromLibrary() and
-    this.(Modifiable).isEffectivelyPublic()
-  }
-}
+private Call aUsage(ExternalEndpoint api) { result.getTarget().getUnboundDeclaration() = api }
 
-private Call aUsage(ExternalApi api) { result.getTarget().getUnboundDeclaration() = api }
-
-from
-  ExternalApi api, string apiName, boolean supported, Call usage, string type, string classification
+from ExternalEndpoint endpoint, boolean supported, Call usage, string type, string classification
 where
-  apiName = api.getApiName() and
-  supported = isSupported(api) and
-  usage = aUsage(api) and
-  type = supportedType(api) and
+  supported = isSupported(endpoint) and
+  usage = aUsage(endpoint) and
+  type = supportedType(endpoint) and
   classification = methodClassification(usage)
-select usage, apiName, supported.toString(), "supported", api.dllName(), api.dllVersion(), type,
-  "type", classification, "classification"
+select usage, endpoint.getNamespace(), endpoint.getTypeName(), endpoint.getName(),
+  endpoint.getParameterTypes(), supported, endpoint.dllName(), endpoint.dllVersion(), type,
+  classification
 `,
   frameworkModeQuery: `/**
- * @name Public methods
- * @description A list of APIs callable by consumers. Excludes test and generated code.
- * @tags telemetry
- * @kind problem
- * @id cs/telemetry/fetch-public-methods
+ * @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 csharp/utils/modeleditor/framework-mode-endpoints
+ * @tags modeleditor endpoints framework-mode
  */
 
-private import csharp
-private import dotnet
-private import semmle.code.csharp.frameworks.Test
-private import AutomodelVsCode
+import csharp
+import FrameworkModeEndpointsQuery
+import ModelEditor
 
-class PublicMethod extends CallableMethod {
-  PublicMethod() { this.fromSource() and not this.getFile() instanceof TestFile }
-}
-
-from PublicMethod publicMethod, string apiName, boolean supported, string type
+from PublicEndpointFromSource endpoint, boolean supported, string type
 where
-  apiName = publicMethod.getApiName() and
-  supported = isSupported(publicMethod) and
-  type = supportedType(publicMethod)
-select publicMethod, apiName, supported.toString(), "supported",
-  publicMethod.getFile().getBaseName(), "library", type, "type", "unknown", "classification"
+  supported = isSupported(endpoint) and
+  type = supportedType(endpoint)
+select endpoint, endpoint.getNamespace(), endpoint.getTypeName(), endpoint.getName(),
+  endpoint.getParameterTypes(), supported, endpoint.getFile().getBaseName(), type
 `,
   dependencies: {
-    "AutomodelVsCode.qll": `/** Provides classes and predicates related to handling APIs for the VS Code extension. */
-
-private import csharp
-private import dotnet
-private import semmle.code.csharp.dispatch.Dispatch
-private import semmle.code.csharp.dataflow.ExternalFlow
-private import semmle.code.csharp.dataflow.FlowSummary
-private import semmle.code.csharp.dataflow.internal.DataFlowImplCommon as DataFlowImplCommon
-private import semmle.code.csharp.dataflow.internal.DataFlowPrivate
+    "ApplicationModeEndpointsQuery.qll": `private import csharp
+private import semmle.code.csharp.dataflow.ExternalFlow as ExternalFlow
 private import semmle.code.csharp.dataflow.internal.DataFlowDispatch as DataFlowDispatch
-private import semmle.code.csharp.dataflow.internal.FlowSummaryImpl as FlowSummaryImpl
+private import semmle.code.csharp.dataflow.internal.DataFlowPrivate
 private import semmle.code.csharp.dataflow.internal.TaintTrackingPrivate
-private import semmle.code.csharp.frameworks.Test
 private import semmle.code.csharp.security.dataflow.flowsources.Remote
-
-pragma[nomagic]
-private predicate isTestNamespace(Namespace ns) {
-  ns.getFullName()
-      .matches([
-          "NUnit.Framework%", "Xunit%", "Microsoft.VisualStudio.TestTools.UnitTesting%", "Moq%"
-        ])
-}
+private import ModelEditor
 
 /**
- * A test library.
+ * A class of effectively public callables in library code.
  */
-class TestLibrary extends RefType {
-  TestLibrary() { isTestNamespace(this.getNamespace()) }
+class ExternalEndpoint extends Endpoint {
+  ExternalEndpoint() { this.fromLibrary() }
+
+  /** Gets a node that is an input to a call to this API. */
+  private ArgumentNode getAnInput() {
+    result
+        .getCall()
+        .(DataFlowDispatch::NonDelegateDataFlowCall)
+        .getATarget(_)
+        .getUnboundDeclaration() = this
+  }
+
+  /** Gets a node that is an output from a call to this API. */
+  private DataFlow::Node getAnOutput() {
+    exists(Call c, DataFlowDispatch::NonDelegateDataFlowCall dc |
+      dc.getDispatchCall().getCall() = c and
+      c.getTarget().getUnboundDeclaration() = this
+    |
+      result = DataFlowDispatch::getAnOutNode(dc, _)
+    )
+  }
+
+  override predicate hasSummary() {
+    Endpoint.super.hasSummary()
+    or
+    defaultAdditionalTaintStep(this.getAnInput(), _)
+  }
+
+  override predicate isSource() {
+    this.getAnOutput() instanceof RemoteFlowSource or ExternalFlow::sourceNode(this.getAnOutput(), _)
+  }
+
+  override predicate isSink() { ExternalFlow::sinkNode(this.getAnInput(), _) }
 }
+`,
+    "FrameworkModeEndpointsQuery.qll": `private import csharp
+private import semmle.code.csharp.frameworks.Test
+private import ModelEditor
+
+/**
+ * A class of effectively public callables from source code.
+ */
+class PublicEndpointFromSource extends Endpoint {
+  PublicEndpointFromSource() { this.fromSource() and not this.getFile() instanceof TestFile }
+
+  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 csharp
+private import semmle.code.csharp.dataflow.FlowSummary
+private import semmle.code.csharp.dataflow.internal.DataFlowPrivate
+private import semmle.code.csharp.dataflow.internal.FlowSummaryImpl as FlowSummaryImpl
+private import semmle.code.csharp.frameworks.Test
 
 /** Holds if the given callable is not worth supporting. */
-private predicate isUninteresting(DotNet::Declaration c) {
+private predicate isUninteresting(Callable c) {
   c.getDeclaringType() instanceof TestLibrary or
   c.(Constructor).isParameterless() or
   c.getDeclaringType() instanceof AnonymousClass
 }
 
 /**
- * An callable method from either the C# Standard Library, a 3rd party library, or from the source.
+ * A callable method or accessor from either the C# Standard Library, a 3rd party library, or from the source.
  */
-class CallableMethod extends DotNet::Declaration {
-  CallableMethod() {
-    this.(Modifiable).isEffectivelyPublic() and
-    not isUninteresting(this)
+class Endpoint extends Callable {
+  Endpoint() {
+    [this.(Modifiable), this.(Accessor).getDeclaration()].isEffectivelyPublic() and
+    not isUninteresting(this) and
+    this.isUnboundDeclaration()
   }
 
   /**
-   * Gets the unbound type, name and parameter types of this API.
-   */
-  bindingset[this]
-  private string getSignature() {
-    result =
-      nestedName(this.getDeclaringType().getUnboundDeclaration()) + "#" + this.getName() + "(" +
-        parameterQualifiedTypeNamesToString(this) + ")"
-  }
-
-  /**
-   * Gets the namespace of this API.
+   * Gets the namespace of this endpoint.
    */
   bindingset[this]
   string getNamespace() { this.getDeclaringType().hasQualifiedName(result, _) }
 
   /**
-   * Gets the namespace and signature of this API.
+   * Gets the unbound type name of this endpoint.
    */
   bindingset[this]
-  string getApiName() { result = this.getNamespace() + "." + this.getSignature() }
+  string getTypeName() { result = nestedName(this.getDeclaringType().getUnboundDeclaration()) }
+
+  /**
+   * Gets the parameter types of this endpoint.
+   */
+  bindingset[this]
+  string getParameterTypes() { result = "(" + parameterQualifiedTypeNamesToString(this) + ")" }
 
   private string getDllName() { result = this.getLocation().(Assembly).getName() }
 
@@ -143,44 +164,17 @@ class CallableMethod extends DotNet::Declaration {
     not exists(this.getDllVersion()) and result = ""
   }
 
-  /** Gets a node that is an input to a call to this API. */
-  private ArgumentNode getAnInput() {
-    result
-        .getCall()
-        .(DataFlowDispatch::NonDelegateDataFlowCall)
-        .getATarget(_)
-        .getUnboundDeclaration() = this
-  }
-
-  /** Gets a node that is an output from a call to this API. */
-  private DataFlow::Node getAnOutput() {
-    exists(
-      Call c, DataFlowDispatch::NonDelegateDataFlowCall dc, DataFlowImplCommon::ReturnKindExt ret
-    |
-      dc.getDispatchCall().getCall() = c and
-      c.getTarget().getUnboundDeclaration() = this
-    |
-      result = ret.getAnOutNode(dc)
-    )
-  }
-
   /** Holds if this API has a supported summary. */
   pragma[nomagic]
-  predicate hasSummary() {
-    this instanceof SummarizedCallable
-    or
-    defaultAdditionalTaintStep(this.getAnInput(), _)
-  }
+  predicate hasSummary() { this instanceof SummarizedCallable }
 
   /** Holds if this API is a known source. */
   pragma[nomagic]
-  predicate isSource() {
-    this.getAnOutput() instanceof RemoteFlowSource or sourceNode(this.getAnOutput(), _)
-  }
+  abstract predicate isSource();
 
   /** Holds if this API is a known sink. */
   pragma[nomagic]
-  predicate isSink() { sinkNode(this.getAnInput(), _) }
+  abstract predicate isSink();
 
   /** Holds if this API is a known neutral. */
   pragma[nomagic]
@@ -195,23 +189,20 @@ class CallableMethod extends DotNet::Declaration {
   }
 }
 
-boolean isSupported(CallableMethod callableMethod) {
-  callableMethod.isSupported() and result = true
-  or
-  not callableMethod.isSupported() and
-  result = false
+boolean isSupported(Endpoint endpoint) {
+  if endpoint.isSupported() then result = true else result = false
 }
 
-string supportedType(CallableMethod method) {
-  method.isSink() and result = "sink"
+string supportedType(Endpoint endpoint) {
+  endpoint.isSink() and result = "sink"
   or
-  method.isSource() and result = "source"
+  endpoint.isSource() and result = "source"
   or
-  method.hasSummary() and result = "summary"
+  endpoint.hasSummary() and result = "summary"
   or
-  method.isNeutral() and result = "neutral"
+  endpoint.isNeutral() and result = "neutral"
   or
-  not method.isSupported() and result = ""
+  not endpoint.isSupported() and result = ""
 }
 
 string methodClassification(Call method) {
@@ -222,18 +213,51 @@ string methodClassification(Call method) {
 }
 
 /**
- * Gets the nested name of the declaration.
+ * Gets the nested name of the type \`t\`.
  *
- * If the declaration is not a nested type, the result is the same as \`getName()\`.
+ * If the type is not a nested type, the result is the same as \`getName()\`.
  * Otherwise the name of the nested type is prefixed with a \`+\` and appended to
  * the name of the enclosing type, which might be a nested type as well.
  */
-private string nestedName(Declaration declaration) {
-  not exists(declaration.getDeclaringType().getUnboundDeclaration()) and
-  result = declaration.getName()
+private string nestedName(Type t) {
+  not exists(t.getDeclaringType().getUnboundDeclaration()) and
+  result = t.getName()
   or
-  nestedName(declaration.getDeclaringType().getUnboundDeclaration()) + "+" + declaration.getName() =
-    result
+  nestedName(t.getDeclaringType().getUnboundDeclaration()) + "+" + t.getName() = result
+}
+
+// Temporary copy of csharp/ql/src/Telemetry/TestLibrary.qll
+
+pragma[nomagic]
+private predicate isTestNamespace(Namespace ns) {
+  ns.getFullName()
+      .matches([
+          "NUnit.Framework%", "Xunit%", "Microsoft.VisualStudio.TestTools.UnitTesting%", "Moq%"
+        ])
+}
+
+/**
+ * A test library.
+ */
+class TestLibrary extends RefType {
+  TestLibrary() { isTestNamespace(this.getNamespace()) }
+}
+
+// Temporary copy of csharp/ql/lib/semmle/code/csharp/dataflow/ExternalFlow.qll
+private import semmle.code.csharp.dataflow.internal.FlowSummaryImplSpecific
+
+/**
+ * A callable where there exists a MaD sink model that applies to it.
+ */
+class SinkCallable extends Callable {
+  SinkCallable() { sinkElement(this, _, _, _) }
+}
+
+/**
+ * A callable where there exists a MaD source model that applies to it.
+ */
+class SourceCallable extends Callable {
+  SourceCallable() { sourceElement(this, _, _, _) }
 }
 `,
   },

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

@@ -2,66 +2,113 @@ import { Query } from "./query";
 
 export const fetchExternalApisQuery: Query = {
   applicationModeQuery: `/**
- * @name Usage of APIs coming from external libraries
- * @description A list of 3rd party APIs used in the codebase. Excludes test and generated code.
- * @tags telemetry
- * @kind problem
- * @id java/telemetry/fetch-external-apis
+ * @name Fetch endpoints for use in the model editor (application mode)
+ * @description A list of 3rd party endpoints (methods) used in the codebase. Excludes test and generated code.
+ * @kind table
+ * @id java/utils/modeleditor/application-mode-endpoints
+ * @tags modeleditor endpoints application-mode
  */
 
-import java
-import AutomodelVsCode
-
-class ExternalApi extends CallableMethod {
-  ExternalApi() { not this.fromSource() }
-}
-
-private Call aUsage(ExternalApi api) { result.getCallee().getSourceDeclaration() = api }
-
-from
-  ExternalApi externalApi, string apiName, boolean supported, Call usage, string type,
-  string classification
-where
-  apiName = externalApi.getApiName() and
-  supported = isSupported(externalApi) and
-  usage = aUsage(externalApi) and
-  type = supportedType(externalApi) and
-  classification = methodClassification(usage)
-select usage, apiName, supported.toString(), "supported", externalApi.jarContainer(),
-  externalApi.jarVersion(), type, "type", classification, "classification"
-`,
-  frameworkModeQuery: `/**
- * @name Public methods
- * @description A list of APIs callable by consumers. Excludes test and generated code.
- * @tags telemetry
- * @kind problem
- * @id java/telemetry/fetch-public-methods
- */
-
-import java
-import AutomodelVsCode
-
-class PublicMethodFromSource extends CallableMethod, ModelApi { }
-
-from PublicMethodFromSource publicMethod, string apiName, boolean supported, string type
-where
-  apiName = publicMethod.getApiName() and
-  supported = isSupported(publicMethod) and
-  type = supportedType(publicMethod)
-select publicMethod, apiName, supported.toString(), "supported",
-  publicMethod.getCompilationUnit().getParentContainer().getBaseName(), "library", type, "type",
-  "unknown", "classification"
-`,
-  dependencies: {
-    "AutomodelVsCode.qll": `/** Provides classes and predicates related to handling APIs for the VS Code extension. */
-
 private import java
-private import semmle.code.java.dataflow.DataFlow
+private import ApplicationModeEndpointsQuery
+private import ModelEditor
+
+private Call aUsage(ExternalEndpoint endpoint) {
+  result.getCallee().getSourceDeclaration() = endpoint
+}
+
+from ExternalEndpoint endpoint, boolean supported, Call usage, string type, string classification
+where
+  supported = isSupported(endpoint) and
+  usage = aUsage(endpoint) and
+  type = supportedType(endpoint) and
+  classification = usageClassification(usage)
+select usage, endpoint.getPackageName(), endpoint.getTypeName(), endpoint.getName(),
+  endpoint.getParameterTypes(), supported, endpoint.jarContainer(), endpoint.jarVersion(), type,
+  classification
+`,
+  frameworkModeQuery: `/**
+ * @name Fetch endpoints for use in the model editor (framework mode)
+ * @description A list of endpoints accessible (methods) for consumers of the library. Excludes test and generated code.
+ * @kind table
+ * @id java/utils/modeleditor/framework-mode-endpoints
+ * @tags modeleditor endpoints framework-mode
+ */
+
+private import java
+private import FrameworkModeEndpointsQuery
+private import ModelEditor
+
+from PublicEndpointFromSource endpoint, boolean supported, string type
+where
+  supported = isSupported(endpoint) and
+  type = supportedType(endpoint)
+select endpoint, endpoint.getPackageName(), endpoint.getTypeName(), endpoint.getName(),
+  endpoint.getParameterTypes(), supported,
+  endpoint.getCompilationUnit().getParentContainer().getBaseName(), type
+`,
+  dependencies: {
+    "ApplicationModeEndpointsQuery.qll": `private import java
 private import semmle.code.java.dataflow.ExternalFlow
 private import semmle.code.java.dataflow.FlowSources
-private import semmle.code.java.dataflow.FlowSummary
 private import semmle.code.java.dataflow.internal.DataFlowPrivate
-private import semmle.code.java.dataflow.internal.FlowSummaryImpl as FlowSummaryImpl
+private import ModelEditor
+
+/**
+ * A class of effectively public callables in library code.
+ */
+class ExternalEndpoint extends Endpoint {
+  ExternalEndpoint() { not this.fromSource() }
+
+  /** Gets a node that is an input to a call to this API. */
+  private DataFlow::Node getAnInput() {
+    exists(Call call | call.getCallee().getSourceDeclaration() = this |
+      result.asExpr().(Argument).getCall() = call or
+      result.(ArgumentNode).getCall().asCall() = call
+    )
+  }
+
+  /** Gets a node that is an output from a call to this API. */
+  private DataFlow::Node getAnOutput() {
+    exists(Call call | call.getCallee().getSourceDeclaration() = this |
+      result.asExpr() = call or
+      result.(DataFlow::PostUpdateNode).getPreUpdateNode().(ArgumentNode).getCall().asCall() = call
+    )
+  }
+
+  override predicate hasSummary() {
+    Endpoint.super.hasSummary()
+    or
+    TaintTracking::localAdditionalTaintStep(this.getAnInput(), _)
+  }
+
+  override predicate isSource() {
+    this.getAnOutput() instanceof RemoteFlowSource or sourceNode(this.getAnOutput(), _)
+  }
+
+  override predicate isSink() { sinkNode(this.getAnInput(), _) }
+}
+`,
+    "FrameworkModeEndpointsQuery.qll": `private import java
+private import semmle.code.java.dataflow.internal.DataFlowPrivate
+private import semmle.code.java.dataflow.internal.FlowSummaryImplSpecific
+private import semmle.code.java.dataflow.internal.ModelExclusions
+private import ModelEditor
+
+/**
+ * A class of effectively public callables from source code.
+ */
+class PublicEndpointFromSource extends Endpoint, ModelApi {
+  override predicate isSource() { sourceElement(this, _, _, _) }
+
+  override predicate isSink() { sinkElement(this, _, _, _) }
+}
+`,
+    "ModelEditor.qll": `/** Provides classes and predicates related to handling APIs for the VS Code extension. */
+
+private import java
+private import semmle.code.java.dataflow.ExternalFlow
+private import semmle.code.java.dataflow.FlowSummary
 private import semmle.code.java.dataflow.TaintTracking
 private import semmle.code.java.dataflow.internal.ModelExclusions
 
@@ -75,17 +122,23 @@ private predicate isUninteresting(Callable c) {
 /**
  * A callable method from either the Standard Library, a 3rd party library or from the source.
  */
-class CallableMethod extends Callable {
-  CallableMethod() { not isUninteresting(this) }
+class Endpoint extends Callable {
+  Endpoint() { not isUninteresting(this) }
 
   /**
-   * Gets information about the external API in the form expected by the MaD modeling framework.
+   * Gets the package name of this endpoint.
    */
-  string getApiName() {
-    result =
-      this.getDeclaringType().getPackage() + "." + this.getDeclaringType().nestedName() + "#" +
-        this.getName() + paramsString(this)
-  }
+  string getPackageName() { result = this.getDeclaringType().getPackage().getName() }
+
+  /**
+   * Gets the type name of this endpoint.
+   */
+  string getTypeName() { result = this.getDeclaringType().nestedName() }
+
+  /**
+   * Gets the parameter types of this endpoint.
+   */
+  string getParameterTypes() { result = paramsString(this) }
 
   private string getJarName() {
     result = this.getCompilationUnit().getParentContainer*().(JarFile).getBaseName()
@@ -113,43 +166,23 @@ class CallableMethod extends Callable {
     not exists(this.getJarVersion()) and result = ""
   }
 
-  /** Gets a node that is an input to a call to this API. */
-  private DataFlow::Node getAnInput() {
-    exists(Call call | call.getCallee().getSourceDeclaration() = this |
-      result.asExpr().(Argument).getCall() = call or
-      result.(ArgumentNode).getCall().asCall() = call
-    )
-  }
-
-  /** Gets a node that is an output from a call to this API. */
-  private DataFlow::Node getAnOutput() {
-    exists(Call call | call.getCallee().getSourceDeclaration() = this |
-      result.asExpr() = call or
-      result.(DataFlow::PostUpdateNode).getPreUpdateNode().(ArgumentNode).getCall().asCall() = call
-    )
-  }
-
   /** Holds if this API has a supported summary. */
   pragma[nomagic]
-  predicate hasSummary() {
-    this = any(SummarizedCallable sc).asCallable() or
-    TaintTracking::localAdditionalTaintStep(this.getAnInput(), _)
-  }
+  predicate hasSummary() { this = any(SummarizedCallable sc).asCallable() }
 
+  /** Holds if this API is a known source. */
   pragma[nomagic]
-  predicate isSource() {
-    this.getAnOutput() instanceof RemoteFlowSource or sourceNode(this.getAnOutput(), _)
-  }
+  abstract predicate isSource();
 
   /** Holds if this API is a known sink. */
   pragma[nomagic]
-  predicate isSink() { sinkNode(this.getAnInput(), _) }
+  abstract predicate isSink();
 
   /** Holds if this API is a known neutral. */
   pragma[nomagic]
   predicate isNeutral() {
-    exists(string namespace, string type, string name, string signature, string kind, string provenance |
-      neutralModel(namespace, type, name, signature, kind, provenance) and
+    exists(string namespace, string type, string name, string signature |
+      neutralModel(namespace, type, name, signature, _, _) and
       this = interpretElement(namespace, type, false, name, signature, "")
     )
   }
@@ -163,108 +196,38 @@ class CallableMethod extends Callable {
   }
 }
 
-boolean isSupported(CallableMethod method) {
-  method.isSupported() and result = true
+boolean isSupported(Endpoint endpoint) {
+  endpoint.isSupported() and result = true
   or
-  not method.isSupported() and result = false
+  not endpoint.isSupported() and result = false
 }
 
-string supportedType(CallableMethod method) {
-  method.isSink() and result = "sink"
+string supportedType(Endpoint endpoint) {
+  endpoint.isSink() and result = "sink"
   or
-  method.isSource() and result = "source"
+  endpoint.isSource() and result = "source"
   or
-  method.hasSummary() and result = "summary"
+  endpoint.hasSummary() and result = "summary"
   or
-  method.isNeutral() and result = "neutral"
+  endpoint.isNeutral() and result = "neutral"
   or
-  not method.isSupported() and result = ""
+  not endpoint.isSupported() and result = ""
 }
 
-string methodClassification(Call method) {
-  isInTestFile(method.getLocation().getFile()) and result = "test"
+string usageClassification(Call usage) {
+  isInTestFile(usage.getLocation().getFile()) and result = "test"
   or
-  method.getFile() instanceof GeneratedFile and result = "generated"
+  usage.getFile() instanceof GeneratedFile and result = "generated"
   or
-  not isInTestFile(method.getLocation().getFile()) and
-  not method.getFile() instanceof GeneratedFile and
+  not isInTestFile(usage.getLocation().getFile()) and
+  not usage.getFile() instanceof GeneratedFile and
   result = "source"
 }
 
-// The below is a copy of https://github.com/github/codeql/blob/249f9f863db1e94e3c46ca85b49fb0ec32f8ca92/java/ql/lib/semmle/code/java/dataflow/internal/ModelExclusions.qll
-// to avoid the use of internal modules.
-/** Holds if the given package \`p\` is a test package. */
-pragma[nomagic]
-private predicate isTestPackage(Package p) {
-  p.getName()
-      .matches([
-          "org.junit%", "junit.%", "org.mockito%", "org.assertj%",
-          "com.github.tomakehurst.wiremock%", "org.hamcrest%", "org.springframework.test.%",
-          "org.springframework.mock.%", "org.springframework.boot.test.%", "reactor.test%",
-          "org.xmlunit%", "org.testcontainers.%", "org.opentest4j%", "org.mockserver%",
-          "org.powermock%", "org.skyscreamer.jsonassert%", "org.rnorth.visibleassertions",
-          "org.openqa.selenium%", "com.gargoylesoftware.htmlunit%", "org.jboss.arquillian.testng%",
-          "org.testng%"
-        ])
-}
-
-/**
- * A test library.
- */
-class TestLibrary extends RefType {
-  TestLibrary() { isTestPackage(this.getPackage()) }
-}
-
-/** Holds if the given file is a test file. */
-private predicate isInTestFile(File file) {
+// Temporarily copied from java/ql/lib/semmle/code/java/dataflow/internal/ModelExclusions.qll
+predicate isInTestFile(File file) {
   file.getAbsolutePath().matches(["%/test/%", "%/guava-tests/%", "%/guava-testlib/%"]) and
-  not file.getAbsolutePath().matches("%/ql/test/%") // allows our test cases to work
-}
-
-/** Holds if the given compilation unit's package is a JDK internal. */
-private predicate isJdkInternal(CompilationUnit cu) {
-  cu.getPackage().getName().matches("org.graalvm%") or
-  cu.getPackage().getName().matches("com.sun%") or
-  cu.getPackage().getName().matches("sun%") or
-  cu.getPackage().getName().matches("jdk%") or
-  cu.getPackage().getName().matches("java2d%") or
-  cu.getPackage().getName().matches("build.tools%") or
-  cu.getPackage().getName().matches("propertiesparser%") or
-  cu.getPackage().getName().matches("org.jcp%") or
-  cu.getPackage().getName().matches("org.w3c%") or
-  cu.getPackage().getName().matches("org.ietf.jgss%") or
-  cu.getPackage().getName().matches("org.xml.sax%") or
-  cu.getPackage().getName().matches("com.oracle%") or
-  cu.getPackage().getName().matches("org.omg%") or
-  cu.getPackage().getName().matches("org.relaxng%") or
-  cu.getPackage().getName() = "compileproperties" or
-  cu.getPackage().getName() = "transparentruler" or
-  cu.getPackage().getName() = "genstubs" or
-  cu.getPackage().getName() = "netscape.javascript" or
-  cu.getPackage().getName() = ""
-}
-
-/** Holds if the given callable is not worth modeling. */
-predicate isUninterestingForModels(Callable c) {
-  isInTestFile(c.getCompilationUnit().getFile()) or
-  isJdkInternal(c.getCompilationUnit()) or
-  c instanceof MainMethod or
-  c instanceof StaticInitializer or
-  exists(FunctionalExpr funcExpr | c = funcExpr.asMethod()) or
-  c.getDeclaringType() instanceof TestLibrary or
-  c.(Constructor).isParameterless()
-}
-
-/**
- * A class that represents all callables for which we might be
- * interested in having a MaD model.
- */
-class ModelApi extends SrcCallable {
-  ModelApi() {
-    this.fromSource() and
-    this.isEffectivelyPublic() and
-    not isUninterestingForModels(this)
-  }
+  not file.getAbsolutePath().matches(["%/ql/test/%", "%/ql/automodel/test/%"]) // allows our test cases to work
 }
 `,
   },

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

@@ -1,18 +1,21 @@
+import { Call, CallClassification } from "../method";
+import { ModeledMethodType } from "../modeled-method";
+
 export type Query = {
   /**
    * The application query.
    *
    * It should select all usages of external APIs, and return the following result pattern:
    * - usage: the usage of the external API. This is an entity.
-   * - apiName: the name of the external API. This is a string.
-   * - supported: whether the external API is modeled. This should be a string representation of a boolean to satify the result pattern for a problem query.
-   * - "supported": a string literal. This is required to make the query a valid problem query.
+   * - packageName: the package name of the external API. This is a string.
+   * - typeName: the type name of the external API. This is a string.
+   * - methodName: the method name of the external API. This is a string.
+   * - methodParameters: the parameters of the external API. This is a string.
+   * - supported: whether the external API is modeled. This is a boolean.
    * - libraryName: the name of the library that contains the external API. This is a string and usually the basename of a file.
    * - libraryVersion: the version of the library that contains the external API. This is a string and can be empty if the version cannot be determined.
    * - type: the modeled kind of the method, either "sink", "source", "summary", or "neutral"
-   * - "type": a string literal. This is required to make the query a valid problem query.
    * - classification: the classification of the use of the method, either "source", "test", "generated", or "unknown"
-   * - "classification: a string literal. This is required to make the query a valid problem query.
    */
   applicationModeQuery: string;
   /**
@@ -21,18 +24,40 @@ export type Query = {
    * It should select all methods that are callable by applications, which is usually all public methods (and constructors).
    * The result pattern should be as follows:
    * - method: the method that is callable by applications. This is an entity.
-   * - apiName: the name of the external API. This is a string.
+   * - packageName: the package name of the method. This is a string.
+   * - typeName: the type name of the method. This is a string.
+   * - methodName: the method name of the method. This is a string.
+   * - methodParameters: the parameters of the method. This is a string.
    * - supported: whether this method is modeled. This should be a string representation of a boolean to satify the result pattern for a problem query.
-   * - "supported": a string literal. This is required to make the query a valid problem query.
-   * - libraryName: an arbitrary string. This is required to make it match the structure of the application query.
-   * - libraryVersion: an arbitrary string. This is required to make it match the structure of the application query.
+   * - libraryName: the name of the file or library that contains the method. This is a string and usually the basename of a file.
    * - type: the modeled kind of the method, either "sink", "source", "summary", or "neutral"
-   * - "type": a string literal. This is required to make the query a valid problem query.
-   * - "unknown": a string literal. This is required to make it match the structure of the application query.
-   * - "classification: a string literal. This is required to make the query a valid problem query.
    */
   frameworkModeQuery: string;
   dependencies?: {
     [filename: string]: string;
   };
 };
+
+export type ApplicationModeTuple = [
+  Call,
+  string,
+  string,
+  string,
+  string,
+  boolean,
+  string,
+  string,
+  ModeledMethodType,
+  CallClassification,
+];
+
+export type FrameworkModeTuple = [
+  Call,
+  string,
+  string,
+  string,
+  string,
+  boolean,
+  string,
+  ModeledMethodType,
+];

extensions/ql-vscode/src/model-editor/shared/modeling-status.ts

@@ -0,0 +1,17 @@
+import { ModeledMethod } from "../modeled-method";
+
+export type ModelingStatus = "unmodeled" | "unsaved" | "saved";
+
+export function getModelingStatus(
+  modeledMethod: ModeledMethod | undefined,
+  methodIsUnsaved: boolean,
+): ModelingStatus {
+  if (modeledMethod) {
+    if (methodIsUnsaved) {
+      return "unsaved";
+    } else if (modeledMethod.type !== "none") {
+      return "saved";
+    }
+  }
+  return "unmodeled";
+}

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

@@ -7,12 +7,13 @@ import {
   extensiblePredicateDefinitions,
 } from "./predicates";
 
-import * as dataSchemaJson from "./data-schema.json";
+import * as modelExtensionFileSchema from "./model-extension-file.schema.json";
 import { Mode } from "./shared/mode";
 import { assertNever } from "../common/helpers-pure";
+import { ModelExtensionFile } from "./model-extension-file";
 
-const ajv = new Ajv({ allErrors: true });
-const dataSchemaValidate = ajv.compile(dataSchemaJson);
+const ajv = new Ajv({ allErrors: true, allowUnionTypes: true });
+const modelExtensionFileSchemaValidate = ajv.compile(modelExtensionFileSchema);
 
 function createDataProperty(
   methods: ModeledMethod[],
@@ -211,24 +212,29 @@ export function createFilenameForPackage(
   return `${prefix}${packageName}${suffix}.yml`;
 }
 
-export function loadDataExtensionYaml(
-  data: any,
-): Record<string, ModeledMethod> | undefined {
-  dataSchemaValidate(data);
+function validateModelExtensionFile(data: unknown): data is ModelExtensionFile {
+  modelExtensionFileSchemaValidate(data);
 
-  if (dataSchemaValidate.errors) {
+  if (modelExtensionFileSchemaValidate.errors) {
     throw new Error(
-      `Invalid data extension YAML: ${dataSchemaValidate.errors
+      `Invalid data extension YAML: ${modelExtensionFileSchemaValidate.errors
         .map((error) => `${error.instancePath} ${error.message}`)
         .join(", ")}`,
     );
   }
 
-  const extensions = data.extensions;
-  if (!Array.isArray(extensions)) {
+  return true;
+}
+
+export function loadDataExtensionYaml(
+  data: unknown,
+): Record<string, ModeledMethod> | undefined {
+  if (!validateModelExtensionFile(data)) {
     return undefined;
   }
 
+  const extensions = data.extensions;
+
   const modeledMethods: Record<string, ModeledMethod> = {};
 
   for (const extension of extensions) {

extensions/ql-vscode/src/packaging/qlpack-file.ts

@@ -0,0 +1,16 @@
+import { SuiteInstruction } from "./suite-instruction";
+
+/**
+ * The qlpack pack file, either in qlpack.yml or in codeql-pack.yml.
+ */
+export interface QlPackFile {
+  name: string;
+  version: string;
+  dependencies?: Record<string, string>;
+  extensionTargets?: Record<string, string>;
+  dbscheme?: string;
+  library?: boolean;
+  defaultSuite?: SuiteInstruction[];
+  defaultSuiteFile?: string;
+  dataExtensions?: string[] | string;
+}

extensions/ql-vscode/src/packaging/qlpack-lock-file.ts

@@ -0,0 +1,8 @@
+/**
+ * The qlpack lock file, either in qlpack.lock.yml or in codeql-pack.lock.yml.
+ */
+export interface QlPackLockFile {
+  lockVersion: string;
+  dependencies?: Record<string, string>;
+  compiled?: boolean;
+}

extensions/ql-vscode/src/packaging/suite-instruction.ts

@@ -0,0 +1,12 @@
+/**
+ * A single entry in a .qls file.
+ */
+export interface SuiteInstruction {
+  qlpack?: string;
+  query?: string;
+  queries?: string;
+  include?: Record<string, string[]>;
+  exclude?: Record<string, string[]>;
+  description?: string;
+  from?: string;
+}

extensions/ql-vscode/src/stories/method-modeling/MethodModeling.stories.tsx

@@ -3,6 +3,7 @@ import * as React from "react";
 import { Meta, StoryFn } from "@storybook/react";
 
 import { MethodModeling as MethodModelingComponent } from "../../view/method-modeling/MethodModeling";
+import { createMethod } from "../../../test/factories/model-editor/method-factories";
 export default {
   title: "Method Modeling/Method Modeling",
   component: MethodModelingComponent,
@@ -12,11 +13,23 @@ const Template: StoryFn<typeof MethodModelingComponent> = (args) => (
   <MethodModelingComponent {...args} />
 );
 
+const method = createMethod();
+
 export const MethodUnmodeled = Template.bind({});
-MethodUnmodeled.args = { modelingStatus: "unmodeled" };
+MethodUnmodeled.args = {
+  method,
+  modelingStatus: "unmodeled",
+};
 
 export const MethodModeled = Template.bind({});
-MethodModeled.args = { modelingStatus: "unsaved" };
+MethodModeled.args = {
+  method,
+
+  modelingStatus: "unsaved",
+};
 
 export const MethodSaved = Template.bind({});
-MethodSaved.args = { modelingStatus: "saved" };
+MethodSaved.args = {
+  method,
+  modelingStatus: "saved",
+};

extensions/ql-vscode/src/stories/method-modeling/MethodModelingInputs.stories.tsx

@@ -0,0 +1,53 @@
+import * as React from "react";
+
+import { Meta, StoryFn } from "@storybook/react";
+
+import { MethodModelingInputs as MethodModelingInputsComponent } from "../../view/method-modeling/MethodModelingInputs";
+import { createMethod } from "../../../test/factories/model-editor/method-factories";
+import { createModeledMethod } from "../../../test/factories/model-editor/modeled-method-factories";
+import { useState } from "react";
+import { ModeledMethod } from "../../model-editor/modeled-method";
+
+export default {
+  title: "Method Modeling/Method Modeling Inputs",
+  component: MethodModelingInputsComponent,
+  argTypes: {
+    modeledMethod: {
+      control: {
+        disable: true,
+      },
+    },
+  },
+} as Meta<typeof MethodModelingInputsComponent>;
+
+const Template: StoryFn<typeof MethodModelingInputsComponent> = (args) => {
+  const [m, setModeledMethod] = useState<ModeledMethod | undefined>(
+    args.modeledMethod,
+  );
+
+  const onChange = (modeledMethod: ModeledMethod) => {
+    setModeledMethod(modeledMethod);
+  };
+
+  return (
+    <MethodModelingInputsComponent
+      {...args}
+      modeledMethod={m}
+      onChange={onChange}
+    />
+  );
+};
+
+const method = createMethod();
+const modeledMethod = createModeledMethod();
+
+export const UnmodeledMethod = Template.bind({});
+UnmodeledMethod.args = {
+  method,
+};
+
+export const FullyModeledMethod = Template.bind({});
+FullyModeledMethod.args = {
+  method,
+  modeledMethod,
+};

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

@@ -3,7 +3,7 @@ import * as React from "react";
 import { Meta, StoryFn } from "@storybook/react";
 
 import { MethodName as MethodNameComponent } from "../../view/model-editor/MethodName";
-import { createMethod } from "../../../test/factories/data-extension/method-factories";
+import { createMethod } from "../../../test/factories/model-editor/method-factories";
 
 export default {
   title: "CodeQL Model Editor/Method Name",

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

@@ -27,7 +27,7 @@ const method: Method = {
   methodName: "open",
   methodParameters: "()",
   supported: false,
-  supportedType: "summary",
+  supportedType: "none",
   usages: [
     {
       label: "open(...)",
@@ -70,30 +70,35 @@ export const Unmodeled = Template.bind({});
 Unmodeled.args = {
   method,
   modeledMethod: undefined,
+  methodCanBeModeled: true,
 };
 
 export const Source = Template.bind({});
 Source.args = {
   method,
   modeledMethod: { ...modeledMethod, type: "source" },
+  methodCanBeModeled: true,
 };
 
 export const Sink = Template.bind({});
 Sink.args = {
   method,
   modeledMethod: { ...modeledMethod, type: "sink" },
+  methodCanBeModeled: true,
 };
 
 export const Summary = Template.bind({});
 Summary.args = {
   method,
   modeledMethod: { ...modeledMethod, type: "summary" },
+  methodCanBeModeled: true,
 };
 
 export const Neutral = Template.bind({});
 Neutral.args = {
   method,
   modeledMethod: { ...modeledMethod, type: "neutral" },
+  methodCanBeModeled: true,
 };
 
 export const AlreadyModeled = Template.bind({});
@@ -107,4 +112,5 @@ ModelingInProgress.args = {
   method,
   modeledMethod,
   modelingInProgress: true,
+  methodCanBeModeled: true,
 };

extensions/ql-vscode/src/variant-analysis/run-remote-query.ts

@@ -37,15 +37,7 @@ import {
 import { QueryLanguage } from "../common/query-language";
 import { tryGetQueryMetadata } from "../codeql-cli/query-metadata";
 import { askForLanguage, findLanguage } from "../codeql-cli/query-language";
-
-interface QlPack {
-  name: string;
-  version: string;
-  library?: boolean;
-  dependencies: { [key: string]: string };
-  defaultSuite?: Array<Record<string, unknown>>;
-  defaultSuiteFile?: string;
-}
+import { QlPackFile } from "../packaging/qlpack-file";
 
 /**
  * Well-known names for the query pack used by the server.
@@ -395,7 +387,7 @@ async function fixPackFile(
       )} file in '${queryPackDir}'`,
     );
   }
-  const qlpack = load(await readFile(packPath, "utf8")) as QlPack;
+  const qlpack = load(await readFile(packPath, "utf8")) as QlPackFile;
 
   updateDefaultSuite(qlpack, packRelativePath);
   removeWorkspaceRefs(qlpack);
@@ -416,7 +408,11 @@ async function injectExtensionPacks(
       )} file in '${queryPackDir}'`,
     );
   }
-  const syntheticQueryPack = load(await readFile(qlpackFile, "utf8")) as QlPack;
+  const syntheticQueryPack = load(
+    await readFile(qlpackFile, "utf8"),
+  ) as QlPackFile;
+
+  const dependencies = syntheticQueryPack.dependencies ?? {};
 
   const extensionPacks = await cliServer.resolveQlpacks(workspaceFolders, true);
   Object.entries(extensionPacks).forEach(([name, paths]) => {
@@ -433,13 +429,16 @@ async function injectExtensionPacks(
     // Add this extension pack as a dependency. It doesn't matter which
     // version we specify, since we are guaranteed that the extension pack
     // is resolved from source at the given path.
-    syntheticQueryPack.dependencies[name] = "*";
+    dependencies[name] = "*";
   });
+
+  syntheticQueryPack.dependencies = dependencies;
+
   await writeFile(qlpackFile, dump(syntheticQueryPack));
   await cliServer.clearCache();
 }
 
-function updateDefaultSuite(qlpack: QlPack, packRelativePath: string) {
+function updateDefaultSuite(qlpack: QlPackFile, packRelativePath: string) {
   delete qlpack.defaultSuiteFile;
   qlpack.defaultSuite = generateDefaultSuite(packRelativePath);
 }
@@ -541,8 +540,12 @@ async function getControllerRepoFromApi(
   }
 }
 
-export function removeWorkspaceRefs(qlpack: QlPack) {
-  for (const [key, value] of Object.entries(qlpack.dependencies || {})) {
+export function removeWorkspaceRefs(qlpack: QlPackFile) {
+  if (!qlpack.dependencies) {
+    return;
+  }
+
+  for (const [key, value] of Object.entries(qlpack.dependencies)) {
     if (value === "${workspace}") {
       qlpack.dependencies[key] = "*";
     }

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

@@ -0,0 +1,17 @@
+import * as React from "react";
+import { styled } from "styled-components";
+import { formatDecimal } from "../../common/number";
+
+const RightAlignedSpan = styled.span`
+  display: inline-block;
+  text-align: right;
+  width: 100%;
+`;
+
+type Props = {
+  value: number;
+};
+
+export const RawNumberValue = ({ value }: Props) => {
+  return <RightAlignedSpan>{formatDecimal(value)}</RightAlignedSpan>;
+};

extensions/ql-vscode/src/view/method-modeling/MethodModeling.tsx

@@ -1,44 +1,67 @@
 import * as React from "react";
 import { styled } from "styled-components";
-import {
-  ModelingStatus,
-  ModelingStatusIndicator,
-} from "../model-editor/ModelingStatusIndicator";
+import { ModelingStatus } from "../../model-editor/shared/modeling-status";
+import { ModelingStatusIndicator } from "../model-editor/ModelingStatusIndicator";
 import { Method } from "../../model-editor/method";
 import { MethodName } from "../model-editor/MethodName";
+import { ModeledMethod } from "../../model-editor/modeled-method";
+import { MethodModelingInputs } from "./MethodModelingInputs";
+import { VSCodeTag } from "@vscode/webview-ui-toolkit/react";
 
 const Container = styled.div`
-  background-color: var(--vscode-peekViewResult-background);
   padding: 0.3rem;
   margin-bottom: 1rem;
+  width: 100%;
 `;
 
 const Title = styled.div`
   padding-bottom: 0.3rem;
-  font-size: 1.2em;
+  font-size: 0.7rem;
+  text-transform: uppercase;
+  display: flex;
+  justify-content: space-between;
 `;
 
 const DependencyContainer = styled.div`
   display: flex;
-  justify-content: space-between;
+  flex-direction: row;
+  align-items: center;
+  gap: 0.5em;
+  background-color: var(--vscode-textBlockQuote-background);
+  border-radius: 0.3rem;
+  border-color: var(--vscode-textBlockQuote-border);
+  padding: 0.5rem;
 `;
 
 export type MethodModelingProps = {
   modelingStatus: ModelingStatus;
   method: Method;
+  modeledMethod: ModeledMethod | undefined;
+  onChange: (modeledMethod: ModeledMethod) => void;
 };
 
 export const MethodModeling = ({
   modelingStatus,
+  modeledMethod,
   method,
+  onChange,
 }: MethodModelingProps): JSX.Element => {
   return (
     <Container>
-      <Title>API or Method</Title>
+      <Title>
+        {method.packageName}
+        {method.libraryVersion && <>@{method.libraryVersion}</>}
+        {modelingStatus === "unsaved" ? <VSCodeTag>Unsaved</VSCodeTag> : null}
+      </Title>
       <DependencyContainer>
-        <MethodName {...method} />
         <ModelingStatusIndicator status={modelingStatus} />
+        <MethodName {...method} />
       </DependencyContainer>
+      <MethodModelingInputs
+        method={method}
+        modeledMethod={modeledMethod}
+        onChange={onChange}
+      />
     </Container>
   );
 };

extensions/ql-vscode/src/view/method-modeling/MethodModelingInputs.tsx

@@ -0,0 +1,66 @@
+import * as React from "react";
+import { styled } from "styled-components";
+import { Method } from "../../model-editor/method";
+import { ModeledMethod } from "../../model-editor/modeled-method";
+import { ModelTypeDropdown } from "../model-editor/ModelTypeDropdown";
+import { ModelInputDropdown } from "../model-editor/ModelInputDropdown";
+import { ModelOutputDropdown } from "../model-editor/ModelOutputDropdown";
+import { ModelKindDropdown } from "../model-editor/ModelKindDropdown";
+
+const Container = styled.div`
+  padding-top: 0.5rem;
+`;
+
+const Input = styled.label``;
+
+const Name = styled.span`
+  display: block;
+  padding-bottom: 0.3rem;
+`;
+
+export type MethodModelingInputsProps = {
+  method: Method;
+  modeledMethod: ModeledMethod | undefined;
+  onChange: (modeledMethod: ModeledMethod) => void;
+};
+
+export const MethodModelingInputs = ({
+  method,
+  modeledMethod,
+  onChange,
+}: MethodModelingInputsProps): JSX.Element => {
+  const inputProps = {
+    method,
+    modeledMethod,
+    onChange,
+  };
+
+  return (
+    <>
+      <Container>
+        <Input>
+          <Name>Model Type</Name>
+          <ModelTypeDropdown {...inputProps} />
+        </Input>
+      </Container>
+      <Container>
+        <Input>
+          <Name>Input</Name>
+          <ModelInputDropdown {...inputProps} />
+        </Input>
+      </Container>
+      <Container>
+        <Input>
+          <Name>Output</Name>
+          <ModelOutputDropdown {...inputProps} />
+        </Input>
+      </Container>
+      <Container>
+        <Input>
+          <Name>Kind</Name>
+          <ModelKindDropdown {...inputProps} />
+        </Input>
+      </Container>
+    </>
+  );
+};

extensions/ql-vscode/src/view/method-modeling/MethodModelingView.tsx

@@ -1,22 +1,48 @@
 import * as React from "react";
-import { useEffect, useState } from "react";
+import { useEffect, useMemo, useState } from "react";
 import { MethodModeling } from "./MethodModeling";
-import { ModelingStatus } from "../model-editor/ModelingStatusIndicator";
+import { getModelingStatus } from "../../model-editor/shared/modeling-status";
 import { Method } from "../../model-editor/method";
 import { ToMethodModelingMessage } from "../../common/interface-types";
 import { assertNever } from "../../common/helpers-pure";
+import { ModeledMethod } from "../../model-editor/modeled-method";
+import { vscode } from "../vscode-api";
 
 export function MethodModelingView(): JSX.Element {
   const [method, setMethod] = useState<Method | undefined>(undefined);
 
+  const [modeledMethod, setModeledMethod] = React.useState<
+    ModeledMethod | undefined
+  >(undefined);
+
+  const [isMethodModified, setIsMethodModified] = useState<boolean>(false);
+
+  const modelingStatus = useMemo(
+    () => getModelingStatus(modeledMethod, isMethodModified),
+    [modeledMethod, isMethodModified],
+  );
+
   useEffect(() => {
     const listener = (evt: MessageEvent) => {
       if (evt.origin === window.origin) {
         const msg: ToMethodModelingMessage = evt.data;
-        if (msg.t === "setMethod") {
-          setMethod(msg.method);
-        } else {
-          assertNever(msg.t);
+        switch (msg.t) {
+          case "setMethod":
+            setMethod(msg.method);
+            break;
+          case "setModeledMethod":
+            setModeledMethod(msg.method);
+            break;
+          case "setMethodModified":
+            setIsMethodModified(msg.isModified);
+            break;
+          case "setSelectedMethod":
+            setMethod(msg.method);
+            setModeledMethod(msg.modeledMethod);
+            setIsMethodModified(msg.isModified);
+            break;
+          default:
+            assertNever(msg);
         }
       } else {
         // sanitize origin
@@ -35,6 +61,19 @@ export function MethodModelingView(): JSX.Element {
     return <>Select method to model</>;
   }
 
-  const modelingStatus: ModelingStatus = "saved";
-  return <MethodModeling modelingStatus={modelingStatus} method={method} />;
+  const onChange = (modeledMethod: ModeledMethod) => {
+    vscode.postMessage({
+      t: "setModeledMethod",
+      method: modeledMethod,
+    });
+  };
+
+  return (
+    <MethodModeling
+      modelingStatus={modelingStatus}
+      method={method}
+      modeledMethod={modeledMethod}
+      onChange={onChange}
+    />
+  );
 }

extensions/ql-vscode/src/view/method-modeling/__tests__/MethodModeling.spec.tsx

@@ -1,18 +1,27 @@
 import * as React from "react";
 import { render as reactRender, screen } from "@testing-library/react";
 import { MethodModeling, MethodModelingProps } from "../MethodModeling";
-import { createMethod } from "../../../../test/factories/data-extension/method-factories";
+import { createMethod } from "../../../../test/factories/model-editor/method-factories";
+import { createModeledMethod } from "../../../../test/factories/model-editor/modeled-method-factories";
 
 describe(MethodModeling.name, () => {
   const render = (props: MethodModelingProps) =>
     reactRender(<MethodModeling {...props} />);
 
   it("renders method modeling panel", () => {
+    const method = createMethod();
+    const modeledMethod = createModeledMethod();
+    const onChange = jest.fn();
+
     render({
       modelingStatus: "saved",
-      method: createMethod(),
+      method,
+      modeledMethod,
+      onChange,
     });
 
-    expect(screen.getByText("API or Method")).toBeInTheDocument();
+    expect(
+      screen.getByText(`${method.packageName}@${method.libraryVersion}`),
+    ).toBeInTheDocument();
   });
 });

extensions/ql-vscode/src/view/method-modeling/__tests__/MethodModelingInputs.spec.tsx

@@ -0,0 +1,100 @@
+import * as React from "react";
+import { render as reactRender, screen } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
+import {
+  MethodModelingInputs,
+  MethodModelingInputsProps,
+} from "../MethodModelingInputs";
+import { createMethod } from "../../../../test/factories/model-editor/method-factories";
+import { createModeledMethod } from "../../../../test/factories/model-editor/modeled-method-factories";
+
+describe(MethodModelingInputs.name, () => {
+  const render = (props: MethodModelingInputsProps) =>
+    reactRender(<MethodModelingInputs {...props} />);
+
+  const method = createMethod();
+  const modeledMethod = createModeledMethod();
+  const onChange = jest.fn();
+
+  it("renders the method modeling inputs", () => {
+    render({
+      method,
+      modeledMethod,
+      onChange,
+    });
+
+    // Check that all the labels are rendered.
+    expect(screen.getByText("Model Type")).toBeInTheDocument();
+    expect(screen.getByText("Input")).toBeInTheDocument();
+    expect(screen.getByText("Output")).toBeInTheDocument();
+    expect(screen.getByText("Kind")).toBeInTheDocument();
+
+    // Check that all the dropdowns are rendered.
+    const comboboxes = screen.getAllByRole("combobox");
+    expect(comboboxes.length).toBe(4);
+    const modelTypeDropdown = screen.getByRole("combobox", {
+      name: "Model type",
+    });
+    expect(modelTypeDropdown).toHaveValue("sink");
+    const modelTypeOptions = modelTypeDropdown.querySelectorAll("option");
+    expect(modelTypeOptions.length).toBe(5);
+  });
+
+  it("allows changing the type", async () => {
+    render({
+      method,
+      modeledMethod,
+      onChange,
+    });
+
+    const modelTypeDropdown = screen.getByRole("combobox", {
+      name: "Model type",
+    });
+
+    await userEvent.selectOptions(modelTypeDropdown, "source");
+
+    expect(onChange).toHaveBeenCalledWith(
+      expect.objectContaining({
+        type: "source",
+      }),
+    );
+  });
+
+  it("sets other dropdowns when model type is changed", () => {
+    const { rerender } = render({
+      method,
+      modeledMethod,
+      onChange,
+    });
+
+    const updatedModeledMethod = createModeledMethod({
+      type: "source",
+    });
+
+    rerender(
+      <MethodModelingInputs
+        method={method}
+        modeledMethod={updatedModeledMethod}
+        onChange={onChange}
+      />,
+    );
+
+    const modelTypeDropdown = screen.getByRole("combobox", {
+      name: "Model type",
+    });
+    const modelInputDropdown = screen.getByRole("combobox", {
+      name: "Input",
+    });
+    const modelOutputDropdown = screen.getByRole("combobox", {
+      name: "Output",
+    });
+    const modelKindDropdown = screen.getByRole("combobox", {
+      name: "Kind",
+    });
+
+    expect(modelTypeDropdown).toHaveValue("source");
+    expect(modelInputDropdown).toHaveValue("-");
+    expect(modelOutputDropdown).toHaveValue("ReturnValue");
+    expect(modelKindDropdown).toHaveValue("local");
+  });
+});

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

@@ -1,56 +0,0 @@
-import * as React from "react";
-import { ChangeEvent, useCallback, useEffect, useMemo } from "react";
-import type { ModeledMethod } from "../../model-editor/modeled-method";
-import { Dropdown } from "../common/Dropdown";
-
-type Props = {
-  kinds: Array<ModeledMethod["kind"]>;
-
-  value: ModeledMethod["kind"] | undefined;
-  disabled?: boolean;
-  onChange: (value: ModeledMethod["kind"]) => void;
-
-  "aria-label"?: string;
-};
-
-export const KindInput = ({
-  kinds,
-  value,
-  disabled,
-  onChange,
-  ...props
-}: Props) => {
-  const options = useMemo(
-    () => kinds.map((kind) => ({ value: kind, label: kind })),
-    [kinds],
-  );
-
-  const handleInput = useCallback(
-    (e: ChangeEvent<HTMLSelectElement>) => {
-      const target = e.target as HTMLSelectElement;
-
-      onChange(target.value as ModeledMethod["kind"]);
-    },
-    [onChange],
-  );
-
-  useEffect(() => {
-    if (value === undefined && kinds.length > 0) {
-      onChange(kinds[0]);
-    }
-
-    if (value !== undefined && !kinds.includes(value)) {
-      onChange(kinds[0]);
-    }
-  }, [value, kinds, onChange]);
-
-  return (
-    <Dropdown
-      value={value}
-      options={options}
-      disabled={disabled}
-      onChange={handleInput}
-      {...props}
-    />
-  );
-};

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

@@ -76,11 +76,7 @@ export type LibraryRowProps = {
   inProgressMethods: InProgressMethods;
   viewState: ModelEditorViewState;
   hideModeledMethods: boolean;
-  onChange: (
-    modelName: string,
-    method: Method,
-    modeledMethod: ModeledMethod,
-  ) => void;
+  onChange: (modeledMethod: ModeledMethod) => void;
   onSaveModelClick: (
     methods: Method[],
     modeledMethods: Record<string, ModeledMethod>,
@@ -166,13 +162,6 @@ export const LibraryRow = ({
     [methods, modeledMethods, onSaveModelClick],
   );
 
-  const onChangeWithModelName = useCallback(
-    (method: Method, modeledMethod: ModeledMethod) => {
-      onChange(title, method, modeledMethod);
-    },
-    [onChange, title],
-  );
-
   const hasUnsavedChanges = useMemo(() => {
     return methods.some((method) => modifiedSignatures.has(method.signature));
   }, [methods, modifiedSignatures]);
@@ -238,7 +227,7 @@ export const LibraryRow = ({
             inProgressMethods={inProgressMethods}
             mode={viewState.mode}
             hideModeledMethods={hideModeledMethods}
-            onChange={onChangeWithModelName}
+            onChange={onChange}
           />
           <SectionDivider />
           <ButtonsContainer>

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

@@ -5,27 +5,22 @@ import {
   VSCodeProgressRing,
 } from "@vscode/webview-ui-toolkit/react";
 import * as React from "react";
-import { ChangeEvent, useCallback, useMemo } from "react";
+import { useCallback } from "react";
 import { styled } from "styled-components";
 import { vscode } from "../vscode-api";
 
 import { Method } from "../../model-editor/method";
-import {
-  ModeledMethod,
-  ModeledMethodType,
-  Provenance,
-} from "../../model-editor/modeled-method";
-import { KindInput } from "./KindInput";
-import { extensiblePredicateDefinitions } from "../../model-editor/predicates";
+import { ModeledMethod } from "../../model-editor/modeled-method";
+import { ModelKindDropdown } from "./ModelKindDropdown";
 import { Mode } from "../../model-editor/shared/mode";
-import { Dropdown } from "../common/Dropdown";
 import { MethodClassifications } from "./MethodClassifications";
-import {
-  ModelingStatus,
-  ModelingStatusIndicator,
-} from "./ModelingStatusIndicator";
+import { getModelingStatus } from "../../model-editor/shared/modeling-status";
+import { ModelingStatusIndicator } from "./ModelingStatusIndicator";
 import { InProgressDropdown } from "./InProgressDropdown";
 import { MethodName } from "./MethodName";
+import { ModelTypeDropdown } from "./ModelTypeDropdown";
+import { ModelInputDropdown } from "./ModelInputDropdown";
+import { ModelOutputDropdown } from "./ModelOutputDropdown";
 
 const ApiOrMethodCell = styled(VSCodeDataGridCell)`
   display: flex;
@@ -52,14 +47,6 @@ const ProgressRing = styled(VSCodeProgressRing)`
   margin-left: auto;
 `;
 
-const modelTypeOptions: Array<{ value: ModeledMethodType; label: string }> = [
-  { value: "none", label: "Unmodeled" },
-  { value: "source", label: "Source" },
-  { value: "sink", label: "Sink" },
-  { value: "summary", label: "Flow summary" },
-  { value: "neutral", label: "Neutral" },
-];
-
 export type MethodRowProps = {
   method: Method;
   methodCanBeModeled: boolean;
@@ -67,7 +54,7 @@ export type MethodRowProps = {
   methodIsUnsaved: boolean;
   modelingInProgress: boolean;
   mode: Mode;
-  onChange: (method: Method, modeledMethod: ModeledMethod) => void;
+  onChange: (modeledMethod: ModeledMethod) => void;
 };
 
 export const MethodRow = (props: MethodRowProps) => {
@@ -83,123 +70,11 @@ export const MethodRow = (props: MethodRowProps) => {
 function ModelableMethodRow(props: MethodRowProps) {
   const { method, modeledMethod, methodIsUnsaved, mode, onChange } = props;
 
-  const argumentsList = useMemo(() => {
-    if (method.methodParameters === "()") {
-      return [];
-    }
-    return method.methodParameters
-      .substring(1, method.methodParameters.length - 1)
-      .split(",");
-  }, [method.methodParameters]);
-
-  const handleTypeInput = useCallback(
-    (e: ChangeEvent<HTMLSelectElement>) => {
-      let newProvenance: Provenance = "manual";
-      if (modeledMethod?.provenance === "df-generated") {
-        newProvenance = "df-manual";
-      } else if (modeledMethod?.provenance === "ai-generated") {
-        newProvenance = "ai-manual";
-      }
-
-      onChange(method, {
-        // If there are no arguments, we will default to "Argument[this]"
-        input: argumentsList.length === 0 ? "Argument[this]" : "Argument[0]",
-        output: "ReturnType",
-        kind: "value",
-        ...modeledMethod,
-        type: e.target.value as ModeledMethodType,
-        provenance: newProvenance,
-        signature: method.signature,
-        packageName: method.packageName,
-        typeName: method.typeName,
-        methodName: method.methodName,
-        methodParameters: method.methodParameters,
-      });
-    },
-    [onChange, method, modeledMethod, argumentsList],
-  );
-  const handleInputInput = useCallback(
-    (e: ChangeEvent<HTMLSelectElement>) => {
-      if (!modeledMethod) {
-        return;
-      }
-
-      const target = e.target as HTMLSelectElement;
-
-      onChange(method, {
-        ...modeledMethod,
-        input: target.value as ModeledMethod["input"],
-      });
-    },
-    [onChange, method, modeledMethod],
-  );
-  const handleOutputInput = useCallback(
-    (e: ChangeEvent<HTMLSelectElement>) => {
-      if (!modeledMethod) {
-        return;
-      }
-
-      const target = e.target as HTMLSelectElement;
-
-      onChange(method, {
-        ...modeledMethod,
-        output: target.value as ModeledMethod["output"],
-      });
-    },
-    [onChange, method, modeledMethod],
-  );
-  const handleKindChange = useCallback(
-    (kind: string) => {
-      if (!modeledMethod) {
-        return;
-      }
-
-      onChange(method, {
-        ...modeledMethod,
-        kind,
-      });
-    },
-    [onChange, method, modeledMethod],
-  );
-
   const jumpToUsage = useCallback(
     () => sendJumpToUsageMessage(method),
     [method],
   );
 
-  const inputOptions = useMemo(
-    () => [
-      { value: "Argument[this]", label: "Argument[this]" },
-      ...argumentsList.map((argument, index) => ({
-        value: `Argument[${index}]`,
-        label: `Argument[${index}]: ${argument}`,
-      })),
-    ],
-    [argumentsList],
-  );
-
-  const outputOptions = useMemo(
-    () => [
-      { value: "ReturnValue", label: "ReturnValue" },
-      { value: "Argument[this]", label: "Argument[this]" },
-      ...argumentsList.map((argument, index) => ({
-        value: `Argument[${index}]`,
-        label: `Argument[${index}]: ${argument}`,
-      })),
-    ],
-    [argumentsList],
-  );
-
-  const showInputCell =
-    modeledMethod?.type && ["sink", "summary"].includes(modeledMethod?.type);
-  const showOutputCell =
-    modeledMethod?.type && ["source", "summary"].includes(modeledMethod?.type);
-  const predicate =
-    modeledMethod?.type && modeledMethod.type !== "none"
-      ? extensiblePredicateDefinitions[modeledMethod.type]
-      : undefined;
-  const showKindCell = predicate?.supportedKinds;
-
   const modelingStatus = getModelingStatus(modeledMethod, methodIsUnsaved);
 
   return (
@@ -235,38 +110,31 @@ function ModelableMethodRow(props: MethodRowProps) {
       {!props.modelingInProgress && (
         <>
           <VSCodeDataGridCell gridColumn={2}>
-            <Dropdown
-              value={modeledMethod?.type ?? "none"}
-              options={modelTypeOptions}
-              onChange={handleTypeInput}
-              aria-label="Model type"
+            <ModelTypeDropdown
+              method={method}
+              modeledMethod={modeledMethod}
+              onChange={onChange}
             />
           </VSCodeDataGridCell>
           <VSCodeDataGridCell gridColumn={3}>
-            <Dropdown
-              value={modeledMethod?.input}
-              options={inputOptions}
-              disabled={!showInputCell}
-              onChange={handleInputInput}
-              aria-label="Input"
+            <ModelInputDropdown
+              method={method}
+              modeledMethod={modeledMethod}
+              onChange={onChange}
             />
           </VSCodeDataGridCell>
           <VSCodeDataGridCell gridColumn={4}>
-            <Dropdown
-              value={modeledMethod?.output}
-              options={outputOptions}
-              disabled={!showOutputCell}
-              onChange={handleOutputInput}
-              aria-label="Output"
+            <ModelOutputDropdown
+              method={method}
+              modeledMethod={modeledMethod}
+              onChange={onChange}
             />
           </VSCodeDataGridCell>
           <VSCodeDataGridCell gridColumn={5}>
-            <KindInput
-              kinds={predicate?.supportedKinds || []}
-              value={modeledMethod?.kind}
-              disabled={!showKindCell}
-              onChange={handleKindChange}
-              aria-label="Kind"
+            <ModelKindDropdown
+              method={method}
+              modeledMethod={modeledMethod}
+              onChange={onChange}
             />
           </VSCodeDataGridCell>
         </>
@@ -311,17 +179,3 @@ function sendJumpToUsageMessage(method: Method) {
     usage: method.usages[0],
   });
 }
-
-function getModelingStatus(
-  modeledMethod: ModeledMethod | undefined,
-  methodIsUnsaved: boolean,
-): ModelingStatus {
-  if (modeledMethod) {
-    if (methodIsUnsaved) {
-      return "unsaved";
-    } else if (modeledMethod.type !== "none") {
-      return "saved";
-    }
-  }
-  return "unmodeled";
-}

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

@@ -123,32 +123,11 @@ export function ModelEditor({
           case "setMethods":
             setMethods(msg.methods);
             break;
-          case "loadModeledMethods":
-            setModeledMethods((oldModeledMethods) => {
-              return {
-                ...msg.modeledMethods,
-                ...oldModeledMethods,
-              };
-            });
+          case "setModeledMethods":
+            setModeledMethods(msg.methods);
             break;
-          case "addModeledMethods":
-            setModeledMethods((oldModeledMethods) => {
-              return {
-                ...msg.modeledMethods,
-                ...Object.fromEntries(
-                  Object.entries(oldModeledMethods).filter(
-                    ([, value]) => value.type !== "none",
-                  ),
-                ),
-              };
-            });
-            setModifiedSignatures(
-              (oldModifiedSignatures) =>
-                new Set([
-                  ...oldModifiedSignatures,
-                  ...Object.keys(msg.modeledMethods),
-                ]),
-            );
+          case "setModifiedMethods":
+            setModifiedSignatures(new Set(msg.methodSignatures));
             break;
           case "setInProgressMethods":
             setInProgressMethods((oldInProgressMethods) =>
@@ -179,19 +158,12 @@ export function ModelEditor({
     [methods],
   );
 
-  const onChange = useCallback(
-    (modelName: string, method: Method, model: ModeledMethod) => {
-      setModeledMethods((oldModeledMethods) => ({
-        ...oldModeledMethods,
-        [method.signature]: model,
-      }));
-      setModifiedSignatures(
-        (oldModifiedSignatures) =>
-          new Set([...oldModifiedSignatures, method.signature]),
-      );
-    },
-    [],
-  );
+  const onChange = useCallback((model: ModeledMethod) => {
+    vscode.postMessage({
+      t: "setModeledMethod",
+      method: model,
+    });
+  }, []);
 
   const onRefreshClick = useCallback(() => {
     vscode.postMessage({
@@ -205,7 +177,6 @@ export function ModelEditor({
       methods,
       modeledMethods,
     });
-    setModifiedSignatures(new Set());
   }, [methods, modeledMethods]);
 
   const onSaveModelClick = useCallback(
@@ -215,13 +186,6 @@ export function ModelEditor({
         methods,
         modeledMethods,
       });
-      setModifiedSignatures((oldModifiedSignatures) => {
-        const newModifiedSignatures = new Set([...oldModifiedSignatures]);
-        for (const method of methods) {
-          newModifiedSignatures.delete(method.signature);
-        }
-        return newModifiedSignatures;
-      });
     },
     [],
   );

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

@@ -0,0 +1,65 @@
+import * as React from "react";
+import { ChangeEvent, useCallback, useMemo } from "react";
+import { Dropdown } from "../common/Dropdown";
+import { ModeledMethod } from "../../model-editor/modeled-method";
+import { Method, getArgumentsList } from "../../model-editor/method";
+
+type Props = {
+  method: Method;
+  modeledMethod: ModeledMethod | undefined;
+  onChange: (modeledMethod: ModeledMethod) => void;
+};
+
+export const ModelInputDropdown = ({
+  method,
+  modeledMethod,
+  onChange,
+}: Props): JSX.Element => {
+  const argumentsList = useMemo(
+    () => getArgumentsList(method.methodParameters),
+    [method.methodParameters],
+  );
+
+  const options = useMemo(
+    () => [
+      { value: "Argument[this]", label: "Argument[this]" },
+      ...argumentsList.map((argument, index) => ({
+        value: `Argument[${index}]`,
+        label: `Argument[${index}]: ${argument}`,
+      })),
+    ],
+    [argumentsList],
+  );
+
+  const enabled = useMemo(
+    () =>
+      modeledMethod?.type && ["sink", "summary"].includes(modeledMethod?.type),
+    [modeledMethod?.type],
+  );
+
+  const handleChange = useCallback(
+    (e: ChangeEvent<HTMLSelectElement>) => {
+      if (!modeledMethod) {
+        return;
+      }
+
+      const target = e.target as HTMLSelectElement;
+
+      onChange({
+        ...modeledMethod,
+        input: target.value,
+      });
+    },
+    [onChange, modeledMethod],
+  );
+
+  return (
+    <Dropdown
+      value={modeledMethod?.input}
+      options={options}
+      disabled={!enabled}
+      onChange={handleChange}
+      aria-label="Input"
+    />
+  );
+};

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

@@ -0,0 +1,84 @@
+import * as React from "react";
+import { ChangeEvent, useCallback, useEffect, useMemo } from "react";
+import type {
+  ModeledMethod,
+  ModeledMethodKind,
+} from "../../model-editor/modeled-method";
+import { Dropdown } from "../common/Dropdown";
+import { Method } from "../../model-editor/method";
+import { extensiblePredicateDefinitions } from "../../model-editor/predicates";
+
+type Props = {
+  method: Method;
+  modeledMethod: ModeledMethod | undefined;
+  onChange: (modeledMethod: ModeledMethod) => void;
+};
+
+export const ModelKindDropdown = ({
+  method,
+  modeledMethod,
+  onChange,
+}: Props) => {
+  const predicate = useMemo(() => {
+    return modeledMethod?.type && modeledMethod.type !== "none"
+      ? extensiblePredicateDefinitions[modeledMethod.type]
+      : undefined;
+  }, [modeledMethod?.type]);
+
+  const kinds = useMemo(() => predicate?.supportedKinds || [], [predicate]);
+
+  const disabled = useMemo(
+    () => !predicate?.supportedKinds,
+    [predicate?.supportedKinds],
+  );
+
+  const options = useMemo(
+    () => kinds.map((kind) => ({ value: kind, label: kind })),
+    [kinds],
+  );
+
+  const onChangeKind = useCallback(
+    (kind: ModeledMethodKind) => {
+      if (!modeledMethod) {
+        return;
+      }
+
+      onChange({
+        ...modeledMethod,
+        kind,
+      });
+    },
+    [modeledMethod, onChange],
+  );
+
+  const handleChange = useCallback(
+    (e: ChangeEvent<HTMLSelectElement>) => {
+      const target = e.target as HTMLSelectElement;
+      const kind = target.value;
+
+      onChangeKind(kind);
+    },
+    [onChangeKind],
+  );
+
+  useEffect(() => {
+    const value = modeledMethod?.kind;
+    if (value === undefined && kinds.length > 0) {
+      onChangeKind(kinds[0]);
+    }
+
+    if (value !== undefined && !kinds.includes(value)) {
+      onChangeKind(kinds[0]);
+    }
+  }, [modeledMethod?.kind, kinds, onChangeKind]);
+
+  return (
+    <Dropdown
+      value={modeledMethod?.kind}
+      options={options}
+      disabled={disabled}
+      onChange={handleChange}
+      aria-label="Kind"
+    />
+  );
+};

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

@@ -0,0 +1,67 @@
+import * as React from "react";
+import { ChangeEvent, useCallback, useMemo } from "react";
+import { Dropdown } from "../common/Dropdown";
+import { ModeledMethod } from "../../model-editor/modeled-method";
+import { Method, getArgumentsList } from "../../model-editor/method";
+
+type Props = {
+  method: Method;
+  modeledMethod: ModeledMethod | undefined;
+  onChange: (modeledMethod: ModeledMethod) => void;
+};
+
+export const ModelOutputDropdown = ({
+  method,
+  modeledMethod,
+  onChange,
+}: Props): JSX.Element => {
+  const argumentsList = useMemo(
+    () => getArgumentsList(method.methodParameters),
+    [method.methodParameters],
+  );
+
+  const options = useMemo(
+    () => [
+      { value: "ReturnValue", label: "ReturnValue" },
+      { value: "Argument[this]", label: "Argument[this]" },
+      ...argumentsList.map((argument, index) => ({
+        value: `Argument[${index}]`,
+        label: `Argument[${index}]: ${argument}`,
+      })),
+    ],
+    [argumentsList],
+  );
+
+  const enabled = useMemo(
+    () =>
+      modeledMethod?.type &&
+      ["source", "summary"].includes(modeledMethod?.type),
+    [modeledMethod?.type],
+  );
+
+  const handleChange = useCallback(
+    (e: ChangeEvent<HTMLSelectElement>) => {
+      if (!modeledMethod) {
+        return;
+      }
+
+      const target = e.target as HTMLSelectElement;
+
+      onChange({
+        ...modeledMethod,
+        output: target.value,
+      });
+    },
+    [onChange, modeledMethod],
+  );
+
+  return (
+    <Dropdown
+      value={modeledMethod?.output}
+      options={options}
+      disabled={!enabled}
+      onChange={handleChange}
+      aria-label="Output"
+    />
+  );
+};

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

@@ -0,0 +1,70 @@
+import * as React from "react";
+import { ChangeEvent, useCallback, useMemo } from "react";
+import { Dropdown } from "../common/Dropdown";
+import {
+  ModeledMethod,
+  ModeledMethodType,
+  Provenance,
+} from "../../model-editor/modeled-method";
+import { Method, getArgumentsList } from "../../model-editor/method";
+
+const options: Array<{ value: ModeledMethodType; label: string }> = [
+  { value: "none", label: "Unmodeled" },
+  { value: "source", label: "Source" },
+  { value: "sink", label: "Sink" },
+  { value: "summary", label: "Flow summary" },
+  { value: "neutral", label: "Neutral" },
+];
+
+type Props = {
+  method: Method;
+  modeledMethod: ModeledMethod | undefined;
+  onChange: (modeledMethod: ModeledMethod) => void;
+};
+
+export const ModelTypeDropdown = ({
+  method,
+  modeledMethod,
+  onChange,
+}: Props): JSX.Element => {
+  const argumentsList = useMemo(
+    () => getArgumentsList(method.methodParameters),
+    [method.methodParameters],
+  );
+
+  const handleChange = useCallback(
+    (e: ChangeEvent<HTMLSelectElement>) => {
+      let newProvenance: Provenance = "manual";
+      if (modeledMethod?.provenance === "df-generated") {
+        newProvenance = "df-manual";
+      } else if (modeledMethod?.provenance === "ai-generated") {
+        newProvenance = "ai-manual";
+      }
+
+      const updatedModeledMethod: ModeledMethod = {
+        // If there are no arguments, we will default to "Argument[this]"
+        input: argumentsList.length === 0 ? "Argument[this]" : "Argument[0]",
+        output: "ReturnValue",
+        kind: "value",
+        type: e.target.value as ModeledMethodType,
+        provenance: newProvenance,
+        signature: method.signature,
+        packageName: method.packageName,
+        typeName: method.typeName,
+        methodName: method.methodName,
+        methodParameters: method.methodParameters,
+      };
+      onChange(updatedModeledMethod);
+    },
+    [onChange, method, modeledMethod, argumentsList],
+  );
+
+  return (
+    <Dropdown
+      value={modeledMethod?.type ?? "none"}
+      options={options}
+      onChange={handleChange}
+      aria-label="Model type"
+    />
+  );
+};

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

@@ -23,7 +23,7 @@ export type ModeledMethodDataGridProps = {
   inProgressMethods: InProgressMethods;
   mode: Mode;
   hideModeledMethods: boolean;
-  onChange: (method: Method, modeledMethod: ModeledMethod) => void;
+  onChange: (modeledMethod: ModeledMethod) => void;
 };
 
 export const ModeledMethodDataGrid = ({

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

@@ -18,11 +18,7 @@ export type ModeledMethodsListProps = {
   inProgressMethods: InProgressMethods;
   viewState: ModelEditorViewState;
   hideModeledMethods: boolean;
-  onChange: (
-    modelName: string,
-    method: Method,
-    modeledMethod: ModeledMethod,
-  ) => void;
+  onChange: (modeledMethod: ModeledMethod) => void;
   onSaveModelClick: (
     methods: Method[],
     modeledMethods: Record<string, ModeledMethod>,

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

@@ -1,8 +1,7 @@
 import * as React from "react";
 import { assertNever } from "../../common/helpers-pure";
 import { Codicon } from "../common/icon/Codicon";
-
-export type ModelingStatus = "unmodeled" | "unsaved" | "saved";
+import { ModelingStatus } from "../../model-editor/shared/modeling-status";
 
 interface Props {
   status: ModelingStatus;

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

@@ -1,62 +0,0 @@
-import * as React from "react";
-import { render, screen } from "@testing-library/react";
-import { KindInput } from "../KindInput";
-import userEvent from "@testing-library/user-event";
-
-describe(KindInput.name, () => {
-  const onChange = jest.fn();
-
-  beforeEach(() => {
-    onChange.mockReset();
-  });
-
-  it("allows changing the kind", async () => {
-    render(
-      <KindInput
-        kinds={["local", "remote"]}
-        value="local"
-        onChange={onChange}
-      />,
-    );
-
-    expect(screen.getByRole("combobox")).toHaveValue("local");
-    await userEvent.selectOptions(screen.getByRole("combobox"), "remote");
-    expect(onChange).toHaveBeenCalledWith("remote");
-  });
-
-  it("resets the kind when changing the supported kinds", () => {
-    const { rerender } = render(
-      <KindInput
-        kinds={["local", "remote"]}
-        value={"local"}
-        onChange={onChange}
-      />,
-    );
-
-    expect(screen.getByRole("combobox")).toHaveValue("local");
-    expect(onChange).not.toHaveBeenCalled();
-
-    rerender(
-      <KindInput
-        kinds={["sql-injection", "log-injection", "url-redirection"]}
-        value="local"
-        onChange={onChange}
-      />,
-    );
-    expect(screen.getByRole("combobox")).toHaveValue("sql-injection");
-    expect(onChange).toHaveBeenCalledWith("sql-injection");
-  });
-
-  it("sets the kind when value is undefined", () => {
-    render(
-      <KindInput
-        kinds={["local", "remote"]}
-        value={undefined}
-        onChange={onChange}
-      />,
-    );
-
-    expect(screen.getByRole("combobox")).toHaveValue("local");
-    expect(onChange).toHaveBeenCalledWith("local");
-  });
-});

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

@@ -1,6 +1,6 @@
 import * as React from "react";
 import { render as reactRender, screen } from "@testing-library/react";
-import { createMethod } from "../../../../test/factories/data-extension/method-factories";
+import { createMethod } from "../../../../test/factories/model-editor/method-factories";
 import { LibraryRow, LibraryRowProps } from "../LibraryRow";
 import { InProgressMethods } from "../../../model-editor/shared/in-progress-methods";
 import { createMockExtensionPack } from "../../../../test/factories/model-editor/extension-pack";

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

@@ -2,7 +2,7 @@ import * as React from "react";
 import { render as reactRender, screen } from "@testing-library/react";
 import { MethodName } from "../MethodName";
 import { Method } from "../../../model-editor/method";
-import { createMethod } from "../../../../test/factories/data-extension/method-factories";
+import { createMethod } from "../../../../test/factories/model-editor/method-factories";
 
 describe(MethodName.name, () => {
   const render = (props: Method) => reactRender(<MethodName {...props} />);

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

@@ -4,7 +4,7 @@ import {
   render as reactRender,
   screen,
 } from "@testing-library/react";
-import { createMethod } from "../../../../test/factories/data-extension/method-factories";
+import { createMethod } from "../../../../test/factories/model-editor/method-factories";
 import { Mode } from "../../../model-editor/shared/mode";
 import { MethodRow, MethodRowProps } from "../MethodRow";
 import { ModeledMethod } from "../../../model-editor/modeled-method";
@@ -66,7 +66,7 @@ describe(MethodRow.name, () => {
     );
 
     expect(onChange).toHaveBeenCalledTimes(1);
-    expect(onChange).toHaveBeenCalledWith(method, {
+    expect(onChange).toHaveBeenCalledWith({
       ...modeledMethod,
       kind: "value",
     });

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

@@ -0,0 +1,94 @@
+import * as React from "react";
+import { render, screen } from "@testing-library/react";
+import { ModelKindDropdown } from "../ModelKindDropdown";
+import userEvent from "@testing-library/user-event";
+import { createMethod } from "../../../../test/factories/model-editor/method-factories";
+import { createModeledMethod } from "../../../../test/factories/model-editor/modeled-method-factories";
+
+describe(ModelKindDropdown.name, () => {
+  const onChange = jest.fn();
+  const method = createMethod();
+
+  beforeEach(() => {
+    onChange.mockReset();
+  });
+
+  it("allows changing the kind", async () => {
+    const modeledMethod = createModeledMethod({
+      type: "source",
+      kind: "local",
+    });
+
+    render(
+      <ModelKindDropdown
+        method={method}
+        modeledMethod={modeledMethod}
+        onChange={onChange}
+      />,
+    );
+
+    expect(screen.getByRole("combobox")).toHaveValue("local");
+    await userEvent.selectOptions(screen.getByRole("combobox"), "remote");
+    expect(onChange).toHaveBeenCalledWith(
+      expect.objectContaining({
+        kind: "remote",
+      }),
+    );
+  });
+
+  it("resets the kind when changing the supported kinds", () => {
+    const method = createMethod();
+    const modeledMethod = createModeledMethod({
+      type: "source",
+      kind: "local",
+    });
+
+    const { rerender } = render(
+      <ModelKindDropdown
+        method={method}
+        modeledMethod={modeledMethod}
+        onChange={onChange}
+      />,
+    );
+
+    expect(screen.getByRole("combobox")).toHaveValue("local");
+    expect(onChange).not.toHaveBeenCalled();
+
+    // Changing the type to sink should update the supported kinds
+    const updatedModeledMethod = createModeledMethod({
+      type: "sink",
+    });
+
+    rerender(
+      <ModelKindDropdown
+        method={method}
+        modeledMethod={updatedModeledMethod}
+        onChange={onChange}
+      />,
+    );
+
+    expect(screen.getByRole("combobox")).toHaveValue("code-injection");
+  });
+
+  it("sets the kind when value is undefined", () => {
+    const method = createMethod();
+    const modeledMethod = createModeledMethod({
+      type: "source",
+    });
+
+    render(
+      <ModelKindDropdown
+        method={method}
+        modeledMethod={modeledMethod}
+        onChange={onChange}
+      />,
+    );
+
+    expect(screen.getByRole("combobox")).toHaveValue("local");
+    expect(onChange).toHaveBeenCalledWith(
+      expect.objectContaining({
+        kind: "local",
+      }),
+    );
+  });
+});

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

@@ -1,6 +1,6 @@
 import * as React from "react";
 import { render as reactRender, screen } from "@testing-library/react";
-import { createMethod } from "../../../../test/factories/data-extension/method-factories";
+import { createMethod } from "../../../../test/factories/model-editor/method-factories";
 import { InProgressMethods } from "../../../model-editor/shared/in-progress-methods";
 import { Mode } from "../../../model-editor/shared/mode";
 import {

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

@@ -1,6 +1,6 @@
 import * as React from "react";
 import { render as reactRender, screen } from "@testing-library/react";
-import { createMethod } from "../../../../test/factories/data-extension/method-factories";
+import { createMethod } from "../../../../test/factories/model-editor/method-factories";
 import { InProgressMethods } from "../../../model-editor/shared/in-progress-methods";
 import { createMockExtensionPack } from "../../../../test/factories/model-editor/extension-pack";
 import { Mode } from "../../../model-editor/shared/mode";

extensions/ql-vscode/src/view/results/RawTableValue.tsx

@@ -2,6 +2,7 @@ import * as React from "react";
 
 import { Location } from "./locations/Location";
 import { CellValue } from "../../common/bqrs-cli-types";
+import { RawNumberValue } from "../common/RawNumberValue";
 
 interface Props {
   value: CellValue;
@@ -9,22 +10,26 @@ interface Props {
   onSelected?: () => void;
 }
 
-export default function RawTableValue(props: Props): JSX.Element {
-  const rawValue = props.value;
-  if (
-    typeof rawValue === "string" ||
-    typeof rawValue === "number" ||
-    typeof rawValue === "boolean"
-  ) {
-    return <Location label={rawValue.toString()} />;
+export default function RawTableValue({
+  value,
+  databaseUri,
+  onSelected,
+}: Props): JSX.Element {
+  switch (typeof value) {
+    case "boolean":
+      return <span>{value.toString()}</span>;
+    case "number":
+      return <RawNumberValue value={value} />;
+    case "string":
+      return <Location label={value.toString()} />;
+    default:
+      return (
+        <Location
+          loc={value.url}
+          label={value.label}
+          databaseUri={databaseUri}
+          onClick={onSelected}
+        />
+      );
   }
-
-  return (
-    <Location
-      loc={rawValue.url}
-      label={rawValue.label}
-      databaseUri={props.databaseUri}
-      onClick={props.onSelected}
-    />
-  );
 }

extensions/ql-vscode/src/view/variant-analysis/RawResultCell.tsx

@@ -0,0 +1,48 @@
+import * as React from "react";
+import { VSCodeLink } from "@vscode/webview-ui-toolkit/react";
+
+import { CellValue } from "../../common/bqrs-cli-types";
+import { sendTelemetry } from "../common/telemetry";
+import { convertNonPrintableChars } from "../../common/text-utils";
+import { tryGetRemoteLocation } from "../../common/bqrs-utils";
+import { RawNumberValue } from "../common/RawNumberValue";
+
+type CellProps = {
+  value: CellValue;
+  fileLinkPrefix: string;
+  sourceLocationPrefix: string;
+};
+
+const sendRawResultsLinkTelemetry = () => sendTelemetry("raw-results-link");
+
+export const RawResultCell = ({
+  value,
+  fileLinkPrefix,
+  sourceLocationPrefix,
+}: CellProps) => {
+  switch (typeof value) {
+    case "boolean":
+      return <span>{value.toString()}</span>;
+    case "number":
+      return <RawNumberValue value={value} />;
+    case "string":
+      return <span>{convertNonPrintableChars(value.toString())}</span>;
+    case "object": {
+      const url = tryGetRemoteLocation(
+        value.url,
+        fileLinkPrefix,
+        sourceLocationPrefix,
+      );
+      const safeLabel = convertNonPrintableChars(value.label);
+      if (url) {
+        return (
+          <VSCodeLink onClick={sendRawResultsLinkTelemetry} href={url}>
+            {safeLabel}
+          </VSCodeLink>
+        );
+      } else {
+        return <span>{safeLabel}</span>;
+      }
+    }
+  }
+};

extensions/ql-vscode/src/view/variant-analysis/RawResultRow.tsx

@@ -0,0 +1,38 @@
+import * as React from "react";
+import { styled } from "styled-components";
+
+import { CellValue } from "../../common/bqrs-cli-types";
+import { RawResultCell } from "./RawResultCell";
+
+const StyledRow = styled.div`
+  border-color: var(--vscode-editor-snippetFinalTabstopHighlightBorder);
+  border-style: solid;
+  justify-content: center;
+  align-items: center;
+  padding: 0.4rem;
+  word-break: break-word;
+`;
+
+type RowProps = {
+  row: CellValue[];
+  fileLinkPrefix: string;
+  sourceLocationPrefix: string;
+};
+
+export const RawResultRow = ({
+  row,
+  fileLinkPrefix,
+  sourceLocationPrefix,
+}: RowProps) => (
+  <>
+    {row.map((cell, cellIndex) => (
+      <StyledRow key={cellIndex}>
+        <RawResultCell
+          value={cell}
+          fileLinkPrefix={fileLinkPrefix}
+          sourceLocationPrefix={sourceLocationPrefix}
+        />
+      </StyledRow>
+    ))}
+  </>
+);

extensions/ql-vscode/src/view/variant-analysis/RawResultsTable.tsx

@@ -1,28 +1,13 @@
 import * as React from "react";
 import { useState } from "react";
 import { styled } from "styled-components";
-import { VSCodeLink } from "@vscode/webview-ui-toolkit/react";
-import {
-  CellValue,
-  RawResultSet,
-  ResultSetSchema,
-} from "../../common/bqrs-cli-types";
-import { tryGetRemoteLocation } from "../../common/bqrs-utils";
+import { RawResultSet, ResultSetSchema } from "../../common/bqrs-cli-types";
 import TextButton from "../common/TextButton";
-import { convertNonPrintableChars } from "../../common/text-utils";
-import { sendTelemetry, useTelemetryOnChange } from "../common/telemetry";
+import { useTelemetryOnChange } from "../common/telemetry";
+import { RawResultRow } from "./RawResultRow";
 
 const numOfResultsInContractedMode = 5;
 
-const StyledRow = styled.div`
-  border-color: var(--vscode-editor-snippetFinalTabstopHighlightBorder);
-  border-style: solid;
-  justify-content: center;
-  align-items: center;
-  padding: 0.4rem;
-  word-break: break-word;
-`;
-
 type TableContainerProps = {
   columnCount: number;
 };
@@ -40,60 +25,6 @@ const TableContainer = styled.div<TableContainerProps>`
   padding: 0.4rem;
 `;
 
-type CellProps = {
-  value: CellValue;
-  fileLinkPrefix: string;
-  sourceLocationPrefix: string;
-};
-
-const sendRawResultsLinkTelemetry = () => sendTelemetry("raw-results-link");
-
-const Cell = ({ value, fileLinkPrefix, sourceLocationPrefix }: CellProps) => {
-  switch (typeof value) {
-    case "string":
-    case "number":
-    case "boolean":
-      return <span>{convertNonPrintableChars(value.toString())}</span>;
-    case "object": {
-      const url = tryGetRemoteLocation(
-        value.url,
-        fileLinkPrefix,
-        sourceLocationPrefix,
-      );
-      const safeLabel = convertNonPrintableChars(value.label);
-      if (url) {
-        return (
-          <VSCodeLink onClick={sendRawResultsLinkTelemetry} href={url}>
-            {safeLabel}
-          </VSCodeLink>
-        );
-      } else {
-        return <span>{safeLabel}</span>;
-      }
-    }
-  }
-};
-
-type RowProps = {
-  row: CellValue[];
-  fileLinkPrefix: string;
-  sourceLocationPrefix: string;
-};
-
-const Row = ({ row, fileLinkPrefix, sourceLocationPrefix }: RowProps) => (
-  <>
-    {row.map((cell, cellIndex) => (
-      <StyledRow key={cellIndex}>
-        <Cell
-          value={cell}
-          fileLinkPrefix={fileLinkPrefix}
-          sourceLocationPrefix={sourceLocationPrefix}
-        />
-      </StyledRow>
-    ))}
-  </>
-);
-
 type RawResultsTableProps = {
   schema: ResultSetSchema;
   results: RawResultSet;
@@ -122,7 +53,7 @@ const RawResultsTable = ({
     <>
       <TableContainer columnCount={schema.columns.length}>
         {results.rows.slice(0, numOfResultsToShow).map((row, rowIndex) => (
-          <Row
+          <RawResultRow
             key={rowIndex}
             row={row}
             fileLinkPrefix={fileLinkPrefix}

extensions/ql-vscode/src/view/variant-analysis/__tests__/AnalyzedRepoItemContent.spec.tsx

@@ -71,6 +71,12 @@ describe(AnalyzedRepoItemContent.name, () => {
             {
               kind: "i",
             },
+            {
+              kind: "s",
+            },
+            {
+              kind: "b",
+            },
           ],
         },
         resultSet: {
@@ -81,9 +87,18 @@ describe(AnalyzedRepoItemContent.name, () => {
               {
                 kind: "i",
               },
+              {
+                kind: "s",
+              },
+              {
+                kind: "b",
+              },
             ],
           },
-          rows: [[60688]],
+          rows: [
+            [60688, "foo", true],
+            [5, "bar", false],
+          ],
         },
         fileLinkPrefix:
           "https://github.com/octodemo/hello-world-1/blob/59a2a6c7d9dde7a6ecb77c2f7e8197d6925c143b",
@@ -92,7 +107,12 @@ describe(AnalyzedRepoItemContent.name, () => {
       },
     });
 
-    expect(screen.getByText("60688")).toBeInTheDocument();
+    expect(screen.getByText("60,688")).toBeInTheDocument();
+    expect(screen.getByText("foo")).toBeInTheDocument();
+    expect(screen.getByText("true")).toBeInTheDocument();
+    expect(screen.getByText("5")).toBeInTheDocument();
+    expect(screen.getByText("bar")).toBeInTheDocument();
+    expect(screen.getByText("false")).toBeInTheDocument();
   });
 
   it("renders the failed state", () => {

extensions/ql-vscode/supported_cli_versions.json

@@ -1,5 +1,5 @@
 [
-  "v2.14.5",
+  "v2.14.6",
   "v2.13.5",
   "v2.12.7",
   "v2.11.6",

extensions/ql-vscode/test/__mocks__/model-editor/modelingStoreMock.ts

@@ -0,0 +1,33 @@
+import { mockedObject } from "../../vscode-tests/utils/mocking.helpers";
+import { ModelingStore } from "../../../src/model-editor/modeling-store";
+
+export function createMockModelingStore({
+  initializeStateForDb = jest.fn(),
+  getStateForActiveDb = jest.fn(),
+  onActiveDbChanged = jest.fn(),
+  onDbClosed = jest.fn(),
+  onMethodsChanged = jest.fn(),
+  onHideModeledMethodsChanged = jest.fn(),
+  onModeledMethodsChanged = jest.fn(),
+  onModifiedMethodsChanged = jest.fn(),
+}: {
+  initializeStateForDb?: ModelingStore["initializeStateForDb"];
+  getStateForActiveDb?: ModelingStore["getStateForActiveDb"];
+  onActiveDbChanged?: ModelingStore["onActiveDbChanged"];
+  onDbClosed?: ModelingStore["onDbClosed"];
+  onMethodsChanged?: ModelingStore["onMethodsChanged"];
+  onHideModeledMethodsChanged?: ModelingStore["onHideModeledMethodsChanged"];
+  onModeledMethodsChanged?: ModelingStore["onModeledMethodsChanged"];
+  onModifiedMethodsChanged?: ModelingStore["onModifiedMethodsChanged"];
+} = {}): ModelingStore {
+  return mockedObject<ModelingStore>({
+    initializeStateForDb,
+    getStateForActiveDb,
+    onActiveDbChanged,
+    onDbClosed,
+    onMethodsChanged,
+    onHideModeledMethodsChanged,
+    onModeledMethodsChanged,
+    onModifiedMethodsChanged,
+  });
+}

extensions/ql-vscode/test/factories/model-editor/modeled-method-factories.ts

@@ -0,0 +1,20 @@
+import { ModeledMethod } from "../../../src/model-editor/modeled-method";
+
+export function createModeledMethod(
+  data: Partial<ModeledMethod> = {},
+): ModeledMethod {
+  return {
+    libraryVersion: "1.6.0",
+    signature: "org.sql2o.Connection#createQuery(String)",
+    packageName: "org.sql2o",
+    typeName: "Connection",
+    methodName: "createQuery",
+    methodParameters: "(String)",
+    type: "sink",
+    input: "Argument[0]",
+    output: "",
+    kind: "jndi-injection",
+    provenance: "manual",
+    ...data,
+  };
+}

extensions/ql-vscode/test/unit-tests/common/split-buffer.test.ts

@@ -0,0 +1,84 @@
+import { LINE_ENDINGS, SplitBuffer } from "../../../src/common/split-stream";
+
+interface Chunk {
+  chunk: string;
+  lines: string[];
+}
+
+function checkLines(
+  buffer: SplitBuffer,
+  expectedLinesForChunk: string[],
+  chunkIndex: number | "end",
+): void {
+  expectedLinesForChunk.forEach((expectedLine, lineIndex) => {
+    const line = buffer.getNextLine();
+    const location = `[chunk ${chunkIndex}, line ${lineIndex}]: `;
+    expect(location + line).toEqual(location + expectedLine);
+  });
+  expect(buffer.getNextLine()).toBeUndefined();
+}
+
+function testSplitBuffer(chunks: Chunk[], endLines: string[]): void {
+  const buffer = new SplitBuffer(LINE_ENDINGS);
+  chunks.forEach((chunk, chunkIndex) => {
+    buffer.addChunk(Buffer.from(chunk.chunk, "utf-8"));
+    checkLines(buffer, chunk.lines, chunkIndex);
+  });
+  buffer.end();
+  checkLines(buffer, endLines, "end");
+}
+
+describe("split buffer", () => {
+  it("should handle a one-chunk string with no terminator", async () => {
+    // Won't return the line until we call `end()`.
+    testSplitBuffer([{ chunk: "some text", lines: [] }], ["some text"]);
+  });
+
+  it("should handle a one-chunk string with a one-byte terminator", async () => {
+    // Won't return the line until we call `end()` because the actual terminator is shorter than the
+    // longest terminator.
+    testSplitBuffer([{ chunk: "some text\n", lines: [] }], ["some text"]);
+  });
+
+  it("should handle a one-chunk string with a two-byte terminator", async () => {
+    testSplitBuffer([{ chunk: "some text\r\n", lines: ["some text"] }], []);
+  });
+
+  it("should handle a multi-chunk string with terminators at the end of each chunk", async () => {
+    testSplitBuffer(
+      [
+        { chunk: "first line\n", lines: [] }, // Waiting for second potential terminator byte
+        { chunk: "second line\r", lines: ["first line"] }, // Waiting for second potential terminator byte
+        { chunk: "third line\r\n", lines: ["second line", "third line"] }, // No wait, because we're at the end
+      ],
+      [],
+    );
+  });
+
+  it("should handle a multi-chunk string with terminators at random offsets", async () => {
+    testSplitBuffer(
+      [
+        { chunk: "first line\nsecond", lines: ["first line"] },
+        {
+          chunk: " line\rthird line",
+          lines: ["second line"],
+        },
+        { chunk: "\r\n", lines: ["third line"] },
+      ],
+      [],
+    );
+  });
+
+  it("should handle a terminator split between chunks", async () => {
+    testSplitBuffer(
+      [
+        { chunk: "first line\r", lines: [] },
+        {
+          chunk: "\nsecond line",
+          lines: ["first line"],
+        },
+      ],
+      ["second line"],
+    );
+  });
+});

extensions/ql-vscode/test/unit-tests/model-editor/shared/modeled-percentage.spec.ts

@@ -1,5 +1,5 @@
 import { calculateModeledPercentage } from "../../../../src/model-editor/shared/modeled-percentage";
-import { createMethod } from "../../../factories/data-extension/method-factories";
+import { createMethod } from "../../../factories/model-editor/method-factories";
 
 describe("calculateModeledPercentage", () => {
   it("when there are no external API usages", () => {