diff --git a/src/Features/LanguageServer/Protocol/Extensions/Extensions.cs b/src/Features/LanguageServer/Protocol/Extensions/Extensions.cs
index d88808f153745..f149a00bbba7e 100644
--- a/src/Features/LanguageServer/Protocol/Extensions/Extensions.cs
+++ b/src/Features/LanguageServer/Protocol/Extensions/Extensions.cs
@@ -88,7 +88,7 @@ private static ImmutableArray<Document> FilterDocumentsByClientName(ImmutableArr
             return documents.FindDocumentInProjectContext(documentIdentifier);
         }
 
-        public static T FindDocumentInProjectContext<T>(this ImmutableArray<T> documents, TextDocumentIdentifier documentIdentifier) where T : TextDocument
+        public static Document FindDocumentInProjectContext(this ImmutableArray<Document> documents, TextDocumentIdentifier documentIdentifier)
         {
             if (documents.Length > 1)
             {
@@ -108,16 +108,10 @@ public static T FindDocumentInProjectContext<T>(this ImmutableArray<T> documents
                     // We were not passed a project context.  This can happen when the LSP powered NavBar is not enabled.
                     // This branch should be removed when we're using the LSP based navbar in all scenarios.
 
-                    // Lookup the active document and determine if any of the documents from the request URI match.
                     var solution = documents.First().Project.Solution;
-                    var service = solution.Workspace.Services.GetRequiredService<IDocumentTrackingService>();
-
-                    var activeDocument = service.GetActiveDocument(solution);
-                    var matchingDocument = documents.FirstOrDefault(d => d.Id == activeDocument?.Id);
-                    if (matchingDocument != null)
-                    {
-                        return matchingDocument;
-                    }
+                    // Lookup which of the linked documents is currently active in the workspace.
+                    var documentIdInCurrentContext = solution.Workspace.GetDocumentIdInCurrentContext(documents.First().Id);
+                    return solution.GetRequiredDocument(documentIdInCurrentContext);
                 }
             }
 
diff --git a/src/Features/LanguageServer/ProtocolUnitTests/Diagnostics/PullDiagnosticTests.cs b/src/Features/LanguageServer/ProtocolUnitTests/Diagnostics/PullDiagnosticTests.cs
index b4ecb81c61e21..44473a03355e2 100644
--- a/src/Features/LanguageServer/ProtocolUnitTests/Diagnostics/PullDiagnosticTests.cs
+++ b/src/Features/LanguageServer/ProtocolUnitTests/Diagnostics/PullDiagnosticTests.cs
@@ -264,21 +264,21 @@ class B {";
             var csproj1Document = testLspServer.GetCurrentSolution().Projects.Where(p => p.Name == "CSProj1").Single().Documents.First();
             var csproj2Document = testLspServer.GetCurrentSolution().Projects.Where(p => p.Name == "CSProj2").Single().Documents.First();
 
-            // Open either of the documents, LSP is just tracking the URI and text.
+            // Open either of the documents via LSP, we're tracking the URI and text.
             await OpenDocumentAsync(testLspServer, csproj1Document);
 
-            // Set the active context to be document from CSProj2
-            var documentTrackingService = (TestDocumentTrackingService)testLspServer.TestWorkspace.Services.GetRequiredService<IDocumentTrackingService>();
-            documentTrackingService.SetActiveDocument(csproj2Document.Id);
+            // This opens all documents in the workspace and ensures buffers are created.
+            testLspServer.TestWorkspace.GetTestDocument(csproj1Document.Id).GetTextBuffer();
 
+            // Set CSProj2 as the active context and get diagnostics.
+            testLspServer.TestWorkspace.SetDocumentContext(csproj2Document.Id);
             var results = await RunGetDocumentPullDiagnosticsAsync(testLspServer, csproj2Document.GetURI());
             Assert.Equal("CS1513", results.Single().Diagnostics.Single().Code);
             var vsDiagnostic = (LSP.VSDiagnostic)results.Single().Diagnostics.Single();
             Assert.Equal("CSProj2", vsDiagnostic.Projects.Single().ProjectName);
 
-            // Set the active context to be document from CSProj1
-            documentTrackingService.SetActiveDocument(csproj1Document.Id);
-
+            // Set CSProj1 as the active context and get diagnostics.
+            testLspServer.TestWorkspace.SetDocumentContext(csproj1Document.Id);
             results = await RunGetDocumentPullDiagnosticsAsync(testLspServer, csproj1Document.GetURI());
             Assert.Equal(2, results.Single().Diagnostics!.Length);
             Assert.All(results.Single().Diagnostics, d => Assert.Equal("CS1513", d.Code));