diff --git a/bridged/src/main/java/dev/ltms/bridged/mcp/BridgeMcp.java b/bridged/src/main/java/dev/ltms/bridged/mcp/BridgeMcp.java index dde9242..0ca4cd2 100644 --- a/bridged/src/main/java/dev/ltms/bridged/mcp/BridgeMcp.java +++ b/bridged/src/main/java/dev/ltms/bridged/mcp/BridgeMcp.java @@ -513,6 +513,9 @@ public final class BridgeMcp { ? "[done — worker finished without a structured bridge_reply; transcript tail follows]\n" + v.reply() : v.reply()); case PENDING -> text("[pending — " + v.detail() + "]"); + case ASKING -> text("[question — worker is waiting for your answer]\n" + v.reply() + + "\n\nAnswer it by calling bridge_send again with turnId=\"" + v.turnId() + + "\" and content set to your answer; the worker resumes the same turn."); case FAILED -> text("[failed — " + v.detail() + "]"); }; } diff --git a/bridged/src/main/java/dev/ltms/bridged/msg/MessageService.java b/bridged/src/main/java/dev/ltms/bridged/msg/MessageService.java index e7758d4..9feff3f 100644 --- a/bridged/src/main/java/dev/ltms/bridged/msg/MessageService.java +++ b/bridged/src/main/java/dev/ltms/bridged/msg/MessageService.java @@ -127,6 +127,8 @@ public final class MessageService { public enum Phase { /** Delegated and in flight — queued for the worker or being worked. */ PENDING, + /** The worker is paused in {@code bridge_ask}; {@link TaskView#reply} and {@link TaskView#turnId} identify it. */ + ASKING, /** The worker's turn finished; {@link TaskView#reply} holds the answer. */ DONE, /** The delegation could not complete (timed out, worker gone, or busy). */ @@ -136,16 +138,27 @@ public final class MessageService { /** * A poll snapshot of an async delegation. * - * @param reply the answer when {@link #phase} is {@link Phase#DONE}, else {@code null} + * @param reply the answer when {@link #phase} is {@link Phase#DONE}, or the question when + * {@link #phase} is {@link Phase#ASKING}; otherwise {@code null} * @param replySource {@code "reply"} (structured {@code bridge_reply}) or {@code "transcript"} * (completion scrape) when {@link Phase#DONE}, else {@code null} - * @param detail a human note (live worker status while pending, or the failure reason) + * @param detail a human note (live worker status while pending, ask state, or failure reason) + * @param turnId correlation id for an {@link Phase#ASKING} ticket, else {@code null} */ - public record TaskView(String ticket, Phase phase, String reply, String replySource, String detail) { + public record TaskView(String ticket, Phase phase, String reply, String replySource, String detail, + String turnId) { } /** An in-flight or finished async delegation, keyed by its ticket. */ - private record Task(String target, CompletableFuture future, long createdNanos) { + private static final class Task { + private final String target; + private final CompletableFuture future = new CompletableFuture<>(); + private final long createdNanos = System.nanoTime(); + private volatile Reply question; + + private Task(String target) { + this.target = target; + } } private final AgentControl agents; @@ -156,6 +169,8 @@ public final class MessageService { private final Metrics metrics; // CB-502: nullable — no registry in unit tests private final ConcurrentHashMap sessionLocks = new ConcurrentHashMap<>(); private final ConcurrentHashMap tasks = new ConcurrentHashMap<>(); + /** One active async task per target, used to retain an ask after its initial send waiter resolves. */ + private final ConcurrentHashMap asyncTasksByTarget = new ConcurrentHashMap<>(); private final AtomicLong ticketSeq = new AtomicLong(); private final ExecutorService asyncExecutor = Executors.newThreadPerTaskExecutor( Thread.ofVirtual().name("bridge-async-", 0).factory()); @@ -218,6 +233,12 @@ public final class MessageService { count(BridgedMetrics.REPLIES, "path", "rendezvous"); return true; // a live send took it — unchanged fast path } + Task task = asyncTasksByTarget.get(session); + if (task != null) { + finishAsyncTask(task, new Reply(Outcome.REPLIED, content)); + count(BridgedMetrics.REPLIES, "path", "rendezvous"); + return true; + } inbox.publish(session, UUID.randomUUID().toString(), content); // A rising inbox share is the signal CB-307 exists to make visible: the worker finished but // nobody was waiting, so delivery now depends on the push loop and a drain. @@ -393,12 +414,14 @@ public final class MessageService { 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); return new AskResult(AskOutcome.ANSWERED, answer); } catch (TimeoutException e) { log.debug("bridge_ask from {} went unanswered in {}ms", workerSession, timeoutMillis); + clearAsyncQuestion(workerSession, ticket.turnId()); return new AskResult(AskOutcome.TIMED_OUT, null); } catch (ExecutionException e) { Throwable cause = e.getCause(); @@ -441,9 +464,12 @@ public final class MessageService { rendezvous.close(workerSession, reply); return new Reply(Outcome.STALE_TURN, null); // lapsed between the lookup and the unblock } + clearAsyncQuestion(workerSession, turnId); try { Rendezvous.Resolution r = reply.get(remainingMillis(deadlineNanos), TimeUnit.MILLISECONDS); - return new Reply(outcomeOf(r.kind()), r.text(), r.turnId()); + Reply result = new Reply(outcomeOf(r.kind()), r.text(), r.turnId()); + finishAsyncTask(workerSession, result); + return result; } catch (TimeoutException e) { // The worker resumed but hasn't replied yet — no completion fallback arms an answered // turn (it never re-entered the injector), so a silent worker rides out the window. @@ -483,9 +509,15 @@ public final class MessageService { */ public String sendAsync(String target, String content, Runnable onAccepted) { String ticket = "task-" + ticketSeq.incrementAndGet(); - CompletableFuture future = CompletableFuture.supplyAsync( - () -> send(target, content, ASYNC_TIMEOUT_MS, onAccepted), asyncExecutor); - tasks.put(ticket, new Task(target, future, System.nanoTime())); + Task task = new Task(target); + tasks.put(ticket, task); + asyncTasksByTarget.put(target, task); + asyncExecutor.submit(() -> { + Reply result = send(target, content, ASYNC_TIMEOUT_MS, onAccepted); + if (result.outcome() != Outcome.QUESTION) { + finishAsyncTask(task, result); + } + }); pruneTerminalTickets(); log.debug("async send {} -> {}", ticket, target); return ticket; @@ -501,27 +533,32 @@ public final class MessageService { if (task == null) { return null; } - CompletableFuture f = task.future(); + CompletableFuture f = task.future; if (!f.isDone()) { - return new TaskView(ticket, Phase.PENDING, null, null, "worker " + liveStatus(task.target())); + Reply question = task.question; + if (question != null) { + return new TaskView(ticket, Phase.ASKING, question.text(), null, + "worker is waiting for your answer", question.turnId()); + } + return new TaskView(ticket, Phase.PENDING, null, null, "worker " + liveStatus(task.target), null); } Reply r; try { r = f.getNow(null); } catch (CompletionException | java.util.concurrent.CancellationException e) { Throwable cause = (e instanceof CompletionException ce && ce.getCause() != null) ? ce.getCause() : e; - return new TaskView(ticket, Phase.FAILED, null, null, cause.getMessage()); + return new TaskView(ticket, Phase.FAILED, null, null, cause.getMessage(), null); } if (r.completed()) { String source = r.outcome() == Outcome.REPLIED ? "reply" : "transcript"; - return new TaskView(ticket, Phase.DONE, r.text(), source, null); + return new TaskView(ticket, Phase.DONE, r.text(), source, null, null); } // A wedged worker (CB-109) carries the error context as its reason; the timeout/busy // outcomes carry none, so fall back to the outcome name. String detail = r.outcome() == Outcome.WORKER_FAILED && r.text() != null ? r.text() : "no reply — " + r.outcome().name().toLowerCase(); - return new TaskView(ticket, Phase.FAILED, null, null, detail); + return new TaskView(ticket, Phase.FAILED, null, null, detail, null); } /** Best-effort live worker status for a pending poll; never throws (a lookup error is just noise). */ @@ -536,7 +573,37 @@ public final class MessageService { /** Drop finished tickets older than the TTL so the registry cannot grow without bound. */ private void pruneTerminalTickets() { long cutoff = System.nanoTime() - TICKET_TTL_NANOS; - tasks.values().removeIf(t -> t.future().isDone() && t.createdNanos() < cutoff); + tasks.values().removeIf(t -> t.future.isDone() && t.createdNanos < cutoff); + } + + /** 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) { + Task task = asyncTasksByTarget.get(target); + if (task != null) { + task.question = new Reply(Outcome.QUESTION, text, turnId); + } + } + + /** Clear an answered or lapsed question, but only when it matches the ticket's current turn. */ + private void clearAsyncQuestion(String target, String turnId) { + Task task = asyncTasksByTarget.get(target); + if (task != null && task.question != null && turnId.equals(task.question.turnId())) { + task.question = null; + } + } + + /** Complete and detach an async ticket after its worker's actual terminal reply. */ + private void finishAsyncTask(Task task, Reply result) { + task.future.complete(result); + asyncTasksByTarget.remove(task.target, task); + } + + /** Complete the active async ticket for a target, if this answer belongs to one. */ + private void finishAsyncTask(String target, Reply result) { + Task task = asyncTasksByTarget.get(target); + if (task != null) { + finishAsyncTask(task, result); + } } /** Release the async executor. */ diff --git a/bridged/src/test/java/dev/ltms/bridged/mcp/BridgeMcpTest.java b/bridged/src/test/java/dev/ltms/bridged/mcp/BridgeMcpTest.java index 0aed9d2..64c8cf9 100644 --- a/bridged/src/test/java/dev/ltms/bridged/mcp/BridgeMcpTest.java +++ b/bridged/src/test/java/dev/ltms/bridged/mcp/BridgeMcpTest.java @@ -132,6 +132,76 @@ class BridgeMcpTest { assertEquals("async LGTM", textOf(polled)); } + @Test + void asyncSendSurfacesAnAskThenKeepsTheTicketForTheFinalReply() throws Exception { + McpSchema.CallToolResult accepted = BridgeMcp.sendAsync(messages, "term_a", "do it", null, Set.of()); + String ticket = textOf(accepted).substring(textOf(accepted).indexOf("ticket=") + "ticket=".length()).trim(); + + long deadline = System.currentTimeMillis() + 3000; + while (!rendezvous.isWaiting("term_a") && System.currentTimeMillis() < deadline) { + Thread.sleep(5); + } + assertTrue(rendezvous.isWaiting("term_a")); + + CompletableFuture ask = CompletableFuture.supplyAsync( + () -> BridgeMcp.ask(messages, "term_a", "which config?", 5000L)); + McpSchema.CallToolResult question = BridgeMcp.poll(messages, ticket, null); + deadline = System.currentTimeMillis() + 3000; + while (!textOf(question).contains("[question") && System.currentTimeMillis() < deadline) { + Thread.sleep(5); + question = BridgeMcp.poll(messages, ticket, null); + } + assertTrue(textOf(question).contains("which config?"), textOf(question)); + String questionText = textOf(question); + String afterTurnId = questionText.substring(questionText.indexOf("turnId=\"") + "turnId=\"".length()); + String turnId = afterTurnId.substring(0, afterTurnId.indexOf('"')); + + CompletableFuture answer = CompletableFuture.supplyAsync( + () -> BridgeMcp.answer(messages, turnId, "config.yaml", 5000L)); + assertEquals("config.yaml", textOf(ask.get(6, TimeUnit.SECONDS))); + + deadline = System.currentTimeMillis() + 3000; + while (!rendezvous.isWaiting("term_a") && System.currentTimeMillis() < deadline) { + Thread.sleep(5); + } + assertTrue(rendezvous.isWaiting("term_a")); + BridgeMcp.reply(messages, "term_a", "done"); + assertEquals("done", textOf(answer.get(6, TimeUnit.SECONDS))); + + McpSchema.CallToolResult done = BridgeMcp.poll(messages, ticket, null); + deadline = System.currentTimeMillis() + 3000; + while (!"done".equals(textOf(done)) && System.currentTimeMillis() < deadline) { + Thread.sleep(5); + done = BridgeMcp.poll(messages, ticket, null); + } + assertEquals("done", textOf(done)); + } + + @Test + void unansweredAsyncAskReturnsTheTicketToPending() throws Exception { + McpSchema.CallToolResult accepted = BridgeMcp.sendAsync(messages, "term_a", "do it", null, Set.of()); + String ticket = textOf(accepted).substring(textOf(accepted).indexOf("ticket=") + "ticket=".length()).trim(); + + long deadline = System.currentTimeMillis() + 3000; + while (!rendezvous.isWaiting("term_a") && System.currentTimeMillis() < deadline) { + Thread.sleep(5); + } + assertTrue(rendezvous.isWaiting("term_a")); + + McpSchema.CallToolResult ask = BridgeMcp.ask(messages, "term_a", "still there?", 50L); + assertTrue(textOf(ask).contains("no answer"), textOf(ask)); + assertTrue(textOf(BridgeMcp.poll(messages, ticket, null)).startsWith("[pending")); + + BridgeMcp.reply(messages, "term_a", "finished after timeout"); + deadline = System.currentTimeMillis() + 3000; + McpSchema.CallToolResult done = BridgeMcp.poll(messages, ticket, null); + while (!"finished after timeout".equals(textOf(done)) && System.currentTimeMillis() < deadline) { + Thread.sleep(5); + done = BridgeMcp.poll(messages, ticket, null); + } + assertEquals("finished after timeout", textOf(done)); + } + @Test void pollUnknownTicketIsAnError() { McpSchema.CallToolResult res = BridgeMcp.poll(messages, "task-999", null);