diff --git a/src/main/database.ts b/src/main/database.ts index 7dc9f034..9bda2123 100644 --- a/src/main/database.ts +++ b/src/main/database.ts @@ -1826,6 +1826,58 @@ export class AppDatabase { ); } + /** + * Find the most recent completed heartbeat task for this agent that can be reused. + * Returns a task with status "done" or "in_review", ordered by most recently updated. + * Returns null if no reusable heartbeat task exists. + */ + findReusableHeartbeatTask(agentId: string, companyId: string): TaskRecord | null { + const row = this.db + .prepare( + `SELECT id, company_id as companyId, project_id as projectId, goal_id as goalId, + parent_id as parentId, title, description, + assignee_agent_id as assigneeAgentId, workspace_id as workspaceId, + status, priority, active_run_id as activeRunId, + session_json as sessionJson, session_display_id as sessionDisplayId, + cost_usd as costUsd, requires_user_review as requiresUserReview, + created_at as createdAt, updated_at as updatedAt + FROM tasks + WHERE company_id = ? AND assignee_agent_id = ? AND title LIKE '[Heartbeat]%' + AND status IN ('done', 'in_review') + ORDER BY updated_at DESC + LIMIT 1`, + ) + .get(companyId, agentId) as TaskRecord | undefined; + if (row) { + row.requiresUserReview = Boolean(row.requiresUserReview); + } + return row ?? null; + } + + /** + * Cancel all duplicate heartbeat tasks per agent within a company, keeping only the most + * recently updated one per agent. Does NOT touch in_progress tasks (they may have active runs). + * Returns the number of tasks cancelled. + */ + cleanupDuplicateHeartbeatTasks(companyId: string): number { + const ts = nowIso(); + const result = this.db + .prepare( + `UPDATE tasks SET status = 'cancelled', updated_at = ? + WHERE company_id = ? AND title LIKE '[Heartbeat]%' + AND status IN ('done', 'in_review', 'todo', 'blocked', 'backlog') + AND id NOT IN ( + SELECT id FROM ( + SELECT id, ROW_NUMBER() OVER (PARTITION BY assignee_agent_id ORDER BY updated_at DESC) as rn + FROM tasks + WHERE company_id = ? AND title LIKE '[Heartbeat]%' AND status != 'cancelled' + ) WHERE rn = 1 + )`, + ) + .run(ts, companyId, companyId); + return Number(result.changes); + } + /** * List tasks from an agent's direct reports that are awaiting review (status = in_review). * Used to populate the "Tasks Awaiting Your Review" section in the agent prompt. diff --git a/tests/unit/heartbeat-coalescing.test.ts b/tests/unit/heartbeat-coalescing.test.ts new file mode 100644 index 00000000..b5f96b62 --- /dev/null +++ b/tests/unit/heartbeat-coalescing.test.ts @@ -0,0 +1,312 @@ +import { join } from "node:path"; +import { tmpdir } from "node:os"; +import { mkdtemp, rm } from "node:fs/promises"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +vi.mock("electron", () => ({ + safeStorage: { + isEncryptionAvailable: () => false, + getSelectedStorageBackend: () => "basic_text", + encryptString: (value: string) => Buffer.from(value, "utf8"), + decryptString: (value: Buffer) => value.toString("utf8"), + }, +})); + +import { AppDatabase } from "../../src/main/database"; + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +function createTestCompany(db: AppDatabase): string { + return db.saveCompany({ + name: "Test Company", + description: "Heartbeat coalescing tests", + status: "active", + }); +} + +function createTestConnector(db: AppDatabase): void { + db.saveConnector({ + id: "codex_local", + command: "codex", + model: "gpt-5", + envBindingText: "", + notes: "Test connector", + }); +} + +function createTestWorkspace(db: AppDatabase, companyId: string): string { + return db.saveWorkspace({ + companyId, + name: "Test Workspace", + localPath: "/tmp/test-workspace", + repoUrl: "https://github.com/example/test", + repoRef: "main", + isPrimary: true, + }); +} + +function createTestAgent(db: AppDatabase, companyId: string, workspaceId: string, name = "Agent"): string { + return db.saveAgent({ + companyId, + name, + role: "IC", + title: name, + status: "active", + reportsTo: null, + connectorId: "codex_local", + workspaceId, + model: null, + capabilities: "Executes tasks", + budgetMonthlyUsd: 100, + }); +} + +function createHeartbeatTask( + db: AppDatabase, + companyId: string, + agentId: string, + status: string, + agentName = "Agent", +): string { + return db.saveTask({ + companyId, + title: `[Heartbeat] ${agentName} — Check for work`, + description: `Automated heartbeat run for ${agentName}.`, + assigneeAgentId: agentId, + status, + priority: "medium", + }); +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +describe("Heartbeat task coalescing", () => { + let tempDir: string; + let db: AppDatabase; + + beforeEach(async () => { + tempDir = await mkdtemp(join(tmpdir(), "agentcompany-heartbeat-")); + db = new AppDatabase(tempDir); + db.init(); + createTestConnector(db); + }); + + afterEach(async () => { + db.close(); + await rm(tempDir, { recursive: true, force: true }); + }); + + // ------------------------------------------------------------------------- + // findReusableHeartbeatTask + // ------------------------------------------------------------------------- + + describe("findReusableHeartbeatTask", () => { + it("returns the most recent done heartbeat task for the agent", () => { + const companyId = createTestCompany(db); + const wsId = createTestWorkspace(db, companyId); + const agentId = createTestAgent(db, companyId, wsId, "CEO"); + + const taskId = createHeartbeatTask(db, companyId, agentId, "done", "CEO"); + + const result = db.findReusableHeartbeatTask(agentId, companyId); + expect(result).not.toBeNull(); + expect(result!.id).toBe(taskId); + expect(result!.status).toBe("done"); + }); + + it("returns null when no heartbeat tasks exist for the agent", () => { + const companyId = createTestCompany(db); + const wsId = createTestWorkspace(db, companyId); + const agentId = createTestAgent(db, companyId, wsId, "Worker"); + + const result = db.findReusableHeartbeatTask(agentId, companyId); + expect(result).toBeNull(); + }); + + it("returns null when all heartbeat tasks are still in_progress", () => { + const companyId = createTestCompany(db); + const wsId = createTestWorkspace(db, companyId); + const agentId = createTestAgent(db, companyId, wsId, "Worker"); + + createHeartbeatTask(db, companyId, agentId, "in_progress", "Worker"); + + const result = db.findReusableHeartbeatTask(agentId, companyId); + expect(result).toBeNull(); + }); + + it("returns the most recently updated heartbeat task when multiple done ones exist", () => { + const companyId = createTestCompany(db); + const wsId = createTestWorkspace(db, companyId); + const agentId = createTestAgent(db, companyId, wsId, "CEO"); + + // Create two done heartbeat tasks + const oldTaskId = createHeartbeatTask(db, companyId, agentId, "done", "CEO"); + const newTaskId = createHeartbeatTask(db, companyId, agentId, "done", "CEO"); + + // Advance the newer one so its updated_at is more recent + db.advanceTaskStatus(newTaskId, "in_review", "Testing"); + db.advanceTaskStatus(newTaskId, "done", "Testing done"); + + const result = db.findReusableHeartbeatTask(agentId, companyId); + expect(result).not.toBeNull(); + expect(result!.id).toBe(newTaskId); + }); + + it("returns in_review task if no done tasks exist", () => { + const companyId = createTestCompany(db); + const wsId = createTestWorkspace(db, companyId); + const agentId = createTestAgent(db, companyId, wsId, "Lead"); + + const taskId = createHeartbeatTask(db, companyId, agentId, "in_review", "Lead"); + + const result = db.findReusableHeartbeatTask(agentId, companyId); + expect(result).not.toBeNull(); + expect(result!.id).toBe(taskId); + }); + + it("does not return heartbeat tasks from other agents", () => { + const companyId = createTestCompany(db); + const wsId = createTestWorkspace(db, companyId); + const agent1 = createTestAgent(db, companyId, wsId, "Agent1"); + const agent2 = createTestAgent(db, companyId, wsId, "Agent2"); + + createHeartbeatTask(db, companyId, agent1, "done", "Agent1"); + + const result = db.findReusableHeartbeatTask(agent2, companyId); + expect(result).toBeNull(); + }); + + it("does not return non-heartbeat tasks", () => { + const companyId = createTestCompany(db); + const wsId = createTestWorkspace(db, companyId); + const agentId = createTestAgent(db, companyId, wsId, "Worker"); + + // Create a regular done task (not a heartbeat) + db.saveTask({ + companyId, + title: "Regular completed task", + description: "Not a heartbeat", + assigneeAgentId: agentId, + status: "done", + priority: "medium", + }); + + const result = db.findReusableHeartbeatTask(agentId, companyId); + expect(result).toBeNull(); + }); + }); + + // ------------------------------------------------------------------------- + // cleanupDuplicateHeartbeatTasks + // ------------------------------------------------------------------------- + + describe("cleanupDuplicateHeartbeatTasks", () => { + it("cancels duplicate heartbeat tasks, keeping only the most recent per agent", () => { + const companyId = createTestCompany(db); + const wsId = createTestWorkspace(db, companyId); + const agentId = createTestAgent(db, companyId, wsId, "CEO"); + + // Create three done heartbeat tasks + const _task1 = createHeartbeatTask(db, companyId, agentId, "done", "CEO"); + const _task2 = createHeartbeatTask(db, companyId, agentId, "done", "CEO"); + const task3 = createHeartbeatTask(db, companyId, agentId, "done", "CEO"); + + // Touch task3 so it becomes the most recent + db.advanceTaskStatus(task3, "in_review", "Touch"); + db.advanceTaskStatus(task3, "done", "Touch done"); + + const cleaned = db.cleanupDuplicateHeartbeatTasks(companyId); + expect(cleaned).toBe(2); + + // The most recent one should still be done + const remaining = db.getTask(task3); + expect(remaining.status).toBe("done"); + + // The older ones should be cancelled + const old1 = db.getTask(_task1); + const old2 = db.getTask(_task2); + expect(old1.status).toBe("cancelled"); + expect(old2.status).toBe("cancelled"); + }); + + it("returns the count of cancelled tasks", () => { + const companyId = createTestCompany(db); + const wsId = createTestWorkspace(db, companyId); + const agentId = createTestAgent(db, companyId, wsId, "Worker"); + + createHeartbeatTask(db, companyId, agentId, "done", "Worker"); + createHeartbeatTask(db, companyId, agentId, "todo", "Worker"); + const newest = createHeartbeatTask(db, companyId, agentId, "done", "Worker"); + // Touch newest so it has the latest updated_at + db.advanceTaskStatus(newest, "in_review", "Touch"); + db.advanceTaskStatus(newest, "done", "Touch done"); + + const cleaned = db.cleanupDuplicateHeartbeatTasks(companyId); + expect(cleaned).toBe(2); + }); + + it("does not cancel in_progress heartbeat tasks", () => { + const companyId = createTestCompany(db); + const wsId = createTestWorkspace(db, companyId); + const agentId = createTestAgent(db, companyId, wsId, "Worker"); + + const activeTask = createHeartbeatTask(db, companyId, agentId, "in_progress", "Worker"); + const doneTask = createHeartbeatTask(db, companyId, agentId, "done", "Worker"); + // Touch doneTask to make it more recent + db.advanceTaskStatus(doneTask, "in_review", "Touch"); + db.advanceTaskStatus(doneTask, "done", "Touch done"); + + const cleaned = db.cleanupDuplicateHeartbeatTasks(companyId); + // Only 0 cancelled because in_progress cannot be cancelled and doneTask is the most recent + expect(cleaned).toBe(0); + + // in_progress task should be untouched + const active = db.getTask(activeTask); + expect(active.status).toBe("in_progress"); + }); + + it("keeps one heartbeat task per agent when company has multiple agents", () => { + const companyId = createTestCompany(db); + const wsId = createTestWorkspace(db, companyId); + const agent1 = createTestAgent(db, companyId, wsId, "Agent1"); + const agent2 = createTestAgent(db, companyId, wsId, "Agent2"); + + // Agent1: 3 heartbeat tasks + createHeartbeatTask(db, companyId, agent1, "done", "Agent1"); + createHeartbeatTask(db, companyId, agent1, "done", "Agent1"); + const a1newest = createHeartbeatTask(db, companyId, agent1, "done", "Agent1"); + db.advanceTaskStatus(a1newest, "in_review", "Touch"); + db.advanceTaskStatus(a1newest, "done", "Touch done"); + + // Agent2: 2 heartbeat tasks + createHeartbeatTask(db, companyId, agent2, "done", "Agent2"); + const a2newest = createHeartbeatTask(db, companyId, agent2, "done", "Agent2"); + db.advanceTaskStatus(a2newest, "in_review", "Touch"); + db.advanceTaskStatus(a2newest, "done", "Touch done"); + + const cleaned = db.cleanupDuplicateHeartbeatTasks(companyId); + // 2 old from agent1 + 1 old from agent2 = 3 + expect(cleaned).toBe(3); + + // Most recent per agent should survive + expect(db.getTask(a1newest).status).toBe("done"); + expect(db.getTask(a2newest).status).toBe("done"); + }); + + it("returns 0 when no duplicate heartbeat tasks exist", () => { + const companyId = createTestCompany(db); + const wsId = createTestWorkspace(db, companyId); + const agentId = createTestAgent(db, companyId, wsId, "Solo"); + + // Only one heartbeat task + createHeartbeatTask(db, companyId, agentId, "done", "Solo"); + + const cleaned = db.cleanupDuplicateHeartbeatTasks(companyId); + expect(cleaned).toBe(0); + }); + }); +});