|
| 1 | +import { |
| 2 | + createChat, |
| 3 | + createDashboardAgentDb, |
| 4 | + getInvestigation, |
| 5 | + settleInvestigationAndCloseCard, |
| 6 | + upsertInvestigationRevision, |
| 7 | + type DashboardAgentDb, |
| 8 | + type DashboardAgentDbClient, |
| 9 | +} from "@internal/dashboard-agent-db"; |
| 10 | +import { |
| 11 | + investigationStateSchema, |
| 12 | + type InvestigationState, |
| 13 | +} from "@internal/dashboard-agent-contracts"; |
| 14 | +import { postgresTest } from "@internal/testcontainers"; |
| 15 | +import type { PrismaClient } from "@trigger.dev/database"; |
| 16 | +import { readdirSync, readFileSync } from "node:fs"; |
| 17 | +import path from "node:path"; |
| 18 | +import { afterEach, describe, expect, vi } from "vitest"; |
| 19 | + |
| 20 | +const ctx = vi.hoisted(() => ({ |
| 21 | + agentDb: undefined as unknown as DashboardAgentDb, |
| 22 | +})); |
| 23 | + |
| 24 | +vi.mock("~/services/dashboardAgentDb.server", () => ({ |
| 25 | + get dashboardAgentDb() { |
| 26 | + return ctx.agentDb; |
| 27 | + }, |
| 28 | +})); |
| 29 | + |
| 30 | +const { sweepDashboardAgentInvestigations, INVESTIGATION_STALE_MS, MAX_SWEEP_ATTEMPTS } = |
| 31 | + await import("~/services/dashboardAgentInvestigationSweep.server"); |
| 32 | + |
| 33 | +async function applyAgentSchema(prisma: PrismaClient) { |
| 34 | + const folder = path.resolve(__dirname, "../../../internal-packages/dashboard-agent-db/drizzle"); |
| 35 | + const migrations = readdirSync(folder) |
| 36 | + .filter((file) => file.endsWith(".sql")) |
| 37 | + .sort(); |
| 38 | + for (const name of migrations) { |
| 39 | + const sql = readFileSync(path.join(folder, name), "utf8"); |
| 40 | + for (const statement of sql.split("--> statement-breakpoint")) { |
| 41 | + const trimmed = statement.trim(); |
| 42 | + if (trimmed.length > 0) await prisma.$executeRawUnsafe(trimmed); |
| 43 | + } |
| 44 | + } |
| 45 | +} |
| 46 | + |
| 47 | +let agentDbClient: DashboardAgentDbClient | undefined; |
| 48 | +let prismaForRaw: PrismaClient | undefined; |
| 49 | + |
| 50 | +async function boot(prisma: PrismaClient, connectionUri: string) { |
| 51 | + await applyAgentSchema(prisma); |
| 52 | + agentDbClient = createDashboardAgentDb(connectionUri, { max: 2 }); |
| 53 | + ctx.agentDb = agentDbClient.db; |
| 54 | + prismaForRaw = prisma; |
| 55 | +} |
| 56 | + |
| 57 | +afterEach(async () => { |
| 58 | + await agentDbClient?.close(); |
| 59 | + agentDbClient = undefined; |
| 60 | +}); |
| 61 | + |
| 62 | +const ORG = "org_poison"; |
| 63 | +const USER = "user_poison"; |
| 64 | + |
| 65 | +function openState(): InvestigationState { |
| 66 | + return investigationStateSchema.parse({ |
| 67 | + outcome: "in_progress", |
| 68 | + severity: "warn", |
| 69 | + confidence: "medium", |
| 70 | + title: "a stuck card", |
| 71 | + headline: "Still checking.", |
| 72 | + progress: "Reading spans", |
| 73 | + checkNext: [], |
| 74 | + hypotheses: [], |
| 75 | + evidence: [], |
| 76 | + }); |
| 77 | +} |
| 78 | + |
| 79 | +async function seedInvestigation(chatId: string, ageMs: number): Promise<string> { |
| 80 | + await createChat(ctx.agentDb, { id: chatId, organizationId: ORG, userId: USER }); |
| 81 | + const created = await upsertInvestigationRevision(ctx.agentDb, { |
| 82 | + chatId, |
| 83 | + projectRef: "proj", |
| 84 | + environmentRef: "env", |
| 85 | + state: openState(), |
| 86 | + }); |
| 87 | + if (!created.ok) throw new Error("fixture investigation not created"); |
| 88 | + await prismaForRaw!.$executeRawUnsafe( |
| 89 | + `update trigger_dashboard_agent.investigations |
| 90 | + set updated_at = now() - ($2 || ' milliseconds')::interval where id = $1`, |
| 91 | + created.id, |
| 92 | + String(ageMs) |
| 93 | + ); |
| 94 | + return created.id; |
| 95 | +} |
| 96 | + |
| 97 | +async function outcomeOf(id: string): Promise<string | undefined> { |
| 98 | + const row = await getInvestigation(ctx.agentDb, { id }); |
| 99 | + return row ? (row.state as { outcome?: string }).outcome : undefined; |
| 100 | +} |
| 101 | + |
| 102 | +const STALE_AGE_MS = INVESTIGATION_STALE_MS + 60_000; |
| 103 | +const OLDER_AGE_MS = STALE_AGE_MS + 60_000; |
| 104 | + |
| 105 | +describe("the investigation sweep with a poison row", () => { |
| 106 | + postgresTest( |
| 107 | + "a row that always fails to settle cannot pin the head and starve a newer row", |
| 108 | + async ({ prisma, postgresContainer }) => { |
| 109 | + await boot(prisma, postgresContainer.getConnectionUri()); |
| 110 | + |
| 111 | + // Poison sorts first (older `updated_at`); renderable is newer. |
| 112 | + const poisonId = await seedInvestigation("chat_poison", OLDER_AGE_MS); |
| 113 | + const renderableId = await seedInvestigation("chat_ok", STALE_AGE_MS); |
| 114 | + |
| 115 | + // Only the poison row's settle throws; the renderable one goes through the real path. |
| 116 | + const settleAndClose = (params: { id: string; chatId: string; note: string }) => { |
| 117 | + if (params.id === poisonId) throw new Error("state isn't renderable"); |
| 118 | + return settleInvestigationAndCloseCard(ctx.agentDb, params); |
| 119 | + }; |
| 120 | + |
| 121 | + // limit 1 forces head contention: without backoff the poison row would win every run. |
| 122 | + // A failed run throws so the job retries, but the attempt is recorded before it does. |
| 123 | + await expect( |
| 124 | + sweepDashboardAgentInvestigations({ limit: 1, settleAndClose }) |
| 125 | + ).rejects.toThrow(); |
| 126 | + expect(await outcomeOf(poisonId)).toBe("in_progress"); |
| 127 | + expect(await outcomeOf(renderableId)).toBe("in_progress"); |
| 128 | + |
| 129 | + // Next run: the poison row now sorts behind the never-attempted renderable one, |
| 130 | + // so the newer row is picked and settled despite the poison row still being stale. |
| 131 | + const second = await sweepDashboardAgentInvestigations({ limit: 1, settleAndClose }); |
| 132 | + expect(second).toMatchObject({ stale: 1, settled: 1, failed: 0 }); |
| 133 | + expect(await outcomeOf(renderableId)).toBe("inconclusive"); |
| 134 | + expect(await outcomeOf(poisonId)).toBe("in_progress"); |
| 135 | + }, |
| 136 | + 30_000 |
| 137 | + ); |
| 138 | + |
| 139 | + postgresTest( |
| 140 | + "after the attempt cap the poison row is abandoned and leaves the queue", |
| 141 | + async ({ prisma, postgresContainer }) => { |
| 142 | + await boot(prisma, postgresContainer.getConnectionUri()); |
| 143 | + const poisonId = await seedInvestigation("chat_poison", STALE_AGE_MS); |
| 144 | + |
| 145 | + const settleAndClose = () => { |
| 146 | + throw new Error("state isn't renderable"); |
| 147 | + }; |
| 148 | + |
| 149 | + // The first MAX_SWEEP_ATTEMPTS-1 runs record a failed attempt and throw; the row stays stale. |
| 150 | + for (let i = 1; i < MAX_SWEEP_ATTEMPTS; i++) { |
| 151 | + await expect(sweepDashboardAgentInvestigations({ settleAndClose })).rejects.toThrow(); |
| 152 | + expect(await outcomeOf(poisonId)).toBe("in_progress"); |
| 153 | + } |
| 154 | + |
| 155 | + // The capped run force-settles the row without the render path, so it leaves the queue. |
| 156 | + const capped = await sweepDashboardAgentInvestigations({ settleAndClose }); |
| 157 | + expect(capped).toMatchObject({ stale: 1, abandoned: 1, failed: 0 }); |
| 158 | + expect(await outcomeOf(poisonId)).toBe("inconclusive"); |
| 159 | + |
| 160 | + // Nothing stale remains, so the poison row is no longer swept. |
| 161 | + const after = await sweepDashboardAgentInvestigations({ settleAndClose }); |
| 162 | + expect(after).toMatchObject({ stale: 0 }); |
| 163 | + }, |
| 164 | + 30_000 |
| 165 | + ); |
| 166 | +}); |
0 commit comments