FutureCallback returns the whole CommandResponseMessage now.

core/src/main/java/org/axonframework/commandhandling/SimpleCommandBus.java

@@ -109,7 +109,7 @@ protected <C> CommandMessage<C> intercept(CommandMessage<C> command) {
      * @param <C>      The type of payload of the command
      * @param <R>      The type of result expected from the command handler
      */
-    protected <C, R> void doDispatch(CommandMessage<C> command, CommandCallback<? super C, R> callback) {
+    protected <C, R> void doDispatch(CommandMessage<C> command, CommandCallback<? super C, ? super R> callback) {
         MessageMonitor.MonitorCallback monitorCallback = messageMonitor.onMessageIngested(command);
 
         MessageHandler<? super CommandMessage<?>> handler = findCommandHandlerFor(command).orElseThrow(() -> {

core/src/main/java/org/axonframework/commandhandling/callbacks/FutureCallback.java

@@ -27,6 +27,7 @@
 import java.util.concurrent.TimeoutException;
 
 import static java.util.Objects.requireNonNull;
+import static org.axonframework.commandhandling.GenericCommandResponseMessage.asCommandResponseMessage;
 
 /**
  * Command Handler Callback that allows the dispatching thread to wait for the result of the callback, using the Future
@@ -37,12 +38,13 @@
  * @author Allard Buijze
  * @since 0.6
  */
-public class FutureCallback<C, R> extends CompletableFuture<R> implements CommandCallback<C, R> {
+public class FutureCallback<C, R> extends CompletableFuture<CommandResponseMessage<? extends R>>
+        implements CommandCallback<C, R> {
 
     @Override
     public void onSuccess(CommandMessage<? extends C> commandMessage,
                           CommandResponseMessage<? extends R> commandResponseMessage) {
-        super.complete(commandResponseMessage.getPayload());
+        super.complete(commandResponseMessage);
     }
 
     @Override
@@ -64,12 +66,12 @@ public void onFailure(CommandMessage commandMessage, Throwable cause) {
      *
      * @see #get()
      */
-    public R getResult() {
+    public CommandResponseMessage<? extends R> getResult() {
         try {
             return get();
         } catch (InterruptedException e) {
             Thread.currentThread().interrupt();
-            return null;
+            return asCommandResponseMessage(null);
         } catch (ExecutionException e) {
             throw asRuntime(e);
         }
@@ -90,14 +92,14 @@ public R getResult() {
      * @param unit    the time unit of the timeout argument
      * @return the result of the command handler execution.
      */
-    public R getResult(long timeout, TimeUnit unit) {
+    public CommandResponseMessage<? extends R> getResult(long timeout, TimeUnit unit) {
         try {
             return get(timeout, unit);
         } catch (InterruptedException e) {
             Thread.currentThread().interrupt();
-            return null;
+            return asCommandResponseMessage(null);
         } catch (TimeoutException e) {
-            return null;
+            return asCommandResponseMessage(null);
         } catch (ExecutionException e) {
             throw asRuntime(e);
         }

core/src/main/java/org/axonframework/commandhandling/gateway/AbstractCommandGateway.java

@@ -70,9 +70,9 @@ protected AbstractCommandGateway(CommandBus commandBus, RetryScheduler retrySche
      * @param callback The callback to notify with the processing result
      * @param <R>      The type of response expected from the command
      */
-    protected <C, R> void send(C command, CommandCallback<? super C, R> callback) {
+    protected <C, R> void send(C command, CommandCallback<? super C, ? super R> callback) {
         CommandMessage<? extends C> commandMessage = processInterceptors(asCommandMessage(command));
-        CommandCallback<? super C, R> commandCallback = callback;
+        CommandCallback<? super C, ? super R> commandCallback = callback;
         if (retryScheduler != null) {
             commandCallback = new RetryingCallback<>(callback, retryScheduler, commandBus);
         }

core/src/main/java/org/axonframework/commandhandling/gateway/CommandGateway.java

@@ -45,7 +45,7 @@ public interface CommandGateway {
      * @param callback The callback to notify when the command has been processed
      * @param <R>      The type of result expected from command execution
      */
-    <C, R> void send(C command, CommandCallback<? super C, R> callback);
+    <C, R> void send(C command, CommandCallback<? super C, ? super R> callback);
 
     /**
      * Sends the given {@code command} and wait for it to execute. The result of the execution is returned when

core/src/main/java/org/axonframework/commandhandling/gateway/CommandGatewayFactory.java

@@ -24,6 +24,7 @@
 import org.axonframework.commandhandling.callbacks.FutureCallback;
 import org.axonframework.common.Assert;
 import org.axonframework.common.annotation.AnnotationUtils;
+import org.axonframework.messaging.Message;
 import org.axonframework.messaging.MessageDispatchInterceptor;
 import org.axonframework.messaging.annotation.MetaDataValue;
 import org.axonframework.messaging.responsetypes.ResponseType;
@@ -459,7 +460,7 @@ public CompletableFuture<R> invoke(Object proxy, Method invokedMethod, Object[]
                 }
                 callbacks.addAll(commandCallbacks);
                 send(command, new CompositeCallback(callbacks));
-                return future;
+                return future.thenApply(Message::getPayload);
             } else {
                 sendAndForget(command);
                 return null;

core/src/main/java/org/axonframework/commandhandling/gateway/DefaultCommandGateway.java

@@ -21,6 +21,7 @@
 import org.axonframework.commandhandling.CommandMessage;
 import org.axonframework.commandhandling.callbacks.FailureLoggingCallback;
 import org.axonframework.commandhandling.callbacks.FutureCallback;
+import org.axonframework.messaging.Message;
 import org.axonframework.messaging.MessageDispatchInterceptor;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -91,7 +92,7 @@ public DefaultCommandGateway(CommandBus commandBus, RetryScheduler retrySchedule
     }
 
     @Override
-    public <C, R> void send(C command, CommandCallback<? super C, R> callback) {
+    public <C, R> void send(C command, CommandCallback<? super C, ? super R> callback) {
         super.send(command, callback);
     }
 
@@ -109,9 +110,9 @@ public <C, R> void send(C command, CommandCallback<? super C, R> callback) {
     @Override
     @SuppressWarnings("unchecked")
     public <R> R sendAndWait(Object command) {
-        FutureCallback<Object, Object> futureCallback = new FutureCallback<>();
+        FutureCallback<Object, R> futureCallback = new FutureCallback<>();
         send(command, futureCallback);
-        return (R) futureCallback.getResult();
+        return futureCallback.getResult().getPayload();
     }
 
     /**
@@ -132,15 +133,15 @@ public <R> R sendAndWait(Object command) {
     @Override
     @SuppressWarnings("unchecked")
     public <R> R sendAndWait(Object command, long timeout, TimeUnit unit) {
-        FutureCallback<Object, Object> futureCallback = new FutureCallback<>();
+        FutureCallback<Object, R> futureCallback = new FutureCallback<>();
         send(command, futureCallback);
-        return (R) futureCallback.getResult(timeout, unit);
+        return futureCallback.getResult(timeout, unit).getPayload();
     }
 
     @Override
     public <R> CompletableFuture<R> send(Object command) {
         FutureCallback<Object, R> callback = new FutureCallback<>();
         send(command, new FailureLoggingCallback<>(logger, callback));
-        return callback;
+        return callback.thenApply(Message::getPayload);
     }
 }

core/src/test/java/org/axonframework/commandhandling/CommandHandlerInterceptorTest.java

@@ -69,7 +69,8 @@ public void setUp() {
     @Test
     public void testInterceptor() {
         commandGateway.sendAndWait(asCommandMessage(new CreateMyAggregateCommand("id")));
-        String result = commandGateway.sendAndWait(asCommandMessage(new UpdateMyAggregateStateCommand("id", "state")));
+        String result = commandGateway
+                .sendAndWait(asCommandMessage(new UpdateMyAggregateStateCommand("id", "state")));
 
         ArgumentCaptor<EventMessage<?>> eventCaptor = ArgumentCaptor.forClass(EventMessage.class);
 

core/src/test/java/org/axonframework/commandhandling/callbacks/FutureCallbackTest.java

@@ -60,7 +60,7 @@ public void testOnSuccess() throws InterruptedException {
         assertTrue(t.isAlive());
         testSubject.onSuccess(COMMAND_MESSAGE, COMMAND_RESPONSE_MESSAGE);
         t.join(THREAD_JOIN_TIMEOUT);
-        assertEquals("Hello world", resultFromParallelThread);
+        assertEquals(COMMAND_RESPONSE_MESSAGE, resultFromParallelThread);
     }
 
     @SuppressWarnings({"ThrowableInstanceNeverThrown"})
@@ -112,6 +112,6 @@ public void testOnSuccessForLimitedTime_InTime() throws InterruptedException {
         testSubject.onSuccess(COMMAND_MESSAGE, COMMAND_RESPONSE_MESSAGE);
         assertTrue(testSubject.isDone());
         t.join(THREAD_JOIN_TIMEOUT);
-        assertEquals("Hello world", resultFromParallelThread);
+        assertEquals(COMMAND_RESPONSE_MESSAGE, resultFromParallelThread);
     }
 }

core/src/test/java/org/axonframework/config/DefaultConfigurerTest.java

@@ -84,7 +84,7 @@ public void defaultConfigurationWithEventSourcing() throws Exception {
 
         FutureCallback<Object, Object> callback = new FutureCallback<>();
         config.commandBus().dispatch(GenericCommandMessage.asCommandMessage("test"), callback);
-        assertEquals("test", callback.get());
+        assertEquals("test", callback.get().getPayload());
         assertNotNull(config.repository(StubAggregate.class));
         assertEquals(2, config.getModules().size());
         assertExpectedModules(config,
@@ -178,7 +178,7 @@ public void testJpaConfigurationWithInitialTransactionManagerJpaRepository() thr
         config.start();
         FutureCallback<Object, Object> callback = new FutureCallback<>();
         config.commandBus().dispatch(GenericCommandMessage.asCommandMessage("test"), callback);
-        assertEquals("test", callback.get());
+        assertEquals("test", callback.get().getPayload());
         assertNotNull(config.repository(StubAggregate.class));
         assertEquals(2, config.getModules().size());
         assertExpectedModules(config,
@@ -202,7 +202,7 @@ public void testJpaConfigurationWithInitialTransactionManagerJpaRepositoryFromCo
         config.start();
         FutureCallback<Object, Object> callback = new FutureCallback<>();
         config.commandBus().dispatch(GenericCommandMessage.asCommandMessage("test"), callback);
-        assertEquals("test", callback.get());
+        assertEquals("test", callback.get().getPayload());
         assertNotNull(config.repository(StubAggregate.class));
         assertTrue(config.getModules()
                          .stream()
@@ -249,7 +249,7 @@ public void testJpaConfigurationWithJpaRepository() throws Exception {
         config.start();
         FutureCallback<Object, Object> callback = new FutureCallback<>();
         config.commandBus().dispatch(GenericCommandMessage.asCommandMessage("test"), callback);
-        assertEquals("test", callback.get());
+        assertEquals("test", callback.get().getPayload());
         assertNotNull(config.repository(StubAggregate.class));
         assertEquals(2, config.getModules().size());
         assertExpectedModules(config,
@@ -273,7 +273,7 @@ public void defaultConfigurationWithMonitors() throws Exception {
 
         FutureCallback<Object, Object> callback = new FutureCallback<>();
         config.commandBus().dispatch(GenericCommandMessage.asCommandMessage("test"), callback);
-        assertEquals("test", callback.get());
+        assertEquals("test", callback.get().getPayload());
         assertEquals(1, defaultMonitor.getMessages().size());
         assertEquals(1, commandBusMonitor.getMessages().size());
     }

distributed-commandbus-jgroups/src/test/java/org/axonframework/jgroups/commandhandling/JGroupsConnectorTest.java

@@ -152,8 +152,8 @@ public void testConnectAndDispatchMessages_Balanced() throws Exception {
             }
             callbacks.add(callback);
         }
-        for (FutureCallback callback : callbacks) {
-            assertEquals("The Reply!", callback.get());
+        for (FutureCallback<?, ?> callback : callbacks) {
+            assertEquals("The Reply!", callback.get().getPayload());
         }
         assertEquals(100, counter1.get() + counter2.get());
         System.out.println("Node 1 got " + counter1.get());
@@ -200,10 +200,10 @@ public void testRingsProperlySynchronized_ChannelAlreadyConnected() throws Excep
         FutureCallback<Object, Object> callback4 = new FutureCallback<>();
         dcb2.dispatch(new GenericCommandMessage<>(1L), callback4);
 
-        assertEquals("The Reply!", callback1.get());
-        assertEquals("The Reply!", callback2.get());
-        assertEquals("The Reply!", callback3.get());
-        assertEquals("The Reply!", callback4.get());
+        assertEquals("The Reply!", callback1.get().getPayload());
+        assertEquals("The Reply!", callback2.get().getPayload());
+        assertEquals("The Reply!", callback3.get().getPayload());
+        assertEquals("The Reply!", callback4.get().getPayload());
 
         assertEquals(connector1.getConsistentHash(), connector2.getConsistentHash());
     }
@@ -311,8 +311,8 @@ public void testConnectAndDispatchMessages_SingleCandidate() throws Exception {
             }
             callbacks.add(callback);
         }
-        for (FutureCallback callback : callbacks) {
-            assertEquals("The Reply!", callback.get());
+        for (FutureCallback<?, ?> callback : callbacks) {
+            assertEquals("The Reply!", callback.get().getPayload());
         }
         assertEquals(100, counter1.get() + counter2.get());
         System.out.println("Node 1 got " + counter1.get());
@@ -357,7 +357,7 @@ public void testUnserializableResponseConvertedToNull() throws Exception {
 
         callback = new FutureCallback<>();
         dcb1.dispatch(new GenericCommandMessage<>("string"), callback);
-        assertNull(callback.getResult());
+        assertNull(callback.getResult().getPayload());
     }
 
     @SuppressWarnings("unchecked")
@@ -395,8 +395,8 @@ public void testConnectAndDispatchMessages_CustomCommandName() throws Exception
             }
             callbacks.add(callback);
         }
-        for (FutureCallback callback : callbacks) {
-            assertEquals("The Reply!", callback.get());
+        for (FutureCallback<?, ?> callback : callbacks) {
+            assertEquals("The Reply!", callback.get().getPayload());
         }
         assertEquals(100, counter1.get() + counter2.get());
         System.out.println("Node 1 got " + counter1.get());