Index: clangd/ClangdServer.h
===================================================================
--- clangd/ClangdServer.h
+++ clangd/ClangdServer.h
@@ -226,20 +226,12 @@
   formatCode(llvm::StringRef Code, PathRef File,
              ArrayRef<tooling::Range> Ranges);
 
-  typedef uint64_t DocVersion;
-
-  void consumeDiagnostics(PathRef File, DocVersion Version,
-                          std::vector<Diag> Diags);
-
   tooling::CompileCommand getCompileCommand(PathRef File);
 
   const GlobalCompilationDatabase &CDB;
   DiagnosticsConsumer &DiagConsumer;
   const FileSystemProvider &FSProvider;
 
-  /// Used to synchronize diagnostic responses for added and removed files.
-  llvm::StringMap<DocVersion> InternalVersion;
-
   Path ResourceDir;
   // The index used to look up symbols. This could be:
   //   - null (all index functionality is optional)
@@ -259,12 +251,6 @@
 
   llvm::Optional<std::string> WorkspaceRoot;
   std::shared_ptr<PCHContainerOperations> PCHs;
-  /// Used to serialize diagnostic callbacks.
-  /// FIXME(ibiryukov): get rid of an extra map and put all version counters
-  /// into CppFile.
-  std::mutex DiagnosticsMutex;
-  /// Maps from a filename to the latest version of reported diagnostics.
-  llvm::StringMap<DocVersion> ReportedDiagnosticVersions;
   // WorkScheduler has to be the last member, because its destructor has to be
   // called before all other members to stop the worker thread that references
   // ClangdServer.
Index: clangd/ClangdServer.cpp
===================================================================
--- clangd/ClangdServer.cpp
+++ clangd/ClangdServer.cpp
@@ -134,19 +134,17 @@
 
 void ClangdServer::addDocument(PathRef File, StringRef Contents,
                                WantDiagnostics WantDiags) {
-  DocVersion Version = ++InternalVersion[File];
   ParseInputs Inputs = {getCompileCommand(File), FSProvider.getFileSystem(),
                         Contents.str()};
-
   Path FileStr = File.str();
   WorkScheduler.update(File, std::move(Inputs), WantDiags,
-                       [this, FileStr, Version](std::vector<Diag> Diags) {
-                         consumeDiagnostics(FileStr, Version, std::move(Diags));
+                       [this, FileStr](std::vector<Diag> Diags) {
+                         DiagConsumer.onDiagnosticsReady(FileStr,
+                                                         std::move(Diags));
                        });
 }
 
 void ClangdServer::removeDocument(PathRef File) {
-  ++InternalVersion[File];
   WorkScheduler.remove(File);
 }
 
@@ -445,24 +443,6 @@
   WorkScheduler.runWithAST("Hover", File, Bind(Action, std::move(CB)));
 }
 
-void ClangdServer::consumeDiagnostics(PathRef File, DocVersion Version,
-                                      std::vector<Diag> Diags) {
-  // We need to serialize access to resulting diagnostics to avoid calling
-  // `onDiagnosticsReady` in the wrong order.
-  std::lock_guard<std::mutex> DiagsLock(DiagnosticsMutex);
-  DocVersion &LastReportedDiagsVersion = ReportedDiagnosticVersions[File];
-
-  // FIXME(ibiryukov): get rid of '<' comparison here. In the current
-  // implementation diagnostics will not be reported after version counters'
-  // overflow. This should not happen in practice, since DocVersion is a
-  // 64-bit unsigned integer.
-  if (Version < LastReportedDiagsVersion)
-    return;
-  LastReportedDiagsVersion = Version;
-
-  DiagConsumer.onDiagnosticsReady(File, std::move(Diags));
-}
-
 tooling::CompileCommand ClangdServer::getCompileCommand(PathRef File) {
   trace::Span Span("GetCompileCommand");
   Optional<tooling::CompileCommand> C = CDB.getCompileCommand(File);
Index: clangd/TUScheduler.cpp
===================================================================
--- clangd/TUScheduler.cpp
+++ clangd/TUScheduler.cpp
@@ -251,6 +251,11 @@
   bool Done;                    /* GUARDED_BY(Mutex) */
   std::deque<Request> Requests; /* GUARDED_BY(Mutex) */
   mutable std::condition_variable RequestsCV;
+  /// Guards a critical section for running the diagnostics callbacks. 
+  std::mutex DiagsMu;
+  /// When set to false, no diagnostics will be reported. Used to ensure we
+  /// don't report any diagnostics after a shutdown of the worker is requested.
+  bool ReportDiagnostics = true; /* GUARDED_BY(DiagMu) */
 };
 
 /// A smart-pointer-like class that points to an active ASTWorker.
