Support pre/post/container/composite actions within composite actions (#1222)

Support Composite Actions with uses: steps

src/Runner.Worker/ActionCommandManager.cs

@@ -295,8 +295,21 @@ namespace GitHub.Runner.Worker
             {
                 throw new Exception("Required field 'name' is missing in ##[save-state] command.");
             }
-
-            context.IntraActionState[stateName] = command.Data;
+            // Embedded steps (composite) keep track of the state at the root level
+            if (context.IsEmbedded)
+            {
+                var id = context.EmbeddedId;
+                if (!context.Root.EmbeddedIntraActionState.ContainsKey(id))
+                {
+                    context.Root.EmbeddedIntraActionState[id] = new Dictionary<string, string>(StringComparer.OrdinalIgnoreCase);
+                }
+                context.Root.EmbeddedIntraActionState[id][stateName] = command.Data;
+            }
+            // Otherwise modify the ExecutionContext
+            else
+            {
+                context.IntraActionState[stateName] = command.Data;
+            }
             context.Debug($"Save intra-action state {stateName} = {command.Data}");
         }
 

src/Runner.Worker/ActionManager.cs

@@ -37,7 +37,10 @@ namespace GitHub.Runner.Worker
     public interface IActionManager : IRunnerService
     {
         Dictionary<Guid, ContainerInfo> CachedActionContainers { get; }
-        Task<PrepareResult> PrepareActionsAsync(IExecutionContext executionContext, IEnumerable<Pipelines.JobStep> steps);
+        Dictionary<Guid, List<Pipelines.ActionStep>> CachedEmbeddedPreSteps { get; }
+        Dictionary<Guid, List<Guid>> CachedEmbeddedStepIds { get; }
+        Dictionary<Guid, Stack<Pipelines.ActionStep>> CachedEmbeddedPostSteps { get; }
+        Task<PrepareResult> PrepareActionsAsync(IExecutionContext executionContext, IEnumerable<Pipelines.JobStep> steps, Guid rootStepId = default(Guid));
         Definition LoadAction(IExecutionContext executionContext, Pipelines.ActionStep action);
     }
 
@@ -48,10 +51,20 @@ namespace GitHub.Runner.Worker
         //81920 is the default used by System.IO.Stream.CopyTo and is under the large object heap threshold (85k).
         private const int _defaultCopyBufferSize = 81920;
         private const string _dotcomApiUrl = "https://api.github.com";
-        private readonly Dictionary<Guid, ContainerInfo> _cachedActionContainers = new Dictionary<Guid, ContainerInfo>();
 
+        private readonly Dictionary<Guid, ContainerInfo> _cachedActionContainers = new Dictionary<Guid, ContainerInfo>();
         public Dictionary<Guid, ContainerInfo> CachedActionContainers => _cachedActionContainers;
-        public async Task<PrepareResult> PrepareActionsAsync(IExecutionContext executionContext, IEnumerable<Pipelines.JobStep> steps)
+
+        private readonly Dictionary<Guid, List<Pipelines.ActionStep>> _cachedEmbeddedPreSteps = new Dictionary<Guid, List<Pipelines.ActionStep>>();
+        public Dictionary<Guid, List<Pipelines.ActionStep>> CachedEmbeddedPreSteps => _cachedEmbeddedPreSteps;
+
+        private readonly Dictionary<Guid, List<Guid>> _cachedEmbeddedStepIds = new Dictionary<Guid, List<Guid>>();
+        public Dictionary<Guid, List<Guid>> CachedEmbeddedStepIds => _cachedEmbeddedStepIds;
+
+        private readonly Dictionary<Guid, Stack<Pipelines.ActionStep>> _cachedEmbeddedPostSteps = new Dictionary<Guid, Stack<Pipelines.ActionStep>>();
+        public Dictionary<Guid, Stack<Pipelines.ActionStep>> CachedEmbeddedPostSteps => _cachedEmbeddedPostSteps;
+
+        public async Task<PrepareResult> PrepareActionsAsync(IExecutionContext executionContext, IEnumerable<Pipelines.JobStep> steps, Guid rootStepId = default(Guid))
         {
             // Assert inputs
             ArgUtil.NotNull(executionContext, nameof(executionContext));
@@ -64,10 +77,30 @@ namespace GitHub.Runner.Worker
                 PreStepTracker = new Dictionary<Guid, IActionRunner>()
             };
             var containerSetupSteps = new List<JobExtensionRunner>();
-            IOUtil.DeleteDirectory(HostContext.GetDirectory(WellKnownDirectory.Actions), executionContext.CancellationToken);
+            var depth = 0;
+            // We are running at the start of a job
+            if (rootStepId == default(Guid))
+            {
+                IOUtil.DeleteDirectory(HostContext.GetDirectory(WellKnownDirectory.Actions), executionContext.CancellationToken);
+            }
+            // We are running mid job due to a local composite action
+            else
+            {
+                if (!_cachedEmbeddedStepIds.ContainsKey(rootStepId))
+                {
+                    _cachedEmbeddedStepIds[rootStepId] = new List<Guid>();
+                    foreach (var compositeStep in steps)
+                    {
+                        var guid = Guid.NewGuid();
+                        compositeStep.Id = guid;
+                        _cachedEmbeddedStepIds[rootStepId].Add(guid);
+                    }
+                }
+                depth = 1;
+            }
             IEnumerable<Pipelines.ActionStep> actions = steps.OfType<Pipelines.ActionStep>();
             executionContext.Output("Prepare all required actions");
