Try add orchestrationid into user-agent using token claim. (#3945)

Author: TingluoHuang

src/Runner.Common/HostContext.cs

@@ -15,6 +15,7 @@
 using GitHub.DistributedTask.Logging;
 using GitHub.Runner.Common.Util;
 using GitHub.Runner.Sdk;
+using GitHub.Services.WebApi.Jwt;
 
 namespace GitHub.Runner.Common
 {
@@ -306,6 +307,36 @@ public void LoadDefaultUserAgents()
                 {
                     _userAgents.Add(new ProductInfoHeaderValue("ClientId", clientId));
                 }
+
+                // for Hosted runner, we can pull orchestrationId from JWT claims of the runner listening token.
+                if (credData != null &&
+                    credData.Scheme == Constants.Configuration.OAuthAccessToken &&
+                    credData.Data.TryGetValue(Constants.Runner.CommandLine.Args.Token, out var accessToken) &&
+                    !string.IsNullOrEmpty(accessToken))
+                {
+                    try
+                    {
+                        var jwt = JsonWebToken.Create(accessToken);
+                        var claims = jwt.ExtractClaims();
+                        var orchestrationId = claims.FirstOrDefault(x => string.Equals(x.Type, "orch_id", StringComparison.OrdinalIgnoreCase))?.Value;
+                        if (string.IsNullOrEmpty(orchestrationId))
+                        {
+                            // fallback to orchid for C# actions-service
+                            orchestrationId = claims.FirstOrDefault(x => string.Equals(x.Type, "orchid", StringComparison.OrdinalIgnoreCase))?.Value;
+                        }
+
+                        if (!string.IsNullOrEmpty(orchestrationId))
+                        {
+                            _trace.Info($"Pull OrchestrationId {orchestrationId} from runner JWT claims");
+                            _userAgents.Insert(0, new ProductInfoHeaderValue("OrchestrationId", orchestrationId));
+                        }
+                    }
+                    catch (Exception ex)
+                    {
+                        _trace.Error("Fail to extract OrchestrationId from runner JWT claims");
+                        _trace.Error(ex);
+                    }
+                }
             }
 
             var runnerFile = GetConfigFile(WellKnownConfigFile.Runner);

src/Runner.Listener/JobDispatcher.cs

@@ -110,7 +110,12 @@ public void Run(Pipelines.AgentJobRequestMessage jobRequestMessage, bool runOnce
             {
                 var jwt = JsonWebToken.Create(accessToken);
                 var claims = jwt.ExtractClaims();
-                orchestrationId = claims.FirstOrDefault(x => string.Equals(x.Type, "orchid", StringComparison.OrdinalIgnoreCase))?.Value;
+                orchestrationId = claims.FirstOrDefault(x => string.Equals(x.Type, "orch_id", StringComparison.OrdinalIgnoreCase))?.Value;
+                if (string.IsNullOrEmpty(orchestrationId))
+                {
+                    orchestrationId = claims.FirstOrDefault(x => string.Equals(x.Type, "orchid", StringComparison.OrdinalIgnoreCase))?.Value;
+                }
+
                 if (!string.IsNullOrEmpty(orchestrationId))
                 {
                     Trace.Info($"Pull OrchestrationId {orchestrationId} from JWT claims");

src/Runner.Worker/JobRunner.cs

@@ -50,8 +50,11 @@ public async Task<TaskResult> RunAsync(AgentJobRequestMessage message, Cancellat
             if (message.Variables.TryGetValue(Constants.Variables.System.OrchestrationId, out VariableValue orchestrationId) &&
                 !string.IsNullOrEmpty(orchestrationId.Value))
             {
-                // make the orchestration id the first item in the user-agent header to avoid get truncated in server log.
-                HostContext.UserAgents.Insert(0, new ProductInfoHeaderValue("OrchestrationId", orchestrationId.Value));
+                if (!HostContext.UserAgents.Any(x => string.Equals(x.Product.Name, "OrchestrationId", StringComparison.OrdinalIgnoreCase)))
+                {
+                    // make the orchestration id the first item in the user-agent header to avoid get truncated in server log.
+                    HostContext.UserAgents.Insert(0, new ProductInfoHeaderValue("OrchestrationId", orchestrationId.Value));
+                }
 
                 // make sure orchestration id is in the user-agent header.
                 VssUtil.InitializeVssClientSettings(HostContext.UserAgents, HostContext.WebProxy);