diff --git a/clang/include/clang/Frontend/ASTUnit.h b/clang/include/clang/Frontend/ASTUnit.h
index ad54016da11d2..27bba8e64859d 100644
--- a/clang/include/clang/Frontend/ASTUnit.h
+++ b/clang/include/clang/Frontend/ASTUnit.h
@@ -629,6 +629,17 @@ class ASTUnit {
     return StoredDiagnostics.end();
   }
 
+  using diags_range = llvm::iterator_range<stored_diag_iterator>;
+  using const_diags_range = llvm::iterator_range<stored_diag_const_iterator>;
+
+  diags_range storedDiagnostics() {
+    return {stored_diag_begin(), stored_diag_end()};
+  }
+
+  const_diags_range storedDiagnostics() const {
+    return {stored_diag_begin(), stored_diag_end()};
+  }
+
   unsigned stored_diag_size() const { return StoredDiagnostics.size(); }
 
   stored_diag_iterator stored_diag_afterDriver_begin() {
diff --git a/clang/include/clang/Tooling/Tooling.h b/clang/include/clang/Tooling/Tooling.h
index 200fb30839a95..9909394495496 100644
--- a/clang/include/clang/Tooling/Tooling.h
+++ b/clang/include/clang/Tooling/Tooling.h
@@ -32,6 +32,7 @@
 #include "clang/AST/ASTConsumer.h"
 #include "clang/Basic/FileManager.h"
 #include "clang/Basic/LLVM.h"
+#include "clang/Frontend/ASTUnit.h"
 #include "clang/Frontend/FrontendAction.h"
 #include "clang/Frontend/PCHContainerOperations.h"
 #include "clang/Tooling/ArgumentsAdjusters.h"
@@ -239,7 +240,8 @@ std::unique_ptr<ASTUnit> buildASTFromCodeWithArgs(
     const FileContentMappings &VirtualMappedFiles = FileContentMappings(),
     DiagnosticConsumer *DiagConsumer = nullptr,
     IntrusiveRefCntPtr<llvm::vfs::FileSystem> BaseFS =
-        llvm::vfs::getRealFileSystem());
+        llvm::vfs::getRealFileSystem(),
+    CaptureDiagsKind CaptureKind = CaptureDiagsKind::None);
 
 /// Utility to run a FrontendAction in a single clang invocation.
 class ToolInvocation {
diff --git a/clang/lib/Tooling/Tooling.cpp b/clang/lib/Tooling/Tooling.cpp
index 45dfdf4360253..1120e33b1f0c3 100644
--- a/clang/lib/Tooling/Tooling.cpp
+++ b/clang/lib/Tooling/Tooling.cpp
@@ -644,9 +644,13 @@ namespace {
 
 class ASTBuilderAction : public ToolAction {
   std::vector<std::unique_ptr<ASTUnit>> &ASTs;
+  CaptureDiagsKind CaptureKind;
 
 public:
-  ASTBuilderAction(std::vector<std::unique_ptr<ASTUnit>> &ASTs) : ASTs(ASTs) {}
+  ASTBuilderAction(
+      std::vector<std::unique_ptr<ASTUnit>> &ASTs,
+      CaptureDiagsKind CaptureDiagnosticsKind = CaptureDiagsKind::None)
+      : ASTs(ASTs), CaptureKind(CaptureDiagnosticsKind) {}
 
   bool runInvocation(std::shared_ptr<CompilerInvocation> Invocation,
                      FileManager *Files,
@@ -658,7 +662,7 @@ class ASTBuilderAction : public ToolAction {
                                             Invocation->getDiagnosticOpts(),
                                             DiagConsumer,
                                             /*ShouldOwnClient=*/false),
-        Files);
+        Files, false, CaptureKind);
     if (!AST)
       return false;
 
@@ -693,9 +697,12 @@ std::unique_ptr<ASTUnit> buildASTFromCodeWithArgs(
     StringRef ToolName, std::shared_ptr<PCHContainerOperations> PCHContainerOps,
     ArgumentsAdjuster Adjuster, const FileContentMappings &VirtualMappedFiles,
     DiagnosticConsumer *DiagConsumer,
-    IntrusiveRefCntPtr<llvm::vfs::FileSystem> BaseFS) {
+    IntrusiveRefCntPtr<llvm::vfs::FileSystem> BaseFS,
+    CaptureDiagsKind CaptureKind) {
   std::vector<std::unique_ptr<ASTUnit>> ASTs;
-  ASTBuilderAction Action(ASTs);
+
+  ASTBuilderAction Action(ASTs, CaptureKind);
+
   auto OverlayFileSystem =
       llvm::makeIntrusiveRefCnt<llvm::vfs::OverlayFileSystem>(
           std::move(BaseFS));
diff --git a/clang/unittests/Sema/HeuristicResolverTest.cpp b/clang/unittests/Sema/HeuristicResolverTest.cpp
index 21aca7a3489b8..883a4e20e40a7 100644
--- a/clang/unittests/Sema/HeuristicResolverTest.cpp
+++ b/clang/unittests/Sema/HeuristicResolverTest.cpp
@@ -8,6 +8,7 @@
 #include "clang/Sema/HeuristicResolver.h"
 #include "clang/ASTMatchers/ASTMatchFinder.h"
 #include "clang/ASTMatchers/ASTMatchers.h"
+#include "clang/Basic/Diagnostic.h"
 #include "clang/Tooling/Tooling.h"
 #include "gmock/gmock-matchers.h"
 #include "gtest/gtest.h"
@@ -41,7 +42,19 @@ template <typename InputNode, typename ParamT, typename InputMatcher,
           typename... OutputMatchers>
 void expectResolution(llvm::StringRef Code, ResolveFnT<ParamT> ResolveFn,
                       const InputMatcher &IM, const OutputMatchers &...OMS) {
-  auto TU = tooling::buildASTFromCodeWithArgs(Code, {"-std=c++23"});
+  auto TU = tooling::buildASTFromCodeWithArgs(
+      Code, {"-std=c++23"}, "input.cc", "clang-tool",
+      std::make_shared<PCHContainerOperations>(),
+      tooling::getClangStripDependencyFileAdjuster(),
+      tooling::FileContentMappings(), nullptr, llvm::vfs::getRealFileSystem(),
+      CaptureDiagsKind::All);
+
+  for (const auto &D : TU->storedDiagnostics()) {
+    EXPECT_TRUE(D.getLevel() < DiagnosticsEngine::Error)
+        << "Unexpected error diagnostic while building AST for test code: "
+        << D.getMessage();
+  }
+
   auto &Ctx = TU->getASTContext();
   auto InputMatches = match(IM, Ctx);
   ASSERT_EQ(1u, InputMatches.size());