evOo-local/packages/adapters/src/artifact-files.test.ts

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<PrismaClient, "artifact">;
  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, "<script></script>");
    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([]);
  });
});