diff --git a/server/src/__tests__/cleanup-removal-service.test.ts b/server/src/__tests__/cleanup-removal-service.test.ts index 25dc11d679b..e65124cc9d3 100644 --- a/server/src/__tests__/cleanup-removal-service.test.ts +++ b/server/src/__tests__/cleanup-removal-service.test.ts @@ -9,6 +9,7 @@ import { createDb, documents, documentRevisions, + heartbeatRunEvents, heartbeatRuns, issueComments, issueDocuments, @@ -42,6 +43,7 @@ describeEmbeddedPostgres("cleanup removal services", () => { }, 20_000); afterEach(async () => { + await db.delete(heartbeatRunEvents); await db.delete(activityLog); await db.delete(issueReadStates); await db.delete(issueComments); @@ -228,4 +230,32 @@ describeEmbeddedPostgres("cleanup removal services", () => { await expect(db.select().from(issueReadStates).where(eq(issueReadStates.companyId, companyId))).resolves.toHaveLength(0); await expect(db.select().from(activityLog).where(eq(activityLog.companyId, companyId))).resolves.toHaveLength(0); }); + + it("removes heartbeat events by run id before deleting company-owned runs", async () => { + const { agentId, companyId, runId } = await seedFixture(); + const otherCompanyId = randomUUID(); + + await db.insert(companies).values({ + id: otherCompanyId, + name: "Other Company", + issuePrefix: `O${otherCompanyId.replace(/-/g, "").slice(0, 6).toUpperCase()}`, + requireBoardApprovalForNewAgents: false, + }); + + await db.insert(heartbeatRunEvents).values({ + companyId: otherCompanyId, + runId, + agentId, + seq: 1, + eventType: "output", + message: "event with mismatched company scope", + }); + + const removed = await companyService(db).remove(companyId); + + expect(removed?.id).toBe(companyId); + await expect(db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, runId))).resolves.toHaveLength(0); + await expect(db.select().from(heartbeatRunEvents).where(eq(heartbeatRunEvents.runId, runId))).resolves.toHaveLength(0); + await expect(db.select().from(companies).where(eq(companies.id, otherCompanyId))).resolves.toHaveLength(1); + }); }); diff --git a/server/src/services/companies.ts b/server/src/services/companies.ts index 17905f11ee0..8704b986cd2 100644 --- a/server/src/services/companies.ts +++ b/server/src/services/companies.ts @@ -265,7 +265,17 @@ export function companyService(db: Db) { remove: (id: string) => db.transaction(async (tx) => { // Delete from child tables in dependency order + const companyRunIds = await tx + .select({ id: heartbeatRuns.id }) + .from(heartbeatRuns) + .where(eq(heartbeatRuns.companyId, id)); + await tx.delete(heartbeatRunEvents).where(eq(heartbeatRunEvents.companyId, id)); + if (companyRunIds.length > 0) { + await tx + .delete(heartbeatRunEvents) + .where(inArray(heartbeatRunEvents.runId, companyRunIds.map((run) => run.id))); + } await tx.delete(agentTaskSessions).where(eq(agentTaskSessions.companyId, id)); await tx.delete(activityLog).where(eq(activityLog.companyId, id)); await tx.delete(heartbeatRuns).where(eq(heartbeatRuns.companyId, id));