Refetch entity selector data when the connection comes back

Reset entity sources and config entries when the websocket reconnects, and drop an unreachable cache guard.

src/common/util/time-cache-function-promise.ts

@@ -36,7 +36,9 @@ export const timeCachePromiseFunc = async <T, H = HomeAssistant>(
       return result.result;
     }
 
-    anyHass[cacheKey] = undefined;
+    if (anyHass[cacheKey] === result) {
+      anyHass[cacheKey] = undefined;
+    }
     return timeCachePromiseFunc(
       cacheKey,
       cacheTime,
@@ -73,7 +75,9 @@ export const timeCachePromiseFunc = async <T, H = HomeAssistant>(
     },
     // On failure, clear cache right away
     () => {
-      anyHass[cacheKey] = undefined;
+      if (anyHass[cacheKey] === cachePromise) {
+        anyHass[cacheKey] = undefined;
+      }
     }
   );
 

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

@@ -1,12 +1,21 @@
+import type { ContextType } from "@lit/context";
 import type { HassEntity } from "home-assistant-js-websocket";
 import type { PropertyValues } from "lit";
 import { html, LitElement, nothing } from "lit";
 import { customElement, property, state } from "lit/decorators";
 import memoizeOne from "memoize-one";
 import { ensureArray } from "../../common/array/ensure-array";
+import { consume } from "../../common/decorators/consume";
 import { fireEvent } from "../../common/dom/fire_event";
 import type { ConfigEntry } from "../../data/config_entries";
 import { getConfigEntries } from "../../data/config_entries";
+import {
+  apiContext,
+  connectionContext,
+  devicesContext,
+  entitiesContext,
+  statesContext,
+} from "../../data/context";
 import { getDeviceIntegrationLookup } from "../../data/device/device_registry";
 import type { EntitySources } from "../../data/entity/entity_sources";
 import { fetchEntitySourcesWithCache } from "../../data/entity/entity_sources";
@@ -15,14 +24,12 @@ import {
   computeCreateDomains,
   filterSelectorEntities,
 } from "../../data/selector";
