Process snapshot tokens (#3135)

* Added Snapshot TemplateToken to AgentJobRequestMessage

* WIP for processing the snapshot token

* Changed snapshot post job step condition to Success, added comments

* Refactored snapshot post-job step

* Added evaluation of snapshot token to retrieve image name

* Added snapshot to workflow schema

* Fixed linter error

* Migrated snapshot logic to new SnapshotOperationProvider

* Fixed linter error

* Fixed linter errors

* Fixed linter error

* Fixed linter errors

* Updated L0 tests

* Fixed linter errors

* Added new JobExtensionL0 tests for snapshot post-job step

* Added JobExtensionL0 test case for snapshot mappings

* Added SnapshotOperationProviderL0 tests

* Enabled nullable types for SnapshotOperationProvider and its tests

* Added more assertions to SnapshotOperationProviderL0 tests

* Fixed linter errors

* Made sure TestHostContexts are disposed of properlyh in SnapshotOperationProviderL0 tests

* Resolved PR comments

* Fixed formatting

* Removed redundant reference

* Addressed PR comments

src/Runner.Worker/JobExtension.cs

@@ -392,6 +392,18 @@ namespace GitHub.Runner.Worker
                         }
                     }
 
+                    // 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));
+                    }
+
                     // Register Job Completed hook if the variable is set
                     var completedHookPath = Environment.GetEnvironmentVariable("ACTIONS_RUNNER_HOOK_JOB_COMPLETED");
                     if (!string.IsNullOrEmpty(completedHookPath))

src/Runner.Worker/SnapshotOperationProvider.cs

@@ -0,0 +1,32 @@
+#nullable enable
+using System.IO;
+using System.Threading.Tasks;
+using GitHub.DistributedTask.Pipelines;
+using GitHub.Runner.Common;
+using GitHub.Runner.Sdk;
+namespace GitHub.Runner.Worker;
+
+[ServiceLocator(Default = typeof(SnapshotOperationProvider))]
+public interface ISnapshotOperationProvider : IRunnerService
+{
+    Task CreateSnapshotRequestAsync(IExecutionContext executionContext, Snapshot snapshotRequest);
+}
+
+public class SnapshotOperationProvider : RunnerService, ISnapshotOperationProvider
+{
+    public 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);
+        }
+
+        IOUtil.SaveObject(snapshotRequest, snapshotRequestFilePath);
+        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.");
+        return Task.CompletedTask;
+    }
+}

src/Sdk/DTPipelines/Pipelines/AgentJobRequestMessage.cs

@@ -43,6 +43,7 @@ namespace GitHub.DistributedTask.Pipelines
             TemplateToken jobOutputs,
             IList<TemplateToken> defaults,
             ActionsEnvironmentReference actionsEnvironment,
+            TemplateToken snapshot,
             String messageType = JobRequestMessageTypes.PipelineAgentJobRequest)
         {
             this.MessageType = messageType;
@@ -57,6 +58,7 @@ namespace GitHub.DistributedTask.Pipelines
             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 @@ namespace GitHub.DistributedTask.Pipelines
             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 @@ namespace GitHub.DistributedTask.Pipelines.ObjectTemplating
         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 @@ namespace GitHub.DistributedTask.Pipelines.ObjectTemplating
         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 @@ namespace GitHub.DistributedTask.Pipelines.ObjectTemplating
             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 @@ namespace GitHub.DistributedTask.Pipelines.ObjectTemplating
             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 @@ namespace GitHub.Runner.Common.Tests.Listener
             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 @@ namespace GitHub.Runner.Common.Tests.Listener
                 },
                 null,
                 new List<TemplateToken>(),
-                new ActionsEnvironmentReference("env")
+                new ActionsEnvironmentReference("env"),
+                null
             );
             return message;
         }

src/Test/L0/Listener/RunnerL0.cs

@@ -42,7 +42,7 @@ namespace GitHub.Runner.Common.Tests.Listener
             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 @@ namespace GitHub.Runner.Common.Tests.Worker
                 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 @@ namespace GitHub.Runner.Common.Tests.Worker
             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,

src/Test/L0/Worker/ExecutionContextL0.cs

@@ -29,7 +29,7 @@ namespace GitHub.Runner.Common.Tests.Worker
                 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,
@@ -106,7 +106,7 @@ namespace GitHub.Runner.Common.Tests.Worker
                 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,
@@ -162,7 +162,7 @@ namespace GitHub.Runner.Common.Tests.Worker
                 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,
