Use generics to statically type key of record in a self-referential manner

The idea is to create a finite state machine helper for my chatbot automation library, so that the people talking to the bot can progress through different phases of a conversation.

I want the consumer of the library to provide a "state machine descriptor" to a function that would instantiate the state machine.

All the behavior is already working. I just want to improve the static typing for the consumers of the library.

This is what I want to accomplish:

const myStateMachine = createStateMachine({
    initialState: "state1",
    states: {
        state1: {
            onMessage: (requester: MessageObj, stateMachine: StateMachineInstance) => {
                requester.reply("hello1");
                stateMachine.setState("state2");
            }
        },
        state2: {
            onMessage: (requester: MessageObj, stateMachine: StateMachineInstance) => {
                requester.reply("hello2");
                stateMachine.setState("state1");
            }
        }
    }
});

Here is the problem I’m facing: the method stateMachine.setState("state2") accepts any string, not just the keys of the states provided in the descriptor. It should accept "state1" | "state2" because those are the states that the state machine would have.

I’ve tried a bunch of different things but most of them result in a typescript error. I just reverted those back to a generic string so that it would compile.

Here is what the types look like at the moment:

type StateId = string;

type State =
    {
        onMessage: (
            requester: MessageObj,
            stateMachineInstance: StateMachineInstance
        ) => any;
    }

type StateMachineDescriptor = {
    initialState: StateId;
    states: {
        [stateId: string]: State,
    }
};

type StateMachineInstance = StateMachineDescriptor & {
    currentState: StateId;
    setState: (newState: StateId) => void;
    reset: () => void;
};

>Solution :

Here is my entire thought process behind writing a possible solution:

Since we need generics to solve this, I’ll go ahead and write the signature of the function using one:

declare function createStateMachine<States extends string>(value: StateMachineDescriptor<States>): StateMachineInstance<States>;

While you could probably make it work with States being an object type, having it represent the state names is a lot easier to work with. The first problem we encounter is TypeScript inferring States from the wrong place. If we write something like this:

type StateMachineDescriptor<States extends string> = {
    initialState: States;
    states: Record<States, State<States>>;
};

Then TypeScript will infer States from initialState, which is wrong since we want initialState to be based off of states, not the other way around. We can prevent inference on initialState using a trick to move it down the list of possible inference sites:

type NoInfer<T> = [T][T extends any ? 0 : never];

Blocking inference at initialState, it works as intended:

type StateMachineDescriptor<States extends string> = {
    initialState: NoInfer<States>;
    states: Record<States, State<States>>;
};

That was really the main problem here. Now all we have to do is add in a generic parameter:

type StateMachineInstance<States extends string> = StateMachineDescriptor<States> & {
    currentState: States;
    setState: (newState: States) => void;
    reset: () => void;
};

type State<States extends string> =
    {
        onMessage: (
            requester: MessageObj,
            stateMachineInstance: StateMachineInstance<States>
        ) => any;
    }

Playground