evOo-local/packages/adapters/src/thread-artifacts.test.ts

import { readFileSync } from "node:fs";
import {
  mkdir,
  mkdtemp,
  readdir,
  readFile,
  realpath,
  rm,
  stat,
  symlink,
  writeFile,
} from "node:fs/promises";
import { tmpdir } from "node:os";
import path from "node:path";
import type {
  AdapterContext,
  ArtifactStore,
  ArtifactStreamPut,
  ComputerRef,
  SandboxProvider,
} from "@rakazo/adapter-kit";
import { ArtifactTooLargeError } from "@rakazo/adapter-kit";
import type { MessageBlock } from "@rakazo/contracts";
import type { PrismaClient } from "@rakazo/db";
import { afterEach, describe, expect, it, vi } from "vitest";
import { byteChunks } from "./byte-streams.js";
import { readComputerFileStream } from "./computer-support.js";
import { DesktopSandboxProvider } from "./desktop-sandbox.js";
import { FakeSandboxProvider } from "./fake-sandbox.js";
import { HostAwareSandbox } from "./host-aware-sandbox.js";
import {
  attachableFileMaxBytes,
  attachWorkspaceFileToThread,
  currentTurnFilesInstruction,
  materializeCurrentTurnFiles,
} from "./thread-artifacts.js";

function byteStream(chunks: Uint8Array[]) {
  return new ReadableStream<Uint8Array>({
    start(controller) {
      for (const chunk of chunks) controller.enqueue(chunk);
      controller.close();
    },
  });
}

