Index: unittests/clangd/ClangdTests.cpp
===================================================================
--- unittests/clangd/ClangdTests.cpp
+++ unittests/clangd/ClangdTests.cpp
@@ -15,6 +15,7 @@
 #include "TestFS.h"
 #include "URI.h"
 #include "clang/Config/config.h"
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/Support/Errc.h"
@@ -1033,6 +1034,100 @@
 }
 #endif
 
+TEST(ClangdTests, AddDocumentTracksDiagnostics) {
+  // We have clients relying on the fact that the Context passed to addDocument
+  // dies around the time we're done computing diagnostics for that particular
+  // request. This test checks this behavior does not break.
+
+  // Used to account for all diagnostic responses.
+  class TestDiagConsumer : public DiagnosticsConsumer {
+  public:
+    ~TestDiagConsumer() {
+      if (PendingDiags)
+        ADD_FAILURE() << "Diagnostics where not consumed";
+    }
+    void onDiagnosticsReady(PathRef, std::vector<Diag>) override {
+      ASSERT_FALSE(PendingDiags) << "Diagnostics were not consumed";
+      PendingDiags = true;
+    }
+
+    LLVM_NODISCARD bool consumeDiags() { return PendingDiags.exchange(false); }
+
+  private:
+    std::atomic<bool> PendingDiags{false};
+  };
+  MockCompilationDatabase CDB;
+  MockFSProvider FS;
+  TestDiagConsumer DiagConsumer;
+  ClangdServer S(CDB, FS, DiagConsumer, ClangdServer::optsForTest());
+
+  auto FooCpp = testPath("foo.cpp");
+  StringRef SourceContents = "int foo() { return 10; }";
+  StringRef OtherSourceContents = "int a;";
+  std::atomic<bool> Called;
+  // New document is added, diagnostics must be produced.
+  Called = false;
+  {
+    WithContextValue SetCtx(llvm::make_scope_exit([&]() {
+      ASSERT_TRUE(DiagConsumer.consumeDiags());
+      Called = true;
+    }));
+    S.addDocument(FooCpp, SourceContents);
+  }
+  ASSERT_TRUE(S.blockUntilIdleForTest());
+  // Add a document with the same contents. The diagnostics will *not* be
+  // produced, but the callback should still happen when processing finishes.
+  Called = false;
+  {
+    WithContextValue SetCtx(llvm::make_scope_exit([&]() {
+      ASSERT_FALSE(DiagConsumer.consumeDiags());
+      Called = true;
+    }));
+    S.addDocument(FooCpp, SourceContents);
+  }
+  ASSERT_TRUE(S.blockUntilIdleForTest());
+  ASSERT_TRUE(Called);
+  // Check the case when contents change, but no diagnostics were requested.
+  Called = false;
+  {
+    WithContextValue SetCtx(llvm::make_scope_exit([&]() {
+      ASSERT_FALSE(DiagConsumer.consumeDiags());
+      Called = true;
+    }));
+    S.addDocument(FooCpp, OtherSourceContents, WantDiagnostics::No);
+  }
+  ASSERT_TRUE(S.blockUntilIdleForTest());
+  ASSERT_TRUE(Called);
+
+  // When ClangdServer gets many Auto and No requests, some of them might be
+  // dropped. But all the contexts should have the lifetime we expect anyway.
+  constexpr int RequestsCount = 10000;
+  std::atomic<int> Callbacks(0);
+  std::atomic<int> ConsumedDiags(0);
+  StringRef Sources[] = {SourceContents, OtherSourceContents};
+  for (int I = 0; I < RequestsCount; ++I) {
+    auto WD = I % 5 == 0 ? WantDiagnostics::No : WantDiagnostics::Auto;
+    WithContextValue SetCtx(llvm::make_scope_exit([&, WD, I]() {
+      EXPECT_EQ(Callbacks, I);
+      ++Callbacks;
+
+      bool HadDiags = DiagConsumer.consumeDiags();
+      if (HadDiags)
+        ++ConsumedDiags;
+      if (WD == WantDiagnostics::No)
+        EXPECT_FALSE(HadDiags);
+      // If WD == WantDiagnostics::Auto, HadDiags might have any of the two
+      // values.
+    }));
+    // Note that we always change the source to make sure we don't hit the
+    // "inputs did not change" optimization.
+    S.addDocument(FooCpp, Sources[I % 2], WD);
+  }
+  ASSERT_TRUE(S.blockUntilIdleForTest());
+  ASSERT_EQ(Callbacks, RequestsCount);
+  log("Test consumed {0} diagnostic responses", ConsumedDiags);
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang