DefaultRequestHandlerTest.java

package io.a2a.server.requesthandlers;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertTrue;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

import io.a2a.server.ServerCallContext;
import io.a2a.server.agentexecution.AgentExecutor;
import io.a2a.server.agentexecution.RequestContext;
import io.a2a.server.auth.UnauthenticatedUser;
import io.a2a.server.events.EventQueue;
import io.a2a.server.events.InMemoryQueueManager;
import io.a2a.server.tasks.InMemoryPushNotificationConfigStore;
import io.a2a.server.tasks.InMemoryTaskStore;
import io.a2a.server.tasks.TaskUpdater;
import io.a2a.spec.JSONRPCError;
import io.a2a.spec.Message;
import io.a2a.spec.MessageSendConfiguration;
import io.a2a.spec.MessageSendParams;
import io.a2a.spec.PushNotificationConfig;
import io.a2a.spec.Task;
import io.a2a.spec.TaskState;
import io.a2a.spec.TaskStatus;
import io.a2a.spec.TextPart;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.Timeout;

/**
 * Comprehensive tests for DefaultRequestHandler, backported from Python's
 * test_default_request_handler.py. These tests cover core functionality that
 * is transport-agnostic and should work across JSON-RPC, gRPC, and REST.
 *
 * Background cleanup and task tracking tests are from Python PR #440 and #472.
 */
public class DefaultRequestHandlerTest {

    private DefaultRequestHandler requestHandler;
    private InMemoryTaskStore taskStore;
    private InMemoryQueueManager queueManager;
    private TestAgentExecutor agentExecutor;
    private ServerCallContext serverCallContext;

    @BeforeEach
    void setUp() {
        taskStore = new InMemoryTaskStore();
        // Pass taskStore as TaskStateProvider to queueManager for task-aware queue management
        queueManager = new InMemoryQueueManager(taskStore);
        agentExecutor = new TestAgentExecutor();

        requestHandler = DefaultRequestHandler.create(
            agentExecutor,
            taskStore,
            queueManager,
            null, // pushConfigStore
            null, // pushSender
            Executors.newCachedThreadPool()
        );

        serverCallContext = new ServerCallContext(UnauthenticatedUser.INSTANCE, Map.of(), Set.of());
    }

    /**
     * Test that multiple blocking messages to the same task work correctly
     * when agent doesn't emit final events (fire-and-forget pattern).
     * This replicates TCK test: test_message_send_continue_task
     */
    @Test
    @Timeout(10)
    void testBlockingMessageContinueTask() throws Exception {
        String taskId = "continue-task-1";
        String contextId = "continue-ctx-1";

        // Configure agent to NOT complete tasks (like TCK fire-and-forget agent)
        agentExecutor.setExecuteCallback((context, queue) -> {
            Task task = context.getTask();
            if (task == null) {
                // First message: create SUBMITTED task
                task = new Task.Builder()
                    .id(context.getTaskId())
                    .contextId(context.getContextId())
                    .status(new TaskStatus(TaskState.SUBMITTED))
                    .build();
            } else {
                // Subsequent messages: emit WORKING task (non-final)
                task = new Task.Builder()
                    .id(context.getTaskId())
                    .contextId(context.getContextId())
                    .status(new TaskStatus(TaskState.WORKING))
                    .build();
            }
            queue.enqueueEvent(task);
            // Don't complete - just return (fire-and-forget)
        });

        // First blocking message - should return SUBMITTED task
        Message message1 = new Message.Builder()
            .messageId("msg-1")
            .role(Message.Role.USER)
            .parts(new TextPart("first message"))
            .taskId(taskId)
            .contextId(contextId)
            .build();

        MessageSendParams params1 = new MessageSendParams(message1, null, null);
        Object result1 = requestHandler.onMessageSend(params1, serverCallContext);

        assertTrue(result1 instanceof Task);
        Task task1 = (Task) result1;
        assertTrue(task1.getId().equals(taskId));
        assertTrue(task1.getStatus().state() == TaskState.SUBMITTED);

        // Second blocking message to SAME taskId - should not hang
        Message message2 = new Message.Builder()
            .messageId("msg-2")
            .role(Message.Role.USER)
            .parts(new TextPart("second message"))
            .taskId(taskId)
            .contextId(contextId)
            .build();

        MessageSendParams params2 = new MessageSendParams(message2, null, null);
        Object result2 = requestHandler.onMessageSend(params2, serverCallContext);

        // Should complete successfully (not timeout)
        assertTrue(result2 instanceof Task);
    }

