Keep grouped queued prompts atomic

apps/server/src/services/threads/queue-drains.ts

@@ -177,7 +177,7 @@ export async function runDueScheduledQueueSweep(
   now: number,
 ): Promise<void> {
   for (const row of listDueScheduledQueuedThreadMessages(deps.db, now)) {
-    await dispatchDueQueuedMessage(deps, row);
+    await dispatchDueQueuedMessage(deps, row, now);
   }
 }
 
@@ -190,6 +190,7 @@ interface QueuedMessageDispatchRef {
 async function dispatchDueQueuedMessage(
   deps: QueueDrainDeps,
   row: QueuedMessageDispatchRef,
+  now: number,
 ): Promise<void> {
   if (isDispatchRequeuedRecently(row.threadId)) {
     // This thread turned an attempt straight back into a queue moments ago.
@@ -205,7 +206,7 @@ async function dispatchDueQueuedMessage(
     // satisfied and every other wait is re-decided from scratch — including
     // the plugin pass, which is what makes a scheduled send still respect a
     // limiter at 9am rather than jumping it.
-    await clearDueWaitAndAttempt(deps, row);
+    await attemptEligibleQueuedMessage(deps, row, now);
   } catch (error) {
     // A background attempt has no caller left to report to, so the row carries
     // the outcome itself — as a `host-offline` wait when the machine is simply
@@ -224,15 +225,18 @@ async function dispatchDueQueuedMessage(
   }
 }
 
-async function clearDueWaitAndAttempt(
+async function attemptEligibleQueuedMessage(
   deps: QueueDrainDeps,
   row: QueuedMessageDispatchRef,
+  eligibility: number | "plugin",
 ): Promise<void> {
-  clearQueuedMessageWait(deps, {
-    queuedMessageId: row.id,
-    threadId: row.threadId,
-  });
   await sendQueuedMessage(deps, {
+    isGroupEligible: (group) =>
+      group.every((member) =>
+        eligibility === "plugin"
+          ? member.waitHolder !== null
+          : member.sendAt !== null && member.sendAt <= eligibility,
+      ),
     mode: "auto",
     queuedMessageId: row.id,
     threadId: row.threadId,
@@ -297,7 +301,7 @@ export async function runRequestedQueueDrain(
     const thread = getThread(deps.db, row.threadId);
     if (!thread || thread.deletedAt !== null) continue;
     try {
-      await clearDueWaitAndAttempt(deps, row);
+      await attemptEligibleQueuedMessage(deps, row, "plugin");
     } catch (error) {
       // Same posture as the due sweep: nobody is listening, so the outcome
       // lands on the row rather than propagating and stopping the walk.

apps/server/src/services/threads/queue-waits.ts

@@ -103,10 +103,6 @@ export function recordQueuedMessageWait(
       { behavior: "immediate" },
     );
   } else {
-    // Hand every claimed row back AND write the lead's wait in one
-    // transaction. Doing it in two would leave the group unclaimed and with no
-    // wait in between, which is exactly the window where the idle drain could
-    // pick up a message a hook has just said must wait.
     row = requeueClaimedQueuedThreadMessages(deps.db, deps.hub, {
       claims: claimed.map((claim) => ({
         id: claim.id,

apps/server/src/services/threads/queued-messages.ts

@@ -85,6 +85,7 @@ import {
 import { validatePromptAttachmentReferences } from "../projects/attachments.js";
 
 interface SendQueuedMessageArgs {
+  isGroupEligible?: Parameters<typeof claimQueuedThreadMessageGroup>[3];
   mode: SendQueuedMessageMode;
   queuedMessageId: string;
   /**
@@ -355,10 +356,12 @@ function claimQueuedThreadMessageForSend(
     deps.db,
     deps.hub,
     args.queuedMessageId,
+    args.isGroupEligible,
   );
   if (claimedQueuedMessages) {
     return claimedQueuedMessages;
   }
+  if (args.isGroupEligible) return [];
 
   const latestQueuedMessage = getQueuedThreadMessage(
     deps.db,
@@ -715,6 +718,12 @@ export async function sendQueuedMessage(
   args: SendQueuedMessageArgs,
 ): Promise<ThreadQueuedMessage> {
   const queuedMessages = claimQueuedThreadMessageForSend(deps, args);
+  if (queuedMessages.length === 0) {
+    const existing = getQueuedThreadMessage(deps.db, args.queuedMessageId);
+    if (!existing)
+      throw new ApiError(404, "invalid_request", "Queued message not found");
+    return toThreadQueuedMessage(existing);
+  }
   const thread = getThread(deps.db, args.threadId);
   if (thread && isManualCompactionActive(deps, thread)) {
     releaseQueuedMessageClaims(deps, queuedMessages);

apps/server/test/threads/requested-queue-drain.test.ts

@@ -1,4 +1,4 @@
-import { listEvents, listQueuedThreadMessages } from "@bb/db";
+import { listEvents, listQueuedThreadMessages, setQueuedThreadMessageGroupBoundary } from "@bb/db";
 import type { PluginHookName } from "@get-bb/plugin-sdk";
 import { afterEach, describe, expect, it, vi } from "vitest";
 import {
@@ -7,6 +7,7 @@ import {
 } from "../../src/services/plugins/plugin-hook-registry.js";
 import {
   requestQueueDrain,
+  runDueScheduledQueueSweep,
   runRequestedQueueDrain,
 } from "../../src/services/threads/queue-drains.js";
 import { acceptThreadSendRequest } from "../../src/services/threads/thread-send-request.js";
@@ -15,6 +16,7 @@ import {
   seedEnvironment,
   seedHostSession,
   seedProjectWithSource,
+  seedQueuedMessage,
   seedThread,
   seedThreadRuntimeState,
   seedTurnStarted,
@@ -40,6 +42,7 @@ function installHooks(registry: HookRegistry): void {
 
 afterEach(() => {
   setPluginHookProvider(undefined);
+  vi.useRealTimers();
 });
 
 function seedRunnableThread(
@@ -84,6 +87,28 @@ function turnRequests(harness: TestAppHarness, threadId: string) {
 }
 
 describe("the requested queue drain", () => {
+  it("does not dispatch a scheduled group tail while its lead is postponed", async () => {
+    await withTestHarness(async (harness) => {
+      vi.useFakeTimers();
+      let attempts = 0;
+      installHooks({ "message.dispatch": [{ pluginId: "limiter", handler: () => ++attempts === 1 ? ({ action: "wait", reason: "At capacity" } as const) : { action: "proceed" } }] });
+      const { thread } = seedRunnableThread(harness, {
+        hostId: "host-scheduled-group",
+        status: "idle",
+      });
+      const lead = seedQueuedMessage(harness.deps, { threadId: thread.id, content: textInput("lead"), waitingOn: { kind: "time" }, sendAt: Date.now() - 2_000 });
+      const tail = seedQueuedMessage(harness.deps, { threadId: thread.id, content: textInput("tail"), waitingOn: { kind: "time" }, sendAt: Date.now() - 1_000 });
+      setQueuedThreadMessageGroupBoundary({ db: harness.db, notifier: harness.deps.hub, threadId: thread.id, expectedGroupedPrefixQueuedMessageIds: [lead.id, tail.id], groupBoundaryQueuedMessageId: tail.id });
+
+      await runDueScheduledQueueSweep(harness.deps, Date.now());
+      vi.advanceTimersByTime(1_001);
+      await runDueScheduledQueueSweep(harness.deps, Date.now());
+
+      expect(attempts).toBe(1);
+      expect(listQueuedThreadMessages(harness.db, thread.id)).toHaveLength(2);
+    });
+  });
+
   it("re-attempts a plugin-queued row once the hook lets it through", async () => {
     // The release path that replaced a plugin releasing its own wait: core
     // re-attempts, the hook re-decides, and a row that is still blocked simply

packages/db/src/data/queued-thread-messages.ts

@@ -780,6 +780,7 @@ export function claimQueuedThreadMessageGroup(
   db: DbConnection,
   notifier: DbNotifier,
   id: string,
+  isGroupEligible?: (rows: readonly QueuedThreadMessageRow[]) => boolean,
 ): ClaimedQueuedThreadMessageRow[] | null {
   const claimedQueuedMessages = db.transaction(
     (tx) => {
@@ -788,12 +789,6 @@ export function claimQueuedThreadMessageGroup(
         return null;
       }
 
-      // A group is a batch that dispatches together, so claiming its head
-      // claims all of it — including members still carrying a wait, because
-      // an explicit claim of the head (send-now, a due schedule, a cleared
-      // plugin wait) is a claim on the batch. Claiming a member that is NOT
-      // its group's head takes that row alone and severs its edges: the user
-      // asked for that message, not for whatever happens to sit around it.
       const queuedMessages = listQueuedThreadMessages(tx, existing.threadId);
       const group =
         partitionQueuedMessageGroups(queuedMessages).find((rows) =>
@@ -802,7 +797,10 @@ export function claimQueuedThreadMessageGroup(
       if (group === null) {
         return null;
       }
-      if (group[0]?.id !== id) {
+      if (isGroupEligible && !isGroupEligible(group)) {
+        return null;
+      }
+      if (!isGroupEligible && group[0]?.id !== id) {
         const now = Date.now();
         clearPreviousQueuedMessageGroupEdgeInTransaction(tx, existing, now);
         clearQueuedMessageGroupEdgeInTransaction(tx, existing, now);
@@ -1074,20 +1072,6 @@ export interface RequeueClaimedQueuedThreadMessagesArgs {
   sendAt: number | null;
 }
 
-/**
- * Hands a claimed group back to the queue and writes the lead row's wait, in
- * ONE transaction.
- *
- * Two statements would leave the rows unclaimed and with no wait in between, which
- * is a window where the idle drain can pick up a message that a gate has just
- * said must wait. Doing both under one immediate transaction closes it: from
- * every other reader's view the group goes straight from "being dispatched" to
- * "waiting on this reason".
- *
- * Returns the queued lead row, or null when the claim no longer holds (the row
- * was deleted, or a stale-claim sweep already reclaimed it) — in which case the
- * caller has nothing left to queue.
- */
 export function requeueClaimedQueuedThreadMessages(
   db: DbConnection,
   notifier: DbNotifier,
@@ -1100,7 +1084,15 @@ export function requeueClaimedQueuedThreadMessages(
       const now = Date.now();
       for (const claim of args.claims) {
         tx.update(queuedThreadMessages)
-          .set({ claimedAt: null, claimToken: null, updatedAt: now })
+          .set({
+            claimedAt: null,
+            claimToken: null,
+            waitingOn: JSON.stringify(args.waitingOn),
+            waitHolder: waitHolderFor(args.waitingOn),
+            sendAt: args.sendAt,
+            failureReason: null,
+            updatedAt: now,
+          })
           .where(
             and(
               eq(queuedThreadMessages.id, claim.id),