Process snapshot tokens

src/Runner.Worker/JobExtension.cs

@@ -17,6 +17,7 @@
 using GitHub.Runner.Common.Util;
 using GitHub.Runner.Sdk;
 using GitHub.Services.Common;
+using Newtonsoft.Json;
 using Pipelines = GitHub.DistributedTask.Pipelines;
 
 namespace GitHub.Runner.Worker
@@ -404,6 +405,18 @@ public async Task<List<IStep>> InitializeJob(IExecutionContext jobContext, Pipel
                                                                           data: (object)jobHookData));
                     }
 
+                    // Register custom image creation post-job step if the "snapshot" token is present in the message.
+                    var snapshotRequest = templateEvaluator.EvaluateJobSnapshotRequest(message.Snapshot, jobContext.ExpressionValues, jobContext.ExpressionFunctions);
+                    if (snapshotRequest != null)
+                    {
+                        var snapshotOperationProvider = HostContext.GetService<ISnapshotOperationProvider>();
+                        jobContext.RegisterPostJobStep(new JobExtensionRunner(
+                            runAsync: (executionContext, _) => snapshotOperationProvider.CreateSnapshotRequestAsync(executionContext, snapshotRequest),
+                            condition: $"{PipelineTemplateConstants.Success}()",
+                            displayName: $"Create custom image",
+                            data: null));
+                    }
+
                     List<IStep> steps = new();
                     steps.AddRange(preJobSteps);
                     steps.AddRange(jobSteps);

src/Runner.Worker/SnapshotOperationProvider.cs

@@ -0,0 +1,35 @@
+#nullable enable
+using System.IO;
+using System.Threading.Tasks;
+using GitHub.DistributedTask.Pipelines;
+using GitHub.Runner.Common;
+using Newtonsoft.Json;
+
+namespace GitHub.Runner.Worker;
+
+[ServiceLocator(Default = typeof(SnapshotOperationProvider))]
+public interface ISnapshotOperationProvider : IRunnerService
+{
+    Task CreateSnapshotRequestAsync(IExecutionContext executionContext, Snapshot snapshotRequest);
+}
+
+public class SnapshotOperationProvider : RunnerService, ISnapshotOperationProvider
+{
+    public async Task CreateSnapshotRequestAsync(IExecutionContext executionContext, Snapshot snapshotRequest)
+    {
+        var snapshotRequestFilePath = Path.Combine(HostContext.GetDirectory(WellKnownDirectory.Root), ".snapshot", "request.json");
+        var snapshotRequestDirectoryPath = Path.GetDirectoryName(snapshotRequestFilePath);
+        if (snapshotRequestDirectoryPath != null)
+        {
+            Directory.CreateDirectory(snapshotRequestDirectoryPath);
+        }
+
+        var snapshotRequestJson = JsonConvert.SerializeObject(snapshotRequest);
+        await File.WriteAllTextAsync(snapshotRequestFilePath, snapshotRequestJson);
+        executionContext.Output($"A snapshot request was created with parameters: {snapshotRequestJson}");
+        executionContext.Output($"Request written to: {snapshotRequestFilePath}");
+        executionContext.Output("This request will be processed after the job completes. You will not receive any feedback on the snapshot process within the workflow logs of this job.");
+        executionContext.Output("If the snapshot process is successful, you should see a new image with the requested name in the list of available custom images when creating a new GitHub-hosted Runner.");
+        await Task.CompletedTask;
+    }
+}

src/Sdk/DTPipelines/Pipelines/AgentJobRequestMessage.cs

@@ -43,6 +43,7 @@ public AgentJobRequestMessage(
             TemplateToken jobOutputs,
             IList<TemplateToken> defaults,
             ActionsEnvironmentReference actionsEnvironment,
+            TemplateToken snapshot,
             String messageType = JobRequestMessageTypes.PipelineAgentJobRequest)
         {
             this.MessageType = messageType;
@@ -57,6 +58,7 @@ public AgentJobRequestMessage(
             this.Workspace = workspaceOptions;
             this.JobOutputs = jobOutputs;
             this.ActionsEnvironment = actionsEnvironment;
+            this.Snapshot = snapshot;
             m_variables = new Dictionary<String, VariableValue>(variables, StringComparer.OrdinalIgnoreCase);
             m_maskHints = new List<MaskHint>(maskHints);
             m_steps = new List<JobStep>(steps);
@@ -237,6 +239,13 @@ public ActionsEnvironmentReference ActionsEnvironment
             set;
         }
 
+        [DataMember(EmitDefaultValue = false)]
+        public TemplateToken Snapshot
+        {
+            get;
+            set;
+        }
+
         /// <summary>
         /// Gets the collection of variables associated with the current context.
         /// </summary>

src/Sdk/DTPipelines/Pipelines/ObjectTemplating/PipelineTemplateConstants.cs

@@ -29,6 +29,7 @@ public sealed class PipelineTemplateConstants
         public const String Id = "id";
         public const String If = "if";
         public const String Image = "image";
+        public const String ImageName = "image-name";
         public const String Include = "include";
         public const String Inputs = "inputs";
         public const String Job = "job";
@@ -60,6 +61,7 @@ public sealed class PipelineTemplateConstants
         public const String Services = "services";
         public const String Shell = "shell";
         public const String Skipped = "skipped";
+        public const String Snapshot = "snapshot";
         public const String StepEnv = "step-env";
         public const String StepIfResult = "step-if-result";
         public const String StepWith = "step-with";

src/Sdk/DTPipelines/Pipelines/ObjectTemplating/PipelineTemplateConverter.cs

@@ -346,6 +346,39 @@ internal static List<KeyValuePair<String, JobContainer>> ConvertToJobServiceCont
             return result;
         }
 
+        internal static Snapshot ConvertToJobSnapshotRequest(TemplateContext context, TemplateToken token)
+        {
+            string imageName = null;
+            if (token is StringToken snapshotStringLiteral)
+            {
+                imageName = snapshotStringLiteral.Value;
+            }
+            else
+            {
+                var snapshotMapping = token.AssertMapping($"{PipelineTemplateConstants.Snapshot}");
+                foreach (var snapshotPropertyPair in snapshotMapping)
+                {
+                    var propertyName = snapshotPropertyPair.Key.AssertString($"{PipelineTemplateConstants.Snapshot} key");
+                    switch (propertyName.Value)
+                    {
+                        case PipelineTemplateConstants.ImageName:
+                            imageName = snapshotPropertyPair.Value.AssertString($"{PipelineTemplateConstants.Snapshot} {propertyName}").Value;
+                            break;
+                        default:
+                            propertyName.AssertUnexpectedValue($"{PipelineTemplateConstants.Snapshot} key");
+                            break;
+                    }
+                }
+            }
+
+            if (String.IsNullOrEmpty(imageName))
+            {
+                return null;
+            }
+
+            return new Snapshot(imageName);
+        }
+
         private static ActionStep ConvertToStep(
             TemplateContext context,
             TemplateToken stepsItem,

src/Sdk/DTPipelines/Pipelines/ObjectTemplating/PipelineTemplateEvaluator.cs

@@ -370,6 +370,32 @@ public IList<KeyValuePair<String, JobContainer>> EvaluateJobServiceContainers(
             return result;
         }
 
+        public Snapshot EvaluateJobSnapshotRequest(TemplateToken token,
+            DictionaryContextData contextData,
+            IList<IFunctionInfo> expressionFunctions)
+        {
+            var result = default(Snapshot);
+
+            if (token != null && token.Type != TokenType.Null)
+            {
+                var context = CreateContext(contextData, expressionFunctions);
+                try
+                {
+                    token = TemplateEvaluator.Evaluate(context, PipelineTemplateConstants.Snapshot, token, 0, null, omitHeader: true);
+                    context.Errors.Check();
+                    result = PipelineTemplateConverter.ConvertToJobSnapshotRequest(context, token);
+                }
+                catch (Exception ex) when (!(ex is TemplateValidationException))
+                {
+                    context.Errors.Add(ex);
+                }
+
+                context.Errors.Check();
+            }
+
+            return result;
+        }
+
         private TemplateContext CreateContext(
             DictionaryContextData contextData,
             IList<IFunctionInfo> expressionFunctions,

src/Sdk/DTPipelines/Pipelines/Snapshot.cs

@@ -0,0 +1,17 @@
+using System;
+using System.Runtime.Serialization;
+
+namespace GitHub.DistributedTask.Pipelines
+{
+    [DataContract]
+    public class Snapshot
+    {
+        public Snapshot(string imageName)
+        {
+            ImageName = imageName;
+        }
+
+        [DataMember(EmitDefaultValue = false)]
+        public String ImageName { get; set; }
+    }
+}

src/Sdk/DTPipelines/workflow-v1.0.json

@@ -71,7 +71,8 @@
           "env": "job-env",
           "outputs": "job-outputs",
           "defaults": "job-defaults",
-          "steps": "steps"
+          "steps": "steps",
+          "snapshot": "snapshot"
         }
       }
     },
@@ -155,6 +156,24 @@
       }
     },
 