describe("current-turn thread files", () => {
  it("removes stored bytes when artifact metadata cannot be created", async () => {
    const remove = vi.fn().mockResolvedValue(undefined);
    const failure = new Error("database unavailable");

    await expect(
      attachWorkspaceFileToThread(
        {
          prisma: {
            artifact: { create: vi.fn().mockRejectedValue(failure) },
          } as unknown as PrismaClient,
          artifacts: {
            describe: () => ({ capabilities: { stream: true } }),
            putStream: vi.fn().mockResolvedValue({ key: "stored-1", size: 3, sha256: "hash" }),
            remove,
          } as unknown as ArtifactStore,
        },
        {
          spaceId: "workspace-1",
          userId: "user-1",
          botId: "bot-1",
          runId: "run-1",
          filePath: "report.pdf",
          content: new Uint8Array([1, 2, 3]),
          operationId: "attach-1",
        },
      ),
    ).rejects.toBe(failure);
    expect(remove).toHaveBeenCalledWith(
      "stored-1",
      expect.objectContaining({ spaceId: "workspace-1", botId: "bot-1" }),
    );
  });

  it("marks a file the bot presents as its result", async () => {
    const deps = {
      prisma: {
        artifact: {
          create: vi.fn(
            async ({ data }: { data: { name: string; mimeType: string; size: number } }) => ({
              id: "artifact-1",
              ...data,
            }),
          ),
        },
      } as unknown as PrismaClient,
      artifacts: {
        describe: () => ({ capabilities: { stream: true } }),
        putStream: vi.fn().mockResolvedValue({ key: "stored-1", size: 9, sha256: "hash" }),
        remove: vi.fn(),
      } as unknown as ArtifactStore,
    };
    const input = {
      spaceId: "workspace-1",
      userId: "user-1",
      botId: "bot-1",
      runId: "run-1",
      filePath: "out/report.md",
      content: new TextEncoder().encode("# Report\n"),
      operationId: "attach-1",
    };

    const presented = await attachWorkspaceFileToThread(deps, { ...input, present: true });
    expect(presented.block).toMatchObject({ kind: "file", name: "report.md", present: true });
    const attached = await attachWorkspaceFileToThread(deps, input);
    expect(attached.block).not.toHaveProperty("present");
  });

  it("copies non-image attachments into the bot workspace and describes their paths", async () => {
    const findMany = vi.fn().mockResolvedValue([
      {
        id: "artifact-1",
        spaceId: "workspace-1",
        botId: "bot-1",
        name: "../quarterly report.pdf",
        mimeType: "application/pdf",
        size: 4,
        storageKey: "stored-1",
      },
    ]);
    const createReadStream = vi.fn(async () =>
      byteStream([new Uint8Array([1, 2]), new Uint8Array([3, 4])]),
    );
    const writeFile = vi.fn().mockResolvedValue(undefined);
    const markWorkspaceDirty = vi.fn();
    const context: AdapterContext & { botId: string } = {
      operationId: "run-1",
      traceId: "run-1",
      spaceId: "workspace-1",
      userId: "user-1",
      botId: "bot-1",
      runId: "run-1",
      signal: new AbortController().signal,
    };
    const computer: ComputerRef = {
      id: "computer-1",
      botId: "bot-1",
      kind: "fake",
      providerRef: "fake-1",
    };
    const blocks: MessageBlock[] = [
      {
        kind: "file",
        artifactId: "artifact-1",
        name: "../quarterly report.pdf",
        mimeType: "application/pdf",
        size: 4,
      },
    ];

    const files = await materializeCurrentTurnFiles(
      {
        prisma: { artifact: { findMany } } as unknown as PrismaClient,
        artifacts: { createReadStream } as unknown as ArtifactStore,
        sandbox: { writeFile } as unknown as SandboxProvider,
      },
      blocks,
      { context, computer, computerMode: "team", markWorkspaceDirty },
    );

    // A provider without streaming writes still receives the whole file as bytes.
    expect(createReadStream).toHaveBeenCalledWith("stored-1", context);
    expect(writeFile).toHaveBeenCalledWith(
      computer,
      {
        path: "bots/bot-1/attachments/artifact-1.pdf",
        content: new Uint8Array([1, 2, 3, 4]),
      },
      context,
    );
    expect(markWorkspaceDirty).toHaveBeenCalledOnce();
    expect(markWorkspaceDirty.mock.invocationCallOrder[0]).toBeLessThan(
      writeFile.mock.invocationCallOrder[0]!,
    );
    expect(files).toEqual([
      {
        artifactId: "artifact-1",
        name: "../quarterly report.pdf",
        mimeType: "application/pdf",
        size: 4,
        path: "attachments/artifact-1.pdf",
      },
    ]);
    expect(currentTurnFilesInstruction(files)).toContain('"attachments/artifact-1.pdf"');
  });

  it("copies images, video and audio too, and keeps files over a remote computer's cap in storage", async () => {
    const mb = 1024 * 1024;
    const findMany = vi.fn().mockResolvedValue([
      { id: "img-1", name: "photo.png", mimeType: "image/png", size: 3, storageKey: "k-img" },
      { id: "aud-1", name: "memo.mp3", mimeType: "audio/mpeg", size: 3, storageKey: "k-aud" },
      { id: "vid-1", name: "clip.mp4", mimeType: "video/mp4", size: 40 * mb, storageKey: "k-vid" },
    ]);
    const createReadStream = vi.fn(async () => byteStream([new Uint8Array([1, 2, 3])]));
    const writeFile = vi.fn().mockResolvedValue(undefined);
    const context: AdapterContext & { botId: string } = {
      operationId: "run-1",
      traceId: "run-1",
      spaceId: "workspace-1",
      userId: "user-1",
      botId: "bot-1",
      signal: new AbortController().signal,
    };
    const blocks: MessageBlock[] = [
      { kind: "image", artifactId: "img-1", name: "photo.png", mimeType: "image/png" },
      { kind: "audio", artifactId: "aud-1", name: "memo.mp3", mimeType: "audio/mpeg", size: 3 },
      {
        kind: "video",
        artifactId: "vid-1",
        name: "clip.mp4",
        mimeType: "video/mp4",
        size: 40 * mb,
      },
      { kind: "text", text: "see these" },
    ];
    const deps = {
      prisma: { artifact: { findMany } } as unknown as PrismaClient,
      artifacts: { createReadStream } as unknown as ArtifactStore,
      sandbox: { writeFile } as unknown as SandboxProvider,
    };

    const remote = await materializeCurrentTurnFiles(deps, blocks, {
      context,
      computer: { id: "computer-1", botId: "bot-1", kind: "docker", providerRef: "docker-1" },
      computerMode: "dedicated",
    });
    expect(remote.map((file) => [file.artifactId, file.path])).toEqual([
      ["img-1", "attachments/img-1.png"],
      ["aud-1", "attachments/aud-1.mp3"],
      ["vid-1", null],
    ]);
    expect(createReadStream).not.toHaveBeenCalledWith("k-vid", context);
    const instruction = currentTurnFilesInstruction(remote);
    expect(instruction).toContain('"attachments/img-1.png"');
    expect(instruction).toContain("too large to copy to your computer");
    expect(instruction).toContain('"clip.mp4" (video/mp4, 41943040 bytes): artifact vid-1');

    writeFile.mockClear();
    const local = await materializeCurrentTurnFiles(deps, blocks, {
      context,
      computer: { id: "computer-2", botId: "bot-1", kind: "desktop", providerRef: "desktop-1" },
      computerMode: "dedicated",
    });
    expect(local.map((file) => file.path)).toEqual([
      "attachments/img-1.png",
      "attachments/aud-1.mp3",
      "attachments/vid-1.mp4",
    ]);
    expect(writeFile).toHaveBeenCalledTimes(3);
    expect(currentTurnFilesInstruction(local)).not.toContain("too large");
  });

  it("skips loading when a turn has no attachments", async () => {
    const findMany = vi.fn();
    const files = await materializeCurrentTurnFiles(
      {
        prisma: { artifact: { findMany } } as unknown as PrismaClient,
        artifacts: {} as ArtifactStore,
        sandbox: {} as SandboxProvider,
      },
      [{ kind: "text", text: "hello" }],
      {
        context: {
          operationId: "run-1",
          traceId: "run-1",
          spaceId: "workspace-1",
          userId: "user-1",
          botId: "bot-1",
          signal: new AbortController().signal,
        },
        computer: {
          id: "computer-1",
          botId: "bot-1",
          kind: "fake",
          providerRef: "fake-1",
        },
        computerMode: "team",
      },
    );

    expect(files).toEqual([]);
    expect(findMany).not.toHaveBeenCalled();
  });

  it("caps attachments by kind when the store streams and at the inline cap otherwise", () => {
    const streaming = { describe: () => ({ capabilities: { stream: true } }) } as ArtifactStore;
    const buffered = { describe: () => ({ capabilities: { stream: false } }) } as ArtifactStore;
    expect(attachableFileMaxBytes(streaming, "video/mp4")).toBe(500 * 1024 * 1024);
    expect(attachableFileMaxBytes(streaming, "text/plain")).toBe(25 * 1024 * 1024);
    expect(attachableFileMaxBytes(buffered, "video/mp4")).toBe(10 * 1024 * 1024);
  });
});

