Return a closed SSE subscriber's undelivered events to the mailbox

When the server closed a subscriber (backpressure, a take-over GET, or a
write failure), the events still queued for it, and the one that tipped
the queue over, were dropped with it: an accepted prompt's result could
be lost although the stream is a mailbox. They now go back to the front
of the mailbox in order and reach the next subscriber. Bytes already
written to a connection that then died can still be lost; only event ids
and Last-Event-ID, which the RFD defers, could close that.

Author: Mark Pollack

acp-streamable-http-jetty/src/main/java/com/agentclientprotocol/sdk/agent/transport/SseOutboundStream.java

@@ -65,16 +65,42 @@ synchronized void push(String payload) {
 			return;
 		}
 		if (subscribers.isEmpty()) {
-			if (replay.size() == mailboxCapacity) {
-				throw new AcpConnectionException(
-						"Outbound SSE replay buffer exceeded " + mailboxCapacity + " events");
+			synchronized (replay) {
+				if (replay.size() == mailboxCapacity) {
+					throw new AcpConnectionException(
+							"Outbound SSE replay buffer exceeded " + mailboxCapacity + " events");
+				}
+				replay.addLast(payload);
 			}
-			replay.addLast(payload);
 			return;
 		}
 		subscribers.forEach(subscriber -> subscriber.send(payload));
 	}
 
