We need to remove this list filtering:
//RoslynDocumentHelper.GetDocumentByDocumentIdAsync
document = (await project.GetSourceGeneratedDocumentsAsync(CancellationToken.None))
.First(d => d.Id.Equals(documentId));
and just use project.GetSourceGeneratedDocumentAsync.
Tracking issue: dotnet/roslyn#71581