convertToAsyncFunction: Fix two-argument then

src/services/codefixes/convertToAsyncFunction.ts

@@ -282,7 +282,7 @@ namespace ts.codefix {
             // We are about to write a 'let' variable declaration, but `transformExpression` for both
             // the try block and catch block will assign to this name. Setting this flag indicates
             // that future assignments should be written as `name = value` instead of `const name = value`.
-            possibleNameForVarDecl.hasBeenDeclared = true;
+            setHasBeenDeclared(possibleNameForVarDecl);
         }
 
         const tryBlock = createBlock(transformExpression(node.expression, transformer, possibleNameForVarDecl));
@@ -319,22 +319,44 @@ namespace ts.codefix {
     function transformThen(node: CallExpression, transformer: Transformer, prevArgName?: SynthBindingName): readonly Statement[] {
         const [onFulfilled, onRejected] = node.arguments;
         const onFulfilledArgumentName = getArgBindingName(onFulfilled, transformer);
-        const transformationBody = getTransformationBody(onFulfilled, prevArgName, onFulfilledArgumentName, node, transformer);
-
         if (onRejected) {
+            // From:
+            //   x().then(onFulfilledArg => { /* on fulfilled */ }, onRejectedArg => { /* on rejected */ })
+            //
+            // To:
+            //   let onFulfilledArg;
+            //   try {
+            //     onFulfilledArg = await x();
+            //   } catch (onRejectedArg) {
+            //     /* on rejected */
+            //     return;
+            //   }
+            //   /* on fulfilled */
             const onRejectedArgumentName = getArgBindingName(onRejected, transformer);
-            const tryBlock = createBlock(transformExpression(node.expression, transformer, onFulfilledArgumentName).concat(transformationBody));
-            const transformationBody2 = getTransformationBody(onRejected, prevArgName, onRejectedArgumentName, node, transformer);
+            const onFulfilledArgType = onFulfilledArgumentName && !transformer.isInJSFile && transformer.checker.getTypeAtLocation(getNode(onFulfilledArgumentName));
+            const onFulfilledArgTypeNode = onFulfilledArgType && !(onFulfilledArgType.flags & TypeFlags.Any) ? transformer.checker.typeToTypeNode(onFulfilledArgType) : undefined;
+            const onFulfilledArgVariableDeclaration = createVariableOrAssignmentOrExpressionStatement(onFulfilledArgumentName, /*rightHandSide*/ undefined, onFulfilledArgTypeNode);
+            const onFulfilledBody = getTransformationBody(onFulfilled, prevArgName, onFulfilledArgumentName, node, transformer);
+            const tryBlock = createBlock(transformExpression(node.expression, transformer, onFulfilledArgumentName));
+            const onRejectedBody = ensureReturnOrThrow(getTransformationBody(onRejected, prevArgName, onRejectedArgumentName, node, transformer));
             const catchArg = onRejectedArgumentName ? isSynthIdentifier(onRejectedArgumentName) ? onRejectedArgumentName.identifier.text : onRejectedArgumentName.bindingPattern : "e";
             const catchVariableDeclaration = createVariableDeclaration(catchArg);
-            const catchClause = createCatchClause(catchVariableDeclaration, createBlock(transformationBody2));
-
-            return [createTry(tryBlock, catchClause, /* finallyBlock */ undefined)];
+            const catchClause = createCatchClause(catchVariableDeclaration, createBlock(onRejectedBody));
+            return onFulfilledArgVariableDeclaration.concat(createTry(tryBlock, catchClause, /* finallyBlock */ undefined), onFulfilledBody);
         }
 
+        const transformationBody = getTransformationBody(onFulfilled, prevArgName, onFulfilledArgumentName, node, transformer);
         return transformExpression(node.expression, transformer, onFulfilledArgumentName).concat(transformationBody);
     }
 
+    function ensureReturnOrThrow(statements: readonly Statement[]): readonly Statement[] {
+        const lastStatement = lastOrUndefined(statements);
+        if (lastStatement?.kind !== SyntaxKind.ReturnStatement && lastStatement?.kind !== SyntaxKind.ThrowStatement) {
+            return statements.concat(createReturn());
+        }
+        return statements;
+    }
+
     /**
      * Transforms the 'x' part of `x.then(...)`, or the 'y()' part of `y().catch(...)`, where 'x' and 'y()' are Promises.
      */
@@ -346,17 +368,19 @@ namespace ts.codefix {
         return createVariableOrAssignmentOrExpressionStatement(prevArgName, createAwait(node), /*typeAnnotation*/ undefined);
     }
 
-    function createVariableOrAssignmentOrExpressionStatement(variableName: SynthBindingName | undefined, rightHandSide: Expression, typeAnnotation: TypeNode | undefined): readonly Statement[] {
+    function createVariableOrAssignmentOrExpressionStatement(variableName: SynthBindingName | undefined, rightHandSide: Expression | undefined, typeAnnotation: TypeNode | undefined): readonly Statement[] {
         if (!variableName || isEmptyBindingName(variableName)) {
             // if there's no argName to assign to, there still might be side effects
-            return [createExpressionStatement(rightHandSide)];
+            return rightHandSide ? [createExpressionStatement(rightHandSide)] : emptyArray;
         }
 
         if (isSynthIdentifier(variableName) && variableName.hasBeenDeclared) {
             // if the variable has already been declared, we don't need "let" or "const"
+            Debug.assertIsDefined(rightHandSide);
             return [createExpressionStatement(createAssignment(getSynthesizedDeepClone(variableName.identifier), rightHandSide))];
         }
 
+        setHasBeenDeclared(variableName);
         return [
             createVariableStatement(
                 /*modifiers*/ undefined,
@@ -365,7 +389,7 @@ namespace ts.codefix {
                         getSynthesizedDeepClone(getNode(variableName)),
                         typeAnnotation,
                         rightHandSide)],
-                    NodeFlags.Const))];
+                    rightHandSide ? NodeFlags.Const : NodeFlags.Let))];
     }
 
     function maybeAnnotateAndReturn(expressionToReturn: Expression | undefined, typeAnnotation: TypeNode | undefined): readonly Statement[] {
@@ -589,6 +613,18 @@ namespace ts.codefix {
         return every(bindingName.elements, isEmptyBindingName);
     }
 
+    function setHasBeenDeclared(bindingName: SynthBindingName) {
+        if (isSynthIdentifier(bindingName)) {
+            setIdentifierHasBeenDeclared(bindingName);
+        }
+        else {
+            bindingName.elements.forEach(setHasBeenDeclared);
+        }
+        function setIdentifierHasBeenDeclared(identifier: SynthIdentifier) {
+            identifier.hasBeenDeclared = true;
+        }
+    }
+
     function getNode(bindingName: SynthBindingName) {
         return isSynthIdentifier(bindingName) ? bindingName.identifier : bindingName.bindingPattern;
     }

src/services/suggestionDiagnostics.ts

@@ -145,20 +145,26 @@ namespace ts {
         }
 
         // ensure all chained calls are valid
-        let currentNode = node.expression;
+        let seenThen = node.expression.name.escapedText === "then";
+        let currentNode: Node = node.expression;
         while (isPromiseHandler(currentNode) || isPropertyAccessExpression(currentNode)) {
-            if (isCallExpression(currentNode) && !currentNode.arguments.every(isFixablePromiseArgument)) {
+            if (isCallExpression(currentNode) && (!currentNode.arguments.every(isFixablePromiseArgument) || isThenCallWithTwoArguments(currentNode) && seenThen)) {
                 return false;
             }
+            seenThen = seenThen || isPromiseHandler(currentNode) && currentNode.expression.name.escapedText === "then";
             currentNode = currentNode.expression;
         }
         return true;
     }
 
-    function isPromiseHandler(node: Node): node is CallExpression {
+    function isPromiseHandler(node: Node): node is CallExpression & { expression: PropertyAccessExpression } {
         return isCallExpression(node) && (hasPropertyAccessExpressionWithName(node, "then") || hasPropertyAccessExpressionWithName(node, "catch"));
     }
 
+    function isThenCallWithTwoArguments(node: CallExpression) {
+        return hasPropertyAccessExpressionWithName(node, "then") && node.arguments.length > 1;
+    }
+
     // should be kept up to date with getTransformationBody in convertToAsyncFunction.ts
     function isFixablePromiseArgument(arg: Expression): boolean {
         switch (arg.kind) {

src/testRunner/unittests/services/convertToAsyncFunction.ts

@@ -1426,6 +1426,48 @@ function [#|get|]() {
         .resolve<APIResponse<{ email: string }>>({ success: true, data: { email: "" } })
         .catch<APIResponse<{ email: string }>>(() => ({ success: false }));
 }
+`);
+
+        _testConvertToAsyncFunction("convertToAsyncFunction_twoArgumentThen", `
+function [#|fSync|]() {
+    return Promise.resolve(0).then(x => {
+        console.log(x); // note: added to illustrate refactor
+        throw new Error('Failure!');
+    }, () => null);
+}
+`);
+
+        _testConvertToAsyncFunction("convertToAsyncFunction_twoArgumentThen_onRejectedNoReturn", `
+function [#|fSync|]() {
+    return Promise.resolve(0).then(x => {
+        console.log(x); // note: added to illustrate refactor
+        throw new Error('Failure!');
+    }, () => {});
+}
+`);
+
+        _testConvertToAsyncFunction("convertToAsyncFunction_twoArgumentThen_onRejectedThrow", `
+function [#|fSync|]() {
+    return Promise.resolve(0).then(x => {
+        console.log(x); // note: added to illustrate refactor
+        throw new Error('Failure!');
+    }, () => {
+        throw new Error('Thrown from onRejected');
+    });
+}
+`);
+
+        _testConvertToAsyncFunctionFailed("convertToAsyncFunction_twoArgumentThen_chainedThen", `
+function [#|fSync|]() {
+    return Promise.resolve(0).then(x => {
+        console.log(x); // note: added to illustrate refactor
+        throw new Error('Failure!');
+    }, () => {
+        throw new Error('Thrown from onRejected');
+    }).then(() => {
+        return 0;
+    });
+}
 `);
     });
 }

tests/baselines/reference/convertToAsyncFunction/convertToAsyncFunction_CatchAndRej.ts

@@ -7,13 +7,15 @@ function /*[#|*/f/*|]*/():Promise<void> {
 
 async function f():Promise<void> {
     try {
+        let result: Response;
         try {
-            const result = await fetch('https://typescriptlang.org');
-            console.log(result);
+            result = await fetch('https://typescriptlang.org');
         }
         catch (rejection) {
             console.log("rejected:", rejection);
+            return;
         }
+        console.log(result);
     }
     catch (err) {
         console.log(err);

tests/baselines/reference/convertToAsyncFunction/convertToAsyncFunction_CatchAndRejRef.ts

@@ -16,13 +16,14 @@ function catch_err(err){
 
 async function f():Promise<void> {
     try {
+        let result;
         try {
-            const result = await fetch('https://typescriptlang.org');
-            return res(result);
+            result = await fetch('https://typescriptlang.org');
         }
         catch (rejection) {
             return rej(rejection);
         }
+        return res(result);
     }
     catch (err) {
         return catch_err(err);

tests/baselines/reference/convertToAsyncFunction/convertToAsyncFunction_Rej.ts

@@ -7,11 +7,13 @@ function /*[#|*/f/*|]*/():Promise<void> {
 // ==ASYNC FUNCTION::Convert to async function==
 
 async function f():Promise<void> {
+    let result: Response;
     try {
-        const result = await fetch('https://typescriptlang.org');
-        console.log(result);
+        result = await fetch('https://typescriptlang.org');
     }
     catch (rejection) {
         console.log("rejected:", rejection);
+        return;
     }
+    console.log(result);
 }

tests/baselines/reference/convertToAsyncFunction/convertToAsyncFunction_RejNoBrackets.ts

@@ -7,11 +7,12 @@ function /*[#|*/f/*|]*/():Promise<void> {
 // ==ASYNC FUNCTION::Convert to async function==
 
 async function f():Promise<void> {
+    let result: Response;
     try {
-        const result = await fetch('https://typescriptlang.org');
-        return console.log(result);
+        result = await fetch('https://typescriptlang.org');
     }
     catch (rejection) {
         return console.log("rejected:", rejection);
     }
+    return console.log(result);
 }

tests/baselines/reference/convertToAsyncFunction/convertToAsyncFunction_RejRef.ts

@@ -13,13 +13,14 @@ function rej(err){
 // ==ASYNC FUNCTION::Convert to async function==
 
 async function f():Promise<void> {
+    let result;
     try {
-        const result = await fetch('https://typescriptlang.org');
-        return res(result);
+        result = await fetch('https://typescriptlang.org');
     }
     catch (err) {
         return rej(err);
     }
+    return res(result);
 }
 function res(result){
     console.log(result);

tests/baselines/reference/convertToAsyncFunction/convertToAsyncFunction_ResRejNoArgsArrow.js

@@ -9,9 +9,9 @@
     async function f() {
         try {
             await Promise.resolve();
-            return 1;
         }
         catch (e) {
             return "a";
         }
+        return 1;
     }

tests/baselines/reference/convertToAsyncFunction/convertToAsyncFunction_ResRejNoArgsArrow.ts

@@ -9,9 +9,9 @@
     async function f() {
         try {
             await Promise.resolve();
-            return 1;
         }
         catch (e) {
             return "a";
         }
+        return 1;
     }

tests/baselines/reference/convertToAsyncFunction/convertToAsyncFunction_twoArgumentThen.js

@@ -0,0 +1,22 @@
+// ==ORIGINAL==
+
+function /*[#|*/fSync/*|]*/() {
+    return Promise.resolve(0).then(x => {
+        console.log(x); // note: added to illustrate refactor
+        throw new Error('Failure!');
+    }, () => null);
+}
+
+// ==ASYNC FUNCTION::Convert to async function==
+
+async function fSync() {
+    let x;
+    try {
+        x = await Promise.resolve(0);
+    }
+    catch (e) {
+        return null;
+    }
+    console.log(x); // note: added to illustrate refactor
+    throw new Error('Failure!');
+}

tests/baselines/reference/convertToAsyncFunction/convertToAsyncFunction_twoArgumentThen.ts

@@ -0,0 +1,22 @@
+// ==ORIGINAL==
+
+function /*[#|*/fSync/*|]*/() {
+    return Promise.resolve(0).then(x => {
+        console.log(x); // note: added to illustrate refactor
+        throw new Error('Failure!');
+    }, () => null);
+}
+
+// ==ASYNC FUNCTION::Convert to async function==
+
+async function fSync() {
+    let x: number;
+    try {
+        x = await Promise.resolve(0);
+    }
+    catch (e) {
+        return null;
+    }
+    console.log(x); // note: added to illustrate refactor
+    throw new Error('Failure!');
+}

tests/baselines/reference/convertToAsyncFunction/convertToAsyncFunction_twoArgumentThen_onRejectedNoReturn.js

@@ -0,0 +1,22 @@
+// ==ORIGINAL==
+
+function /*[#|*/fSync/*|]*/() {
+    return Promise.resolve(0).then(x => {
+        console.log(x); // note: added to illustrate refactor
+        throw new Error('Failure!');
+    }, () => {});
+}
+
+// ==ASYNC FUNCTION::Convert to async function==
+
+async function fSync() {
+    let x;
+    try {
+        x = await Promise.resolve(0);
+    }
+    catch (e) {
+        return;
+    }
+    console.log(x); // note: added to illustrate refactor
+    throw new Error('Failure!');
+}

tests/baselines/reference/convertToAsyncFunction/convertToAsyncFunction_twoArgumentThen_onRejectedNoReturn.ts

@@ -0,0 +1,22 @@
+// ==ORIGINAL==
+
+function /*[#|*/fSync/*|]*/() {
+    return Promise.resolve(0).then(x => {
+        console.log(x); // note: added to illustrate refactor
+        throw new Error('Failure!');
+    }, () => {});
+}
+
+// ==ASYNC FUNCTION::Convert to async function==
+
+async function fSync() {
+    let x: number;
+    try {
+        x = await Promise.resolve(0);
+    }
+    catch (e) {
+        return;
+    }
+    console.log(x); // note: added to illustrate refactor
+    throw new Error('Failure!');
+}