+	/**
+	 * Puts events a closed subscriber never wrote back at the front of the mailbox, in
+	 * order, so the next subscriber receives them. Takes only the mailbox lock, so a
+	 * subscriber may call it while holding its own lock; a push that lands in the
+	 * mailbox meanwhile is newer and correctly stays behind them. Bytes already written
+	 * to a connection that then died can still be lost: without event ids and
+	 * {@code Last-Event-ID} (deferred by the RFD) nothing can know they did not arrive.
+	 */
+	private void requeue(List<String> unsent) {
+		if (unsent.isEmpty() || closed.get()) {
+			return;
+		}
+		synchronized (replay) {
+			if (replay.size() + unsent.size() > mailboxCapacity) {
+				logger.warn("Dropping {} undelivered SSE events: mailbox full ({} events)", unsent.size(),
+						mailboxCapacity);
+				return;
+			}
+			for (int i = unsent.size() - 1; i >= 0; i--) {
+				replay.addFirst(unsent.get(i));
+			}
+		}
+	}
+
 	synchronized void subscribe(AsyncContext asyncContext, HttpServletResponse response) throws IOException {
 		if (closed.get()) {
 			// DELETE may close the connection after GET has started async processing.
@@ -93,10 +119,14 @@ synchronized void subscribe(AsyncContext asyncContext, HttpServletResponse respo
 		SseSubscriber subscriber = new SseSubscriber(this, asyncContext, response);
 		subscribers.add(subscriber);
 		subscriber.start();
-		for (String payload : new ArrayList<>(replay)) {
+		List<String> retained;
+		synchronized (replay) {
+			retained = new ArrayList<>(replay);
+			replay.clear();
+		}
+		for (String payload : retained) {
 			subscriber.send(payload);
 		}
-		replay.clear();
 		subscriber.drain();
 	}
 
@@ -114,7 +144,9 @@ synchronized void close() {
 		if (closed.compareAndSet(false, true)) {
 			subscribers.forEach(SseSubscriber::close);
 			subscribers.clear();
-			replay.clear();
+			synchronized (replay) {
+				replay.clear();
+			}
 		}
 	}
 
@@ -127,7 +159,11 @@ private static final class SseSubscriber implements AsyncListener, WriteListener
 
 		private final ServletOutputStream output;
 
-		private final ArrayDeque<byte[]> pendingEvents = new ArrayDeque<>();
+		/** Queued writes; {@code payload} is null for comments (open, keep-alive), which are not requeued. */
+		private record Pending(byte[] bytes, String payload) {
+		}
+
+		private final ArrayDeque<Pending> pendingEvents = new ArrayDeque<>();
 
 		private final AtomicBoolean closed = new AtomicBoolean(false);
 
@@ -141,7 +177,7 @@ private static final class SseSubscriber implements AsyncListener, WriteListener
 
 		synchronized void start() {
 			asyncContext.addListener(this);
-			pendingEvents.addLast(SSE_OPEN_COMMENT);
+			pendingEvents.addLast(new Pending(SSE_OPEN_COMMENT, null));
 			output.setWriteListener(this);
 		}
 
@@ -152,30 +188,32 @@ synchronized void send(String payload) {
 			if (pendingEvents.size() == parent.maxPendingSseEvents) {
 				logger.warn("Closing backpressured SSE subscriber after {} pending events",
 						parent.maxPendingSseEvents);
+				// The event that did not fit goes back to the mailbox with the queue.
+				pendingEvents.addLast(new Pending(null, payload));
 				close();
 				return;
 			}
-			pendingEvents.addLast(("data: " + payload + "\n\n").getBytes(StandardCharsets.UTF_8));
+			pendingEvents.addLast(new Pending(("data: " + payload + "\n\n").getBytes(StandardCharsets.UTF_8), payload));
 			drain();
 		}
 
 		synchronized void sendKeepAlive() {
 			if (closed.get() || !pendingEvents.isEmpty()) {
 				return;
 			}
-			pendingEvents.addLast(SSE_KEEP_ALIVE_COMMENT);
+			pendingEvents.addLast(new Pending(SSE_KEEP_ALIVE_COMMENT, null));
 			drain();
 		}
 
 		synchronized void drain() {
 			try {
 				flushIfReady();
 				while (!closed.get() && output.isReady()) {
-					byte[] event = pendingEvents.pollFirst();
+					Pending event = pendingEvents.pollFirst();
 					if (event == null) {
 						break;
 					}
-					output.write(event);
+					output.write(event.bytes());
 					flushPending = true;
 				}
 				flushIfReady();
@@ -225,9 +263,17 @@ public void onStartAsync(AsyncEvent event) {
 		void close() {
 			if (closed.compareAndSet(false, true)) {
 				parent.remove(this);
+				List<String> unsent = new ArrayList<>();
 				synchronized (this) {
+					for (Pending event : pendingEvents) {
+						if (event.payload() != null) {
+							unsent.add(event.payload());
+						}
+					}
 					pendingEvents.clear();
 				}
+				// Undelivered events are not lost with the subscriber (mailbox guarantee).
+				parent.requeue(unsent);
 				try {
 					asyncContext.complete();
 				}

acp-streamable-http-jetty/src/test/java/com/agentclientprotocol/sdk/agent/transport/SseOutboundStreamTest.java

@@ -60,7 +60,7 @@ void mailboxOverflowThrows() {
 	}
 
 	@Test
-	void backpressuredSubscriberIsClosedAndLaterEventsGoToTheMailbox() throws IOException {
+	void backpressuredSubscriberIsClosedAndItsUndeliveredEventsReachTheNextSubscriber() throws IOException {
 		SseOutboundStream stream = new SseOutboundStream(4, 2);
 
 		// Never ready: the open comment and every event stay queued on the subscriber.
@@ -70,7 +70,7 @@ void backpressuredSubscriberIsClosedAndLaterEventsGoToTheMailbox() throws IOExce
 		verify(slow.asyncContext, never()).complete();
 
 		// The queue now holds maxPendingSseEvents (open comment + one event): the next
-		// event closes the subscriber and is discarded with its queue.
+		// event closes the subscriber; its unsent events go back to the mailbox.
 		stream.push("\"overflow\"");
 		verify(slow.asyncContext).complete();
 
@@ -79,7 +79,9 @@ void backpressuredSubscriberIsClosedAndLaterEventsGoToTheMailbox() throws IOExce
 		assertThat(slow.output.written()).isEmpty();
 
 		Attached next = attach(stream, true);
-		assertThat(next.output.written()).isEqualTo(OPEN + "data: \"later\"\n\n");
+		assertThat(next.output.written())
+			.as("nothing the slow subscriber never wrote is lost, and order is kept")
+			.isEqualTo(OPEN + "data: \"queued\"\n\n" + "data: \"overflow\"\n\n" + "data: \"later\"\n\n");
 		verify(next.asyncContext, never()).complete();
 	}