@W-17388345 - add context support to ssrv2 component lifecycle

packages/@lwc/integration-karma/scripts/karma-plugins/hydration-tests.js

@@ -279,12 +279,6 @@ function createHCONFIG2JSPreprocessor(config, logger, emitter) {
             const location = path.relative(basePath, filePath);
             log.error('Error processing “%s”\n\n%s\n', location, error.stack || error.message);
             done(error, null);
-        } finally {
-            if (requiredFeatureFlags) {
-                requiredFeatureFlags.forEach((featureFlag) => {
-                    lwcSsr.setFeatureFlagForTest(featureFlag, false);
-                });
-            }
         }
     };
 }

packages/@lwc/integration-karma/test-hydration/context/x/main/main.js

@@ -2,7 +2,8 @@ import { LightningElement, api } from 'lwc';
 import { defineMalformedContext } from 'x/contextManager';
 export default class Root extends LightningElement {
     @api showTree = false;
-    malformedContext = defineMalformedContext()();
+    // Only test in CSR right now as SSR throws which prevents content from being rendered. There is additional fixtures ssr coverage for this case.
+    malformedContext = typeof window !== 'undefined' ? defineMalformedContext()() : undefined;
 
     connectedCallback() {
         this.showTree = true;

packages/@lwc/integration-karma/test-hydration/context/x/tooMuchContext/tooMuchContext.js

@@ -3,5 +3,9 @@ import { LightningElement } from 'lwc';
 
 export default class TooMuchContext extends LightningElement {
     context = grandparentContextFactory('grandparent provided value');
-    tooMuch = grandparentContextFactory('this world is not big enough for me');
+    // Only test in CSR right now as it throws in SSR. There is additional fixtures ssr coverage for this case.
+    tooMuch =
+        typeof window !== 'undefined'
+            ? grandparentContextFactory('this world is not big enough for me')
+            : undefined;
 }

packages/@lwc/ssr-runtime/package.json

@@ -50,6 +50,7 @@
     "devDependencies": {
         "@lwc/shared": "8.20.0",
         "@lwc/engine-core": "8.20.0",
+        "@lwc/features": "8.20.0",
         "observable-membrane": "2.0.0"
     }
 }

packages/@lwc/ssr-runtime/src/__tests__/context.spec.ts

@@ -0,0 +1,174 @@
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+import * as shared from '@lwc/shared';
+import { connectContext, ContextBinding } from '../context';
+import { LightningElement, SYMBOL__CONTEXT_VARIETIES } from '../lightning-element';
+import { getContextfulStack } from '../wire';
+import type { Signal } from '@lwc/signals';
+import type { ContextProvidedCallback } from '@lwc/shared';
+
+// Mock the wire module
+vi.mock('../wire', () => ({
+    getContextfulStack: vi.fn(),
+}));
+
+// Mock @lwc/shared
+vi.mock('@lwc/shared', () => ({
+    isTrustedContext: vi.fn(),
+    getContextKeys: vi.fn(),
+    isUndefined: vi.fn(),
+    keys: vi.fn(),
+    entries: vi.fn().mockReturnValue([]),
+    AriaAttrNameToPropNameMap: {},
+    defineProperties: vi.fn(),
+    assign: vi.fn(),
+    ArrayFilter: {
+        call: vi.fn(),
+    },
+}));
+
+describe('context', () => {
+    let mockContextKeys: { connectContext: string };
+    let mockContextfulStack: LightningElement[];
+    let mockContextVariety: object;
+    let mockContextSignal: Signal<unknown>;
+    let mockContextProvidedCallback: ContextProvidedCallback;
+
+    beforeEach(() => {
+        // Reset mocks
+        vi.clearAllMocks();
+
+        // Setup mock context keys
+        mockContextKeys = {
+            connectContext: 'connectContext',
+        };
+
+        // Setup mock contextful stack
+        mockContextfulStack = [new LightningElement({ tagName: 'div' })];
+        (getContextfulStack as any).mockReturnValue(mockContextfulStack);
+
+        // Setup mock context variety and signal
+        mockContextVariety = {};
+        mockContextSignal = { value: undefined } as Signal<unknown>;
+        mockContextProvidedCallback = vi.fn();
+
+        // Mock global context keys
+        (global as any).__LWC_CONTEXT_KEYS__ = mockContextKeys;
+
+        // Mock @lwc/shared functions
+        (shared.getContextKeys as any).mockReturnValue(mockContextKeys);
+        (shared.isUndefined as any).mockImplementation((val: unknown) => val === undefined);
+        (shared.keys as any).mockReturnValue(['contextful']);
+        (shared.isTrustedContext as any).mockReturnValue(true);
+        (shared.ArrayFilter.call as any).mockReturnValue(['contextful']);
+    });
+
+    describe('connectContext', () => {
+        it('should do nothing if context keys are undefined', () => {
+            (shared.getContextKeys as any).mockReturnValue(undefined);
+            const component = new LightningElement({ tagName: 'div' });
+            connectContext(component);
+            expect(getContextfulStack).not.toHaveBeenCalled();
+        });
+
+        it('should do nothing if no contextful keys are found', () => {
+            (shared.ArrayFilter.call as any).mockReturnValue([]);
+            const component = new LightningElement({ tagName: 'div' });
+            connectContext(component);
+            expect(getContextfulStack).not.toHaveBeenCalled();
+        });
+
+        it('should connect context when contextful keys are found', () => {
+            const component = new LightningElement({ tagName: 'div' });
+            const mockContextful = {
+                [mockContextKeys.connectContext]: vi.fn(),
+            };
+            Object.defineProperty(component, 'contextful', {
+                value: mockContextful,
+                configurable: true,
+            });
+
+            connectContext(component);
+
+            expect(mockContextful[mockContextKeys.connectContext]).toHaveBeenCalled();
+        });
+
+        it('should throw error in development when connecting context fails', () => {
+            const originalNodeEnv = process.env.NODE_ENV;
+            process.env.NODE_ENV = 'development';
+
+            const component = new LightningElement({ tagName: 'div' });
+            const mockContextful = {
+                [mockContextKeys.connectContext]: vi.fn().mockImplementation(() => {
+                    throw new Error('Connection failed');
+                }),
+            };
+            Object.defineProperty(component, 'contextful', {
+                value: mockContextful,
+                configurable: true,
+            });
+
+            expect(() => connectContext(component)).toThrow(
+                'Attempted to connect to trusted context'
+            );
+
+            process.env.NODE_ENV = originalNodeEnv;
+        });
+    });
+
+    describe('ContextBinding', () => {
+        it('should provide context successfully', () => {
+            const component = new LightningElement({ tagName: 'div' });
+            const binding = new ContextBinding(component);
+
+            binding.provideContext(mockContextVariety, mockContextSignal);
+
+            expect(component[SYMBOL__CONTEXT_VARIETIES].has(mockContextVariety)).toBe(true);
+            expect(component[SYMBOL__CONTEXT_VARIETIES].get(mockContextVariety)).toBe(
+                mockContextSignal
+            );
+        });
+
+        it('should not allow multiple contexts of the same variety in development', () => {
+            const originalNodeEnv = process.env.NODE_ENV;
+            process.env.NODE_ENV = 'development';
+
+            const component = new LightningElement({ tagName: 'div' });
+            const binding = new ContextBinding(component);
+
+            binding.provideContext(mockContextVariety, mockContextSignal);
+            expect(() => {
+                binding.provideContext(mockContextVariety, mockContextSignal);
+            }).toThrow('Multiple contexts of the same variety were provided');
+
+            process.env.NODE_ENV = originalNodeEnv;
+        });
+
+        it('should consume context from ancestor', () => {
+            const component = new LightningElement({ tagName: 'div' });
+            const binding = new ContextBinding(component);
+
+            // Setup ancestor with context
+            const ancestor = mockContextfulStack[0];
+            ancestor[SYMBOL__CONTEXT_VARIETIES] = new Map([
+                [mockContextVariety, mockContextSignal],
+            ]);
+
+            binding.consumeContext(mockContextVariety, mockContextProvidedCallback);
+
+            expect(mockContextProvidedCallback).toHaveBeenCalledWith(mockContextSignal);
+        });
+
+        it('should not call callback if context is not found in ancestors', () => {
+            const component = new LightningElement({ tagName: 'div' });
+            const binding = new ContextBinding(component);
+
+            // Setup empty ancestor
+            const ancestor = mockContextfulStack[0];
+            ancestor[SYMBOL__CONTEXT_VARIETIES] = new Map();
+
+            binding.consumeContext(mockContextVariety, mockContextProvidedCallback);
+
+            expect(mockContextProvidedCallback).not.toHaveBeenCalled();
+        });
+    });
+});

packages/@lwc/ssr-runtime/src/context.ts

@@ -0,0 +1,90 @@
+/*
+ * 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 {
+    type ContextProvidedCallback,
+    type ContextBinding as IContextBinding,
+    isTrustedContext,
+    getContextKeys,
+    isUndefined,
+    keys,
+    ArrayFilter,
+} from '@lwc/shared';
+import { getContextfulStack } from './wire';
+import { type LightningElement, SYMBOL__CONTEXT_VARIETIES } from './lightning-element';
+import type { Signal } from '@lwc/signals';
+
+class ContextBinding<C extends LightningElement> implements IContextBinding<LightningElement> {
+    component: C;
+
+    constructor(component: C) {
+        this.component = component;
+    }
+
+    provideContext<V extends object>(
+        contextVariety: V,
+        providedContextSignal: Signal<unknown>
+    ): void {
+        const contextVarieties = this.component[SYMBOL__CONTEXT_VARIETIES];
+        if (contextVarieties.has(contextVariety)) {
+            if (process.env.NODE_ENV !== 'production') {
+                throw new Error('Multiple contexts of the same variety were provided.');
+            }
+            return;
+        }
+        contextVarieties.set(contextVariety, providedContextSignal);
+    }
+
+    consumeContext<V extends object>(
+        contextVariety: V,
+        contextProvidedCallback: ContextProvidedCallback
+    ): void {
+        const contextfulStack = getContextfulStack(this.component);
+        for (const ancestor of contextfulStack) {
+            // If the ancestor has the specified context variety, consume it and stop searching
+            const ancestorContextVarieties = ancestor[SYMBOL__CONTEXT_VARIETIES];
+            if (ancestorContextVarieties.has(contextVariety)) {
+                contextProvidedCallback(ancestorContextVarieties.get(contextVariety));
+                break;
+            }
+        }
+    }
+}
+
+export { ContextBinding };
+
+export function connectContext(le: LightningElement) {
+    const contextKeys = getContextKeys();
+
+    if (isUndefined(contextKeys)) {
+        return;
+    }
+
+    const { connectContext } = contextKeys;
+
+    const enumerableKeys = keys(le);
+    const contextfulKeys = ArrayFilter.call(enumerableKeys, (enumerableKey) =>
+        isTrustedContext((le as any)[enumerableKey])
+    );
+
+    if (contextfulKeys.length === 0) {
+        return;
+    }
+
+    try {
+        for (let i = 0; i < contextfulKeys.length; i++) {
+            (le as any)[contextfulKeys[i]][connectContext](new ContextBinding(le));
+        }
+    } catch (err: any) {
+        if (process.env.NODE_ENV !== 'production') {
+            throw new Error(
+                `Attempted to connect to trusted context but received the following error: ${
+                    err.message
+                }`
+            );
+        }
+    }
+}

packages/@lwc/ssr-runtime/src/index.ts

@@ -14,8 +14,13 @@ export {
     sanitizeHtmlContent,
     normalizeClass,
     normalizeTabIndex,
+    setContextKeys,
+    setTrustedSignalSet,
+    setTrustedContextSet,
 } from '@lwc/shared';
 
+export { setFeatureFlag, setFeatureFlagForTest } from '@lwc/features';
+
 export { ClassList } from './class-list';
 export {
     LightningElement,

packages/@lwc/ssr-runtime/src/lightning-element.ts

@@ -29,13 +29,16 @@ import { ClassList } from './class-list';
 import { mutationTracker } from './mutation-tracker';
 import { descriptors as reflectionDescriptors } from './reflection';
 import { getReadOnlyProxy } from './get-read-only-proxy';
+import { connectContext } from './context';
 import type { Attributes, Properties } from './types';
 import type { Stylesheets } from '@lwc/shared';
+import type { Signal } from '@lwc/signals';
 
 type EventListenerOrEventListenerObject = unknown;
 type AddEventListenerOptions = unknown;
 type EventListenerOptions = unknown;
 type ShadowRoot = unknown;
+type ContextVarieties = Map<unknown, Signal<unknown>>;
 
 export type LightningElementConstructor = typeof LightningElement;
 
@@ -46,6 +49,7 @@ interface PropsAvailableAtConstruction {
 export const SYMBOL__SET_INTERNALS = Symbol('set-internals');
 export const SYMBOL__GENERATE_MARKUP = Symbol('generate-markup');
 export const SYMBOL__DEFAULT_TEMPLATE = Symbol('default-template');
+export const SYMBOL__CONTEXT_VARIETIES = Symbol('context-varieties');
 
 export class LightningElement implements PropsAvailableAtConstruction {
     static renderMode?: 'light' | 'shadow';
@@ -73,6 +77,7 @@ export class LightningElement implements PropsAvailableAtConstruction {
     #props!: Properties;
     #attrs!: Attributes;
     #classList: ClassList | null = null;
+    [SYMBOL__CONTEXT_VARIETIES]: ContextVarieties = new Map();
 
     constructor(propsAvailableAtConstruction: PropsAvailableAtConstruction & Properties) {
         assign(this, propsAvailableAtConstruction);
@@ -82,6 +87,11 @@ export class LightningElement implements PropsAvailableAtConstruction {
         this.#props = props;
         this.#attrs = attrs;
 
+        if (lwcRuntimeFlags.ENABLE_EXPERIMENTAL_SIGNALS) {
+            // Setup context before connected callback is executed
+            connectContext(this);
+        }
+
         // Class should be set explicitly to avoid it being overridden by connectedCallback classList mutation.
         if (attrs.class) {
             this.className = attrs.class;

packages/@lwc/ssr-runtime/src/stubs.ts

@@ -42,12 +42,6 @@ export function registerTemplate(..._: unknown[]): never {
 export function sanitizeAttribute(..._: unknown[]): never {
     throw new Error('sanitizeAttribute cannot be used in SSR context.');
 }
-export function setFeatureFlag(..._: unknown[]): never {
-    throw new Error('setFeatureFlag cannot be used in SSR context.');
-}
-export function setFeatureFlagForTest(..._: unknown[]): never {
-    throw new Error('setFeatureFlagForTest cannot be used in SSR context.');
-}
 export function swapComponent(..._: unknown[]): never {
     throw new Error('swapComponent cannot be used in SSR context.');
 }
@@ -66,12 +60,6 @@ export function unwrap(..._: unknown[]): never {
 export function wire(..._: unknown[]): never {
     throw new Error('@wire cannot be used in SSR context.');
 }
-export function setContextKeys(..._: unknown[]): never {
-    throw new Error('@setContextKeys cannot be used in SSR context.');
-}
-export function setTrustedContextSet(..._: unknown[]): never {
-    throw new Error('setTrustedContextSet cannot be used in SSR context.');
-}
 
 export const renderer = {
     isSyntheticShadowDefined: false,