-            var result = await PrepareActionsRecursiveAsync(executionContext, state, actions, 0);
+            var result = await PrepareActionsRecursiveAsync(executionContext, state, actions, depth, rootStepId);
             if (state.ImagesToPull.Count > 0)
             {
                 foreach (var imageToPull in result.ImagesToPull)
@@ -103,7 +136,7 @@ namespace GitHub.Runner.Worker
             return new PrepareResult(containerSetupSteps, result.PreStepTracker);
         }
 
-        private async Task<PrepareActionsState> PrepareActionsRecursiveAsync(IExecutionContext executionContext, PrepareActionsState state, IEnumerable<Pipelines.ActionStep> actions, Int32 depth = 0)
+        private async Task<PrepareActionsState> PrepareActionsRecursiveAsync(IExecutionContext executionContext, PrepareActionsState state, IEnumerable<Pipelines.ActionStep> actions, Int32 depth = 0, Guid parentStepId = default(Guid))
         {
             ArgUtil.NotNull(executionContext, nameof(executionContext));
             if (depth > Constants.CompositeActionsMaxDepth)
@@ -187,24 +220,45 @@ namespace GitHub.Runner.Worker
                             state.ImagesToBuildInfo[setupInfo.Container.ActionRepository] = setupInfo.Container;
                         }
                     }
-                    else if(setupInfo != null && setupInfo.Steps != null && setupInfo.Steps.Count > 0)
+                    else if (setupInfo != null && setupInfo.Steps != null && setupInfo.Steps.Count > 0)
                     {
-                        state = await PrepareActionsRecursiveAsync(executionContext, state, setupInfo.Steps, depth + 1);
+                        state = await PrepareActionsRecursiveAsync(executionContext, state, setupInfo.Steps, depth + 1, action.Id);
                     }
                     var repoAction = action.Reference as Pipelines.RepositoryPathReference;
                     if (repoAction.RepositoryType != Pipelines.PipelineConstants.SelfAlias)
                     {
                         var definition = LoadAction(executionContext, action);
-                        // TODO: Support pre's in composite actions
-                        if (definition.Data.Execution.HasPre && depth < 1)
+                        if (definition.Data.Execution.HasPre)
                         {
-                            var actionRunner = HostContext.CreateService<IActionRunner>();
-                            actionRunner.Action = action;
-                            actionRunner.Stage = ActionRunStage.Pre;
-                            actionRunner.Condition = definition.Data.Execution.InitCondition;
-
                             Trace.Info($"Add 'pre' execution for {action.Id}");
-                            state.PreStepTracker[action.Id] = actionRunner;
+                            // Root Step
+                            if (depth < 1)
+                            {
+                                var actionRunner = HostContext.CreateService<IActionRunner>();
+                                actionRunner.Action = action;
+                                actionRunner.Stage = ActionRunStage.Pre;
+                                actionRunner.Condition = definition.Data.Execution.InitCondition;
+                                state.PreStepTracker[action.Id] = actionRunner;
+                            }
+                            // Embedded Step
+                            else 
+                            {
+                                if (!_cachedEmbeddedPreSteps.ContainsKey(parentStepId))
+                                {
+                                    _cachedEmbeddedPreSteps[parentStepId] = new List<Pipelines.ActionStep>();
+                                }
+                                _cachedEmbeddedPreSteps[parentStepId].Add(action);
+                            }
+                        }
+
+                        if (definition.Data.Execution.HasPost && depth > 0)
+                        {
+                            if (!_cachedEmbeddedPostSteps.ContainsKey(parentStepId))
+                            {
+                                // If we haven't done so already, add the parent to the post steps
+                                _cachedEmbeddedPostSteps[parentStepId] = new Stack<Pipelines.ActionStep>();
+                            }
+                            _cachedEmbeddedPostSteps[parentStepId].Push(action);
                         }
                     }
                 }
@@ -355,6 +409,29 @@ namespace GitHub.Runner.Worker
                         Trace.Verbose($"Details: {StringUtil.ConvertToJson(compositeAction?.Steps)}");
                         Trace.Info($"Load: {compositeAction.Outputs?.Count ?? 0} number of outputs");
                         Trace.Info($"Details: {StringUtil.ConvertToJson(compositeAction?.Outputs)}");
