From 78d534b0763c09eea301aeb0bf8a84dd4db7ac84 Mon Sep 17 00:00:00 2001 From: Aniruddha Adak Date: Sun, 23 Aug 2026 04:48:12 +0530 Subject: [PATCH] feat(sandbox): stream file downloads end to end --- .changeset/stream-sandbox-file-downloads.md | 6 + packages/trueforge-core/src/core/index.ts | 2 + .../core/sandbox/provider/DaytonaProvider.ts | 52 +++- .../src/core/sandbox/provider/Provider.ts | 28 +- .../sandbox/provider/TFYSandboxProvider.ts | 83 ++++-- .../sandbox/provider/boundedFileStream.ts | 49 ++++ .../core/sandbox/daytonaDownload.test.ts | 112 ++++++++ .../provider/boundedFileStream.test.ts | 89 ++++++ .../provider/sandboxProviderContractSuite.ts | 45 ++- .../core/sandbox/provider/tfyDownload.test.ts | 134 +++++++++ packages/trueforge/src/apis/turns.ts | 26 +- .../src/sandbox/local/core/hostRun.ts | 262 +++++++++++++++--- .../local/provider/LocalSandboxProvider.ts | 69 +++-- .../tests/sandbox/local/smoke.test.ts | 19 +- .../unit/apis/sandboxFileDownload.test.ts | 105 ++++++- 15 files changed, 969 insertions(+), 112 deletions(-) create mode 100644 .changeset/stream-sandbox-file-downloads.md create mode 100644 packages/trueforge-core/src/core/sandbox/provider/boundedFileStream.ts create mode 100644 packages/trueforge-core/tests/core/sandbox/daytonaDownload.test.ts create mode 100644 packages/trueforge-core/tests/core/sandbox/provider/boundedFileStream.test.ts create mode 100644 packages/trueforge-core/tests/core/sandbox/provider/tfyDownload.test.ts diff --git a/.changeset/stream-sandbox-file-downloads.md b/.changeset/stream-sandbox-file-downloads.md new file mode 100644 index 000000000..c116a32da --- /dev/null +++ b/.changeset/stream-sandbox-file-downloads.md @@ -0,0 +1,6 @@ +--- +"@truefoundry/trueforge-core": patch +"@truefoundry/trueforge": patch +--- + +Stream sandbox file downloads end to end. `SandboxProvider.downloadFile` now returns `{ size, stream }` so the download route can send bytes incrementally (bounded by chunk size, not file size), keep the exact `Content-Length`, stop provider reads when the client disconnects, and re-enforce the size cap while streaming. diff --git a/packages/trueforge-core/src/core/index.ts b/packages/trueforge-core/src/core/index.ts index 5ebbd311a..560fc5199 100644 --- a/packages/trueforge-core/src/core/index.ts +++ b/packages/trueforge-core/src/core/index.ts @@ -142,6 +142,7 @@ export type { CodeModeLogger } from './sandbox/codeMode/CodeModeDispatcher'; export type { CodeModeClientInstall, CodeModeTransport } from './sandbox/codeMode/CodeModeTransport'; export { CodeModeErrorSourceSchema, CodeModeReplySchema, CodeModeRequestSchema } from './sandbox/codeMode/types'; export type { CodeModeErrorSource, CodeModeReply, CodeModeRequest } from './sandbox/codeMode/types'; +export { boundedFileStream } from './sandbox/provider/boundedFileStream'; export { DaytonaSandboxProvider } from './sandbox/provider/DaytonaProvider'; export type { DaytonaSandboxProviderOptions } from './sandbox/provider/DaytonaProvider'; export { absolutizeRelativeExecEnv } from './sandbox/provider/execEnv'; @@ -154,6 +155,7 @@ export type { SandboxBuildMetadata, SandboxBuildStatus, SandboxExecParams, + SandboxFileDownload, SandboxInit, SandboxProvider, } from './sandbox/provider/Provider'; diff --git a/packages/trueforge-core/src/core/sandbox/provider/DaytonaProvider.ts b/packages/trueforge-core/src/core/sandbox/provider/DaytonaProvider.ts index 1fba52e16..896438e9e 100644 --- a/packages/trueforge-core/src/core/sandbox/provider/DaytonaProvider.ts +++ b/packages/trueforge-core/src/core/sandbox/provider/DaytonaProvider.ts @@ -6,6 +6,9 @@ import { randomUUID } from 'node:crypto'; import { join } from 'node:path/posix'; import type { Logger } from 'winston'; import { extractErrorLogFields } from '../../util/errorLogFields'; +import type { CodeModeTransport } from '../codeMode/CodeModeTransport'; +import { CodeModeNatsTransport } from '../codeMode/nats/CodeModeNatsTransport'; +import { DEFAULT_PREVIEW_URL_EXPIRY_SECONDS, DEFAULT_SANDBOX_NATS_WS_PORT } from '../constants'; import { SandboxFileNotFoundError, SandboxFileTooLargeError, @@ -13,10 +16,15 @@ import { SandboxPathIsDirectoryError, validateSandboxOwnedByTenant, } from '../SandboxErrors'; -import type { CodeModeTransport } from '../codeMode/CodeModeTransport'; -import { CodeModeNatsTransport } from '../codeMode/nats/CodeModeNatsTransport'; -import { DEFAULT_PREVIEW_URL_EXPIRY_SECONDS, DEFAULT_SANDBOX_NATS_WS_PORT } from '../constants'; -import type { ExecResult, SandboxBuild, SandboxExecParams, SandboxFileInfo, SandboxProvider } from './Provider'; +import { boundedFileStream } from './boundedFileStream'; +import type { + ExecResult, + SandboxBuild, + SandboxExecParams, + SandboxFileDownload, + SandboxFileInfo, + SandboxProvider, +} from './Provider'; const SANDBOX_NOT_FOUND_STATUS = 404; /** Another replica already registered this build name; its create is the one that counts. */ @@ -415,9 +423,15 @@ export class DaytonaSandboxProvider implements SandboxProvider { return { size: details.size, isDir: details.isDir }; } - async downloadFile(params: { sandboxId: string; path: string }): Promise { + async downloadFile(params: { + sandboxId: string; + path: string; + signal?: AbortSignal | undefined; + }): Promise { return context.with(suppressTracing(context.active()), async () => { try { + // Stat + stream-open happen here so domain errors (404 / dir / too large) reject the + // promise while the route can still answer with a JSON status instead of a broken body. return await this.executeWithSandboxRecovery(params.sandboxId, async () => { const { sandbox } = await this.getOrCreateSandbox(params.sandboxId); @@ -429,7 +443,33 @@ export class DaytonaSandboxProvider implements SandboxProvider { throw new SandboxFileTooLargeError(params.path, info.size, this.fileMaxBytesForDownload); } - return await sandbox.fs.downloadFile(params.path); + const nodeStream = await sandbox.fs.downloadFileStream(params.path); + if (params.signal !== undefined) { + params.signal.addEventListener( + 'abort', + () => { + nodeStream.destroy(new Error(`Download of ${params.path} aborted by the client`)); + }, + { once: true }, + ); + } + return { + size: info.size, + stream: boundedFileStream({ + path: params.path, + maxBytes: this.fileMaxBytesForDownload, + chunks: async function* () { + try { + for await (const chunk of nodeStream) { + yield chunk; + } + } finally { + // No-op on natural completion; stops the transfer when cancelled or errored. + nodeStream.destroy(); + } + }, + }), + }; }); } catch (e: unknown) { if (e instanceof SandboxPathIsDirectoryError || e instanceof SandboxFileTooLargeError) { diff --git a/packages/trueforge-core/src/core/sandbox/provider/Provider.ts b/packages/trueforge-core/src/core/sandbox/provider/Provider.ts index 6801c239d..95aae7b65 100644 --- a/packages/trueforge-core/src/core/sandbox/provider/Provider.ts +++ b/packages/trueforge-core/src/core/sandbox/provider/Provider.ts @@ -64,6 +64,21 @@ export interface SandboxBuild { metadata: SandboxBuildMetadata | null; } +/** An opened sandbox file ready to be streamed to a client. */ +export interface SandboxFileDownload { + /** + * Byte count from the provider's pre-download stat. Drives the response Content-Length; + * providers enforce this same bound again while streaming in case the file changed. + */ + size: number; + /** + * File bytes, single consumption. Domain errors raised before the first byte propagate to the + * caller as a rejected promise (mappable HTTP statuses); anything failing after that surfaces + * as an errored stream — a truncated body the client can detect against Content-Length. + */ + stream: ReadableStream; +} + export interface SandboxProvider { /** Stable provider kind used in fancy sandbox ids and carry-forward (plain string). */ readonly type: string; @@ -93,8 +108,17 @@ export interface SandboxProvider { getSkillsDir(sandboxId: string): string; /** Path the git skill downloader script is written to before it runs. */ getGitDownloaderPath(sandboxId: string): string; - /** Downloads a file from the sandbox as a Buffer. Throws SandboxFileNotFoundError / SandboxNotAvailableError / SandboxPathIsDirectoryError / SandboxFileTooLargeError. */ - downloadFile(params: { sandboxId: string; path: string }): Promise; + /** + * Opens a file for streaming download. Implementations stat first (path checks, is-dir, + * max-size) so domain errors reject before any byte moves, then hand back an incremental + * stream — peak memory stays bounded by chunk size rather than file size. + * `signal` aborts provider reads once the HTTP client goes away mid-stream. + */ + downloadFile(params: { + sandboxId: string; + path: string; + signal?: AbortSignal | undefined; + }): Promise; /** Uploads a file to the sandbox. */ uploadFile(params: { sandboxId: string; remotePath: string; content: Buffer }): Promise; diff --git a/packages/trueforge-core/src/core/sandbox/provider/TFYSandboxProvider.ts b/packages/trueforge-core/src/core/sandbox/provider/TFYSandboxProvider.ts index 4dc6a05f2..b9196043f 100644 --- a/packages/trueforge-core/src/core/sandbox/provider/TFYSandboxProvider.ts +++ b/packages/trueforge-core/src/core/sandbox/provider/TFYSandboxProvider.ts @@ -14,6 +14,7 @@ import { SandboxPathIsDirectoryError, validateSandboxOwnedByTenant, } from '../SandboxErrors'; +import { boundedFileStream } from './boundedFileStream'; import { absolutizeRelativeExecEnv } from './execEnv'; import { ensureExecSuccess, @@ -21,6 +22,7 @@ import { type ExecResult, type SandboxBuild, type SandboxExecParams, + type SandboxFileDownload, type SandboxFileInfo, type SandboxProvider, } from './Provider'; @@ -28,6 +30,8 @@ import { const DEFAULT_TIMEOUT_SECONDS = 60; // Buffer for network latency + response processing on top of the server-side timeout. const CLIENT_TIMEOUT_BUFFER_SECONDS = 5; +/** Chunk size for streamed downloads — bounds per-exec memory while amortizing round-trips. */ +const TFY_DOWNLOAD_CHUNK_BYTES = 2 * 1024 * 1024; const TFY_MCP_CLIENT_BIN = 'mcp-client/bin'; @@ -107,6 +111,11 @@ export class TFYSandboxProvider implements SandboxProvider { * git and `mcp-client` work after `cd` — resolved from `pwd` / `$PATH`, not a hardcoded prefix. */ async exec(params: SandboxExecParams): Promise { + return this.runInSandbox(params); + } + + /** Shared exec body: tenant check, layout discovery, env absolutization, then the HTTP exec. */ + private async runInSandbox(params: SandboxExecParams & { signal?: AbortSignal | undefined }): Promise { validateSandboxOwnedByTenant({ sandboxId: params.sandboxId, tenantName: this.tenantName }); const { root, inheritedPath } = await this.sandboxLayout(params.sandboxId); const incoming = params.env ?? {}; @@ -116,7 +125,6 @@ export class TFYSandboxProvider implements SandboxProvider { }); return this.postExec({ ...params, env }); } - private async sandboxLayout(sandboxId: string): Promise<{ root: string; inheritedPath: string }> { const cached = this.sandboxLayouts.get(sandboxId); if (cached !== undefined) { @@ -137,7 +145,7 @@ export class TFYSandboxProvider implements SandboxProvider { return parsed; } - private async postExec(params: SandboxExecParams): Promise { + private async postExec(params: SandboxExecParams & { signal?: AbortSignal | undefined }): Promise { return context.with(suppressTracing(context.active()), async (): Promise => { // Honor the per-call timeout (e.g. the longer skill-download timeout) for both the server-side // request timeout and the client abort timer; fall back to the provider default otherwise. @@ -155,13 +163,16 @@ export class TFYSandboxProvider implements SandboxProvider { const timer = setTimeout(() => { controller.abort(); }, clientTimeoutMs); + // One composite signal so either the timeout or a caller abort stops the in-flight fetch. + const requestSignal = + params.signal === undefined ? controller.signal : AbortSignal.any([controller.signal, params.signal]); try { const response = await fetch(`${this.serverUrl}/exec`, { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify(body), - signal: controller.signal, + signal: requestSignal, }); if (!response.ok) { @@ -173,6 +184,10 @@ export class TFYSandboxProvider implements SandboxProvider { const result = (await response.json()) as ExecResult; return result; } catch (e: unknown) { + if (params.signal?.aborted === true) { + this.logger.info('Sandbox exec aborted by the caller'); + return { success: false, error: 'Sandbox exec aborted by the caller' }; + } if (e instanceof Error && e.name === 'AbortError') { this.logger.error(`Sandbox exec timed out after ${String(timeoutSeconds)}s`, extractErrorLogFields(e)); return { success: false, error: `Sandbox exec timed out after ${String(timeoutSeconds)}s` }; @@ -203,7 +218,11 @@ export class TFYSandboxProvider implements SandboxProvider { return { size: parsed.size, isDir: parsed.type === 'directory' }; } - async downloadFile(params: { sandboxId: string; path: string }): Promise { + async downloadFile(params: { + sandboxId: string; + path: string; + signal?: AbortSignal | undefined; + }): Promise { const info = await this.getFileInfo(params); if (info.isDir) { throw new SandboxPathIsDirectoryError(params.path); @@ -212,19 +231,49 @@ export class TFYSandboxProvider implements SandboxProvider { throw new SandboxFileTooLargeError(params.path, info.size, this.fileMaxBytesForDownload); } - const result = await this.exec({ - sandboxId: params.sandboxId, - command: `base64 -w0 ${shellEscape(params.path)}`, - }); - - if (!result.success) { - throw new Error(`Failed to download file: ${result.error}`); - } - if (result.response.exitCode !== 0) { - throw new SandboxFileNotFoundError(params.path); - } - - return Buffer.from(result.response.result.trim(), 'base64'); + // The sandbox server has no file endpoint, so bytes are pulled as base64 through sequential + // execs over fixed byte windows (`tail -c +N | head -c M`): memory stays bounded per chunk, + // the client receives bytes before the whole file is read, and an abort stops further execs. + const { sandboxId, path, signal } = params; + const execWindow = (offset: number, length: number): Promise => + this.runInSandbox({ + sandboxId, + command: `tail -c +${String(offset + 1)} ${shellEscape(path)} | head -c ${String(length)} | base64 -w0`, + ...(signal === undefined ? {} : { signal }), + }); + let offset = 0; + const chunks = async function* (): AsyncGenerator { + while (offset < info.size) { + const length = Math.min(TFY_DOWNLOAD_CHUNK_BYTES, info.size - offset); + const result = await execWindow(offset, length); + if (!result.success) { + throw new Error(`Failed to download file: ${result.error}`); + } + if (result.response.exitCode !== 0) { + // First window failing means the stat target is gone (or unreadable); later failures + // mean the file mutated mid-download, which no HTTP status can express. + if (offset === 0) { + throw new SandboxFileNotFoundError(path); + } + throw new Error(`Sandbox file changed during download: ${path}`); + } + const decoded = Buffer.from(result.response.result.trim(), 'base64'); + if (decoded.byteLength !== length) { + throw new Error(`Sandbox file changed during download: ${path}`); + } + offset += length; + yield decoded; + } + }; + + return { + size: info.size, + stream: boundedFileStream({ + path, + maxBytes: this.fileMaxBytesForDownload, + chunks, + }), + }; } async uploadFile(params: { sandboxId: string; remotePath: string; content: Buffer }): Promise { diff --git a/packages/trueforge-core/src/core/sandbox/provider/boundedFileStream.ts b/packages/trueforge-core/src/core/sandbox/provider/boundedFileStream.ts new file mode 100644 index 000000000..431fed8f0 --- /dev/null +++ b/packages/trueforge-core/src/core/sandbox/provider/boundedFileStream.ts @@ -0,0 +1,49 @@ +import { SandboxFileTooLargeError } from '../SandboxErrors'; + +export interface BoundedFileStreamParams { + /** User-visible sandbox path used in cap-violation errors. */ + path: string; + /** Hard byte ceiling; also enforced here because the pre-download stat can go stale. */ + maxBytes: number; + /** + * Provider-specific incremental source, created fresh per call so the generator's own + * `finally` cleanup runs on completion, source errors, and consumer cancellation alike. + */ + chunks: () => AsyncGenerator; +} + +/** + * Wraps incremental file chunks into a single-consumption web stream that enforces the + * download cap while streaming. Consumer cancellation (`stream.cancel()`) returns the + * underlying generator early, which runs its cleanup — providers stop their reads there. + */ +export function boundedFileStream(params: BoundedFileStreamParams): ReadableStream { + const iterator = params.chunks(); + let streamed = 0; + return new ReadableStream({ + async pull(controller) { + let next: IteratorResult; + try { + next = await iterator.next(); + } catch (error) { + controller.error(error instanceof Error ? error : new Error(String(error))); + return; + } + if (next.done) { + controller.close(); + return; + } + streamed += next.value.byteLength; + if (streamed > params.maxBytes) { + // Held back rather than enqueued: the client never sees past-the-cap bytes. + await iterator.return(undefined); + controller.error(new SandboxFileTooLargeError(params.path, streamed, params.maxBytes)); + return; + } + controller.enqueue(next.value); + }, + cancel() { + void iterator.return(undefined); + }, + }); +} diff --git a/packages/trueforge-core/tests/core/sandbox/daytonaDownload.test.ts b/packages/trueforge-core/tests/core/sandbox/daytonaDownload.test.ts new file mode 100644 index 000000000..b4b28bb5e --- /dev/null +++ b/packages/trueforge-core/tests/core/sandbox/daytonaDownload.test.ts @@ -0,0 +1,112 @@ +import type { Sandbox } from '@daytona/sdk'; +import { Daytona, DaytonaError } from '@daytona/sdk'; +import { Readable } from 'node:stream'; +import { DaytonaSandboxProvider } from '../../../src/core/sandbox/provider/DaytonaProvider'; +import { + SandboxFileNotFoundError, + SandboxFileTooLargeError, + SandboxPathIsDirectoryError, +} from '../../../src/core/sandbox/SandboxErrors'; +import { makeSilentLogger } from '../harnessMocks'; + +const NOT_FOUND_STATUS = 404; +const API_URL = 'https://daytona.test/api'; +const SANDBOX_ID = 'test-tenant.dl-sandbox'; + +interface FsMockOverrides { + details?: { size: number; isDir: boolean }; + fileNotFoundError?: boolean; + stream?: Readable; +} + +/** Provider wired to a fake SDK client whose `get` resolves a stub sandbox with mocked fs. */ +function makeProviderWithFs(overrides: FsMockOverrides): DaytonaSandboxProvider { + const client = new Daytona({ apiKey: 'dtn-test', useDeprecatedPolling: true }); + const sandbox = { + state: 'started', + refreshData: jest.fn(), + start: jest.fn(), + fs: { + getFileDetails: jest.fn(() => { + if (overrides.fileNotFoundError === true) { + throw new DaytonaError('no such file', NOT_FOUND_STATUS); + } + return Promise.resolve(overrides.details ?? { size: 0, isDir: false }); + }), + downloadFileStream: jest.fn(() => Promise.resolve(overrides.stream ?? Readable.from([]))), + }, + }; + jest.spyOn(client, 'get').mockResolvedValue(sandbox as unknown as Sandbox); + + return new DaytonaSandboxProvider({ + client, + apiKey: 'dtn-test', + apiUrl: API_URL, + tenantName: 'test-tenant', + sandboxImage: 'registry.example.com/sandbox:029ea5ff', + timeoutMs: 1000, + autoStopIntervalInMinutes: 5, + autoArchiveIntervalInMinutes: 60, + autoDeleteIntervalInMinutes: 7200, + fileMaxBytesForDownload: 1024, + logger: makeSilentLogger(), + }); +} + +async function drain(download: Awaited>): Promise { + const reader = download.stream.getReader(); + const parts: Uint8Array[] = []; + while (true) { + const next = await reader.read(); + if (next.done) { + break; + } + parts.push(next.value); + } + return Buffer.concat(parts.map(part => Buffer.from(part))); +} + +describe('DaytonaSandboxProvider downloadFile streaming', () => { + it('streams the file incrementally and reports the stat size', async () => { + const provider = makeProviderWithFs({ + details: { size: 11, isDir: false }, + stream: Readable.from([Buffer.from('hello '), Buffer.from('world')]), + }); + + const download = await provider.downloadFile({ sandboxId: SANDBOX_ID, path: 'report.txt' }); + + expect(download.size).toBe(11); + expect((await drain(download)).toString()).toBe('hello world'); + }); + + it('rejects directories and oversized files before opening the stream', async () => { + const directoryProvider = makeProviderWithFs({ details: { size: 10, isDir: true } }); + await expect(directoryProvider.downloadFile({ sandboxId: SANDBOX_ID, path: 'dir' })).rejects.toBeInstanceOf( + SandboxPathIsDirectoryError, + ); + + const oversizedProvider = makeProviderWithFs({ details: { size: 2048, isDir: false } }); + await expect(oversizedProvider.downloadFile({ sandboxId: SANDBOX_ID, path: 'big.bin' })).rejects.toBeInstanceOf( + SandboxFileTooLargeError, + ); + }); + + it('maps a missing file to SandboxFileNotFoundError', async () => { + const provider = makeProviderWithFs({ fileNotFoundError: true }); + + await expect(provider.downloadFile({ sandboxId: SANDBOX_ID, path: 'gone.txt' })).rejects.toBeInstanceOf( + SandboxFileNotFoundError, + ); + }); + + it('enforces the cap again while streaming in case the file grew after stat', async () => { + // Stat lies small; the actual stream exceeds fileMaxBytesForDownload (1024). + const provider = makeProviderWithFs({ + details: { size: 8, isDir: false }, + stream: Readable.from([Buffer.alloc(700, 1), Buffer.alloc(700, 2)]), + }); + + const download = await provider.downloadFile({ sandboxId: SANDBOX_ID, path: 'grew.bin' }); + await expect(drain(download)).rejects.toBeInstanceOf(SandboxFileTooLargeError); + }); +}); diff --git a/packages/trueforge-core/tests/core/sandbox/provider/boundedFileStream.test.ts b/packages/trueforge-core/tests/core/sandbox/provider/boundedFileStream.test.ts new file mode 100644 index 000000000..8fe70cc32 --- /dev/null +++ b/packages/trueforge-core/tests/core/sandbox/provider/boundedFileStream.test.ts @@ -0,0 +1,89 @@ +import { boundedFileStream } from '../../../../src/core/sandbox/provider/boundedFileStream'; +import { SandboxFileTooLargeError } from '../../../../src/core/sandbox/SandboxErrors'; + +async function drain(stream: ReadableStream): Promise<{ parts: Uint8Array[]; error?: unknown }> { + const reader = stream.getReader(); + const parts: Uint8Array[] = []; + try { + while (true) { + const next = await reader.read(); + if (next.done) { + return { parts }; + } + parts.push(next.value); + } + } catch (error) { + return { parts, error }; + } finally { + reader.releaseLock(); + } +} + +function sourceOf(chunks: Buffer[], hooks?: { onFinally?: () => void }): () => AsyncGenerator { + return async function* () { + try { + for (const chunk of chunks) { + yield chunk; + } + } finally { + hooks?.onFinally?.(); + } + }; +} + +describe('boundedFileStream', () => { + it('forwards chunks in order without copying', async () => { + const chunks = [Buffer.from('he'), Buffer.from('ll'), Buffer.from('o')]; + const stream = boundedFileStream({ path: 'f.bin', maxBytes: 100, chunks: sourceOf(chunks) }); + + const { parts, error } = await drain(stream); + + expect(error).toBeUndefined(); + expect(parts.map(part => Buffer.from(part).toString())).toEqual(['he', 'll', 'o']); + }); + + it('errors with SandboxFileTooLargeError when streamed bytes exceed the cap', async () => { + const chunks = [Buffer.alloc(60, 1), Buffer.alloc(60, 2)]; + const stream = boundedFileStream({ path: 'big.bin', maxBytes: 100, chunks: sourceOf(chunks) }); + + const { parts, error } = await drain(stream); + + // The offending chunk is withheld: the client never sees past-the-cap bytes. + expect(parts).toHaveLength(1); + expect(error).toBeInstanceOf(SandboxFileTooLargeError); + expect((error as SandboxFileTooLargeError).fileSize).toBe(120); + }); + + it('propagates source failures as stream errors', async () => { + const stream = boundedFileStream({ + path: 'f.bin', + maxBytes: 100, + chunks: async function* () { + yield Buffer.from('ok'); + throw new Error('backend exploded'); + }, + }); + + const { parts, error } = await drain(stream); + + expect(parts.map(part => Buffer.from(part).toString())).toEqual(['ok']); + expect(error).toBeInstanceOf(Error); + expect((error as Error).message).toBe('backend exploded'); + }); + + it('runs source cleanup when the consumer cancels mid-stream', async () => { + let cleanedUp = false; + const chunks = [Buffer.from('first'), Buffer.from('second'), Buffer.from('third')]; + const stream = boundedFileStream({ + path: 'f.bin', + maxBytes: 100, + chunks: sourceOf(chunks, { onFinally: () => (cleanedUp = true) }), + }); + + const reader = stream.getReader(); + await reader.read(); + await reader.cancel(); + + expect(cleanedUp).toBe(true); + }); +}); diff --git a/packages/trueforge-core/tests/core/sandbox/provider/sandboxProviderContractSuite.ts b/packages/trueforge-core/tests/core/sandbox/provider/sandboxProviderContractSuite.ts index ca93ddb7c..55d15f6f4 100644 --- a/packages/trueforge-core/tests/core/sandbox/provider/sandboxProviderContractSuite.ts +++ b/packages/trueforge-core/tests/core/sandbox/provider/sandboxProviderContractSuite.ts @@ -1,5 +1,6 @@ import { basename, isAbsolute, join } from 'node:path'; -import type { SandboxProvider } from '../../../../src/core/sandbox/provider/Provider'; +import { SandboxFileNotFoundError, SandboxPathIsDirectoryError } from '../../../../src/core/sandbox/SandboxErrors'; +import type { SandboxFileDownload, SandboxProvider } from '../../../../src/core/sandbox/provider/Provider'; import { ensureExecSuccess } from '../../../../src/core/sandbox/provider/Provider'; export interface SandboxProviderContractFixture { @@ -7,6 +8,18 @@ export interface SandboxProviderContractFixture { dispose: () => Promise; } +/** Drains a download stream into one Buffer so callers can assert on exact bytes. */ +export async function collectDownload(download: SandboxFileDownload): Promise { + const reader = download.stream.getReader(); + const chunks: Uint8Array[] = []; + let next = await reader.read(); + while (!next.done) { + chunks.push(next.value); + next = await reader.read(); + } + return Buffer.concat(chunks.map(chunk => Buffer.from(chunk))); +} + /** * SandboxProvider contract suite — factory-injected so backends can reuse it. * Does not exercise `createCodeModeTransport` (NATS vs UDS is covered by the Code Mode transport suite). @@ -135,19 +148,41 @@ export function runSandboxProviderContractSuite( } }); - it('upload then download round-trips bytes', async () => { + it('upload then download round-trips bytes incrementally', async () => { const { sandboxId } = await fixture.provider.createSandbox(); - const payload = Buffer.from('upload-download-contract\n', 'utf8'); + // Large enough that the transport must deliver several chunks (progressive delivery). + const payload = Buffer.alloc(1024 * 1024); + for (let index = 0; index < payload.length; index += 1) { + payload[index] = index % 251; + } await fixture.provider.uploadFile({ sandboxId, remotePath: 'roundtrip.bin', content: payload, }); - const downloaded = await fixture.provider.downloadFile({ + const download = await fixture.provider.downloadFile({ sandboxId, path: 'roundtrip.bin', }); - expect(Buffer.compare(downloaded, payload)).toBe(0); + expect(download.size).toBe(payload.length); + const collected = await collectDownload(download); + expect(Buffer.compare(collected, payload)).toBe(0); + }); + + it('download reports missing files with SandboxFileNotFoundError', async () => { + const { sandboxId } = await fixture.provider.createSandbox(); + await expect(fixture.provider.downloadFile({ sandboxId, path: 'no-such-file.bin' })).rejects.toBeInstanceOf( + SandboxFileNotFoundError, + ); + }); + + it('download rejects directories with SandboxPathIsDirectoryError', async () => { + const { sandboxId } = await fixture.provider.createSandbox(); + const mkdir = await fixture.provider.exec({ sandboxId, command: 'mkdir -p a-directory' }); + ensureExecSuccess(mkdir); + await expect(fixture.provider.downloadFile({ sandboxId, path: 'a-directory' })).rejects.toBeInstanceOf( + SandboxPathIsDirectoryError, + ); }); }); } diff --git a/packages/trueforge-core/tests/core/sandbox/provider/tfyDownload.test.ts b/packages/trueforge-core/tests/core/sandbox/provider/tfyDownload.test.ts new file mode 100644 index 000000000..9547288d3 --- /dev/null +++ b/packages/trueforge-core/tests/core/sandbox/provider/tfyDownload.test.ts @@ -0,0 +1,134 @@ +import type { ExecResult } from '../../../../src/core/sandbox/provider/Provider'; +import { TFYSandboxProvider } from '../../../../src/core/sandbox/provider/TFYSandboxProvider'; +import { SandboxFileNotFoundError } from '../../../../src/core/sandbox/SandboxErrors'; +import { makeSilentLogger } from '../../harnessMocks'; + +const TENANT = 'acme'; +const SANDBOX_ID = `${TENANT}.00000000-0000-0000-0000-000000000001`; +/** Mirrors TFY_DOWNLOAD_CHUNK_BYTES so tests can stage multi-window payloads. */ +const CHUNK_BYTES = 2 * 1024 * 1024; + +function makeProvider(fileMaxBytesForDownload: number): TFYSandboxProvider { + return new TFYSandboxProvider({ + serverUrl: 'http://sandbox.example', + natsBridgeUrl: 'ws://nats.example', + tenantName: TENANT, + fileMaxBytesForDownload, + logger: makeSilentLogger(), + }); +} + +function execOk(exitCode: number, result: string): ExecResult { + return { success: true, response: { exitCode, result } }; +} + +/** Stages a queue of /exec JSON responses served by the mocked fetch, recording commands. */ +function mockSandboxServer(responses: ExecResult[]): { commands: () => string[] } { + const seen: string[] = []; + const queue = [...responses]; + jest.spyOn(globalThis, 'fetch').mockImplementation(async (_input, init) => { + const body = JSON.parse(String(init?.body)) as { command: string }; + seen.push(body.command); + return new Response(JSON.stringify(queue.shift()), { status: 200 }); + }); + return { commands: () => seen }; +} + +async function drain(download: Awaited>): Promise { + const reader = download.stream.getReader(); + const parts: Uint8Array[] = []; + while (true) { + const next = await reader.read(); + if (next.done) { + break; + } + parts.push(next.value); + } + return Buffer.concat(parts.map(part => Buffer.from(part))); +} + +describe('TFYSandboxProvider downloadFile streaming', () => { + it('pulls the file in byte windows and reassembles the exact bytes', async () => { + // CHUNK_BYTES + 1 bytes forces two windows (second carries the final byte). + const payload = Buffer.alloc(CHUNK_BYTES + 1); + for (let index = 0; index < payload.length; index += 1) { + payload[index] = index % 253; + } + const server = mockSandboxServer([ + execOk(0, '/sandbox\n/usr/local/bin:/usr/bin:/bin'), + execOk(0, JSON.stringify({ size: payload.length, type: 'regular file' })), + execOk(0, payload.subarray(0, CHUNK_BYTES).toString('base64')), + execOk(0, payload.subarray(CHUNK_BYTES).toString('base64')), + ]); + const provider = makeProvider(16 * 1024 * 1024); + + const download = await provider.downloadFile({ sandboxId: SANDBOX_ID, path: 'out/report.bin' }); + + expect(download.size).toBe(payload.length); + expect(Buffer.compare(await drain(download), payload)).toBe(0); + + const commands = server.commands(); + expect(commands[1]).toMatch(/^stat -L --printf=/); + expect(commands[2]).toContain('tail -c +1'); + expect(commands[2]).toContain(`head -c ${String(CHUNK_BYTES)}`); + expect(commands[3]).toContain(`tail -c +${String(CHUNK_BYTES + 1)}`); + expect(commands).toHaveLength(4); + }); + + it('returns an empty stream for an empty file', async () => { + mockSandboxServer([ + execOk(0, '/sandbox\n/usr/local/bin:/usr/bin:/bin'), + execOk(0, JSON.stringify({ size: 0, type: 'regular file' })), + ]); + const provider = makeProvider(1024); + + const download = await provider.downloadFile({ sandboxId: SANDBOX_ID, path: 'empty.txt' }); + + expect(download.size).toBe(0); + expect((await drain(download)).length).toBe(0); + }); + + it('rejects missing files and directories before any chunk is pulled', async () => { + const missing = mockSandboxServer([execOk(0, '/sandbox\n/usr/local/bin:/usr/bin:/bin'), execOk(1, '')]); + const missingProvider = makeProvider(1024); + await expect(missingProvider.downloadFile({ sandboxId: SANDBOX_ID, path: 'gone.txt' })).rejects.toBeInstanceOf( + SandboxFileNotFoundError, + ); + expect(missing.commands()).toHaveLength(2); + + // Fresh provider so its layout probe runs and consumes this server's first response. + const directory = mockSandboxServer([ + execOk(0, '/sandbox\n/usr/local/bin:/usr/bin:/bin'), + execOk(0, JSON.stringify({ size: 10, type: 'directory' })), + ]); + const directoryProvider = makeProvider(1024); + await expect(directoryProvider.downloadFile({ sandboxId: SANDBOX_ID, path: 'a-dir' })).rejects.toMatchObject({ + name: 'SandboxPathIsDirectoryError', + }); + expect(directory.commands()).toHaveLength(2); + }); + + it('fails mid-stream when a window comes back short (file mutated after stat)', async () => { + mockSandboxServer([ + execOk(0, '/sandbox\n/usr/local/bin:/usr/bin:/bin'), + execOk(0, JSON.stringify({ size: 10, type: 'regular file' })), + execOk(0, Buffer.from('short').toString('base64')), + ]); + const provider = makeProvider(1024); + + const download = await provider.downloadFile({ sandboxId: SANDBOX_ID, path: 'shrank.txt' }); + await expect(drain(download)).rejects.toThrow(/changed during download/); + }); + + it('maps a first-window command failure to SandboxFileNotFoundError', async () => { + mockSandboxServer([ + execOk(0, '/sandbox\n/usr/local/bin:/usr/bin:/bin'), + execOk(0, JSON.stringify({ size: 100, type: 'regular file' })), + execOk(127, ''), + ]); + const provider = makeProvider(1024); + + const download = await provider.downloadFile({ sandboxId: SANDBOX_ID, path: 'vanished.txt' }); + await expect(drain(download)).rejects.toBeInstanceOf(SandboxFileNotFoundError); + }); +}); diff --git a/packages/trueforge/src/apis/turns.ts b/packages/trueforge/src/apis/turns.ts index 26ec91bac..4cfecb079 100644 --- a/packages/trueforge/src/apis/turns.ts +++ b/packages/trueforge/src/apis/turns.ts @@ -68,17 +68,6 @@ export function toWireTurn(record: TurnRecordWithoutSnapshot): Turn { }; } -/** - * Copies into a standalone ArrayBuffer for the response body: a pooled Buffer's - * backing store is shared and typed as ArrayBufferLike, which is not a body. - * Bounded by SANDBOX_FILE_MAX_BYTES_FOR_DOWNLOAD. - */ -function toArrayBuffer(content: Buffer): ArrayBuffer { - const buffer = new ArrayBuffer(content.byteLength); - new Uint8Array(buffer).set(content); - return buffer; -} - /** * Builds the Content-Disposition telling the browser to save the file under its sandbox name. * @@ -439,11 +428,18 @@ export function createTurnsRouter(deps: TurnsRouterDeps) { return c.json({ error: { message: 'No sandbox provider configured' } }, 412); } - // TODO: stream the body instead of buffering the whole file in memory. - const content = await provider.downloadFile({ sandboxId: rawSandboxId(sandboxId), path }); - return c.body(toArrayBuffer(content), 200, { + // Stat + stream-open errors reject above with a JSON status; bytes flow incrementally so + // peak memory tracks chunk size, and `req.raw.signal` stops provider reads on disconnect. + // Once the body has started, mid-stream failures can only surface as a truncated response, + // detectable by the client against Content-Length. + const download = await provider.downloadFile({ + sandboxId: rawSandboxId(sandboxId), + path, + signal: c.req.raw.signal, + }); + return c.body(download.stream, 200, { 'Content-Type': 'application/octet-stream', - 'Content-Length': String(content.byteLength), + 'Content-Length': String(download.size), 'Content-Disposition': toContentDisposition(path), 'Cache-Control': 'private, no-store', }); diff --git a/packages/trueforge/src/sandbox/local/core/hostRun.ts b/packages/trueforge/src/sandbox/local/core/hostRun.ts index 2453713fc..8626673c1 100644 --- a/packages/trueforge/src/sandbox/local/core/hostRun.ts +++ b/packages/trueforge/src/sandbox/local/core/hostRun.ts @@ -445,51 +445,49 @@ export function killExecTree(child: ChildProcess | undefined): void { } } -/** - * Run one SRT-wrapped command. Code Mode UDS (if any) is supplied via `env.TFY_MCP_SOCK` - * from {@link CodeModeUdsTransport.start}. - */ -export async function runSupervisorSession(params: { +function ignoreStreamError( + stream: + | { + on: (event: 'error', cb: (err: Error) => void) => void; + } + | null + | undefined, +): void { + stream?.on('error', () => undefined); +} + +/** Everything needed to spawn one SRT-wrapped command; shared by buffered and streaming runs. */ +interface SandboxSpawnParams { sandboxRootPath: string; command: string; /** Absolute shell path used to wrap the command string (from isSupported). */ shell: string; /** Platform policy for allowRead / PATH (from isSupported). */ platform: LocalSandboxPlatform; - cwd?: string; - env?: Record; - /** Optional stdin bytes for the sandboxed command (e.g. upload payload). */ - stdin?: Buffer; - /** Host-visible pid of the sandboxed process-group leader (after spawn). */ - onChildSpawn?: (pid: number) => void; - /** Hard wall-clock limit for the sandboxed command; caller must choose deliberately. */ - timeoutMs: number; -}): Promise { - const { - sandboxRootPath, - command: rawCommand, - shell, - platform, - cwd = sandboxRootPath, - env, - stdin, - onChildSpawn, - timeoutMs, - } = params; - const command = wrapSandboxCommand({ platform, command: rawCommand }); + cwd?: string | undefined; + env?: Record | undefined; + /** Process-group leader (Unix) so {@link killExecTree} can tear down the whole tree. */ + detached: boolean; + /** Pipe stdin for the sandboxed command (e.g. upload payload); 'ignore' otherwise. */ + pipeStdin: boolean; + onChildSpawn?: ((pid: number) => void) | undefined; +} + +async function spawnSandboxedCommand(params: SandboxSpawnParams): Promise { + const command = wrapSandboxCommand({ platform: params.platform, command: params.command }); const wrap = await SandboxManager.wrapWithSandboxArgv( command, - shell, + params.shell, { - filesystem: filesystemPolicy({ sandboxRootPath, platform }), + filesystem: filesystemPolicy({ sandboxRootPath: params.sandboxRootPath, platform: params.platform }), network: { allowedDomains: allowedNetworkDomains(), deniedDomains: [], }, }, undefined, - sandboxRootPath, + params.sandboxRootPath, { commandId: randomUUID(), commandText: command }, ); @@ -501,20 +499,70 @@ export async function runSupervisorSession(params: { // Curated env only — do not spread wrap.env (it can carry ambient host secrets). // Code Mode sock path (TFY_MCP_SOCK) is expected in `env` when the caller starts a transport. const childEnv: NodeJS.ProcessEnv = { - ...commandEnv({ sandboxRootPath, platform, ...(env === undefined ? {} : { extra: env }) }), + ...commandEnv({ + sandboxRootPath: params.sandboxRootPath, + platform: params.platform, + ...(params.env === undefined ? {} : { extra: params.env }), + }), }; const child = spawn(argv0, argvRest, { - cwd, + cwd: params.cwd ?? params.sandboxRootPath, env: childEnv, shell: false, // Detached process groups break stdin forwarding for upload (`cat` via pipe) under Jest. - detached: stdin === undefined && process.platform !== 'win32', - stdio: [stdin === undefined ? 'ignore' : 'pipe', 'pipe', 'pipe'], + detached: params.detached && process.platform !== 'win32', + stdio: [params.pipeStdin ? 'pipe' : 'ignore', 'pipe', 'pipe'], }); if (child.pid !== undefined) { - onChildSpawn?.(child.pid); + params.onChildSpawn?.(child.pid); } + return child; +} + +/** + * Run one SRT-wrapped command. Code Mode UDS (if any) is supplied via `env.TFY_MCP_SOCK` + * from {@link CodeModeUdsTransport.start}. + */ +export async function runSupervisorSession(params: { + sandboxRootPath: string; + command: string; + /** Absolute shell path used to wrap the command string (from isSupported). */ + shell: string; + /** Platform policy for allowRead / PATH (from isSupported). */ + platform: LocalSandboxPlatform; + cwd?: string; + env?: Record; + /** Optional stdin bytes for the sandboxed command (e.g. upload payload). */ + stdin?: Buffer; + /** Host-visible pid of the sandboxed process-group leader (after spawn). */ + onChildSpawn?: (pid: number) => void; + /** Hard wall-clock limit for the sandboxed command; caller must choose deliberately. */ + timeoutMs: number; +}): Promise { + const { + sandboxRootPath, + command: rawCommand, + shell, + platform, + cwd = sandboxRootPath, + env, + stdin, + onChildSpawn, + timeoutMs, + } = params; + + const child = await spawnSandboxedCommand({ + sandboxRootPath, + command: rawCommand, + shell, + platform, + cwd, + env, + detached: stdin === undefined, + pipeStdin: stdin !== undefined, + onChildSpawn, + }); if (stdin !== undefined) { const stdinStream = child.stdin; if (stdinStream === null) { @@ -541,17 +589,6 @@ export async function runSupervisorSession(params: { let timedOut = false; let closed = false; - const ignoreStreamError = ( - stream: - | { - on: (event: 'error', cb: (err: Error) => void) => void; - } - | null - | undefined, - ): void => { - stream?.on('error', () => undefined); - }; - const appendOutput = (stream: 'stdout' | 'stderr', chunk: Buffer): void => { bufferedOutput += chunk.length; if (bufferedOutput > MAX_OUTPUT_BYTES) { @@ -610,3 +647,138 @@ export async function runSupervisorSession(params: { }); }); } + +/** + * Spawns one SRT-wrapped command whose raw stdout is consumed incrementally. + * Unlike {@link runSupervisorSession}, bytes are handed over as they arrive (binary-safe) + * instead of being buffered and utf8-decoded — sized for file downloads, where peak memory + * must track chunk size rather than file size. + * + * Wrap/spawn failures throw here so callers can still respond with a proper status before + * any byte moves downstream; consume the returned stream promptly. + */ +export async function startSupervisorStdoutStream(params: { + sandboxRootPath: string; + command: string; + shell: string; + platform: LocalSandboxPlatform; + cwd?: string; + env?: Record; + /** Hard wall-clock limit; expiry kills the process group and fails the stream. */ + timeoutMs: number; +}): Promise { + const child = await spawnSandboxedCommand({ + sandboxRootPath: params.sandboxRootPath, + command: params.command, + shell: params.shell, + platform: params.platform, + cwd: params.cwd, + detached: true, + pipeStdin: false, + }); + + let stderrText = ''; + let protocolError: string | undefined; + let timedOut = false; + let done = false; + const pending: Buffer[] = []; + let wake: (() => void) | undefined; + + const timer = setTimeout(() => { + timedOut = true; + killExecTree(child); + }, params.timeoutMs); + + const finish = (): void => { + done = true; + const w = wake; + wake = undefined; + w?.(); + }; + + ignoreStreamError(child.stderr); + child.stderr?.on('data', (chunk: Buffer) => { + if (stderrText.length >= MAX_OUTPUT_BYTES) { + protocolError = `stderr exceeded ${String(MAX_OUTPUT_BYTES)} bytes`; + killExecTree(child); + return; + } + stderrText += chunk.toString('utf8'); + }); + child.on('error', () => undefined); + child.stdout?.on('data', (chunk: Buffer) => { + pending.push(chunk); + const w = wake; + wake = undefined; + w?.(); + }); + // Node emits close after error for spawned processes, so close alone drives termination. + const exitCodePromise = new Promise(resolve => { + child.on('close', code => { + finish(); + resolve(typeof code === 'number' ? code : timedOut ? 1 : 0); + }); + }); + + async function* chunks(): AsyncGenerator { + try { + // Drain queued bytes first; block only when the queue is empty and the process runs on. + while (pending.length > 0 || !done) { + const chunk = pending.shift(); + if (chunk !== undefined) { + yield chunk; + continue; + } + await new Promise(resolveWake => { + wake = resolveWake; + }); + } + const exitCode = await exitCodePromise; + if (protocolError !== undefined || exitCode !== 0) { + const detail = [protocolError, stderrText.trim()] + .filter(part => part !== undefined && part.length > 0) + .join(' '); + throw new SupervisorStreamError( + detail.length > 0 ? detail : `sandboxed command exited ${String(exitCode)}`, + exitCode, + ); + } + } finally { + done = true; + clearTimeout(timer); + // Early consumer return lands here too; a completed process is an ESRCH no-op. + killExecTree(child); + SandboxManager.cleanupAfterCommand(); + } + } + + return { + abort: () => { + killExecTree(child); + }, + chunks, + }; +} + +/** Failure of a streamed sandboxed command after a successful spawn (nonzero/timeout/protocol). */ +export class SupervisorStreamError extends Error { + readonly exitCode: number; + + constructor(message: string, exitCode: number) { + super(message); + this.name = 'SupervisorStreamError'; + this.exitCode = exitCode; + } +} + +/** A spawned sandboxed command exposing its raw stdout for incremental consumption. */ +export interface SupervisedStdoutStream { + /** + * Yields stdout Buffers as they arrive; single consumption. Fails with + * {@link SupervisorStreamError} when the command ends non-zero / timed out / protocol-error. + * Returning early (or calling {@link SupervisedStdoutStream.abort}) kills the process group. + */ + chunks(): AsyncGenerator; + /** Force-stops the underlying process group (e.g. when the HTTP client disconnects). */ + abort: () => void; +} diff --git a/packages/trueforge/src/sandbox/local/provider/LocalSandboxProvider.ts b/packages/trueforge/src/sandbox/local/provider/LocalSandboxProvider.ts index a82c6e635..ff824c81d 100644 --- a/packages/trueforge/src/sandbox/local/provider/LocalSandboxProvider.ts +++ b/packages/trueforge/src/sandbox/local/provider/LocalSandboxProvider.ts @@ -6,10 +6,12 @@ import type { ExecResult, SandboxBuild, SandboxExecParams, + SandboxFileDownload, SandboxProvider, } from '@truefoundry/trueforge-core/core'; import { absolutizeRelativeExecEnv, + boundedFileStream, SandboxFileNotFoundError, SandboxFileTooLargeError, SandboxNotAvailableError, @@ -42,6 +44,8 @@ import { runSupervisorSession, SANDBOX_VENV_DIR, srtHostBinaryNames, + startSupervisorStdoutStream, + SupervisorStreamError, type LocalSandboxPlatform, type SessionResult, } from '../core/hostRun.js'; @@ -446,15 +450,6 @@ export class LocalSandboxProvider implements SandboxProvider { return this.pythonC(code, relPath); } - private base64EncodeCommand(relPath: string): string { - const code = [ - 'import base64, sys', - 'p = sys.argv[1]', - 'sys.stdout.write(base64.b64encode(open(p, "rb").read()).decode("ascii"))', - ].join('\n'); - return this.pythonC(code, relPath); - } - buildImage(): Promise { return Promise.resolve(LocalSandboxProvider.readyBuild); } @@ -677,7 +672,11 @@ export class LocalSandboxProvider implements SandboxProvider { return 'git_downloader.py'; } - async downloadFile(params: { sandboxId: string; path: string }): Promise { + async downloadFile(params: { + sandboxId: string; + path: string; + signal?: AbortSignal | undefined; + }): Promise { this.ensureSandboxRoot(params.sandboxId); await this.ensureSrt(); await this.ensureVenv(params.sandboxId); @@ -691,18 +690,50 @@ export class LocalSandboxProvider implements SandboxProvider { if (info.size > this.fileMaxBytesForDownload) { throw new SandboxFileTooLargeError(params.path, info.size, this.fileMaxBytesForDownload); } - const result = await this.runSandboxCommand({ + + // Raw bytes stream out under the same SRT wrap as exec. `/bin/cat` is absolute like the + // upload command: the jail PATH prefers Homebrew, where a stale Cellar cat can be ENOENT. + const supervised = await startSupervisorStdoutStream({ sandboxRootPath, - command: this.base64EncodeCommand(relPath), + command: `/bin/cat ${shellEscape(relPath)}`, + shell: this.support.shell, + platform: this.support.platform, + timeoutMs: this.defaultExecTimeoutSeconds * 1000, }); - if (result.exitCode !== 0) { - throw new SandboxFileNotFoundError(params.path); - } - const buf = Buffer.from(result.stdoutText.trim(), 'base64'); - if (buf.length > this.fileMaxBytesForDownload) { - throw new SandboxFileTooLargeError(params.path, buf.length, this.fileMaxBytesForDownload); + if (params.signal !== undefined) { + params.signal.addEventListener( + 'abort', + () => { + supervised.abort(); + }, + { once: true }, + ); } - return buf; + + let emitted = 0; + const chunks = async function* (): AsyncGenerator { + try { + for await (const chunk of supervised.chunks()) { + emitted += chunk.byteLength; + yield chunk; + } + } catch (error) { + // Failing before any byte mirrors the buffered behavior: the stat raced a deletion. + if (error instanceof SupervisorStreamError && emitted === 0) { + throw new SandboxFileNotFoundError(params.path); + } + throw error; + } + }; + + return { + size: info.size, + stream: boundedFileStream({ + path: params.path, + maxBytes: this.fileMaxBytesForDownload, + chunks, + }), + }; } /** Payload on stdin so large uploads stay off argv. */ diff --git a/packages/trueforge/tests/sandbox/local/smoke.test.ts b/packages/trueforge/tests/sandbox/local/smoke.test.ts index 9c01b5907..800fd894a 100644 --- a/packages/trueforge/tests/sandbox/local/smoke.test.ts +++ b/packages/trueforge/tests/sandbox/local/smoke.test.ts @@ -1686,10 +1686,16 @@ async function main(): Promise { remotePath: 'uploads/hello.txt', content: Buffer.from('upload-ok\n'), }); - const downloaded = await provider.downloadFile({ + const download = await provider.downloadFile({ sandboxId, path: 'uploads/hello.txt', }); + const downloadedChunks: Uint8Array[] = []; + for await (const chunk of download.stream) { + downloadedChunks.push(chunk); + } + const downloaded = Buffer.concat(downloadedChunks.map(chunk => Buffer.from(chunk))); + assert.equal(download.size, 'upload-ok\n'.length); assert.equal(downloaded.toString('utf8'), 'upload-ok\n'); const catUpload = await provider.exec({ sandboxId, @@ -1867,7 +1873,16 @@ async function main(): Promise { let downloadLeaked = false; try { const leaked = await provider.downloadFile({ sandboxId, path: 'api-escape-dl' }); - downloadLeaked = leaked.toString('utf8').includes('host-secret-should-not-leak'); + const reader = leaked.stream.getReader(); + const parts: Uint8Array[] = []; + while (true) { + const next = await reader.read(); + if (next.done) break; + parts.push(next.value); + } + downloadLeaked = Buffer.concat(parts.map(part => Buffer.from(part))) + .toString('utf8') + .includes('host-secret-should-not-leak'); } catch { // deny / throw is fine; host check below still runs } diff --git a/packages/trueforge/tests/unit/apis/sandboxFileDownload.test.ts b/packages/trueforge/tests/unit/apis/sandboxFileDownload.test.ts index fd80dee47..433b2a21d 100644 --- a/packages/trueforge/tests/unit/apis/sandboxFileDownload.test.ts +++ b/packages/trueforge/tests/unit/apis/sandboxFileDownload.test.ts @@ -1,7 +1,10 @@ import { OpenAPIHono } from '@hono/zod-openapi'; import type { AgentSpec } from '@truefoundry/trueforge-core/agent-session'; import { AgentSpecSchema, Sessions } from '@truefoundry/trueforge-core/agent-session'; +import type { SandboxProvider } from '@truefoundry/trueforge-core/core'; +import { SandboxFileNotFoundError } from '@truefoundry/trueforge-core/core'; import { createLogger } from 'winston'; +import { makeCreateTurnInput } from '../../../../trueforge-core/tests/agent-session/testHelpers'; import { TENANT_ID } from '../../../src/apis/sessions'; import { createTurnsRouter, toContentDisposition } from '../../../src/apis/turns'; import { LOCAL_USER_CONTEXT } from '../../../src/auth/identity'; @@ -16,6 +19,14 @@ import { SqliteSkillStore } from '../../../src/db/sqlite/skill-store/SqliteSkill import { SqliteOAuthTokenStore } from '../../../src/db/sqlite/token-store/SqliteOAuthTokenStore'; import { ActiveTurnRegistry } from '../../../src/runtime/activeTurns'; import { EventSubscriptionRegistry } from '../../../src/runtime/event-subscription'; +import { resolveSandboxProvider } from '../../../src/runtime/sessionResources'; + +jest.mock('../../../src/runtime/sessionResources', () => ({ + ...jest.requireActual('../../../src/runtime/sessionResources'), + resolveSandboxProvider: jest.fn(), +})); + +const mockedResolveSandboxProvider = jest.mocked(resolveSandboxProvider); /** Parsed rather than built literally, so config defaults match what the create route stores. */ function agentSpec(): AgentSpec { @@ -50,7 +61,7 @@ async function buildApp() { }), ); - return { app, sessions }; + return { app, sessions, sessionStore }; } function downloadUrl({ @@ -121,6 +132,98 @@ describe('GET /{session_id}/turns/{turn_id}/download-sandbox-file', () => { }); }); +describe('download-sandbox-file streaming', () => { + const RAW_SANDBOX_ID = 'raw-sbx-id'; + + /** Session with one turn whose snapshot carries a sandbox id, so the route reaches the provider. */ + async function buildSessionWithSandboxTurn(): Promise<{ app: OpenAPIHono; sessionId: string }> { + const { app, sessions, sessionStore } = await buildApp(); + const session = await sessions.create({ + tenant_id: TENANT_ID, + session_id: 'with-sbx-turn', + created_by: LOCAL_USER_CONTEXT.userRef, + agent: { type: 'inline', spec: agentSpec() }, + }); + await sessionStore.createTurn(makeCreateTurnInput({ sessionId: session.session_id, turnId: 'turn-1' })); + await sessionStore.patchSandboxInfo({ + session_id: session.session_id, + turn_id: 'turn-1', + sandbox_info: { sandbox_id: RAW_SANDBOX_ID }, + }); + return { app, sessionId: session.session_id }; + } + + afterEach(() => { + jest.resetAllMocks(); + }); + + it('streams bytes to the client progressively with exact headers', async () => { + const encoder = new TextEncoder(); + let releaseSecondChunk: (() => void) | undefined; + const secondChunkGated = new Promise(resolve => { + releaseSecondChunk = resolve; + }); + const downloadFile = jest.fn(async () => ({ + size: 11, + stream: new ReadableStream({ + async start(controller) { + controller.enqueue(encoder.encode('part-one-')); + await secondChunkGated; + controller.enqueue(encoder.encode('tail')); + controller.close(); + }, + }), + })); + mockedResolveSandboxProvider.mockResolvedValue({ type: 'test', downloadFile } as unknown as SandboxProvider); + + const { app, sessionId } = await buildSessionWithSandboxTurn(); + // The handler returns as soon as headers exist; the body keeps streaming after this resolves. + const response = await app.request(downloadUrl({ sessionId, path: '/workspace/report.pdf' })); + + expect(response.status).toBe(200); + expect(response.headers.get('Content-Type')).toBe('application/octet-stream'); + expect(response.headers.get('Content-Length')).toBe('11'); + expect(response.headers.get('Content-Disposition')).toBe(toContentDisposition('/workspace/report.pdf')); + expect(response.headers.get('Cache-Control')).toBe('private, no-store'); + expect(downloadFile).toHaveBeenCalledWith({ + sandboxId: RAW_SANDBOX_ID, + path: '/workspace/report.pdf', + signal: expect.anything(), + }); + + const body = response.body; + if (!body) { + throw new Error('response has no body'); + } + const reader = body.getReader(); + // The producer is parked on secondChunkGated: receiving the first chunk here proves + // bytes reach the client before the whole file has been produced. + const first = await reader.read(); + expect(Buffer.from(first.value).toString()).toBe('part-one-'); + + releaseSecondChunk?.(); + const second = await reader.read(); + const done = await reader.read(); + expect(done.done).toBe(true); + expect(Buffer.from(first.value).toString() + Buffer.from(second.value).toString()).toBe('part-one-tail'); + }); + + it('maps provider SandboxError rejections to their status before any body starts', async () => { + mockedResolveSandboxProvider.mockResolvedValue({ + type: 'test', + downloadFile: jest.fn(async () => { + throw new SandboxFileNotFoundError('/workspace/gone.txt'); + }), + } as unknown as SandboxProvider); + + const { app, sessionId } = await buildSessionWithSandboxTurn(); + const response = await app.request(downloadUrl({ sessionId, path: '/workspace/gone.txt' })); + + expect(response.status).toBe(404); + expect(await response.json()).toEqual({ error: { message: 'File not found: /workspace/gone.txt' } }); + }); +}); + describe('toContentDisposition', () => { it('names the file after the last path segment', () => { expect(toContentDisposition('/workspace/out/report.pdf')).toBe(`attachment; filename*=UTF-8''report.pdf`);