diff --git a/test/SourceKit/Diagnostics/cancel_diags.swift b/test/SourceKit/Diagnostics/cancel_diags.swift
new file mode 100644
index 0000000000000..7b4281d240693
--- /dev/null
+++ b/test/SourceKit/Diagnostics/cancel_diags.swift
@@ -0,0 +1,7 @@
+// RUN: not %sourcekitd-test -req=diags %s -print-raw-response -id=diag -async -- %s == -cancel=diag 2>&1 | %FileCheck %s
+
+func foo(x: Invalid1, y: Invalid2) {
+    x / y / x / y / x / y / x / y
+}
+
+// CHECK: error response (Request Cancelled)
diff --git a/tools/SourceKit/lib/SwiftLang/SwiftASTManager.cpp b/tools/SourceKit/lib/SwiftLang/SwiftASTManager.cpp
index 03f4454413ad1..4b2e4c7e84e02 100644
--- a/tools/SourceKit/lib/SwiftLang/SwiftASTManager.cpp
+++ b/tools/SourceKit/lib/SwiftLang/SwiftASTManager.cpp
@@ -1069,6 +1069,7 @@ ASTUnitRef ASTBuildOperation::buildASTUnit(std::string &Error) {
     }
     return nullptr;
   }
+  CompIns.getASTContext().CancellationFlag = CancellationFlag;
   registerIDERequestFunctions(CompIns.getASTContext().evaluator);
   if (TracedOp.enabled()) {
     TracedOp.start(TraceInfo);
@@ -1112,9 +1113,19 @@ ASTUnitRef ASTBuildOperation::buildASTUnit(std::string &Error) {
     // don't block any other AST processing for the same SwiftInvocation.
 
     if (auto SF = CompIns.getPrimarySourceFile()) {
+      if (CancellationFlag->load(std::memory_order_relaxed)) {
+        return nullptr;
+      }
+      // Disable cancellation while performing SILGen. If the cancellation flag
+      // is set, type checking performed during SILGen checks the cancellation
+      // flag and might thus fail, which SILGen cannot handle.
+      llvm::SaveAndRestore<std::shared_ptr<std::atomic<bool>>> DisableCancellationDuringSILGen(CompIns.getASTContext().CancellationFlag, nullptr);
       SILOptions SILOpts = Invocation.getSILOptions();
       auto &TC = CompIns.getSILTypes();
       std::unique_ptr<SILModule> SILMod = performASTLowering(*SF, TC, SILOpts);
+      if (CancellationFlag->load(std::memory_order_relaxed)) {
+        return nullptr;
+      }
       runSILDiagnosticPasses(*SILMod);
     }
   }
diff --git a/tools/SourceKit/lib/SwiftLang/SwiftSourceDocInfo.cpp b/tools/SourceKit/lib/SwiftLang/SwiftSourceDocInfo.cpp
index c3bbb16d51f3b..6adf5eedd1444 100644
--- a/tools/SourceKit/lib/SwiftLang/SwiftSourceDocInfo.cpp
+++ b/tools/SourceKit/lib/SwiftLang/SwiftSourceDocInfo.cpp
@@ -1621,6 +1621,10 @@ static void computeDiagnostics(
       auto Diagnostics = DiagConsumer.getDiagnosticsForBuffer(BufferID);
       Receiver(RequestResult<DiagnosticsResult>::fromResult(Diagnostics));
     }
+
+    void cancelled() override {
+      Receiver(RequestResult<DiagnosticsResult>::cancelled());
+    }
   };
 
   auto Consumer = std::make_shared<DiagnosticsConsumer>(std::move(Receiver));
diff --git a/unittests/SourceKit/SwiftLang/CursorInfoTest.cpp b/unittests/SourceKit/SwiftLang/CursorInfoTest.cpp
index 6c113f363d235..bdc872a538add 100644
--- a/unittests/SourceKit/SwiftLang/CursorInfoTest.cpp
+++ b/unittests/SourceKit/SwiftLang/CursorInfoTest.cpp
@@ -424,8 +424,7 @@ TEST_F(CursorInfoTest, CursorInfoMustWaitDueTokenRace) {
   EXPECT_EQ(strlen("fog"), Info.Length);
 }
 
-// Disabled until we re-enable cancellation (rdar://91251055)
-TEST_F(CursorInfoTest, DISABLED_CursorInfoCancelsPreviousRequest) {
+TEST_F(CursorInfoTest, CursorInfoCancelsPreviousRequest) {
   // TODO: This test case relies on the following snippet being slow to type 
   // check so that the first cursor info request takes longer to execute than it 
   // takes time to schedule the second request. If that is fixed, we need to 
@@ -475,8 +474,7 @@ TEST_F(CursorInfoTest, DISABLED_CursorInfoCancelsPreviousRequest) {
     llvm::report_fatal_error("Did not receive a resonse for the first request");
 }
 
-// Disabled until we re-enable cancellation (rdar://91251055)
-TEST_F(CursorInfoTest, DISABLED_CursorInfoCancellation) {
+TEST_F(CursorInfoTest, CursorInfoCancellation) {
   // TODO: This test case relies on the following snippet being slow to type
   // check so that the first cursor info request takes longer to execute than it
   // takes time to schedule the second request. If that is fixed, we need to