Feature-flagged support for JobContext.CheckRunId (#3811)

Author: pje

src/Runner.Common/Constants.cs

@@ -163,6 +163,7 @@ public static class Features
                 public static readonly string LogTemplateErrorsAsDebugMessages = "DistributedTask.LogTemplateErrorsAsDebugMessages";
                 public static readonly string UseContainerPathForTemplate = "DistributedTask.UseContainerPathForTemplate";
                 public static readonly string AllowRunnerContainerHooks = "DistributedTask.AllowRunnerContainerHooks";
+                public static readonly string AddCheckRunIdToJobContext = "actions_add_check_run_id_to_job_context";
             }
 
             public static readonly string InternalTelemetryIssueDataKey = "_internal_telemetry";

src/Runner.Worker/ExecutionContext.cs

@@ -862,7 +862,21 @@ public void InitializeJob(Pipelines.AgentJobRequestMessage message, Cancellation
 
             ExpressionValues["secrets"] = Global.Variables.ToSecretsContext();
             ExpressionValues["runner"] = new RunnerContext();
-            ExpressionValues["job"] = new JobContext();
+
+            Trace.Info("Initializing Job context");
+            var jobContext = new JobContext();
+            if (Global.Variables.GetBoolean(Constants.Runner.Features.AddCheckRunIdToJobContext) ?? false)
+            {
+                ExpressionValues.TryGetValue("job", out var jobDictionary);
+                if (jobDictionary != null)
+                {
+                    foreach (var pair in jobDictionary.AssertDictionary("job"))
+                    {
+                        jobContext[pair.Key] = pair.Value;
+                    }
+                }
+            }
+            ExpressionValues["job"] = jobContext;
 
             Trace.Info("Initialize GitHub context");
             var githubAccessToken = new StringContextData(Global.Variables.Get("system.github.token"));

src/Runner.Worker/JobContext.cs

@@ -1,4 +1,4 @@
-using GitHub.DistributedTask.Pipelines.ContextData;
+using GitHub.DistributedTask.Pipelines.ContextData;
 using GitHub.Runner.Common.Util;
 using GitHub.Runner.Common;
 
@@ -56,5 +56,31 @@ public DictionaryContextData Container
                 }
             }
         }
+
+        public double? CheckRunId
+        {
+            get
+            {
+                if (this.TryGetValue("check_run_id", out var value) && value is NumberContextData number)
+                {
+                    return number.Value;
+                }
+                else
+                {
+                    return null;
+                }
+            }
+            set
+            {
+                if (value.HasValue)
+                {
+                    this["check_run_id"] = new NumberContextData(value.Value);
+                }
+                else
+                {
+                    this["check_run_id"] = null;
+                }
+            }
+        }
     }
 }

src/Test/L0/Worker/ExecutionContextL0.cs

@@ -1168,6 +1168,77 @@ public void ActionVariables_SecretsPrecedenceForDebugUsingVars()
             }
         }
 
