feat(store): enforce exact return types in on() and updater() callbacks

modules/component-store/spec/types/component-store.types.spec.ts

@@ -345,5 +345,201 @@ describe('ComponentStore types', () => {
         componentStore.updater((state, v: string) => ({ ...state }))(number$);
       });
     });
+
+    describe('catches excess properties', () => {
+      it('when extra property is returned with spread', () => {
+        const componentStore = new ComponentStore({
+          prop: 'init',
+          prop2: 'yeah!',
+        });
+        componentStore.updater(
+          // @ts-expect-error updater callback return type must exactly match the state type. Remove excess properties.
+          (state, v: string) => ({ ...state, extraProp: 'bad' })
+        );
+      });
+
+      it('when extra property is returned with explicit object', () => {
+        const componentStore = new ComponentStore({
+          prop: 'init',
+          prop2: 'yeah!',
+        });
+        componentStore.updater(
+          // @ts-expect-error updater callback return type must exactly match the state type. Remove excess properties.
+          (state, v: string) => ({
+            prop: v,
+            prop2: state.prop2,
+            extraProp: 'bad',
+          })
+        );
+      });
+
+      it('when extra property is returned from void updater', () => {
+        const componentStore = new ComponentStore({
+          prop: 'init',
+          prop2: 'yeah!',
+        });
+        componentStore.updater(
+          // @ts-expect-error updater callback return type must exactly match the state type. Remove excess properties.
+          (state) => ({ ...state, extraProp: true })
+        );
+      });
+
+      it('when required property is missing', () => {
+        const componentStore = new ComponentStore({
+          prop: 'init',
+          prop2: 'yeah!',
+        });
+        componentStore.updater(
+          // @ts-expect-error Property 'prop2' is missing in type '{ prop: string; }'
+          (state, v: string) => ({ prop: v })
+        );
+      });
+
+      it('when property has wrong type', () => {
+        const componentStore = new ComponentStore({
+          prop: 'init',
+          prop2: 'yeah!',
+        });
+        componentStore.updater(
+          // @ts-expect-error Type 'number' is not assignable to type 'string'
+          (state, v: string) => ({ ...state, prop: 123 })
+        );
+      });
+
+      it('allows spread with override', () => {
+        const componentStore = new ComponentStore({
+          prop: 'init',
+          prop2: 'yeah!',
+        });
+        const sub = componentStore.updater((state, v: string) => ({
+          ...state,
+          prop: v,
+        }))('test');
+        expectTypeOf(sub).toEqualTypeOf<Subscription>();
+      });
+
+      it('allows full explicit return matching all state keys', () => {
+        const componentStore = new ComponentStore({
+          prop: 'init',
+          prop2: 'yeah!',
+        });
+        const sub = componentStore.updater((state, v: string) => ({
+          prop: v,
+          prop2: state.prop2,
+        }))('test');
+        expectTypeOf(sub).toEqualTypeOf<Subscription>();
+      });
+
+      it('allows void updater with spread return', () => {
+        const componentStore = new ComponentStore({
+          prop: 'init',
+          prop2: 'yeah!',
+        });
+        const v = componentStore.updater((state) => ({
+          ...state,
+          prop: 'updated',
+        }))();
+        expectTypeOf(v).toBeVoid();
+      });
+
+      it('allows direct state return', () => {
+        const componentStore = new ComponentStore({
+          prop: 'init',
+          prop2: 'yeah!',
+        });
+        const v = componentStore.updater((state) => state)();
+        expectTypeOf(v).toBeVoid();
+      });
+    });
+
+    describe('with a generic state type parameter', () => {
+      class GenericStore<T extends { id: string }> extends ComponentStore<T> {
+        // Keys guaranteed by the constraint can be overridden via spread.
+        readonly setIdViaSpread = this.updater((state, id: string) => ({
+          ...state,
+          id,
+        }));
+
+        readonly returnState = this.updater((state) => state);
+
+        readonly copyState = this.updater((state) => ({ ...state }));
+
+        readonly replaceViaDirectReturn = this.updater(
+          (_state, next: T) => next
+        );
+
+        readonly setIdViaAssertion = this.updater(
+          (state, id: string) => ({ ...state, id }) as T
+        );
+
+        // Keys not guaranteed by the constraint are excess properties.
+        readonly setExtra = this.updater(
+          // @ts-expect-error updater callback return type must exactly match the state type. Remove excess properties.
+          (state, extra: string) => ({ ...state, extra })
+        );
+      }
+
+      class ListStore<Item> extends ComponentStore<{
+        items: Item[];
+        loading: boolean;
+      }> {
+        readonly setItems = this.updater((state, items: Item[]) => ({
+          ...state,
+          items,
+          loading: false,
+        }));
+      }
+
+      class ExtendableStore<T extends object> extends ComponentStore<
+        { loading: boolean } & T
+      > {
+        readonly setLoading = this.updater((state) => ({
+          ...state,
+          loading: true,
+        }));
+      }
+
+      it('allows overriding known keys and catches excess keys', () => {
+        expectTypeOf(GenericStore).toBeConstructibleWith({ id: '1' });
+        expectTypeOf(ListStore).toBeConstructibleWith({
+          items: [],
+          loading: false,
+        });
+        expectTypeOf(ExtendableStore).toBeConstructibleWith({ loading: false });
+      });
+    });
+
+    describe('with union state', () => {
+      type FetchState =
+        { status: 'idle' } | { status: 'loaded'; data: string[] };
+
+      it('allows switching to a union member with its own keys', () => {
+        const componentStore = new ComponentStore<FetchState>({
+          status: 'idle',
+        });
+        componentStore.updater((_, data: string[]) => ({
+          status: 'loaded' as const,
+          data,
+        }));
+      });
+
+      it('catches keys that exist on no union member', () => {
+        const componentStore = new ComponentStore<FetchState>({
+          status: 'idle',
+        });
+        componentStore.updater(
+          // @ts-expect-error updater callback return type must exactly match the state type. Remove excess properties.
+          () => ({ status: 'idle' as const, bogus: true })
+        );
+      });
+    });
+
+    it('catches excess properties returned from one branch of a conditional', () => {
+      const componentStore = new ComponentStore({ prop: 'init' });
+      componentStore.updater(
+        // @ts-expect-error updater callback return type must exactly match the state type. Remove excess properties.
+        (state, flag: boolean) => (flag ? { ...state, extra: true } : state)
+      );
+    });
   });
 });

