EventConsumer.java

package io.a2a.server.events;

import java.util.concurrent.Flow;

import io.a2a.spec.A2AServerException;
import io.a2a.spec.Event;
import io.a2a.spec.Message;
import io.a2a.spec.Task;
import io.a2a.spec.TaskStatusUpdateEvent;
import mutiny.zero.BackpressureStrategy;
import mutiny.zero.TubeConfiguration;
import mutiny.zero.ZeroPublisher;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class EventConsumer {
    private static final Logger LOGGER = LoggerFactory.getLogger(EventConsumer.class);
    private final EventQueue queue;
    private Throwable error;

    private static final String ERROR_MSG = "Agent did not return any response";
    private static final int NO_WAIT = -1;
    private static final int QUEUE_WAIT_MILLISECONDS = 500;

    public EventConsumer(EventQueue queue) {
        this.queue = queue;
        LOGGER.debug("EventConsumer created with queue {}", System.identityHashCode(queue));
    }

    public Event consumeOne() throws A2AServerException, EventQueueClosedException {
        Event event = queue.dequeueEvent(NO_WAIT);
        if (event == null) {
            throw new A2AServerException(ERROR_MSG, new InternalError(ERROR_MSG));
        }
        return event;
    }

    public Flow.Publisher<Event> consumeAll() {
        TubeConfiguration conf = new TubeConfiguration()
                .withBackpressureStrategy(BackpressureStrategy.BUFFER)
                .withBufferSize(256);
        return ZeroPublisher.create(conf, tube -> {
            boolean completed = false;
            try {
                while (true) {
                    if (error != null) {
                        completed = true;
                        tube.fail(error);
                        return;
                    }
                    // We use a timeout when waiting for an event from the queue.
                    // This is required because it allows the loop to check if
                    // `self._exception` has been set by the `agent_task_callback`.
                    // Without the timeout, loop might hang indefinitely if no events are
                    // enqueued by the agent and the agent simply threw an exception

                    // TODO the callback mentioned above seems unused in the Python 0.2.1 tag
                    Event event;
                    try {
                        event = queue.dequeueEvent(QUEUE_WAIT_MILLISECONDS);
                        if (event == null) {
                            continue;
                        }
                        if (event instanceof Throwable thr) {
                            tube.fail(thr);
                            return;
                        }
                        tube.send(event);
                    } catch (EventQueueClosedException e) {
                        completed = true;
                        tube.complete();
                        return;
                    } catch (Throwable t) {
                        tube.fail(t);
                        return;
                    }

                    boolean isFinalEvent = false;
                    if (event instanceof TaskStatusUpdateEvent tue && tue.isFinal()) {
                        isFinalEvent = true;
                    } else if (event instanceof Message) {
                        isFinalEvent = true;
                    } else if (event instanceof Task task) {
                        switch (task.getStatus().state()) {
                            case COMPLETED:
                            case CANCELED:
                            case FAILED:
                            case REJECTED:
                            case UNKNOWN:
                                isFinalEvent = true;
                        }
                    }

                    if (isFinalEvent) {
                        queue.close();
                        break;
                    }
                }
            } finally {
                if (!completed) {
                    tube.complete();
                }
            }
        });
    }

    public EnhancedRunnable.DoneCallback createAgentRunnableDoneCallback() {
        return agentRunnable -> {
            if (agentRunnable.getError() != null) {
                error = agentRunnable.getError();
            }
        };
    }

    public void close() {
        // Close the queue to stop the polling loop in consumeAll()
        // This will cause EventQueueClosedException and exit the while(true) loop
        LOGGER.debug("EventConsumer closing queue {}", System.identityHashCode(queue));
        queue.close();
    }
}