chore: Ensure that the TaskStateProcessor passed to the ResultAggregator is never null

Signed-off-by: Jeff Mesnil <[email protected]>

Author: jmesnil

server-common/src/main/java/io/a2a/server/tasks/ResultAggregator.java

@@ -3,6 +3,7 @@
 import static io.a2a.server.util.async.AsyncUtils.consumer;
 import static io.a2a.server.util.async.AsyncUtils.createTubeConfig;
 import static io.a2a.server.util.async.AsyncUtils.processor;
+import static io.a2a.util.Assert.checkNotNullParam;
 
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
@@ -34,6 +35,9 @@ public class ResultAggregator {
     private volatile Message message;
 
     public ResultAggregator(TaskManager taskManager, Message message, Executor executor, TaskStateProcessor stateProcessor) {
+        checkNotNullParam("taskManager", taskManager);
+        checkNotNullParam("executor", executor);
+        checkNotNullParam("stateProcessor", stateProcessor);
         this.taskManager = taskManager;
         this.message = message;
         this.executor = executor;
@@ -218,10 +222,8 @@ public EventTypeAndInterrupt consumeAndBreakOnInterrupt(EventConsumer consumer,
                                     LOGGER.debug("Persisted final task state after background consumption: {}", taskIdForLogging());
                                 }
                                 // Remove task from state processor after final persistence
-                                if (stateProcessor != null) {
-                                    stateProcessor.removeTask(finalTask.getId());
-                                    LOGGER.debug("Removed task {} from state processor after background consumption", finalTask.getId());
-                                }
+                                stateProcessor.removeTask(finalTask.getId());
+                                LOGGER.debug("Removed task {} from state processor after background consumption", finalTask.getId());
                             }
                         }
                         completionFuture.complete(null);

server-common/src/test/java/io/a2a/server/tasks/ResultAggregatorTest.java

@@ -49,7 +49,7 @@ public class ResultAggregatorTest {
     @BeforeEach
     void setUp() {
         MockitoAnnotations.openMocks(this);
-        aggregator = new ResultAggregator(mockTaskManager, null, testExecutor, null);
+        aggregator = new ResultAggregator(mockTaskManager, null, testExecutor, new TaskStateProcessor());
     }
 
     // Helper methods for creating sample data
@@ -75,7 +75,7 @@ private Task createSampleTask(String taskId, TaskState statusState, String conte
     @Test
     void testConstructorWithMessage() {
         Message initialMessage = createSampleMessage("initial", "msg1", Message.Role.USER);
-        ResultAggregator aggregatorWithMessage = new ResultAggregator(mockTaskManager, initialMessage, testExecutor, null);
+        ResultAggregator aggregatorWithMessage = new ResultAggregator(mockTaskManager, initialMessage, testExecutor, new TaskStateProcessor());
 
         // Test that the message is properly stored by checking getCurrentResult
         assertEquals(initialMessage, aggregatorWithMessage.getCurrentResult());
@@ -86,7 +86,7 @@ void testConstructorWithMessage() {
     @Test
     void testGetCurrentResultWithMessageSet() {
         Message sampleMessage = createSampleMessage("hola", "msg1", Message.Role.USER);
-        ResultAggregator aggregatorWithMessage = new ResultAggregator(mockTaskManager, sampleMessage, testExecutor, null);
+        ResultAggregator aggregatorWithMessage = new ResultAggregator(mockTaskManager, sampleMessage, testExecutor, new TaskStateProcessor());
 
         EventKind result = aggregatorWithMessage.getCurrentResult();
 
@@ -121,7 +121,7 @@ void testConstructorStoresTaskManagerCorrectly() {
 
     @Test
     void testConstructorWithNullMessage() {
-        ResultAggregator aggregatorWithNullMessage = new ResultAggregator(mockTaskManager, null, testExecutor, null);
+        ResultAggregator aggregatorWithNullMessage = new ResultAggregator(mockTaskManager, null, testExecutor, new TaskStateProcessor());
         Task expectedTask = createSampleTask("null_msg_task", TaskState.WORKING, "ctx1");
         when(mockTaskManager.getTask()).thenReturn(expectedTask);
 
@@ -181,7 +181,7 @@ void testMultipleGetCurrentResultCalls() {
     void testGetCurrentResultWithMessageTakesPrecedence() {
         // Test that when both message and task are available, message takes precedence
         Message message = createSampleMessage("priority message", "pri1", Message.Role.USER);
-        ResultAggregator messageAggregator = new ResultAggregator(mockTaskManager, message, testExecutor, null);
+        ResultAggregator messageAggregator = new ResultAggregator(mockTaskManager, message, testExecutor, new TaskStateProcessor());
 
         // Even if we set up the task manager to return something, message should take precedence
         Task task = createSampleTask("should_not_be_returned", TaskState.WORKING, "ctx1");