feat: bake context into lwc framework

packages/@lwc/engine-core/src/framework/context.ts

@@ -0,0 +1,23 @@
+/*
+ * Copyright (c) 2024, salesforce.com, inc.
+ * All rights reserved.
+ * SPDX-License-Identifier: MIT
+ * For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/MIT
+ */
+import { ContextEventName, getContextKeys } from '@lwc/shared';
+import type { Signal } from '@lwc/signals';
+
+export type ContextProvidedCallback = (contextSignal: Signal<unknown>) => void;
+export class ContextRequestEvent extends CustomEvent<{
+    key: symbol;
+    contextVariety: unknown;
+    callback: ContextProvidedCallback;
+}> {
+    constructor(detail: { contextVariety: unknown; callback: ContextProvidedCallback }) {
+        super(ContextEventName, {
+            bubbles: true,
+            composed: true,
+            detail: { ...detail, key: getContextKeys().contextEventKey },
+        });
+    }
+}

packages/@lwc/engine-core/src/framework/main.ts

@@ -72,5 +72,5 @@ export { default as wire } from './decorators/wire';
 export { readonly } from './readonly';
 
 export { setFeatureFlag, setFeatureFlagForTest } from '@lwc/features';
-export { setTrustedSignalSet } from '@lwc/shared';
+export { setContextKeys, setTrustedSignalSet } from '@lwc/shared';
 export type { Stylesheet, Stylesheets } from '@lwc/shared';

packages/@lwc/engine-core/src/framework/vm.ts