describe("streamed media on a local computer", () => {
  const CHUNK = 64 * 1024;
  /** 8 MiB, far above the one chunk the copy may hold at a time. */
  const CHUNKS = 128;
  const roots: string[] = [];
  const context: AdapterContext & { botId: string } = {
    operationId: "run-1",
    traceId: "run-1",
    spaceId: "workspace-1",
    userId: "user-1",
    botId: "bot-1",
    signal: new AbortController().signal,
  };
  const videoBlock: MessageBlock = {
    kind: "video",
    artifactId: "vid-1",
    name: "clip.mp4",
    mimeType: "video/mp4",
    size: CHUNK * CHUNKS,
  };

  afterEach(async () => {
    vi.restoreAllMocks();
    await Promise.all(roots.splice(0).map((root) => rm(root, { recursive: true, force: true })));
  });

  async function localComputer() {
    const root = await realpath(await mkdtemp(path.join(tmpdir(), "rakazo-media-stream-")));
    roots.push(root);
    const desktop = new DesktopSandboxProvider({ root });
    // The production wrapper: desktop computers route to the host provider.
    const sandbox = new HostAwareSandbox(new FakeSandboxProvider(), desktop, async () => true);
    const computer = await desktop.provision({ botId: "bot-1", homePath: "/unused" }, context);
    return { root, sandbox, computer, home: computer.providerRef! };
  }

  /** A stored video served one chunk per read, like the local store's file stream. */
  function storedVideo(onPull: (index: number) => Promise<void> = async () => {}) {
    let index = 0;
    const cancel = vi.fn();
    const stream = new ReadableStream<Uint8Array>(
      {
        async pull(controller) {
          if (index === CHUNKS) {
            controller.close();
            return;
          }
          await onPull(index);
          controller.enqueue(new Uint8Array(CHUNK).fill(index % 251));
          index += 1;
        },
        cancel,
      },
      { highWaterMark: 0 },
    );
    const get = vi.fn();
    const deps = (sandbox: SandboxProvider) => ({
      prisma: {
        artifact: {
          findMany: vi.fn().mockResolvedValue([
            {
              id: "vid-1",
              name: "clip.mp4",
              mimeType: "video/mp4",
              size: CHUNK * CHUNKS,
              storageKey: "k-vid",
            },
          ]),
        },
      } as unknown as PrismaClient,
      artifacts: { get, createReadStream: vi.fn(async () => stream) } as unknown as ArtifactStore,
      sandbox,
    });
    return { deps, get, cancel };
  }

  it("copies a large video to disk chunk by chunk instead of reading it whole", async () => {
    const { sandbox, computer, home } = await localComputer();
    const target = path.join(home, "attachments/vid-1.mp4");
    let mostAhead = 0;
    const video = storedVideo(async (index) => {
      const written = await stat(target).then(
        (info) => info.size,
        () => 0,
      );
      mostAhead = Math.max(mostAhead, index * CHUNK - written);
    });

    const files = await materializeCurrentTurnFiles(video.deps(sandbox), [videoBlock], {
      context,
      computer,
      computerMode: "dedicated",
    });

    expect(files.map((file) => file.path)).toEqual(["attachments/vid-1.mp4"]);
    expect(video.get).not.toHaveBeenCalled();
    // Every chunk reached the disk before the next one was read from the store.
    expect(mostAhead).toBe(0);
    const copied = await readFile(target);
    expect(copied.byteLength).toBe(CHUNK * CHUNKS);
    expect(copied[100 * CHUNK]).toBe(100);
  });

  it("removes a partly copied video, even one replacing an earlier copy, when the run aborts", async () => {
    const { sandbox, computer, home } = await localComputer();
    const target = path.join(home, "attachments/vid-1.mp4");
    await mkdir(path.dirname(target));
    await writeFile(target, "earlier copy");
    const controller = new AbortController();
    const video = storedVideo(async (index) => {
      if (index === 3) controller.abort(new Error("run cancelled"));
    });

    await expect(
      materializeCurrentTurnFiles(video.deps(sandbox), [videoBlock], {
        context: { ...context, signal: controller.signal },
        computer,
        computerMode: "dedicated",
      }),
    ).rejects.toThrow("run cancelled");
    await expect(stat(target)).rejects.toMatchObject({ code: "ENOENT" });
    expect(video.cancel).toHaveBeenCalled();
  });

  it("keeps streamed copies inside the workspace and releases the stored file", async () => {
    const { root, sandbox, computer, home } = await localComputer();
    const outside = path.join(root, "outside");
    await mkdir(outside);
    await symlink(outside, path.join(home, "attachments"), "junction");
    const video = storedVideo();

    await expect(
      materializeCurrentTurnFiles(video.deps(sandbox), [videoBlock], {
        context,
        computer,
        computerMode: "dedicated",
      }),
    ).rejects.toThrow("Path escapes the computer workspace");
    expect(await readdir(outside)).toEqual([]);
    expect(video.cancel).toHaveBeenCalled();
  });

  it("attaches a large local file by streaming it into the store", async () => {
    const { sandbox, computer, home } = await localComputer();
    await mkdir(path.join(home, "out"));
    await writeFile(path.join(home, "out/render.mp4"), new Uint8Array(CHUNK * CHUNKS + 5).fill(7));
    const readWhole = vi.spyOn(DesktopSandboxProvider.prototype, "readFile");
    let largestChunk = 0;
    let stored = 0;
    const putStream = vi.fn(async (artifact: ArtifactStreamPut) => {
      for await (const chunk of byteChunks(artifact.stream)) {
        largestChunk = Math.max(largestChunk, chunk.byteLength);
        stored += chunk.byteLength;
      }
      return { key: "stored-1", size: stored, sha256: "hash" };
    });

    const content = await readComputerFileStream(sandbox, computer, "out/render.mp4", context, {
      maxBytes: 500 * 1024 * 1024,
    });
    const attached = await attachWorkspaceFileToThread(
      {
        prisma: {
          artifact: {
            create: vi.fn(async ({ data }: { data: Record<string, unknown> }) => ({
              id: "artifact-1",
              ...data,
            })),
          },
        } as unknown as PrismaClient,
        artifacts: {
          describe: () => ({ capabilities: { stream: true } }),
          putStream,
          remove: vi.fn(),
        } as unknown as ArtifactStore,
      },
      {
        spaceId: "workspace-1",
        userId: "user-1",
        botId: "bot-1",
        runId: "run-1",
        filePath: "out/render.mp4",
        content,
        operationId: "attach-1",
      },
    );

    expect(attached.block).toMatchObject({ kind: "video", name: "render.mp4" });
    expect(stored).toBe(CHUNK * CHUNKS + 5);
    expect(largestChunk).toBeLessThanOrEqual(256 * 1024);
    expect(readWhole).not.toHaveBeenCalled();
  });

  it("refuses streamed reads above the cap or outside the workspace", async () => {
    const { root, sandbox, computer, home } = await localComputer();
    await writeFile(path.join(home, "clip.mp4"), new Uint8Array(10));
    await writeFile(path.join(root, "secret.mp4"), new Uint8Array(10));
    await symlink(path.join(root, "secret.mp4"), path.join(home, "escape.mp4"));

    await expect(
      readComputerFileStream(sandbox, computer, "clip.mp4", context, { maxBytes: 9 }),
    ).rejects.toThrow("computer file exceeds 9 bytes");
    await expect(readComputerFileStream(sandbox, computer, "escape.mp4", context)).rejects.toThrow(
      "Path escapes the computer workspace",
    );
    await expect(
      readComputerFileStream(sandbox, computer, "../secret.mp4", context),
    ).rejects.toThrow();
  });

  it("releases an attached stream that is not stored and keeps the limit message", async () => {
    const deps = (putStream: ArtifactStore["putStream"]) => ({
      prisma: { artifact: { create: vi.fn() } } as unknown as PrismaClient,
      artifacts: {
        describe: () => ({ capabilities: { stream: true } }),
        putStream,
        remove: vi.fn(),
      } as unknown as ArtifactStore,
    });
    const input = {
      spaceId: "workspace-1",
      userId: "user-1",
      botId: "bot-1",
      runId: "run-1",
      operationId: "attach-1",
    };
    const unread = vi.fn();
    const putStream = vi.fn();

    await expect(
      attachWorkspaceFileToThread(deps(putStream), {
        ...input,
        filePath: "notes.unknown-type",
        content: new ReadableStream({ cancel: unread }),
      }),
    ).rejects.toThrow("Unsupported attachment type");
    expect(unread).toHaveBeenCalled();
    expect(putStream).not.toHaveBeenCalled();

    await expect(
      attachWorkspaceFileToThread(
        deps(vi.fn().mockRejectedValue(new ArtifactTooLargeError(500 * 1024 * 1024))),
        { ...input, filePath: "render.mp4", content: byteStream([new Uint8Array(1)]) },
      ),
    ).rejects.toThrow("file exceeds the 500 MiB attachment limit");
  });
});

describe("executor attachments", () => {
  it("records the thread's group on every file the executor attaches", () => {
    // Driving render_plot's PNG fallback needs a full run; this guards the wiring that once
    // dropped groupId, so the chart could not be opened in group conversations.
    const source = readFileSync(new URL("./executor.ts", import.meta.url), "utf8");
    const calls = source.split("attachWorkspaceFileToThread(").slice(1);
    expect(calls.length).toBeGreaterThanOrEqual(2);
    for (const call of calls) {
      const input = call.slice(0, call.indexOf("operationId:"));
      expect(input).toContain("groupId: thread.groupId ?? undefined");
    }
  });
});