@@ -216,7 +216,7 @@ namespace GitHub.Runner.Common.Tests.Worker
                 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,
@@ -271,7 +271,7 @@ namespace GitHub.Runner.Common.Tests.Worker
                 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,
@@ -322,7 +322,7 @@ namespace GitHub.Runner.Common.Tests.Worker
                 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,
@@ -373,7 +373,7 @@ namespace GitHub.Runner.Common.Tests.Worker
                 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,
@@ -471,7 +471,7 @@ namespace GitHub.Runner.Common.Tests.Worker
                 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,
@@ -555,7 +555,7 @@ namespace GitHub.Runner.Common.Tests.Worker
                 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,
@@ -610,7 +610,7 @@ namespace GitHub.Runner.Common.Tests.Worker
                 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,
@@ -653,7 +653,7 @@ namespace GitHub.Runner.Common.Tests.Worker
                 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,
@@ -717,7 +717,7 @@ namespace GitHub.Runner.Common.Tests.Worker
                 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,
@@ -781,7 +781,7 @@ namespace GitHub.Runner.Common.Tests.Worker
                 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,
@@ -969,7 +969,7 @@ namespace GitHub.Runner.Common.Tests.Worker
                 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,
@@ -1014,7 +1014,7 @@ namespace GitHub.Runner.Common.Tests.Worker
                 TimelineReference timeline = new TimelineReference();
                 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,
@@ -1057,7 +1057,7 @@ namespace GitHub.Runner.Common.Tests.Worker
                 TimelineReference timeline = new TimelineReference();
                 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/JobExtensionL0.cs

@@ -4,6 +4,8 @@ using System.Linq;
 using System.Runtime.CompilerServices;
 using System.Threading;
 using System.Threading.Tasks;
+using GitHub.DistributedTask.ObjectTemplating.Tokens;
+using GitHub.DistributedTask.Pipelines.ObjectTemplating;
 using GitHub.DistributedTask.WebApi;
 using GitHub.Runner.Worker;
 using Moq;
@@ -25,6 +27,9 @@ namespace GitHub.Runner.Common.Tests.Worker
         private Mock<IContainerOperationProvider> _containerProvider;
         private Mock<IDiagnosticLogManager> _diagnosticLogManager;
         private Mock<IJobHookProvider> _jobHookProvider;
+        private Mock<ISnapshotOperationProvider> _snapshotOperationProvider;
+
+        private Pipelines.Snapshot _requestedSnapshot;
 
         private CancellationTokenSource _tokenSource;
         private TestHostContext CreateTestContext([CallerMemberName] String testName = "")
@@ -41,7 +46,16 @@ namespace GitHub.Runner.Common.Tests.Worker
             _directoryManager.Setup(x => x.PrepareDirectory(It.IsAny<IExecutionContext>(), It.IsAny<Pipelines.WorkspaceOptions>()))
                              .Returns(new TrackingConfig() { PipelineDirectory = "runner", WorkspaceDirectory = "runner/runner" });
             _jobHookProvider = new Mock<IJobHookProvider>();
+            _snapshotOperationProvider = new Mock<ISnapshotOperationProvider>();
 
+            _requestedSnapshot = null;
+            _snapshotOperationProvider
+                .Setup(p => p.CreateSnapshotRequestAsync(It.IsAny<IExecutionContext>(), It.IsAny<Pipelines.Snapshot>()))
+                .Returns((IExecutionContext _, object data) =>
+                {
+                    _requestedSnapshot = data as Pipelines.Snapshot;
+                    return Task.CompletedTask;
+                });
             IActionRunner step1 = new ActionRunner();
             IActionRunner step2 = new ActionRunner();
             IActionRunner step3 = new ActionRunner();
@@ -100,7 +114,7 @@ namespace GitHub.Runner.Common.Tests.Worker
             };
 
             Guid jobId = Guid.NewGuid();
-            _message = 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(), steps, null, null, null, null);
+            _message = 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(), steps, null, null, null, null, null);
             GitHubContext github = new();
             github["repository"] = new Pipelines.ContextData.StringContextData("actions/runner");
             github["secret_source"] = new Pipelines.ContextData.StringContextData("Actions");
@@ -125,6 +139,7 @@ namespace GitHub.Runner.Common.Tests.Worker
             hc.SetSingleton(_directoryManager.Object);
             hc.SetSingleton(_diagnosticLogManager.Object);
             hc.SetSingleton(_jobHookProvider.Object);
