From 15cc12dba67c1f2e72a1317cc8f081e24147492a Mon Sep 17 00:00:00 2001 From: Dotta Date: Fri, 30 Oct 2026 17:06:23 -0500 Subject: [PATCH] fix(heartbeat): scope and batch dispatch wakeups by owner Co-Authored-By: Paperclip --- doc/plans/2026-10-30-run-dispatch-sleep.md | 2 +- .../heartbeat-run-dispatch-lifecycle.test.ts | 34 ++++++++++--- .../heartbeat-run-dispatch-work.test.ts | 50 ++++++++++++++++--- .../modules/run-dispatch/adapters/postgres.ts | 2 + .../modules/wake-queue/adapters/postgres.ts | 4 +- server/src/services/agent-configuration.ts | 2 +- server/src/services/agent-lifecycle.ts | 2 +- server/src/services/budgets.ts | 8 +-- .../chat-run-milestone-notifications.ts | 2 +- server/src/services/companies.ts | 4 +- .../execution-control-notifications.ts | 6 +-- server/src/services/heartbeat.ts | 4 +- .../src/services/heartbeat/dispatch-work.ts | 25 ++++++++-- server/src/services/heartbeat/queue.ts | 5 +- server/src/services/heartbeat/retries.ts | 4 +- .../src/services/heartbeat/run-lifecycle.ts | 4 +- .../issue-review-state-notifications.ts | 2 +- server/src/services/issues.ts | 14 +++--- .../native-provider-capacity-retry.ts | 2 +- server/src/services/recovery/service.ts | 4 +- .../services/run-dispatch-notifications.ts | 9 ++-- ...eduled-retry-cancellation-notifications.ts | 2 +- 22 files changed, 136 insertions(+), 55 deletions(-) diff --git a/doc/plans/2026-10-30-run-dispatch-sleep.md b/doc/plans/2026-10-30-run-dispatch-sleep.md index 82dd8692043..e5cd0a9e014 100644 --- a/doc/plans/2026-10-30-run-dispatch-sleep.md +++ b/doc/plans/2026-10-30-run-dispatch-sleep.md @@ -12,7 +12,7 @@ invalidate retries notify the retry worker in their transaction. Promotion uses the existing gate and compare-and-set, and its committed write wakes ordinary queued dispatch. A queued run behind occupied capacity waits for the owning run's lifecycle transition. Configuration and authority changes also wake it. -An idle queue has no recurring read or timer-driven dispatch attempt. +An idle queue has no recurring read or timer-driven dispatch attempt. Notifications batch for one second and carry the changed company to queued dispatch. Known retry deadlines still preempt batching. Cancellation-owned retries stay with their cancellation worker and do not retain a promotion deadline. The supported writer contract is service writes through the shared database connection scope, including caller-owned and nested transactions. Existing diff --git a/server/src/__tests__/heartbeat-run-dispatch-lifecycle.test.ts b/server/src/__tests__/heartbeat-run-dispatch-lifecycle.test.ts index 3d6566e543d..b6fbb7353de 100644 --- a/server/src/__tests__/heartbeat-run-dispatch-lifecycle.test.ts +++ b/server/src/__tests__/heartbeat-run-dispatch-lifecycle.test.ts @@ -1,3 +1,4 @@ +import { idleWorkSnapshot } from "../services/task-admission.js"; import { PgDialect } from "drizzle-orm/pg-core"; import { afterEach, describe, expect, it, vi } from "vitest"; import { installDatabaseWorkSignals } from "../../../packages/db/src/work-signals.js"; @@ -28,7 +29,7 @@ function fixture(input: { due?: number; suppressed?: boolean; draining?: boolean if (state.due !== null && state.due <= Date.now()) state.due = null; }); const next = vi.fn(async () => state.due); - const dispatch = vi.fn(async () => {}); + const dispatch = vi.fn(async (_scope?: { companyId: string }) => {}); const coordinator = createDeliveryWorkCoordinator({ owner: db, canRun: () => !state.draining && !state.standby, canReconcile: () => !state.standby, onError: vi.fn() }); const worker = registerRunDispatchWork(coordinator, { @@ -37,8 +38,8 @@ function fixture(input: { due?: number; suppressed?: boolean; draining?: boolean }, () => !state.draining && !state.standby); cleanup.push(() => coordinator.stop()); return { db, state, promote, next, dispatch, coordinator, worker, - notify: () => db.transaction(tx => notifyRunDispatch(tx)), - settle: () => vi.advanceTimersByTimeAsync(0) }; + notify: (scope?: { companyId: string }) => db.transaction(tx => notifyRunDispatch(tx, scope)), + settle: () => vi.advanceTimersByTimeAsync(1_000) }; } describe("run dispatch obligations", () => { @@ -67,20 +68,39 @@ describe("run dispatch obligations", () => { const before = f.next.mock.calls.length; await vi.advanceTimersByTimeAsync(20 * 60_000); expect(f.next).toHaveBeenCalledTimes(before); }); + it("batches repeated tenant writes and passes only their company to queued dispatch", async () => { + const f = fixture(); await f.worker.ready; + for (let index = 0; index < 100; index++) { + await f.notify({ companyId: "company-a" }); + await vi.advanceTimersByTimeAsync(10); + } + expect(f.promote.mock.calls.length).toBeLessThanOrEqual(2); + expect(f.dispatch.mock.calls.length).toBeLessThanOrEqual(2); + await f.settle(); + expect(f.dispatch).toHaveBeenLastCalledWith({ companyId: "company-a" }); + expect(f.dispatch.mock.calls.slice(1).every(([scope]) => scope?.companyId === "company-a")).toBe(true); + expect(f.coordinator.nextWakeAt()).toBeNull(); + const count = f.dispatch.mock.calls.length; + await vi.advanceTimersByTimeAsync(20 * 60_000); expect(f.dispatch).toHaveBeenCalledTimes(count); + }); it("remembers the first producer arriving during drain without owner SQL", async () => { const f = fixture(); await f.worker.ready; f.state.draining = true; await f.notify(); await f.settle(); await vi.advanceTimersByTimeAsync(90_000); expect(f.promote).toHaveBeenCalledOnce(); expect(f.next).toHaveBeenCalledOnce(); + expect(idleWorkSnapshot().active).toBe(0); f.state.draining = false; await vi.advanceTimersByTimeAsync(30_000); expect(f.promote).toHaveBeenCalledTimes(2); expect(f.dispatch).toHaveBeenCalledTimes(2); expect(f.coordinator.nextWakeAt()).toBeNull(); }); - it("retains failed work across drain and settles after recovery", async () => { - const f = fixture(); f.state.fail = true; await f.worker.ready; + it("releases a failed producer hold during drain and retains recovery for resume", async () => { + const f = fixture(); await f.worker.ready; + f.state.fail = true; await f.notify(); await f.settle(); + expect(f.promote).toHaveBeenCalledTimes(2); f.state.draining = true; await vi.advanceTimersByTimeAsync(90_000); - expect(f.promote).toHaveBeenCalledOnce(); expect(f.next).not.toHaveBeenCalled(); + expect(f.promote).toHaveBeenCalledTimes(2); expect(f.next).toHaveBeenCalledOnce(); + expect(idleWorkSnapshot().active).toBe(0); f.state.fail = false; f.state.draining = false; await vi.advanceTimersByTimeAsync(30_000); - expect(f.promote).toHaveBeenCalledTimes(2); expect(f.coordinator.nextWakeAt()).toBeNull(); + expect(f.promote).toHaveBeenCalledTimes(3); expect(f.coordinator.nextWakeAt()).toBeNull(); }); it("preserves a notification during a nonoverlapping pass", async () => { const f = fixture(); await f.worker.ready; diff --git a/server/src/__tests__/heartbeat-run-dispatch-work.test.ts b/server/src/__tests__/heartbeat-run-dispatch-work.test.ts index 4a616887479..bcc8ed23972 100644 --- a/server/src/__tests__/heartbeat-run-dispatch-work.test.ts +++ b/server/src/__tests__/heartbeat-run-dispatch-work.test.ts @@ -1,3 +1,4 @@ +import { idleWorkSnapshot } from "../services/task-admission.js"; import { randomUUID } from "node:crypto"; import { and, eq, sql } from "drizzle-orm"; import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; @@ -76,9 +77,11 @@ describe.skipIf(!support.supported)("run dispatch production wiring (controlled return db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, id)).then(rows => rows[0]); } async function completed(id: string) { - await vi.waitFor(async () => expect((await status(id))?.status).toBe("succeeded"), { timeout: 10_000 }); + await vi.waitFor(async () => expect(["running", "succeeded"]).toContain((await status(id))?.status), { timeout: 10_000 }); await heartbeat.drainActiveRunExecutions(); - expect((await status(id)).stdoutExcerpt).toContain("dispatch-proof"); + const completed = await status(id); + expect(completed.status).toBe("succeeded"); + expect(completed.stdoutExcerpt).toContain("dispatch-proof"); } it("keeps fresh empty, configured idle and retained terminal history quiet", async () => { await start(); await quiet(); @@ -87,10 +90,10 @@ describe.skipIf(!support.supported)("run dispatch production wiring (controlled }); it("ordinary retry creation wakes the worker, respects the exact deadline, dispatches and settles", async () => { const f = await fixture(); await start(); await quiet(); - const result = await heartbeat.scheduleBoundedRetry(f.runId, { delayMs: 2_000, random: () => 0 }); + const result = await heartbeat.scheduleBoundedRetry(f.runId, { delayMs: 4_000, random: () => 0 }); expect(result.outcome).toBe("scheduled"); if (result.outcome !== "scheduled") throw new Error("Expected retry creation"); - await vi.waitFor(() => expect(coordinator!.nextWakeAt()).toBe(result.dueAt.getTime())); + await vi.waitFor(() => expect(coordinator!.nextWakeAt()).toBe(result.dueAt.getTime()), { timeout: 5_000 }); expect((await status(result.run.id)).status).toBe("scheduled_retry"); await new Promise(resolve => setTimeout(resolve, 100)); expect(Date.now()).toBeLessThan(result.dueAt.getTime()); @@ -110,7 +113,7 @@ describe.skipIf(!support.supported)("run dispatch production wiring (controlled const f = await fixture(); await start(); const scheduled = await heartbeat.scheduleBoundedRetry(f.runId, { delayMs: 600_000 }); if (scheduled.outcome !== "scheduled") throw new Error("Expected retry creation"); - await vi.waitFor(() => expect(coordinator!.nextWakeAt()).toBe(scheduled.dueAt.getTime())); + await vi.waitFor(() => expect(coordinator!.nextWakeAt()).toBe(scheduled.dueAt.getTime()), { timeout: 5_000 }); await heartbeat.cancelRun(scheduled.run.id, "operator stop", { suppressImmediateRecovery: true }); await quiet(); expect((await status(scheduled.run.id)).status).toBe("cancelled"); }); @@ -118,7 +121,14 @@ describe.skipIf(!support.supported)("run dispatch production wiring (controlled const f = await fixture(); await start(); await quiet(); draining = true; const scheduled = await heartbeat.scheduleBoundedRetry(f.runId, { delayMs: 200 }); if (scheduled.outcome !== "scheduled") throw new Error("Expected retry creation"); - await new Promise(resolve => setTimeout(resolve, 400)); + await vi.waitFor(() => expect(idleWorkSnapshot().active).toBe(0)); + const count = reads.mock.calls.length; + vi.useFakeTimers({ toFake: ["Date", "setTimeout", "clearTimeout"] }); + try { + await vi.advanceTimersByTimeAsync(90_000); + expect(reads).toHaveBeenCalledTimes(count); + expect(idleWorkSnapshot().active).toBe(0); + } finally { vi.useRealTimers(); } expect((await status(scheduled.run.id)).status).toBe("scheduled_retry"); draining = false; worker!.wake(); await completed(scheduled.run.id); await quiet(); @@ -151,13 +161,39 @@ describe.skipIf(!support.supported)("run dispatch production wiring (controlled await blocker.select().from(agents).where(eq(agents.id, f.agentId)).for("update"); writing = agentService(db).clearError(f.agentId); void writing.catch(() => {}); - await vi.waitFor(() => expect(databaseWorkPending(db, DELIVERY_QUEUES.queuedRunDispatch)).toBe(true)); + await new Promise(resolve => setTimeout(resolve, 100)); + expect(databaseWorkPending(db, DELIVERY_QUEUES.queuedRunDispatch)).toBe(false); expect((await status(queued.id)).status).toBe("queued"); expect((await db.select().from(agents).where(eq(agents.id, f.agentId)))[0].status).toBe("error"); }); } finally { await writing; } await completed(queued.id); await quiet(); }); + it("leaves cancellation-owned retries quiet even when their old promotion time is overdue", async () => { + const f = await fixture(); + const scheduled = await heartbeat.scheduleBoundedRetry(f.runId, { now: new Date(Date.now() - 60_000), delayMs: 1_000 }); + if (scheduled.outcome !== "scheduled") throw new Error("Expected retry creation"); + await db.update(heartbeatRuns).set({ resultJson: { startupCancellation: { requestedAt: new Date().toISOString() }, + cancellation: { expected: true, reason: "operator stop", source: "operator" } } }).where(eq(heartbeatRuns.id, scheduled.run.id)); + await start(); await quiet(); + expect((await status(scheduled.run.id)).status).toBe("scheduled_retry"); + expect(await heartbeat.promoteDueScheduledRetries()).toEqual({ promoted: 0, runIds: [] }); + await quiet(); + }); + it("rejects an error clear without publishing work when the lifecycle is no longer ready", async () => { + const f = await fixture(); await start(); await quiet(); + await db.update(agents).set({ status: "error", lifecycleState: "paused" }).where(eq(agents.id, f.agentId)); + vi.useFakeTimers({ toFake: ["Date", "setTimeout", "clearTimeout"] }); + try { + await expect(agentService(db).clearError(f.agentId)).rejects.toThrow("Only agents in error status"); + expect(coordinator!.nextWakeAt()).toBeNull(); + const count = reads.mock.calls.length; + await vi.advanceTimersByTimeAsync(60_000); + expect(reads).toHaveBeenCalledTimes(count); + } finally { vi.useRealTimers(); } + expect((await db.select().from(agents).where(eq(agents.id, f.agentId)))[0].status).toBe("error"); + await quiet(); + }); it("keeps a foreign company's future retry intact while concurrent workers claim one due retry", async () => { const selected = await fixture(), foreign = await fixture(); await start(); const later = await heartbeat.scheduleBoundedRetry(foreign.runId, { delayMs: 600_000 }); diff --git a/server/src/modules/run-dispatch/adapters/postgres.ts b/server/src/modules/run-dispatch/adapters/postgres.ts index 63559e7701d..b05f94e1a4b 100644 --- a/server/src/modules/run-dispatch/adapters/postgres.ts +++ b/server/src/modules/run-dispatch/adapters/postgres.ts @@ -514,6 +514,7 @@ export function createPostgresRunDispatchAdapter( async function nextScheduledRetryAt(input: { cutoff: Date | null }): Promise { const [row] = await db.select({ at: heartbeatRuns.scheduledRetryAt }).from(heartbeatRuns) .where(and(eq(heartbeatRuns.status, "scheduled_retry"), + sql`jsonb_typeof(${heartbeatRuns.resultJson}->'startupCancellation') is distinct from 'object'`, input.cutoff ? gte(heartbeatRuns.createdAt, input.cutoff) : undefined)) .orderBy(asc(heartbeatRuns.scheduledRetryAt)).limit(1); return row?.at?.getTime() ?? null; @@ -526,6 +527,7 @@ export function createPostgresRunDispatchAdapter( .where( and( eq(heartbeatRuns.status, "scheduled_retry"), + sql`jsonb_typeof(${heartbeatRuns.resultJson}->'startupCancellation') is distinct from 'object'`, lte(heartbeatRuns.scheduledRetryAt, input.now), input.cutoff ? gte(heartbeatRuns.createdAt, input.cutoff) : undefined, ), diff --git a/server/src/modules/wake-queue/adapters/postgres.ts b/server/src/modules/wake-queue/adapters/postgres.ts index c773b903dec..1df4420615f 100644 --- a/server/src/modules/wake-queue/adapters/postgres.ts +++ b/server/src/modules/wake-queue/adapters/postgres.ts @@ -638,7 +638,7 @@ function buildTransaction(tx: Db, deps: WakeQueuePostgresAdapterDeps, db: Db, ru .returning() .then((rows) => rows[0]); - await notifyRunDispatch(tx); + await notifyRunDispatch(tx, { companyId }); const queuedRun = await tx .insert(heartbeatRuns) .values({ @@ -717,7 +717,7 @@ function buildTransaction(tx: Db, deps: WakeQueuePostgresAdapterDeps, db: Db, ru .returning() .then((rows) => rows[0]); - await notifyRunDispatch(tx); + await notifyRunDispatch(tx, { companyId }); const queuedRun = await tx .insert(heartbeatRuns) .values({ diff --git a/server/src/services/agent-configuration.ts b/server/src/services/agent-configuration.ts index 88a2c636bdb..dc10efaa082 100644 --- a/server/src/services/agent-configuration.ts +++ b/server/src/services/agent-configuration.ts @@ -84,7 +84,7 @@ export async function prepareAgentConfiguration(db: Db, existing: AgentConfigura existing.adapterConfig, ); } - await notifyRunDispatch(db); + await notifyRunDispatch(db, { companyId: existing.companyId }); if (normalizedPatch.runtimeConfig !== undefined || normalizedPatch.reportsTo !== undefined) { await notifyDeliveryWork(db, DELIVERY_QUEUES.agentTimers); } diff --git a/server/src/services/agent-lifecycle.ts b/server/src/services/agent-lifecycle.ts index 700c6ba8ffa..59357fcc5de 100644 --- a/server/src/services/agent-lifecycle.ts +++ b/server/src/services/agent-lifecycle.ts @@ -40,7 +40,7 @@ export function createAgentLifecycleEffects(_hooks: BudgetServiceHooks = {}): li notifyWork: tx => notifyDeliveryWork(tx, DELIVERY_QUEUES.agentLifecycle), async beforeStatusChange(tx, companyId, agentId, before, after) { if (before !== after) { - await notifyRunDispatch(tx); + await notifyRunDispatch(tx, { companyId }); await notifyDeliveryWork(tx, DELIVERY_QUEUES.agentTimers); } if (!isAgentStatusInvokable(before) && isAgentStatusInvokable(after)) { diff --git a/server/src/services/budgets.ts b/server/src/services/budgets.ts index 86d00c4d108..186253b19bb 100644 --- a/server/src/services/budgets.ts +++ b/server/src/services/budgets.ts @@ -288,7 +288,7 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica } if (policy.scopeType === "project") { - await notifyRunDispatch(db); + await notifyRunDispatch(db, { companyId: policy.companyId }); await db .update(projects) .set({ @@ -302,7 +302,7 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica const scope = await resolveScopeRecord(db, "company", policy.scopeId); if (scope.paused) return; - await notifyRunDispatch(db); + await notifyRunDispatch(db, { companyId: policy.companyId }); await notifyCompanyRetentionWork(db, policy.companyId); await notifyDeliveryWork(db, DELIVERY_QUEUES.agentTimers); await notifyDeliveryWork(db, DELIVERY_QUEUES.routineSchedules); @@ -349,7 +349,7 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica } if (policy.scopeType === "project") { - await notifyRunDispatch(db); + await notifyRunDispatch(db, { companyId: policy.companyId }); await db .update(projects) .set({ @@ -364,7 +364,7 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica await notifyDeliveryWork(db, DELIVERY_QUEUES.skillSourceSync); const scope = await resolveScopeRecord(db, "company", policy.scopeId); if (scope.pauseReason !== "budget") return; - await notifyRunDispatch(db); + await notifyRunDispatch(db, { companyId: policy.companyId }); await notifyCompanyRetentionWork(db, policy.companyId); await notifyDeliveryWork(db, DELIVERY_QUEUES.agentTimers); await notifyDeliveryWork(db, DELIVERY_QUEUES.routineSchedules); diff --git a/server/src/services/chat-run-milestone-notifications.ts b/server/src/services/chat-run-milestone-notifications.ts index 159055c10d3..524936395b5 100644 --- a/server/src/services/chat-run-milestone-notifications.ts +++ b/server/src/services/chat-run-milestone-notifications.ts @@ -12,7 +12,7 @@ import { DELIVERY_QUEUES, notifyDeliveryWork } from "./delivery-work-notificatio export type ChatRunMilestoneScope = { companyId: string; issueId?: string; runId?: string }; export async function notifyChatRunMilestones(tx: object, scope: ChatRunMilestoneScope) { - if (!scope.issueId && !scope.runId) await notifyRunDispatch(tx); + if (!scope.issueId && !scope.runId) await notifyRunDispatch(tx, scope); if (scope.issueId || scope.runId) { await notifyIssueReviewState(tx, { companyId: scope.companyId, issueId: scope.issueId, runId: scope.issueId ? undefined : scope.runId }); diff --git a/server/src/services/companies.ts b/server/src/services/companies.ts index cc10ffa3bb8..a9a8f59998c 100644 --- a/server/src/services/companies.ts +++ b/server/src/services/companies.ts @@ -333,7 +333,7 @@ export function companyService(db: Db, _budgetHooks: BudgetServiceHooks = {}) { const renamedPrefix = await resolveRenamedIssuePrefix(tx, id, companyPatch); if (companyPatch.status !== undefined && companyPatch.status !== existing.status) { - await notifyRunDispatch(tx); + await notifyRunDispatch(tx, { companyId: id }); await notifyDeliveryWork(tx, DELIVERY_QUEUES.agentTimers); await notifyDeliveryWork(tx, DELIVERY_QUEUES.skillSourceSync); await notifyDeliveryWork(tx, DELIVERY_QUEUES.routineSchedules); @@ -490,7 +490,7 @@ export function companyService(db: Db, _budgetHooks: BudgetServiceHooks = {}) { const wasAlreadyArchived = existing.status === "archived"; if (!wasAlreadyArchived) { - await notifyRunDispatch(tx); + await notifyRunDispatch(tx, { companyId: id }); await notifyDeliveryWork(tx, DELIVERY_QUEUES.agentTimers); await notifyDeliveryWork(tx, DELIVERY_QUEUES.skillSourceSync); await notifyDeliveryWork(tx, DELIVERY_QUEUES.routineSchedules); diff --git a/server/src/services/execution-control-notifications.ts b/server/src/services/execution-control-notifications.ts index 6fef650fe42..39b1d7765c6 100644 --- a/server/src/services/execution-control-notifications.ts +++ b/server/src/services/execution-control-notifications.ts @@ -1,9 +1,9 @@ -import { notifyRunDispatch } from "./run-dispatch-notifications.js"; +import { notifyRunDispatch, type RunDispatchScope } from "./run-dispatch-notifications.js"; import type { Db } from "@paperclipai/db"; import { DELIVERY_QUEUES, notifyDeliveryWork, type DeliveryQueue } from "./delivery-work-notifications.js"; -export async function notifyExecutionRunWork(transaction: object) { - await notifyRunDispatch(transaction); +export async function notifyExecutionRunWork(transaction: object, scope?: RunDispatchScope) { + await notifyRunDispatch(transaction, scope); for (const queue of [DELIVERY_QUEUES.executionStatus, DELIVERY_QUEUES.executionFinalization, DELIVERY_QUEUES.executionReplacement, DELIVERY_QUEUES.executionDisposition, DELIVERY_QUEUES.issueMonitors]) { await notifyDeliveryWork(transaction, queue); diff --git a/server/src/services/heartbeat.ts b/server/src/services/heartbeat.ts index 0772b920023..bbce8ba8f5b 100644 --- a/server/src/services/heartbeat.ts +++ b/server/src/services/heartbeat.ts @@ -1335,8 +1335,8 @@ export function heartbeatService( return heartbeatQueue.dispatchPendingNativeStatusWakeups(...args); } - function resumeQueuedRunsAfterDrain() { - const promise = heartbeatQueue.resumeQueuedRunsAfterDrain(); + function resumeQueuedRunsAfterDrain(...args: Parameters) { + const promise = heartbeatQueue.resumeQueuedRunsAfterDrain(...args); activeWakeupPromises.add(promise); void promise.then( () => { activeWakeupPromises.delete(promise); }, diff --git a/server/src/services/heartbeat/dispatch-work.ts b/server/src/services/heartbeat/dispatch-work.ts index 3291008dd9a..5b1f728688a 100644 --- a/server/src/services/heartbeat/dispatch-work.ts +++ b/server/src/services/heartbeat/dispatch-work.ts @@ -1,3 +1,4 @@ +import type { RunDispatchScope } from "../run-dispatch-notifications.js"; import type { createDeliveryWorkCoordinator } from "../delivery-work-coordinator.js"; import { registerDeadlineWork } from "../deadline-work.js"; import { DELIVERY_QUEUES } from "../delivery-work-notifications.js"; @@ -7,7 +8,7 @@ export function registerRunDispatchWork( owner: { promoteDueScheduledRetries(): Promise; nextScheduledRetryAt(): Promise; - resumeQueuedRunsAfterDrain(): Promise; + resumeQueuedRunsAfterDrain(scope?: RunDispatchScope): Promise; schedulingSuppressed(): Promise; }, canRun: () => boolean, @@ -16,6 +17,7 @@ export function registerRunDispatchWork( const retries = registerDeadlineWork(coordinator, DELIVERY_QUEUES.scheduledRetryDispatch, { retryMs: 30_000, continuationMs: 1_000, + notificationCoalesceMs: 1_000, canRun, run: async signal => { suppressed = await owner.schedulingSuppressed(); @@ -23,17 +25,34 @@ export function registerRunDispatchWork( }, next: () => suppressed ? Promise.resolve(null) : owner.nextScheduledRetryAt(), }); + const scopes = new Set(["*"]); const queue = registerDeadlineWork(coordinator, DELIVERY_QUEUES.queuedRunDispatch, { retryMs: 30_000, + notificationCoalesceMs: 1_000, canRun, + onSettled: hints => { for (const hint of hints ?? ["*"]) scopes.add(hint); }, run: async signal => { - if (!signal.aborted && canRun()) await owner.resumeQueuedRunsAfterDrain(); + const requested = scopes.has("*") ? ["*"] : [...scopes]; + if (requested[0] === "*") scopes.clear(); + else for (const hint of requested) scopes.delete(hint); + for (let index = 0; index < requested.length; index++) { + if (signal.aborted || !canRun()) { + for (const hint of requested.slice(index)) scopes.add(hint); + return; + } + try { + await owner.resumeQueuedRunsAfterDrain(requested[index] === "*" ? undefined : JSON.parse(requested[index]) as RunDispatchScope); + } catch (error) { + for (const hint of requested.slice(index)) scopes.add(hint); + throw error; + } + } }, next: async () => null, }); return { ready: Promise.all([retries.ready, queue.ready]), - wake() { retries.wake(); queue.wake(); }, + wake() { scopes.add("*"); retries.wake(); queue.wake(); }, async stop() { await Promise.all([retries.stop(), queue.stop()]); }, }; } diff --git a/server/src/services/heartbeat/queue.ts b/server/src/services/heartbeat/queue.ts index 376530b2f0a..64ca0b84ffd 100644 --- a/server/src/services/heartbeat/queue.ts +++ b/server/src/services/heartbeat/queue.ts @@ -1973,7 +1973,7 @@ export function createHeartbeatQueue(db: Db, dependencies: HeartbeatQueueDepende // Release admission promptly without first scanning unrelated recovery work. // Every run still passes the ordinary per-agent claim and dispatch gates. - async function resumeQueuedRunsAfterDrain() { + async function resumeQueuedRunsAfterDrain(scope?: { companyId: string }) { if ((await getSchedulingSuppression()).suppressed) return; const cutoff = await getWorktreeExecutionCutoff(); const queuedRuns = await db @@ -1983,6 +1983,7 @@ export function createHeartbeatQueue(db: Db, dependencies: HeartbeatQueueDepende .where( and( eq(heartbeatRuns.status, "queued"), + scope ? eq(heartbeatRuns.companyId, scope.companyId) : undefined, eq(companies.status, "active"), cutoff ? gte(heartbeatRuns.createdAt, cutoff) : undefined, ), @@ -4633,7 +4634,7 @@ export function createHeartbeatQueue(db: Db, dependencies: HeartbeatQueueDepende } const queueOutcome = await db.transaction(async (tx) => { - await notifyRunDispatch(tx); + await notifyRunDispatch(tx, { companyId: agent.companyId }); await tx.execute( sql`select id from agents where id = ${agentId} and company_id = ${agent.companyId} for update`, ); diff --git a/server/src/services/heartbeat/retries.ts b/server/src/services/heartbeat/retries.ts index 5f4df5ed74e..32333c0ebec 100644 --- a/server/src/services/heartbeat/retries.ts +++ b/server/src/services/heartbeat/retries.ts @@ -998,7 +998,7 @@ export function createHeartbeatRetries(db: Db, dependencies: HeartbeatRetryDepen retryContextSnapshot.previousRunId = continuation.previousRunId; } - await notifyRunDispatch(tx); + await notifyRunDispatch(tx, { companyId: run.companyId }); const wakeupRequest = await tx .insert(agentWakeupRequests) .values({ @@ -1756,7 +1756,7 @@ export function createHeartbeatRetries(db: Db, dependencies: HeartbeatRetryDepen }; const updated = await db.transaction(async (tx) => { - await notifyRunDispatch(tx); + await notifyRunDispatch(tx, { companyId: scheduled.run.companyId }); const row = await tx .update(heartbeatRuns) .set({ diff --git a/server/src/services/heartbeat/run-lifecycle.ts b/server/src/services/heartbeat/run-lifecycle.ts index 07bb54dfdf7..9d3b30367d2 100644 --- a/server/src/services/heartbeat/run-lifecycle.ts +++ b/server/src/services/heartbeat/run-lifecycle.ts @@ -1045,7 +1045,7 @@ export function createHeartbeatLifecycle(db: Db, dependencies: HeartbeatLifecycl if (previous && isHeartbeatRunTerminalStatus(status)) await notifyRetainedNativeSessionCleanup(tx, { companyId: previous.companyId, runId: previous.id }); if (isHeartbeatRunTerminalStatus(status)) await notifyDeliveryWork(tx, DELIVERY_QUEUES.accounting, companyId); - await notifyExecutionRunWork(tx); + await notifyExecutionRunWork(tx, companyId ? { companyId } : undefined); if (previous) await notifyChatRunMilestones(tx, { companyId: previous.companyId, runId: previous.id, issueId: previous.nativeIssueId ?? previous.issueId ?? readNonEmptyString(parseObject(previous.contextSnapshot).issueId) ?? undefined }); return write(tx); @@ -2448,7 +2448,7 @@ export function createHeartbeatLifecycle(db: Db, dependencies: HeartbeatLifecycl .returning() .then((rows) => rows[0]); - await notifyRunDispatch(tx); + await notifyRunDispatch(tx, { companyId: run.companyId }); const queuedRun = await tx .insert(heartbeatRuns) .values({ diff --git a/server/src/services/issue-review-state-notifications.ts b/server/src/services/issue-review-state-notifications.ts index b769be64056..d9d4d756816 100644 --- a/server/src/services/issue-review-state-notifications.ts +++ b/server/src/services/issue-review-state-notifications.ts @@ -5,7 +5,7 @@ import { notifyNativeReviewResponse } from "./native-review-response-notificatio import { notifyNativeStatusReassessment } from "./native-status-reassessment-notifications.js"; export async function notifyIssueReviewState(tx: object, scope: { companyId: string; issueId?: string; runId?: string; agentId?: string }) { - await notifyRunDispatch(tx); + await notifyRunDispatch(tx, scope); await notifyDeliveryWork(tx, DELIVERY_QUEUES.issueMonitors); await notifyExternalConversationSettlement(tx, scope); await notifyNativeReviewResponse(tx, scope); diff --git a/server/src/services/issues.ts b/server/src/services/issues.ts index a7c8bb55023..ade48bf0cc2 100644 --- a/server/src/services/issues.ts +++ b/server/src/services/issues.ts @@ -10606,7 +10606,7 @@ export function issueService(db: Db) { : values.parentId ? `issue:${values.parentId}` : undefined); } if (values.monitorNextCheckAt) await notifyDeliveryWork(tx, DELIVERY_QUEUES.issueMonitors); - await notifyRunDispatch(tx); + await notifyRunDispatch(tx, { companyId: values.companyId }); await notifyDeliveryWork(tx, DELIVERY_QUEUES.agentTimers); const [issue] = await tx.insert(issues).values(values).returning(); await attachOwnedDraftImages(tx, issue, issue.description, { userId: issueData.createdByUserId, agentId: issueData.createdByAgentId }); @@ -11495,7 +11495,7 @@ export function issueService(db: Db) { } if (patch.status !== undefined || patch.assigneeAgentId !== undefined || patch.assigneeUserId !== undefined || patch.hiddenAt !== undefined || patch.conversationAgentId !== undefined) { - await notifyRunDispatch(tx); + await notifyRunDispatch(tx, { companyId: receiptExisting.companyId }); await notifyDeliveryWork(tx, DELIVERY_QUEUES.agentTimers); } if (patch.monitorNextCheckAt !== undefined || (receiptExisting.monitorNextCheckAt && @@ -11960,7 +11960,7 @@ export function issueService(db: Db) { } await notifyDecisionTargetIssueRemoval(tx, id); await notifyDeliveryWork(tx, DELIVERY_QUEUES.issueMonitors); - await notifyRunDispatch(tx); + await notifyRunDispatch(tx, owner ? { companyId: owner.companyId } : undefined); await notifyDeliveryWork(tx, DELIVERY_QUEUES.agentTimers); removedIssue = await tx .delete(issues) @@ -12089,7 +12089,7 @@ export function issueService(db: Db) { : isNull(issues.executionRunId); const updated = await db.transaction(async tx => { await notifyDeliveryWork(tx, DELIVERY_QUEUES.issueMonitors); - await notifyRunDispatch(tx); + await notifyRunDispatch(tx, { companyId: issueCompany.companyId }); await notifyDeliveryWork(tx, DELIVERY_QUEUES.agentTimers); return tx .update(issues) @@ -12219,7 +12219,7 @@ export function issueService(db: Db) { const staleExecutionRunId = current.executionRunId; const adopted = await db.transaction(async tx => { await notifyDeliveryWork(tx, DELIVERY_QUEUES.issueMonitors); - await notifyRunDispatch(tx); + await notifyRunDispatch(tx, { companyId: issueCompany.companyId }); await notifyDeliveryWork(tx, DELIVERY_QUEUES.agentTimers); return tx .update(issues) @@ -12475,7 +12475,7 @@ export function issueService(db: Db) { await notifyTerminalWorkspaceWork(tx, `issue:${id}`); await notifyNativeStatusReassessment(tx, { companyId: existing.companyId, issueId: id }); await notifyDeliveryWork(tx, DELIVERY_QUEUES.issueMonitors); - await notifyRunDispatch(tx); + await notifyRunDispatch(tx, { companyId: existing.companyId }); await notifyDeliveryWork(tx, DELIVERY_QUEUES.agentTimers); const updated = await tx .update(issues) @@ -12536,7 +12536,7 @@ export function issueService(db: Db) { await notifyUnsafeWorkspaceExportRecovery(tx, { companyId: existing.companyId, issueId: id }); } await notifyDeliveryWork(tx, DELIVERY_QUEUES.issueMonitors); - await notifyRunDispatch(tx); + await notifyRunDispatch(tx, { companyId: existing.companyId }); await notifyDeliveryWork(tx, DELIVERY_QUEUES.agentTimers); const updated = await tx .update(issues) diff --git a/server/src/services/native-runtime/native-provider-capacity-retry.ts b/server/src/services/native-runtime/native-provider-capacity-retry.ts index 7d007c383dd..83314ff9c51 100644 --- a/server/src/services/native-runtime/native-provider-capacity-retry.ts +++ b/server/src/services/native-runtime/native-provider-capacity-retry.ts @@ -50,7 +50,7 @@ export async function materializeNativeProviderCapacityRetry(input: { scheduledRetryAt: dueAt.toISOString(), executionRetryAccounting: accountingForScheduledRetry(run, NATIVE_PROVIDER_OVERLOADED_CODE, attempt), }; - await notifyRunDispatch(input.tx); + await notifyRunDispatch(input.tx, { companyId: run.companyId }); const [wake] = await input.tx.insert(agentWakeupRequests).values({ companyId: run.companyId, agentId: run.agentId, source: "automation", triggerDetail: "system", reason: NATIVE_PROVIDER_OVERLOADED_CODE, status: "queued", runId: retryId, diff --git a/server/src/services/recovery/service.ts b/server/src/services/recovery/service.ts index 4393e0ad5cf..fffc2699dbf 100644 --- a/server/src/services/recovery/service.ts +++ b/server/src/services/recovery/service.ts @@ -2717,7 +2717,7 @@ export function recoveryService( const now = new Date(); const retryAt = readProviderQuotaRetryAt(input.latestRun, now); return database.transaction(async (tx) => { - await notifyRunDispatch(tx); + await notifyRunDispatch(tx, { companyId: input.issue.companyId }); const wakeup = await tx .insert(agentWakeupRequests) .values({ @@ -3318,7 +3318,7 @@ export function recoveryService( created = Boolean(enqueuedRun); } else { scheduledRun = await db.transaction(async (tx) => { - await notifyRunDispatch(tx); + await notifyRunDispatch(tx, { companyId: input.issue.companyId }); const wakeup = await tx .insert(agentWakeupRequests) .values({ diff --git a/server/src/services/run-dispatch-notifications.ts b/server/src/services/run-dispatch-notifications.ts index 9b1bb175894..2e04e6e7531 100644 --- a/server/src/services/run-dispatch-notifications.ts +++ b/server/src/services/run-dispatch-notifications.ts @@ -1,6 +1,9 @@ import { DELIVERY_QUEUES, notifyDeliveryWork } from "./delivery-work-notifications.js"; -export async function notifyRunDispatch(transaction: object) { - await notifyDeliveryWork(transaction, DELIVERY_QUEUES.scheduledRetryDispatch); - await notifyDeliveryWork(transaction, DELIVERY_QUEUES.queuedRunDispatch); +export type RunDispatchScope = { companyId: string }; + +export async function notifyRunDispatch(transaction: object, scope?: RunDispatchScope) { + const hint = scope ? JSON.stringify(scope) : "*"; + await notifyDeliveryWork(transaction, DELIVERY_QUEUES.scheduledRetryDispatch, hint); + await notifyDeliveryWork(transaction, DELIVERY_QUEUES.queuedRunDispatch, hint); } diff --git a/server/src/services/scheduled-retry-cancellation-notifications.ts b/server/src/services/scheduled-retry-cancellation-notifications.ts index c5548996607..24e57de66cb 100644 --- a/server/src/services/scheduled-retry-cancellation-notifications.ts +++ b/server/src/services/scheduled-retry-cancellation-notifications.ts @@ -4,6 +4,6 @@ import { DELIVERY_QUEUES, notifyDeliveryWork } from "./delivery-work-notificatio export type ScheduledRetryCancellationScope = { companyId: string; runId?: string }; export async function notifyScheduledRetryCancellation(transaction: object, scope?: ScheduledRetryCancellationScope) { - await notifyRunDispatch(transaction); + await notifyRunDispatch(transaction, scope); await notifyDeliveryWork(transaction, DELIVERY_QUEUES.scheduledRetryCancellation, scope ? JSON.stringify(scope) : "*"); }