Typescript union type of functions weirdly becomes an intersection type of the arguments

I'm trying to model my data in Typescript in the following manner. Main focus is the type MessageHandler that maps the type in a message with a callback that accepts that message. I followed the handbook chapter for mapped types

type BaseMessageBody = {
	msg_id?: number;
	in_reply_to?: number;
};

type BaseMessage<TBody> = {
	src: string;
	dest: string;
	body: BaseMessageBody & TBody;
};

type BroadcastMessage = BaseMessage<{
	type: 'broadcast';
	message: number;
}>;

type InitMessage = BaseMessage<{
	type: 'init';
	node_id: string;
	node_ids: string[];
}>;

type Message = BroadcastMessage | InitMessage;

type Body = Message['body'];

type MessageHandler = {
	[M in Message as M['body']['type']]?: (message: M) => void;
};

So far so good. The expanded type of MessageHandler is

type MessageHandler = {
    broadcast?: ((message: BroadcastMessage) => void) | undefined;
    init?: ((message: InitMessage) => void) | undefined;
}

But when I try to actually use the type like below:

const handlers: MessageHandler = {};

export const handle = (message: Message) => {
	if (message.body.type === 'init') {
		console.log('Initialized');
	}
	const handler = handlers[message.body.type];
	if (!handler) {
		console.warn('Unable to handle type', message.body.type);
		return;
	}
	handler(message); //Error here
};

I get the following error. Somehow the handler type has transformed to const handler: (message: BroadcastMessage & InitMessage) => void

error: TS2345 [ERROR]: Argument of type 'Message' is not assignable to parameter of type 'BroadcastMessage & InitMessage'.
  Type 'BroadcastMessage' is not assignable to type 'BroadcastMessage & InitMessage'.
    Type 'BroadcastMessage' is not assignable to type 'InitMessage'.
      Types of property 'body' are incompatible.
        Type 'BaseMessageBody & { type: "broadcast"; message: number; }' is not assignable to type 'BaseMessageBody & { type: "init"; node_id: string; node_ids: string[]; }'.
          Type 'BaseMessageBody & { type: "broadcast"; message: number; }' is missing the following properties from type '{ type: "init"; node_id: string; node_ids: string[]; }': node_id, node_ids
        handler(message);
                ~~~~~~~

There are slightly related questions on stack overflow but I wasn't able to resolve my problem by following them. Here is the playground with all the code.

TypeScript doesn't have direct support for what I call "correlated unions" as described in microsoft/TypeScript#30581. The compiler can only type check a block of code like handler(message) once. If the type of handler is a union type of functions like ((message: BroadcastMessage) => void) | ((message: InitMessage) => void) and the type of message is a union of arguments like BroadcastMessage | InitMessage, the compiler can't see that as safe. After all, for arbitrary handler and message variables of those types, it could be a mistake to allow the call; if message were of type InitMessage but handler were of type (message: BroadcastMessage) => void, you'd have a problem. The only safe way to call a union of functions is with an intersection of its arguments, not a union of its arguments. A union argument could turn out to be the wrong type for the union function parameter.

In your case it is, of course, impossible for that failure to occur, because the type of handler and the type of message are correlated due to them coming from the same source. But the only way to see that would be if the compiler could analyze handler(message) once for each possible narrowing of message. But it doesn't do that. So you get that error.

If you just want to suppress the error and move on, you can use a type assertion:

handler(message as BroadcastMessage & InitMessage); // 🤷‍♂️

That's technically a lie, but it is easy. But it doesn't mean the compiler sees what you're doing as type safe; if you accidentally wrote something like handler(broadcastMessageOnly as BroadcastMessage & InitMessage) (assuming broadcastMessageOnly is of type BroadcastMessage and not InitMessage) the compiler wouldn't catch the mistake. But that might not matter, as long as you're confident that you've implemented it right.

If you care about the compiler verifying type safety here, then the recommended approach to dealing with correlated unions is to refactor away from unions and toward generic indexes into simple key-value object types or mapped types over such object types. This technique is described in detail in microsoft/TypeScript#47109. For your example, the relevant changes look like this:

First let's make the basic key-value type we're going to build from:

interface MessageMap {
  broadcast: { message: number };
  init: { node_id: string; node_ids: string[] };
}

Then you can redefine your Message type to take a particular key as a type argument:

type Message<K extends keyof MessageMap = keyof MessageMap> =
  { [P in K]: BaseMessage<MessageMap[P] & { type: P }> }[K]

If you need your original message types you can recover them:

type BroadcastMessage = Message<"broadcast">;
// type BroadcastMessage = { src: string; dest: string; 
//  body: BaseMessageBody & { message: number; } & { type: "broadcast"; };
// }

type InitMessage = Message<"init">;
// type InitMessage = { src: string; dest: string; body: BaseMessageBody & 
//  { node_id: string; node_ids: string[]; } & { type: "init"; };
// }

And since Message<K> has a default type argument corresponding to the full union of keys, just Message by itself is equivalent to your original version:

type MessageTest = Message
// type MessageTest = { src: string; dest: string; 
//  body: BaseMessageBody & { message: number; } & { type: "broadcast"; };
// } | { src: string; dest: string; body: BaseMessageBody & 
//  { node_id: string; node_ids: string[]; } & { type: "init"; };
// }But 

And MessageHandler can be written in terms of MessageMap also:

type MessageHandler = {
  [K in keyof MessageMap]?: (message: Message<K>) => void;
};

Finally, you make handle a generic function that accepts a Message<K>, and the error goes away:

export const handle = <K extends keyof MessageMap>(message: Message<K>) => {

  if (message.body.type === 'init') {
    console.log('Initialized');
  }
  const handler = handlers[message.body.type];
  if (!handler) {
    console.warn('Unable to handle type', message.body.type);
    return;
  }
  handler(message); // okay
  // const handler: (message: Message<K>) => void
};

By the time you call handler, the compiler sees it as type (message: Message<K>) => void. It's no longer a union of functions; its a single function with a parameter type of Message<K>. And since message's type is also Message<K>, the call is allowed.

Is it worth refactoring to this form? If you're confident that your original version works and will continue to work, it's certainly going to be easier to just assert and move on. If you're less confident, then maybe the refactoring here is worth the effort. It depends on your use cases.

Playground link to code