diff --git a/server/src/services/agent-cost-history-deletion.test.ts b/server/src/services/agent-cost-history-deletion.test.ts new file mode 100644 index 00000000000..923298a48f6 --- /dev/null +++ b/server/src/services/agent-cost-history-deletion.test.ts @@ -0,0 +1,61 @@ +import { randomUUID } from "node:crypto"; +import { eq } from "drizzle-orm"; +import { agents, companies, costEvents, createDb, financeEvents, heartbeatRuns } from "@paperclipai/db"; +import { afterAll, beforeAll, describe, expect, it } from "vitest"; +import { startEmbeddedPostgresTestDatabase } from "../__tests__/helpers/embedded-postgres.js"; +import { createAgentLifecycle } from "./agent-lifecycle.js"; + +describe("agent purge with retained cost history", () => { + let temporary: Awaited>; + let db: ReturnType; + beforeAll(async () => { + temporary = await startEmbeddedPostgresTestDatabase("agent-cost-deletion-"); + db = createDb(temporary.connectionString); + }); + afterAll(async () => { await temporary?.cleanup(); }); + + async function fixture() { + const companyId = randomUUID(), agentId = randomUUID(); + await db.insert(companies).values({ id: companyId, name: "Cost deletion", issuePrefix: `C${companyId.slice(0, 6)}` }); + await db.insert(agents).values({ id: agentId, companyId, name: "Billed worker", status: "terminated", lifecycleState: "terminated" }); + const [run] = await db.insert(heartbeatRuns).values({ companyId, agentId, status: "succeeded" }).returning(); + return { companyId, agentId, run: run! }; + } + const cost = (companyId: string, values: Partial) => db.insert(costEvents).values({ + companyId, usageKind: "decision", provider: "anthropic", model: "claude", costCents: 1, occurredAt: new Date(), ...values, + }).returning(); + + async function expectRetained(f: Awaited>) { + await expect(createAgentLifecycle(db).purgeAgent(f.agentId)).rejects.toMatchObject({ status: 409, + details: { code: "agent_cost_history_retained" } }); + expect(await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, f.run.id))).toHaveLength(1); + expect(await db.select().from(agents).where(eq(agents.id, f.agentId))).toMatchObject([{ lifecycleState: "terminated" }]); + } + + it("refuses before cleanup when a cost event references one of the agent's runs", async () => { + const f = await fixture(); + const [event] = await cost(f.companyId, { heartbeatRunId: f.run.id }); + await expectRetained(f); + expect(await db.select().from(costEvents).where(eq(costEvents.id, event!.id))).toEqual([event]); + }); + + it("refuses when a cost event is attributed to the agent", async () => { + const f = await fixture(); + await cost(f.companyId, { usageKind: "agent", agentId: f.agentId }); + await expectRetained(f); + }); + + it("refuses when a finance event references one of the agent's runs", async () => { + const f = await fixture(); + const [event] = await db.insert(financeEvents).values({ companyId: f.companyId, heartbeatRunId: f.run.id, + eventKind: "usage", biller: "anthropic", amountCents: 1, occurredAt: new Date() }).returning(); + await expectRetained(f); + expect(await db.select().from(financeEvents).where(eq(financeEvents.id, event!.id))).toEqual([event]); + }); + + it("still purges an agent whose runs have no cost history", async () => { + const f = await fixture(); + expect(await createAgentLifecycle(db).purgeAgent(f.agentId)).toMatchObject({ id: f.agentId }); + expect(await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, f.run.id))).toHaveLength(0); + }); +}); diff --git a/server/src/services/agent-deletion.ts b/server/src/services/agent-deletion.ts index b4252d3b9a1..a4ae5b16aa7 100644 --- a/server/src/services/agent-deletion.ts +++ b/server/src/services/agent-deletion.ts @@ -10,7 +10,7 @@ import { notifyInstructionCommit } from "./instruction-commit-notifications.js"; import { notifyScheduledRetryCancellation } from "./scheduled-retry-cancellation-notifications.js"; import { and, eq, or, sql } from "drizzle-orm"; import { agents, agentApiKeys, agentRuntimeState, agentTaskSessions, agentWakeupRequests, activityLog, - budgetReservations, heartbeatRunEvents, heartbeatRuns, issueExecutionDecisions, issues, issueComments, + budgetReservations, costEvents, financeEvents, heartbeatRunEvents, heartbeatRuns, issueExecutionDecisions, issues, issueComments, caseHandoffs, principalPermissionGrants, type Db } from "@paperclipai/db"; import { conflict } from "../errors.js"; import { issueThreadInteractionService } from "./issue-thread-interactions.js"; @@ -47,6 +47,16 @@ export async function deleteAgentDependencies(tx: Db, companyId: string, id: str if (fastResponseHold) throw conflict("Wait for active fast responses or resolve their unknown charges in Costs before deleting this agent", { code: "agent_fast_response_accounting_pending", }); + const agentRuns = sql`select ${heartbeatRuns.id} from ${heartbeatRuns} where ${heartbeatRuns.agentId} = ${id}`; + const [retainedCost] = await tx.select({ id: costEvents.id }).from(costEvents).where(and( + eq(costEvents.companyId, companyId), or(eq(costEvents.agentId, id), sql`${costEvents.heartbeatRunId} in (${agentRuns})`), + )).limit(1); + const [retainedFinance] = await tx.select({ id: financeEvents.id }).from(financeEvents).where(and( + eq(financeEvents.companyId, companyId), or(eq(financeEvents.agentId, id), sql`${financeEvents.heartbeatRunId} in (${agentRuns})`), + )).limit(1); + if (retainedCost || retainedFinance) throw conflict("Cost history requires this agent record. Keep the agent terminated instead of permanently deleting it", { + code: "agent_cost_history_retained", + }); await notifyAttentionRetentionWork(tx, { companyId }); await notifyExternalConversationSettlement(tx, { companyId }); await issueThreadInteractionService(tx).cancelPendingForDeletedAddressee(companyId, id); @@ -63,7 +73,7 @@ export async function deleteAgentDependencies(tx: Db, companyId: string, id: str await tx.delete(activityLog).where( or( eq(activityLog.agentId, id), - sql`${activityLog.runId} in (select ${heartbeatRuns.id} from ${heartbeatRuns} where ${heartbeatRuns.agentId} = ${id})`, + sql`${activityLog.runId} in (${agentRuns})`, ), ); await tx.delete(issueExecutionDecisions).where(eq(issueExecutionDecisions.actorAgentId, id));