Exit ephemeral runners on broker acknowledge job-not-found

src/Runner.Listener/Runner.cs

@@ -495,6 +495,8 @@ private async Task<int> RunAsync(RunnerSettings settings, bool runOnce = false,
 
                 // Should we try to cleanup ephemeral runners
                 bool runOnceJobCompleted = false;
+                bool singleUseRunnerConsumed = false;
+                bool isSingleUseRunner = settings.Ephemeral || (runOnce && !returnRunOnceJobResult);
                 bool skipSessionDeletion = false;
                 bool restartSession = false; // Flag to indicate session restart
                 bool restartSessionPending = false;
@@ -687,14 +689,20 @@ private async Task<int> RunAsync(RunnerSettings settings, bool runOnce = false,
                                 else
                                 {
                                     var messageRef = StringUtil.ConvertFromJson<RunnerJobRequestRef>(message.Body);
-
+                                    
                                     // Acknowledge (best-effort)
                                     if (messageRef.ShouldAcknowledge) // Temporary feature flag
                                     {
                                         try
                                         {
                                             await _listener.AcknowledgeMessageAsync(messageRef.RunnerRequestId, messageQueueLoopTokenSource.Token);
                                         }
+                                        catch (RunnerRequestJobNotFoundException) when (isSingleUseRunner)
+                                        {
+                                            Trace.Info($"Acknowledge returned job-not-found for single-use runner request '{messageRef.RunnerRequestId}'. Exiting runner.");
+                                            singleUseRunnerConsumed = true;
+                                            return Constants.Runner.ReturnCode.Success;
+                                        }
                                         catch (Exception ex)
                                         {
                                             Trace.Error($"Best-effort acknowledge failed for request '{messageRef.RunnerRequestId}'");
@@ -859,7 +867,7 @@ await configUpdater.UpdateRunnerConfigAsync(
 
                     messageQueueLoopTokenSource.Dispose();
 
-                    if (settings.Ephemeral && runOnceJobCompleted)
+                    if (settings.Ephemeral && (runOnceJobCompleted || singleUseRunnerConsumed))
                     {
                         configManager.DeleteLocalRunnerConfig();
                     }

src/Sdk/WebApi/WebApi/BrokerHttpClient.cs

@@ -256,11 +256,24 @@ public async Task AcknowledgeRunnerRequestAsync(
                     default:
                         break;
                 }
+
+                if (IsAcknowledgeJobNotFound(result.StatusCode, brokerError))
+                {
+                    throw new RunnerRequestJobNotFoundException(brokerError.Message);
+                }
             }
 
             throw new Exception($"Failed to acknowledge runner request. Request to {requestUri} failed with status: {result.StatusCode}. Error message {result.Error}");
         }
 
+        private static bool IsAcknowledgeJobNotFound(HttpStatusCode statusCode, BrokerError brokerError)
+        {
+            return statusCode == HttpStatusCode.NotFound &&
+                brokerError?.StatusCode == (int)HttpStatusCode.NotFound &&
+                !string.Equals(brokerError.ErrorKind, BrokerErrorKind.RunnerNotFound, StringComparison.Ordinal) &&
+                string.Equals(brokerError.Message, "Job not found", StringComparison.OrdinalIgnoreCase);
+        }
+
         private static bool TryParseErrorBody(string errorBody, out BrokerError error)
         {
             if (!string.IsNullOrEmpty(errorBody))

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,200 @@ public async Task TestRunnerJobRequestMessageFromRunService()
             }
         }
 
+        [Fact]
+        [Trait("Level", "L0")]
+        [Trait("Category", "Runner")]
+        public async Task TestRunOnceRunnerJobRequestMessageFromRunServiceExitsOnAcknowledgeJobNotFound()
+        {
+            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")]
+        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 GitHub.Services.WebApi.Jwt;
+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",
+                StatusCode = (int)HttpStatusCode.NotFound,
+                Message = "Job not found",
+            });
+
+        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);
+        }
+    }
+}