Fixed the issue with contextually typed return type of an async function when unions are involved

Author: Andarist

src/compiler/checker.ts

@@ -26043,12 +26043,10 @@ namespace ts {
                             return undefined;
                         }
                         contextualReturnType = iterationTypes.returnType;
-                        // falls through to unwrap Promise for AsyncGenerators
                     }
 
                     if (functionFlags & FunctionFlags.Async) { // Async function or AsyncGenerator function
-                        // Get the awaited type without the `Awaited<T>` alias
-                        const contextualAwaitedType = mapType(contextualReturnType, getAwaitedTypeNoAlias);
+                        const contextualAwaitedType = mapType(contextualReturnType, functionFlags & FunctionFlags.Generator ? getAwaitedTypeNoAlias : getAwaitedTypeOfPromise);
                         return contextualAwaitedType && getUnionType([contextualAwaitedType, createPromiseLikeType(contextualAwaitedType)]);
                     }
 

tests/baselines/reference/contextuallyTypeAsyncFunctionReturnTypeFromUnion.errors.txt

@@ -20,26 +20,26 @@ createMachine<{ count: number }>({
 >createMachine<{ count: number }>({  services: {    test: async () => Promise.reject("some err"),    async test2() {      return Promise.reject("some err");    },  },}) : void
 >createMachine : <T>(implementations: { services: Record<string, () => Promise<T> | StateMachine<T>>; }) => void
 >count : number
->{  services: {    test: async () => Promise.reject("some err"),    async test2() {      return Promise.reject("some err");    },  },} : { services: { test: () => Promise<{ count: number; } | StateMachine<{ count: number; }>>; test2(): Promise<{ count: number; } | StateMachine<{ count: number; }>>; }; }
+>{  services: {    test: async () => Promise.reject("some err"),    async test2() {      return Promise.reject("some err");    },  },} : { services: { test: () => Promise<{ count: number; }>; test2(): Promise<{ count: number; }>; }; }
 
   services: {
->services : { test: () => Promise<{ count: number; } | StateMachine<{ count: number; }>>; test2(): Promise<{ count: number; } | StateMachine<{ count: number; }>>; }
->{    test: async () => Promise.reject("some err"),    async test2() {      return Promise.reject("some err");    },  } : { test: () => Promise<{ count: number; } | StateMachine<{ count: number; }>>; test2(): Promise<{ count: number; } | StateMachine<{ count: number; }>>; }
+>services : { test: () => Promise<{ count: number; }>; test2(): Promise<{ count: number; }>; }
+>{    test: async () => Promise.reject("some err"),    async test2() {      return Promise.reject("some err");    },  } : { test: () => Promise<{ count: number; }>; test2(): Promise<{ count: number; }>; }
 
     test: async () => Promise.reject("some err"),
->test : () => Promise<{ count: number; } | StateMachine<{ count: number; }>>
->async () => Promise.reject("some err") : () => Promise<{ count: number; } | StateMachine<{ count: number; }>>
->Promise.reject("some err") : Promise<{ count: number; } | StateMachine<{ count: number; }>>
+>test : () => Promise<{ count: number; }>
+>async () => Promise.reject("some err") : () => Promise<{ count: number; }>
+>Promise.reject("some err") : Promise<{ count: number; }>
 >Promise.reject : <T = never>(reason?: any) => Promise<T>
 >Promise : PromiseConstructor
 >reject : <T = never>(reason?: any) => Promise<T>
 >"some err" : "some err"
 
     async test2() {
->test2 : () => Promise<{ count: number; } | StateMachine<{ count: number; }>>
+>test2 : () => Promise<{ count: number; }>
 
       return Promise.reject("some err");
->Promise.reject("some err") : Promise<{ count: number; } | StateMachine<{ count: number; }>>
+>Promise.reject("some err") : Promise<{ count: number; }>
 >Promise.reject : <T = never>(reason?: any) => Promise<T>
 >Promise : PromiseConstructor
 >reject : <T = never>(reason?: any) => Promise<T>
@@ -55,8 +55,8 @@ function fn1(): () => Promise<{ count: number }> | StateMachine<{ count: number
 >count : number
 
   return async () => Promise.reject('some err')
->async () => Promise.reject('some err') : () => Promise<{ count: number; } | StateMachine<{ count: number; }>>
->Promise.reject('some err') : Promise<{ count: number; } | StateMachine<{ count: number; }>>
+>async () => Promise.reject('some err') : () => Promise<{ count: number; }>
+>Promise.reject('some err') : Promise<{ count: number; }>
 >Promise.reject : <T = never>(reason?: any) => Promise<T>
 >Promise : PromiseConstructor
 >reject : <T = never>(reason?: any) => Promise<T>