+
+                        if (CachedEmbeddedPreSteps.TryGetValue(action.Id, out var preSteps))
+                        {
+                            compositeAction.PreSteps = preSteps;
+                        }
+
+                        if (CachedEmbeddedPostSteps.TryGetValue(action.Id, out var postSteps))
+                        {
+                            compositeAction.PostSteps = postSteps;
+                        }
+
+                        if (_cachedEmbeddedStepIds.ContainsKey(action.Id))
+                        {
+                            for (var i = 0; i < compositeAction.Steps.Count; i++)
+                            {
+                                // Store Id's for later load actions
+                                compositeAction.Steps[i].Id = _cachedEmbeddedStepIds[action.Id][i];
+                                if (string.IsNullOrEmpty(executionContext.Global.Variables.Get("DistributedTask.EnableCompositeActions")) && compositeAction.Steps[i].Reference.Type != Pipelines.ActionSourceType.Script)
+                                {
+                                    throw new Exception("`uses:` keyword is not currently supported.");
+                                }
+                            }
+                        }
                     }
                     else
                     {
@@ -928,20 +1005,30 @@ namespace GitHub.Runner.Worker
                 }
                 else if (actionDefinitionData.Execution.ExecutionType == ActionExecutionType.Composite)
                 {
-                    // TODO: we need to generate unique Id's for composite steps
                     Trace.Info($"Loading Composite steps");
                     var compositeAction = actionDefinitionData.Execution as CompositeActionExecutionData;
                     setupInfo.Steps = compositeAction.Steps;
 
+                    // cache steps ids if not done so already
+                    if (!_cachedEmbeddedStepIds.ContainsKey(repositoryAction.Id))
+                    {
+                        _cachedEmbeddedStepIds[repositoryAction.Id] = new List<Guid>();
+                        foreach (var compositeStep in compositeAction.Steps)
+                        {
+                            var guid = Guid.NewGuid();
+                            compositeStep.Id = guid;
+                            _cachedEmbeddedStepIds[repositoryAction.Id].Add(guid);
+                        }
+                    }
+
+                    // TODO: remove once we remove the DistributedTask.EnableCompositeActions FF
                     foreach (var step in compositeAction.Steps)
                     {
-                        step.Id = Guid.NewGuid();
                         if (string.IsNullOrEmpty(executionContext.Global.Variables.Get("DistributedTask.EnableCompositeActions")) && step.Reference.Type != Pipelines.ActionSourceType.Script)
                         {
                             throw new Exception("`uses:` keyword is not currently supported.");
                         }
                     }
-
                     return setupInfo;
                 }
                 else
@@ -1102,9 +1189,11 @@ namespace GitHub.Runner.Worker
     public sealed class CompositeActionExecutionData : ActionExecutionData
     {
         public override ActionExecutionType ExecutionType => ActionExecutionType.Composite;
-        public override bool HasPre => false;
-        public override bool HasPost => false;
+        public override bool HasPre => PreSteps.Count > 0;
+        public override bool HasPost => PostSteps.Count > 0;
+        public List<Pipelines.ActionStep> PreSteps { get; set; }
         public List<Pipelines.ActionStep> Steps { get; set; }
+        public Stack<Pipelines.ActionStep> PostSteps { get; set; }
         public MappingToken Outputs { get; set; }
     }
 

src/Runner.Worker/ActionManifestManager.cs

@@ -480,6 +480,10 @@ namespace GitHub.Runner.Worker
                         return new CompositeActionExecutionData()
                         {
                             Steps = steps.Cast<Pipelines.ActionStep>().ToList(),
+                            PreSteps = new List<Pipelines.ActionStep>(),
+                            PostSteps = new Stack<Pipelines.ActionStep>(),
+                            InitCondition = "always()",
+                            CleanupCondition = "always()",
                             Outputs = outputs
                         };
                     }

src/Runner.Worker/ActionRunner.cs

@@ -88,6 +88,27 @@ namespace GitHub.Runner.Worker
             {
                 ExecutionContext.Warning($"`pre` execution is not supported for local action from '{repoAction.Path}'");
             }
+            List<JobExtensionRunner> localActionContainerSetupSteps = null;
+            // Handle Composite Local Actions
+            // Need to download and expand the tree of referenced actions
+            if (handlerData.ExecutionType == ActionExecutionType.Composite &&
+                handlerData is CompositeActionExecutionData compositeHandlerData &&
+                Stage == ActionRunStage.Main &&
+                Action.Reference is Pipelines.RepositoryPathReference localAction &&
+                string.Equals(localAction.RepositoryType, Pipelines.PipelineConstants.SelfAlias, StringComparison.OrdinalIgnoreCase))
+            {
+                var actionManager = HostContext.GetService<IActionManager>();
+                var prepareResult = await actionManager.PrepareActionsAsync(ExecutionContext, compositeHandlerData.Steps, ExecutionContext.Id);
+
+                // Reload definition since post may exist now (from embedded steps that were JIT downloaded)
+                definition = taskManager.LoadAction(ExecutionContext, Action);
+                ArgUtil.NotNull(definition, nameof(definition));
+                handlerData = definition.Data?.Execution;
+                ArgUtil.NotNull(handlerData, nameof(handlerData));
+
+                // Save container setup steps so we can reference them later
+                localActionContainerSetupSteps = prepareResult.ContainerSetupSteps;
+            }
 
             // The action has post cleanup defined.
             // we need to create timeline record for them and add them to the step list that StepRunner is using