@@ -11,7 +11,9 @@ import {
     assert,
     create,
     defineProperty,
+    getPrototypeOf,
     getOwnPropertyNames,
+    keys,
     isArray,
     isFalse,
     isFunction,
@@ -20,6 +22,8 @@ import {
     isTrue,
     isUndefined,
     flattenStylesheets,
+    getContextKeys,
+    ContextEventName,
 } from '@lwc/shared';
 
 import { addErrorComponentStack } from '../shared/error';
@@ -49,6 +53,8 @@ import { flushMutationLogsForVM, getAndFlushMutationLogs } from './mutation-logg
 import { connectWireAdapters, disconnectWireAdapters, installWireAdapters } from './wiring';
 import { VNodeType, isVFragment } from './vnodes';
 import { isReportingEnabled, report, ReportingEventId } from './reporting';
+import { type ContextProvidedCallback, ContextRequestEvent } from './context';
+
 import type { VNodes, VCustomElement, VNode, VBaseElement, VStaticPartElement } from './vnodes';
 import type { ReactiveObserver } from './mutation-tracker';
 import type {
@@ -60,6 +66,7 @@ import type { ComponentDef } from './def';
 import type { Template } from './template';
 import type { HostNode, HostElement, RendererAPI } from './renderer';
 import type { Stylesheet, Stylesheets, APIVersion } from '@lwc/shared';
+import type { Signal } from '@lwc/signals';
 
 type ShadowRootMode = 'open' | 'closed';
 
@@ -699,6 +706,9 @@ export function runConnectedCallback(vm: VM) {
     if (hasWireAdapters(vm)) {
         connectWireAdapters(vm);
     }
+    // Setup context before connected callback is executed
+    setupContext(vm);
+
     const { connectedCallback } = vm.def;
     if (!isUndefined(connectedCallback)) {
         logOperationStart(OperationId.ConnectedCallback, vm);
@@ -740,10 +750,109 @@ export function runConnectedCallback(vm: VM) {
     }
 }
 
+function setupContext(vm: VM) {
+    const contextKeys = getContextKeys();
+
+    if (!contextKeys) {
+        return;
+    }
+
+    const { connectContext, contextEventKey } = contextKeys;
+    const { component } = vm;
+    const enumerableKeys = keys(getPrototypeOf(component));
+    const contextfulFieldsOrProps = enumerableKeys.filter(
+        (propName) => (component as any)[propName]?.[connectContext]
+    );
+
+    if (contextfulFieldsOrProps.length === 0) {
+        return;
+    }
+
+    let isProvidingContext = false;
+    const providedContextVarieties = new Map<unknown, Signal<unknown>>();
+    const contextRuntimeAdapter = {
+        isServerSide: false,
+        component,
+        provideContext<T extends object>(
+            contextVariety: T,
+            providedContextSignal: Signal<unknown>
+        ): void {
+            if (!isProvidingContext) {
+                isProvidingContext = true;
+
+                component.addEventListener(ContextEventName, (event: any) => {
+                    if (
+                        event.detail.key === contextEventKey &&
+                        providedContextVarieties.has(event.detail.contextVariety)
+                    ) {
+                        event.stopImmediatePropagation();
+                        const providedContextSignal = providedContextVarieties.get(
+                            event.detail.contextVariety
+                        );
+                        event.detail.callback(providedContextSignal);
+                    }
+                });
+            }
+
+            let multipleContextWarningShown = false;
+
+            if (providedContextVarieties.has(contextVariety)) {
+                if (!multipleContextWarningShown) {
+                    multipleContextWarningShown = true;
+                    logError(
+                        'Multiple contexts of the same variety were provided. Only the first context will be used.'
+                    );
+                }
+                return;
+            }
+
+            providedContextVarieties.set(contextVariety, providedContextSignal);
+        },
+        consumeContext<T extends object>(
+            contextVariety: T,
+            contextProvidedCallback: ContextProvidedCallback
+        ): void {
+            const event = new ContextRequestEvent({
+                contextVariety,
+                callback: contextProvidedCallback,
+            });
+
+            component.dispatchEvent(event);
+        },
+    };
+
+    for (const contextfulFieldsOrProp of contextfulFieldsOrProps) {
+        (component as any)[contextfulFieldsOrProp][connectContext](contextRuntimeAdapter);
+    }
+}
+
 function hasWireAdapters(vm: VM): boolean {
     return getOwnPropertyNames(vm.def.wire).length > 0;
 }
 
+function cleanupContext(vm: VM) {
+    const contextKeys = getContextKeys();
+
+    if (!contextKeys) {
+        return;
+    }
+
+    const { disconnectContext } = contextKeys;
+    const { component } = vm;
+    const enumerableKeys = keys(getPrototypeOf(component));
+    const contextfulFieldsOrProps = enumerableKeys.filter(
+        (propName) => (component as any)[propName]?.[disconnectContext]
+    );
+
+    if (contextfulFieldsOrProps.length === 0) {
+        return;
+    }
+
+    for (const contextfulField of contextfulFieldsOrProps) {
+        (component as any)[contextfulField][disconnectContext](component);
+    }
+}
+
 function runDisconnectedCallback(vm: VM) {
     if (process.env.NODE_ENV !== 'production') {
         assert.isTrue(vm.state !== VMState.disconnected, `${vm} must be inserted.`);
@@ -767,6 +876,7 @@ function runDisconnectedCallback(vm: VM) {
 
         logOperationEnd(OperationId.DisconnectedCallback, vm);
     }
+    cleanupContext(vm);
 }
 
 function runChildNodesDisconnectedCallback(vm: VM) {

packages/@lwc/engine-dom/src/index.ts

@@ -30,6 +30,7 @@ export {
     isComponentConstructor,
     parseFragment,
     parseSVGFragment,
+    setContextKeys,
     setTrustedSignalSet,
     swapComponent,
     swapStyle,

packages/@lwc/integration-karma/helpers/test-utils.js

@@ -709,6 +709,18 @@ window.TestUtils = (function (lwc, jasmine, beforeAll) {
         signalValidator.add(signal);
     }
 
+    const contextKeys = {
+        connectContext: Symbol('connectContext'),
+        disconnectContext: Symbol('disconnectContext'),
+        contextEventKey: Symbol('contextEventKey'),
+    };
+
+    lwc.setContextKeys(contextKeys);
+
+    function getContextKeys() {
+        return contextKeys;
+    }
+
     return {
         clearRegister,
         extractDataIds,
@@ -736,6 +748,7 @@ window.TestUtils = (function (lwc, jasmine, beforeAll) {
         catchUnhandledRejectionsAndErrors,
         addTrustedSignal,
         expectEquivalentDOM,
+        getContextKeys,
         ...apiFeatures,
     };
 })(LWC, jasmine, beforeAll);

packages/@lwc/integration-karma/test/component/LightningElement.connectedCallback/index.spec.js

@@ -4,6 +4,7 @@ import { customElementCallbackReactionErrorListener } from 'test-utils';
 import Test from 'x/test';
 import ConnectedCallbackThrow from 'x/connectedCallbackThrow';
 import XSlottedParent from 'x/slottedParent';
+import ContextParent from 'x/contextParent';
 
 function testConnectSlot(name, fn) {
     it(`should invoke the connectedCallback the root element is added in the DOM via ${name}`, () => {
@@ -63,3 +64,14 @@ describe('addEventListner in `connectedCallback`', () => {
         });
     });
 });
+
+describe('context', () => {
+    it('connects contextful fields when running connectedCallback', () => {
+        const elm = createElement('x-context-parent', { is: ContextParent });
+        document.body.appendChild(elm);
+        const child = elm.shadowRoot.querySelector('x-context-child');
+
+        expect(child).toBeDefined();
+        expect(child.shadowRoot.querySelector('p').textContent).toBe('Test Child: foo');
+    });
+});

packages/@lwc/integration-karma/test/component/LightningElement.connectedCallback/x/contextChild/contextChild.html

@@ -0,0 +1,3 @@
+<template>
+    <p>Test Child: {name}</p>
+</template>

packages/@lwc/integration-karma/test/component/LightningElement.connectedCallback/x/contextChild/contextChild.js

@@ -0,0 +1,10 @@
+import { LightningElement } from 'lwc';
+import { consumeStateFactory } from 'x/state';
+
+export default class TestChildSymbol extends LightningElement {
+    randomChild = consumeStateFactory();
+
+    get name() {
+        return this.randomChild.value.context?.value?.name ?? 'not available';
+    }
+}

packages/@lwc/integration-karma/test/component/LightningElement.connectedCallback/x/contextParent/contextParent.html

@@ -0,0 +1,4 @@
+<template>
+    <x-context-child>
+    </x-context-child>
+</template>

packages/@lwc/integration-karma/test/component/LightningElement.connectedCallback/x/contextParent/contextParent.js

@@ -0,0 +1,6 @@
+import { LightningElement } from 'lwc';
+import { nameStateFactory } from 'x/state';
+
+export default class ContextParent extends LightningElement {
+    random = nameStateFactory();
+}