+            hc.SetSingleton(_snapshotOperationProvider.Object);
             hc.EnqueueInstance<IPagingLogger>(_logger.Object); // JobExecutionContext
             hc.EnqueueInstance<IPagingLogger>(_logger.Object); // job start hook
             hc.EnqueueInstance<IPagingLogger>(_logger.Object); // Initial Job
@@ -443,5 +458,80 @@ namespace GitHub.Runner.Common.Tests.Worker
                 Assert.Equal(0, _jobEc.PostJobSteps.Count);
             }
         }
+
+        [Fact]
+        [Trait("Level", "L0")]
+        [Trait("Category", "Worker")]
+        public async Task EnsureNoSnapshotPostJobStep()
+        {
+            using (TestHostContext hc = CreateTestContext())
+            {
+                var jobExtension = new JobExtension();
+                jobExtension.Initialize(hc);
+
+                _actionManager.Setup(x => x.PrepareActionsAsync(It.IsAny<IExecutionContext>(), It.IsAny<IEnumerable<Pipelines.JobStep>>(), It.IsAny<Guid>()))
+                    .Returns(Task.FromResult(new PrepareResult(new List<JobExtensionRunner>(), new Dictionary<Guid, IActionRunner>())));
+
+                _message.Snapshot = null;
+                await jobExtension.InitializeJob(_jobEc, _message);
+
+                var postJobSteps = _jobEc.PostJobSteps;
+                Assert.Equal(0, postJobSteps.Count);
+            }
+        }
+
+        [Fact]
+        [Trait("Level", "L0")]
+        [Trait("Category", "Worker")]
+        public Task EnsureSnapshotPostJobStepForStringToken()
+        {
+            var snapshot = new Pipelines.Snapshot("TestImageNameFromStringToken");
+            var imageNameValueStringToken = new StringToken(null, null, null, snapshot.ImageName);
+            return EnsureSnapshotPostJobStepForToken(imageNameValueStringToken, snapshot);
+        }
+
+        [Fact]
+        [Trait("Level", "L0")]
+        [Trait("Category", "Worker")]
+        public Task EnsureSnapshotPostJobStepForMappingToken()
+        {
+            var snapshot = new Pipelines.Snapshot("TestImageNameFromMappingToken");
+            var imageNameValueStringToken = new StringToken(null, null, null, snapshot.ImageName);
+            var mappingToken = new MappingToken(null, null, null)
+            {
+                { new StringToken(null,null,null, PipelineTemplateConstants.ImageName), imageNameValueStringToken }
+            };
+
+            return EnsureSnapshotPostJobStepForToken(mappingToken, snapshot);
+        }
+
+        private async Task EnsureSnapshotPostJobStepForToken(TemplateToken snapshotToken, Pipelines.Snapshot expectedSnapshot)
+        {
+            using (TestHostContext hc = CreateTestContext())
+            {
+                var jobExtension = new JobExtension();
+                jobExtension.Initialize(hc);
+
+                _actionManager.Setup(x => x.PrepareActionsAsync(It.IsAny<IExecutionContext>(), It.IsAny<IEnumerable<Pipelines.JobStep>>(), It.IsAny<Guid>()))
+                    .Returns(Task.FromResult(new PrepareResult(new List<JobExtensionRunner>(), new Dictionary<Guid, IActionRunner>())));
+
+                _message.Snapshot = snapshotToken;
+
+                await jobExtension.InitializeJob(_jobEc, _message);
+
+                var postJobSteps = _jobEc.PostJobSteps;
+
+                Assert.Equal(1, postJobSteps.Count);
+                var snapshotStep = postJobSteps.First();
+                Assert.Equal("Create custom image", snapshotStep.DisplayName);
+                Assert.Equal($"{PipelineTemplateConstants.Success}()", snapshotStep.Condition);
+
+                // Run the mock snapshot step, so we can verify it was executed with the expected snapshot object.
+                await snapshotStep.RunAsync();
+
+                Assert.NotNull(_requestedSnapshot);
+                Assert.Equal(expectedSnapshot.ImageName, _requestedSnapshot.ImageName);
+            }
+        }
     }
 }

src/Test/L0/Worker/JobRunnerL0.cs

@@ -101,6 +101,7 @@ namespace GitHub.Runner.Common.Tests.Worker
                 testName,
                 testName, 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,
                 new ActionsEnvironmentReference("staging"),
