diff --git a/server/src/__tests__/deferred-execution-receipt-work.test.ts b/server/src/__tests__/deferred-execution-receipt-work.test.ts index 76fe58ed5dd..c74076e0568 100644 --- a/server/src/__tests__/deferred-execution-receipt-work.test.ts +++ b/server/src/__tests__/deferred-execution-receipt-work.test.ts @@ -21,7 +21,7 @@ import { startEmbeddedPostgresTestDatabase } from "./helpers/embedded-postgres.j import { createDeliveryWorkCoordinator } from "../services/delivery-work-coordinator.js"; import { notifyDeferredExecutionEligibility, notifyDeferredExecutionWork, notifyNativeStatusWakeWork, notifyTaskRetrySettlement } from "../services/deferred-execution-notifications.js"; import { registerDeferredReceiptWork } from "../services/heartbeat/deferred-receipt-work.js"; -import { heartbeatService } from "../services/heartbeat.js"; +import { heartbeatService, startTaskDrain, stopTaskDrain } from "../services/heartbeat.js"; import { requestTaskRetry } from "../services/task-retry.js"; import { issueService } from "../services/issues.js"; import { environmentService } from "../services/environments.js"; @@ -376,6 +376,85 @@ describe("deferred receipt durable dispatch", () => { expect(await db.select().from(heartbeatRuns).where(sql`${heartbeatRuns.contextSnapshot}->>'taskRetryRequestId' = ${receipt.requestId}`)).toHaveLength(1); await quiet(); }); + it("dispatches a task Retry held by a deploy drain as soon as the drain stops", async () => { + const f = await fixture(); + await heartbeat.registerDeferredReceiptWork(coordinator, () => !drain).ready; + await heartbeat.registerRunDispatchWork(coordinator, () => !drain).ready; + await heartbeat.registerScheduledRetryCancellationWork(coordinator, () => !drain).ready; + startTaskDrain({ ttlMs: 20 * 60_000 }); + let requestId: string; + try { + const receipt = await requestTaskRetry({ db, companyId: f.companyId, issueId: f.issue.id, actorId: f.actorId, + observedRunId: f.source.id, idempotencyKey: randomUUID() }); + expect(receipt.accepted).toBe(true); + requestId = receipt.requestId!; + await heartbeat.resumeTaskRetry(f.companyId, requestId); + await vi.advanceTimersByTimeAsync(1000); + const [held] = await db.select().from(agentWakeupRequests).where(eq(agentWakeupRequests.id, requestId)); + expect(held).toMatchObject({ status: "deferred_issue_execution", runId: null }); + } finally { stopTaskDrain(); } + let runId: string | undefined; + await waitForReal(async () => { + const [wake] = await db.select().from(agentWakeupRequests).where(eq(agentWakeupRequests.id, requestId)); + expect(wake.runId, JSON.stringify(wake)).toBeTruthy(); runId = wake.runId!; + }, { timeout: 5_000 }); + await waitForReal(async () => { + await heartbeat.drainActiveRunExecutions(); + const [run] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, runId!)); + expect(run.status, JSON.stringify({ error: run.error, code: run.errorCode })).toBe("succeeded"); + }, { timeout: 15_000 }); + await finish(f.issue.id); + await quiet(); + }); + it("dispatches an operator Retry ahead of a backlog of passive waits", async () => { + const f = await fixture(); + await db.insert(agentWakeupRequests).values(Array.from({ length: 120 }, (_, index) => ({ + id: `00000000-0000-4000-8000-${String(index).padStart(12, "0")}`, companyId: f.companyId, agentId: f.agent.id, + source: "automation", status: "deferred_issue_execution", requestedByActorType: "system", + payload: { executionWait: { reason: "recovery" } }, updatedAt: new Date(Date.now() - 60_000), + }))); + const [retry] = await db.insert(agentWakeupRequests).values({ id: "ffffffff-ffff-4fff-bfff-ffffffffffff", + companyId: f.companyId, agentId: f.agent.id, source: "on_demand", status: "deferred_issue_execution", + reason: "task_retry", requestedByActorType: "user", requestedByActorId: f.actorId, + payload: { issueId: f.issue.id } }).returning(); + const order: string[] = []; + const resume = vi.fn(async (scope: { requestId?: string }) => { + order.push(scope.requestId!); + await db.update(agentWakeupRequests).set({ status: "skipped" }).where(eq(agentWakeupRequests.id, scope.requestId!)); + }); + await registerDeferredReceiptWork(db, coordinator, { resume, settleTaskChat: vi.fn(), createNativeBatch: () => vi.fn(), + schedulingSuppressed: async () => false }, () => !drain).ready; + expect(order[0]).toBe(retry.id); + await waitForReal(() => expect(order).toHaveLength(121), { timeout: 30_000 }); + await quiet(); + }); + it("reloads a held Retry without a pending notification as soon as a drain ends, right after a full scan", async () => { + const f = await fixture(); + let refresh: (() => void) | undefined; + const resume = vi.fn(async (scope: { requestId?: string }) => { + await db.update(agentWakeupRequests).set({ status: "skipped" }).where(eq(agentWakeupRequests.id, scope.requestId!)); + }); + await registerDeferredReceiptWork(db, coordinator, { resume, settleTaskChat: vi.fn(), createNativeBatch: () => vi.fn(), + schedulingSuppressed: async () => false, subscribeRefresh: listener => { refresh = listener; return () => {}; } }, + () => !drain).ready; + const [retry] = await db.insert(agentWakeupRequests).values({ companyId: f.companyId, agentId: f.agent.id, + source: "on_demand", status: "deferred_issue_execution", reason: "task_retry", requestedByActorType: "user", + requestedByActorId: f.actorId, payload: { issueId: f.issue.id } }).returning(); + refresh!(); + await waitForReal(() => expect(resume).toHaveBeenCalledWith(expect.objectContaining({ requestId: retry.id })), + { timeout: 5_000 }); + await quiet(); + }); + it("removes its task-drain refresh listener when the coordinator stops", async () => { + let refresh: (() => void) | undefined; + const unsubscribe = vi.fn(); + await registerDeferredReceiptWork(db, coordinator, { resume: vi.fn(), settleTaskChat: vi.fn(), + createNativeBatch: () => vi.fn(), schedulingSuppressed: async () => false, + subscribeRefresh: listener => { refresh = listener; return unsubscribe; } }, () => !drain).ready; + expect(refresh).toBeTypeOf("function"); + await coordinator.stop(); + expect(unsubscribe).toHaveBeenCalledTimes(1); + }); async function blockedRetry(gate: "budget" | "project budget" | "paused agent", conversation = false) { const f = await fixture(), lifecycle = createAgentLifecycle(db), budgets = budgetService(db); if (conversation) await db.update(issues).set({ status: "in_review", conversationAgentId: f.agent.id, diff --git a/server/src/services/heartbeat.ts b/server/src/services/heartbeat.ts index f04965ec2cc..7724e412126 100644 --- a/server/src/services/heartbeat.ts +++ b/server/src/services/heartbeat.ts @@ -561,7 +561,8 @@ import { resolveAndRetainRunTrustPreset } from "./run-trust-preset.js"; import { effectiveProjectWorkspacePolicy, sourcelessProjectWorkspacePolicy } from "./project-workspace-provider.js"; import type { PluginWorkerManager } from "./plugin-worker-manager.js"; -import { computeTaskDrain, applyTaskDrain, startTaskDrain, stopTaskDrain, readTaskDrain } from "./task-admission.js"; +import { computeTaskDrain, applyTaskDrain, startTaskDrain, stopTaskDrain, readTaskDrain, subscribeTaskDrainChanges, + subscribeTaskDrainExpiry } from "./task-admission.js"; const execFile = promisify(execFileCallback); const UNSUCCESSFUL_HEARTBEAT_RUN_TERMINAL_STATUSES = [ @@ -4515,6 +4516,10 @@ export function heartbeatService( settleTaskChat: scope => reconcileTaskChatRetryReceipts(db, { ...scope, settleOnly: true }), createNativeBatch: heartbeatQueue.createNativeStatusWakeBatch, schedulingSuppressed: async () => (await getSchedulingSuppression(true)).suppressed, + subscribeRefresh: refresh => { + const unsubscribe = [subscribeTaskDrainChanges(refresh), subscribeTaskDrainExpiry(refresh)]; + return () => { for (const stop of unsubscribe) stop(); }; + }, }, () => canRun() && readTaskDrain(new Date()) === null), readIdleWatchdogSchedule: taskWatchdogs.readIdleSchedule, registerIssueStateRecoveryWork: (coordinator: Parameters[1], canRun: () => boolean) => diff --git a/server/src/services/heartbeat/deferred-receipt-work.ts b/server/src/services/heartbeat/deferred-receipt-work.ts index 5a3e3bddd64..bbf49a50739 100644 --- a/server/src/services/heartbeat/deferred-receipt-work.ts +++ b/server/src/services/heartbeat/deferred-receipt-work.ts @@ -7,7 +7,7 @@ import type { DeferredExecutionScope, DeferredReceiptRecoveryOptions } from "../ import { registerScopedDeadlineWork } from "../scoped-deadline-work.js"; import { dependencyResolvedWakePayloadSql } from "../issue-dependency-wakeups.js"; -type Entry = { id: string; companyId: string; dueAt: number; immediate: boolean }; +type Entry = { id: string; companyId: string; dueAt: number; immediate: boolean; operatorRetry?: boolean }; type NativeIntentOptions = DeferredExecutionScope & { onDeferredDeadline(at: number): void; onError(error: unknown): void; @@ -21,9 +21,15 @@ export function registerDeferredReceiptWork( settleTaskChat(scope: DeferredExecutionScope): Promise; createNativeBatch(): (scope: NativeIntentOptions) => Promise; schedulingSuppressed(): Promise; + subscribeRefresh?(refresh: () => void): () => void; }, canRun: () => boolean, ) { + let unsubscribeRefresh: (() => void) | undefined; + function releaseRefresh() { + unsubscribeRefresh?.(); + unsubscribeRefresh = undefined; + } function register(kind: "deferred" | "task_chat" | "native") { const cursors = new Map(); let enabled = false; @@ -45,23 +51,28 @@ export function registerDeferredReceiptWork( // receipts already checked; actual wait/claim deadlines still replace. mergeEntry: (pending, loaded) => pending.immediate && loaded.immediate ? { ...loaded, dueAt: Math.min(pending.dueAt, loaded.dueAt) } : loaded, + priority: entry => entry.operatorRetry ? 1 : 0, + onStop: releaseRefresh, load: async scope => { const key = scope ? JSON.stringify({ companyId: scope.companyId, requestId: scope.requestId }) : "*"; if (!enabled) { cursors.delete(key); return []; } const cursor = cursors.get(key); - const rows = await db.select({ id: agentWakeupRequests.id, companyId: agentWakeupRequests.companyId, + const operatorRetry = sql`coalesce(${agentWakeupRequests.reason}, '') = 'task_retry' + and coalesce(${agentWakeupRequests.requestedByActorType}, '') = 'user'`; + const columns = { id: agentWakeupRequests.id, companyId: agentWakeupRequests.companyId, status: agentWakeupRequests.status, updatedAt: agentWakeupRequests.updatedAt, claimedAt: agentWakeupRequests.claimedAt, error: agentWakeupRequests.error, + operatorRetry: sql`${operatorRetry}`, delayed: sql`${agentWakeupRequests.payload}->'queuedCommentInterrupt' is not null or ${agentWakeupRequests.payload}->'executionWait' is not null`, interruptedLegacy: sql`exists (select 1 from heartbeat_runs r where r.company_id = ${agentWakeupRequests.companyId} and r.agent_id = ${agentWakeupRequests.agentId} and r.runtime_mode = 'legacy' and r.status = 'cancelled' and r.result_json->>'queuedCommentInterruptQueueId' = ${agentWakeupRequests.id}::text)`, - }).from(agentWakeupRequests).where(and( + }; + const eligible = and( scope ? eq(agentWakeupRequests.companyId, scope.companyId) : undefined, scope?.requestId ? eq(agentWakeupRequests.id, scope.requestId) : undefined, - cursor ? gt(agentWakeupRequests.id, cursor) : undefined, kind === "native" ? and( eq(agentWakeupRequests.requestedByActorType, "system"), eq(agentWakeupRequests.requestedByActorId, "native-status-committer"), @@ -80,17 +91,25 @@ export function registerDeferredReceiptWork( and r.runtime_mode = 'legacy' and r.status = 'cancelled' and r.result_json->>'queuedCommentInterruptQueueId' = ${agentWakeupRequests.id}::text)`, )), - )).orderBy(asc(agentWakeupRequests.id)).limit(100); + ); + const prioritized = kind === "deferred" && !cursor + ? await db.select(columns).from(agentWakeupRequests).where(and(eligible, operatorRetry)) + .orderBy(asc(agentWakeupRequests.id)).limit(100) + : []; + const rows = await db.select(columns).from(agentWakeupRequests).where(and(eligible, + cursor ? gt(agentWakeupRequests.id, cursor) : undefined)) + .orderBy(asc(agentWakeupRequests.id)).limit(100); const complete = rows.length < 100; if (complete) cursors.delete(key); else cursors.set(key, rows.at(-1)!.id); const now = Date.now(); - return { rows: rows.map(row => { + return { rows: [...prioritized, ...rows].map(row => { const deadline = kind === "native" && row.status === "claimed" && row.claimedAt ? row.claimedAt.getTime() + 60_000 : kind === "native" && row.error ? row.updatedAt.getTime() + 30_000 : kind === "deferred" && row.delayed && !row.interruptedLegacy ? row.updatedAt.getTime() + 30_000 : null; const immediate = deadline === null || deadline <= now; - return { id: row.id, companyId: row.companyId, dueAt: immediate ? now : deadline, immediate }; + return { id: row.id, companyId: row.companyId, dueAt: immediate ? now : deadline, immediate, + operatorRetry: kind === "deferred" && row.operatorRetry }; }), complete }; }, check: async entry => { @@ -132,6 +151,7 @@ export function registerDeferredReceiptWork( }, canRun); } const workers = [register("deferred"), register("task_chat"), register("native")]; + unsubscribeRefresh = owner.subscribeRefresh?.(() => { for (const worker of workers) worker.refresh(); }); return { ready: Promise.all(workers.map(worker => worker.ready)), wake() { for (const worker of workers) worker.invalidate(); }, async stop() { await Promise.all(workers.map(worker => worker.stop())); } }; diff --git a/server/src/services/scoped-deadline-work.ts b/server/src/services/scoped-deadline-work.ts index b6613ac54c1..117632b52c4 100644 --- a/server/src/services/scoped-deadline-work.ts +++ b/server/src/services/scoped-deadline-work.ts @@ -32,7 +32,9 @@ export function registerScopedDeadlineWork( load(scope?: Scope, signal?: AbortSignal): Promise; // Owners may retain pending queue priority while refreshing eligibility. mergeEntry?(pending: Entry, loaded: Entry): Entry; + priority?(entry: Entry): number; check(entry: Entry, signal?: AbortSignal): Promise; + onStop?(): void; }, canRun: () => boolean, ) { @@ -73,7 +75,7 @@ export function registerScopedDeadlineWork( const worker = coordinator.register(options.queue, { retryMs: 30_000, pausePassiveMaintenance: true, - onStop: () => { stopped = true; revision++; }, + onStop: () => { stopped = true; revision++; options.onStop?.(); }, onSettled: hints => { revision++; for (const hint of hints ?? ["*"]) { dirty.add(hint); versions.set(hint, (versions.get(hint) ?? 0) + 1); } }, @@ -134,7 +136,8 @@ export function registerScopedDeadlineWork( const unavailable = dirtyScopes(); const due = [...entries.values()].filter(entry => !unavailable.some(scope => options.matches(entry, scope)) && Math.max(entry.dueAt, failures.get(entry.id)?.retryAt ?? 0) <= Date.now()) - .sort((a, b) => a.dueAt - b.dueAt || a.id.localeCompare(b.id)); + .sort((a, b) => (options.priority?.(b) ?? 0) - (options.priority?.(a) ?? 0) + || a.dueAt - b.dueAt || a.id.localeCompare(b.id)); const batch = takeCompanyWorkBatch(due, options.batchSize ?? 25, lastCompany); lastCompany = batch.lastCompany; // An owner may share more than one resource (for example a task and an // endpoint). Merge intersecting keys before dispatch so waiting siblings @@ -275,6 +278,11 @@ export function registerScopedDeadlineWork( nextRunAt: next, nextRunAtOnError: next, }); + function invalidate(hint = "*") { + revision++; + dirty.add(hint); versions.set(hint, (versions.get(hint) ?? 0) + 1); + worker.wake(); + } return { ...worker, reconciliationSnapshot: (): DeadlineReconciliationSnapshot => ({ revision: `${ownerRevision}:${revision}`, @@ -282,10 +290,10 @@ export function registerScopedDeadlineWork( && failures.size === 0 && scopeFailures.size === 0, }), wake() { revision++; worker.wake(); }, - invalidate(hint = "*") { - revision++; - dirty.add(hint); versions.set(hint, (versions.get(hint) ?? 0) + 1); - worker.wake(); + invalidate, + refresh(hint = "*") { + reconstructAfter.delete(hint); + invalidate(hint); }, }; }