modules/component-store/spec/types/regression.types.spec.ts

@@ -30,4 +30,26 @@ describe('regression component-store', () => {
       }
     }
   });
+
+  describe('updater exact return type', () => {
+    it('should work with state containing optional properties', () => {
+      const store = new ComponentStore<{ req: string; opt?: number }>({
+        req: 'a',
+      });
+      store.updater((state) => ({ req: 'b' }))();
+    });
+
+    it('should work with state containing index signature', () => {
+      const store = new ComponentStore<{ [key: string]: number }>({});
+      store.updater((state, v: number) => ({ ...state, newKey: v }))(5);
+    });
+
+    it('should catch excess properties with concrete state type', () => {
+      const store = new ComponentStore<{ name: string }>({ name: 'test' });
+      store.updater(
+        // @ts-expect-error updater callback return type must exactly match the state type. Remove excess properties.
+        (state, v: string) => ({ ...state, name: v, extra: true })
+      );
+    });
+  });
 });

modules/component-store/src/component-store.ts

@@ -40,6 +40,33 @@ import {
 import { isOnStateInitDefined, isOnStoreInitDefined } from './lifecycle_hooks';
 import { toSignal } from '@angular/core/rxjs-interop';
 
+const excessPropertiesAreNotAllowedMsg =
+  'updater callback return type must exactly match the state type. Remove excess properties.';
+type ExcessPropertiesAreNotAllowed = typeof excessPropertiesAreNotAllowedMsg;
+
+type KeysOfUnion<T> = T extends unknown ? keyof T : never;
+// Keys a returned object may carry: any key of any union member, plus numeric keys for string index signatures.
+type AllowedStateKeys<S> =
+  keyof S | KeysOfUnion<S> | (string extends keyof S ? number : never);
+// `K` is inferred from the returned object's keys and constrained to AllowedStateKeys<S>. Constraint checks use
+// assignability, which respects generic constraints (e.g. `T extends { loading: boolean }`). When a key is not
+// allowed, inference falls back to the constraint and the leftover keys are typed as the error message.
+type ExactStateReturn<R, S, K, Msg extends string> = [R] extends [
+  null | undefined,
+]
+  ? unknown
+  : // Fast path: no key outside the state's keys (also resolves for identical generic types, e.g. `return state`)
+    [Exclude<KeysOfUnion<R>, AllowedStateKeys<S>>] extends [never]
+    ? unknown
+    : { [P in K & PropertyKey]?: unknown } & {
+        [
+          P in Exclude<
+            KeysOfUnion<R>,
+            K | AllowedStateKeys<S>
+          > as `${Msg} Excess property: ${P & string}`
+        ]: never;
+      };
+
 export interface SelectConfig<T = unknown> {
   debounce?: boolean;
   equal?: ValueEqualityFn<T>;
@@ -132,7 +159,16 @@ export class ComponentStore<T extends object> implements OnDestroy {
     ReturnType = OriginType extends void
       ? () => void
       : (observableOrValue: ValueType | Observable<ValueType>) => Subscription,
-  >(updaterFn: (state: T, value: OriginType) => T): ReturnType {
+    // Captures the actual return type to enforce exact state shape
+    R extends T = T,
+    // Keys of the returned object, vetted against the state's keys via the constraint
+    K extends AllowedStateKeys<T> = never,
+  >(
+    updaterFn: (
+      state: T,
+      value: OriginType
+    ) => R & ExactStateReturn<R, T, K, ExcessPropertiesAreNotAllowed>
+  ): ReturnType {
     return ((
       observableOrValue?: OriginType | Observable<OriginType>
     ): Subscription => {
@@ -210,9 +246,7 @@ export class ComponentStore<T extends object> implements OnDestroy {
    */
   patchState(
     partialStateOrUpdaterFn:
-      | Partial<T>
-      | Observable<Partial<T>>
-      | ((state: T) => Partial<T>)
+      Partial<T> | Observable<Partial<T>> | ((state: T) => Partial<T>)
   ): void {
     const patchedState =
       typeof partialStateOrUpdaterFn === 'function'
@@ -379,9 +413,8 @@ export class ComponentStore<T extends object> implements OnDestroy {
     // This type quickly became part of effect 'API'
     ProvidedType = void,
     // The actual origin$ type, which could be unknown, when not specified
-    OriginType extends
-      | Observable<ProvidedType>
-      | unknown = Observable<ProvidedType>,
+    OriginType extends Observable<ProvidedType> | unknown =
+      Observable<ProvidedType>,
     // Unwrapped actual type of the origin$ Observable, after default was applied
     ObservableType = OriginType extends Observable<infer A> ? A : never,
     // Return either an optional callback or a function requiring specific types as inputs