|
10 | 10 | import static org.mockito.Mockito.mock; |
11 | 11 | import static org.mockito.Mockito.when; |
12 | 12 |
|
| 13 | +import java.time.Duration; |
13 | 14 | import java.time.Instant; |
14 | 15 | import java.util.Arrays; |
15 | 16 | import java.util.List; |
|
24 | 25 | import java.util.concurrent.Future; |
25 | 26 | import java.util.concurrent.TimeUnit; |
26 | 27 | import java.util.concurrent.TimeoutException; |
| 28 | +import java.util.concurrent.atomic.AtomicBoolean; |
27 | 29 | import java.util.concurrent.atomic.AtomicReference; |
28 | 30 |
|
29 | 31 | import org.a2aproject.sdk.server.ServerCallContext; |
@@ -1750,6 +1752,104 @@ public void cancel(RequestContext context, AgentEmitter emitter) { |
1750 | 1752 | assertEquals(TaskState.TASK_STATE_FAILED, storedTask.status().state()); |
1751 | 1753 | } |
1752 | 1754 |
|
| 1755 | + /** |
| 1756 | + * A blocking call returns to the handler on the agent's first event, and the handler closes the |
| 1757 | + * caller's queue as soon as the agent returns. A final event the event bus distributes after that |
| 1758 | + * must still reach the queue, which was told to expect it; otherwise the consumer waits out its |
| 1759 | + * grace period for it and the call fails or returns late. |
| 1760 | + */ |
| 1761 | + @Test |
| 1762 | + void blockingSendStillReturnsCompletedTaskWhenFinalEventArrivesLate() throws Exception { |
| 1763 | + AtomicReference<EventQueue> callerQueue = new AtomicReference<>(); |
| 1764 | + useQueueManager(new InMemoryQueueManager((TaskStateProvider) taskStore, mainEventBus) { |
| 1765 | + @Override |
| 1766 | + public EventQueue createOrTap(String taskId) { |
| 1767 | + EventQueue queue = super.createOrTap(taskId); |
| 1768 | + callerQueue.compareAndSet(null, queue); |
| 1769 | + return queue; |
| 1770 | + } |
| 1771 | + }); |
| 1772 | + AtomicBoolean callerQueueClosedInTime = new AtomicBoolean(); |
| 1773 | + holdEventsAfterWorkingUntilClosed(callerQueue, callerQueueClosedInTime); |
| 1774 | + agentExecutorExecute = (context, emitter) -> { |
| 1775 | + emitter.startWork(); |
| 1776 | + emitter.complete(); |
| 1777 | + }; |
| 1778 | + |
| 1779 | + EventKind result = requestHandler.onMessageSend(blockingParams("msg-blocking-final-event"), NULL_CONTEXT); |
| 1780 | + |
| 1781 | + assertTrue(callerQueueClosedInTime.get(), "the caller's queue was not closed while the final event was held"); |
| 1782 | + Task task = assertInstanceOf(Task.class, result); |
| 1783 | + assertEquals(TaskState.TASK_STATE_COMPLETED, task.status().state()); |
| 1784 | + } |
| 1785 | + |
| 1786 | + private void useQueueManager(InMemoryQueueManager manager) { |
| 1787 | + EventQueueUtil.stop(mainEventBusProcessor); |
| 1788 | + queueManager = manager; |
| 1789 | + mainEventBusProcessor = new MainEventBusProcessor(mainEventBus, taskStore, |
| 1790 | + PushNotificationSender.NOOP, queueManager); |
| 1791 | + EventQueueUtil.start(mainEventBusProcessor); |
| 1792 | + requestHandler = DefaultRequestHandler.builder() |
| 1793 | + .agentExecutor(executor) |
| 1794 | + .taskStore(taskStore) |
| 1795 | + .queueManager(queueManager) |
| 1796 | + .pushConfigStore(pushConfigStore) |
| 1797 | + .pushNotificationsEnabled(true) |
| 1798 | + .mainEventBusProcessor(mainEventBusProcessor) |
| 1799 | + .executor(internalExecutor) |
| 1800 | + .eventConsumerExecutor(internalExecutor) |
| 1801 | + .authorizationRequired(false) |
| 1802 | + .build(); |
| 1803 | + } |
| 1804 | + |
| 1805 | + private void holdEventsAfterWorkingUntilClosed(AtomicReference<EventQueue> callerQueue, AtomicBoolean closedInTime) { |
| 1806 | + mainEventBusProcessor.setCallback(new MainEventBusProcessorCallback() { |
| 1807 | + @Override |
| 1808 | + public void onEventProcessed(String taskId, Event event) { |
| 1809 | + if (event instanceof TaskStatusUpdateEvent statusUpdate |
| 1810 | + && statusUpdate.status().state() == TaskState.TASK_STATE_WORKING) { |
| 1811 | + // Hold the agent's next event in the bus until the handler has closed the caller's queue |
| 1812 | + closedInTime.set(awaitClosed(callerQueue, Duration.ofSeconds(5))); |
| 1813 | + } |
| 1814 | + } |
| 1815 | + |
| 1816 | + @Override |
| 1817 | + public void onTaskFinalized(String taskId) { |
| 1818 | + } |
| 1819 | + }); |
| 1820 | + } |
| 1821 | + |
| 1822 | + private static boolean awaitClosed(AtomicReference<EventQueue> queue, Duration timeout) { |
| 1823 | + long deadline = System.nanoTime() + timeout.toNanos(); |
| 1824 | + while (System.nanoTime() < deadline) { |
| 1825 | + EventQueue current = queue.get(); |
| 1826 | + if (current != null && current.isClosed()) { |
| 1827 | + return true; |
| 1828 | + } |
| 1829 | + try { |
| 1830 | + Thread.sleep(10); |
| 1831 | + } catch (InterruptedException e) { |
| 1832 | + Thread.currentThread().interrupt(); |
| 1833 | + return false; |
| 1834 | + } |
| 1835 | + } |
| 1836 | + return false; |
| 1837 | + } |
| 1838 | + |
| 1839 | + private static MessageSendParams blockingParams(String messageId) { |
| 1840 | + return MessageSendParams.builder() |
| 1841 | + .message(Message.builder() |
| 1842 | + .messageId(messageId) |
| 1843 | + .role(Message.Role.ROLE_USER) |
| 1844 | + .parts(new TextPart("hello")) |
| 1845 | + .build()) |
| 1846 | + .configuration(MessageSendConfiguration.builder() |
| 1847 | + .returnImmediately(false) |
| 1848 | + .acceptedOutputModes(List.of()) |
| 1849 | + .build()) |
| 1850 | + .build(); |
| 1851 | + } |
| 1852 | + |
1753 | 1853 | @Test |
1754 | 1854 | void interruptedTerminalEnqueueStillReportsAgentFailure() throws Exception { |
1755 | 1855 | assertInterruptedTerminalEnqueueReportsError( |
|
0 commit comments