From f04ce8593659195648fc8a33800abb1a03d4d687 Mon Sep 17 00:00:00 2001
From: rax-it <rakshit.verma@salesforce.com>
Date: Fri, 15 Nov 2024 11:25:51 -0800
Subject: [PATCH] fix: cleanup on disconnect

---
 .../contextParentDetached.html                   |  4 +++-
 .../contextParentDetached.js                     |  7 +++++++
 .../src/modules/x/contextRoot/context.spec.js    | 16 ++++++++++++++--
 packages/@lwc/state/src/standalone-context.ts    | 12 ++++++++++--
 4 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/example/src/modules/x/contextParentDetached/contextParentDetached.html b/example/src/modules/x/contextParentDetached/contextParentDetached.html
index a4e3490..bb2d86f 100644
--- a/example/src/modules/x/contextParentDetached/contextParentDetached.html
+++ b/example/src/modules/x/contextParentDetached/contextParentDetached.html
@@ -1,3 +1,5 @@
 <template>
-  <x-context-child-detached></x-context-child-detached>
+  <template lwc:if={showChild}>
+    <x-context-child-detached></x-context-child-detached>
+  </template>
 </template>
diff --git a/example/src/modules/x/contextParentDetached/contextParentDetached.js b/example/src/modules/x/contextParentDetached/contextParentDetached.js
index 938dba1..d92044a 100644
--- a/example/src/modules/x/contextParentDetached/contextParentDetached.js
+++ b/example/src/modules/x/contextParentDetached/contextParentDetached.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;
+  }
 }
diff --git a/example/src/modules/x/contextRoot/context.spec.js b/example/src/modules/x/contextRoot/context.spec.js
index 5c1f417..ea60893 100644
--- a/example/src/modules/x/contextRoot/context.spec.js
+++ b/example/src/modules/x/contextRoot/context.spec.js
@@ -83,13 +83,13 @@ describe('context', () => {
     const contextParent = querySelectorDeep('x-context-parent');
 
     // Current active subscriptions
-    // 1. <x-context-parent>
+    // 1. <x-context-child-sibling>
     // 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>
+    // from component: <x-context-child-sibling>
     // subscription of context child and grand child
     // should be removed after unsubscribing
     contextParent.hideChild = true;
@@ -125,4 +125,16 @@ describe('context', () => {
     const childWithDetachedFromContext = querySelectorDeep('.child-content-detached', el);
     expect(childWithDetachedFromContext.innerText).to.include('parentFoo');
   });
+
+  it('standalone context is unsubscribed once the child disconnects', async () => {
+    const el = await clientSideRender(parentEl, componentPath, {});
+    const contextParentDetached = querySelectorDeep('x-context-parent-detached', el);
+
+    expect(contextParentDetached.parentState.subscribers.size).toBe(1);
+
+    contextParentDetached.hideChild = true;
+    await freshRender();
+
+    expect(contextParentDetached.parentState.subscribers.size).toBe(0);
+  });
 });
diff --git a/packages/@lwc/state/src/standalone-context.ts b/packages/@lwc/state/src/standalone-context.ts
index c09bd10..8a78434 100644
--- a/packages/@lwc/state/src/standalone-context.ts
+++ b/packages/@lwc/state/src/standalone-context.ts
@@ -1,4 +1,4 @@
-import { connectContext } from './shared.js';
+import { connectContext, disconnectContext } from './shared.js';
 import { SignalBaseClass, type Signal } from '@lwc/signals';
 import type { ExposedUpdater } from './types.js';
 import type { ContextRuntimeAdapter } from './runtime-interface.js';
@@ -13,7 +13,6 @@ class ConsumedContextSignal<StateShape extends ValidStateShape>
 {
   private desiredStateDef: ValidStateDef<StateShape>;
   private _value: StateShape | null = null;
-  // Currently unused. Should be called once `disconnectContext` is implemented.
   private unsubscribe: () => void = () => {};
 
   constructor(stateDef: ValidStateDef<StateShape>) {
@@ -44,6 +43,15 @@ class ConsumedContextSignal<StateShape extends ValidStateShape>
       },
     );
   }
+
+  [disconnectContext](_componentId: ContextRuntimeAdapter<object>['component']) {
+    // Unlike the state manager's fromContext which can subscribe to multiple
+    // ancestor contexts simultaneously, this standalone version only subscribes
+    // to a single context at a time. Therefore, we don't need to use componentId
+    // to track subscriptions.
+    this.unsubscribe();
+    this.unsubscribe = () => {};
+  }
 }
 
 export const fromContext = <