diff --git a/src/main/services/heartbeat-scheduler.ts b/src/main/services/heartbeat-scheduler.ts index 41dbe1e0..9b410bd7 100644 --- a/src/main/services/heartbeat-scheduler.ts +++ b/src/main/services/heartbeat-scheduler.ts @@ -13,7 +13,7 @@ export class HeartbeatScheduler { /** Tasks already logged as budget-spent, so the park is announced once, not every 5s tick. */ private spentTasksLogged = new Set(); - constructor(private readonly db: DatabaseSync, private readonly cooldowns: ProviderCooldowns) {} + constructor(private readonly db: DatabaseSync, private readonly cooldowns: ProviderCooldowns, private readonly now: () => number) {} /** * Heartbeat: keeps every open assigned task of a heartbeat-enabled agent @@ -33,8 +33,10 @@ export class HeartbeatScheduler { */ enqueueHeartbeatRuns(): string[] { const window = currentProviderUsageWindow(); + const nowIso = new Date(this.now()).toISOString(); + // An agent's heartbeat interval paces it: one heartbeat run per interval, its most important task first. const rows = this.db.prepare(` - select t.id as task_id, t.org_id as org_id, p.id as provider_id + select t.id as task_id, t.org_id as org_id, p.id as provider_id, a.id as agent_id, a.heartbeat_interval_sec as interval_sec from tasks t join agents a on a.id = t.assignee_agent_id and a.org_id = t.org_id join providers p on p.id = a.model_provider_id and p.org_id = t.org_id @@ -47,13 +49,20 @@ export class HeartbeatScheduler { and coalesce(u.request_count, 0) < p.monthly_request_limit and not exists (select 1 from task_runs r where r.task_id = t.id and r.state in ('queued', 'running')) and not exists (select 1 from tasks c where c.parent_task_id = t.id and c.status <> 'done') - `).all(window.start) as Array<{ task_id: string; org_id: string; provider_id: string }>; + and (a.heartbeat_interval_sec <= 0 or not exists ( + select 1 from task_runs hb where hb.agent_id = a.id and hb.trigger = 'heartbeat' and hb.state <> 'superseded' + and hb.created_at > strftime('%Y-%m-%dT%H:%M:%fZ', ?, '-' || a.heartbeat_interval_sec || ' seconds') + )) + order by case t.priority when 'high' then 0 when 'normal' then 1 else 2 end, t.created_at + `).all(window.start, nowIso) as Array<{ task_id: string; org_id: string; provider_id: string; agent_id: string; interval_sec: number }>; this.observeQuotaBlocks(window.start); const cooling = new Set(this.cooldowns.cooledProviderIds()); const attemptMessages = this.db.prepare("select error_message from task_runs where task_id = ? and state <> 'superseded'"); const created: string[] = []; + const pacedAgents = new Set(); for (const row of rows) { if (cooling.has(row.provider_id)) continue; + if (row.interval_sec > 0 && pacedAgents.has(row.agent_id)) continue; const attempts = (attemptMessages.all(row.task_id) as Array<{ error_message: string | null }>) .filter((run) => !isTransientProviderFailure(run.error_message)).length; if (attempts >= TASK_RETRY_BUDGET) { @@ -62,6 +71,7 @@ export class HeartbeatScheduler { } try { created.push(this.enqueueRun(row.org_id, row.task_id)); + pacedAgents.add(row.agent_id); } catch { // Skip tasks that cannot be prepared yet (e.g. provider not fully configured). } @@ -120,7 +130,7 @@ export class HeartbeatScheduler { const workspace = loadWorkspace(this.db, task.workspace_id); const isolation = createIsolationContext(this.db, orgId, task, agent, provider, workspace); const runId = createId("run"); - const now = new Date().toISOString(); + const now = new Date(this.now()).toISOString(); this.db.prepare(` insert into task_runs( id, org_id, task_id, agent_id, provider_id, provider_request_id, diff --git a/src/main/services/task-execution-service.ts b/src/main/services/task-execution-service.ts index f14aa0aa..799197c6 100644 --- a/src/main/services/task-execution-service.ts +++ b/src/main/services/task-execution-service.ts @@ -78,7 +78,7 @@ export class TaskExecutionService { this.organization = new OrganizationService(db); this.chat = new ChatService(db); this.cooldowns = new ProviderCooldowns(now); - this.heartbeat = new HeartbeatScheduler(db, this.cooldowns); + this.heartbeat = new HeartbeatScheduler(db, this.cooldowns, now); } async runTask(input: RunTaskInput): Promise { diff --git a/tests/integration/task-dependencies.test.ts b/tests/integration/task-dependencies.test.ts index 58458fec..24380198 100644 --- a/tests/integration/task-dependencies.test.ts +++ b/tests/integration/task-dependencies.test.ts @@ -79,7 +79,8 @@ describe("task dependencies", () => { expect(waiting).toHaveLength(1); expect(queuedTasks()).toEqual([childId]); - h.connection.db.prepare("update task_runs set state = 'completed' where state = 'queued'").run(); + // Past the CEO's heartbeat interval, so pacing does not hold the parent back. + h.connection.db.prepare("update task_runs set state = 'completed', created_at = '2000-01-01T00:00:00.000Z' where state = 'queued'").run(); h.organization.updateTaskStatus(h.orgId, childId, TASK_STATUS.DONE); h.connection.transaction(() => h.execution.enqueueHeartbeatRuns()); expect(queuedTasks()).toEqual([parentId]); diff --git a/tests/unit/task-execution-service.test.ts b/tests/unit/task-execution-service.test.ts index 0187f202..347094e2 100644 --- a/tests/unit/task-execution-service.test.ts +++ b/tests/unit/task-execution-service.test.ts @@ -654,6 +654,50 @@ describe("task execution service", () => { connection.close(); }); + test("each agent's heartbeat interval paces how often its work is queued, most important task first", () => { + const connection = createConnection(); + const organization = new OrganizationService(connection.db); + const providers = new ProviderService(connection.db, providerAdapters, testVault); + let clock = Date.now(); + const execution = new TaskExecutionService(connection.db, providers, undefined, undefined, () => clock); + let orgId = ""; + const tasks: Record = {}; + connection.transaction(() => { + orgId = organization.createOrganization({ organizationName: "Acme", employeeName: "Dana Lee", mission: "m", goal: "g" }); + const providerId = openAiProviderId(connection, orgId); + providers.saveCredential({ orgId, providerId, secret: "sk-live-provider", model: "gpt-5.4" }); + for (const [name, interval] of [["Slow", 600], ["Eager", 0]] as const) { + const agentId = organization.createAgent({ + orgId, name, role: "Engineer", department: "Engineering", modelProviderId: providerId, modelId: "gpt-5.4", + budgetCents: 10000, managerAgentId: null, heartbeatEnabled: true, heartbeatIntervalSec: interval, + }); + for (const priority of [TASK_PRIORITY.LOW, TASK_PRIORITY.HIGH]) { + tasks[`${name}-${priority}`] = organization.createTask({ + orgId, title: `${name} ${priority}`, description: "d", status: TASK_STATUS.BACKLOG, assigneeAgentId: agentId, + risk: RISK_LEVEL.LOW, projectId: null, goalId: null, workspaceId: null, priority, requiresFilesystem: false, + }); + } + } + }); + const queuedTasks = () => (connection.db.prepare("select task_id from task_runs where state = 'queued'").all() as Array<{ task_id: string }>).map((row) => row.task_id).sort(); + const finishQueued = () => connection.db.prepare("update task_runs set state = 'completed' where state = 'queued'").run(); + + connection.transaction(() => execution.enqueueHeartbeatRuns()); + expect(queuedTasks()).toEqual([tasks["Eager-high"], tasks["Eager-low"], tasks["Slow-high"]].sort()); + finishQueued(); + + clock += 300_000; + connection.transaction(() => execution.enqueueHeartbeatRuns()); + expect(queuedTasks()).toEqual([tasks["Eager-high"], tasks["Eager-low"]].sort()); + finishQueued(); + + clock += 301_000; + connection.transaction(() => execution.enqueueHeartbeatRuns()); + expect(queuedTasks()).toContain(tasks["Slow-high"]); + expect(queuedTasks()).not.toContain(tasks["Slow-low"]); + connection.close(); + }); + test("an open task with a completed run keeps getting sessions until done or the budget is spent", () => { const connection = createConnection(); const organization = new OrganizationService(connection.db);