+        [Fact]
+        [Trait("Level", "L0")]
+        [Trait("Category", "Worker")]
+        public void InitializeJob_HydratesJobContextWithCheckRunId()
+        {
+            using (TestHostContext hc = CreateTestContext())
+            {
+                // Arrange: Create a job request message and make sure the feature flag is enabled
+                var variables = new Dictionary<string, VariableValue>()
+                {
+                    [Constants.Runner.Features.AddCheckRunIdToJobContext] = new VariableValue("true"),
+                };
+                var jobRequest = new Pipelines.AgentJobRequestMessage(new TaskOrchestrationPlanReference(), new TimelineReference(), Guid.NewGuid(), "some job name", "some job name", null, null, null, variables, new List<MaskHint>(), new Pipelines.JobResources(), new Pipelines.ContextData.DictionaryContextData(), new Pipelines.WorkspaceOptions(), new List<Pipelines.ActionStep>(), null, null, null, null, null);
+                var pagingLogger = new Moq.Mock<IPagingLogger>();
+                var jobServerQueue = new Moq.Mock<IJobServerQueue>();
+                hc.EnqueueInstance(pagingLogger.Object);
+                hc.SetSingleton(jobServerQueue.Object);
+                var ec = new Runner.Worker.ExecutionContext();
+                ec.Initialize(hc);
+
+                // Arrange: Add check_run_id to the job context
+                var jobContext = new Pipelines.ContextData.DictionaryContextData();
+                jobContext["check_run_id"] = new NumberContextData(123456);
+                jobRequest.ContextData["job"] = jobContext;
+                jobRequest.ContextData["github"] = new Pipelines.ContextData.DictionaryContextData();
+
+                // Act
+                ec.InitializeJob(jobRequest, CancellationToken.None);
+
+                // Assert
+                Assert.NotNull(ec.JobContext);
+                Assert.Equal(123456, ec.JobContext.CheckRunId);
+            }
+        }
+
+        // TODO: this test can be deleted when `AddCheckRunIdToJobContext` is fully rolled out
+        [Fact]
+        [Trait("Level", "L0")]
+        [Trait("Category", "Worker")]
+        public void InitializeJob_HydratesJobContextWithCheckRunId_FeatureFlagDisabled()
+        {
+            using (TestHostContext hc = CreateTestContext())
+            {
+                // Arrange: Create a job request message and make sure the feature flag is disabled
+                var variables = new Dictionary<string, VariableValue>()
+                {
+                    [Constants.Runner.Features.AddCheckRunIdToJobContext] = new VariableValue("false"),
+                };
+                var jobRequest = new Pipelines.AgentJobRequestMessage(new TaskOrchestrationPlanReference(), new TimelineReference(), Guid.NewGuid(), "some job name", "some job name", null, null, null, variables, new List<MaskHint>(), new Pipelines.JobResources(), new Pipelines.ContextData.DictionaryContextData(), new Pipelines.WorkspaceOptions(), new List<Pipelines.ActionStep>(), null, null, null, null, null);
+                var pagingLogger = new Moq.Mock<IPagingLogger>();
+                var jobServerQueue = new Moq.Mock<IJobServerQueue>();
+                hc.EnqueueInstance(pagingLogger.Object);
+                hc.SetSingleton(jobServerQueue.Object);
+                var ec = new Runner.Worker.ExecutionContext();
+                ec.Initialize(hc);
+
+                // Arrange: Add check_run_id to the job context
+                var jobContext = new Pipelines.ContextData.DictionaryContextData();
+                jobContext["check_run_id"] = new NumberContextData(123456);
+                jobRequest.ContextData["job"] = jobContext;
+                jobRequest.ContextData["github"] = new Pipelines.ContextData.DictionaryContextData();
+
+                // Act
+                ec.InitializeJob(jobRequest, CancellationToken.None);
+
+                // Assert
+                Assert.NotNull(ec.JobContext);
+                Assert.Null(ec.JobContext.CheckRunId); // with the feature flag disabled we should not have added a CheckRunId to the JobContext
+            }
+        }
+
         private bool ExpressionValuesAssertEqual(DictionaryContextData expect, DictionaryContextData actual)
         {
             foreach (var key in expect.Keys.ToList())

src/Test/L0/Worker/JobContextL0.cs

@@ -0,0 +1,38 @@
+using System;
+using GitHub.DistributedTask.Pipelines.ContextData;
+using GitHub.Runner.Worker;
+using Xunit;
+
+namespace GitHub.Runner.Common.Tests.Worker
+{
+    public class JobContextL0
+    {
+        [Fact]
+        public void CheckRunId_SetAndGet_WorksCorrectly()
+        {
+            var ctx = new JobContext();
+            ctx.CheckRunId = 12345;
+            Assert.Equal(12345, ctx.CheckRunId);
+            Assert.True(ctx.TryGetValue("check_run_id", out var value));
+            Assert.IsType<NumberContextData>(value);
+            Assert.Equal(12345, ((NumberContextData)value).Value);
+        }
+
+        [Fact]
+        public void CheckRunId_NotSet_ReturnsNull()
+        {
+            var ctx = new JobContext();
+            Assert.Null(ctx.CheckRunId);
+            Assert.False(ctx.TryGetValue("check_run_id", out var value));
+        }
+
+        [Fact]
+        public void CheckRunId_SetNull_RemovesKey()
+        {
+            var ctx = new JobContext();
+            ctx.CheckRunId = 12345;
+            ctx.CheckRunId = null;
+            Assert.Null(ctx.CheckRunId);
+        }
+    }
+}