+    "snapshot": {
+      "one-of": [
+        "non-empty-string",
+        "snapshot-mapping"
+      ]
+    },
+
+    "snapshot-mapping": {
+      "mapping": {
+        "properties": {
+          "image-name": {
+            "type": "non-empty-string",
+            "required": true
+          }
+        }
+      }
+    },
+
     "runs-on": {
       "context": [
         "github",

src/Test/L0/Listener/JobDispatcherL0.cs

@@ -41,7 +41,7 @@ private Pipelines.AgentJobRequestMessage CreateJobRequestMessage()
             TaskOrchestrationPlanReference plan = new();
             TimelineReference timeline = null;
             Guid jobId = Guid.NewGuid();
-            var result = new Pipelines.AgentJobRequestMessage(plan, timeline, jobId, "someJob", "someJob", null, null, null, new Dictionary<string, VariableValue>(), new List<MaskHint>(), new Pipelines.JobResources(), new Pipelines.ContextData.DictionaryContextData(), new Pipelines.WorkspaceOptions(), new List<Pipelines.ActionStep>(), null, null, null, null);
+            var result = new Pipelines.AgentJobRequestMessage(plan, timeline, jobId, "someJob", "someJob", null, null, null, new Dictionary<string, VariableValue>(), new List<MaskHint>(), new Pipelines.JobResources(), new Pipelines.ContextData.DictionaryContextData(), new Pipelines.WorkspaceOptions(), new List<Pipelines.ActionStep>(), null, null, null, null, null);
             result.ContextData["github"] = new Pipelines.ContextData.DictionaryContextData();
             return result;
         }
@@ -806,7 +806,8 @@ private static AgentJobRequestMessage GetAgentJobRequestMessage()
                 },
                 null,
                 new List<TemplateToken>(),
-                new ActionsEnvironmentReference("env")
+                new ActionsEnvironmentReference("env"),
+                null
             );
             return message;
         }

src/Test/L0/Listener/RunnerL0.cs

@@ -42,7 +42,7 @@ private Pipelines.AgentJobRequestMessage CreateJobRequestMessage(string jobName)
             TaskOrchestrationPlanReference plan = new();
             TimelineReference timeline = null;
             Guid jobId = Guid.NewGuid();
-            return new Pipelines.AgentJobRequestMessage(plan, timeline, jobId, "test", "test", null, null, null, new Dictionary<string, VariableValue>(), new List<MaskHint>(), new Pipelines.JobResources(), new Pipelines.ContextData.DictionaryContextData(), new Pipelines.WorkspaceOptions(), new List<Pipelines.ActionStep>(), null, null, null, null);
+            return new Pipelines.AgentJobRequestMessage(plan, timeline, jobId, "test", "test", null, null, null, new Dictionary<string, VariableValue>(), new List<MaskHint>(), new Pipelines.JobResources(), new Pipelines.ContextData.DictionaryContextData(), new Pipelines.WorkspaceOptions(), new List<Pipelines.ActionStep>(), null, null, null, null, null);
         }
 
         private JobCancelMessage CreateJobCancelMessage()

src/Test/L0/Worker/ActionCommandManagerL0.cs

@@ -232,7 +232,7 @@ public void EchoProcessCommandDebugOn()
                 TimelineReference timeline = new();
                 Guid jobId = Guid.NewGuid();
                 string jobName = "some job name";
-                var jobRequest = new Pipelines.AgentJobRequestMessage(plan, timeline, jobId, jobName, jobName, null, null, null, new Dictionary<string, VariableValue>(), new List<MaskHint>(), new Pipelines.JobResources(), new Pipelines.ContextData.DictionaryContextData(), new Pipelines.WorkspaceOptions(), new List<Pipelines.ActionStep>(), null, null, null, null);
+                var jobRequest = new Pipelines.AgentJobRequestMessage(plan, timeline, jobId, jobName, jobName, null, null, null, new Dictionary<string, VariableValue>(), new List<MaskHint>(), new Pipelines.JobResources(), new Pipelines.ContextData.DictionaryContextData(), new Pipelines.WorkspaceOptions(), new List<Pipelines.ActionStep>(), null, null, null, null, null);
                 jobRequest.Resources.Repositories.Add(new Pipelines.RepositoryResource()
                 {
                     Alias = Pipelines.PipelineConstants.SelfAlias,

src/Test/L0/Worker/CreateStepSummaryCommandL0.cs

@@ -193,7 +193,7 @@ private TestHostContext Setup([CallerMemberName] string name = "")
             TimelineReference timeline = new();
             Guid jobId = Guid.NewGuid();
             string jobName = "Summary Job";
-            var jobRequest = new Pipelines.AgentJobRequestMessage(plan, timeline, jobId, jobName, jobName, null, null, null, new Dictionary<string, VariableValue>(), new List<MaskHint>(), new Pipelines.JobResources(), new Pipelines.ContextData.DictionaryContextData(), new Pipelines.WorkspaceOptions(), new List<Pipelines.ActionStep>(), null, null, null, null);
+            var jobRequest = new Pipelines.AgentJobRequestMessage(plan, timeline, jobId, jobName, jobName, null, null, null, new Dictionary<string, VariableValue>(), new List<MaskHint>(), new Pipelines.JobResources(), new Pipelines.ContextData.DictionaryContextData(), new Pipelines.WorkspaceOptions(), new List<Pipelines.ActionStep>(), null, null, null, null, null);
             jobRequest.Resources.Repositories.Add(new Pipelines.RepositoryResource()
             {
                 Alias = Pipelines.PipelineConstants.SelfAlias,