From c2307b3c9260c4a52b86f6c264f502f71d8596a3 Mon Sep 17 00:00:00 2001
From: Luke Tomlinson <luketomlinson@github.com>
Date: Thu, 23 Mar 2023 08:02:21 -0700
Subject: [PATCH] .

---
 src/Runner.Listener/MessageListener.cs |  2 +-
 src/Runner.Listener/Runner.cs          | 14 +++++++++++++-
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/src/Runner.Listener/MessageListener.cs b/src/Runner.Listener/MessageListener.cs
index 80c0227ea..cb8edb607 100644
--- a/src/Runner.Listener/MessageListener.cs
+++ b/src/Runner.Listener/MessageListener.cs
@@ -16,7 +16,7 @@ using GitHub.Services.OAuth;
 
 namespace GitHub.Runner.Listener
 {
-    [ServiceLocator(Default = typeof(BrokerMessageListener))]
+    [ServiceLocator(Default = typeof(MessageListener))]
     public interface IMessageListener : IRunnerService
     {
         Task<Boolean> CreateSessionAsync(CancellationToken token);
diff --git a/src/Runner.Listener/Runner.cs b/src/Runner.Listener/Runner.cs
index c727f1b38..e95379ad3 100644
--- a/src/Runner.Listener/Runner.cs
+++ b/src/Runner.Listener/Runner.cs
@@ -339,13 +339,25 @@ namespace GitHub.Runner.Listener
             }
         }
 
+        private IMessageListener GetMesageListener(RunnerSettings settings)
+        {
+            if (settings.UseV2Flow)
+            {
+                var brokerListener = new BrokerMessageListener();
+                brokerListener.Initialize(HostContext);
+                return brokerListener;
+            }
+
+            return HostContext.GetService<IMessageListener>();
+        }
+
         //create worker manager, create message listener and start listening to the queue
         private async Task<int> RunAsync(RunnerSettings settings, bool runOnce = false)
         {
             try
             {
                 Trace.Info(nameof(RunAsync));
-                _listener = HostContext.GetService<IMessageListener>();
+                _listener = GetMesageListener(settings);
                 if (!await _listener.CreateSessionAsync(HostContext.RunnerShutdownToken))
                 {
                     return Constants.Runner.ReturnCode.TerminatedError;