From 49ea508ee121dc07aa5e3948f714307845fe0850 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Tue, 28 Apr 2026 20:26:29 -0400 Subject: [PATCH] fix(scheduler): drain pending wakeup requests --- src/main/database/wakeup.ts | 11 ++ src/main/orchestration/scheduler.ts | 41 ++++++- tests/unit/scheduler-heartbeat.test.ts | 143 +++++++++++++++++++++++++ 3 files changed, 191 insertions(+), 4 deletions(-) diff --git a/src/main/database/wakeup.ts b/src/main/database/wakeup.ts index 8dde389a..2e83f592 100644 --- a/src/main/database/wakeup.ts +++ b/src/main/database/wakeup.ts @@ -159,6 +159,17 @@ export function buildWakeupMethods( return promoted; }, + listPendingWakeupAgents(companyId: string): Array<{ agentId: string }> { + return readAll( + `SELECT agent_id AS agentId + FROM wakeup_requests + WHERE company_id = ? AND status IN ('queued', 'deferred') + GROUP BY agent_id + ORDER BY MIN(requested_at) ASC`, + companyId, + ); + }, + listWakeupRequests(companyId: string, filters?: { status?: string; agentId?: string; limit?: number }): Array<{ id: string; companyId: string; agentId: string; source: string; triggerDetail: string | null; status: string; coalescedCount: number; runId: string | null; requestedAt: string; diff --git a/src/main/orchestration/scheduler.ts b/src/main/orchestration/scheduler.ts index eb1d1fe1..458828b5 100644 --- a/src/main/orchestration/scheduler.ts +++ b/src/main/orchestration/scheduler.ts @@ -8,6 +8,7 @@ import type { OrchestrationDependencies } from "./types"; type SchedulerPhase = "lock_recovery" | "housekeeping" | "idle_automation" | "heartbeat_cadence"; const TICK_INTERVAL_MS = 10_000; +const PENDING_WAKE_RETRY_INTERVAL_MS = 120_000; const PHASE_INTERVAL_MS: Record = { lock_recovery: 60_000, housekeeping: 300_000, @@ -36,6 +37,7 @@ export function createScheduler( let unlockListener: (() => void) | null = null; const phaseLastRunAt = new Map(); const scheduleLastFiredByCompany = new Map(); + const pendingWakeAttemptByAgent = new Map(); const taskWakeAttemptByAgent = new Map(); const bootstrapAttemptByAgent = new Map(); @@ -46,6 +48,10 @@ export function createScheduler( return true; } + function agentAttemptKey(companyId: string, agentId: string): string { + return `${companyId}:${agentId}`; + } + function runLockRecoveryPhase() { if (!hasWorkerReady()) return; try { @@ -206,13 +212,14 @@ export function createScheduler( } } - function runCompanyBootstrap(company: { id: string; name: string }, now: number) { + function runCompanyBootstrap(company: { id: string; name: string }, now: number, wokenAgentIds: Set) { const goalCount = db.core.entity.countActiveGoalsForCompany(company.id); const advancingCount = db.core.tasks.countAdvancingTasksForCompany(company.id); if (goalCount === 0 || advancingCount > 0) return; const ceo = db.core.agents.getCeoForCompany(company.id); if (!ceo) return; + if (wokenAgentIds.has(ceo.id)) return; if (hasActiveRun({ agentId: ceo.id, companyId: company.id })) return; const lastHeartbeat = ceo.lastHeartbeatAt ? new Date(ceo.lastHeartbeatAt).getTime() : 0; @@ -232,9 +239,34 @@ export function createScheduler( scheduleLastFiredByCompany.set(companyId, now); } - function runAssignedTaskWakeCadence(companyId: string, now: number) { + function runPendingWakeupCadence(companyId: string, now: number): Set { + const wokenAgentIds = new Set(); + const pendingWakeupAgents = db.automation.wakeup.listPendingWakeupAgents(companyId); + for (const { agentId } of pendingWakeupAgents) { + if (hasActiveRun({ agentId, companyId })) continue; + if (deps.checkpointEngine?.shouldPauseHeartbeat(companyId, agentId)) continue; + + const attemptKey = agentAttemptKey(companyId, agentId); + const lastAttempt = pendingWakeAttemptByAgent.get(attemptKey) ?? 0; + if (lastAttempt > 0 && now - lastAttempt < PENDING_WAKE_RETRY_INTERVAL_MS) continue; + + const trigger = db.automation.wakeup.getPreferredWakeTrigger(agentId, companyId); + if (!trigger) continue; + + pendingWakeAttemptByAgent.set(attemptKey, now); + const runId = wakeAgentIfPossible(agentId, companyId, trigger); + if (runId) { + pendingWakeAttemptByAgent.delete(attemptKey); + wokenAgentIds.add(agentId); + } + } + return wokenAgentIds; + } + + function runAssignedTaskWakeCadence(companyId: string, now: number, wokenAgentIds: Set) { const activeAgents = db.core.agents.listActiveAgentsForCompany(companyId); for (const activeAgent of activeAgents) { + if (wokenAgentIds.has(activeAgent.id)) continue; if (hasActiveRun({ agentId: activeAgent.id, companyId })) continue; if (deps.checkpointEngine?.shouldPauseHeartbeat(companyId, activeAgent.id)) continue; @@ -258,8 +290,9 @@ export function createScheduler( for (const company of activeCompanies) { try { runScheduleAutomation(company.id, now); - runAssignedTaskWakeCadence(company.id, now); - runCompanyBootstrap(company, now); + const wokenAgentIds = runPendingWakeupCadence(company.id, now); + runAssignedTaskWakeCadence(company.id, now, wokenAgentIds); + runCompanyBootstrap(company, now, wokenAgentIds); } catch (error) { logAppError(logger, error, { kind: "recoverable", diff --git a/tests/unit/scheduler-heartbeat.test.ts b/tests/unit/scheduler-heartbeat.test.ts index f4299d3b..1b56ca40 100644 --- a/tests/unit/scheduler-heartbeat.test.ts +++ b/tests/unit/scheduler-heartbeat.test.ts @@ -199,6 +199,149 @@ describe("unified scheduler heartbeat cadence", () => { expect(wakeAgentIfPossible).toHaveBeenCalledWith(agentId, companyId, "timer"); }); + it("wakes an agent with a queued wakeup request on heartbeat cadence", () => { + const companyId = createCompany(db); + createConnector(db); + const workspaceId = createWorkspace(db, companyId, tempDir); + const agentId = createAgent(db, companyId, workspaceId); + db.automation.wakeup.enqueueWakeupRequest( + companyId, + agentId, + "subtask_completed", + "Queued subtask completion wake", + ); + const wakeAgentIfPossible = vi.fn(); + const scheduler = createScheduler( + createDeps(db), + () => false, + wakeAgentIfPossible, + ); + + try { + scheduler.startAllSchedulers(); + vi.advanceTimersByTime(60_000); + } finally { + scheduler.stopAllSchedulers(); + } + + expect(wakeAgentIfPossible).toHaveBeenCalledWith(agentId, companyId, "subtask_completed"); + }); + + it("throttles queued wakeup recovery attempts", () => { + const companyId = createCompany(db); + createConnector(db); + const workspaceId = createWorkspace(db, companyId, tempDir); + const agentId = createAgent(db, companyId, workspaceId); + db.automation.wakeup.enqueueWakeupRequest( + companyId, + agentId, + "subtask_completed", + "Queued subtask completion wake", + ); + const wakeAgentIfPossible = vi.fn(); + const scheduler = createScheduler( + createDeps(db), + () => false, + wakeAgentIfPossible, + ); + + try { + scheduler.startAllSchedulers(); + vi.advanceTimersByTime(60_000); + expect(wakeAgentIfPossible).toHaveBeenCalledTimes(1); + + vi.advanceTimersByTime(60_000); + expect(wakeAgentIfPossible).toHaveBeenCalledTimes(1); + + vi.advanceTimersByTime(60_000); + expect(wakeAgentIfPossible).toHaveBeenCalledTimes(2); + } finally { + scheduler.stopAllSchedulers(); + } + }); + + it("clears queued wakeup throttle after a started run", () => { + const companyId = createCompany(db); + createConnector(db); + const workspaceId = createWorkspace(db, companyId, tempDir); + const agentId = createAgent(db, companyId, workspaceId); + db.automation.wakeup.enqueueWakeupRequest( + companyId, + agentId, + "subtask_completed", + "Queued subtask completion wake", + ); + let runCount = 0; + const wakeAgentIfPossible = vi.fn((wokenAgentId: string, wokenCompanyId: string) => { + runCount += 1; + const runId = `run-${runCount}`; + db.automation.wakeup.completePendingWakeupRequests(wokenAgentId, wokenCompanyId, runId); + return runId; + }); + const scheduler = createScheduler( + createDeps(db), + () => false, + wakeAgentIfPossible, + ); + + try { + scheduler.startAllSchedulers(); + vi.advanceTimersByTime(60_000); + expect(wakeAgentIfPossible).toHaveBeenCalledTimes(1); + + db.automation.wakeup.enqueueWakeupRequest( + companyId, + agentId, + "approval_resolved", + "Queued approval resolution wake", + ); + vi.advanceTimersByTime(60_000); + } finally { + scheduler.stopAllSchedulers(); + } + + expect(wakeAgentIfPossible).toHaveBeenCalledTimes(2); + expect(wakeAgentIfPossible).toHaveBeenLastCalledWith(agentId, companyId, "approval_resolved"); + }); + + it("lets assigned task wake run after a queued wakeup attempt fails", () => { + const companyId = createCompany(db); + createConnector(db); + const workspaceId = createWorkspace(db, companyId, tempDir); + const agentId = createAgent(db, companyId, workspaceId); + db.automation.wakeup.enqueueWakeupRequest( + companyId, + agentId, + "subtask_completed", + "Queued subtask completion wake", + ); + db.core.tasks.saveTask({ + companyId, + title: "Continue after failed queued wake", + description: "The assigned task cadence can recover the agent in the same tick.", + assigneeAgentId: agentId, + workspaceId, + priority: "high", + status: "todo", + }); + const wakeAgentIfPossible = vi.fn(() => null); + const scheduler = createScheduler( + createDeps(db), + () => false, + wakeAgentIfPossible, + ); + + try { + scheduler.startAllSchedulers(); + vi.advanceTimersByTime(60_000); + } finally { + scheduler.stopAllSchedulers(); + } + + expect(wakeAgentIfPossible).toHaveBeenNthCalledWith(1, agentId, companyId, "subtask_completed"); + expect(wakeAgentIfPossible).toHaveBeenNthCalledWith(2, agentId, companyId, "timer"); + }); + it("bootstraps the CEO through the wake pipeline", () => { const companyId = createCompany(db); createConnector(db);