@@ -401,10 +406,17 @@
       }
     }
 
-    // We only need to build the AST if diagnostics were requested.
+    // We only need to build the AST if diagnostics were requested and the
+    // callers are still interested in this ASTWorker (it's not shutting down).
     if (WantDiags == WantDiagnostics::No)
       return;
 
+    {
+      std::lock_guard<std::mutex> Lock(DiagsMu);
+      if (!ReportDiagnostics)
+        return;
+    }
+
     // Get the AST for diagnostics.
     Optional<std::unique_ptr<ParsedAST>> AST = IdleASTs.take(this);
     if (!AST) {
@@ -417,7 +429,11 @@
     // spam us with updates.
     // Note *AST can still be null if buildAST fails.
     if (*AST) {
-      OnUpdated((*AST)->getDiagnostics());
+      {
+        std::lock_guard<std::mutex> Lock(DiagsMu);
+        if (ReportDiagnostics)
+          OnUpdated((*AST)->getDiagnostics());
+      }
       trace::Span Span("Running main AST callback");
       Callbacks.onMainAST(FileName, **AST);
       DiagsWereReported = true;
@@ -514,6 +530,10 @@
     assert(!Done && "stop() called twice");
     Done = true;
   }
+  {
+    std::lock_guard<std::mutex> Lock(DiagsMu);
+    ReportDiagnostics = false;
+  }
   RequestsCV.notify_all();
 }
 
Index: unittests/clangd/ClangdTests.cpp
===================================================================
--- unittests/clangd/ClangdTests.cpp
+++ unittests/clangd/ClangdTests.cpp
@@ -748,7 +748,8 @@
         BlockingRequests[RequestIndex]();
       }
     }
-  } // Wait for ClangdServer to shutdown before proceeding.
+    ASSERT_TRUE(Server.blockUntilIdleForTest());
+  }
 
   // Check some invariants about the state of the program.
   std::vector<FileStat> Stats = DiagConsumer.takeFileStats();
Index: unittests/clangd/TUSchedulerTests.cpp
===================================================================
--- unittests/clangd/TUSchedulerTests.cpp
+++ unittests/clangd/TUSchedulerTests.cpp
@@ -124,6 +124,8 @@
     S.update(Path, getInputs(Path, "auto (produces)"), WantDiagnostics::Auto,
              [&](std::vector<Diag> Diags) { ++CallbackCount; });
     Ready.notify();
+
+    ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
   }
   EXPECT_EQ(2, CallbackCount);
 }
@@ -147,6 +149,8 @@
     std::this_thread::sleep_for(std::chrono::seconds(2));
     S.update(Path, getInputs(Path, "auto (shut down)"), WantDiagnostics::Auto,
              [&](std::vector<Diag> Diags) { ++CallbackCount; });
+
+    ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
   }
   EXPECT_EQ(2, CallbackCount);
 }
@@ -304,6 +308,7 @@
         }
       }
     }
+    ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
   } // TUScheduler destructor waits for all operations to finish.
 
   std::lock_guard<std::mutex> Lock(Mut);