I am creating a state management solution for building a form wizard (multi-step form) in react and typescript. I am super close in getting a fully type-safe configuration helper. Almost all of the type inference is working except for the getNextStep property. I've included a code sandbox link so you can inspect the types. Hovering over the 'data' argument on line 44, the type is 'any' where it should be the inferred zod type of the schema property on the same object. For example, in Step1, the type for data in getNextStep should resolve to { userId: string }. The types are setup this way so that the component props can also be inferred based on which step is being accessed. An example can be seen in the Test component at the bottom. I need to be able to infer the correct component props so I can render the component for the step in a type-safe way. Any help is appreciated here.

CodeSanbox Link

import { Context, ComponentType, createContext, useContext } from "react";
import { z } from "zod";

// Step type with explicit generic for component props
type Step<
  TKey = string,
  TSchema extends z.ZodType = z.ZodType,
  TProps = never
> = {
  schema: TSchema;
  component: ComponentType<TProps>;
  getNextStep?: (data: z.infer<TSchema>) => TKey;
};

type Config<TKey extends string = string> = {
  [K in TKey]: Step<Exclude<TKey, K>>;
};

type InferData<TConfig> = {
  [K in keyof TConfig]: TConfig[K] extends { schema?: z.ZodType }
    ? z.infer<NonNullable<TConfig[K]["schema"]>>
    : unknown;
};

function createConfig<TConfig extends Config<keyof TConfig & string>>(
  config: TConfig
) {
  const Context = createContext<TConfig | null>(null);

  const useConfig = createUseConfig(Context);
  const useConfigStep = createUseConfigStep(Context);

  return Object.freeze({
    useConfig,
    useConfigStep,
    config,
  });
}

const { useConfig, useConfigStep, config } = createConfig({
  Step1: {
    component: ({ Id }: { Id: string }) => null,
    schema: z.object({ userId: z.string() }),
    getNextStep: (data) => "Step2",
  },
  Step2: {
    component: ({ userId }: { userId: string }) => null,
    schema: z.unknown(),
  },
});

type Data = InferData<typeof config>;

function createUseConfig<TConfig extends Config<keyof TConfig & string>>(
  Context: Context<TConfig | null>
) {
  return function useConfig() {
    const store = useContext(Context);

    if (!store) {
      throw new Error(
        "useFormWizardStore must be used within a FormWizardProvider"
      );
    }

    return store;
  };
}

function createUseConfigStep<TConfig extends Config<keyof TConfig & string>>(
  Context: Context<TConfig | null>
) {
  return function useConfigStep<TKey extends keyof TConfig & string>(
    step: TKey
  ) {
    const store = useContext(Context);

    if (!store) {
      throw new Error(
        "useFormWizardStore must be used within a FormWizardProvider"
      );
    }

    return store[step];
  };
}

function Test() {
  const { component: Component } = useConfigStep("Step1");
  return <Component Id="" />;
}