Handle TimeoutException During Agent Replacement in agent configuration

src/Agent.Listener/Configuration/ConfigurationManager.cs

@@ -23,6 +23,7 @@
 using Newtonsoft.Json;
 using Microsoft.Win32;
 using Microsoft.VisualStudio.Services.Agent.Listener.Telemetry;
+using Microsoft.TeamFoundation.Test.WebApi;
 
 namespace Microsoft.VisualStudio.Services.Agent.Listener.Configuration
 {
@@ -45,6 +46,8 @@ public sealed class ConfigurationManager : AgentService, IConfigurationManager
 
         private const string VsTelemetryRegPath = @"SOFTWARE\Microsoft\VisualStudio\Telemetry\PersistentPropertyBag\c57a9efce9b74de382d905a89852db71";
         private const string VsTelemetryRegKey = "IsPipelineAgent";
+        private const int _maxRetries = 3;
+        private const int _delaySeconds = 2;
 
         public override void Initialize(IHostContext hostContext)
         {
@@ -252,18 +255,15 @@ public async Task ConfigureAsync(CommandSettings command)
                     {
                         // Update existing agent with new PublicKey, agent version and SystemCapabilities.
                         agent = UpdateExistingAgent(agent, publicKey, systemCapabilities);
-
-                        try
+                        agent = await UpdateAgentWithRetryAsync<TaskAgent>(
+                            () => agentProvider.UpdateAgentAsync(agentSettings, agent, command),
+                            command
+                          );
+                        if (agent != null)
                         {
-                            agent = await agentProvider.UpdateAgentAsync(agentSettings, agent, command);
                             _term.WriteLine(StringUtil.Loc("AgentReplaced"));
                             break;
                         }
-                        catch (Exception e) when (!command.Unattended())
-                        {
-                            _term.WriteError(e);
-                            _term.WriteError(StringUtil.Loc("FailedToReplaceAgent"));
-                        }
                     }
                     else if (command.Unattended())
                     {
@@ -455,6 +455,58 @@ public async Task ConfigureAsync(CommandSettings command)
             }
         }
 
+        private async Task<T> UpdateAgentWithRetryAsync<T>(
+            Func<Task<T>> operation,
+            CommandSettings command)
+        {
+            int attempt = 0;
+            while (true)
+            {
+                try
+                {
+                    return await operation();
+                }
+                catch (Exception e) when (
+                    e is TimeoutException ||
+                    e is TaskCanceledException ||
+                    (e is OperationCanceledException && !(e is TaskCanceledException))
+                )
+                {
+                    attempt++;
+                    if (command.Unattended())
+                    {
+                        if (attempt >= _maxRetries)
+                        {
+                            _term.WriteError(e);
+                            _term.WriteError(StringUtil.Loc("FailedToReplaceAgent"));
+                            Trace.Error($"{operation.Method.Name} failed after maximum retries. Exception: {e}");
+                            throw new InvalidOperationException(StringUtil.Loc("FailedToReplaceAgent"), e);
+                        }
+                        else
+                        {
+                            Trace.Info($"Retrying operation, Attempt: '{attempt}'.");
+                            int backoff = _delaySeconds * (int)Math.Pow(2, attempt - 1);
+                            _term.WriteLine(StringUtil.Loc("RetryingReplaceAgent", attempt, _maxRetries, backoff));
+                            await Task.Delay(TimeSpan.FromSeconds(backoff));
+                        }
+                    }
+                    else
+                    {
+                        _term.WriteError(e);
+                        _term.WriteError(StringUtil.Loc("FailedToReplaceAgent"));
+                        break;
+                    }
+                }
+                catch (Exception e)
+                {
+                    _term.WriteError(e);
+                    _term.WriteError(StringUtil.Loc("FailedToReplaceAgent"));
+                    break;
+                }
+            }
+            return default(T); 
+        }
+
         public async Task UnconfigureAsync(CommandSettings command)
         {
             ArgUtil.NotNull(command, nameof(command));
@@ -695,19 +747,18 @@ public async Task ReAuthAsync(CommandSettings command)
                     }
                     else
                     {
-                        try
+                        agent.Authorization = new TaskAgentAuthorization
                         {
-                            agent.Authorization = new TaskAgentAuthorization
-                            {
-                                PublicKey = new TaskAgentPublicKey(publicKey.Exponent, publicKey.Modulus),
-                            };
-                            agent = await agentProvider.UpdateAgentAsync(agentSettings, agent, command);
-                            _term.WriteLine(StringUtil.Loc("AgentReplaced"));
-                        }
-                        catch (Exception e) when (!command.Unattended())
+                            PublicKey = new TaskAgentPublicKey(publicKey.Exponent, publicKey.Modulus),
+                        };
+                        agent = await UpdateAgentWithRetryAsync<TaskAgent>(
+                            () => agentProvider.UpdateAgentAsync(agentSettings, agent, command),
+                            command
+                          );
+                        if (agent != null)
                         {
-                            _term.WriteError(e);
-                            _term.WriteError(StringUtil.Loc("FailedToReplaceAgent"));
+                            _term.WriteLine(StringUtil.Loc("AgentReplaced"));
+                            break;
                         }
                     }
 
@@ -838,7 +889,6 @@ private TaskAgent UpdateExistingAgent(TaskAgent agent, RSAParameters publicKey,
             {
                 agent.SystemCapabilities[capability.Key] = capability.Value ?? string.Empty;
             }
-
             return agent;
         }
 

src/Misc/layoutbin/en-US/strings.json

@@ -511,6 +511,7 @@
   "RestartMessage": "Restart the machine to launch agent and for autologon settings to take effect.",
   "ReStreamLogsToFilesError": "You cannot use --disableloguploads and --reStreamLogsToFiles at the same time!",
   "RetryCountLimitExceeded": "The maximum allowed number of attempts is {0} but got {1}. Retry attempts count will be decreased to {0}.",
+  "RetryingReplaceAgent": "Retrying to replace agent (attempt {0} of {1}). Waiting {2} seconds before next attempt...",
   "RMApiFailure": "Api {0} failed with an error code {1}",
   "RMArtifactContainerDetailsInvalidError": "The artifact does not have valid container details: {0}",
   "RMArtifactContainerDetailsNotFoundError": "The artifact does not contain container details: {0}",