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

import path from "node:path";
import type {
  AdapterContext,
  ArtifactStore,
  ComputerRef,
  SandboxProvider,
} from "@rakazo/adapter-kit";
import type { ComputerMode, MessageBlock } from "@rakazo/contracts";
import { ATTACHMENT_MAX_BYTES } from "@rakazo/contracts";
import {
  attachmentExtensionForMimeType,
  inferAttachmentMimeType,
  messageBlockForArtifact,
} from "@rakazo/core";
import type { PrismaClient } from "@rakazo/db";
import { artifactMaxBytes, storeArtifactStream } from "./artifact-files.js";
import {
  assertTeamWorkspacePathAllowed,
  computerFileTransferMaxBytes,
  resolveBotWorkspacePath,
} from "./computer-support.js";

export type MaterializedThreadFile = {
  artifactId: string;
  name: string;
  mimeType: string;
  size: number;
  /** Workspace path, or null when the file is too large to copy to this computer. */
  path: string | null;
};

type AttachmentBlock = Extract<MessageBlock, { kind: "image" | "video" | "audio" | "file" }>;

const ATTACHMENT_BLOCK_KINDS: ReadonlySet<string> = new Set(["image", "video", "audio", "file"]);

/** Media up to its kind cap when the store streams; otherwise the inline attachment cap. */
export function attachableFileMaxBytes(artifacts: ArtifactStore, mimeType: string): number {
  return artifacts.describe().capabilities.stream
    ? artifactMaxBytes(mimeType)
    : ATTACHMENT_MAX_BYTES;
}

export function formatMiB(bytes: number): string {
  return `${Math.floor(bytes / (1024 * 1024))} MiB`;
}

export async function attachWorkspaceFileToThread(
  deps: {
    prisma: PrismaClient;
    artifacts: ArtifactStore;
  },
  input: {
    spaceId: string;
    userId: string;
    botId: string;
    groupId?: string;
    runId: string;
    filePath: string;
    bytes: Uint8Array;
    operationId: string;
  },
): Promise<{ artifactId: string; block: AttachmentBlock }> {
  const name = path.basename(input.filePath) || input.filePath;
  const mimeType = inferAttachmentMimeType(name);
  if (!mimeType) {
    throw new Error(`Unsupported attachment type for ${name}`);
  }
  const maxBytes = attachableFileMaxBytes(deps.artifacts, mimeType);
  if (input.bytes.byteLength > maxBytes) {
    throw new Error(`file exceeds the ${formatMiB(maxBytes)} attachment limit`);
  }
  const bytes = input.bytes;
  const row = await storeArtifactStream(deps, {
    spaceId: input.spaceId,
    userId: input.userId,
    botId: input.botId,
    groupId: input.groupId ?? null,
    runId: input.runId,
    name,
    mimeType,
    stream: (async function* () {
      yield bytes;
    })(),
    maxBytes,
    operationId: input.operationId,
  });
  return { artifactId: row.id, block: messageBlockForArtifact(row) };
}

/**
 * Copies this turn's attachments (files, images, video and audio) into the bot's workspace so its
 * tools can open them. Files above the computer's transfer cap stay in storage and are listed
 * with their artifact id instead.
 */
export async function materializeCurrentTurnFiles(
  deps: {
    prisma: PrismaClient;
    artifacts: ArtifactStore;
    sandbox: SandboxProvider;
  },
  blocks: MessageBlock[] | undefined,
  input: {
    context: AdapterContext & { botId: string };
    computer: ComputerRef;
    computerMode: ComputerMode;
    teamBotIds?: ReadonlySet<string>;
    markWorkspaceDirty?: () => void;
  },
): Promise<MaterializedThreadFile[]> {
  const attachmentBlocks = blocks?.filter((block): block is AttachmentBlock =>
    ATTACHMENT_BLOCK_KINDS.has(block.kind),
  );
  if (!attachmentBlocks?.length) return [];

  const rows = await deps.prisma.artifact.findMany({
    where: {
      id: { in: attachmentBlocks.map((block) => block.artifactId) },
      spaceId: input.context.spaceId,
      userId: input.context.userId,
    },
  });
  const byId = new Map(rows.map((row) => [row.id, row]));
  const transferMaxBytes = computerFileTransferMaxBytes(input.computer);
  const materialized: MaterializedThreadFile[] = [];

  for (const block of attachmentBlocks) {
    const row = byId.get(block.artifactId);
    if (!row) throw new Error(`Attached file is unavailable: ${block.name}`);
    const entry = { artifactId: row.id, name: row.name, mimeType: row.mimeType, size: row.size };
    if (row.size > Math.min(transferMaxBytes, artifactMaxBytes(row.mimeType))) {
      materialized.push({ ...entry, path: null });
      continue;
    }
    const bytes = await deps.artifacts.get(row.storageKey, input.context);
    const relativePath = `attachments/${row.id}${attachmentExtensionForMimeType(row.mimeType)}`;
    const storedPath = resolveBotWorkspacePath(
      input.computerMode,
      input.context.botId,
      relativePath,
    );
    if (input.teamBotIds) assertTeamWorkspacePathAllowed(storedPath, input.teamBotIds);
    input.markWorkspaceDirty?.();
    await deps.sandbox.writeFile(
      input.computer,
      {
        path: storedPath,
        content: bytes,
      },
      input.context,
    );
    materialized.push({ ...entry, path: relativePath });
  }

  return materialized;
}

export function currentTurnFilesInstruction(files: readonly MaterializedThreadFile[]): string {
  const copied = files.filter((file) => file.path !== null);
  const kept = files.filter((file) => file.path === null);
  const describe = (file: MaterializedThreadFile) =>
    `- ${JSON.stringify(file.name)} (${file.mimeType}, ${file.size} bytes)`;
  return [
    ...(copied.length
      ? [
          "Current-turn attached files are available on your computer. Inspect them before answering:",
          ...copied.map((file) => `${describe(file)}: ${JSON.stringify(file.path)}`),
        ]
      : []),
    ...(kept.length
      ? [
          "These current-turn attachments are too large to copy to your computer. Refer to them by artifact id; do not guess their contents:",
          ...kept.map((file) => `${describe(file)}: artifact ${file.artifactId}`),
        ]
      : []),
  ].join("\n");
}