import { randomUUID } from "node:crypto"; import { afterAll, afterEach, beforeAll, describe, expect, it } from "vitest"; import { eq } from "drizzle-orm"; import { agents, companies, createDb, environmentLeases, environments, heartbeatRuns } from "@paperclipai/db"; import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase, } from "./helpers/embedded-postgres.js"; import { environmentService } from "../services/environments.ts"; const embeddedPostgresSupport = await getEmbeddedPostgresTestSupport(); const describeEmbeddedPostgres = embeddedPostgresSupport.supported ? describe : describe.skip; if (!embeddedPostgresSupport.supported) { console.warn( `Skipping embedded Postgres environment service tests on this host: ${embeddedPostgresSupport.reason ?? "unsupported environment"}`, ); } describeEmbeddedPostgres("environmentService leases", () => { let stopDb: (() => Promise) | null = null; let db!: ReturnType; let svc!: ReturnType; beforeAll(async () => { const started = await startEmbeddedPostgresTestDatabase("environment-service"); stopDb = started.stop; db = createDb(started.connectionString); svc = environmentService(db); }); afterEach(async () => { await db.delete(environmentLeases); await db.delete(heartbeatRuns); await db.delete(agents); await db.delete(environments); await db.delete(companies); }); afterAll(async () => { await stopDb?.(); }); async function seedEnvironment() { const companyId = randomUUID(); const agentId = randomUUID(); const environmentId = randomUUID(); const runId = randomUUID(); await db.insert(companies).values({ id: companyId, name: "Acme", status: "active", createdAt: new Date(), updatedAt: new Date(), }); await db.insert(agents).values({ id: agentId, companyId, name: "CodexCoder", role: "engineer", status: "active", adapterType: "codex_local", adapterConfig: {}, runtimeConfig: {}, permissions: {}, createdAt: new Date(), updatedAt: new Date(), }); await db.insert(environments).values({ id: environmentId, companyId, name: "Local", driver: "local", status: "active", config: {}, createdAt: new Date(), updatedAt: new Date(), }); await db.insert(heartbeatRuns).values({ id: runId, companyId, agentId, invocationSource: "manual", status: "running", createdAt: new Date(), updatedAt: new Date(), }); return { companyId, agentId, environmentId, runId }; } it("acquires and releases a lease for a run", async () => { const { companyId, environmentId, runId } = await seedEnvironment(); const lease = await svc.acquireLease({ companyId, environmentId, heartbeatRunId: runId, metadata: { driver: "local" }, }); expect(lease.status).toBe("active"); expect(lease.heartbeatRunId).toBe(runId); const released = await svc.releaseLease(lease.id); expect(released?.status).toBe("released"); expect(released?.releasedAt).not.toBeNull(); }); it("releases all active leases for a run without touching unrelated rows", async () => { const { companyId, agentId, environmentId, runId } = await seedEnvironment(); const otherRunId = randomUUID(); await db.insert(heartbeatRuns).values({ id: otherRunId, companyId, agentId, invocationSource: "manual", status: "running", createdAt: new Date(), updatedAt: new Date(), }); const targetLease = await svc.acquireLease({ companyId, environmentId, heartbeatRunId: runId, }); const otherLease = await svc.acquireLease({ companyId, environmentId, heartbeatRunId: otherRunId, }); const released = await svc.releaseLeasesForRun(runId); expect(released.map((lease) => lease.id)).toEqual([targetLease.id]); const stillActive = await svc.listLeases(environmentId, { status: "active" }); expect(stillActive.map((lease) => lease.id)).toEqual([otherLease.id]); }); it("creates and then reuses the default local environment for a company", async () => { const companyId = randomUUID(); await db.insert(companies).values({ id: companyId, name: "Acme", status: "active", createdAt: new Date(), updatedAt: new Date(), }); const created = await svc.ensureLocalEnvironment(companyId); const reused = await svc.ensureLocalEnvironment(companyId); expect(created.driver).toBe("local"); expect(reused.id).toBe(created.id); const rows = await db.select().from(environments).where(eq(environments.companyId, companyId)); expect(rows).toHaveLength(1); expect(rows[0]?.name).toBe("Local"); }); it("leaves an existing default local environment untouched", async () => { const companyId = randomUUID(); await db.insert(companies).values({ id: companyId, name: "Acme", status: "active", createdAt: new Date(), updatedAt: new Date(), }); const archivedAt = new Date("2025-01-01T00:00:00.000Z"); const [existing] = await db .insert(environments) .values({ companyId, name: "Archived Local", description: "Operator-managed local environment", driver: "local", status: "archived", config: { shell: "zsh" }, metadata: { owner: "operator" }, createdAt: archivedAt, updatedAt: archivedAt, }) .returning(); const ensured = await svc.ensureLocalEnvironment(companyId); expect(ensured.id).toBe(existing?.id); expect(ensured.name).toBe("Archived Local"); expect(ensured.status).toBe("archived"); expect(ensured.metadata).toEqual({ owner: "operator" }); const rows = await db.select().from(environments).where(eq(environments.companyId, companyId)); expect(rows).toHaveLength(1); expect(rows[0]?.updatedAt.toISOString()).toBe(archivedAt.toISOString()); }); it("deduplicates concurrent default local environment creation", async () => { const companyId = randomUUID(); await db.insert(companies).values({ id: companyId, name: "Acme", status: "active", createdAt: new Date(), updatedAt: new Date(), }); const results = await Promise.all( Array.from({ length: 8 }, () => svc.ensureLocalEnvironment(companyId)), ); expect(new Set(results.map((environment) => environment.id)).size).toBe(1); const rows = await db.select().from(environments).where(eq(environments.companyId, companyId)); expect(rows).toHaveLength(1); expect(rows[0]?.driver).toBe("local"); expect(rows[0]?.status).toBe("active"); }); });