fix(reply): refresh followup drain callbacks

This commit is contained in:
Vincent Koc
2026-03-23 09:45:21 -07:00
committed by Peter Steinberger
parent bcaadc39ea
commit a35dcf608e
4 changed files with 65 additions and 3 deletions

View File

@@ -215,13 +215,25 @@ export async function runReplyAgent(params: {
queueMode: resolvedQueue.mode,
});
const queuedRunFollowupTurn = createFollowupRunner({
opts,
typing,
typingMode,
sessionEntry: activeSessionEntry,
sessionStore: activeSessionStore,
sessionKey,
storePath,
defaultModel,
agentCfgContextTokens,
});
if (activeRunQueueAction === "drop") {
typing.cleanup();
return undefined;
}
if (activeRunQueueAction === "enqueue-followup") {
enqueueFollowupRun(queueKey, followupRun, resolvedQueue);
enqueueFollowupRun(queueKey, followupRun, resolvedQueue, "message-id", queuedRunFollowupTurn);
await touchActiveSessionEntry();
typing.cleanup();
return undefined;

View File

@@ -22,6 +22,13 @@ const FOLLOWUP_RUN_CALLBACKS = resolveGlobalMap<string, (run: FollowupRun) => Pr
FOLLOWUP_DRAIN_CALLBACKS_KEY,
);
export function rememberFollowupDrainCallback(
key: string,
runFollowup: (run: FollowupRun) => Promise<void>,
): void {
FOLLOWUP_RUN_CALLBACKS.set(key, runFollowup);
}
export function clearFollowupDrainCallback(key: string): void {
FOLLOWUP_RUN_CALLBACKS.delete(key);
}
@@ -78,7 +85,7 @@ export function scheduleFollowupDrain(
}
// Cache callback only when a drain actually starts. Avoid keeping stale
// callbacks around from finalize calls where no queue work is pending.
FOLLOWUP_RUN_CALLBACKS.set(key, runFollowup);
rememberFollowupDrainCallback(key, runFollowup);
void (async () => {
try {
const collectState = { forceIndividualCollect: false };

View File

@@ -1,6 +1,6 @@
import { resolveGlobalDedupeCache } from "../../../infra/dedupe.js";
import { applyQueueDropPolicy, shouldSkipQueueItem } from "../../../utils/queue-helpers.js";
import { kickFollowupDrainIfIdle } from "./drain.js";
import { kickFollowupDrainIfIdle, rememberFollowupDrainCallback } from "./drain.js";
import { getExistingFollowupQueue, getFollowupQueue } from "./state.js";
import type { FollowupRun, QueueDedupeMode, QueueSettings } from "./types.js";
@@ -59,6 +59,7 @@ export function enqueueFollowupRun(
run: FollowupRun,
settings: QueueSettings,
dedupeMode: QueueDedupeMode = "message-id",
runFollowup?: (run: FollowupRun) => Promise<void>,
): boolean {
const queue = getFollowupQueue(key, settings);
const recentMessageIdKey = dedupeMode !== "none" ? buildRecentMessageIdKey(run, key) : undefined;
@@ -92,6 +93,9 @@ export function enqueueFollowupRun(
if (recentMessageIdKey) {
RECENT_QUEUE_MESSAGE_IDS.check(recentMessageIdKey);
}
if (runFollowup) {
rememberFollowupDrainCallback(key, runFollowup);
}
// If drain finished and deleted the queue before this item arrived, a new queue
// object was created (draining: false) but nobody scheduled a drain for it.
// Use the cached callback to restart the drain now.

View File

@@ -1496,6 +1496,45 @@ describe("followup queue drain restart after idle window", () => {
expect(calls[1]?.prompt).toBe("after-idle");
});
it("restarts an idle drain with the newest followup callback", async () => {
const key = `test-idle-window-fresh-callback-${Date.now()}`;
const settings: QueueSettings = { mode: "followup", debounceMs: 0, cap: 50 };
const staleCalls: FollowupRun[] = [];
const freshCalls: FollowupRun[] = [];
const firstProcessed = createDeferred<void>();
const secondProcessed = createDeferred<void>();
const staleFollowup = async (run: FollowupRun) => {
staleCalls.push(run);
if (staleCalls.length === 1) {
firstProcessed.resolve();
}
};
const freshFollowup = async (run: FollowupRun) => {
freshCalls.push(run);
secondProcessed.resolve();
};
enqueueFollowupRun(key, createRun({ prompt: "before-idle" }), settings);
scheduleFollowupDrain(key, staleFollowup);
await firstProcessed.promise;
await new Promise<void>((resolve) => setImmediate(resolve));
enqueueFollowupRun(
key,
createRun({ prompt: "after-idle" }),
settings,
"message-id",
freshFollowup,
);
await secondProcessed.promise;
expect(staleCalls).toHaveLength(1);
expect(staleCalls[0]?.prompt).toBe("before-idle");
expect(freshCalls).toHaveLength(1);
expect(freshCalls[0]?.prompt).toBe("after-idle");
});
it("restarts an idle drain across distinct enqueue and drain module instances", async () => {
const drainA = await importFreshModule<typeof import("./queue/drain.js")>(
import.meta.url,