diff --git a/lib/Frontend/ASTMerge.cpp b/lib/Frontend/ASTMerge.cpp
index d51647b..f2de09a 100644
--- a/lib/Frontend/ASTMerge.cpp
+++ b/lib/Frontend/ASTMerge.cpp
@@ -32,6 +32,8 @@
 
 void ASTMergeAction::ExecuteAction() {
   CompilerInstance &CI = getCompilerInstance();
+  CI.getDiagnostics().getClient()->BeginSourceFile(
+                                         CI.getASTContext().getLangOptions());
   CI.getDiagnostics().SetArgToStringFn(&FormatASTNodeDiagnosticArgument,
                                        &CI.getASTContext());
   for (unsigned I = 0, N = ASTFiles.size(); I != N; ++I) {
@@ -68,8 +70,8 @@
     delete Unit;
   }
 
-
-  return AdaptedAction->ExecuteAction();
+  AdaptedAction->ExecuteAction();
+  CI.getDiagnostics().getClient()->EndSourceFile();
 }
 
 void ASTMergeAction::EndSourceFileAction() {