import { createHash } from "node:crypto"; import { mkdtemp, readdir, rm, symlink, truncate, writeFile } from "node:fs/promises"; import { tmpdir } from "node:os"; import path from "node:path"; import { ArtifactTooLargeError } from "@rakazo/adapter-kit"; import { ARTIFACT_MAX_BYTES_BY_KIND } from "@rakazo/contracts"; import { AttachmentValidationError } from "@rakazo/core"; import type { PrismaClient } from "@rakazo/db"; import { afterEach, describe, expect, it, vi } from "vitest"; import { storeArtifactFromFile } from "./artifact-files.js"; import { LocalArtifactStore } from "./artifacts.js"; const dirs: string[] = []; afterEach(async () => { await Promise.all(dirs.splice(0).map((dir) => rm(dir, { recursive: true, force: true }))); }); async function setup( create = vi.fn(async ({ data }: { data: object }) => ({ id: "art-1", ...data })), ) { const root = await mkdtemp(path.join(tmpdir(), "rakazo-artifact-files-")); dirs.push(root); const store = new LocalArtifactStore(path.join(root, "data")); const prisma = { artifact: { create } } as unknown as Pick; return { root, store, prisma, create }; } const owner = { spaceId: "space-1", userId: "user-1", botId: "bot-1", groupId: null, runId: "run-1", }; describe("storeArtifactFromFile", () => { it("streams the file into the store and records size, hash and media metadata", async () => { const { root, store, prisma, create } = await setup(); const file = path.join(root, "out.mp4"); await writeFile(file, "not really a video"); const row = await storeArtifactFromFile( { prisma, artifacts: store }, { ...owner, path: file, name: "out.mp4", mimeType: "video/mp4", width: 1080, height: 1920, durationMs: 1500, parentArtifactId: "art-0", }, ); expect(create).toHaveBeenCalledWith({ data: expect.objectContaining({ ...owner, name: "out.mp4", mimeType: "video/mp4", size: 18, hash: createHash("sha256").update("not really a video").digest("hex"), width: 1080, height: 1920, durationMs: 1500, parentArtifactId: "art-0", }), }); const key = (row as unknown as { storageKey: string }).storageKey; const context = { spaceId: "space-1", signal: new AbortController().signal } as never; expect(Buffer.from(await store.get(key, context)).toString()).toBe("not really a video"); }); it("rejects disallowed types, oversized files and symlinks before storing anything", async () => { const { root, store, prisma, create } = await setup(); const html = path.join(root, "page.html"); await writeFile(html, ""); await expect( storeArtifactFromFile( { prisma, artifacts: store }, { ...owner, path: html, name: "page.html", mimeType: "text/html" }, ), ).rejects.toBeInstanceOf(AttachmentValidationError); const huge = path.join(root, "huge.png"); await writeFile(huge, ""); await truncate(huge, ARTIFACT_MAX_BYTES_BY_KIND.image + 1); await expect( storeArtifactFromFile( { prisma, artifacts: store }, { ...owner, path: huge, name: "huge.png", mimeType: "image/png" }, ), ).rejects.toBeInstanceOf(ArtifactTooLargeError); const link = path.join(root, "link.txt"); await writeFile(path.join(root, "secret.txt"), "secret"); await symlink(path.join(root, "secret.txt"), link); await expect( storeArtifactFromFile( { prisma, artifacts: store }, { ...owner, path: link, name: "link.txt", mimeType: "text/plain" }, ), ).rejects.toThrow(); expect(create).not.toHaveBeenCalled(); }); it("removes the stored blob when the artifact row cannot be written", async () => { const { root, store, prisma } = await setup( vi.fn(async () => { throw new Error("db down"); }), ); const file = path.join(root, "notes.txt"); await writeFile(file, "hello"); await expect( storeArtifactFromFile( { prisma, artifacts: store }, { ...owner, path: file, name: "notes.txt", mimeType: "text/plain" }, ), ).rejects.toThrow("db down"); expect(await readdir(path.join(root, "data", "artifacts", "space-1"))).toEqual([]); }); });