WIP

src/components/ha-selector/ha-selector-action.ts

@@ -1,23 +1,18 @@
-import { consume, ContextProvider } from "@lit/context";
-import type { UnsubscribeFunc } from "home-assistant-js-websocket";
+import { consume } from "@lit/context";
 import { css, html, LitElement, nothing } from "lit";
 import { customElement, property, query, state } from "lit/decorators";
 import memoizeOne from "memoize-one";
 import { fullEntitiesContext } from "../../data/context";
-import {
-  subscribeEntityRegistry,
-  type EntityRegistryEntry,
-} from "../../data/entity/entity_registry";
+import type { EntityRegistryEntry } from "../../data/entity/entity_registry";
 import type { Action } from "../../data/script";
 import { migrateAutomationAction } from "../../data/script";
 import type { ActionSelector } from "../../data/selector";
-import { SubscribeMixin } from "../../mixins/subscribe-mixin";
 import "../../panels/config/automation/action/ha-automation-action";
 import type HaAutomationAction from "../../panels/config/automation/action/ha-automation-action";
 import type { HomeAssistant } from "../../types";
 
 @customElement("ha-selector-action")
-export class HaActionSelector extends SubscribeMixin(LitElement) {
+export class HaActionSelector extends LitElement {
   @property({ attribute: false }) public hass!: HomeAssistant;
 
   @property({ type: Boolean }) public narrow = false;
@@ -34,13 +29,9 @@ export class HaActionSelector extends SubscribeMixin(LitElement) {
   @consume({ context: fullEntitiesContext, subscribe: true })
   _entityReg: EntityRegistryEntry[] | undefined;
 
-  @state() private _entitiesContext;
-
   @query("ha-automation-action")
   private _actionElement?: HaAutomationAction;
 
-  protected hassSubscribeRequiredHostProps = ["_entitiesContext"];
-
   private _actions = memoizeOne((action: Action | undefined) => {
     if (!action) {
       return [];
@@ -48,23 +39,6 @@ export class HaActionSelector extends SubscribeMixin(LitElement) {
     return migrateAutomationAction(action);
   });
 
-  protected firstUpdated() {
-    if (!this._entityReg) {
-      this._entitiesContext = new ContextProvider(this, {
-        context: fullEntitiesContext,
-        initialValue: [],
-      });
-    }
-  }
-
-  public hassSubscribe(): UnsubscribeFunc[] {
-    return [
-      subscribeEntityRegistry(this.hass.connection!, (entities) => {
-        this._entitiesContext.setValue(entities);
-      }),
-    ];
-  }
-
   public expandAll() {
     this._actionElement?.expandAll();
   }

src/panels/config/automation/manual-automation-editor.ts

@@ -1,6 +1,5 @@
-import { ContextProvider } from "@lit/context";
 import { mdiContentSave } from "@mdi/js";
-import type { HassEntity, UnsubscribeFunc } from "home-assistant-js-websocket";
+import type { HassEntity } from "home-assistant-js-websocket";
 import { load } from "js-yaml";
 import type { CSSResultGroup, PropertyValues } from "lit";
 import { css, html, LitElement, nothing } from "lit";
@@ -47,13 +46,7 @@ import {
   isTrigger,
   normalizeAutomationConfig,
 } from "../../../data/automation";
-import {
-  subscribeAndProcessConfigEntries,
-  type ConfigEntry,
-} from "../../../data/config_entries";
-import { configEntriesContext } from "../../../data/context";
 import { getActionType, type Action } from "../../../data/script";
-import { SubscribeMixin } from "../../../mixins/subscribe-mixin";
 import type { HomeAssistant, ValueChangedEvent } from "../../../types";
 import { showToast } from "../../../util/toast";
 import "./action/ha-automation-action";
@@ -86,7 +79,7 @@ const automationConfigStruct = union([
 export const SIDEBAR_DEFAULT_WIDTH = 500;
 
 @customElement("manual-automation-editor")
-export class HaManualAutomationEditor extends SubscribeMixin(LitElement) {
+export class HaManualAutomationEditor extends LitElement {
   @property({ attribute: false }) public hass!: HomeAssistant;
 
   @property({ attribute: "is-wide", type: Boolean }) public isWide = false;
@@ -123,11 +116,6 @@ export class HaManualAutomationEditor extends SubscribeMixin(LitElement) {
     HaAutomationAction | HaAutomationCondition
   >;
 
-  private _configEntries = new ContextProvider(this, {
-    context: configEntriesContext,
-    initialValue: [],
-  });
-
   private _prevSidebarWidthPx?: number;
 
   public connectedCallback() {
@@ -135,18 +123,6 @@ export class HaManualAutomationEditor extends SubscribeMixin(LitElement) {
     window.addEventListener("paste", this._handlePaste);
   }
 
-  public hassSubscribe(): Promise<UnsubscribeFunc>[] {
-    return [
-      subscribeAndProcessConfigEntries(
-        this.hass,
-        (message: ConfigEntry[]) => {
-          this._configEntries.setValue(message);
-        },
-        undefined
-      ),
-    ];
-  }
-
   public disconnectedCallback() {
     window.removeEventListener("paste", this._handlePaste);
     super.disconnectedCallback();

src/panels/config/ha-panel-config.ts

@@ -1,4 +1,3 @@
-import { ContextProvider } from "@lit/context";
 import {
   mdiAccount,
   mdiBackupRestore,
@@ -34,24 +33,19 @@ import {
   mdiZigbee,
   mdiZWave,
 } from "@mdi/js";
-import type { UnsubscribeFunc } from "home-assistant-js-websocket";
 import type { PropertyValues } from "lit";
 import { customElement, property, state } from "lit/decorators";
 import { isComponentLoaded } from "../../common/config/is_component_loaded";
 import { listenMediaQuery } from "../../common/dom/media_query";
 import type { CloudStatus } from "../../data/cloud";
 import { fetchCloudStatus } from "../../data/cloud";
-import { fullEntitiesContext, labelsContext } from "../../data/context";
 import {
   entityRegistryByEntityId,
   entityRegistryById,
-  subscribeEntityRegistry,
 } from "../../data/entity/entity_registry";
-import { subscribeLabelRegistry } from "../../data/label/label_registry";
 import type { RouterOptions } from "../../layouts/hass-router-page";
 import { HassRouterPage } from "../../layouts/hass-router-page";
 import type { PageNavigation } from "../../layouts/hass-tabs-subpage";
-import { SubscribeMixin } from "../../mixins/subscribe-mixin";
 import type { HomeAssistant, Route } from "../../types";
 
 declare global {
@@ -499,34 +493,13 @@ export const configSections: Record<string, PageNavigation[]> = {
 };
 
 @customElement("ha-panel-config")
-class HaPanelConfig extends SubscribeMixin(HassRouterPage) {
+class HaPanelConfig extends HassRouterPage {
   @property({ attribute: false }) public hass!: HomeAssistant;
 
   @property({ type: Boolean }) public narrow = false;
 
   @property({ attribute: false }) public route!: Route;
 
-  private _entitiesContext = new ContextProvider(this, {
-    context: fullEntitiesContext,
-    initialValue: [],
-  });
-
-  private _labelsContext = new ContextProvider(this, {
-    context: labelsContext,
-    initialValue: [],
-  });
-
-  public hassSubscribe(): UnsubscribeFunc[] {
-    return [
-      subscribeEntityRegistry(this.hass.connection!, (entities) => {
-        this._entitiesContext.setValue(entities);
-      }),
-      subscribeLabelRegistry(this.hass.connection!, (labels) => {
-        this._labelsContext.setValue(labels);
-      }),
-    ];
-  }
-
   protected routerOptions: RouterOptions = {
     defaultPage: "dashboard",
     routes: {

src/panels/config/script/manual-script-editor.ts

@@ -1,6 +1,4 @@
-import { ContextProvider } from "@lit/context";
 import { mdiContentSave, mdiHelpCircleOutline } from "@mdi/js";
-import type { UnsubscribeFunc } from "home-assistant-js-websocket";
 import { load } from "js-yaml";
 import type { CSSResultGroup, PropertyValues } from "lit";
 import { css, html, LitElement, nothing } from "lit";
@@ -37,8 +35,6 @@ import type {
   ActionSidebarConfig,
   SidebarConfig,
 } from "../../../data/automation";
-import { subscribeAndProcessConfigEntries } from "../../../data/config_entries";
-import { configEntriesContext } from "../../../data/context";
 import type {
   Action,
   Fields,
@@ -50,7 +46,6 @@ import {
   MODES,
   normalizeScriptConfig,
 } from "../../../data/script";
-import { SubscribeMixin } from "../../../mixins/subscribe-mixin";
 import type { HomeAssistant, ValueChangedEvent } from "../../../types";
 import { documentationUrl } from "../../../util/documentation-url";
 import { showToast } from "../../../util/toast";
@@ -75,7 +70,7 @@ const scriptConfigStruct = object({
 });
 
 @customElement("manual-script-editor")
-export class HaManualScriptEditor extends SubscribeMixin(LitElement) {
+export class HaManualScriptEditor extends LitElement {
   @property({ attribute: false }) public hass!: HomeAssistant;
 
   @property({ attribute: "is-wide", type: Boolean }) public isWide = false;
@@ -113,11 +108,6 @@ export class HaManualScriptEditor extends SubscribeMixin(LitElement) {
     HaAutomationAction | HaScriptFields
   >;
 
-  private _configEntries = new ContextProvider(this, {
-    context: configEntriesContext,
-    initialValue: [],
-  });
-
   private _openFields = false;
 
   private _prevSidebarWidthPx?: number;
@@ -139,14 +129,6 @@ export class HaManualScriptEditor extends SubscribeMixin(LitElement) {
     });
   }
 
-  public hassSubscribe(): Promise<UnsubscribeFunc>[] {
-    return [
-      subscribeAndProcessConfigEntries(this.hass, (configEntries) => {
-        this._configEntries.setValue(configEntries);
-      }),
-    ];
-  }
-
   protected updated(changedProps) {
     if (this._openFields && changedProps.has("config")) {
       this._openFields = false;

src/state/context-mixin.ts

@@ -1,12 +1,17 @@
 import { ContextProvider } from "@lit/context";
-import type { UnsubscribeFunc } from "home-assistant-js-websocket";
+import {
+  ConfigEntryStream,
+  type ConfigEntryUpdate,
+} from "../data/config_entries";
 import {
   areasContext,
   configContext,
+  configEntriesContext,
   connectionContext,
   devicesContext,
   entitiesContext,
   floorsContext,
+  fullEntitiesContext,
   labelsContext,
   localeContext,
   localizeContext,
@@ -17,16 +22,16 @@ import {
   userContext,
   userDataContext,
 } from "../data/context";
+import { subscribeEntityRegistry } from "../data/entity/entity_registry";
 import { subscribeLabelRegistry } from "../data/label/label_registry";
 import type { Constructor, HomeAssistant } from "../types";
 import type { HassBaseEl } from "./hass-base-mixin";
+import { LazyContextProvider } from "./lazy-context-provider";
 
 export const contextMixin = <T extends Constructor<HassBaseEl>>(
   superClass: T
 ) =>
   class extends superClass {
-    private _unsubscribeLabels?: UnsubscribeFunc;
-
     private __contextProviders: Record<
       string,
       ContextProvider<any> | undefined
@@ -97,9 +102,30 @@ export const contextMixin = <T extends Constructor<HassBaseEl>>(
         context: floorsContext,
         initialValue: this.hass ? this.hass.floors : this._pendingHass.floors,
       }),
-      labels: new ContextProvider(this, {
+    };
+
+    private __lazyContextProviders = {
+      labels: new LazyContextProvider(this, {
         context: labelsContext,
-        initialValue: [],
+        subscribeFn: (connection, setValue) =>
+          subscribeLabelRegistry(connection, setValue),
+      }),
+      fullEntities: new LazyContextProvider(this, {
+        context: fullEntitiesContext,
+        subscribeFn: (connection, setValue) =>
+          subscribeEntityRegistry(connection, setValue),
+      }),
+      configEntries: new LazyContextProvider(this, {
+        context: configEntriesContext,
+        subscribeFn: (connection, setValue) => {
+          const stream = new ConfigEntryStream();
+          return connection.subscribeMessage<ConfigEntryUpdate[]>(
+            (messages) => {
+              setValue(stream.processMessage(messages));
+            },
+            { type: "config_entries/subscribe" }
+          );
+        },
       }),
     };
 
@@ -111,12 +137,11 @@ export const contextMixin = <T extends Constructor<HassBaseEl>>(
         }
       }
 
-      this._unsubscribeLabels = subscribeLabelRegistry(
-        this.hass!.connection!,
-        (labels) => {
-          this.__contextProviders.labels!.setValue(labels);
-        }
-      );
+      // Provide connection to lazy providers so they can subscribe on demand
+      const connection = this.hass!.connection!;
+      for (const provider of Object.values(this.__lazyContextProviders)) {
+        provider.setConnection(connection);
+      }
     }
 
     protected _updateHass(obj: Partial<HomeAssistant>) {
@@ -130,6 +155,8 @@ export const contextMixin = <T extends Constructor<HassBaseEl>>(
 
     public disconnectedCallback() {
       super.disconnectedCallback();
-      this._unsubscribeLabels?.();
+      for (const provider of Object.values(this.__lazyContextProviders)) {
+        provider.unsubscribe();
+      }
     }
   };

src/state/lazy-context-provider.ts

@@ -0,0 +1,188 @@
+import type { Context, ContextType } from "@lit/context";
+import { ContextProvider } from "@lit/context";
+import type { Connection, UnsubscribeFunc } from "home-assistant-js-websocket";
+import type { ReactiveElement } from "lit";
+
+type SubscribeFn<T> = (
+  connection: Connection,
+  setValue: (value: T) => void
+) => UnsubscribeFunc | Promise<UnsubscribeFunc>;
+
+/**
+ * A context provider that defers its data subscription until the first
+ * consumer requests the context. This avoids unnecessary WebSocket
+ * subscriptions for data that may never be needed.
+ *
+ * Consumers that request the context before data has loaded will have
+ * their callbacks buffered and flushed once the first value arrives.
+ */
+export class LazyContextProvider<
+  C extends Context<unknown, unknown>,
+  T extends ContextType<C> = ContextType<C>,
+> {
+  private _provider: ContextProvider<C>;
+
+  private _context: C;
+
+  private _loaded = false;
+
+  private _subscribing = false;
+
+  private _connection?: Connection;
+
+  private _unsubscribe?: UnsubscribeFunc;
+
+  private _subscribeFn: SubscribeFn<T>;
+
+  private _pendingCallbacks: {
+    callback: (value: T, unsubscribe?: () => void) => void;
+    consumerHost: Element;
+    subscribe?: boolean;
+  }[] = [];
+
+  constructor(
+    private _host: ReactiveElement,
+    options: { context: C; subscribeFn: SubscribeFn<T> }
+  ) {
+    this._context = options.context;
+    this._subscribeFn = options.subscribeFn;
+
+    // Listen for context-request events BEFORE the ContextProvider does,
+    // so we can intercept requests when data hasn't loaded yet.
+    this._host.addEventListener(
+      "context-request",
+      this._onContextRequest as EventListener
+    );
+
+    // Create the underlying ContextProvider without an initial value.
+    // The provider's internal value will be undefined until data loads.
+    this._provider = new ContextProvider(this._host, {
+      context: options.context,
+    });
+  }
+
+  /**
+   * Set the connection reference. Called from hassConnected().
+   * Does not start subscribing -- that only happens when a consumer
+   * requests the context.
+   */
+  setConnection(connection: Connection): void {
+    this._connection = connection;
+
+    // If we were already subscribed (reconnection scenario), re-subscribe
+    if (this._loaded) {
+      this._unsubscribe?.();
+      this._startSubscription();
+    }
+
+    // If there were pending callbacks waiting for a connection, start now
+    if (this._pendingCallbacks.length > 0 && !this._subscribing) {
+      this._startSubscription();
+    }
+  }
+
+  /**
+   * Clean up the subscription.
+   */
+  unsubscribe(): void {
+    if (this._unsubscribe) {
+      this._unsubscribe();
+      this._unsubscribe = undefined;
+    }
+    this._loaded = false;
+    this._subscribing = false;
+  }
+
+  private _onContextRequest = (ev: Event): void => {
+    const contextEvent = ev as Event & {
+      context: unknown;
+      callback: (value: T, unsubscribe?: () => void) => void;
+      contextTarget?: Element;
+      subscribe?: boolean;
+    };
+
+    // Only handle requests for our context
+    if (contextEvent.context !== this._context) {
+      return;
+    }
+
+    // Don't intercept if data is already loaded --
+    // let the ContextProvider handle it normally
+    if (this._loaded) {
+      return;
+    }
+
+    const consumerHost =
+      contextEvent.contextTarget ?? (ev.composedPath()[0] as Element);
+
+    // Don't self-register
+    if (consumerHost === this._host) {
+      return;
+    }
+
+    // Intercept: stop propagation so the inner provider doesn't
+    // call back with undefined
+    ev.stopPropagation();
+
+    // Buffer this callback
+    this._pendingCallbacks.push({
+      callback: contextEvent.callback,
+      consumerHost,
+      subscribe: contextEvent.subscribe,
+    });
+
+    // Trigger the subscription if not already in progress
+    if (!this._subscribing && this._connection) {
+      this._startSubscription();
+    }
+  };
+
+  private _startSubscription(): void {
+    if (!this._connection || this._subscribing) {
+      return;
+    }
+    this._subscribing = true;
+
+    const result = this._subscribeFn(this._connection, (value: T) => {
+      this._loaded = true;
+      this._subscribing = false;
+
+      // Set the value on the real provider -- this updates all future consumers
+      this._provider.setValue(value as ContextType<C>);
+
+      // Flush any pending callbacks that were buffered before data loaded
+      this._flushPendingCallbacks();
+    });
+
+    // Handle async unsubscribe (Promise<UnsubscribeFunc>)
+    if (result instanceof Promise) {
+      result.then((unsub) => {
+        this._unsubscribe = unsub;
+      });
+    } else {
+      this._unsubscribe = result;
+    }
+  }
+
+  private _flushPendingCallbacks(): void {
+    if (this._pendingCallbacks.length === 0) {
+      return;
+    }
+
+    const pending = [...this._pendingCallbacks];
+    this._pendingCallbacks = [];
+
+    // Re-add each pending callback to the provider now that it has data
+    for (const { callback, consumerHost, subscribe } of pending) {
+      (
+        this._provider as unknown as {
+          addCallback: (
+            callback: (value: T, unsubscribe?: () => void) => void,
+            consumerHost: Element,
+            subscribe?: boolean
+          ) => void;
+        }
+      ).addCallback(callback, consumerHost, subscribe);
+    }
+  }
+}