From 14856e63bc026c3e17916a99623d80e92e2a6efb Mon Sep 17 00:00:00 2001
From: Tingluo Huang <tingluohuang@github.com>
Date: Wed, 16 Jul 2025 14:11:09 -0400
Subject: [PATCH] Try add orchestrationid into user-agent using token claim.
 (#3945)

---
 src/Runner.Common/HostContext.cs     | 31 ++++++++++++++++++++++++++++
 src/Runner.Listener/JobDispatcher.cs |  7 ++++++-
 src/Runner.Worker/JobRunner.cs       |  7 +++++--
 3 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/src/Runner.Common/HostContext.cs b/src/Runner.Common/HostContext.cs
index 73ca108ae..ffb08684a 100644
--- a/src/Runner.Common/HostContext.cs
+++ b/src/Runner.Common/HostContext.cs
@@ -15,6 +15,7 @@ using System.Threading.Tasks;
 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 @@ namespace GitHub.Runner.Common
                 {
                     _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);
diff --git a/src/Runner.Listener/JobDispatcher.cs b/src/Runner.Listener/JobDispatcher.cs
index b908c8b07..f98204b42 100644
--- a/src/Runner.Listener/JobDispatcher.cs
+++ b/src/Runner.Listener/JobDispatcher.cs
@@ -110,7 +110,12 @@ namespace GitHub.Runner.Listener
             {
                 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");
diff --git a/src/Runner.Worker/JobRunner.cs b/src/Runner.Worker/JobRunner.cs
index 324474706..1af7d8b3b 100644
--- a/src/Runner.Worker/JobRunner.cs
+++ b/src/Runner.Worker/JobRunner.cs
@@ -50,8 +50,11 @@ namespace GitHub.Runner.Worker
             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);