diff --git a/apps/api/src/artifacts.ts b/apps/api/src/artifacts.ts index b0804a0..9dec9c4 100644 --- a/apps/api/src/artifacts.ts +++ b/apps/api/src/artifacts.ts @@ -11,6 +11,12 @@ import { } from "@rakazo/core"; import { IsolationError, type PrismaClient } from "@rakazo/db"; +/** + * Base64 reads load the whole file into memory; larger media is served by + * GET /api/artifacts/:id instead. + */ +export const ARTIFACT_INLINE_READ_MAX_BYTES = 25 * 1024 * 1024; + function adapterContext(actor: Actor, botId: string, operationId: string) { return { operationId, @@ -131,6 +137,9 @@ async function readArtifact( }, contextBotId: string, ) { + if (row.size > ARTIFACT_INLINE_READ_MAX_BYTES) { + throw new AttachmentValidationError("This file is too large to load inline."); + } const bytes = await artifacts.get( row.storageKey, adapterContext(actor, contextBotId, `artifact-get:${row.id}`), diff --git a/apps/api/src/router.ts b/apps/api/src/router.ts index 2baf21e..cc14276 100644 --- a/apps/api/src/router.ts +++ b/apps/api/src/router.ts @@ -4640,24 +4640,26 @@ export function createRouter(deps: RouterDeps) { } }), get: authed.artifacts.get.handler(async ({ context, input }) => { - if (input.groupId) { - const group = await groupRepos.getGroupTarget(context.actor, input.groupId); - const contextBotId = group.members[0]?.bot.id; - if (!contextBotId) throw new IsolationError(); - return getSpaceArtifact(deps, context.actor, { - artifactId: input.artifactId, - groupId: input.groupId, - contextBotId, - }); - } - await repos.getBot(context.actor, input.botId!); try { + if (input.groupId) { + const group = await groupRepos.getGroupTarget(context.actor, input.groupId); + const contextBotId = group.members[0]?.bot.id; + if (!contextBotId) throw new IsolationError(); + return await getSpaceArtifact(deps, context.actor, { + artifactId: input.artifactId, + groupId: input.groupId, + contextBotId, + }); + } + await repos.getBot(context.actor, input.botId!); return await getOwnedArtifact(deps, context.actor, { botId: input.botId!, artifactId: input.artifactId, }); } catch (error) { - if (error instanceof IsolationError) throw error; + if (error instanceof AttachmentValidationError) { + throw new ORPCError("BAD_REQUEST", { message: error.message }); + } throw error; } }), diff --git a/packages/adapters/src/builtin-tools.ts b/packages/adapters/src/builtin-tools.ts index 0ff52fa..043aa9d 100644 --- a/packages/adapters/src/builtin-tools.ts +++ b/packages/adapters/src/builtin-tools.ts @@ -144,7 +144,7 @@ export const builtinAgentTools: ConnectorTool[] = [ { name: "attach_file", description: - "Attach a workspace file from this bot's home to the chat thread as an image or common file. The file stays in place; users can open it from the message.", + "Attach a workspace file from this bot's home to the chat thread as an image, video, audio or common file. The file stays in place; users can open or play it from the message.", inputSchema: { type: "object", properties: { path: { type: "string" } }, diff --git a/packages/adapters/src/computer-support.ts b/packages/adapters/src/computer-support.ts index 62deb83..a6e6519 100644 --- a/packages/adapters/src/computer-support.ts +++ b/packages/adapters/src/computer-support.ts @@ -17,6 +17,18 @@ export function toComputerRef(computer: { }; } +/** + * Remote computers move files as base64 JSON (the Docker supervisor accepts 16 MiB requests), so + * files copied to or read from them stay below this size. Local computers have no transfer cap. + */ +export const REMOTE_COMPUTER_FILE_TRANSFER_MAX_BYTES = 12 * 1024 * 1024; + +export function computerFileTransferMaxBytes(computer: Pick): number { + return computer.kind === "desktop" || computer.kind === "fake" + ? Number.POSITIVE_INFINITY + : REMOTE_COMPUTER_FILE_TRANSFER_MAX_BYTES; +} + const EMPTY_PNG = Uint8Array.from( Buffer.from( "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=", diff --git a/packages/adapters/src/executor.ts b/packages/adapters/src/executor.ts index 5cb9181..8b8f8a6 100644 --- a/packages/adapters/src/executor.ts +++ b/packages/adapters/src/executor.ts @@ -197,6 +197,7 @@ import { withComputerScreenAvailability } from "./computer-screens.js"; import { assertTeamShellCommandAllowed, assertTeamWorkspacePathAllowed, + computerFileTransferMaxBytes, displayBotWorkspacePath, filterTeamWorkspaceEntries, resolveBotWorkspaceCwd, @@ -368,8 +369,10 @@ import { } from "./takeover-resume.js"; import { getActiveTeachingSession, parsePlaybook } from "./teaching-session.js"; import { + attachableFileMaxBytes, attachWorkspaceFileToThread, currentTurnFilesInstruction, + formatMiB, materializeCurrentTurnFiles, } from "./thread-artifacts.js"; import { advanceToolCallLoopGuard } from "./tool-loop.js"; @@ -2366,6 +2369,7 @@ export function createRunExecutor(deps: ExecutorDeps) { spaceId: run.spaceId, userId: run.userId, botId: bot.id, + groupId: thread.groupId ?? undefined, runId: run.id, filePath: outPath, bytes: png, @@ -2391,18 +2395,23 @@ export function createRunExecutor(deps: ExecutorDeps) { return finish({ error: "artifact storage unavailable", path: filePath }); } const storedPath = resolveStoredWorkspacePath(filePath); - let bytes: Uint8Array; - try { - bytes = await deps.sandbox.readFile(computer, storedPath, context, { - maxBytes: ATTACHMENT_MAX_BYTES, - }); - } catch { - return finish({ error: "file not found or unreadable", path: filePath }); - } const mimeType = inferAttachmentMimeType(filePath); if (!mimeType) { return finish({ error: "unsupported attachment type", path: filePath }); } + const maxBytes = Math.min( + attachableFileMaxBytes(deps.artifacts, mimeType), + computerFileTransferMaxBytes(computer), + ); + let bytes: Uint8Array; + try { + bytes = await deps.sandbox.readFile(computer, storedPath, context, { maxBytes }); + } catch { + return finish({ + error: `file not found, unreadable or larger than ${formatMiB(maxBytes)}`, + path: filePath, + }); + } try { const attached = await attachWorkspaceFileToThread( { prisma: deps.prisma, artifacts: deps.artifacts }, @@ -3886,7 +3895,7 @@ export function createRunExecutor(deps: ExecutorDeps) { item.blocks, context.signal, ); - workspaceCheckpoint.markFiles(files); + workspaceCheckpoint.markFiles(files.filter((file) => file.path)); const filesInstruction = currentTurnFilesInstruction(files); return { id: item.id, diff --git a/packages/adapters/src/history-compaction.test.ts b/packages/adapters/src/history-compaction.test.ts index ef8ad64..b130201 100644 --- a/packages/adapters/src/history-compaction.test.ts +++ b/packages/adapters/src/history-compaction.test.ts @@ -498,7 +498,7 @@ describe("compactHistory", () => { const [request] = harness.runtime.run.mock.calls[0]!; expect(request.prompt).toContain("[file: plan.pdf (application/pdf, 123 bytes)]"); - expect(request.prompt).toContain("[image: diagram.png]"); + expect(request.prompt).toContain("[image: diagram.png (artifact artifact-2)]"); }); it("compacts locally without a semantic memory provider", async () => { diff --git a/packages/adapters/src/reply-context.test.ts b/packages/adapters/src/reply-context.test.ts index 136dea7..9b0db83 100644 --- a/packages/adapters/src/reply-context.test.ts +++ b/packages/adapters/src/reply-context.test.ts @@ -55,11 +55,11 @@ describe("reply context", () => { ...target, blocks: [ { kind: "text", text: "" }, - { kind: "image", name: "example.png" }, + { kind: "image", artifactId: "artifact-1", name: "example.png" }, ], }); const context = await loadReplyContext(prisma, "thread-1", "user-reply"); - expect(context).toContain("[image: example.png]"); + expect(context).toContain("[image: example.png (artifact artifact-1)]"); expect(context).toContain("\\u003c/reply_target\\u003e"); expect(context?.match(/<\/reply_target>/g)).toHaveLength(1); }); diff --git a/packages/adapters/src/thread-artifacts.test.ts b/packages/adapters/src/thread-artifacts.test.ts index 5d164bd..fb79536 100644 --- a/packages/adapters/src/thread-artifacts.test.ts +++ b/packages/adapters/src/thread-artifacts.test.ts @@ -8,6 +8,7 @@ import type { MessageBlock } from "@rakazo/contracts"; import type { PrismaClient } from "@rakazo/db"; import { describe, expect, it, vi } from "vitest"; import { + attachableFileMaxBytes, attachWorkspaceFileToThread, currentTurnFilesInstruction, materializeCurrentTurnFiles, @@ -25,7 +26,8 @@ describe("current-turn thread files", () => { artifact: { create: vi.fn().mockRejectedValue(failure) }, } as unknown as PrismaClient, artifacts: { - put: vi.fn().mockResolvedValue({ id: "stored-1", hash: "hash" }), + describe: () => ({ capabilities: { stream: true } }), + putStream: vi.fn().mockResolvedValue({ key: "stored-1", size: 3, sha256: "hash" }), remove, } as unknown as ArtifactStore, }, @@ -111,6 +113,7 @@ describe("current-turn thread files", () => { ); expect(files).toEqual([ { + artifactId: "artifact-1", name: "../quarterly report.pdf", mimeType: "application/pdf", size: 4, @@ -120,7 +123,73 @@ describe("current-turn thread files", () => { expect(currentTurnFilesInstruction(files)).toContain('"attachments/artifact-1.pdf"'); }); - it("does not load images as computer files", async () => { + 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 get = vi.fn().mockResolvedValue(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: { get } 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(get).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( { @@ -128,14 +197,7 @@ describe("current-turn thread files", () => { artifacts: {} as ArtifactStore, sandbox: {} as SandboxProvider, }, - [ - { - kind: "image", - artifactId: "image-1", - name: "photo.png", - mimeType: "image/png", - }, - ], + [{ kind: "text", text: "hello" }], { context: { operationId: "run-1", @@ -158,4 +220,12 @@ describe("current-turn thread files", () => { 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); + }); }); diff --git a/packages/adapters/src/thread-artifacts.ts b/packages/adapters/src/thread-artifacts.ts index 72bc432..8d82aef 100644 --- a/packages/adapters/src/thread-artifacts.ts +++ b/packages/adapters/src/thread-artifacts.ts @@ -1,4 +1,3 @@ -import { createHash } from "node:crypto"; import path from "node:path"; import type { AdapterContext, @@ -12,18 +11,39 @@ import { attachmentExtensionForMimeType, inferAttachmentMimeType, messageBlockForArtifact, - validateAttachmentMimeType, } from "@rakazo/core"; import type { PrismaClient } from "@rakazo/db"; -import { assertTeamWorkspacePathAllowed, resolveBotWorkspacePath } from "./computer-support.js"; +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; - path: string; + /** Workspace path, or null when the file is too large to copy to this computer. */ + path: string | null; }; +type AttachmentBlock = Extract; + +const ATTACHMENT_BLOCK_KINDS: ReadonlySet = 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) + : Math.min(ATTACHMENT_MAX_BYTES, artifactMaxBytes(mimeType)); +} + +export function formatMiB(bytes: number): string { + return `${Math.floor(bytes / (1024 * 1024))} MiB`; +} + export async function attachWorkspaceFileToThread( deps: { prisma: PrismaClient; @@ -39,60 +59,38 @@ export async function attachWorkspaceFileToThread( bytes: Uint8Array; operationId: string; }, -): Promise<{ - artifactId: string; - block: Extract; -}> { +): 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}`); } - validateAttachmentMimeType(mimeType); - if (input.bytes.byteLength > ATTACHMENT_MAX_BYTES) { - throw new Error("file exceeds the 10 MiB attachment limit"); + const maxBytes = attachableFileMaxBytes(deps.artifacts, mimeType); + if (input.bytes.byteLength > maxBytes) { + throw new Error(`file exceeds the ${formatMiB(maxBytes)} attachment limit`); } - - const context = { - operationId: input.operationId, - traceId: input.operationId, + const bytes = input.bytes; + const row = await storeArtifactStream(deps, { spaceId: input.spaceId, userId: input.userId, botId: input.botId, - signal: new AbortController().signal, - }; - const stored = await deps.artifacts.put({ name, mimeType, bytes: input.bytes }, context); - const hash = createHash("sha256").update(input.bytes).digest("hex"); - const row = await deps.prisma.artifact - .create({ - data: { - spaceId: input.spaceId, - userId: input.userId, - botId: input.botId, - groupId: input.groupId, - runId: input.runId, - name, - mimeType, - size: input.bytes.byteLength, - hash, - storageKey: stored.id, - }, - }) - .catch(async (error) => { - await deps.artifacts.remove(stored.id, context).catch(() => undefined); - throw error; - }); - return { - artifactId: row.id, - block: messageBlockForArtifact({ - id: row.id, - name: row.name, - mimeType: row.mimeType, - size: row.size, - }), - }; + groupId: input.groupId ?? null, + runId: input.runId, + name, + mimeType, + stream: (async function* () { + yield bytes; + })(), + maxBytes, + }); + 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; @@ -108,28 +106,31 @@ export async function materializeCurrentTurnFiles( markWorkspaceDirty?: () => void; }, ): Promise { - const fileBlocks = blocks?.filter( - (block): block is Extract => block.kind === "file", + const attachmentBlocks = blocks?.filter((block): block is AttachmentBlock => + ATTACHMENT_BLOCK_KINDS.has(block.kind), ); - if (!fileBlocks?.length) return []; + if (!attachmentBlocks?.length) return []; const rows = await deps.prisma.artifact.findMany({ where: { - id: { in: fileBlocks.map((block) => block.artifactId) }, + 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 fileBlocks) { + for (const block of attachmentBlocks) { const row = byId.get(block.artifactId); if (!row) throw new Error(`Attached file is unavailable: ${block.name}`); - const bytes = await deps.artifacts.get(row.storageKey, input.context); - if (bytes.byteLength > ATTACHMENT_MAX_BYTES) { - throw new Error(`Attached file exceeds the 10 MiB limit: ${row.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, @@ -146,24 +147,29 @@ export async function materializeCurrentTurnFiles( }, input.context, ); - materialized.push({ - name: row.name, - mimeType: row.mimeType, - size: row.size, - path: relativePath, - }); + materialized.push({ ...entry, path: relativePath }); } return materialized; } export function currentTurnFilesInstruction(files: readonly MaterializedThreadFile[]): string { - if (!files.length) return ""; + 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 [ - "Current-turn attached files are available on your computer. Inspect them before answering:", - ...files.map( - (file) => - `- ${JSON.stringify(file.name)} (${file.mimeType}, ${file.size} bytes): ${JSON.stringify(file.path)}`, - ), + ...(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"); } diff --git a/packages/core/src/attachments.test.ts b/packages/core/src/attachments.test.ts index 6db2e57..b39c7d4 100644 --- a/packages/core/src/attachments.test.ts +++ b/packages/core/src/attachments.test.ts @@ -44,7 +44,7 @@ describe("attachment helpers", () => { size: 99, }, ]), - ).toBe("hello\n[image: shot.png]\n[file: brief.pdf (application/pdf, 99 bytes)]"); + ).toBe("hello\n[image: shot.png (artifact a1)]\n[file: brief.pdf (application/pdf, 99 bytes)]"); }); it("infers attachment mime types from extensions", () => { diff --git a/packages/core/src/attachments.ts b/packages/core/src/attachments.ts index b4a6759..9953cb8 100644 --- a/packages/core/src/attachments.ts +++ b/packages/core/src/attachments.ts @@ -119,12 +119,11 @@ export function blocksToAgentHistoryText(blocks: MessageBlock[]): string { .map((block) => { if (block.kind === "text") return block.text; if (block.kind === "chart") return `[chart: ${block.name}]`; - if (block.kind === "image") return `[image: ${block.name}]`; if (block.kind === "file") { return `[file: ${block.name} (${block.mimeType}, ${block.size} bytes)]`; } // Media keeps the artifact id so a later turn can edit or export what it refers to. - if (block.kind === "video" || block.kind === "audio") { + if (block.kind === "image" || block.kind === "video" || block.kind === "audio") { return `[${block.kind}: ${block.name} (artifact ${block.artifactId})]`; } if (block.kind === "media") { diff --git a/packages/db/src/events.test.ts b/packages/db/src/events.test.ts index e8ca709..afc0e39 100644 --- a/packages/db/src/events.test.ts +++ b/packages/db/src/events.test.ts @@ -1739,7 +1739,7 @@ describe("claimSteering", () => { { id: "steer-1", messageId: "message-1", - text: "First\n[image: chart.png]", + text: "First\n[image: chart.png (artifact artifact-1)]", blocks: [ { kind: "text", text: "First" }, { diff --git a/packages/testkit/src/artifact-routes.test.ts b/packages/testkit/src/artifact-routes.test.ts index bc33ca4..3b6a3e6 100644 --- a/packages/testkit/src/artifact-routes.test.ts +++ b/packages/testkit/src/artifact-routes.test.ts @@ -283,6 +283,19 @@ describeRoutes("artifact HTTP routes", () => { const me = await rpc<{ spaceId: string }>(app, owner, "me"); const spaceDir = path.join(dataDir, "artifacts", me.spaceId); expect(readdirSafe(spaceDir)).toEqual([]); + + // Video far above the inline cap uploads, streams by range, and is refused as base64. + const large = Buffer.alloc(26 * 1024 * 1024, 1); + large.write("head", 0); + const uploaded = await upload(app, owner, { botId: bot.id }, "long.mp4", "video/mp4", large); + expect(uploaded.status).toBe(201); + const largeVideo = (await uploaded.json()) as { id: string }; + const head = await get(app, owner, largeVideo.id, { range: "bytes=0-3" }); + expect(head.status).toBe(206); + expect(await head.text()).toBe("head"); + await expect( + rpc(app, owner, "artifacts/get", { botId: bot.id, artifactId: largeVideo.id }), + ).rejects.toThrow(/400: .*too large to load inline/); }); });