+                null,
                 messageType: messageType);
             message.Variables[Constants.Variables.System.Culture] = "en-US";
             message.Resources.Endpoints.Add(new ServiceEndpoint()

src/Test/L0/Worker/SnapshotOperationProviderL0.cs

@@ -0,0 +1,78 @@
+#nullable enable
+using System;
+using System.IO;
+using System.Runtime.CompilerServices;
+using GitHub.DistributedTask.Pipelines;
+using GitHub.Runner.Sdk;
+using GitHub.Runner.Worker;
+using Moq;
+using Xunit;
+
+namespace GitHub.Runner.Common.Tests.Worker;
+
+public class SnapshotOperationProviderL0
+{
+    private Mock<IExecutionContext>? _ec;
+    private SnapshotOperationProvider? _snapshotOperationProvider;
+    private string? _snapshotRequestFilePath;
+    private string? _snapshotRequestDirectoryPath;
+
+    [Theory]
+    [InlineData(true)]
+    [InlineData(false)]
+    [Trait("Level", "L0")]
+    [Trait("Category", "Worker")]
+    public async void CreateSnapshotRequestAsync(bool shouldSnapshotDirectoryAlreadyExist)
+    {
+        using (TestHostContext testHostContext = CreateTestHostContext())
+        {
+            //Arrange
+            Setup(testHostContext, shouldSnapshotDirectoryAlreadyExist);
+            var expectedSnapshot = new Snapshot(Guid.NewGuid().ToString());
+
+            //Act
+            await _snapshotOperationProvider!.CreateSnapshotRequestAsync(_ec!.Object, expectedSnapshot);
+
+            //Assert
+            var actualSnapshot = IOUtil.LoadObject<Snapshot>(_snapshotRequestFilePath);
+            Assert.NotNull(actualSnapshot);
+            Assert.Equal(expectedSnapshot.ImageName, actualSnapshot!.ImageName);
+            _ec.Verify(ec => ec.Write(null, $"Request written to: {_snapshotRequestFilePath}"), Times.Once);
+            _ec.Verify(ec => ec.Write(null, "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."), Times.Once);
+            _ec.Verify(ec => ec.Write(null, "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."), Times.Once);
+            _ec.VerifyNoOtherCalls();
+        }
+    }
+
+    private void Setup(IHostContext hostContext, bool shouldSnapshotDirectoryAlreadyExist)
+    {
+        _ec = new Mock<IExecutionContext>();
+        _snapshotOperationProvider = new SnapshotOperationProvider();
+        _snapshotOperationProvider.Initialize(hostContext);
+        _snapshotRequestFilePath = Path.Combine(hostContext.GetDirectory(WellKnownDirectory.Root), ".snapshot", "request.json");
+        _snapshotRequestDirectoryPath = Path.GetDirectoryName(_snapshotRequestFilePath);
+
+        if (_snapshotRequestDirectoryPath != null)
+        {
+            // Clean up any existing the snapshot directory and its contents before starting the test. 
+            if (Directory.Exists(_snapshotRequestDirectoryPath))
+            {
+                Directory.Delete(_snapshotRequestDirectoryPath, true);
+            }
+
+            if (shouldSnapshotDirectoryAlreadyExist)
+            {
+                // Create a fresh snapshot directory if it's required for the test. 
+                Directory.CreateDirectory(_snapshotRequestDirectoryPath);
+            }
+        }
+    }
+
+    private TestHostContext CreateTestHostContext([CallerMemberName] string testName = "")
+    {
+        var testHostContext = new TestHostContext(this, testName);
+        _ec = new Mock<IExecutionContext>();
+        _ec.Object.Initialize(testHostContext);
+        return testHostContext;
+    }
+}

src/Test/L0/Worker/WorkerL0.cs

@@ -67,7 +67,7 @@ namespace GitHub.Runner.Common.Tests.Worker
                     new Pipelines.ContextData.DictionaryContextData()
                 },
             };
-            var jobRequest = new Pipelines.AgentJobRequestMessage(plan, timeline, JobId, jobName, jobName, new StringToken(null, null, null, "ubuntu"), sidecarContainers, null, variables, new List<MaskHint>(), resources, context, null, actions, null, null, null, null);
+            var jobRequest = new Pipelines.AgentJobRequestMessage(plan, timeline, JobId, jobName, jobName, new StringToken(null, null, null, "ubuntu"), sidecarContainers, null, variables, new List<MaskHint>(), resources, context, null, actions, null, null, null, null, null);
             return jobRequest;
         }