Merge remote-tracking branch 'origin/main' into robertbrignull/upgrade_msw

2023-10-03 09:36:03 +02:00
parent e15b7681db 558b9329c5
commit 8107bf7cb1
26 changed files with 664 additions and 196 deletions
--- a/docs/test-plan.md
+++ b/docs/test-plan.md
@@ -173,6 +173,8 @@ Note that this test requires the feature flag: `codeQL.model.llmGeneration`

 #### Test Case 4: Model as dependency

+Note that this test requires the feature flag: `codeQL.model.flowGeneration`
+
 1. Click "Model as dependency"
   - Check that grouping are now per package (e.g. `com.alipay.sofa.rraft.option` or `com.google.protobuf`)
 2. Click "Generate".
--- a/extensions/ql-vscode/CHANGELOG.md
+++ b/extensions/ql-vscode/CHANGELOG.md
@@ -2,6 +2,15 @@

 ## [UNRELEASED]

+- Fix a bug where the query server was restarted twice after configuration changes. [#2884](https://github.com/github/vscode-codeql/pull/2884).
+
+## 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)
--- a/extensions/ql-vscode/package-lock.json
+++ b/extensions/ql-vscode/package-lock.json
@@ -1,12 +1,12 @@
 {
  "name": "vscode-codeql",
-  "version": "1.9.1",
+  "version": "1.9.2",
  "lockfileVersion": 2,
  "requires": true,
  "packages": {
    "": {
      "name": "vscode-codeql",
-      "version": "1.9.1",
+      "version": "1.9.2",
      "hasInstallScript": true,
      "license": "MIT",
      "dependencies": {
--- a/extensions/ql-vscode/package.json
+++ b/extensions/ql-vscode/package.json
@@ -4,7 +4,7 @@
  "description": "CodeQL for Visual Studio Code",
  "author": "GitHub",
  "private": true,
-  "version": "1.9.1",
+  "version": "1.9.2",
  "publisher": "GitHub",
  "license": "MIT",
  "icon": "media/VS-marketplace-CodeQL-icon.png",
@@ -1077,7 +1077,7 @@
        },
        {
          "submenu": "codeQLDatabases.languages",
-          "when": "view == codeQLDatabases && config.codeQL.canary",
+          "when": "view == codeQLDatabases && config.codeQL.canary && config.codeQL.showLanguageFilter",
          "group": "2_databases@0"
        },
        {
@@ -1870,11 +1870,11 @@
      "codeQLDatabases.languages": [
        {
          "command": "codeQLDatabases.displayAllLanguages",
-          "when": "codeQLDatabases.languageFilter != All"
+          "when": "codeQLDatabases.languageFilter"
        },
        {
          "command": "codeQLDatabases.displayAllLanguagesSelected",
-          "when": "codeQLDatabases.languageFilter == All"
+          "when": "!codeQLDatabases.languageFilter"
        },
        {
          "command": "codeQLDatabases.displayCpp",
--- a/extensions/ql-vscode/src/common/interface-types.ts
+++ b/extensions/ql-vscode/src/common/interface-types.ts
@@ -598,8 +598,7 @@ export type FromModelEditorMessage =
  | SetModeledMethodMessage;

 export type FromMethodModelingMessage =
-  | TelemetryMessage
-  | UnhandledErrorMessage
+  | CommonFromViewMessages
  | SetModeledMethodMessage;

 interface SetMethodMessage {
--- a/extensions/ql-vscode/src/common/query-language.ts
+++ b/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;
+}
--- a/extensions/ql-vscode/src/common/vscode/abstract-webview-view-provider.ts
+++ b/extensions/ql-vscode/src/common/vscode/abstract-webview-view-provider.ts
@@ -0,0 +1,85 @@
+import * as vscode from "vscode";
+import { Uri, WebviewViewProvider } from "vscode";
+import { WebviewKind, WebviewMessage, getHtmlForWebview } from "./webview-html";
+import { Disposable } from "../disposable-object";
+import { App } from "../app";
+
+export abstract class AbstractWebviewViewProvider<
+  ToMessage extends WebviewMessage,
+  FromMessage extends WebviewMessage,
+> implements WebviewViewProvider
+{
+  protected webviewView: vscode.WebviewView | undefined = undefined;
+  private disposables: Disposable[] = [];
+
+  constructor(
+    private readonly app: App,
+    private readonly webviewKind: WebviewKind,
+  ) {}
+
+  /**
+   * This is called when a view first becomes visible. This may happen when the view is
+   * first loaded or when the user hides and then shows a view again.
+   */
+  public resolveWebviewView(
+    webviewView: vscode.WebviewView,
+    _context: vscode.WebviewViewResolveContext,
+    _token: vscode.CancellationToken,
+  ) {
+    webviewView.webview.options = {
+      enableScripts: true,
+      localResourceRoots: [Uri.file(this.app.extensionPath)],
+    };
+
+    const html = getHtmlForWebview(
+      this.app,
+      webviewView.webview,
+      this.webviewKind,
+      {
+        allowInlineStyles: true,
+        allowWasmEval: false,
+      },
+    );
+
+    webviewView.webview.html = html;
+
+    this.webviewView = webviewView;
+
+    webviewView.webview.onDidReceiveMessage(async (msg) => this.onMessage(msg));
+    webviewView.onDidDispose(() => this.dispose());
+  }
+
+  protected get isShowingView() {
+    return this.webviewView?.visible ?? false;
+  }
+
+  protected async postMessage(msg: ToMessage): Promise<void> {
+    await this.webviewView?.webview.postMessage(msg);
+  }
+
+  protected dispose() {
+    while (this.disposables.length > 0) {
+      const disposable = this.disposables.pop()!;
+      disposable.dispose();
+    }
+
+    this.webviewView = undefined;
+  }
+
+  protected push<T extends Disposable>(obj: T): T {
+    if (obj !== undefined) {
+      this.disposables.push(obj);
+    }
+    return obj;
+  }
+
+  protected abstract onMessage(msg: FromMessage): Promise<void>;
+
+  /**
+   * This is called when a view first becomes visible. This may happen when the view is
+   * first loaded or when the user hides and then shows a view again.
+   */
+  protected onWebViewLoaded(): void {
+    // Do nothing by default.
+  }
+}
--- a/extensions/ql-vscode/src/common/vscode/abstract-webview.ts
+++ b/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;
  }
 }
--- a/extensions/ql-vscode/src/databases/local-databases-ui.ts
+++ b/extensions/ql-vscode/src/databases/local-databases-ui.ts
@@ -51,7 +51,8 @@ import {
  createMultiSelectionCommand,
  createSingleSelectionCommand,
 } from "../common/vscode/selection-commands";
-import { QueryLanguage } from "../common/query-language";
+import { QueryLanguage, tryGetQueryLanguage } from "../common/query-language";
+import { LanguageContextStore } from "../language-context-store";

 enum SortOrder {
  NameAsc = "NameAsc",
@@ -60,8 +61,6 @@ enum SortOrder {
  DateAddedDesc = "DateAddedDesc",
 }

-type LanguageFilter = QueryLanguage | "All";
-
 /**
 * Tree data provider for the databases view.
 */
@@ -70,14 +69,16 @@ class DatabaseTreeDataProvider
  implements TreeDataProvider<DatabaseItem>
 {
  private _sortOrder = SortOrder.NameAsc;
-  private _languageFilter = "All" as LanguageFilter;

  private readonly _onDidChangeTreeData = this.push(
    new EventEmitter<DatabaseItem | undefined>(),
  );
  private currentDatabaseItem: DatabaseItem | undefined;

-  constructor(private databaseManager: DatabaseManager) {
+  constructor(
+    private databaseManager: DatabaseManager,
+    private languageContext: LanguageContextStore,
+  ) {
    super();

    this.currentDatabaseItem = databaseManager.currentDatabaseItem;
@@ -92,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> {
@@ -137,11 +143,9 @@ class DatabaseTreeDataProvider
    if (element === undefined) {
      // Filter items by language
      const displayItems = this.databaseManager.databaseItems.filter((item) => {
-        if (this.languageFilter === "All") {
-          return true;
-        } else {
-          return item.language === this.languageFilter;
-        }
+        return this.languageContext.shouldInclude(
+          tryGetQueryLanguage(item.language),
+        );
      });

      // Sort items
@@ -178,15 +182,6 @@ class DatabaseTreeDataProvider
    this._sortOrder = newSortOrder;
    this._onDidChangeTreeData.fire(undefined);
  }
-
-  public get languageFilter() {
-    return this._languageFilter;
-  }
-
-  public set languageFilter(newLanguageFilter: LanguageFilter) {
-    this._languageFilter = newLanguageFilter;
-    this._onDidChangeTreeData.fire(undefined);
-  }
 }

 /** Gets the first element in the given list, if any, or undefined if the list is empty or undefined. */
@@ -223,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,
@@ -230,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", {
@@ -269,7 +265,7 @@ export class DatabaseUI extends DisposableObject {
      "codeQLDatabases.sortByName": this.handleSortByName.bind(this),
      "codeQLDatabases.sortByDateAdded": this.handleSortByDateAdded.bind(this),
      "codeQLDatabases.displayAllLanguages":
-        this.handleChangeLanguageFilter.bind(this, "All"),
+        this.handleClearLanguageFilter.bind(this),
      "codeQLDatabases.displayCpp": this.handleChangeLanguageFilter.bind(
        this,
        QueryLanguage.Cpp,
@@ -303,7 +299,7 @@ export class DatabaseUI extends DisposableObject {
        QueryLanguage.Swift,
      ),
      "codeQLDatabases.displayAllLanguagesSelected":
-        this.handleChangeLanguageFilter.bind(this, "All"),
+        this.handleClearLanguageFilter.bind(this),
      "codeQLDatabases.displayCppSelected":
        this.handleChangeLanguageFilter.bind(this, QueryLanguage.Cpp),
      "codeQLDatabases.displayCsharpSelected":
@@ -612,13 +608,12 @@ export class DatabaseUI extends DisposableObject {
    }
  }

-  private async handleChangeLanguageFilter(languageFilter: LanguageFilter) {
-    this.treeDataProvider.languageFilter = languageFilter;
-    await this.app.commands.execute(
-      "setContext",
-      "codeQLDatabases.languageFilter",
-      languageFilter,
-    );
+  private async handleClearLanguageFilter() {
+    await this.languageContext.clearLanguageContext();
+  }
+
+  private async handleChangeLanguageFilter(languageFilter: QueryLanguage) {
+    await this.languageContext.setLanguageContext(languageFilter);
  }

  private async handleUpgradeCurrentDatabase(): Promise<void> {
--- a/extensions/ql-vscode/src/extension.ts
+++ b/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,17 +775,22 @@ 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,
  );
  ctx.subscriptions.push(databaseUI);

-  QueriesModule.initialize(app, cliServer);
+  QueriesModule.initialize(app, languageContext, cliServer);

  void extLogger.log("Initializing evaluator log viewer.");
  const evalLogViewer = new EvalLogViewer();
--- a/extensions/ql-vscode/src/language-context-store.ts
+++ b/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;
+  }
+}
--- a/extensions/ql-vscode/src/local-queries/results-view.ts
+++ b/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 = [];
+  }
 }
--- a/extensions/ql-vscode/src/model-editor/method-modeling/method-modeling-panel.ts
+++ b/extensions/ql-vscode/src/model-editor/method-modeling/method-modeling-panel.ts
@@ -12,7 +12,6 @@ export class MethodModelingPanel extends DisposableObject {
    super();

    this.provider = new MethodModelingViewProvider(app, modelingStore);
-    this.push(this.provider);
    this.push(
      window.registerWebviewViewProvider(
        MethodModelingViewProvider.viewType,
--- a/extensions/ql-vscode/src/model-editor/method-modeling/method-modeling-view-provider.ts
+++ b/extensions/ql-vscode/src/model-editor/method-modeling/method-modeling-view-provider.ts
@@ -1,82 +1,51 @@
-import * as vscode from "vscode";
-import { Uri, WebviewViewProvider } from "vscode";
-import { getHtmlForWebview } from "../../common/vscode/webview-html";
-import { FromMethodModelingMessage } from "../../common/interface-types";
+import {
+  FromMethodModelingMessage,
+  ToMethodModelingMessage,
+} from "../../common/interface-types";
 import { telemetryListener } from "../../common/vscode/telemetry";
 import { showAndLogExceptionWithTelemetry } from "../../common/logging/notifications";
 import { extLogger } from "../../common/logging/vscode/loggers";
 import { App } from "../../common/app";
 import { redactableError } from "../../common/errors";
 import { Method } from "../method";
-import { DisposableObject } from "../../common/disposable-object";
 import { ModelingStore } from "../modeling-store";
+import { AbstractWebviewViewProvider } from "../../common/vscode/abstract-webview-view-provider";

-export class MethodModelingViewProvider
-  extends DisposableObject
-  implements WebviewViewProvider
-{
+export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
+  ToMethodModelingMessage,
+  FromMethodModelingMessage
+> {
  public static readonly viewType = "codeQLMethodModeling";

-  private webviewView: vscode.WebviewView | undefined = undefined;
-
  private method: Method | undefined = undefined;

  constructor(
-    private readonly app: App,
+    app: App,
    private readonly modelingStore: ModelingStore,
  ) {
-    super();
+    super(app, "method-modeling");
  }

-  /**
-   * This is called when a view first becomes visible. This may happen when the view is
-   * first loaded or when the user hides and then shows a view again.
-   */
-  public resolveWebviewView(
-    webviewView: vscode.WebviewView,
-    _context: vscode.WebviewViewResolveContext,
-    _token: vscode.CancellationToken,
-  ) {
-    webviewView.webview.options = {
-      enableScripts: true,
-      localResourceRoots: [Uri.file(this.app.extensionPath)],
-    };
-
-    const html = getHtmlForWebview(
-      this.app,
-      webviewView.webview,
-      "method-modeling",
-      {
-        allowInlineStyles: true,
-        allowWasmEval: false,
-      },
-    );
-
-    webviewView.webview.html = html;
-
-    webviewView.webview.onDidReceiveMessage(async (msg) => this.onMessage(msg));
-
-    this.webviewView = webviewView;
-
-    this.setInitialState(webviewView);
+  protected override onWebViewLoaded(): void {
+    this.setInitialState();
    this.registerToModelingStoreEvents();
  }

  public async setMethod(method: Method): Promise<void> {
    this.method = method;

-    if (this.webviewView) {
-      await this.webviewView.webview.postMessage({
+    if (this.isShowingView) {
+      await this.postMessage({
        t: "setMethod",
        method,
      });
    }
  }

-  private setInitialState(webviewView: vscode.WebviewView): void {
+  private setInitialState(): void {
    const selectedMethod = this.modelingStore.getSelectedMethodDetails();
    if (selectedMethod) {
-      void webviewView.webview.postMessage({
+      void this.postMessage({
        t: "setSelectedMethod",
        method: selectedMethod.method,
        modeledMethod: selectedMethod.modeledMethod,
@@ -85,8 +54,28 @@ export class MethodModelingViewProvider
    }
  }

-  private async onMessage(msg: FromMethodModelingMessage): Promise<void> {
+  protected override async onMessage(
+    msg: FromMethodModelingMessage,
+  ): Promise<void> {
    switch (msg.t) {
+      case "viewLoaded":
+        this.onWebViewLoaded();
+        break;
+
+      case "telemetry":
+        telemetryListener?.sendUIInteraction(msg.action);
+        break;
+
+      case "unhandledError":
+        void showAndLogExceptionWithTelemetry(
+          extLogger,
+          telemetryListener,
+          redactableError(
+            msg.error,
+          )`Unhandled error in method modeling view: ${msg.error.message}`,
+        );
+        break;
+
      case "setModeledMethod": {
        const activeState = this.modelingStore.getStateForActiveDb();
        if (!activeState) {
@@ -98,56 +87,48 @@ export class MethodModelingViewProvider
        );
        break;
      }
-
-      case "telemetry": {
-        telemetryListener?.sendUIInteraction(msg.action);
-        break;
-      }
-      case "unhandledError":
-        void showAndLogExceptionWithTelemetry(
-          extLogger,
-          telemetryListener,
-          redactableError(
-            msg.error,
-          )`Unhandled error in method modeling view: ${msg.error.message}`,
-        );
-        break;
    }
  }

  private registerToModelingStoreEvents(): void {
-    this.modelingStore.onModeledMethodsChanged(async (e) => {
-      if (this.webviewView && e.isActiveDb) {
-        const modeledMethod = e.modeledMethods[this.method?.signature ?? ""];
-        if (modeledMethod) {
+    this.push(
+      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.push(
+      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: "setModeledMethod",
-            method: modeledMethod,
+            t: "setMethodModified",
+            isModified,
          });
        }
-      }
-    });
+      }),
+    );

-    this.modelingStore.onModifiedMethodsChanged(async (e) => {
-      if (this.webviewView && e.isActiveDb && this.method) {
-        const isModified = e.modifiedMethods.has(this.method.signature);
-        await this.webviewView.webview.postMessage({
-          t: "setMethodModified",
-          isModified,
-        });
-      }
-    });
-
-    this.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,
-        });
-      }
-    });
+    this.push(
+      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,
+          });
+        }
+      }),
+    );
  }
 }
--- a/extensions/ql-vscode/src/model-editor/methods-usage/methods-usage-data-provider.ts
+++ b/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) {
--- a/extensions/ql-vscode/src/model-editor/methods-usage/methods-usage-panel.ts
+++ b/extensions/ql-vscode/src/model-editor/methods-usage/methods-usage-panel.ts
@@ -8,6 +8,7 @@ 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;
@@ -33,8 +34,16 @@ export class MethodsUsagePanel extends DisposableObject {
    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;
@@ -73,6 +82,14 @@ export class MethodsUsagePanel extends DisposableObject {
        }
      }),
    );
+
+    this.push(
+      this.modelingStore.onModifiedMethodsChanged(async (event) => {
+        if (event.isActiveDb) {
+          await this.handleStateChangeEvent();
+        }
+      }),
+    );
  }

  private async handleStateChangeEvent(): Promise<void> {
@@ -82,6 +99,8 @@ export class MethodsUsagePanel extends DisposableObject {
        activeState.methods,
        activeState.databaseItem,
        activeState.hideModeledMethods,
+        activeState.modeledMethods,
+        activeState.modifiedMethodSignatures,
      );
    }
  }
--- a/extensions/ql-vscode/src/model-editor/modeling-store.ts
+++ b/extensions/ql-vscode/src/model-editor/modeling-store.ts
@@ -158,7 +158,7 @@ export class ModelingStore extends DisposableObject {
    const dbState = this.getState(dbItem);
    const dbUri = dbItem.databaseUri.toString();

-    dbState.methods = methods;
+    dbState.methods = [...methods];

    this.onMethodsChangedEventEmitter.fire({
      methods,
@@ -204,13 +204,15 @@ export class ModelingStore extends DisposableObject {
    methods: Record<string, ModeledMethod>,
  ) {
    this.changeModeledMethods(dbItem, (state) => {
-      state.modeledMethods = methods;
+      state.modeledMethods = { ...methods };
    });
  }

  public updateModeledMethod(dbItem: DatabaseItem, method: ModeledMethod) {
    this.changeModeledMethods(dbItem, (state) => {
-      state.modeledMethods[method.signature] = method;
+      const newModeledMethods = { ...state.modeledMethods };
+      newModeledMethods[method.signature] = method;
+      state.modeledMethods = newModeledMethods;
    });
  }

@@ -219,7 +221,7 @@ export class ModelingStore extends DisposableObject {
    methodSignatures: Set<string>,
  ) {
    this.changeModifiedMethods(dbItem, (state) => {
-      state.modifiedMethodSignatures = methodSignatures;
+      state.modifiedMethodSignatures = new Set(methodSignatures);
    });
  }

@@ -228,9 +230,11 @@ export class ModelingStore extends DisposableObject {
    methodSignatures: Iterable<string>,
  ) {
    this.changeModifiedMethods(dbItem, (state) => {
-      for (const signature of methodSignatures) {
-        state.modifiedMethodSignatures.add(signature);
-      }
+      const newModifiedMethods = new Set([
+        ...state.modifiedMethodSignatures,
+        ...methodSignatures,
+      ]);
+      state.modifiedMethodSignatures = newModifiedMethods;
    });
  }

@@ -243,9 +247,11 @@ export class ModelingStore extends DisposableObject {
    methodSignatures: string[],
  ) {
    this.changeModifiedMethods(dbItem, (state) => {
-      methodSignatures.forEach((signature) => {
-        state.modifiedMethodSignatures.delete(signature);
-      });
+      const newModifiedMethods = Array.from(
+        state.modifiedMethodSignatures,
+      ).filter((s) => !methodSignatures.includes(s));
+
+      state.modifiedMethodSignatures = new Set(newModifiedMethods);
    });
  }

--- a/extensions/ql-vscode/src/queries-panel/queries-module.ts
+++ b/extensions/ql-vscode/src/queries-panel/queries-module.ts
@@ -6,6 +6,7 @@ import { DisposableObject } from "../common/disposable-object";
 import { QueriesPanel } from "./queries-panel";
 import { QueryDiscovery } from "./query-discovery";
 import { QueryPackDiscovery } from "./query-pack-discovery";
+import { LanguageContextStore } from "../language-context-store";

 export class QueriesModule extends DisposableObject {
  private queriesPanel: QueriesPanel | undefined;
@@ -16,16 +17,21 @@ export class QueriesModule extends DisposableObject {

  public static initialize(
    app: App,
+    languageContext: LanguageContextStore,
    cliServer: CodeQLCliServer,
  ): QueriesModule {
    const queriesModule = new QueriesModule(app);
    app.subscriptions.push(queriesModule);

-    queriesModule.initialize(app, cliServer);
+    queriesModule.initialize(app, languageContext, cliServer);
    return queriesModule;
  }

-  private initialize(app: App, cliServer: CodeQLCliServer): void {
+  private initialize(
+    app: App,
+    langauageContext: LanguageContextStore,
+    cliServer: CodeQLCliServer,
+  ): void {
    // Currently, we only want to expose the new panel when we are in canary mode
    // and the user has enabled the "Show queries panel" flag.
    if (!isCanary() || !showQueriesPanel()) {
@@ -38,8 +44,9 @@ export class QueriesModule extends DisposableObject {
    void queryPackDiscovery.initialRefresh();

    const queryDiscovery = new QueryDiscovery(
-      app.environment,
+      app,
      queryPackDiscovery,
+      langauageContext,
    );
    this.push(queryDiscovery);
    void queryDiscovery.initialRefresh();
--- a/extensions/ql-vscode/src/queries-panel/query-discovery.ts
+++ b/extensions/ql-vscode/src/queries-panel/query-discovery.ts
@@ -1,6 +1,6 @@
 import { dirname, basename, normalize, relative } from "path";
 import { Event } from "vscode";
-import { EnvironmentContext } from "../common/app";
+import { App } from "../common/app";
 import {
  FileTreeDirectory,
  FileTreeLeaf,
@@ -11,6 +11,8 @@ import { FilePathDiscovery } from "../common/vscode/file-path-discovery";
 import { containsPath } from "../common/files";
 import { getOnDiskWorkspaceFoldersObjects } from "../common/vscode/workspace-folders";
 import { QueryLanguage } from "../common/query-language";
+import { LanguageContextStore } from "../language-context-store";
+import { AppEvent, AppEventEmitter } from "../common/events";

 const QUERY_FILE_EXTENSION = ".ql";

@@ -31,24 +33,36 @@ export class QueryDiscovery
  extends FilePathDiscovery<Query>
  implements QueryDiscoverer
 {
+  public readonly onDidChangeQueries: AppEvent<void>;
+  private readonly onDidChangeQueriesEmitter: AppEventEmitter<void>;
+
  constructor(
-    private readonly env: EnvironmentContext,
+    private readonly app: App,
    private readonly queryPackDiscovery: QueryPackDiscoverer,
+    private readonly languageContext: LanguageContextStore,
  ) {
    super("Query Discovery", `**/*${QUERY_FILE_EXTENSION}`);

+    // Set up event emitters
+    this.onDidChangeQueriesEmitter = this.push(app.createEventEmitter<void>());
+    this.onDidChangeQueries = this.onDidChangeQueriesEmitter.event;
+
+    // Handlers
    this.push(
      this.queryPackDiscovery.onDidChangeQueryPacks(
        this.recomputeAllData.bind(this),
      ),
    );
-  }
-
-  /**
-   * Event that fires when the set of queries in the workspace changes.
-   */
-  public get onDidChangeQueries(): Event<void> {
-    return this.onDidChangePathData;
+    this.push(
+      this.onDidChangePathData(() => {
+        this.onDidChangeQueriesEmitter.fire();
+      }),
+    );
+    this.push(
+      this.languageContext.onLanguageContextChanged(() => {
+        this.onDidChangeQueriesEmitter.fire();
+      }),
+    );
  }

  /**
@@ -64,8 +78,10 @@ export class QueryDiscovery

    const roots = [];
    for (const workspaceFolder of getOnDiskWorkspaceFoldersObjects()) {
-      const queriesInRoot = pathData.filter((query) =>
-        containsPath(workspaceFolder.uri.fsPath, query.path),
+      const queriesInRoot = pathData.filter(
+        (query) =>
+          containsPath(workspaceFolder.uri.fsPath, query.path) &&
+          this.languageContext.shouldInclude(query.language),
      );
      if (queriesInRoot.length === 0) {
        continue;
@@ -73,7 +89,7 @@ export class QueryDiscovery
      const root = new FileTreeDirectory<string>(
        workspaceFolder.uri.fsPath,
        workspaceFolder.name,
-        this.env,
+        this.app.environment,
      );
      for (const query of queriesInRoot) {
        const dirName = dirname(normalize(relative(root.path, query.path)));
--- a/extensions/ql-vscode/src/query-server/server-process.ts
+++ b/extensions/ql-vscode/src/query-server/server-process.ts
@@ -26,6 +26,7 @@ export class ServerProcess implements Disposable {
    this.connection.end();
    this.child.stdin!.end();
    this.child.stderr!.destroy();
+    this.child.removeAllListeners();
    // TODO kill the process if it doesn't terminate after a certain time limit.

    // On Windows, we usually have to terminate the process before closing its stdout.
--- a/extensions/ql-vscode/src/view/method-modeling/MethodModeling.tsx
+++ b/extensions/ql-vscode/src/view/method-modeling/MethodModeling.tsx
@@ -27,10 +27,12 @@ const DependencyContainer = styled.div`
  flex-direction: row;
  align-items: center;
  gap: 0.5em;
-  background-color: var(--vscode-textBlockQuote-background);
+  background-color: var(--vscode-editor-background);
+  border: 0.05rem solid var(--vscode-panelSection-border);
  border-radius: 0.3rem;
-  border-color: var(--vscode-textBlockQuote-border);
  padding: 0.5rem;
+  word-wrap: break-word;
+  word-break: break-all;
 `;

 export type MethodModelingProps = {
--- a/extensions/ql-vscode/test/mocks/appMock.ts
+++ b/extensions/ql-vscode/test/mocks/appMock.ts
@@ -56,7 +56,7 @@ class MockAppEventEmitter<T> implements AppEventEmitter<T> {

  constructor() {
    this.event = () => {
-      return {} as Disposable;
+      return new MockAppEvent();
    };
  }

@@ -69,7 +69,17 @@ class MockAppEventEmitter<T> implements AppEventEmitter<T> {
  }
 }

-export function createMockEnvironmentContext(): EnvironmentContext {
+class MockAppEvent implements Disposable {
+  public fire(): void {
+    // no-op
+  }
+
+  public dispose() {
+    // no-op
+  }
+}
+
+function createMockEnvironmentContext(): EnvironmentContext {
  return {
    language: "en-US",
  };
--- a/extensions/ql-vscode/test/vscode-tests/minimal-workspace/queries-panel/query-discovery.test.ts
+++ b/extensions/ql-vscode/test/vscode-tests/minimal-workspace/queries-panel/query-discovery.test.ts
@@ -3,8 +3,8 @@ import {
  QueryDiscovery,
  QueryPackDiscoverer,
 } from "../../../../src/queries-panel/query-discovery";
-import { createMockEnvironmentContext } from "../../../__mocks__/appMock";
-import { dirname, join } from "path";
+import { createMockApp } from "../../../__mocks__/appMock";
+import { basename, dirname, join } from "path";
 import * as tmp from "tmp";
 import {
  FileTreeDirectory,
@@ -13,6 +13,7 @@ import {
 import { mkdirSync, writeFileSync } from "fs";
 import { QueryLanguage } from "../../../../src/common/query-language";
 import { sleep } from "../../../../src/common/time";
+import { LanguageContextStore } from "../../../../src/language-context-store";

 describe("Query pack discovery", () => {
  let tmpDir: string;
@@ -20,7 +21,10 @@ describe("Query pack discovery", () => {

  let workspacePath: string;

-  const env = createMockEnvironmentContext();
+  const app = createMockApp({});
+  const env = app.environment;
+
+  const languageContext = new LanguageContextStore(app);

  const onDidChangeQueryPacks = new EventEmitter<void>();
  let queryPackDiscoverer: QueryPackDiscoverer;
@@ -45,7 +49,7 @@ describe("Query pack discovery", () => {
      getLanguageForQueryFile: () => QueryLanguage.Java,
      onDidChangeQueryPacks: onDidChangeQueryPacks.event,
    };
-    discovery = new QueryDiscovery(env, queryPackDiscoverer);
+    discovery = new QueryDiscovery(app, queryPackDiscoverer, languageContext);
  });

  afterEach(() => {
@@ -160,6 +164,52 @@ describe("Query pack discovery", () => {
        ]),
      ]);
    });
+
+    it("should respect the language context filter", async () => {
+      makeTestFile(join(workspacePath, "query1.ql"));
+      makeTestFile(join(workspacePath, "query2.ql"));
+
+      queryPackDiscoverer.getLanguageForQueryFile = (path) => {
+        if (basename(path) === "query1.ql") {
+          return QueryLanguage.Java;
+        } else {
+          return QueryLanguage.Python;
+        }
+      };
+
+      await discovery.initialRefresh();
+
+      // Set the language to python-only
+      await languageContext.setLanguageContext(QueryLanguage.Python);
+
+      expect(discovery.buildQueryTree()).toEqual([
+        new FileTreeDirectory(workspacePath, "workspace", env, [
+          new FileTreeLeaf(
+            join(workspacePath, "query2.ql"),
+            "query2.ql",
+            "python",
+          ),
+        ]),
+      ]);
+
+      // Clear the language context filter
+      await languageContext.clearLanguageContext();
+
+      expect(discovery.buildQueryTree()).toEqual([
+        new FileTreeDirectory(workspacePath, "workspace", env, [
+          new FileTreeLeaf(
+            join(workspacePath, "query1.ql"),
+            "query1.ql",
+            "java",
+          ),
+          new FileTreeLeaf(
+            join(workspacePath, "query2.ql"),
+            "query2.ql",
+            "python",
+          ),
+        ]),
+      ]);
+    });
  });

  describe("recomputeAllQueryLanguages", () => {
--- a/extensions/ql-vscode/test/vscode-tests/no-workspace/databases/local-databases-ui.test.ts
+++ b/extensions/ql-vscode/test/vscode-tests/no-workspace/databases/local-databases-ui.test.ts
@@ -99,6 +99,11 @@ describe("local-databases-ui", () => {
          /**/
        },
      } as any,
+      {
+        onLanguageContextChanged: () => {
+          /**/
+        },
+      } as any,
      {} as any,
      storageDir,
      storageDir,
--- a/extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/methods-usage/methods-usage-data-provider.test.ts
+++ b/extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/methods-usage/methods-usage-data-provider.test.ts
@@ -7,6 +7,7 @@ import {
  createUsage,
 } from "../../../../factories/model-editor/method-factories";
 import { mockedObject } from "../../../utils/mocking.helpers";
+import { ModeledMethod } from "../../../../../src/model-editor/modeled-method";

 describe("MethodsUsageDataProvider", () => {
  const mockCliServer = mockedObject<CodeQLCliServer>({});
@@ -19,17 +20,31 @@ describe("MethodsUsageDataProvider", () => {
  describe("setState", () => {
    const hideModeledMethods = false;
    const methods: Method[] = [];
+    const modeledMethods: Record<string, ModeledMethod> = {};
+    const modifiedMethodSignatures: Set<string> = new Set();
    const dbItem = mockedObject<DatabaseItem>({
      getSourceLocationPrefix: () => "test",
    });

    it("should not emit onDidChangeTreeData event when state has not changed", async () => {
-      await dataProvider.setState(methods, dbItem, hideModeledMethods);
+      await dataProvider.setState(
+        methods,
+        dbItem,
+        hideModeledMethods,
+        modeledMethods,
+        modifiedMethodSignatures,
+      );

      const onDidChangeTreeDataListener = jest.fn();
      dataProvider.onDidChangeTreeData(onDidChangeTreeDataListener);

-      await dataProvider.setState(methods, dbItem, hideModeledMethods);
+      await dataProvider.setState(
+        methods,
+        dbItem,
+        hideModeledMethods,
+        modeledMethods,
+        modifiedMethodSignatures,
+      );

      expect(onDidChangeTreeDataListener).not.toHaveBeenCalled();
    });
@@ -37,12 +52,24 @@ describe("MethodsUsageDataProvider", () => {
    it("should emit onDidChangeTreeData event when methods has changed", async () => {
      const methods2: Method[] = [];

-      await dataProvider.setState(methods, dbItem, hideModeledMethods);
+      await dataProvider.setState(
+        methods,
+        dbItem,
+        hideModeledMethods,
+        modeledMethods,
+        modifiedMethodSignatures,
+      );

      const onDidChangeTreeDataListener = jest.fn();
      dataProvider.onDidChangeTreeData(onDidChangeTreeDataListener);

-      await dataProvider.setState(methods2, dbItem, hideModeledMethods);
+      await dataProvider.setState(
+        methods2,
+        dbItem,
+        hideModeledMethods,
+        modeledMethods,
+        modifiedMethodSignatures,
+      );

      expect(onDidChangeTreeDataListener).toHaveBeenCalledTimes(1);
    });
@@ -52,23 +79,97 @@ describe("MethodsUsageDataProvider", () => {
        getSourceLocationPrefix: () => "test",
      });

-      await dataProvider.setState(methods, dbItem, hideModeledMethods);
+      await dataProvider.setState(
+        methods,
+        dbItem,
+        hideModeledMethods,
+        modeledMethods,
+        modifiedMethodSignatures,
+      );

      const onDidChangeTreeDataListener = jest.fn();
      dataProvider.onDidChangeTreeData(onDidChangeTreeDataListener);

-      await dataProvider.setState(methods, dbItem2, hideModeledMethods);
+      await dataProvider.setState(
+        methods,
+        dbItem2,
+        hideModeledMethods,
+        modeledMethods,
+        modifiedMethodSignatures,
+      );

      expect(onDidChangeTreeDataListener).toHaveBeenCalledTimes(1);
    });

    it("should emit onDidChangeTreeData event when hideModeledMethods has changed", async () => {
-      await dataProvider.setState(methods, dbItem, hideModeledMethods);
+      await dataProvider.setState(
+        methods,
+        dbItem,
+        hideModeledMethods,
+        modeledMethods,
+        modifiedMethodSignatures,
+      );

      const onDidChangeTreeDataListener = jest.fn();
      dataProvider.onDidChangeTreeData(onDidChangeTreeDataListener);

-      await dataProvider.setState(methods, dbItem, !hideModeledMethods);
+      await dataProvider.setState(
+        methods,
+        dbItem,
+        !hideModeledMethods,
+        modeledMethods,
+        modifiedMethodSignatures,
+      );
+
+      expect(onDidChangeTreeDataListener).toHaveBeenCalledTimes(1);
+    });
+
+    it("should emit onDidChangeTreeData event when modeled methods has changed", async () => {
+      const modeledMethods2: Record<string, ModeledMethod> = {};
+
+      await dataProvider.setState(
+        methods,
+        dbItem,
+        hideModeledMethods,
+        modeledMethods,
+        modifiedMethodSignatures,
+      );
+
+      const onDidChangeTreeDataListener = jest.fn();
+      dataProvider.onDidChangeTreeData(onDidChangeTreeDataListener);
+
+      await dataProvider.setState(
+        methods,
+        dbItem,
+        hideModeledMethods,
+        modeledMethods2,
+        modifiedMethodSignatures,
+      );
+
+      expect(onDidChangeTreeDataListener).toHaveBeenCalledTimes(1);
+    });
+
+    it("should emit onDidChangeTreeData event when modified method signatures has changed", async () => {
+      const modifiedMethodSignatures2: Set<string> = new Set();
+
+      await dataProvider.setState(
+        methods,
+        dbItem,
+        hideModeledMethods,
+        modeledMethods,
+        modifiedMethodSignatures,
+      );
+
+      const onDidChangeTreeDataListener = jest.fn();
+      dataProvider.onDidChangeTreeData(onDidChangeTreeDataListener);
+
+      await dataProvider.setState(
+        methods,
+        dbItem,
+        hideModeledMethods,
+        modeledMethods,
+        modifiedMethodSignatures2,
+      );

      expect(onDidChangeTreeDataListener).toHaveBeenCalledTimes(1);
    });
@@ -79,12 +180,24 @@ describe("MethodsUsageDataProvider", () => {
      });
      const methods2: Method[] = [];

-      await dataProvider.setState(methods, dbItem, hideModeledMethods);
+      await dataProvider.setState(
+        methods,
+        dbItem,
+        hideModeledMethods,
+        modeledMethods,
+        modifiedMethodSignatures,
+      );

      const onDidChangeTreeDataListener = jest.fn();
      dataProvider.onDidChangeTreeData(onDidChangeTreeDataListener);

-      await dataProvider.setState(methods2, dbItem2, !hideModeledMethods);
+      await dataProvider.setState(
+        methods2,
+        dbItem2,
+        !hideModeledMethods,
+        modeledMethods,
+        modifiedMethodSignatures,
+      );

      expect(onDidChangeTreeDataListener).toHaveBeenCalledTimes(1);
    });
@@ -100,6 +213,9 @@ describe("MethodsUsageDataProvider", () => {
    });

    const methods: Method[] = [supportedMethod, unsupportedMethod];
+    const modeledMethods: Record<string, ModeledMethod> = {};
+    const modifiedMethodSignatures: Set<string> = new Set();
+
    const dbItem = mockedObject<DatabaseItem>({
      getSourceLocationPrefix: () => "test",
    });
@@ -117,13 +233,25 @@ describe("MethodsUsageDataProvider", () => {

    it("should show all methods if hideModeledMethods is false and looking at the root", async () => {
      const hideModeledMethods = false;
-      await dataProvider.setState(methods, dbItem, hideModeledMethods);
+      await dataProvider.setState(
+        methods,
+        dbItem,
+        hideModeledMethods,
+        modeledMethods,
+        modifiedMethodSignatures,
+      );
      expect(dataProvider.getChildren().length).toEqual(2);
    });

    it("should filter methods if hideModeledMethods is true and looking at the root", async () => {
      const hideModeledMethods = true;
-      await dataProvider.setState(methods, dbItem, hideModeledMethods);
+      await dataProvider.setState(
+        methods,
+        dbItem,
+        hideModeledMethods,
+        modeledMethods,
+        modifiedMethodSignatures,
+      );
      expect(dataProvider.getChildren().length).toEqual(1);
    });
  });
--- a/extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/methods-usage/methods-usage-panel.test.ts
+++ b/extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/methods-usage/methods-usage-panel.test.ts
@@ -10,6 +10,7 @@ import {
 } from "../../../../factories/model-editor/method-factories";
 import { ModelingStore } from "../../../../../src/model-editor/modeling-store";
 import { createMockModelingStore } from "../../../../__mocks__/model-editor/modelingStoreMock";
+import { ModeledMethod } from "../../../../../src/model-editor/modeled-method";

 describe("MethodsUsagePanel", () => {
  const mockCliServer = mockedObject<CodeQLCliServer>({});
@@ -20,6 +21,8 @@ describe("MethodsUsagePanel", () => {
  describe("setState", () => {
    const hideModeledMethods = false;
    const methods: Method[] = [createMethod()];
+    const modeledMethods: Record<string, ModeledMethod> = {};
+    const modifiedMethodSignatures: Set<string> = new Set();

    it("should update the tree view with the correct batch number", async () => {
      const mockTreeView = {
@@ -30,7 +33,13 @@ describe("MethodsUsagePanel", () => {
      const modelingStore = createMockModelingStore();

      const panel = new MethodsUsagePanel(modelingStore, mockCliServer);
-      await panel.setState(methods, dbItem, hideModeledMethods);
+      await panel.setState(
+        methods,
+        dbItem,
+        hideModeledMethods,
+        modeledMethods,
+        modifiedMethodSignatures,
+      );

      expect(mockTreeView.badge?.value).toBe(1);
    });
@@ -41,6 +50,8 @@ describe("MethodsUsagePanel", () => {
    let modelingStore: ModelingStore;

    const hideModeledMethods: boolean = false;
+    const modeledMethods: Record<string, ModeledMethod> = {};
+    const modifiedMethodSignatures: Set<string> = new Set();
    const usage = createUsage();

    beforeEach(() => {
@@ -60,7 +71,13 @@ describe("MethodsUsagePanel", () => {
      ];

      const panel = new MethodsUsagePanel(modelingStore, mockCliServer);
-      await panel.setState(methods, dbItem, hideModeledMethods);
+      await panel.setState(
+        methods,
+        dbItem,
+        hideModeledMethods,
+        modeledMethods,
+        modifiedMethodSignatures,
+      );

      await panel.revealItem(usage);

@@ -70,7 +87,13 @@ describe("MethodsUsagePanel", () => {
    it("should do nothing if usage cannot be found", async () => {
      const methods = [createMethod({})];
      const panel = new MethodsUsagePanel(modelingStore, mockCliServer);
-      await panel.setState(methods, dbItem, hideModeledMethods);
+      await panel.setState(
+        methods,
+        dbItem,
+        hideModeledMethods,
+        modeledMethods,
+        modifiedMethodSignatures,
+      );

      await panel.revealItem(usage);