CB-577: model async question ownership race

bridged/src/main/java/dev/ltms/bridged/msg/MessageService.java

@@ -173,6 +173,9 @@ public final class MessageService {
     private final ConcurrentHashMap<String, Task> tasks = new ConcurrentHashMap<>();
     /** Async tasks that have accepted delivery for a target. */
     private final ConcurrentHashMap<String, Set<Task>> asyncTasksByTarget = new ConcurrentHashMap<>();
+    /** Async task that owns each exact forward rendezvous waiter. */
+    private final ConcurrentHashMap<CompletableFuture<Rendezvous.Resolution>, Task> asyncTasksByWaiter =
+            new ConcurrentHashMap<>();
     /** Async tickets paused on a specific {@code bridge_ask} turn. */
     private final ConcurrentHashMap<String, Task> asyncTasksByTurn = new ConcurrentHashMap<>();
     private final AtomicLong ticketSeq = new AtomicLong();
@@ -385,6 +388,9 @@ public final class MessageService {
             // failed send leaves no stale waiter behind.
             CompletableFuture<Rendezvous.Resolution> reply = rendezvous.open(target);
             try {
+                if (task != null) {
+                    asyncTasksByWaiter.put(reply, task);
+                }
                 // The send has won the lock; the accepted-delivery hook records delegator ownership
                 // here (CB-548). It runs BEFORE enqueue so a throwing hook — onAccepted is now a
                 // public callback — fails the send without queuing a message that would orphan.
@@ -408,6 +414,7 @@ public final class MessageService {
                     throw new IllegalStateException("interrupted awaiting reply from " + target, e);
                 }
             } finally {
+                asyncTasksByWaiter.remove(reply);
                 rendezvous.close(target, reply);
             }
         } finally {
@@ -433,11 +440,15 @@ public final class MessageService {
         if (ticket.fresh()) {
             // Register the reverse waiter first, then surface the question — so the answer, which can
             // arrive the instant the primary reacts, always finds an open waiter to resolve.
+            CompletableFuture<Rendezvous.Resolution> waiter = rendezvous.currentWaiter(workerSession);
+            Task task = markAsyncQuestion(waiter, question, ticket.turnId());
             if (!rendezvous.resolveQuestion(workerSession, question, ticket.turnId())) {
+                if (task != null) {
+                    clearAsyncQuestion(ticket.turnId(), true);
+                }
                 rendezvous.closeAsk(ticket.turnId());
                 return new AskResult(AskOutcome.NO_WAITER, null); // no primary is blocked on this worker
             }
-            markAsyncQuestion(workerSession, question, ticket.turnId());
         }
         try {
             String answer = ticket.answer().get(timeoutMillis, TimeUnit.MILLISECONDS);
@@ -613,17 +624,14 @@ public final class MessageService {
     }
 
     /** Record the active question for an async ticket; blocking sends have no entry and stay unchanged. */
-    private void markAsyncQuestion(String target, String text, String turnId) {
-        Set<Task> targetTasks = asyncTasksByTarget.get(target);
-        Task task = targetTasks == null ? null : targetTasks.stream()
-                .filter(candidate -> !candidate.future.isDone())
-                .findFirst()
-                .orElse(null);
+    private Task markAsyncQuestion(CompletableFuture<Rendezvous.Resolution> waiter, String text, String turnId) {
+        Task task = waiter == null ? null : asyncTasksByWaiter.get(waiter);
         if (task != null) {
             task.question = new Reply(Outcome.QUESTION, text, turnId);
             task.turnId = turnId;
             asyncTasksByTurn.put(turnId, task);
         }
+        return task;
     }
 
     /** Clear an answered or lapsed question, but only when it matches the ticket's current turn. */
@@ -634,6 +642,7 @@ public final class MessageService {
             if (forgetTurn) {
                 asyncTasksByTurn.remove(turnId, task);
                 task.turnId = null;
+                untrackAsyncTarget(task);
             }
         }
     }

bridged/src/test/java/dev/ltms/bridged/msg/MessageServiceTest.java

@@ -10,6 +10,9 @@ import dev.ltms.bridged.inject.Injector;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
+import java.lang.reflect.Field;
+import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 
@@ -676,6 +679,53 @@ class MessageServiceTest {
         assertEquals(MessageService.Outcome.REPLIED, answer.get(5, TimeUnit.SECONDS).outcome());
     }
 
+    @Test
+    void unansweredAsyncQuestionReturnsTheTicketToPendingAndReleasesItsTarget() throws Exception {
+        String ticket = messages.sendAsync(T, "task that asks");
+        awaitWaiting();
+        injectDelivery();
+
+        assertEquals(MessageService.AskOutcome.TIMED_OUT,
+                messages.ask(T, "which config?", 200).outcome());
+        assertEquals(MessageService.Phase.PENDING, messages.poll(ticket).phase(),
+                "only the question wait ended; the delegated turn may still finish");
+
+        String next = messages.sendAsync(T, "next task");
+        awaitWaiting();
+        assertTrue(rendezvous.resolve(T, "done"));
+        assertEquals(MessageService.Phase.DONE, awaitTicketPhase(next, MessageService.Phase.DONE).phase());
+    }
+
+    @Test
+    @SuppressWarnings("unchecked")
+    void asyncQuestionBelongsToTheTaskThatOwnsItsForwardWaiter() throws Exception {
+        String first = messages.sendAsync(T, "first task");
+        awaitWaiting();
+
+        // Model the resolveQuestion/markAsyncQuestion race: another accepted task reached the target set.
+        Field byTargetField = MessageService.class.getDeclaredField("asyncTasksByTarget");
+        byTargetField.setAccessible(true);
+        Map<String, Set<Object>> byTarget = (Map<String, Set<Object>>) byTargetField.get(messages);
+        Set<Object> targetTasks = byTarget.get(T);
+        Class<?> taskClass = Class.forName(MessageService.class.getName() + "$Task");
+        var constructor = taskClass.getDeclaredConstructor(String.class);
+        constructor.setAccessible(true);
+        targetTasks.clear();
+        targetTasks.add(constructor.newInstance(T));
+
+        CompletableFuture<MessageService.AskResult> ask =
+                CompletableFuture.supplyAsync(() -> messages.ask(T, "which config?", 5000));
+        assertEquals(MessageService.Phase.ASKING, awaitTicketPhase(first, MessageService.Phase.ASKING).phase());
+
+        MessageService.TaskView asking = messages.poll(first);
+        CompletableFuture<MessageService.Reply> answer = CompletableFuture.supplyAsync(
+                () -> messages.answer(asking.turnId(), "config.yaml", 5000));
+        assertEquals("config.yaml", ask.get(5, TimeUnit.SECONDS).answer());
+        awaitWaiting();
+        assertTrue(rendezvous.resolve(T, "done"));
+        assertEquals(MessageService.Outcome.REPLIED, answer.get(5, TimeUnit.SECONDS).outcome());
+    }
+
     private void assertFailedTicket(String ticket, String reason) throws Exception {
         MessageService.TaskView view = awaitTicketPhase(ticket, MessageService.Phase.FAILED);
         assertEquals(reason, view.detail());