Implement codefix for adding missing await in return statements enclosed by try statements

src/compiler/diagnosticMessages.json

@@ -7255,6 +7255,10 @@
         "category": "Suggestion",
         "code": 80010
     },
+    "This may need `await` keyword. Otherwise the enclosing try statement won't handle this.": {
+        "category": "Suggestion",
+        "code": 80011
+    },
 
     "Add missing 'super()' call": {
         "category": "Message",
@@ -8245,6 +8249,14 @@
         "category": "Message",
         "code": 95197
     },
+    "Add missing await.": {
+        "category": "Message",
+        "code": 95198
+    },
+    "Add all missing awaits.": {
+        "category": "Message",
+        "code": 95199
+    },
 
     "No value exists in scope for the shorthand property '{0}'. Either declare one or provide an initializer.": {
         "category": "Error",

src/services/_namespaces/ts.codefix.ts

@@ -5,6 +5,7 @@ export * from "../codefixes/addConvertToUnknownForNonOverlappingTypes.js";
 export * from "../codefixes/addEmptyExportDeclaration.js";
 export * from "../codefixes/addMissingAsync.js";
 export * from "../codefixes/addMissingAwait.js";
+export * from "../codefixes/addMissingAwaitInReturn.js";
 export * from "../codefixes/addMissingConst.js";
 export * from "../codefixes/addMissingDeclareProperty.js";
 export * from "../codefixes/addMissingInvocationForDecorator.js";

src/services/codefixes/addMissingAwaitInReturn.ts

@@ -0,0 +1,41 @@
+import {
+    codeFixAll,
+    createCodeFixAction,
+    registerCodeFix,
+} from "../_namespaces/ts.codefix.js";
+import {
+    CodeFixContext,
+    Debug,
+    Diagnostics,
+    factory,
+    getSynthesizedDeepClone,
+    getTokenAtPosition,
+    isReturnStatement,
+    SourceFile,
+    textChanges,
+} from "../_namespaces/ts.js";
+
+const fixId = "addMissingAwaitInReturn";
+const errorCodes = [Diagnostics.This_may_need_await_keyword_Otherwise_the_enclosing_try_statement_won_t_handle_this.code];
+
+registerCodeFix({
+    errorCodes,
+    getCodeActions(context: CodeFixContext) {
+        const changes = textChanges.ChangeTracker.with(context, t => makeChange(t, context.sourceFile, context.span.start));
+        return [createCodeFixAction(fixId, changes, Diagnostics.Add_missing_await, fixId, Diagnostics.Add_all_missing_awaits)];
+    },
+    fixIds: [fixId],
+    getAllCodeActions: context => codeFixAll(context, errorCodes, (changes, diag) => makeChange(changes, diag.file, diag.start)),
+});
+
+function makeChange(changeTracker: textChanges.ChangeTracker, sourceFile: SourceFile, pos: number) {
+    const token = getTokenAtPosition(sourceFile, pos);
+    Debug.assertNode(token.parent, isReturnStatement);
+    Debug.assertIsDefined(token.parent.expression);
+    const expression = token.parent.expression;
+    changeTracker.replaceNode(
+        sourceFile,
+        expression,
+        factory.createAwaitExpression(getSynthesizedDeepClone(expression, /*includeTrivia*/ true)),
+    );
+}

src/services/suggestionDiagnostics.ts

@@ -15,6 +15,7 @@ import {
     ExpressionStatement,
     Extension,
     fileExtensionIsOneOf,
+    findAncestor,
     forEachReturnStatement,
     FunctionDeclaration,
     FunctionExpression,
@@ -28,19 +29,22 @@ import {
     Identifier,
     importFromModuleSpecifier,
     isAsyncFunction,
+    isAwaitExpression,
     isBinaryExpression,
     isBlock,
     isCallExpression,
     isExportAssignment,
     isFunctionDeclaration,
     isFunctionExpression,
     isFunctionLike,
+    isFunctionLikeDeclaration,
     isIdentifier,
     isPropertyAccessExpression,
     isRequireCall,
     isReturnStatement,
     isSourceFileJS,
     isStringLiteral,
+    isTryStatement,
     isVariableDeclaration,
     isVariableStatement,
     MethodDeclaration,
@@ -52,6 +56,7 @@ import {
     PropertyAccessExpression,
     ReturnStatement,
     skipAlias,
+    skipParentheses,
     some,
     SourceFile,
     SyntaxKind,
@@ -132,6 +137,9 @@ export function computeSuggestionDiagnostics(sourceFile: SourceFile, program: Pr
         if (canBeConvertedToAsync(node)) {
             addConvertToAsyncFunctionDiagnostics(node, checker, diags);
         }
+        if (isFunctionLikeDeclaration(node) && isAsyncFunction(node)) {
+            addMissingAwaitInReturnDiagnostics(node, checker, diags);
+        }
         node.forEachChild(check);
     }
 }
@@ -190,6 +198,31 @@ function isConvertibleFunction(node: FunctionLikeDeclaration, checker: TypeCheck
         returnsPromise(node, checker);
 }
 
+function addMissingAwaitInReturnDiagnostics(node: FunctionLikeDeclaration, checker: TypeChecker, diags: DiagnosticWithLocation[]) {
+    if (!node.body || !isBlock(node.body)) {
+        return;
+    }
+
+    forEachReturnStatement(node.body, statement => {
+        if (!statement.expression) {
+            return;
+        }
+        const expression = skipParentheses(statement.expression);
+        if (isAwaitExpression(expression)) {
+            return;
+        }
+        const type = checker.getTypeAtLocation(expression);
+        if (type.isUnion() ? type.types.every(t => !checker.getPromisedTypeOfPromise(t)) : !checker.getPromisedTypeOfPromise(type)) {
+            return;
+        }
+        const ancestor = findAncestor(statement, n => n === node || isTryStatement(n));
+        if (!ancestor || !isTryStatement(ancestor)) {
+            return;
+        }
+        diags.push(createDiagnosticForNode(statement, Diagnostics.This_may_need_await_keyword_Otherwise_the_enclosing_try_statement_won_t_handle_this));
+    });
+}
+
 /** @internal */
 export function returnsPromise(node: FunctionLikeDeclaration, checker: TypeChecker): boolean {
     const signature = checker.getSignatureFromDeclaration(node);

tests/cases/fourslash/codeFixAddMissingAwaitInReturn1.ts

@@ -0,0 +1,19 @@
+/// <reference path="fourslash.ts" />
+
+// @strict: true
+// @target: esnext
+// @lib: esnext
+
+//// async function inner() {
+////   if (Math.random() > 0.5) {
+////     throw new Error("Ooops");
+////   }
+////   return 42;
+//// }
+////
+//// async function outer() {
+////   return inner();
+//// }
+
+verify.getSuggestionDiagnostics([]);
+verify.not.codeFixAvailable(ts.Diagnostics.Add_missing_await.message);

tests/cases/fourslash/codeFixAddMissingAwaitInReturn10.ts

@@ -0,0 +1,14 @@
+/// <reference path="fourslash.ts" />
+
+// @strict: true
+// @target: esnext
+// @lib: esnext
+
+//// async function inner() {
+////   return 42;
+//// }
+////
+//// const outer = async () => inner();
+
+verify.getSuggestionDiagnostics([]);
+verify.not.codeFixAvailable(ts.Diagnostics.Add_missing_await.message);

tests/cases/fourslash/codeFixAddMissingAwaitInReturn11.ts

@@ -0,0 +1,61 @@
+/// <reference path="fourslash.ts" />
+
+// @strict: true
+// @target: esnext
+// @lib: esnext
+
+//// async function inner() {
+////   return 42;
+//// }
+////
+//// async function outer(v: number) {
+////   try {
+////     if (v > 30) {
+////       [|return|] inner();
+////     }
+////     if (v > 20) {
+////       return await inner();
+////     }
+////     if (v > 10) {
+////       [|return|] inner();
+////     }
+////     [|return|] inner();
+////   } catch (e) {}
+//// }
+
+verify.getSuggestionDiagnostics([{
+    message: ts.Diagnostics.This_may_need_await_keyword_Otherwise_the_enclosing_try_statement_won_t_handle_this.message,
+    code: 80011,
+    range: test.ranges()[0],
+}, {
+    message: ts.Diagnostics.This_may_need_await_keyword_Otherwise_the_enclosing_try_statement_won_t_handle_this.message,
+    code: 80011,
+    range: test.ranges()[1],
+}, {
+    message: ts.Diagnostics.This_may_need_await_keyword_Otherwise_the_enclosing_try_statement_won_t_handle_this.message,
+    code: 80011,
+    range: test.ranges()[2],
+}]);
+verify.codeFixAll({
+    fixId: "addMissingAwaitInReturn",
+    fixAllDescription: ts.Diagnostics.Add_all_missing_awaits.message,
+    newFileContent:
+`async function inner() {
+  return 42;
+}
+
+async function outer(v: number) {
+  try {
+    if (v > 30) {
+      return await inner();
+    }
+    if (v > 20) {
+      return await inner();
+    }
+    if (v > 10) {
+      return await inner();
+    }
+    return await inner();
+  } catch (e) {}
+}`,
+});

tests/cases/fourslash/codeFixAddMissingAwaitInReturn12.ts

@@ -0,0 +1,41 @@
+/// <reference path="fourslash.ts" />
+
+// @strict: true
+// @target: esnext
+// @lib: esnext
+
+//// async function inner() {
+////   if (Math.random() > 0.5) {
+////     throw new Error("Ooops");
+////   }
+////   return 42;
+//// }
+////
+//// async function outer() {
+////   try {
+////     [|return|] Math.random() > 0.5 ? inner() : 5;
+////   } catch (e) {}
+//// }
+
+verify.getSuggestionDiagnostics([{
+    message: ts.Diagnostics.This_may_need_await_keyword_Otherwise_the_enclosing_try_statement_won_t_handle_this.message,
+    code: 80011,
+    range: test.ranges()[0],
+}]);
+verify.codeFix({
+    description: ts.Diagnostics.Add_missing_await.message,
+    index: 0,
+    newFileContent:
+`async function inner() {
+  if (Math.random() > 0.5) {
+    throw new Error("Ooops");
+  }
+  return 42;
+}
+
+async function outer() {
+  try {
+    return await (Math.random() > 0.5 ? inner() : 5);
+  } catch (e) {}
+}`,
+});

tests/cases/fourslash/codeFixAddMissingAwaitInReturn2.ts

@@ -0,0 +1,41 @@
+/// <reference path="fourslash.ts" />
+
+// @strict: true
+// @target: esnext
+// @lib: esnext
+
+//// async function inner() {
+////   if (Math.random() > 0.5) {
+////     throw new Error("Ooops");
+////   }
+////   return 42;
+//// }
+////
+//// async function outer() {
+////   try {
+////     [|return|] inner();
+////   } catch (e) {}
+//// }
+
+verify.getSuggestionDiagnostics([{
+    message: ts.Diagnostics.This_may_need_await_keyword_Otherwise_the_enclosing_try_statement_won_t_handle_this.message,
+    code: 80011,
+    range: test.ranges()[0],
+}]);
+verify.codeFix({
+    description: ts.Diagnostics.Add_missing_await.message,
+    index: 0,
+    newFileContent:
+`async function inner() {
+  if (Math.random() > 0.5) {
+    throw new Error("Ooops");
+  }
+  return 42;
+}
+
+async function outer() {
+  try {
+    return await inner();
+  } catch (e) {}
+}`,
+});

tests/cases/fourslash/codeFixAddMissingAwaitInReturn3.ts

@@ -0,0 +1,41 @@
+/// <reference path="fourslash.ts" />
+
+// @strict: true
+// @target: esnext
+// @lib: esnext
+
+//// async function inner() {
+////   if (Math.random() > 0.5) {
+////     throw new Error("Ooops");
+////   }
+////   return 42;
+//// }
+////
+//// async function outer() {
+////   try {
+////     [|return|] inner();
+////   } catch {}
+//// }
+
+verify.getSuggestionDiagnostics([{
+    message: ts.Diagnostics.This_may_need_await_keyword_Otherwise_the_enclosing_try_statement_won_t_handle_this.message,
+    code: 80011,
+    range: test.ranges()[0],
+}]);
+verify.codeFix({
+    description: ts.Diagnostics.Add_missing_await.message,
+    index: 0,
+    newFileContent:
+`async function inner() {
+  if (Math.random() > 0.5) {
+    throw new Error("Ooops");
+  }
+  return 42;
+}
+
+async function outer() {
+  try {
+    return await inner();
+  } catch {}
+}`,
+});