    /**
     * Test that background cleanup tasks are properly tracked and cleared.
     * Backported from Python test: test_background_cleanup_task_is_tracked_and_cleared
     */
    @Test
    @Timeout(10)
    void testBackgroundCleanupTaskIsTrackedAndCleared() throws Exception {
        String taskId = "track-task-1";
        String contextId = "track-ctx-1";

        // Create a task that will trigger background cleanup
        Task task = new Task.Builder()
            .id(taskId)
            .contextId(contextId)
            .status(new TaskStatus(TaskState.SUBMITTED))
            .build();

        taskStore.save(task);

        Message message = new Message.Builder()
            .messageId("msg-track")
            .role(Message.Role.USER)
            .parts(new TextPart("test message"))
            .taskId(taskId)
            .contextId(contextId)
            .build();

        MessageSendParams params = new MessageSendParams(message, null, null);

        // Set up agent to finish quickly so cleanup runs
        CountDownLatch agentStarted = new CountDownLatch(1);
        CountDownLatch allowAgentFinish = new CountDownLatch(1);

        agentExecutor.setExecuteCallback((context, queue) -> {
            agentStarted.countDown();
            try {
                allowAgentFinish.await(5, TimeUnit.SECONDS);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
        });

        // Start streaming (this will create background tasks)
        var streamingResult = requestHandler.onMessageSendStream(params, serverCallContext);

        // Wait for agent to start
        assertTrue(agentStarted.await(5, TimeUnit.SECONDS), "Agent should start");

        // Allow agent to finish, which should trigger cleanup
        allowAgentFinish.countDown();

        // Give some time for background tasks to be tracked and cleaned up
        Thread.sleep(1000);

        // Background tasks should eventually be cleared
        // Note: We can't directly access the backgroundTasks field without reflection,
        // but the test verifies the mechanism doesn't hang or leak tasks
        assertTrue(true, "Background cleanup completed without hanging");
    }

    /**
     * Test that client disconnect triggers background cleanup and producer continues.
     * Backported from Python test: test_on_message_send_stream_client_disconnect_triggers_background_cleanup_and_producer_continues
     */
    @Test
    @Timeout(10)
    void testStreamingClientDisconnectTriggersBackgroundCleanup() throws Exception {
        String taskId = "disc-task-1";
        String contextId = "disc-ctx-1";

        Message message = new Message.Builder()
            .messageId("mid")
            .role(Message.Role.USER)
            .parts(new TextPart("test message"))
            .taskId(taskId)
            .contextId(contextId)
            .build();

        MessageSendParams params = new MessageSendParams(message, null, null);

        // Agent should start and then wait
        CountDownLatch agentStarted = new CountDownLatch(1);
        CountDownLatch allowAgentFinish = new CountDownLatch(1);
        AtomicBoolean agentCompleted = new AtomicBoolean(false);

        agentExecutor.setExecuteCallback((context, queue) -> {
            agentStarted.countDown();
            try {
                allowAgentFinish.await(10, TimeUnit.SECONDS);
                agentCompleted.set(true);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
        });

        // Start streaming
        var streamingResult = requestHandler.onMessageSendStream(params, serverCallContext);

        // Wait for agent to start
        assertTrue(agentStarted.await(5, TimeUnit.SECONDS), "Agent should start executing");

        // Simulate client disconnect by not consuming the stream
        // In real scenarios, the client would close the connection

        // Agent should still be running (not finished immediately on "disconnect")
        Thread.sleep(500);
        assertTrue(agentExecutor.isExecuting(), "Producer should still be running after simulated disconnect");

        // Allow agent to finish
        allowAgentFinish.countDown();

        // Wait a bit for completion
        Thread.sleep(1000);

        assertTrue(agentCompleted.get(), "Agent should have completed execution");
    }

    /**
     * Test that resubscription works after client disconnect.
     * Backported from Python test: test_stream_disconnect_then_resubscribe_receives_future_events
     */
    @Test
    @Timeout(15)
    void testStreamDisconnectThenResubscribeReceivesFutureEvents() throws Exception {
        String taskId = "reconn-task-1";
        String contextId = "reconn-ctx-1";

        // Create initial task
        Task initialTask = new Task.Builder()
            .id(taskId)
            .contextId(contextId)
            .status(new TaskStatus(TaskState.WORKING))
            .build();

        taskStore.save(initialTask);

        Message message = new Message.Builder()
            .messageId("msg-reconn")
            .role(Message.Role.USER)
            .parts(new TextPart("test message"))
            .taskId(taskId)
            .contextId(contextId)
            .build();

        MessageSendParams params = new MessageSendParams(message, null, null);

        // Set up agent to emit events with controlled timing
        CountDownLatch agentStarted = new CountDownLatch(1);
        CountDownLatch allowSecondEvent = new CountDownLatch(1);
        CountDownLatch allowFinish = new CountDownLatch(1);

        agentExecutor.setExecuteCallback((context, queue) -> {
            agentStarted.countDown();

            // Emit first event
            Task firstEvent = new Task.Builder()
                .id(taskId)
                .contextId(contextId)
                .status(new TaskStatus(TaskState.WORKING))
                .build();
            queue.enqueueEvent(firstEvent);

            // Wait for permission to emit second event
            try {
                allowSecondEvent.await(10, TimeUnit.SECONDS);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                return;
            }

            // Emit second event
            Task secondEvent = new Task.Builder()
                .id(taskId)
                .contextId(contextId)
                .status(new TaskStatus(TaskState.COMPLETED))
                .build();
            queue.enqueueEvent(secondEvent);

            try {
                allowFinish.await(10, TimeUnit.SECONDS);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
        });

        // Start streaming and simulate getting first event then disconnecting
        var streamingResult = requestHandler.onMessageSendStream(params, serverCallContext);

        // Wait for agent to start and emit first event
        assertTrue(agentStarted.await(5, TimeUnit.SECONDS), "Agent should start");

        // Simulate client disconnect (in real scenario, client would close connection)
        // The background cleanup should keep the producer running

        // Now try to resubscribe to the task
        io.a2a.spec.TaskIdParams resubParams = new io.a2a.spec.TaskIdParams(taskId);

        // Allow agent to emit second event
        allowSecondEvent.countDown();

        // Try resubscription - this should work because queue is still alive
        var resubResult = requestHandler.onResubscribeToTask(resubParams, serverCallContext);
        // If we get here without exception, resubscription worked
        assertTrue(true, "Resubscription succeeded");

        // Clean up
        allowFinish.countDown();
    }

    /**
     * Test that task state is persisted to task store after client disconnect.
     * Backported from Python test: test_disconnect_persists_final_task_to_store
     */
    @Test
    @Timeout(15)
    void testDisconnectPersistsFinalTaskToStore() throws Exception {
        String taskId = "persist-task-1";
        String contextId = "persist-ctx-1";

        Message message = new Message.Builder()
            .messageId("msg-persist")
            .role(Message.Role.USER)
            .parts(new TextPart("test message"))
            .taskId(taskId)
            .contextId(contextId)
            .build();

        MessageSendParams params = new MessageSendParams(message, null, null);

        // Agent that completes after some delay
        CountDownLatch agentStarted = new CountDownLatch(1);
        CountDownLatch allowCompletion = new CountDownLatch(1);

        agentExecutor.setExecuteCallback((context, queue) -> {
            agentStarted.countDown();

            // Emit working status
            Task workingTask = new Task.Builder()
                .id(taskId)
                .contextId(contextId)
                .status(new TaskStatus(TaskState.WORKING))
                .build();
            queue.enqueueEvent(workingTask);

            try {
                allowCompletion.await(10, TimeUnit.SECONDS);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                return;
            }

            // Emit final completed status
            Task completedTask = new Task.Builder()
                .id(taskId)
                .contextId(contextId)
                .status(new TaskStatus(TaskState.COMPLETED))
                .build();
            queue.enqueueEvent(completedTask);
        });

        // Start streaming and simulate client disconnect
        var streamingResult = requestHandler.onMessageSendStream(params, serverCallContext);

        // Wait for agent to start
        assertTrue(agentStarted.await(5, TimeUnit.SECONDS), "Agent should start");

        // Simulate client disconnect by not consuming the stream further
        // In real scenarios, the reactive stream would be cancelled

        // Allow agent to complete in background
        allowCompletion.countDown();

        // Give time for background processing to persist the final state
        Thread.sleep(2000);

        // Verify the final task state was persisted despite client disconnect
        Task persistedTask = taskStore.get(taskId);
        if (persistedTask != null) {
            // If task was persisted, it should have the final state
            assertTrue(
                persistedTask.getStatus().state() == TaskState.COMPLETED ||
                persistedTask.getStatus().state() == TaskState.WORKING,
                "Task should be persisted with working or completed state, got: " + persistedTask.getStatus().state()
            );
        }
        // Note: In some architectures, the task might not be persisted if the
        // background consumption isn't implemented. This test documents the expected behavior.
    }

    /**
     * Test that blocking message call waits for agent to finish and returns complete Task
     * even when agent does fire-and-forget (emits non-final state and returns).
     *
     * Expected behavior:
     * 1. Agent emits WORKING state with artifacts
     * 2. Agent's execute() method returns WITHOUT emitting final state
     * 3. Blocking onMessageSend() should wait for agent execution to complete
     * 4. Blocking onMessageSend() should wait for all queued events to be processed
     * 5. Returned Task should have WORKING state with all artifacts included
     *
     * This tests fire-and-forget pattern with blocking calls.
     */
    @Test
    @Timeout(15)
    void testBlockingFireAndForgetReturnsNonFinalTask() throws Exception {
        String taskId = "blocking-fire-forget-task";
        String contextId = "blocking-fire-forget-ctx";

        Message message = new Message.Builder()
            .messageId("msg-blocking-fire-forget")
            .role(Message.Role.USER)
            .parts(new TextPart("test message"))
            .taskId(taskId)
            .contextId(contextId)
            .build();

        MessageSendConfiguration config = new MessageSendConfiguration.Builder()
                .blocking(true)
                .build();

        MessageSendParams params = new MessageSendParams(message, config, null);

        // Agent that does fire-and-forget: emits WORKING with artifact but never completes
        agentExecutor.setExecuteCallback((context, queue) -> {
            TaskUpdater updater = new TaskUpdater(context, queue);

            // Start work (WORKING state)
            updater.startWork();

            // Add artifact
            updater.addArtifact(
                List.of(new TextPart("Fire and forget artifact", null)),
                "artifact-1", "FireForget", null);

            // Agent returns WITHOUT calling updater.complete()
            // Task stays in WORKING state (non-final)
        });

        // Call blocking onMessageSend - should wait for agent to finish
        Object result = requestHandler.onMessageSend(params, serverCallContext);

        // The returned result should be a Task in WORKING state with artifact
        assertTrue(result instanceof Task, "Result should be a Task");
        Task returnedTask = (Task) result;

        // Verify task is in WORKING state (non-final, fire-and-forget)
        assertEquals(TaskState.WORKING, returnedTask.getStatus().state(),
            "Returned task should be WORKING (fire-and-forget), got: " + returnedTask.getStatus().state());

        // Verify artifacts are included in the returned task
        assertNotNull(returnedTask.getArtifacts(),
            "Returned task should have artifacts");
        assertTrue(returnedTask.getArtifacts().size() >= 1,
            "Returned task should have at least 1 artifact, got: " +
            returnedTask.getArtifacts().size());
    }

    /**
     * Test that non-blocking message call returns immediately and persists all events in background.
     *
     * Expected behavior:
     * 1. Non-blocking call returns immediately with first event (WORKING state)
     * 2. Agent continues running in background and produces more events
     * 3. Background consumption continues and persists all events to TaskStore
     * 4. Final task state (COMPLETED) is persisted in background
     */
    @Test
    @Timeout(15)
    void testNonBlockingMessagePersistsAllEventsInBackground() throws Exception {
        String taskId = "blocking-persist-task";
        String contextId = "blocking-persist-ctx";

        Message message = new Message.Builder()
            .messageId("msg-nonblocking-persist")
            .role(Message.Role.USER)
            .parts(new TextPart("test message"))
            .taskId(taskId)
            .contextId(contextId)
            .build();

        // Set blocking=false for non-blocking behavior
        MessageSendConfiguration config = new MessageSendConfiguration.Builder()
                .blocking(false)
                .build();

        MessageSendParams params = new MessageSendParams(message, config, null);

        // Agent that produces multiple events with delays
        CountDownLatch agentStarted = new CountDownLatch(1);
        CountDownLatch firstEventEmitted = new CountDownLatch(1);
        CountDownLatch allowCompletion = new CountDownLatch(1);

        agentExecutor.setExecuteCallback((context, queue) -> {
            agentStarted.countDown();

            // Emit first event (WORKING state)
            Task workingTask = new Task.Builder()
                .id(taskId)
                .contextId(contextId)
                .status(new TaskStatus(TaskState.WORKING))
                .build();
            queue.enqueueEvent(workingTask);
            firstEventEmitted.countDown();

            // Sleep to ensure the non-blocking call has returned before we emit more events
            try {
                Thread.sleep(1000);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                return;
            }

            // Wait for permission to complete
            try {
                allowCompletion.await(10, TimeUnit.SECONDS);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                return;
            }

            // Emit final event (COMPLETED state)
            // This event should be persisted to TaskStore in background
            Task completedTask = new Task.Builder()
                .id(taskId)
                .contextId(contextId)
                .status(new TaskStatus(TaskState.COMPLETED))
                .build();
            queue.enqueueEvent(completedTask);
        });

        // Call non-blocking onMessageSend
        Object result = requestHandler.onMessageSend(params, serverCallContext);

        // Assertion 1: The immediate result should be the first event (WORKING)
        assertTrue(result instanceof Task, "Result should be a Task");
        Task immediateTask = (Task) result;
        assertEquals(TaskState.WORKING, immediateTask.getStatus().state(),
            "Non-blocking should return immediately with WORKING state, got: " + immediateTask.getStatus().state());

        // At this point, the non-blocking call has returned, but the agent is still running

        // Allow the agent to emit the final COMPLETED event
        allowCompletion.countDown();

        // Assertion 2: Poll for the final task state to be persisted in background
        // Use polling loop instead of fixed sleep for faster and more reliable test
        long timeoutMs = 5000;
        long startTime = System.currentTimeMillis();
        Task persistedTask = null;
        boolean completedStateFound = false;

        while (System.currentTimeMillis() - startTime < timeoutMs) {
            persistedTask = taskStore.get(taskId);
            if (persistedTask != null && persistedTask.getStatus().state() == TaskState.COMPLETED) {
                completedStateFound = true;
                break;
            }
            Thread.sleep(100);  // Poll every 100ms
        }

        assertTrue(persistedTask != null, "Task should be persisted to store");
        assertTrue(
            completedStateFound,
            "Final task state should be COMPLETED (background consumption should have processed it), got: " +
            (persistedTask != null ? persistedTask.getStatus().state() : "null") +
            " after " + (System.currentTimeMillis() - startTime) + "ms"
        );
    }

    /**
     * Test the BIG idea: MainQueue stays open for non-final tasks even when all children close.
     * This enables fire-and-forget tasks and late resubscriptions.
     */
    @Test
    @Timeout(15)
    void testMainQueueStaysOpenForNonFinalTasks() throws Exception {
        String taskId = "fire-and-forget-task";
        String contextId = "fire-ctx";

        // Create initial task in WORKING state (non-final)
        Task initialTask = new Task.Builder()
            .id(taskId)
            .contextId(contextId)
            .status(new TaskStatus(TaskState.WORKING))
            .build();
        taskStore.save(initialTask);

        Message message = new Message.Builder()
            .messageId("msg-fire")
            .role(Message.Role.USER)
            .parts(new TextPart("fire and forget"))
            .taskId(taskId)
            .contextId(contextId)
            .build();

        MessageSendParams params = new MessageSendParams(message, null, null);

        // Agent that emits WORKING status but never completes (fire-and-forget pattern)
        CountDownLatch agentStarted = new CountDownLatch(1);
        CountDownLatch allowAgentFinish = new CountDownLatch(1);

        agentExecutor.setExecuteCallback((context, queue) -> {
            agentStarted.countDown();

            // Emit WORKING status (non-final)
            Task workingTask = new Task.Builder()
                .id(taskId)
                .contextId(contextId)
                .status(new TaskStatus(TaskState.WORKING))
                .build();
            queue.enqueueEvent(workingTask);

            // Don't emit final state - just wait and finish
            try {
                allowAgentFinish.await(10, TimeUnit.SECONDS);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
            // Agent finishes WITHOUT emitting final task state
        });

        // Start streaming
        var streamingResult = requestHandler.onMessageSendStream(params, serverCallContext);

        // Wait for agent to start and emit WORKING event
        assertTrue(agentStarted.await(5, TimeUnit.SECONDS), "Agent should start");

        // Give time for WORKING event to be processed
        Thread.sleep(500);

        // Simulate client disconnect - this closes the ChildQueue
        // but MainQueue should stay open because task is non-final

        // Allow agent to finish
        allowAgentFinish.countDown();

        // Give time for agent to finish and cleanup to run
        Thread.sleep(2000);

        // THE BIG IDEA TEST: Resubscription should work because MainQueue is still open
        // Even though:
        // 1. The original ChildQueue closed (client disconnected)
        // 2. The agent finished executing
        // 3. Task is still in non-final WORKING state
        // Therefore: MainQueue should still be open for resubscriptions

        io.a2a.spec.TaskIdParams resubParams = new io.a2a.spec.TaskIdParams(taskId);
        var resubResult = requestHandler.onResubscribeToTask(resubParams, serverCallContext);

        // If we get here without exception, the BIG idea works!
        assertTrue(true, "Resubscription succeeded - MainQueue stayed open for non-final task");
    }

    /**
     * Test that MainQueue DOES close when task is finalized.
     * This ensures Level 2 protection doesn't prevent cleanup of completed tasks.
     */
    @Test
    @Timeout(15)
    void testMainQueueClosesForFinalizedTasks() throws Exception {
        String taskId = "completed-task";
        String contextId = "completed-ctx";

        // Create initial task in COMPLETED state (already finalized)
        Task completedTask = new Task.Builder()
            .id(taskId)
            .contextId(contextId)
            .status(new TaskStatus(TaskState.COMPLETED))
            .build();
        taskStore.save(completedTask);

        // Create a queue for this task
        EventQueue mainQueue = queueManager.createOrTap(taskId);
        assertTrue(mainQueue != null, "Queue should be created");

        // Close the child queue (simulating client disconnect)
        mainQueue.close();

        // Give time for cleanup callback to run
        Thread.sleep(1000);

        // Since the task is finalized (COMPLETED), the MainQueue should be removed from the map
        // This tests that Level 2 protection (childClosing check) allows cleanup for finalized tasks
        EventQueue queue = queueManager.get(taskId);
        assertTrue(queue == null || queue.isClosed(),
            "Queue for finalized task should be null or closed");
    }

    /**
     * Test that blocking message call returns a Task with ALL artifacts included.
     * This reproduces the reported bug: blocking call returns before artifacts are processed.
     *
     * Expected behavior:
     * 1. Agent emits multiple artifacts via TaskUpdater
     * 2. Blocking onMessageSend() should wait for ALL events to be processed
     * 3. Returned Task should have all artifacts included in COMPLETED state
     *
     * Bug manifestation:
     * - onMessageSend() returns after first event
     * - Artifacts are still being processed in background
     * - Returned Task is incomplete
     */
    @Test
    @Timeout(15)
    void testBlockingCallReturnsCompleteTaskWithArtifacts() throws Exception {
        String taskId = "blocking-artifacts-task";
        String contextId = "blocking-artifacts-ctx";

        Message message = new Message.Builder()
            .messageId("msg-blocking-artifacts")
            .role(Message.Role.USER)
            .parts(new TextPart("test message"))
            .taskId(taskId)
            .contextId(contextId)
            .build();

        MessageSendConfiguration config = new MessageSendConfiguration.Builder()
                .blocking(true)
                .build();

        MessageSendParams params = new MessageSendParams(message, config, null);

        // Agent that uses TaskUpdater to emit multiple artifacts (like real agents do)
        agentExecutor.setExecuteCallback((context, queue) -> {
            TaskUpdater updater = new TaskUpdater(context, queue);

            // Start work (WORKING state)
            updater.startWork();

            // Add first artifact
            updater.addArtifact(
                List.of(new TextPart("First artifact", null)),
                "artifact-1", "First", null);

            // Add second artifact
            updater.addArtifact(
                List.of(new TextPart("Second artifact", null)),
                "artifact-2", "Second", null);

            // Complete the task
            updater.complete();
        });

        // Call blocking onMessageSend - should wait for ALL events
        Object result = requestHandler.onMessageSend(params, serverCallContext);

        // The returned result should be a Task with ALL artifacts
        assertTrue(result instanceof Task, "Result should be a Task");
        Task returnedTask = (Task) result;

        // Verify task is completed
        assertEquals(TaskState.COMPLETED, returnedTask.getStatus().state(),
            "Returned task should be COMPLETED");

        // Verify artifacts are included in the returned task
        assertNotNull(returnedTask.getArtifacts(),
            "Returned task should have artifacts");
        assertTrue(returnedTask.getArtifacts().size() >= 2,
            "Returned task should have at least 2 artifacts, got: " +
            returnedTask.getArtifacts().size());
    }

    /**
     * Test that pushNotificationConfig from SendMessageConfiguration is stored for NEW tasks
     * in non-streaming (blocking) mode. This reproduces the bug from issue #84.
     *
     * Expected behavior:
     * 1. Client sends message with pushNotificationConfig in SendMessageConfiguration
     * 2. Agent creates a new task
     * 3. pushNotificationConfig should be stored in PushNotificationConfigStore
     * 4. Config should be retrievable via getInfo()
     */
    @Test
    @Timeout(10)
    void testBlockingMessageStoresPushNotificationConfigForNewTask() throws Exception {
        String taskId = "push-config-blocking-new-task";
        String contextId = "push-config-ctx";

        // Create test config store
        InMemoryPushNotificationConfigStore pushConfigStore = new InMemoryPushNotificationConfigStore();

        // Re-create request handler with pushConfigStore
        requestHandler = DefaultRequestHandler.create(
            agentExecutor,
            taskStore,
            queueManager,
            pushConfigStore, // Add push config store
            null, // pushSender
            Executors.newCachedThreadPool()
        );

        // Create push notification config
        PushNotificationConfig pushConfig = new PushNotificationConfig.Builder()
            .id("config-1")
            .url("https://example.com/webhook")
            .token("test-token-123")
            .build();

        // Create message with pushNotificationConfig
        Message message = new Message.Builder()
            .messageId("msg-push-config")
            .role(Message.Role.USER)
            .parts(new TextPart("test message"))
            .taskId(taskId)
            .contextId(contextId)
            .build();

        MessageSendConfiguration config = new MessageSendConfiguration.Builder()
            .blocking(true)
            .pushNotificationConfig(pushConfig)
            .build();

        MessageSendParams params = new MessageSendParams(message, config, null);

        // Agent creates a new task
        agentExecutor.setExecuteCallback((context, queue) -> {
            TaskUpdater updater = new TaskUpdater(context, queue);
            updater.submit();  // Creates new task in SUBMITTED state
            updater.complete();
        });

        // Call blocking onMessageSend
        Object result = requestHandler.onMessageSend(params, serverCallContext);

        // Verify result is a task
        assertTrue(result instanceof Task, "Result should be a Task");
        Task returnedTask = (Task) result;
        assertEquals(taskId, returnedTask.getId());

        // THE KEY ASSERTION: Verify pushNotificationConfig was stored
        List<PushNotificationConfig> storedConfigs = pushConfigStore.getInfo(taskId);
        assertNotNull(storedConfigs, "Push notification config should be stored for new task");
        assertEquals(1, storedConfigs.size(),
            "Should have exactly 1 push config stored");
        assertEquals("config-1", storedConfigs.get(0).id());
        assertEquals("https://example.com/webhook", storedConfigs.get(0).url());
    }

    /**
     * Test that pushNotificationConfig is stored for EXISTING tasks.
     * This verifies the initMessageSend logic works correctly.
     */
    @Test
    @Timeout(10)
    void testMessageStoresPushNotificationConfigForExistingTask() throws Exception {
        String taskId = "push-config-existing-task";
        String contextId = "push-config-existing-ctx";

        // Create test config store
        InMemoryPushNotificationConfigStore pushConfigStore = new InMemoryPushNotificationConfigStore();

        // Re-create request handler with pushConfigStore
        requestHandler = DefaultRequestHandler.create(
            agentExecutor,
            taskStore,
            queueManager,
            pushConfigStore, // Add push config store
            null, // pushSender
            Executors.newCachedThreadPool()
        );

        // Create EXISTING task in store
        Task existingTask = new Task.Builder()
            .id(taskId)
            .contextId(contextId)
            .status(new TaskStatus(TaskState.WORKING))
            .build();
        taskStore.save(existingTask);

        // Create push notification config
        PushNotificationConfig pushConfig = new PushNotificationConfig.Builder()
            .id("config-existing-1")
            .url("https://example.com/existing-webhook")
            .token("existing-token-789")
            .build();

        Message message = new Message.Builder()
            .messageId("msg-push-existing")
            .role(Message.Role.USER)
            .parts(new TextPart("update existing task"))
            .taskId(taskId)
            .contextId(contextId)
            .build();

        MessageSendConfiguration config = new MessageSendConfiguration.Builder()
            .blocking(true)
            .pushNotificationConfig(pushConfig)
            .build();

        MessageSendParams params = new MessageSendParams(message, config, null);

        // Agent updates the existing task
        agentExecutor.setExecuteCallback((context, queue) -> {
            TaskUpdater updater = new TaskUpdater(context, queue);
            updater.addArtifact(
                List.of(new TextPart("update artifact", null)),
                "artifact-1", "Update", null);
            updater.complete();
        });

        // Call blocking onMessageSend
        Object result = requestHandler.onMessageSend(params, serverCallContext);

        // Verify result
        assertTrue(result instanceof Task, "Result should be a Task");

        // Verify pushNotificationConfig was stored (initMessageSend path)
        List<PushNotificationConfig> storedConfigs = pushConfigStore.getInfo(taskId);
        assertNotNull(storedConfigs,
            "Push notification config should be stored for existing task");
        assertEquals(1, storedConfigs.size(),
            "Should have exactly 1 push config stored");
        assertEquals("config-existing-1", storedConfigs.get(0).id());
        assertEquals("https://example.com/existing-webhook", storedConfigs.get(0).url());
    }

    /**
     * Simple test agent executor that allows controlling execution timing
     */
    private static class TestAgentExecutor implements AgentExecutor {
        private ExecuteCallback executeCallback;
        private volatile boolean executing = false;

        interface ExecuteCallback {
            void call(RequestContext context, EventQueue queue) throws JSONRPCError;
        }

        void setExecuteCallback(ExecuteCallback callback) {
            this.executeCallback = callback;
        }

        boolean isExecuting() {
            return executing;
        }

        @Override
        public void execute(RequestContext context, EventQueue eventQueue) throws JSONRPCError {
            executing = true;
            try {
                if (executeCallback != null) {
                    // Custom callback is responsible for emitting events
                    executeCallback.call(context, eventQueue);
                } else {
                    // No custom callback - emit default completion event
                    Task completedTask = new Task.Builder()
                        .id(context.getTaskId())
                        .contextId(context.getContextId())
                        .status(new TaskStatus(TaskState.COMPLETED))
                        .build();
                    eventQueue.enqueueEvent(completedTask);
                }

            } finally {
                executing = false;
            }
        }

        @Override
        public void cancel(RequestContext context, EventQueue eventQueue) throws JSONRPCError {
            // Simple cancel implementation
            executing = false;
        }
    }
}