@@ -249,7 +270,8 @@ namespace GitHub.Runner.Worker
                             inputs,
                             environment,
                             ExecutionContext.Global.Variables,
-                            actionDirectory: definition.Directory);
+                            actionDirectory: definition.Directory,
+                            localActionContainerSetupSteps: localActionContainerSetupSteps);
 
             // Print out action details
             handler.PrintActionDetails(Stage);

src/Runner.Worker/ExecutionContext.cs

@@ -36,6 +36,7 @@ namespace GitHub.Runner.Worker
     public interface IExecutionContext : IRunnerService
     {
         Guid Id { get; }
+        Guid EmbeddedId { get; }
         string ScopeName { get; }
         string ContextName { get; }
         Task ForceCompleted { get; }
@@ -58,6 +59,10 @@ namespace GitHub.Runner.Worker
 
         // Only job level ExecutionContext has PostJobSteps
         Stack<IStep> PostJobSteps { get; }
+        HashSet<Guid> EmbeddedStepsWithPostRegistered{ get; }
+
+        // Keep track of embedded steps states
+        Dictionary<Guid, Dictionary<string, string>> EmbeddedIntraActionState { get; }
 
         bool EchoOnActionCommand { get; set; }
 
@@ -68,8 +73,8 @@ namespace GitHub.Runner.Worker
         // Initialize
         void InitializeJob(Pipelines.AgentJobRequestMessage message, CancellationToken token);
         void CancelToken();
-        IExecutionContext CreateChild(Guid recordId, string displayName, string refName, string scopeName, string contextName, Dictionary<string, string> intraActionState = null, int? recordOrder = null, IPagingLogger logger = null, bool isEmbedded = false, CancellationTokenSource cancellationTokenSource = null);
-        IExecutionContext CreateEmbeddedChild(string scopeName, string contextName);
+        IExecutionContext CreateChild(Guid recordId, string displayName, string refName, string scopeName, string contextName, Dictionary<string, string> intraActionState = null, int? recordOrder = null, IPagingLogger logger = null, bool isEmbedded = false, CancellationTokenSource cancellationTokenSource = null, Guid embeddedId = default(Guid));
+        IExecutionContext CreateEmbeddedChild(string scopeName, string contextName, Guid embeddedId, Dictionary<string, string> intraActionState = null);
 
         // logging
         long Write(string tag, string message);
@@ -132,6 +137,7 @@ namespace GitHub.Runner.Worker
         private long _totalThrottlingDelayInMilliseconds = 0;
 
         public Guid Id => _record.Id;
+        public Guid EmbeddedId { get; private set; }
         public string ScopeName { get; private set; }
         public string ContextName { get; private set; }
         public Task ForceCompleted => _forceCompleted.Task;
@@ -155,6 +161,11 @@ namespace GitHub.Runner.Worker
         // Only job level ExecutionContext has StepsWithPostRegistered
         public HashSet<Guid> StepsWithPostRegistered { get; private set; }
 
+        // Only job level ExecutionContext has EmbeddedStepsWithPostRegistered
+        public HashSet<Guid> EmbeddedStepsWithPostRegistered { get; private set; }
+
+        public Dictionary<Guid, Dictionary<string, string>> EmbeddedIntraActionState { get; private set; }
+
         public bool EchoOnActionCommand { get; set; }
 
         // An embedded execution context shares the same record ID, record name, and logger
@@ -245,13 +256,15 @@ namespace GitHub.Runner.Worker
 
         public void RegisterPostJobStep(IStep step)
         {
-            // TODO: Remove when we support composite post job steps
             if (this.IsEmbedded)
             {
-                throw new Exception("Composite actions do not currently support post steps");
-
-            }
-            if (step is IActionRunner actionRunner && !Root.StepsWithPostRegistered.Add(actionRunner.Action.Id))
+                if (step is IActionRunner actionRunner && !Root.EmbeddedStepsWithPostRegistered.Add(actionRunner.Action.Id))
+                {
+                    Trace.Info($"'post' of '{actionRunner.DisplayName}' already push to child post step stack.");
+                }
+                return;
+            } 
+            else if (step is IActionRunner actionRunner && !Root.StepsWithPostRegistered.Add(actionRunner.Action.Id))
             {
                 Trace.Info($"'post' of '{actionRunner.DisplayName}' already push to post step stack.");
                 return;
@@ -261,7 +274,7 @@ namespace GitHub.Runner.Worker
             Root.PostJobSteps.Push(step);
         }
 
-        public IExecutionContext CreateChild(Guid recordId, string displayName, string refName, string scopeName, string contextName, Dictionary<string, string> intraActionState = null, int? recordOrder = null, IPagingLogger logger = null, bool isEmbedded = false, CancellationTokenSource cancellationTokenSource = null)
+        public IExecutionContext CreateChild(Guid recordId, string displayName, string refName, string scopeName, string contextName, Dictionary<string, string> intraActionState = null, int? recordOrder = null, IPagingLogger logger = null, bool isEmbedded = false, CancellationTokenSource cancellationTokenSource = null, Guid embeddedId = default(Guid))
         {
             Trace.Entering();
 
@@ -270,6 +283,7 @@ namespace GitHub.Runner.Worker
             child.Global = Global;
             child.ScopeName = scopeName;
             child.ContextName = contextName;
+            child.EmbeddedId = embeddedId;
             if (intraActionState == null)
             {
                 child.IntraActionState = new Dictionary<string, string>(StringComparer.OrdinalIgnoreCase);
@@ -317,9 +331,9 @@ namespace GitHub.Runner.Worker
         /// An embedded execution context shares the same record ID, record name, logger,
         /// and a linked cancellation token.
         /// </summary>
-        public IExecutionContext CreateEmbeddedChild(string scopeName, string contextName)
+        public IExecutionContext CreateEmbeddedChild(string scopeName, string contextName, Guid embeddedId, Dictionary<string, string> intraActionState = null)
         {
-            return Root.CreateChild(_record.Id, _record.Name, _record.Id.ToString("N"), scopeName, contextName, logger: _logger, isEmbedded: true, cancellationTokenSource: CancellationTokenSource.CreateLinkedTokenSource(_cancellationTokenSource.Token));
+            return Root.CreateChild(_record.Id, _record.Name, _record.Id.ToString("N"), scopeName, contextName, logger: _logger, isEmbedded: true, cancellationTokenSource: CancellationTokenSource.CreateLinkedTokenSource(_cancellationTokenSource.Token), intraActionState: intraActionState, embeddedId: embeddedId);
         }
 
         public void Start(string currentOperation = null)
@@ -375,7 +389,7 @@ namespace GitHub.Runner.Worker
 
             _logger.End();
 
-            // Skip if generated context name. Generated context names start with "__". After M271-ish the server will never send an empty context name.
+            // Skip if generated context name. Generated context names start with "__". After 3.2 the server will never send an empty context name.
             if (!string.IsNullOrEmpty(ContextName) && !ContextName.StartsWith("__", StringComparison.Ordinal))
             {
                 Global.StepsContext.SetOutcome(ScopeName, ContextName, (Outcome ?? Result ?? TaskResult.Succeeded).ToActionResult());
@@ -438,7 +452,7 @@ namespace GitHub.Runner.Worker
         {
             ArgUtil.NotNullOrEmpty(name, nameof(name));
 
-            // Skip if generated context name. Generated context names start with "__". After M271-ish the server will never send an empty context name.
+            // Skip if generated context name. Generated context names start with "__". After 3.2 the server will never send an empty context name.
             if (string.IsNullOrEmpty(ContextName) || ContextName.StartsWith("__", StringComparison.Ordinal))
             {
                 reference = null;
@@ -683,6 +697,12 @@ namespace GitHub.Runner.Worker
             // StepsWithPostRegistered for job ExecutionContext
             StepsWithPostRegistered = new HashSet<Guid>();
 
+            // EmbeddedStepsWithPostRegistered for job ExecutionContext
+            EmbeddedStepsWithPostRegistered = new HashSet<Guid>();
+
+            // EmbeddedIntraActionState for job ExecutionContext
+            EmbeddedIntraActionState = new Dictionary<Guid, Dictionary<string,string>>();
+
             // Job timeline record.
             InitializeTimelineRecord(
                 timelineId: message.Timeline.Id,

src/Runner.Worker/Handlers/CompositeActionHandler.cs

@@ -34,7 +34,32 @@ namespace GitHub.Runner.Worker.Handlers
             Trace.Entering();
             ArgUtil.NotNull(ExecutionContext, nameof(ExecutionContext));
             ArgUtil.NotNull(Inputs, nameof(Inputs));
-            ArgUtil.NotNull(Data.Steps, nameof(Data.Steps));
+
+            List<Pipelines.ActionStep> steps;
+
+            if (stage == ActionRunStage.Pre)
+            {
+                ArgUtil.NotNull(Data.PreSteps, nameof(Data.PreSteps));
+                steps = Data.PreSteps;
+            } 
+            else if (stage == ActionRunStage.Post)
+            {
+                ArgUtil.NotNull(Data.PostSteps, nameof(Data.PostSteps));
+                steps = Data.PostSteps.ToList();
+                // Only register post steps for steps that actually ran
+                foreach (var step in steps)
+                {
+                    if (!ExecutionContext.Root.EmbeddedStepsWithPostRegistered.Contains(step.Id))
+                    {
+                        steps.Remove(step);
+                    }
+                }
+            }  
+            else
+            {
+                ArgUtil.NotNull(Data.Steps, nameof(Data.Steps));
+                steps = Data.Steps;
+            }
 
             try
             {
@@ -45,7 +70,7 @@ namespace GitHub.Runner.Worker.Handlers
                     inputsData[i.Key] = new StringContextData(i.Value);
                 }
 
-                // Temporary hack until after M271-ish. After M271-ish the server will never send an empty
+                // Temporary hack until after 3.2. After 3.2 the server will never send an empty
                 // context name. Generated context names start with "__"
                 var childScopeName = ExecutionContext.GetFullyQualifiedContextName();
                 if (string.IsNullOrEmpty(childScopeName))
@@ -55,13 +80,28 @@ namespace GitHub.Runner.Worker.Handlers
 
                 // Create embedded steps
                 var embeddedSteps = new List<IStep>();
-                foreach (Pipelines.ActionStep stepData in Data.Steps)
+
+                 // If we need to setup containers beforehand, do it
+                // only relevant for local composite actions that need to JIT download/setup containers
+                if (LocalActionContainerSetupSteps != null && LocalActionContainerSetupSteps.Count > 0)
+                {
+                    foreach(var step in LocalActionContainerSetupSteps)
+                    {
+                        ArgUtil.NotNull(step, step.DisplayName);
+                        var stepId = $"__{Guid.NewGuid()}";
+                        step.ExecutionContext = ExecutionContext.CreateEmbeddedChild(childScopeName, stepId, Guid.NewGuid());
+                        embeddedSteps.Add(step);
+                    }
+                }
+
+                foreach (Pipelines.ActionStep stepData in steps)
                 {
                     var step = HostContext.CreateService<IActionRunner>();
                     step.Action = stepData;
                     step.Stage = stage;
                     step.Condition = stepData.Condition;
-                    step.ExecutionContext = ExecutionContext.CreateEmbeddedChild(childScopeName, stepData.ContextName);
+                    ExecutionContext.Root.EmbeddedIntraActionState.TryGetValue(step.Action.Id, out var intraActionState);
+                    step.ExecutionContext = ExecutionContext.CreateEmbeddedChild(childScopeName, stepData.ContextName, step.Action.Id, intraActionState: intraActionState);
                     step.ExecutionContext.ExpressionValues["inputs"] = inputsData;
                     step.ExecutionContext.ExpressionValues["steps"] = ExecutionContext.Global.StepsContext.GetScope(childScopeName);
 
@@ -84,7 +124,6 @@ namespace GitHub.Runner.Worker.Handlers
                 ExecutionContext.ExpressionValues["inputs"] = inputsData;
                 ExecutionContext.ExpressionValues["steps"] = ExecutionContext.Global.StepsContext.GetScope(childScopeName);
                 ProcessOutputs();
-                ExecutionContext.Global.StepsContext.ClearScope(childScopeName);
             }
             catch (Exception ex)
             {
@@ -178,16 +217,17 @@ namespace GitHub.Runner.Worker.Handlers
                     }
                 }
 
-                var actionStep = step as IActionRunner;
-
                 try
                 {
-                    // Evaluate and merge embedded-step env
-                    var templateEvaluator = step.ExecutionContext.ToPipelineTemplateEvaluator();
-                    var actionEnvironment = templateEvaluator.EvaluateStepEnvironment(actionStep.Action.Environment, step.ExecutionContext.ExpressionValues, step.ExecutionContext.ExpressionFunctions, Common.Util.VarUtil.EnvironmentVariableKeyComparer);
-                    foreach (var env in actionEnvironment)
+                    if (step is IActionRunner actionStep)
                     {
-                        envContext[env.Key] = new StringContextData(env.Value ?? string.Empty);
+                        // Evaluate and merge embedded-step env
+                        var templateEvaluator = step.ExecutionContext.ToPipelineTemplateEvaluator();
+                        var actionEnvironment = templateEvaluator.EvaluateStepEnvironment(actionStep.Action.Environment, step.ExecutionContext.ExpressionValues, step.ExecutionContext.ExpressionFunctions, Common.Util.VarUtil.EnvironmentVariableKeyComparer);
+                        foreach (var env in actionEnvironment)
+                        {
+                            envContext[env.Key] = new StringContextData(env.Value ?? string.Empty);
+                        }
                     }
                 }
                 catch (Exception ex)

src/Runner.Worker/Handlers/Handler.cs

@@ -20,6 +20,7 @@ namespace GitHub.Runner.Worker.Handlers
         IStepHost StepHost { get; set; }
         Dictionary<string, string> Inputs { get; set; }
         string ActionDirectory { get; set; }
+        List<JobExtensionRunner> LocalActionContainerSetupSteps { get; set; }
         Task RunAsync(ActionRunStage stage);
         void PrintActionDetails(ActionRunStage stage);
     }
@@ -41,6 +42,7 @@ namespace GitHub.Runner.Worker.Handlers
         public IStepHost StepHost { get; set; }
         public Dictionary<string, string> Inputs { get; set; }
         public string ActionDirectory { get; set; }
+        public List<JobExtensionRunner> LocalActionContainerSetupSteps { get; set; }
 
         public virtual void PrintActionDetails(ActionRunStage stage)
         {

src/Runner.Worker/Handlers/HandlerFactory.cs

@@ -19,7 +19,8 @@ namespace GitHub.Runner.Worker.Handlers
             Dictionary<string, string> inputs,
             Dictionary<string, string> environment,
             Variables runtimeVariables,
-            string actionDirectory);
+            string actionDirectory,
+            List<JobExtensionRunner> localActionContainerSetupSteps);
     }
 
     public sealed class HandlerFactory : RunnerService, IHandlerFactory
@@ -32,7 +33,8 @@ namespace GitHub.Runner.Worker.Handlers
             Dictionary<string, string> inputs,
             Dictionary<string, string> environment,
             Variables runtimeVariables,
-            string actionDirectory)
+            string actionDirectory,
+            List<JobExtensionRunner> localActionContainerSetupSteps)
         {
             // Validate args.
             Trace.Entering();
@@ -84,6 +86,7 @@ namespace GitHub.Runner.Worker.Handlers
             handler.StepHost = stepHost;
             handler.Inputs = inputs;
             handler.ActionDirectory = actionDirectory;
+            handler.LocalActionContainerSetupSteps = localActionContainerSetupSteps;
             return handler;
         }
     }

