Exit ephemeral runners on broker acknowledge job-not-found (#4540)

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

src/Runner.Listener/Runner.cs

@@ -687,7 +687,7 @@ namespace GitHub.Runner.Listener
                                 else
                                 {
                                     var messageRef = StringUtil.ConvertFromJson<RunnerJobRequestRef>(message.Body);
-
+                                    
                                     // Acknowledge (best-effort)
                                     if (messageRef.ShouldAcknowledge) // Temporary feature flag
                                     {
@@ -695,6 +695,12 @@ namespace GitHub.Runner.Listener
                                         {
                                             await _listener.AcknowledgeMessageAsync(messageRef.RunnerRequestId, messageQueueLoopTokenSource.Token);
                                         }
+                                        catch (RunnerRequestJobNotFoundException) when (settings.Ephemeral)
+                                        {
+                                            Trace.Info($"Acknowledge returned job-not-found for ephemeral runner request '{messageRef.RunnerRequestId}'. Exiting runner.");
+                                            runOnceJobCompleted = true;
+                                            return Constants.Runner.ReturnCode.Success;
+                                        }
                                         catch (Exception ex)
                                         {
                                             Trace.Error($"Best-effort acknowledge failed for request '{messageRef.RunnerRequestId}'");

src/Sdk/RSWebApi/Contracts/BrokerErrorKind.cs

@@ -5,6 +5,7 @@ namespace GitHub.Actions.RunService.WebApi
     [DataContract]
     public class BrokerErrorKind
     {
+        public const string AcknowledgeJobNotFound = "AcknowledgeJobNotFound";
         public const string RunnerNotFound = "RunnerNotFound";
         public const string RunnerVersionTooOld = "RunnerVersionTooOld";
         public const string HostedRunnerDeprovisioned = "HostedRunnerDeprovisioned";

src/Sdk/WebApi/WebApi/BrokerHttpClient.cs

@@ -251,10 +251,10 @@ namespace GitHub.Actions.RunService.WebApi
             {
                 switch (brokerError.ErrorKind)
                 {
+                    case BrokerErrorKind.AcknowledgeJobNotFound:
+                        throw new RunnerRequestJobNotFoundException(brokerError.Message);
                     case BrokerErrorKind.RunnerNotFound:
                         throw new RunnerNotFoundException(brokerError.Message);
-                    default:
-                        break;
                 }
             }
 

src/Sdk/WebApi/WebApi/Exceptions/RunnerRequestJobNotFoundException.cs

@@ -0,0 +1,23 @@
+using System;
+
+namespace GitHub.Services.WebApi
+{
+    [Serializable]
+    public sealed class RunnerRequestJobNotFoundException : Exception
+    {
+        public RunnerRequestJobNotFoundException()
+            : base()
+        {
+        }
+
+        public RunnerRequestJobNotFoundException(String message)
+            : base(message)
+        {
+        }
+
+        public RunnerRequestJobNotFoundException(String message, Exception innerException)
+            : base(message, innerException)
+        {
+        }
+    }
+}

src/Test/L0/Listener/RunnerL0.cs

@@ -871,6 +871,202 @@ namespace GitHub.Runner.Common.Tests.Listener
             }
         }
 
+        [Fact]
+        [Trait("Level", "L0")]
+        [Trait("Category", "Runner")]
+        public async Task TestEphemeralRunnerJobRequestMessageFromRunServiceExitsOnAcknowledgeJobNotFound()
+        {
+            using (var hc = new TestHostContext(this))
+            {
+                //Arrange
+                var runner = new Runner.Listener.Runner();
+                hc.SetSingleton<IConfigurationManager>(_configurationManager.Object);
+                hc.SetSingleton<IJobNotification>(_jobNotification.Object);
+                hc.SetSingleton<IMessageListener>(_messageListener.Object);
+                hc.SetSingleton<IPromptManager>(_promptManager.Object);
+                hc.SetSingleton<IRunnerServer>(_runnerServer.Object);
+                hc.SetSingleton<IConfigurationStore>(_configStore.Object);
+                hc.SetSingleton<ISelfUpdater>(_updater.Object);
+                hc.SetSingleton<ICredentialManager>(_credentialManager.Object);
+                hc.EnqueueInstance<IErrorThrottler>(_acquireJobThrottler.Object);
+                hc.EnqueueInstance<IRunServer>(_runServer.Object);
+                hc.EnqueueInstance<IJobDispatcher>(_jobDispatcher.Object);
+
+                runner.Initialize(hc);
+                var settings = new RunnerSettings
+                {
+                    PoolId = 43242,
+                    AgentId = 5678,
+                    Ephemeral = true,
+                    ServerUrl = "https://github.com",
+                };
+
+                var message = new TaskAgentMessage()
+                {
+                    Body = JsonUtility.ToString(new RunnerJobRequestRef() { BillingOwnerId = "github", RunnerRequestId = "999", RunServiceUrl = "https://run-service.com", ShouldAcknowledge = true }),
+                    MessageId = 4234,
+                    MessageType = JobRequestMessageTypes.RunnerJobRequest
+                };
+
+                var messages = new Queue<TaskAgentMessage>();
+                messages.Enqueue(message);
+                _configurationManager.Setup(x => x.LoadSettings())
+                    .Returns(settings);
+                _configurationManager.Setup(x => x.IsConfigured())
+                    .Returns(true);
+                _messageListener.Setup(x => x.CreateSessionAsync(It.IsAny<CancellationToken>()))
+                    .Returns(Task.FromResult<CreateSessionResult>(CreateSessionResult.Success));
+                _messageListener.Setup(x => x.GetNextMessageAsync(It.IsAny<CancellationToken>()))
+                    .Returns(async (CancellationToken token) =>
+                        {
+                            if (0 == messages.Count)
+                            {
+                                await Task.Delay(2000, token);
+                            }
+
+                            return messages.Dequeue();
+                        });
+                _messageListener.Setup(x => x.AcknowledgeMessageAsync("999", It.IsAny<CancellationToken>()))
+                    .ThrowsAsync(new RunnerRequestJobNotFoundException("Job not found"));
+                _messageListener.Setup(x => x.DeleteSessionAsync())
+                    .Returns(Task.CompletedTask);
+                _messageListener.Setup(x => x.DeleteMessageAsync(It.IsAny<TaskAgentMessage>()))
+                    .Returns(Task.CompletedTask);
+                _jobNotification.Setup(x => x.StartClient(It.IsAny<String>()))
+                    .Callback(() =>
+                    {
+
+                    });
+
+                _configStore.Setup(x => x.IsServiceConfigured()).Returns(false);
+
+                //Act
+                var command = new CommandSettings(hc, new string[] { "run" });
+                Task<int> runnerTask = runner.ExecuteCommand(command);
+
+                //Assert
+                await Task.WhenAny(runnerTask, Task.Delay(30000));
+
+                Assert.True(runnerTask.IsCompleted, $"{nameof(runner.ExecuteCommand)} timed out.");
+                Assert.True(!runnerTask.IsFaulted, runnerTask.Exception?.ToString());
+                if (runnerTask.IsCompleted)
+                {
+                    Assert.Equal(Constants.Runner.ReturnCode.Success, await runnerTask);
+                }
+
+                _messageListener.Verify(x => x.AcknowledgeMessageAsync("999", It.IsAny<CancellationToken>()), Times.Once());
+                _jobDispatcher.Verify(x => x.Run(It.IsAny<Pipelines.AgentJobRequestMessage>(), It.IsAny<bool>()), Times.Never());
+                _runServer.Verify(x => x.GetJobMessageAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<CancellationToken>()), Times.Never());
+                _credentialManager.Verify(x => x.LoadCredentials(true), Times.Never());
+                _messageListener.Verify(x => x.DeleteSessionAsync(), Times.Once());
+                _messageListener.Verify(x => x.DeleteMessageAsync(It.IsAny<TaskAgentMessage>()), Times.Once());
+                _configurationManager.Verify(x => x.DeleteLocalRunnerConfig(), Times.Once());
+            }
+        }
+
+        [Fact]
+        [Trait("Level", "L0")]
+        [Trait("Category", "Runner")]
+        [Trait("Level", "L0")]
+        [Trait("Category", "Runner")]
+        public async Task TestRunnerJobRequestMessageFromRunServiceContinuesOnAcknowledgeJobNotFoundForPersistentRunner()
+        {
+            using (var hc = new TestHostContext(this))
+            {
+                //Arrange
+                var runner = new Runner.Listener.Runner();
+                hc.SetSingleton<IConfigurationManager>(_configurationManager.Object);
+                hc.SetSingleton<IJobNotification>(_jobNotification.Object);
+                hc.SetSingleton<IMessageListener>(_messageListener.Object);
+                hc.SetSingleton<IPromptManager>(_promptManager.Object);
+                hc.SetSingleton<IRunnerServer>(_runnerServer.Object);
+                hc.SetSingleton<IConfigurationStore>(_configStore.Object);
+                hc.SetSingleton<ISelfUpdater>(_updater.Object);
+                hc.SetSingleton<ICredentialManager>(_credentialManager.Object);
+                hc.EnqueueInstance<IErrorThrottler>(_acquireJobThrottler.Object);
+                hc.EnqueueInstance<IRunServer>(_runServer.Object);
+                hc.EnqueueInstance<IJobDispatcher>(_jobDispatcher.Object);
+
+                runner.Initialize(hc);
+                var settings = new RunnerSettings
+                {
+                    PoolId = 43242,
+                    AgentId = 5678,
+                    Ephemeral = false,
+                    ServerUrl = "https://github.com",
+                };
+
+                var message = new TaskAgentMessage()
+                {
+                    Body = JsonUtility.ToString(new RunnerJobRequestRef() { BillingOwnerId = "github", RunnerRequestId = "999", RunServiceUrl = "https://run-service.com", ShouldAcknowledge = true }),
+                    MessageId = 4234,
+                    MessageType = JobRequestMessageTypes.RunnerJobRequest
+                };
+
+                var messages = new Queue<TaskAgentMessage>();
+                messages.Enqueue(message);
+                var signalWorkerStarted = new SemaphoreSlim(0, 1);
+                _configurationManager.Setup(x => x.LoadSettings())
+                    .Returns(settings);
+                _configurationManager.Setup(x => x.IsConfigured())
+                    .Returns(true);
+                _messageListener.Setup(x => x.CreateSessionAsync(It.IsAny<CancellationToken>()))
+                    .Returns(Task.FromResult<CreateSessionResult>(CreateSessionResult.Success));
+                _messageListener.Setup(x => x.GetNextMessageAsync(It.IsAny<CancellationToken>()))
+                    .Returns(async (CancellationToken token) =>
+                        {
+                            if (0 == messages.Count)
+                            {
+                                await Task.Delay(2000, token);
+                            }
+
+                            return messages.Dequeue();
+                        });
+                _messageListener.Setup(x => x.AcknowledgeMessageAsync("999", It.IsAny<CancellationToken>()))
+                    .ThrowsAsync(new RunnerRequestJobNotFoundException("Job not found"));
+                _messageListener.Setup(x => x.DeleteSessionAsync())
+                    .Returns(Task.CompletedTask);
+                _messageListener.Setup(x => x.DeleteMessageAsync(It.IsAny<TaskAgentMessage>()))
+                    .Returns(Task.CompletedTask);
+                _jobNotification.Setup(x => x.StartClient(It.IsAny<String>()))
+                    .Callback(() =>
+                    {
+
+                    });
+                _runServer.Setup(x => x.GetJobMessageAsync("999", "github", It.IsAny<CancellationToken>()))
+                    .Returns(Task.FromResult(CreateJobRequestMessage("test")));
+                _credentialManager.Setup(x => x.LoadCredentials(true)).Returns(new VssCredentials());
+                _jobDispatcher.Setup(x => x.Run(It.IsAny<Pipelines.AgentJobRequestMessage>(), false))
+                    .Callback(() =>
+                    {
+                        signalWorkerStarted.Release();
+                    });
+
+                _configStore.Setup(x => x.IsServiceConfigured()).Returns(false);
+
+                //Act
+                var command = new CommandSettings(hc, new string[] { "run" });
+                Task<int> runnerTask = runner.ExecuteCommand(command);
+
+                //Assert
+                if (!await signalWorkerStarted.WaitAsync(2000))
+                {
+                    Assert.Fail($"{nameof(_jobDispatcher.Object.Run)} was not invoked.");
+                }
+
+                hc.ShutdownRunner(ShutdownReason.UserCancelled);
+                await Task.WhenAny(runnerTask, Task.Delay(2000));
+
+                Assert.True(runnerTask.IsCompleted, $"{nameof(runner.ExecuteCommand)} timed out.");
+                Assert.True(runnerTask.IsCanceled);
+                _messageListener.Verify(x => x.AcknowledgeMessageAsync("999", It.IsAny<CancellationToken>()), Times.Once());
+                _jobDispatcher.Verify(x => x.Run(It.IsAny<Pipelines.AgentJobRequestMessage>(), false), Times.Once());
+                _runServer.Verify(x => x.GetJobMessageAsync("999", "github", It.IsAny<CancellationToken>()), Times.Once());
+                _credentialManager.Verify(x => x.LoadCredentials(true), Times.Once());
+                _configurationManager.Verify(x => x.DeleteLocalRunnerConfig(), Times.Never());
+            }
+        }
+
         [Fact]
         [Trait("Level", "L0")]
         [Trait("Category", "Runner")]

src/Test/L0/Sdk/RSWebApi/BrokerHttpClientL0.cs

@@ -0,0 +1,76 @@
+using System;
+using System.Net;
+using System.Net.Http;
+using System.Text;
+using System.Threading;
+using System.Threading.Tasks;
+using GitHub.DistributedTask.WebApi;
+using GitHub.Runner.Sdk;
+using GitHub.Services.WebApi;
+using Xunit;
+
+namespace GitHub.Actions.RunService.WebApi.Tests;
+
+public sealed class BrokerHttpClientL0
+{
+    [Fact]
+    public async Task AcknowledgeRunnerRequestAsyncThrowsRunnerRequestJobNotFoundException()
+    {
+        using var client = CreateClient(
+            HttpStatusCode.NotFound,
+            new BrokerError
+            {
+                Source = "actions-broker-listener",
+                ErrorKind = BrokerErrorKind.AcknowledgeJobNotFound,
+                StatusCode = (int)HttpStatusCode.NotFound,
+                Message = "Assigned job no longer exists",
+            });
+
+        await Assert.ThrowsAsync<RunnerRequestJobNotFoundException>(() =>
+            client.AcknowledgeRunnerRequestAsync("runner-request", Guid.NewGuid(), "2.0.0", TaskAgentStatus.Online, cancellationToken: CancellationToken.None));
+    }
+
+    [Fact]
+    public async Task AcknowledgeRunnerRequestAsyncKeepsRunnerNotFoundClassification()
+    {
+        using var client = CreateClient(
+            HttpStatusCode.NotFound,
+            new BrokerError
+            {
+                Source = "actions-broker-listener",
+                ErrorKind = BrokerErrorKind.RunnerNotFound,
+                StatusCode = (int)HttpStatusCode.NotFound,
+                Message = "Runner not found",
+            });
+
+        await Assert.ThrowsAsync<RunnerNotFoundException>(() =>
+            client.AcknowledgeRunnerRequestAsync("runner-request", Guid.NewGuid(), "2.0.0", TaskAgentStatus.Online, cancellationToken: CancellationToken.None));
+    }
+
+    private static BrokerHttpClient CreateClient(HttpStatusCode statusCode, BrokerError brokerError)
+    {
+        return new BrokerHttpClient(
+            new Uri("https://broker.actions.githubusercontent.com/"),
+            new StaticResponseHandler(new HttpResponseMessage(statusCode)
+            {
+                Content = new StringContent(JsonUtility.ToString(brokerError), Encoding.UTF8, "application/json"),
+            }),
+            disposeHandler: true);
+    }
+
+    private sealed class StaticResponseHandler : HttpMessageHandler
+    {
+        private readonly HttpResponseMessage _response;
+
+        public StaticResponseHandler(HttpResponseMessage response)
+        {
+            _response = response;
+        }
+
+        protected override Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
+        {
+            _response.RequestMessage = request;
+            return Task.FromResult(_response);
+        }
+    }
+}