diff --git a/apps/server/src/services/threads/queue-drains.ts b/apps/server/src/services/threads/queue-drains.ts index 330fe86198..78f8ef9d9a 100644 --- a/apps/server/src/services/threads/queue-drains.ts +++ b/apps/server/src/services/threads/queue-drains.ts @@ -177,7 +177,7 @@ export async function runDueScheduledQueueSweep( now: number, ): Promise { 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 { 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 { - 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. diff --git a/apps/server/src/services/threads/queue-waits.ts b/apps/server/src/services/threads/queue-waits.ts index 7cc0c59024..847026639d 100644 --- a/apps/server/src/services/threads/queue-waits.ts +++ b/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, diff --git a/apps/server/src/services/threads/queued-messages.ts b/apps/server/src/services/threads/queued-messages.ts index 040a6bb45a..4131f9d69c 100644 --- a/apps/server/src/services/threads/queued-messages.ts +++ b/apps/server/src/services/threads/queued-messages.ts @@ -85,6 +85,7 @@ import { import { validatePromptAttachmentReferences } from "../projects/attachments.js"; interface SendQueuedMessageArgs { + isGroupEligible?: Parameters[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 { 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); diff --git a/apps/server/test/threads/requested-queue-drain.test.ts b/apps/server/test/threads/requested-queue-drain.test.ts index 660e20c953..a3eecc190d 100644 --- a/apps/server/test/threads/requested-queue-drain.test.ts +++ b/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 diff --git a/packages/db/src/data/queued-thread-messages.ts b/packages/db/src/data/queued-thread-messages.ts index 75b387898d..d0feb5d00a 100644 --- a/packages/db/src/data/queued-thread-messages.ts +++ b/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),