src/Test/L0/Worker/ActionManagerL0.cs

@@ -877,6 +877,7 @@ namespace GitHub.Runner.Common.Tests.Worker
                         }
                     }
                 };
+                _hc.EnqueueInstance<IActionRunner>(new Mock<IActionRunner>().Object);
 
                 //Act
                 var steps = (await _actionManager.PrepareActionsAsync(_ec.Object, actions)).ContainerSetupSteps;
@@ -967,8 +968,7 @@ namespace GitHub.Runner.Common.Tests.Worker
                 var result =  await _actionManager.PrepareActionsAsync(_ec.Object, actions);
 
                 //Assert
-                // TODO: Update this test
-                Assert.Equal(0, result.PreStepTracker.Count);
+                Assert.Equal(1, result.PreStepTracker.Count);
 
             }
             finally

src/Test/L0/Worker/ActionRunnerL0.cs

@@ -61,8 +61,8 @@ namespace GitHub.Runner.Common.Tests.Worker
             _actionRunner.Action = action;
 
             Dictionary<string, string> finialInputs = new Dictionary<string, string>();
-            _handlerFactory.Setup(x => x.Create(It.IsAny<IExecutionContext>(), It.IsAny<ActionStepDefinitionReference>(), It.IsAny<IStepHost>(), It.IsAny<ActionExecutionData>(), It.IsAny<Dictionary<string, string>>(), It.IsAny<Dictionary<string, string>>(), It.IsAny<Variables>(), It.IsAny<string>()))
-                           .Callback((IExecutionContext executionContext, Pipelines.ActionStepDefinitionReference actionReference, IStepHost stepHost, ActionExecutionData data, Dictionary<string, string> inputs, Dictionary<string, string> environment, Variables runtimeVariables, string taskDirectory) =>
+            _handlerFactory.Setup(x => x.Create(It.IsAny<IExecutionContext>(), It.IsAny<ActionStepDefinitionReference>(), It.IsAny<IStepHost>(), It.IsAny<ActionExecutionData>(), It.IsAny<Dictionary<string, string>>(), It.IsAny<Dictionary<string, string>>(), It.IsAny<Variables>(), It.IsAny<string>(), It.IsAny<List<JobExtensionRunner>>()))
+                           .Callback((IExecutionContext executionContext, Pipelines.ActionStepDefinitionReference actionReference, IStepHost stepHost, ActionExecutionData data, Dictionary<string, string> inputs, Dictionary<string, string> environment, Variables runtimeVariables, string taskDirectory, List<JobExtensionRunner> localActionContainerSetupSteps) =>
                            {
                                finialInputs = inputs;
                            })
@@ -107,8 +107,8 @@ namespace GitHub.Runner.Common.Tests.Worker
             _actionRunner.Action = action;
 
             Dictionary<string, string> finialInputs = new Dictionary<string, string>();
-            _handlerFactory.Setup(x => x.Create(It.IsAny<IExecutionContext>(), It.IsAny<ActionStepDefinitionReference>(), It.IsAny<IStepHost>(), It.IsAny<ActionExecutionData>(), It.IsAny<Dictionary<string, string>>(), It.IsAny<Dictionary<string, string>>(), It.IsAny<Variables>(), It.IsAny<string>()))
-                           .Callback((IExecutionContext executionContext, Pipelines.ActionStepDefinitionReference actionReference, IStepHost stepHost, ActionExecutionData data, Dictionary<string, string> inputs, Dictionary<string, string> environment, Variables runtimeVariables, string taskDirectory) =>
+            _handlerFactory.Setup(x => x.Create(It.IsAny<IExecutionContext>(), It.IsAny<ActionStepDefinitionReference>(), It.IsAny<IStepHost>(), It.IsAny<ActionExecutionData>(), It.IsAny<Dictionary<string, string>>(), It.IsAny<Dictionary<string, string>>(), It.IsAny<Variables>(), It.IsAny<string>(), It.IsAny<List<JobExtensionRunner>>()))
+                           .Callback((IExecutionContext executionContext, Pipelines.ActionStepDefinitionReference actionReference, IStepHost stepHost, ActionExecutionData data, Dictionary<string, string> inputs, Dictionary<string, string> environment, Variables runtimeVariables, string taskDirectory, List<JobExtensionRunner> localActionContainerSetupSteps) =>
                            {
                                finialInputs = inputs;
                            })
@@ -308,8 +308,8 @@ namespace GitHub.Runner.Common.Tests.Worker
             _actionRunner.Action = action;
 
             Dictionary<string, string> finialInputs = new Dictionary<string, string>();
-            _handlerFactory.Setup(x => x.Create(It.IsAny<IExecutionContext>(), It.IsAny<ActionStepDefinitionReference>(), It.IsAny<IStepHost>(), It.IsAny<ActionExecutionData>(), It.IsAny<Dictionary<string, string>>(), It.IsAny<Dictionary<string, string>>(), It.IsAny<Variables>(), It.IsAny<string>()))
-                           .Callback((IExecutionContext executionContext, Pipelines.ActionStepDefinitionReference actionReference, IStepHost stepHost, ActionExecutionData data, Dictionary<string, string> inputs, Dictionary<string, string> environment, Variables runtimeVariables, string taskDirectory) =>
+            _handlerFactory.Setup(x => x.Create(It.IsAny<IExecutionContext>(), It.IsAny<ActionStepDefinitionReference>(), It.IsAny<IStepHost>(), It.IsAny<ActionExecutionData>(), It.IsAny<Dictionary<string, string>>(), It.IsAny<Dictionary<string, string>>(), It.IsAny<Variables>(), It.IsAny<string>(), It.IsAny<List<JobExtensionRunner>>()))
+                           .Callback((IExecutionContext executionContext, Pipelines.ActionStepDefinitionReference actionReference, IStepHost stepHost, ActionExecutionData data, Dictionary<string, string> inputs, Dictionary<string, string> environment, Variables runtimeVariables, string taskDirectory, List<JobExtensionRunner> localActionContainerSetupSteps) =>
                            {
                                finialInputs = inputs;
                            })
@@ -359,8 +359,8 @@ namespace GitHub.Runner.Common.Tests.Worker
             _actionRunner.Action = action;
 
             Dictionary<string, string> finialInputs = new Dictionary<string, string>();
-            _handlerFactory.Setup(x => x.Create(It.IsAny<IExecutionContext>(), It.IsAny<ActionStepDefinitionReference>(), It.IsAny<IStepHost>(), It.IsAny<ActionExecutionData>(), It.IsAny<Dictionary<string, string>>(), It.IsAny<Dictionary<string, string>>(), It.IsAny<Variables>(), It.IsAny<string>()))
-                           .Callback((IExecutionContext executionContext, Pipelines.ActionStepDefinitionReference actionReference, IStepHost stepHost, ActionExecutionData data, Dictionary<string, string> inputs, Dictionary<string, string> environment, Variables runtimeVariables, string taskDirectory) =>
+            _handlerFactory.Setup(x => x.Create(It.IsAny<IExecutionContext>(), It.IsAny<ActionStepDefinitionReference>(), It.IsAny<IStepHost>(), It.IsAny<ActionExecutionData>(), It.IsAny<Dictionary<string, string>>(), It.IsAny<Dictionary<string, string>>(), It.IsAny<Variables>(), It.IsAny<string>(), It.IsAny<List<JobExtensionRunner>>()))
+                           .Callback((IExecutionContext executionContext, Pipelines.ActionStepDefinitionReference actionReference, IStepHost stepHost, ActionExecutionData data, Dictionary<string, string> inputs, Dictionary<string, string> environment, Variables runtimeVariables, string taskDirectory, List<JobExtensionRunner> localActionContainerSetupSteps) =>
                            {
                                finialInputs = inputs;
                            })

src/Test/L0/Worker/JobExtensionL0.cs

@@ -141,7 +141,7 @@ namespace GitHub.Runner.Common.Tests.Worker
                 var jobExtension = new JobExtension();
                 jobExtension.Initialize(hc);
 
-                _actionManager.Setup(x => x.PrepareActionsAsync(It.IsAny<IExecutionContext>(), It.IsAny<IEnumerable<Pipelines.JobStep>>()))
+                _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>())));
 
                 List<IStep> result = await jobExtension.InitializeJob(_jobEc, _message);
@@ -176,7 +176,7 @@ namespace GitHub.Runner.Common.Tests.Worker
                 var jobExtension = new JobExtension();
                 jobExtension.Initialize(hc);
 
-                _actionManager.Setup(x => x.PrepareActionsAsync(It.IsAny<IExecutionContext>(), It.IsAny<IEnumerable<Pipelines.JobStep>>()))
+                _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 JobExtensionRunner(null, "", "prepare1", null), new JobExtensionRunner(null, "", "prepare2", null) }, new Dictionary<Guid, IActionRunner>())));
 
                 List<IStep> result = await jobExtension.InitializeJob(_jobEc, _message);