diff --git a/packages/adapter-kit/src/interfaces.ts b/packages/adapter-kit/src/interfaces.ts index c432310..3968d14 100644 --- a/packages/adapter-kit/src/interfaces.ts +++ b/packages/adapter-kit/src/interfaces.ts @@ -51,6 +51,7 @@ import type { PageBrowserCommand, PageBrowserResult, PortableFile, + PortableFileStream, ProcessEvent, SandboxCapabilities, ScreenRequest, @@ -145,6 +146,21 @@ export interface SandboxProvider { options?: { maxBytes?: number }, ): Promise; writeFile(computer: ComputerRef, file: PortableFile, context: AdapterContext): Promise; + /** + * Bounded-memory variants of readFile and writeFile for large media on computers without a + * transfer cap. Same path rules as their byte versions; a failed or aborted write leaves no file. + */ + readFileStream?( + computer: ComputerRef, + path: string, + context: AdapterContext, + options?: { maxBytes?: number }, + ): Promise>; + writeFileStream?( + computer: ComputerRef, + file: PortableFileStream, + context: AdapterContext, + ): Promise; exportWorkspace(computer: ComputerRef, context: AdapterContext): AsyncIterable; importWorkspace( computer: ComputerRef, diff --git a/packages/adapter-kit/src/types.ts b/packages/adapter-kit/src/types.ts index 67669e5..3b29a16 100644 --- a/packages/adapter-kit/src/types.ts +++ b/packages/adapter-kit/src/types.ts @@ -52,6 +52,13 @@ export interface PortableFile { executable?: boolean; } +/** A workspace file whose bytes arrive as a stream instead of one buffer. */ +export interface PortableFileStream { + path: string; + stream: ReadableStream | AsyncIterable; + executable?: boolean; +} + export interface ComputerRef { id: string; botId: string; diff --git a/packages/adapters/src/artifacts.ts b/packages/adapters/src/artifacts.ts index 6143df4..ee8d59c 100644 --- a/packages/adapters/src/artifacts.ts +++ b/packages/adapters/src/artifacts.ts @@ -1,6 +1,5 @@ import { createHash, randomUUID } from "node:crypto"; import { constants } from "node:fs"; -import type { FileHandle } from "node:fs/promises"; import { lstat, mkdir, open, rm } from "node:fs/promises"; import path from "node:path"; import type { @@ -14,9 +13,9 @@ import type { StoredArtifactStream, } from "@rakazo/adapter-kit"; import { ArtifactTooLargeError } from "@rakazo/adapter-kit"; +import { byteChunks, fileByteStream, writeFully } from "./byte-streams.js"; const O_NOFOLLOW = constants.O_NOFOLLOW ?? 0; -const READ_CHUNK_BYTES = 256 * 1024; /** Space ids and artifact keys are generated ids; anything else could address another path. */ const SAFE_SEGMENT = /^[A-Za-z0-9][A-Za-z0-9_-]{0,127}$/; @@ -155,72 +154,6 @@ export class LocalArtifactStore implements ArtifactStore { } } -async function* byteChunks( - source: ReadableStream | AsyncIterable, -): AsyncIterable { - if (!("getReader" in source)) { - yield* source; - return; - } - const reader = source.getReader(); - try { - for (;;) { - const { done, value } = await reader.read(); - if (done) return; - yield value; - } - } finally { - // Stop the producer (e.g. an upload body) when the store gives up early, without waiting on - // a producer that may never confirm. - void reader.cancel().catch(() => undefined); - reader.releaseLock(); - } -} - -async function writeFully(handle: FileHandle, chunk: Uint8Array) { - let offset = 0; - while (offset < chunk.byteLength) { - const { bytesWritten } = await handle.write(chunk, offset, chunk.byteLength - offset); - offset += bytesWritten; - } -} - -/** Streams [start, end] of an open file and closes it when done, failed or cancelled. */ -function fileByteStream(handle: FileHandle, start: number, end: number) { - let position = start; - let closed = false; - const close = async () => { - if (closed) return; - closed = true; - await handle.close().catch(() => undefined); - }; - return new ReadableStream({ - async pull(controller) { - try { - const length = Math.min(end - position + 1, READ_CHUNK_BYTES); - if (length <= 0) { - await close(); - controller.close(); - return; - } - const buffer = new Uint8Array(length); - const { bytesRead } = await handle.read(buffer, 0, length, position); - if (bytesRead === 0) throw new Error("Artifact ended before the requested range"); - position += bytesRead; - controller.enqueue(bytesRead === length ? buffer : buffer.subarray(0, bytesRead)); - if (position > end) { - await close(); - controller.close(); - } - } catch (error) { - await close(); - controller.error(error); - } - }, - cancel: close, - }); -} - export class CapturingNotificationProvider implements NotificationProvider { readonly sent: NotificationMessage[] = []; diff --git a/packages/adapters/src/byte-streams.ts b/packages/adapters/src/byte-streams.ts new file mode 100644 index 0000000..b7e2da5 --- /dev/null +++ b/packages/adapters/src/byte-streams.ts @@ -0,0 +1,91 @@ +import type { FileHandle } from "node:fs/promises"; + +type ByteSource = ReadableStream | AsyncIterable; + +const READ_CHUNK_BYTES = 256 * 1024; + +export async function* byteChunks(source: ByteSource): AsyncIterable { + if (!("getReader" in source)) { + yield* source; + return; + } + const reader = source.getReader(); + try { + for (;;) { + const { done, value } = await reader.read(); + if (done) return; + yield value; + } + } finally { + // Stop the producer (e.g. an upload body) when the consumer gives up early, without waiting + // on a producer that may never confirm. + void reader.cancel().catch(() => undefined); + reader.releaseLock(); + } +} + +/** Releases a stream nobody will read, so its producer (e.g. an open file) is closed. */ +export function discardByteStream(source: ByteSource): void { + if ("cancel" in source) void source.cancel().catch(() => undefined); +} + +/** Collects a stream into one buffer, for destinations that only accept bytes. */ +export async function collectBytes(source: ByteSource): Promise { + const chunks: Uint8Array[] = []; + let size = 0; + for await (const chunk of byteChunks(source)) { + chunks.push(chunk); + size += chunk.byteLength; + } + const bytes = new Uint8Array(size); + let offset = 0; + for (const chunk of chunks) { + bytes.set(chunk, offset); + offset += chunk.byteLength; + } + return bytes; +} + +export async function writeFully(handle: FileHandle, chunk: Uint8Array) { + let offset = 0; + while (offset < chunk.byteLength) { + const { bytesWritten } = await handle.write(chunk, offset, chunk.byteLength - offset); + offset += bytesWritten; + } +} + +/** Streams [start, end] of an open file and closes it when done, failed or cancelled. */ +export function fileByteStream(handle: FileHandle, start: number, end: number) { + let position = start; + let closed = false; + const close = async () => { + if (closed) return; + closed = true; + await handle.close().catch(() => undefined); + }; + return new ReadableStream({ + async pull(controller) { + try { + const length = Math.min(end - position + 1, READ_CHUNK_BYTES); + if (length <= 0) { + await close(); + controller.close(); + return; + } + const buffer = new Uint8Array(length); + const { bytesRead } = await handle.read(buffer, 0, length, position); + if (bytesRead === 0) throw new Error("File ended before the requested range"); + position += bytesRead; + controller.enqueue(bytesRead === length ? buffer : buffer.subarray(0, bytesRead)); + if (position > end) { + await close(); + controller.close(); + } + } catch (error) { + await close(); + controller.error(error); + } + }, + cancel: close, + }); +} diff --git a/packages/adapters/src/computer-support.ts b/packages/adapters/src/computer-support.ts index a6e6519..00948e0 100644 --- a/packages/adapters/src/computer-support.ts +++ b/packages/adapters/src/computer-support.ts @@ -1,7 +1,15 @@ import { createHash } from "node:crypto"; import path from "node:path"; -import type { ComputerAction, ComputerObservation, ComputerRef } from "@rakazo/adapter-kit"; +import type { + AdapterContext, + ComputerAction, + ComputerObservation, + ComputerRef, + PortableFileStream, + SandboxProvider, +} from "@rakazo/adapter-kit"; import type { ComputerMode } from "@rakazo/contracts"; +import { collectBytes, discardByteStream } from "./byte-streams.js"; export function toComputerRef(computer: { homeKey: string; @@ -29,6 +37,48 @@ export function computerFileTransferMaxBytes(computer: Pick : REMOTE_COMPUTER_FILE_TRANSFER_MAX_BYTES; } +/** Streams a file into the workspace; providers without streaming writes receive its bytes. */ +export async function writeComputerFileStream( + sandbox: SandboxProvider, + computer: ComputerRef, + file: PortableFileStream, + context: AdapterContext, +): Promise { + try { + if (sandbox.writeFileStream) { + await sandbox.writeFileStream(computer, file, context); + return; + } + const content = await collectBytes(file.stream); + await sandbox.writeFile( + computer, + { path: file.path, content, ...(file.executable ? { executable: true } : {}) }, + context, + ); + } catch (error) { + discardByteStream(file.stream); + throw error; + } +} + +/** Opens a workspace file as a stream; providers without streaming reads return it in one chunk. */ +export async function readComputerFileStream( + sandbox: SandboxProvider, + computer: ComputerRef, + filePath: string, + context: AdapterContext, + options?: { maxBytes?: number }, +): Promise> { + if (sandbox.readFileStream) return sandbox.readFileStream(computer, filePath, context, options); + const bytes = await sandbox.readFile(computer, filePath, context, options); + return new ReadableStream({ + start(controller) { + controller.enqueue(bytes); + controller.close(); + }, + }); +} + const EMPTY_PNG = Uint8Array.from( Buffer.from( "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=", diff --git a/packages/adapters/src/desktop-sandbox-write-containment.test.ts b/packages/adapters/src/desktop-sandbox-write-containment.test.ts index 313a3d7..46b2dcb 100644 --- a/packages/adapters/src/desktop-sandbox-write-containment.test.ts +++ b/packages/adapters/src/desktop-sandbox-write-containment.test.ts @@ -265,7 +265,7 @@ describe("desktop sandbox write containment without O_NOFOLLOW", () => { lstatRace.afterRealpath = async (inspected) => { if (path.basename(inspected) !== "notes") return; notesRealpaths += 1; - // After openContainedWorkspaceFile's parent realpath, point the pathname at an + // After writeContainedWorkspaceFile's parent realpath, point the pathname at an // outside directory. Without binding the opened parent inode to a fresh // contained realpath, later checks could validate a restored inside path // while the handle still referred outside. @@ -313,6 +313,74 @@ describe("desktop sandbox write containment without O_NOFOLLOW", () => { expect(await readFile(outside, "utf8")).toBe("outside-before"); }); + it("applies the same containment to streamed writes", async () => { + const { root, desktop, computer } = await fixture("streamed-escapes"); + const outsideFile = path.join(root, "outside.txt"); + const outsideDirectory = path.join(root, "outside-directory"); + await writeFile(outsideFile, "outside-before"); + await mkdir(outsideDirectory); + await symlink(outsideFile, path.join(computer.providerRef, "final-link.txt")); + await link(outsideFile, path.join(computer.providerRef, "hard-link.txt")); + await symlink( + outsideDirectory, + path.join(computer.providerRef, "escape-directory"), + "junction", + ); + const streamed = (filePath: string) => + desktop.writeFileStream( + computer, + { + path: filePath, + stream: (async function* () { + yield new TextEncoder().encode("after"); + })(), + }, + ctx, + ); + + for (const filePath of [ + "final-link.txt", + "hard-link.txt", + "escape-directory/outside.txt", + "escape-directory/nested/outside.txt", + "../outside.txt", + ]) { + await expect(streamed(filePath)).rejects.toThrow(); + } + expect(await readFile(outsideFile, "utf8")).toBe("outside-before"); + await expect(readFile(path.join(outsideDirectory, "outside.txt"))).rejects.toMatchObject({ + code: "ENOENT", + }); + await expect(readFile(path.join(outsideDirectory, "nested"))).rejects.toMatchObject({ + code: "ENOENT", + }); + + await streamed("notes/result.txt"); + expect(await readFile(path.join(computer.providerRef, "notes/result.txt"), "utf8")).toBe( + "after", + ); + }); + + it("removes a streamed file whose source fails part-way", async () => { + const { desktop, computer } = await fixture("streamed-failure"); + const target = path.join(computer.providerRef, "clip.mp4"); + + await expect( + desktop.writeFileStream( + computer, + { + path: "clip.mp4", + stream: (async function* () { + yield new Uint8Array(1024); + throw new Error("store read failed"); + })(), + }, + ctx, + ), + ).rejects.toThrow("store read failed"); + expect(await pathExists(target)).toBe(false); + }); + it("still creates and replaces ordinary workspace files", async () => { const { desktop, computer } = await fixture("ordinary-files"); diff --git a/packages/adapters/src/desktop-sandbox.ts b/packages/adapters/src/desktop-sandbox.ts index f4c1638..8559e48 100644 --- a/packages/adapters/src/desktop-sandbox.ts +++ b/packages/adapters/src/desktop-sandbox.ts @@ -23,12 +23,14 @@ import type { ComputerRef, ControlLeaseRef, PortableFile, + PortableFileStream, ProcessEvent, SandboxProvider, ScreenRequest, ScreenSession, } from "@rakazo/adapter-kit"; import { boundedSandboxCommandTimeoutMs } from "@rakazo/core"; +import { byteChunks, fileByteStream } from "./byte-streams.js"; import { applyPlaceholderAction, boundedComputerActions, @@ -232,21 +234,51 @@ export class DesktopSandboxProvider implements SandboxProvider { return new Uint8Array(await readFile(target)); } + async readFileStream( + computer: ComputerRef, + filePath: string, + _context?: AdapterContext, + options?: { maxBytes?: number }, + ) { + const box = this.requiredBox(computer); + const target = await localWorkspaceTarget(box.home, filePath, true); + const handle = await open(target, constants.O_RDONLY | O_NOFOLLOW); + try { + const info = await handle.stat(); + if (!info.isFile()) throw new Error("computer file not found"); + if (options?.maxBytes !== undefined && info.size > options.maxBytes) { + throw new Error(`computer file exceeds ${options.maxBytes} bytes`); + } + return fileByteStream(handle, 0, info.size - 1); + } catch (error) { + await handle.close(); + throw error; + } + } + async writeFile(computer: ComputerRef, file: PortableFile) { const box = this.requiredBox(computer); const target = await localWorkspaceTarget(box.home, file.path, false); - const handle = await openContainedWorkspaceFile( + await writeContainedWorkspaceFile(box.home, target, Boolean(file.executable), (handle) => + handle.writeFile(file.content), + ); + } + + async writeFileStream(computer: ComputerRef, file: PortableFileStream, context: AdapterContext) { + const box = this.requiredBox(computer); + const target = await localWorkspaceTarget(box.home, file.path, false); + await writeContainedWorkspaceFile( box.home, target, - file.executable ? 0o700 : 0o600, + Boolean(file.executable), + async (handle) => { + // Each write continues at the file position, so memory stays at one chunk. + for await (const chunk of byteChunks(file.stream)) { + context.signal.throwIfAborted(); + await handle.writeFile(chunk); + } + }, ); - try { - await handle.truncate(0); - await handle.writeFile(file.content); - if (file.executable) await handle.chmod(0o700); - } finally { - await handle.close(); - } } async *exportWorkspace(computer: ComputerRef): AsyncIterable { @@ -416,7 +448,17 @@ async function localWorkspaceTarget(home: string, relative: string, mustExist: b return resolved; } -async function openContainedWorkspaceFile(home: string, target: string, mode: number) { +/** + * Opens `target` without following links outside the workspace, truncates it and runs `write`. + * A file that fails to write is removed again, so no partial file stays behind. + */ +async function writeContainedWorkspaceFile( + home: string, + target: string, + executable: boolean, + write: (handle: ContainedHandle) => Promise, +) { + const mode = executable ? 0o700 : 0o600; const resolvedHome = await realpath(home); const parentPath = path.dirname(target); const name = path.basename(target); @@ -434,6 +476,7 @@ async function openContainedWorkspaceFile(home: string, target: string, mode: nu let handle: ContainedHandle | undefined; let created = false; + let writing = false; let openedPath = path.join(parentReal, name); try { // Bind the parent handle to a still-contained realpath. Otherwise a swap @@ -496,9 +539,15 @@ async function openContainedWorkspaceFile(home: string, target: string, mode: nu throw new Error("Path escapes the computer workspace"); } await assertContainedFileHandle(handle, resolvedHome, parentHandle, containedParent, name); - return handle; + writing = true; + await handle.truncate(0); + await write(handle); + if (executable) await handle.chmod(0o700); + const written = handle; + handle = undefined; + await written.close(); } catch (error) { - if (created && handle) { + if ((created || writing) && handle) { const createdStat = await handle.stat({ bigint: true }).catch(() => undefined); let cleanupPath = openedPath; if (useWin32Relative) { @@ -510,7 +559,7 @@ async function openContainedWorkspaceFile(home: string, target: string, mode: nu } await handle.close().catch(() => undefined); handle = undefined; - // Only unlink if the path still names the inode we created. + // Only unlink if the path still names the inode we created or started writing. if (createdStat) { const present = await lstat(cleanupPath, { bigint: true }).catch(() => undefined); if (present && present.dev === createdStat.dev && present.ino === createdStat.ino) { diff --git a/packages/adapters/src/executor.ts b/packages/adapters/src/executor.ts index 5953f06..80ed17f 100644 --- a/packages/adapters/src/executor.ts +++ b/packages/adapters/src/executor.ts @@ -208,6 +208,7 @@ import { computerFileTransferMaxBytes, displayBotWorkspacePath, filterTeamWorkspaceEntries, + readComputerFileStream, resolveBotWorkspaceCwd, resolveBotWorkspacePath, teamBotWorkspaceDirectory, @@ -2410,7 +2411,7 @@ export function createRunExecutor(deps: ExecutorDeps) { groupId: thread.groupId ?? undefined, runId: run.id, filePath: outPath, - bytes: png, + content: png, operationId: executionId, }, ); @@ -2441,9 +2442,12 @@ export function createRunExecutor(deps: ExecutorDeps) { attachableFileMaxBytes(deps.artifacts, mimeType), computerFileTransferMaxBytes(computer), ); - let bytes: Uint8Array; + let content: ReadableStream; try { - bytes = await deps.sandbox.readFile(computer, storedPath, context, { maxBytes }); + // Streamed, so a large video from a local computer is never held in memory whole. + content = await readComputerFileStream(deps.sandbox, computer, storedPath, context, { + maxBytes, + }); } catch { return finish({ error: `file not found, unreadable or larger than ${formatMiB(maxBytes)}`, @@ -2460,8 +2464,9 @@ export function createRunExecutor(deps: ExecutorDeps) { groupId: thread.groupId ?? undefined, runId: run.id, filePath, - bytes, + content, operationId: executionId, + signal: context.signal, present: args.present === true, }, ); diff --git a/packages/adapters/src/host-aware-sandbox.ts b/packages/adapters/src/host-aware-sandbox.ts index 08574e1..d5533ef 100644 --- a/packages/adapters/src/host-aware-sandbox.ts +++ b/packages/adapters/src/host-aware-sandbox.ts @@ -6,11 +6,13 @@ import type { ComputerRef, ControlLeaseRef, PortableFile, + PortableFileStream, ProcessEvent, SandboxProvider, ScreenRequest, } from "@rakazo/adapter-kit"; import type { PrismaClient } from "@rakazo/db"; +import { readComputerFileStream, writeComputerFileStream } from "./computer-support.js"; import { DesktopSandboxProvider } from "./desktop-sandbox.js"; import { createSandboxProvider, type SandboxProviderOptions } from "./sandbox-factory.js"; @@ -141,6 +143,19 @@ export class HostAwareSandbox implements SandboxProvider { return this.route(computer).writeFile(computer, file, context); } + readFileStream( + computer: ComputerRef, + path: string, + context: AdapterContext, + options?: { maxBytes?: number }, + ) { + return readComputerFileStream(this.route(computer), computer, path, context, options); + } + + writeFileStream(computer: ComputerRef, file: PortableFileStream, context: AdapterContext) { + return writeComputerFileStream(this.route(computer), computer, file, context); + } + exportWorkspace(computer: ComputerRef, context: AdapterContext) { return this.route(computer).exportWorkspace(computer, context); } diff --git a/packages/adapters/src/media-runtime.ts b/packages/adapters/src/media-runtime.ts index 5e26917..d699a40 100644 --- a/packages/adapters/src/media-runtime.ts +++ b/packages/adapters/src/media-runtime.ts @@ -598,7 +598,12 @@ export function createMediaRuntime( }; /** Copies a stored artifact to `file`: a clone of the local file when possible, else a stream. */ - const copyArtifact = async (job: MediaJob, storageKey: string, file: string) => { + const copyArtifact = async ( + job: MediaJob, + storageKey: string, + file: string, + signal: AbortSignal, + ) => { const context = artifactContext(job, `media-input:${job.id}`); const local = await deps.artifacts.localPath?.(storageKey, context).catch(() => undefined); if (local) { @@ -609,6 +614,8 @@ export function createMediaRuntime( await pipeline( Readable.fromWeb(stream as NodeReadableStream), createWriteStream(file, { flags: "wx", mode: 0o600 }), + // A cancelled job stops copying; the job's temp dir (and the partial file) is removed after. + { signal }, ); }; @@ -722,7 +729,7 @@ export function createMediaRuntime( attachmentExtensionForMimeType(row.mimeType) || (/^\.[a-z0-9]{1,8}$/i.test(path.extname(row.name)) ? path.extname(row.name) : ""); const file = path.join(tempDir, `input-${row.id}${extension.toLowerCase()}`); - await copyArtifact(job, row.storageKey, file); + await copyArtifact(job, row.storageKey, file, signal); const input: MediaJobInputFile = { artifactId: row.id, path: file, diff --git a/packages/adapters/src/thread-artifacts.test.ts b/packages/adapters/src/thread-artifacts.test.ts index d3d355e..075d14b 100644 --- a/packages/adapters/src/thread-artifacts.test.ts +++ b/packages/adapters/src/thread-artifacts.test.ts @@ -1,13 +1,33 @@ 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 { describe, expect, it, vi } from "vitest"; +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, @@ -15,6 +35,15 @@ import { materializeCurrentTurnFiles, } from "./thread-artifacts.js"; +function byteStream(chunks: Uint8Array[]) { + return new ReadableStream({ + 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); @@ -38,7 +67,7 @@ describe("current-turn thread files", () => { botId: "bot-1", runId: "run-1", filePath: "report.pdf", - bytes: new Uint8Array([1, 2, 3]), + content: new Uint8Array([1, 2, 3]), operationId: "attach-1", }, ), @@ -73,7 +102,7 @@ describe("current-turn thread files", () => { botId: "bot-1", runId: "run-1", filePath: "out/report.md", - bytes: new TextEncoder().encode("# Report\n"), + content: new TextEncoder().encode("# Report\n"), operationId: "attach-1", }; @@ -95,7 +124,9 @@ describe("current-turn thread files", () => { storageKey: "stored-1", }, ]); - const get = vi.fn().mockResolvedValue(new Uint8Array([1, 2, 3, 4])); + 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 } = { @@ -126,14 +157,15 @@ describe("current-turn thread files", () => { const files = await materializeCurrentTurnFiles( { prisma: { artifact: { findMany } } as unknown as PrismaClient, - artifacts: { get } as unknown as ArtifactStore, + artifacts: { createReadStream } as unknown as ArtifactStore, sandbox: { writeFile } as unknown as SandboxProvider, }, blocks, { context, computer, computerMode: "team", markWorkspaceDirty }, ); - expect(get).toHaveBeenCalledWith("stored-1", context); + // A provider without streaming writes still receives the whole file as bytes. + expect(createReadStream).toHaveBeenCalledWith("stored-1", context); expect(writeFile).toHaveBeenCalledWith( computer, { @@ -165,7 +197,7 @@ describe("current-turn thread files", () => { { 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 createReadStream = vi.fn(async () => byteStream([new Uint8Array([1, 2, 3])])); const writeFile = vi.fn().mockResolvedValue(undefined); const context: AdapterContext & { botId: string } = { operationId: "run-1", @@ -189,7 +221,7 @@ describe("current-turn thread files", () => { ]; const deps = { prisma: { artifact: { findMany } } as unknown as PrismaClient, - artifacts: { get } as unknown as ArtifactStore, + artifacts: { createReadStream } as unknown as ArtifactStore, sandbox: { writeFile } as unknown as SandboxProvider, }; @@ -203,7 +235,7 @@ describe("current-turn thread files", () => { ["aud-1", "attachments/aud-1.mp3"], ["vid-1", null], ]); - expect(get).not.toHaveBeenCalledWith("k-vid", context); + 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"); @@ -265,6 +297,254 @@ describe("current-turn thread files", () => { }); }); +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 = async () => {}) { + let index = 0; + const cancel = vi.fn(); + const stream = new ReadableStream( + { + 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 }) => ({ + 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 diff --git a/packages/adapters/src/thread-artifacts.ts b/packages/adapters/src/thread-artifacts.ts index ef976d3..0903f7b 100644 --- a/packages/adapters/src/thread-artifacts.ts +++ b/packages/adapters/src/thread-artifacts.ts @@ -5,6 +5,7 @@ import type { ComputerRef, SandboxProvider, } from "@rakazo/adapter-kit"; +import { ArtifactTooLargeError } from "@rakazo/adapter-kit"; import type { ComputerMode, MessageBlock } from "@rakazo/contracts"; import { ATTACHMENT_MAX_BYTES } from "@rakazo/contracts"; import { @@ -14,10 +15,12 @@ import { } from "@rakazo/core"; import type { PrismaClient } from "@rakazo/db"; import { artifactMaxBytes, storeArtifactStream } from "./artifact-files.js"; +import { discardByteStream } from "./byte-streams.js"; import { assertTeamWorkspacePathAllowed, computerFileTransferMaxBytes, resolveBotWorkspacePath, + writeComputerFileStream, } from "./computer-support.js"; export type MaterializedThreadFile = { @@ -56,35 +59,47 @@ export async function attachWorkspaceFileToThread( groupId?: string; runId: string; filePath: string; - bytes: Uint8Array; + /** A stream (e.g. from readComputerFileStream) is stored without buffering the file. */ + content: Uint8Array | ReadableStream; operationId: string; + signal?: AbortSignal; present?: boolean; }, ): Promise<{ artifactId: string; block: AttachmentBlock }> { + const content = input.content; + const stream = + content instanceof Uint8Array + ? (async function* () { + yield content; + })() + : content; const name = path.basename(input.filePath) || input.filePath; - const mimeType = inferAttachmentMimeType(name); - if (!mimeType) { - throw new Error(`Unsupported attachment type for ${name}`); + let row: Awaited>; + try { + const mimeType = inferAttachmentMimeType(name); + if (!mimeType) { + throw new Error(`Unsupported attachment type for ${name}`); + } + row = await storeArtifactStream(deps, { + spaceId: input.spaceId, + userId: input.userId, + botId: input.botId, + groupId: input.groupId ?? null, + runId: input.runId, + name, + mimeType, + stream, + maxBytes: attachableFileMaxBytes(deps.artifacts, mimeType), + operationId: input.operationId, + signal: input.signal, + }); + } catch (error) { + discardByteStream(stream); + if (error instanceof ArtifactTooLargeError) { + throw new Error(`file exceeds the ${formatMiB(error.maxBytes)} attachment limit`); + } + throw error; } - 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, - }); const block = messageBlockForArtifact(row); const presentable = block.kind === "image" || block.kind === "file"; return { @@ -137,7 +152,6 @@ export async function materializeCurrentTurnFiles( 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, @@ -145,13 +159,13 @@ export async function materializeCurrentTurnFiles( relativePath, ); if (input.teamBotIds) assertTeamWorkspacePathAllowed(storedPath, input.teamBotIds); + // Streamed, so a large video copied to a local computer never sits in memory whole. + const stream = await deps.artifacts.createReadStream(row.storageKey, input.context); input.markWorkspaceDirty?.(); - await deps.sandbox.writeFile( + await writeComputerFileStream( + deps.sandbox, input.computer, - { - path: storedPath, - content: bytes, - }, + { path: storedPath, stream }, input.context, ); materialized.push({ ...entry, path: relativePath });