From 13754b587904f703eceb3406674a48a69d6ed1a4 Mon Sep 17 00:00:00 2001
From: salaboy <Salaboy@gmail.com>
Date: Mon, 18 Aug 2025 09:45:54 +0100
Subject: [PATCH] add support for custom status

Signed-off-by: salaboy <Salaboy@gmail.com>
---
 .../java/io/dapr/workflows/WorkflowContext.java   |  8 ++++++++
 .../workflows/runtime/DefaultWorkflowContext.java | 15 ++++++++++++---
 .../workflows/DefaultWorkflowContextTest.java     | 14 ++++++++++++++
 3 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/sdk-workflows/src/main/java/io/dapr/workflows/WorkflowContext.java b/sdk-workflows/src/main/java/io/dapr/workflows/WorkflowContext.java
index 9994cb779..8608e9693 100644
--- a/sdk-workflows/src/main/java/io/dapr/workflows/WorkflowContext.java
+++ b/sdk-workflows/src/main/java/io/dapr/workflows/WorkflowContext.java
@@ -526,4 +526,12 @@ public interface WorkflowContext {
   default UUID newUuid() {
     throw new RuntimeException("No implementation found.");
   }
+
+  /**
+   * Set a custom status to a workflow execution.
+   *
+   * @param status to be set to the current execution
+   */
+  void setCustomStatus(Object status);
+  
 }
diff --git a/sdk-workflows/src/main/java/io/dapr/workflows/runtime/DefaultWorkflowContext.java b/sdk-workflows/src/main/java/io/dapr/workflows/runtime/DefaultWorkflowContext.java
index d11e1fe77..d12346ca0 100644
--- a/sdk-workflows/src/main/java/io/dapr/workflows/runtime/DefaultWorkflowContext.java
+++ b/sdk-workflows/src/main/java/io/dapr/workflows/runtime/DefaultWorkflowContext.java
@@ -59,7 +59,7 @@ public class DefaultWorkflowContext implements WorkflowContext {
    * @throws IllegalArgumentException if context or logger is null
    */
   public DefaultWorkflowContext(TaskOrchestrationContext context, Logger logger)
-      throws IllegalArgumentException {
+          throws IllegalArgumentException {
     if (context == null) {
       throw new IllegalArgumentException("Context cannot be null");
     }
@@ -114,7 +114,7 @@ public class DefaultWorkflowContext implements WorkflowContext {
    */
   @Override
   public <V> Task<V> waitForExternalEvent(String name, Duration timeout, Class<V> dataType)
-      throws TaskCanceledException {
+          throws TaskCanceledException {
     return this.innerContext.waitForExternalEvent(name, timeout, dataType);
   }
 
@@ -130,7 +130,7 @@ public class DefaultWorkflowContext implements WorkflowContext {
    * @param timeout the amount of time to wait before canceling the returned
    *                {@code Task}
    * @return a new {@link Task} that completes when the external event is received
-   *         or when {@code timeout} expires
+   *     or when {@code timeout} expires
    * @throws TaskCanceledException if the specified {@code timeout} value expires
    *                               before the event is received
    */
@@ -294,4 +294,13 @@ public class DefaultWorkflowContext implements WorkflowContext {
       return workflowTaskRetryHandler.handle(workflowRetryContext);
     };
   }
+
+  /**
+   * Set custom status to a workflow execution.
+   *
+   * @param status to set to the execution
+   */
+  public void setCustomStatus(Object status) {
+    innerContext.setCustomStatus(status);
+  }
 }
diff --git a/sdk-workflows/src/test/java/io/dapr/workflows/DefaultWorkflowContextTest.java b/sdk-workflows/src/test/java/io/dapr/workflows/DefaultWorkflowContextTest.java
index f0eaa9ee1..837bf4c7d 100644
--- a/sdk-workflows/src/test/java/io/dapr/workflows/DefaultWorkflowContextTest.java
+++ b/sdk-workflows/src/test/java/io/dapr/workflows/DefaultWorkflowContextTest.java
@@ -135,6 +135,11 @@ public class DefaultWorkflowContextTest {
       @Override
       public void continueAsNew(Object input, boolean preserveUnprocessedEvents) {
       }
+
+      @Override
+      public void setCustomStatus(Object status) {
+
+      }
     };
   }
 
@@ -403,6 +408,15 @@ public class DefaultWorkflowContextTest {
     verify(mockInnerContext, times(1)).callSubOrchestrator(expectedName, expectedInput, null, null, String.class);
   }
 
+  @Test
+  public void setCustomStatusWorkflow() {
+    String customStatus = "CustomStatus";
+
+    context.setCustomStatus(customStatus);
+    verify(mockInnerContext, times(1)).setCustomStatus(customStatus);
+
+  }
+
   @Test
   public void newUuidTest() {
     context.newUuid();