-import type { HomeAssistant } from "../../types";
+import type { HomeAssistant, HomeAssistantConnection } from "../../types";
 import "../entity/ha-entities-picker";
 import "../entity/ha-entity-picker";
 
 @customElement("ha-selector-entity")
 export class HaEntitySelector extends LitElement {
-  @property({ attribute: false }) public hass!: HomeAssistant;
-
   @property({ attribute: false }) public selector!: EntitySelector;
 
   @state() private _entitySources?: EntitySources;
@@ -43,15 +50,31 @@ export class HaEntitySelector extends LitElement {
 
   @state() private _createDomains: string[] | undefined;
 
+  @consume({ context: apiContext, subscribe: true })
+  private _api!: ContextType<typeof apiContext>;
+
+  @state()
+  @consume({ context: connectionContext, subscribe: true })
+  private _connection?: ContextType<typeof connectionContext>;
+
+  @consume({ context: statesContext, subscribe: true })
+  private _states!: ContextType<typeof statesContext>;
+
+  @consume({ context: entitiesContext, subscribe: true })
+  private _entities!: ContextType<typeof entitiesContext>;
+
+  @consume({ context: devicesContext, subscribe: true })
+  private _devices!: ContextType<typeof devicesContext>;
+
   private _deviceIntegrationLookup = memoizeOne(
     (
-      entities: HomeAssistant["entities"],
-      devices: HomeAssistant["devices"],
+      entities: HomeAssistant["entities"] | null,
+      devices: HomeAssistant["devices"] | null,
       configEntries?: ConfigEntry[]
     ) =>
       getDeviceIntegrationLookup(
-        Object.values(entities),
-        Object.values(devices),
+        Object.values(entities ?? {}),
+        Object.values(devices ?? {}),
         configEntries
       )
   );
@@ -125,12 +148,19 @@ export class HaEntitySelector extends LitElement {
     `;
   }
 
-  protected updated(changedProps: PropertyValues<this>): void {
+  protected updated(changedProps: PropertyValues): void {
     super.updated(changedProps);
 
-    // The connection changed (e.g. reconnect); refetch config entries.
-    const oldHass = changedProps.get("hass");
-    if (oldHass && oldHass.connection !== this.hass.connection) {
+    // The connection changed or came back after a drop; refetch.
+    const oldConnection = changedProps.get("_connection") as
+      HomeAssistantConnection | undefined;
+    if (
+      oldConnection &&
+      this._connection &&
+      (oldConnection.connection !== this._connection.connection ||
+        (!oldConnection.connected && this._connection.connected))
+    ) {
+      this._entitySources = undefined;
       this._fetchedConfigEntries = false;
       this._configEntries = undefined;
     }
@@ -138,14 +168,18 @@ export class HaEntitySelector extends LitElement {
     const needs = this._dataNeeds(this.selector);
 
     if (needs.entitySources && !this._entitySources) {
-      fetchEntitySourcesWithCache(this.hass).then((sources) => {
-        this._entitySources = sources;
-      });
+      fetchEntitySourcesWithCache(this._api.callWS, this._states ?? {})
+        .then((sources) => {
+          this._entitySources = sources;
+        })
+        .catch(() => {
+          this._entitySources = {};
+        });
     }
 
     if (needs.configEntries && !this._fetchedConfigEntries) {
       this._fetchedConfigEntries = true;
-      getConfigEntries(this.hass)
+      getConfigEntries(this._api)
         .then((entries) => {
           this._configEntries = entries;
         })
@@ -169,8 +203,8 @@ export class HaEntitySelector extends LitElement {
     }
     const deviceIntegrationLookup = this._dataNeeds(this.selector).configEntries
       ? this._deviceIntegrationLookup(
-          this.hass.entities,
-          this.hass.devices,
+          this._entities,
+          this._devices,
           this._configEntries
         )
       : undefined;
@@ -180,8 +214,8 @@ export class HaEntitySelector extends LitElement {
         filter,
         entity,
         this._entitySources,
-        this.hass.entities,
-        this.hass.devices,
+        this._entities ?? undefined,
+        this._devices ?? undefined,
         deviceIntegrationLookup
       )
     );

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

@@ -69,9 +69,11 @@ export class HaTargetSelector extends LitElement {
       this._hasIntegration(this.selector) &&
       !this._entitySources
     ) {
-      fetchEntitySourcesWithCache(this.hass).then((sources) => {
-        this._entitySources = sources;
-      });
+      fetchEntitySourcesWithCache(this.hass.callWS, this.hass.states).then(
+        (sources) => {
+          this._entitySources = sources;
+        }
+      );
     }
     if (changedProperties.has("selector")) {
       this._createDomains = computeCreateDomains(this.selector);

src/components/target-picker/dialog/dialog-target-details.ts

@@ -135,7 +135,10 @@ class DialogTargetDetails extends LitElement implements HassDialog {
 
   private async _loadEntitySources(): Promise<void> {
     try {
-      this._entitySources = await fetchEntitySourcesWithCache(this.hass);
+      this._entitySources = await fetchEntitySourcesWithCache(
+        this.hass.callWS,
+        this.hass.states
+      );
     } catch (err) {
       // eslint-disable-next-line no-console
       console.error("Failed to load entity sources for target details", err);

src/data/entity/entity_sources.ts

@@ -1,3 +1,4 @@
+import type { HassEntities } from "home-assistant-js-websocket";
 import { timeCachePromiseFunc } from "../../common/util/time-cache-function-promise";
 import type { HomeAssistant } from "../../types";
 
@@ -7,20 +8,25 @@ interface EntitySource {
 
 export type EntitySources = Record<string, EntitySource>;
 
-const fetchEntitySources = (
-  hass: Pick<HomeAssistant, "callWS">
-): Promise<EntitySources> => hass.callWS({ type: "entity/source" });
+const caches = new WeakMap<HomeAssistant["callWS"], Record<string, unknown>>();
 
 export const fetchEntitySourcesWithCache = (
-  hass: Pick<HomeAssistant, "callWS" | "states">
-): Promise<EntitySources> =>
-  timeCachePromiseFunc(
+  callWS: HomeAssistant["callWS"],
+  states: HassEntities
+): Promise<EntitySources> => {
+  let cache = caches.get(callWS);
+  if (!cache) {
+    cache = {};
+    caches.set(callWS, cache);
+  }
+  return timeCachePromiseFunc(
     "_entitySources",
     // cache for 30 seconds
     30000,
-    fetchEntitySources,
+    () => callWS<EntitySources>({ type: "entity/source" }),
     // We base the cache on number of states. If number of states
     // changes we force a refresh
-    (hass2) => Object.keys(hass2.states).length,
-    hass
+    () => Object.keys(states).length,
+    cache
   );
+};

src/dialogs/more-info/components/update/ha-more-info-update-backup.ts

@@ -145,10 +145,10 @@ export class HaMoreInfoUpdateBackup extends LitElement {
     let type: UpdateType | undefined = undefined;
 
     try {
-      this._entitySources = await fetchEntitySourcesWithCache({
-        callWS: this._api.callWS,
-        states: this._states,
-      });
+      this._entitySources = await fetchEntitySourcesWithCache(
+        this._api.callWS,
+        this._states
+      );
       type = getUpdateType(this.stateObj!, this._entitySources!);
     } catch (err) {
       // ignore error, because the generic backup option remains available

src/panels/config/core/ha-config-section-updates.ts

@@ -96,7 +96,10 @@ class HaConfigSectionUpdates extends SubscribeMixin(LitElement) {
 
   private async _loadEntitySources() {
     try {
-      this._entitySources = await fetchEntitySourcesWithCache(this.hass);
+      this._entitySources = await fetchEntitySourcesWithCache(
+        this.hass.callWS,
+        this.hass.states
+      );
     } catch (_err) {
       // Non-fatal: grouping falls back to entity registry platform lookup.
     }

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

@@ -1164,9 +1164,11 @@ export class HaConfigEntities extends LitElement {
 
   protected firstUpdated() {
     this._fetchExposedEntities();
-    fetchEntitySourcesWithCache(this.hass).then((sources) => {
-      this._entitySources = sources;
-    });
+    fetchEntitySourcesWithCache(this.hass.callWS, this.hass.states).then(
+      (sources) => {
+        this._entitySources = sources;
+      }
+    );
   }
 
   // Links, like the ones on a device or integration page, say in the URL what

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

@@ -1252,7 +1252,7 @@ ${rejected
 
   private async _fetchEntitySources() {
     const [entitySources, fetchedManifests] = await Promise.all([
-      fetchEntitySourcesWithCache(this.hass),
+      fetchEntitySourcesWithCache(this.hass.callWS, this.hass.states),
       fetchIntegrationManifests(this.hass),
     ]);
 

src/panels/config/integrations/ha-config-integration-page.ts

@@ -268,7 +268,10 @@ class HaConfigIntegrationPage extends SubscribeMixin(LitElement) {
   }
 
   private async _fetchEntitySources() {
-    const entitySources = await fetchEntitySourcesWithCache(this.hass);
+    const entitySources = await fetchEntitySourcesWithCache(
+      this.hass.callWS,
+      this.hass.states
+    );
 
     const entitiesByDomain = {};
 

src/panels/config/integrations/ha-config-integrations-dashboard.ts

@@ -857,7 +857,10 @@ class HaConfigIntegrationsDashboard extends KeyboardShortcutMixin(
   };
 
   private async _fetchEntitySources() {
-    const entitySources = await fetchEntitySourcesWithCache(this.hass);
+    const entitySources = await fetchEntitySourcesWithCache(
+      this.hass.callWS,
+      this.hass.states
+    );
 
     const entitiesByDomain = {};
 

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

@@ -460,9 +460,11 @@ class HaPanelHistory extends LitElement {
     super.firstUpdated(changedProps);
     // The filters label themselves with keys from the config panel.
     this.hass.loadFragmentTranslation("config");
-    fetchEntitySourcesWithCache(this.hass).then((sources) => {
-      this._entitySources = sources;
-    });
+    fetchEntitySourcesWithCache(this.hass.callWS, this.hass.states).then(
+      (sources) => {
+        this._entitySources = sources;
+      }
+    );
     const searchParams = extractSearchParamsObject();
     if (searchParams.back === "1" && history.length > 1) {
       this._showBack = true;

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

@@ -370,9 +370,11 @@ export class HaPanelLogbook extends LitElement {
     this.hass.loadBackendTranslation("title");
     // The filters label themselves with keys from the config panel.
     this.hass.loadFragmentTranslation("config");
-    fetchEntitySourcesWithCache(this.hass).then((sources) => {
-      this._entitySources = sources;
-    });
+    fetchEntitySourcesWithCache(this.hass.callWS, this.hass.states).then(
+      (sources) => {
+        this._entitySources = sources;
+      }
+    );
 
     const searchParams = extractSearchParamsObject();
     if (searchParams.back === "1" && history.length > 1) {

test/common/util/time-cache-function-promise.test.ts

@@ -176,3 +176,66 @@ describe("timeCachePromiseFunc cache ownership", () => {
     }
   });
 });
+
+describe("timeCachePromiseFunc overlapping callers", () => {
+  const flush = () =>
+    new Promise((resolve) => {
+      setTimeout(resolve, 0);
+    });
+
+  const setup = () => {
+    const resolvers: ((value: string) => void)[] = [];
+    const fetchData = vi.fn(
+      () =>
+        new Promise<string>((resolve) => {
+          resolvers.push(resolve);
+        })
+    );
+    const cache = {};
+    const call = (key: number) =>
+      timeCachePromiseFunc("_testOverlap", 30_000, fetchData, () => key, cache);
+    return { resolvers, fetchData, call };
+  };
+
+  it("starts one refetch when callers with the same new key wait on an in-flight request", async () => {
+    const { resolvers, fetchData, call } = setup();
+
+    call(1);
+    const first = call(2);
+    const second = call(2);
+
+    resolvers[0]("old");
+    await flush();
+    expect(fetchData).toHaveBeenCalledTimes(2);
+
+    resolvers[1]("new");
+    await expect(first).resolves.toBe("new");
+    await expect(second).resolves.toBe("new");
+
+    await expect(call(2)).resolves.toBe("new");
+    expect(fetchData).toHaveBeenCalledTimes(2);
+  });
+
+  it("does not let an older request overwrite a newer cached result", async () => {
+    const { resolvers, fetchData, call } = setup();
+
+    call(1);
+    const keyTwo = call(2);
+    const keyThree = call(3);
+
+    resolvers[0]("one");
+    await flush();
+    expect(fetchData).toHaveBeenCalledTimes(2);
+
+    resolvers[1]("two");
+    await expect(keyTwo).resolves.toBe("two");
+    await flush();
+    expect(fetchData).toHaveBeenCalledTimes(3);
+
+    resolvers[2]("three");
+    await expect(keyThree).resolves.toBe("three");
+
+    await expect(call(3)).resolves.toBe("three");
+    expect(fetchData).toHaveBeenCalledTimes(3);
+  });
+});

test/data/entity/entity_sources.test.ts

@@ -0,0 +1,65 @@
+import type { HassEntities } from "home-assistant-js-websocket";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { fetchEntitySourcesWithCache } from "../../../src/data/entity/entity_sources";
+import type { HomeAssistant } from "../../../src/types";
+
+const states = { "light.kitchen": {} } as unknown as HassEntities;
+
+const mockCallWS = () =>
+  vi.fn().mockResolvedValue({
+    "light.kitchen": { domain: "hue" },
+  }) as unknown as HomeAssistant["callWS"] & ReturnType<typeof vi.fn>;
+
+describe("fetchEntitySourcesWithCache", () => {
+  beforeEach(() => {
+    vi.useFakeTimers();
+  });
+
+  afterEach(() => {
+    vi.useRealTimers();
+  });
+
+  it("reuses the result for the same callWS within 30 seconds", async () => {
+    const callWS = mockCallWS();
+
+    await fetchEntitySourcesWithCache(callWS, states);
+    await vi.advanceTimersByTimeAsync(29_000);
+    await fetchEntitySourcesWithCache(callWS, { ...states });
+
+    expect(callWS).toHaveBeenCalledTimes(1);
+    expect(callWS).toHaveBeenCalledWith({ type: "entity/source" });
+  });
+
+  it("refetches after 30 seconds", async () => {
+    const callWS = mockCallWS();
+
+    await fetchEntitySourcesWithCache(callWS, states);
+    await vi.advanceTimersByTimeAsync(30_000);
+    await fetchEntitySourcesWithCache(callWS, states);
+
+    expect(callWS).toHaveBeenCalledTimes(2);
+  });
+
+  it("refetches when the number of states changes", async () => {
+    const callWS = mockCallWS();
+
+    await fetchEntitySourcesWithCache(callWS, states);
+    await fetchEntitySourcesWithCache(callWS, {
+      ...states,
+      "light.hall": {},
+    } as unknown as HassEntities);
+
+    expect(callWS).toHaveBeenCalledTimes(2);
+  });
+
+  it("keeps a separate cache per callWS", async () => {
+    const first = mockCallWS();
+    const second = mockCallWS();
+
+    await fetchEntitySourcesWithCache(first, states);
+    await fetchEntitySourcesWithCache(second, states);
+
+    expect(first).toHaveBeenCalledTimes(1);
+    expect(second).toHaveBeenCalledTimes(1);
+  });
+});