fix(state): unsub to subscribed context on disconnect

example/src/modules/x/contextParent/contextParent.html

@@ -1,5 +1,7 @@
 <template>
   <h3>Parent</h3>
-  <x-context-child></x-context-child>
+  <template lwc:if={showChild}>
+    <x-context-child></x-context-child>
+  </template>
   <x-context-child-sibling></x-context-child-sibling>
 </template>

example/src/modules/x/contextParent/contextParent.js

@@ -5,4 +5,11 @@ import parentStateFactory from 'x/parentState';
 export default class ContextParent extends ContextfulLightningElement {
   @api
   parentState = parentStateFactory('parentFoo');
+
+  @api
+  hideChild = false;
+
+  get showChild() {
+    return !this.hideChild;
+  }
 }

example/src/modules/x/contextRoot/context.spec.js

@@ -76,4 +76,24 @@ describe('context', () => {
     expect(grandChildContentParent.innerText).to.include('parentFoo');
     expect(grandChildContentChild.innerText).to.include('bar');
   });
+
+  it('disconnecting child removes its subscriptions to parent state context', async () => {
+    const el = await clientSideRender(parentEl, componentPath, {});
+    const contextParent = querySelectorDeep('x-context-parent');
+
+    // Current active subscriptions
+    // 1. <x-context-parent>
+    // 2. <x-context-child>
+    // 3. <x-context-grand-child>
+    expect(contextParent.parentState.subscribers.size).toBe(3);
+
+    // Only active subscription to parent State will be
+    // LWC framework subscribing to component <x-context-parent>
+    // subscription of context child and grand child
+    // should be removed after unsubscribing
+    contextParent.hideChild = true;
+    await freshRender();
+
+    expect(contextParent.parentState.subscribers.size).toBe(1);
+  });
 });

packages/@lwc/state/src/contextful-lwc.ts

@@ -1,14 +1,18 @@
 import type { Signal } from '@lwc/signals';
 import { LightningElement } from 'lwc';
 import { ContextRequestEvent, type ContextProvidedCallback, symbolContextKey } from './event.js';
-import { connectContext } from './shared.js';
+import { connectContext, disconnectContext } from './shared.js';
 import type { ContextRuntimeAdapter } from './runtime-interface.js';
 
 export class ContextfulLightningElement extends LightningElement {
   connectedCallback(): void {
     this.setupContextReactivity();
   }
 
+  disconnectedCallback(): void {
+    this.cleanupContext();
+  }
+
   private setupContextReactivity() {
     const contextfulFields = Object.getOwnPropertyNames(Object.getPrototypeOf(this)).filter(
       (propName) => this[propName]?.[connectContext],
@@ -24,6 +28,7 @@ export class ContextfulLightningElement extends LightningElement {
 
     const contextRuntimeAdapter: ContextRuntimeAdapter<LightningElement> = {
       isServerSide: false,
+      component: this,
 
       provideContext<T extends object>(
         contextVariety: T,
@@ -74,4 +79,18 @@ export class ContextfulLightningElement extends LightningElement {
       this[contextfulField][connectContext](contextRuntimeAdapter);
     }
   }
+
+  private cleanupContext() {
+    const contextfulFields = Object.getOwnPropertyNames(Object.getPrototypeOf(this)).filter(
+      (propName) => this[propName]?.[disconnectContext],
+    );
+
+    if (contextfulFields.length === 0) {
+      return;
+    }
+
+    for (const contextfulField of contextfulFields) {
+      this[contextfulField][disconnectContext](this);
+    }
+  }
 }

packages/@lwc/state/src/index.ts

@@ -1,6 +1,6 @@
 import { type Signal, SignalBaseClass } from '@lwc/signals';
 import type { ContextRuntimeAdapter } from './runtime-interface.js';
-import { connectContext } from './shared.js';
+import { connectContext, disconnectContext } from './shared.js';
 import type {
   Computer,
   DefineState,
@@ -147,6 +147,7 @@ export const defineState: DefineState = <
       private contextConsumptionQueue: Array<
         (runtimeAdapter: ContextRuntimeAdapter<object>) => void
       > = [];
+      private contextUnsubscribes = new WeakMap<object, Array<() => void>>();
 
       constructor() {
         super();
@@ -177,18 +178,18 @@ export const defineState: DefineState = <
               (providedContextSignal: Signal<T>) => {
                 // Make sure the local signal initially shares the same value as the provided context signal.
                 localContextSignal[atomSetter](providedContextSignal.value);
-                // TODO: capture this unsubscribe in a map somewhere so that when the state manager disconnects
-                //       from the DOM, we can disconnect the context as well.
-                const _unsubscribe = providedContextSignal.subscribe(() => {
+                const unsub = providedContextSignal.subscribe(() => {
                   localContextSignal[atomSetter](providedContextSignal.value);
                 });
+
+                if (!this.contextUnsubscribes.has(runtimeAdapter.component)) {
+                  this.contextUnsubscribes.set(runtimeAdapter.component, []);
+                }
+                this.contextUnsubscribes.get(runtimeAdapter.component).push(unsub);
               },
             );
           });
 
-          // TODO: if this.runtimeAdapter is null but is not-null in the future, we'll need to connect
-          //       to context in the same way we have done above.
-
           return localContextSignal;
         };
 
@@ -224,6 +225,18 @@ export const defineState: DefineState = <
         // we have something concrete to work with, write a test, and then work backwards
       }
 
+      [disconnectContext](componentId: ContextRuntimeAdapter<object>['component']) {
+        const unsubArray = this.contextUnsubscribes.get(componentId);
+
+        if (!unsubArray) {
+          return;
+        }
+
+        while (unsubArray.length !== 0) {
+          unsubArray.pop()();
+        }
+      }
+
       private shareableContext(): ContextAtomSignal<unknown> {
         const contextAtom = new ContextAtomSignal<unknown>(undefined);
 

packages/@lwc/state/src/runtime-interface.ts

@@ -4,6 +4,7 @@ type ContextProvidedCallback = (contextSignal: Signal<unknown>) => void;
 
 export interface ContextRuntimeAdapter<T extends object> {
   isServerSide: boolean;
+  component: object;
   provideContext<T extends object>(contextVariety: T, providedContextSignal: Signal<unknown>): void;
   consumeContext<T extends object>(
     contextVariety: T,

packages/@lwc/state/src/shared.ts

@@ -1 +1,2 @@
 export const connectContext = Symbol('connectContext');
+export const disconnectContext = Symbol('disconnectContext');