diff --git a/temporal-test-server/src/main/java/io/temporal/internal/testservice/TestWorkflowMutableStateImpl.java b/temporal-test-server/src/main/java/io/temporal/internal/testservice/TestWorkflowMutableStateImpl.java
index 21d420c79..2db86eefb 100644
--- a/temporal-test-server/src/main/java/io/temporal/internal/testservice/TestWorkflowMutableStateImpl.java
+++ b/temporal-test-server/src/main/java/io/temporal/internal/testservice/TestWorkflowMutableStateImpl.java
@@ -553,28 +553,7 @@ public void completeWorkflowTask(
                     || request.getForceCreateNewWorkflowTask())) {
               scheduleWorkflowTask(ctx);
             }
-            if (completed) {
-              updates.forEach(
-                  (k, updateStateMachine) -> {
-                    if (!(updateStateMachine.getState() == StateMachines.State.COMPLETED
-                        || updateStateMachine.getState() == StateMachines.State.FAILED)) {
-                      updateStateMachine.action(
-                          Action.COMPLETE,
-                          ctx,
-                          Message.newBuilder()
-                              .setBody(
-                                  Any.pack(
-                                      Response.newBuilder()
-                                          .setOutcome(
-                                              Outcome.newBuilder()
-                                                  .setFailure(FAILED_UPDATE_ON_WF_COMPLETION)
-                                                  .build())
-                                          .build()))
-                              .build(),
-                          workflowTaskCompletedId);
-                    }
-                  });
-            }
+
             workflowTaskStateMachine.getData().bufferedEvents.clear();
             Map<String, ConsistentQuery> queries = data.consistentQueryRequests;
             Map<String, WorkflowQueryResult> queryResultsMap = request.getQueryResultsMap();
@@ -1705,6 +1684,27 @@ private void processWorkflowCompletionCallbacks(RequestContext ctx) {
       return;
     }
 
+    updates.forEach(
+        (k, updateStateMachine) -> {
+          if (!(updateStateMachine.getState() == StateMachines.State.COMPLETED
+              || updateStateMachine.getState() == StateMachines.State.FAILED)) {
+            updateStateMachine.action(
+                Action.COMPLETE,
+                ctx,
+                Message.newBuilder()
+                    .setBody(
+                        Any.pack(
+                            Response.newBuilder()
+                                .setOutcome(
+                                    Outcome.newBuilder()
+                                        .setFailure(FAILED_UPDATE_ON_WF_COMPLETION)
+                                        .build())
+                                .build()))
+                    .build(),
+                completionEvent.get().getEventId());
+          }
+        });
+
     for (Callback cb : startRequest.getCompletionCallbacksList()) {
       if (!cb.hasNexus()) {
         // test server only supports nexus callbacks currently