From 287442ca7edc79c744743ad5e25fbfb64611fca7 Mon Sep 17 00:00:00 2001
From: Sipke Schoorstra <sipkeschoorstra@outlook.com>
Date: Mon, 11 Sep 2023 21:58:27 +0200
Subject: [PATCH] Update middleware to also inspect memory block for output
 values

---
 .../Contexts/ActivityExecutionContext.cs      | 38 ++++++++++++++-----
 .../PersistActivityExecutionLogMiddleware.cs  | 15 +++++++-
 2 files changed, 43 insertions(+), 10 deletions(-)

diff --git a/src/modules/Elsa.Workflows.Core/Contexts/ActivityExecutionContext.cs b/src/modules/Elsa.Workflows.Core/Contexts/ActivityExecutionContext.cs
index 59c2ec0dd..1e2464b81 100644
--- a/src/modules/Elsa.Workflows.Core/Contexts/ActivityExecutionContext.cs
+++ b/src/modules/Elsa.Workflows.Core/Contexts/ActivityExecutionContext.cs
@@ -470,15 +470,9 @@ public class ActivityExecutionContext : IExecutionContext
     /// <exception cref="InvalidOperationException">The memory block does not exist.</exception>
     public object? Get(MemoryBlockReference blockReference)
     {
-        var memoryBlock = GetMemoryBlock(blockReference);
-
-        if (memoryBlock != null)
-            return memoryBlock.Value;
-
-        if (blockReference is Literal literal)
-            return literal.Value;
-
-        throw new InvalidOperationException($"The memory block '{blockReference}' does not exist.");
+        return !TryGet(blockReference, out var value)
+            ? throw new InvalidOperationException($"The memory block '{blockReference}' does not exist.")
+            : value;
     }
 
     /// <summary>
@@ -493,6 +487,32 @@ public class ActivityExecutionContext : IExecutionContext
         return value != default ? value.ConvertTo<T>() : default;
     }
 
+    /// <summary>
+    /// Tries to get the value of the specified memory block.
+    /// </summary>
+    /// <param name="blockReference">The memory block reference.</param>
+    /// <param name="value">The memory block value.</param>
+    /// <returns>True if the memory block exists, false otherwise.</returns>
+    public bool TryGet(MemoryBlockReference blockReference, out object? value)
+    {
+        var memoryBlock = GetMemoryBlock(blockReference);
+
+        if (memoryBlock != null)
+        {
+            value = memoryBlock.Value;
+            return true;
+        }
+
+        if (blockReference is Literal literal)
+        {
+            value = literal.Value;
+            return true;
+        }
+
+        value = null;
+        return false;
+    }
+
     /// <summary>
     /// Sets a value at the specified memory block.
     /// </summary>
diff --git a/src/modules/Elsa.Workflows.Runtime/Middleware/Workflows/PersistActivityExecutionLogMiddleware.cs b/src/modules/Elsa.Workflows.Runtime/Middleware/Workflows/PersistActivityExecutionLogMiddleware.cs
index 5e116c277..f3b3181a7 100644
--- a/src/modules/Elsa.Workflows.Runtime/Middleware/Workflows/PersistActivityExecutionLogMiddleware.cs
+++ b/src/modules/Elsa.Workflows.Runtime/Middleware/Workflows/PersistActivityExecutionLogMiddleware.cs
@@ -1,6 +1,7 @@
 using Elsa.Extensions;
 using Elsa.Mediator.Contracts;
 using Elsa.Workflows.Core;
+using Elsa.Workflows.Core.Models;
 using Elsa.Workflows.Core.Pipelines.WorkflowExecution;
 using Elsa.Workflows.Core.State;
 using Elsa.Workflows.Runtime.Contracts;
@@ -48,7 +49,19 @@ public class PersistActivityExecutionLogMiddleware : WorkflowExecutionMiddleware
             var expressionExecutionContext = activityExecutionContext.ExpressionExecutionContext;
             var activityDescriptor = activityExecutionContext.ActivityDescriptor;
             var outputDescriptors = activityDescriptor.Outputs;
-            var outputs = outputDescriptors.ToDictionary(x => x.Name, x => activity.GetOutput(expressionExecutionContext, x.Name)!);
+            
+            var outputs = outputDescriptors.ToDictionary(x => x.Name, x =>
+            {
+                var cachedValue = activity.GetOutput(expressionExecutionContext, x.Name);
+                
+                if(cachedValue != default)
+                    return cachedValue;
+                
+                if(x.ValueGetter(activity) is Output output && activityExecutionContext.TryGet(output.MemoryBlockReference(), out var outputValue))
+                    return outputValue;
+                
+                return default;
+            });
             
             return new ActivityExecutionRecord
             {