From 2ceeff33dea96a92d2a4b5d2cd5269f196e25a52 Mon Sep 17 00:00:00 2001 From: Jalil Date: Fri, 13 Nov 2026 17:38:11 -0800 Subject: [PATCH] perf(headless-computer): send files straight to and from the bucket A person's files reached their computer through the runner: one stat per file, then each file read whole from the bucket into the runner and written through the provider's file API, 30 seconds at most. Which files were already copied was only remembered in memory, so every new runner process asked about every file again. Outputs came back the same way, each read through the runner after every command. With the S3 store the computer now moves the bytes itself. The runner signs a link per file, for one object and one method, that stops working within minutes; the computer runs curl over all of them in one command. The links are written to a file that is removed when the command ends, never to a command line, a log or what the model sees. An output is sent to a holding object. The runner measures what arrived, applies the per-file and per-conversation limits, and only then copies it into place inside the bucket, so a kept file is never larger than its row says and a version that doesn't fit leaves the kept one untouched. What is already on a computer is read in one command: the folder's names and sizes, and a record kept in /workspace/.jobs/.copied.json. The record only counts while the file is still there, because a replacement computer carries .jobs over but not /workspace/files. A disk store has no links, and a file that doesn't arrive by link goes through the runner as before. After a transfer in which nothing arrives, links are skipped for ten minutes. Signed-off-by: Jalil --- ee/apps/headless-runner/README.md | 3 +- .../scripts/saved-file-links.test.ts | 180 +++++++++ ee/apps/headless-runner/src/blobs.ts | 37 +- ee/apps/headless-runner/src/saved-files.ts | 64 +++- .../src/computer-files.test.ts | 351 ++++++++++++++++++ ee/packages/headless-computer/src/computer.ts | 303 ++++++++++++--- ee/packages/headless-computer/src/types.ts | 15 + 7 files changed, 900 insertions(+), 53 deletions(-) create mode 100644 ee/apps/headless-runner/scripts/saved-file-links.test.ts create mode 100644 ee/packages/headless-computer/src/computer-files.test.ts diff --git a/ee/apps/headless-runner/README.md b/ee/apps/headless-runner/README.md index f4108f1c5b8..d9a51d0add4 100644 --- a/ee/apps/headless-runner/README.md +++ b/ee/apps/headless-runner/README.md @@ -262,7 +262,7 @@ The runner is a private service with one caller (Den), which authorizes each per - **Opt-in per session.** Configuring files or a computer makes them *available*; a session only gets them when it is created (or updated) with `files: true` / `computer: true`. Slack replies and Automations don't ask, so they keep their exact behavior: no file tools, no computer, no uploads, and the preview route never starts a computer for them. - **Files belong to one session.** Every read, download, preview and delete looks a file up by session id *and* file id, and bytes are stored under `sessions//`. Deleting a session deletes its bytes and its computer. - **Limits.** `HEADLESS_FILES_MAX_BYTES` (100 MB) per file, enforced while the upload streams in (with or without `Content-Length`), and on `save_file` and computer outputs. `HEADLESS_FILES_MAX_SESSION_BYTES` (5 GB) per session in total. Past either, the upload gets `413`, and the model is told why so it can say so. -- **Storage credentials.** Use a store or bucket for this runner alone, with credentials scoped to it. Neither the storage credentials nor the Freestyle key ever enter a computer. +- **Storage credentials.** Use a store or bucket for this runner alone, with credentials scoped to it. Neither the storage credentials nor the Freestyle key ever enter a computer. With an S3 store a computer does receive links: each reaches one object with one method, carries a signature rather than the credentials, and stops working within minutes. - **Computers.** One VM per session, created only when that session first runs a command; it holds no credentials and has outbound internet only. ## Computer (optional) @@ -275,6 +275,7 @@ The runner is a private service with one caller (Den), which authorizes each per | `look { paths }` | Shows up to 4 files from the VM like any tool file: images and PDFs as model input, Office files and text as text | - **Files:** uploads are copied into `/workspace/files`; files written to `/workspace/out` become saved files after each command. A new version of an out file updates the same saved file (same id, newer `updatedAt`). +- **How the bytes move:** with `HEADLESS_FILES=s3` the computer fetches uploads from the bucket and sends outputs to it itself, with `curl` over those links, so files never pass through the runner; the computer must be able to reach the bucket's endpoint. An output goes to a holding object first and is kept only after the runner has measured it against the limits. With a disk store, or when a transfer doesn't arrive, files go through the runner as before. Which uploads are already on a computer is recorded there, in `/workspace/.jobs/.copied.json`. - **Previews:** slide decks, documents and PDFs get page images (LibreOffice, then `pdftoppm`), served at `GET /v1/sessions/:id/saved-files/:fileId/preview` and `/preview/:page`. - **Only for sessions with `computer: true`.** Others never see the tools, and nothing starts a VM for them. - **Lifecycle:** one VM per session, found by a slug derived from the session id. Paused `HEADLESS_COMPUTER_PAUSE_SECONDS` (300) after the last turn unless a `background` job runs; deleted after `HEADLESS_COMPUTER_KEEP_DAYS` (14) unused, or with the session. diff --git a/ee/apps/headless-runner/scripts/saved-file-links.test.ts b/ee/apps/headless-runner/scripts/saved-file-links.test.ts new file mode 100644 index 00000000000..b0f8689aa1b --- /dev/null +++ b/ee/apps/headless-runner/scripts/saved-file-links.test.ts @@ -0,0 +1,180 @@ +/** Links for moving kept files without holding them (src/blobs.ts, src/saved-files.ts): what a link is, and what is kept of what arrives. */ +import assert from "node:assert/strict" +import { mkdtempSync, rmSync } from "node:fs" +import { tmpdir } from "node:os" +import { join } from "node:path" +import { test } from "node:test" +import { diskBlobStore, s3BlobStore } from "../src/blobs.js" +import { nodeSqlite } from "../src/node-sqlite.js" +import { SavedFileLimitError, SavedFiles } from "../src/saved-files.js" +import { Store } from "../src/store.js" + +const SECRET = "secret-key-never-leaves-the-runner" +const bytesOf = (text: string) => new Uint8Array(Buffer.from(text)) + +/** An S3-compatible bucket in memory, reached through `fetch`: enough of GET, PUT, HEAD, DELETE and CopyObject. */ +function bucket() { + const objects = new Map; type: string }>() + const keyOf = (path: string) => decodeURIComponent(path).replace(/^\/files\//, "") + const fetcher: typeof fetch = async (input, init) => { + const request = input instanceof Request ? input : new Request(input, init) + const key = keyOf(new URL(request.url).pathname) + const found = objects.get(key) + if (request.method === "PUT") { + const source = request.headers.get("x-amz-copy-source") + const from = source ? objects.get(keyOf(source)) : null + if (source && !from) return new Response("NoSuchKey", { status: 404 }) + objects.set(key, { bytes: from ? from.bytes : new Uint8Array(await request.arrayBuffer()), type: request.headers.get("content-type") ?? "" }) + return new Response(from ? "" : null, { status: 200 }) + } + if (request.method === "DELETE") { + objects.delete(key) + return new Response(null, { status: 204 }) + } + if (!found) return new Response(null, { status: 404 }) + if (request.method === "HEAD") return new Response(null, { status: 200, headers: { "content-length": String(found.bytes.byteLength) } }) + return new Response(found.bytes, { status: 200 }) + } + const blobs = s3BlobStore({ endpoint: "https://store.example", region: "auto", bucket: "files", accessKeyId: "key-id", secretAccessKey: SECRET, forcePathStyle: true, fetch: fetcher }) + /** What the computer does with a link: one plain request, no credentials of its own. */ + const send = (url: string, body: string) => fetcher(url, { method: "PUT", body }) + return { objects, blobs, send, fetcher } +} + +function conversation(limits = { maxFileBytes: 1_000, maxSessionBytes: 5_000 }) { + const store = new Store(nodeSqlite(":memory:")) + const session = store.createSession({ owner: "test-owner", files: true }).id + const remote = bucket() + return { store, session, files: new SavedFiles(store, remote.blobs, limits), ...remote } +} + +test("a link reaches one object with one method for a limited time, and carries no secret", async () => { + const { blobs } = bucket() + assert.ok(blobs.link) + const get = new URL(await blobs.link("GET", "sessions/s1/fl_a", 240)) + assert.equal(`${get.origin}${get.pathname}`, "https://store.example/files/sessions/s1/fl_a") + assert.equal(get.searchParams.get("X-Amz-Expires"), "240") + assert.equal(get.searchParams.get("X-Amz-SignedHeaders"), "host") + assert.match(get.searchParams.get("X-Amz-Signature") ?? "", /^[0-9a-f]{64}$/) + assert.ok(!get.href.includes(SECRET)) + // Another method, object or lifetime needs another signature: one link can't be turned into another. + const signatures = new Set(await Promise.all([ + blobs.link("PUT", "sessions/s1/fl_a", 240), + blobs.link("GET", "sessions/s1/fl_b", 240), + blobs.link("GET", "sessions/s1/fl_a", 241), + ].map(async (link) => new URL(await link).searchParams.get("X-Amz-Signature")))) + assert.equal(signatures.size, 3) + assert.ok(!signatures.has(get.searchParams.get("X-Amz-Signature"))) +}) + +test("a file sent to its link is measured, then kept under its own key with its type", async () => { + const { store, session, files, objects, send, fetcher } = conversation() + try { + const place = await files.receive(session, { name: "chart.png", size: 5, source: "agent", expiresSeconds: 120 }) + assert.ok(place) + // The link writes a holding object, never a kept file. + const holding = decodeURIComponent(new URL(place.url).pathname).replace("/files/", "") + assert.match(holding, new RegExp(`^sessions/${session}/incoming-[0-9a-f]{32}$`)) + assert.equal((await send(place.url, "12345")).status, 200) + assert.deepEqual(files.list(session), []) + + const file = await place.finish() + assert.ok(file) + assert.deepEqual({ name: file.name, size: file.size, mediaType: file.mediaType, source: file.source }, { name: "chart.png", size: 5, mediaType: "image/png", source: "agent" }) + assert.deepEqual([...objects.keys()], [`sessions/${session}/${file.id}`]) + assert.equal(objects.get(`sessions/${session}/${file.id}`)?.type, "image/png") + assert.equal(Buffer.from((await files.read(session, file.id))?.bytes ?? []).toString(), "12345") + + // The person's own file is fetched through a link to that one object. + const link = await files.downloadUrl(session, file.id, 120) + assert.equal(new URL(link ?? "").pathname, `/files/sessions/${session}/${file.id}`) + assert.equal(await (await fetcher(link ?? "")).text(), "12345") + assert.equal(await files.downloadUrl(session, "fl_missing", 120), null) + } finally { + store.close() + } +}) + +test("what arrives is held to the limits, whatever size was announced", async () => { + const { store, session, files, objects, send } = conversation({ maxFileBytes: 10, maxSessionBytes: 12 }) + try { + // Too large as announced: no link at all. + await assert.rejects(files.receive(session, { name: "big.bin", size: 20, source: "agent", expiresSeconds: 120 }), (error) => error instanceof SavedFileLimitError && error.code === "file_too_large") + + // Announced small, sent large: measured after it arrived, and not kept. + const sly = await files.receive(session, { name: "sly.bin", size: 4, source: "agent", expiresSeconds: 120 }) + assert.ok(sly) + await send(sly.url, "x".repeat(11)) + await assert.rejects(sly.finish(), (error) => error instanceof SavedFileLimitError && error.code === "file_too_large") + assert.deepEqual(files.list(session), []) + assert.equal(objects.size, 0) + + // The conversation's total counts what arrived too. + const kept = await files.add(session, { name: "notes.txt", bytes: bytesOf("12345678"), source: "agent" }) + const extra = await files.receive(session, { name: "extra.txt", size: 3, source: "agent", expiresSeconds: 120 }) + assert.ok(extra) + await send(extra.url, "12345") + await assert.rejects(extra.finish(), (error) => error instanceof SavedFileLimitError && error.code === "files_full") + + // A new version that doesn't fit leaves the kept version as it was. + const next = await files.receive(session, { name: "notes.txt", size: 8, source: "agent", replaces: kept.id, expiresSeconds: 120 }) + assert.ok(next) + await send(next.url, "x".repeat(11)) + await assert.rejects(next.finish(), SavedFileLimitError) + assert.deepEqual(files.list(session).map((file) => [file.id, file.size]), [[kept.id, 8]]) + assert.equal(Buffer.from((await files.read(session, kept.id))?.bytes ?? []).toString(), "12345678") + assert.deepEqual([...objects.keys()], [`sessions/${session}/${kept.id}`]) + } finally { + store.close() + } +}) + +test("a new version sent to a link replaces the file in place; when nothing arrives, nothing changes", async () => { + const { store, session, files, objects, send } = conversation() + try { + const first = await files.add(session, { name: "deck.pptx", bytes: bytesOf("version one"), source: "agent" }) + await files.putPreview(session, first.id, [bytesOf("page")], { width: 10, height: 10 }) + + const nothing = await files.receive(session, { name: "deck.pptx", size: 11, source: "agent", replaces: first.id, expiresSeconds: 120 }) + assert.equal(await nothing?.finish(), null) + assert.equal(Buffer.from((await files.read(session, first.id))?.bytes ?? []).toString(), "version one") + assert.ok(await files.readPreview(session, first.id)) + + const place = await files.receive(session, { name: "deck.pptx", size: 11, source: "agent", replaces: first.id, expiresSeconds: 120 }) + assert.ok(place) + await send(place.url, "the second version") + const second = await place.finish() + // Same file, new bytes and size, and its old preview is gone with the old bytes. + assert.deepEqual([second?.id, second?.size], [first.id, 18]) + assert.deepEqual(files.list(session).map((file) => [file.id, file.size]), [[first.id, 18]]) + assert.equal(Buffer.from((await files.read(session, first.id))?.bytes ?? []).toString(), "the second version") + assert.equal(await files.readPreview(session, first.id), null) + assert.deepEqual([...objects.keys()], [`sessions/${session}/${first.id}`]) + + // The file it was to replace is gone by the time it arrives: kept as a new file, like `replace` then `add`. + const late = await files.receive(session, { name: "deck.pptx", size: 5, source: "agent", replaces: first.id, expiresSeconds: 120 }) + assert.ok(late) + await files.delete(session, first.id) + await send(late.url, "third") + const third = await late.finish() + assert.ok(third && third.id !== first.id) + assert.deepEqual(files.list(session).map((file) => [file.id, file.size]), [[third.id, 5]]) + } finally { + store.close() + } +}) + +test("a disk store gives no links, so files keep going through the runner", async () => { + const directory = mkdtempSync(join(tmpdir(), "saved-file-links-")) + const store = new Store(nodeSqlite(":memory:")) + try { + const session = store.createSession({ owner: "test-owner", files: true }).id + const files = new SavedFiles(store, diskBlobStore(directory), { maxFileBytes: 1_000, maxSessionBytes: 5_000 }) + const file = await files.add(session, { name: "notes.txt", bytes: bytesOf("on disk"), source: "user" }) + assert.equal(await files.downloadUrl(session, file.id, 120), null) + assert.equal(await files.receive(session, { name: "out.txt", size: 3, source: "agent", expiresSeconds: 120 }), null) + } finally { + store.close() + rmSync(directory, { recursive: true, force: true }) + } +}) diff --git a/ee/apps/headless-runner/src/blobs.ts b/ee/apps/headless-runner/src/blobs.ts index de2ea4f7ac4..0130a8b1132 100644 --- a/ee/apps/headless-runner/src/blobs.ts +++ b/ee/apps/headless-runner/src/blobs.ts @@ -22,6 +22,16 @@ export type BlobStore = { * large bodies (storage/), which deletes only old objects nothing references. Optional for test stores. */ list?(prefix: string): Promise> + /** + * Moving bytes without holding them (S3 only; a disk store has none of these). `link` is an address for one object + * and one method that stops working after `expiresSeconds`: it carries a signature, never the store's credentials, + * so it can be handed to a conversation's computer. `size` and `copy` let the runner check and keep what arrived. + */ + link?(method: "GET" | "PUT", key: string, expiresSeconds: number): Promise + /** An object's size in bytes; null when it does not exist. */ + size?(key: string): Promise + /** Copies an object inside the store, giving the copy this content type. */ + copy?(from: string, to: string, contentType: string): Promise } /** Keys are runner-generated (`sessions//`); this only guards against a key escaping the folder. */ @@ -104,11 +114,11 @@ export function s3BlobStore(options: S3Options): BlobStore { if (options.forcePathStyle) return `${base.origin}/${encodeURIComponent(options.bucket)}/${path}` return `${base.protocol}//${options.bucket}.${base.host}/${path}` } - const send = async (method: string, key: string, init: { body?: Uint8Array; contentType?: string } = {}) => { + const send = async (method: string, key: string, init: { body?: Uint8Array; contentType?: string; headers?: Record } = {}) => { const request = await client.sign(urlOf(key), { method, ...(init.body ? { body: init.body } : {}), - headers: init.contentType ? { "content-type": init.contentType } : {}, + headers: { ...init.headers, ...(init.contentType ? { "content-type": init.contentType } : {}) }, }) return (options.fetch ?? fetch)(request, { signal: AbortSignal.timeout(S3_TIMEOUT_MS) }) } @@ -137,6 +147,29 @@ export function s3BlobStore(options: S3Options): BlobStore { if (!response.ok && response.status !== 404) throw await fail("delete", response) await response.body?.cancel() }, + // A presigned URL: the signature is in the query, bound to this object, this method and the expiry. + async link(method, key, expiresSeconds) { + const url = new URL(urlOf(key)) + url.searchParams.set("X-Amz-Expires", String(Math.max(1, Math.round(expiresSeconds)))) + return (await client.sign(url.toString(), { method, aws: { signQuery: true } })).url + }, + async size(key) { + const response = await send("HEAD", key) + await response.body?.cancel() + if (response.status === 404) return null + const length = response.headers.get("content-length") + if (!response.ok || length === null || !/^\d+$/.test(length)) throw new Error(`s3_head_${response.status}`) + return Number(length) + }, + // CopyObject; the copy's metadata is replaced so it gets this content type, whatever the source was sent with. + async copy(from, to, contentType) { + const source = `/${encodeURIComponent(options.bucket)}/${safeKey(from).split("/").map(encodeURIComponent).join("/")}` + const response = await send("PUT", to, { contentType, headers: { "x-amz-copy-source": source, "x-amz-metadata-directive": "REPLACE" } }) + if (!response.ok) throw await fail("copy", response) + // A copy can fail after the 200 was sent: the error is then in the body. + const body = await response.text() + if (//.test(body)) throw new Error(`s3_copy_failed: ${body.slice(0, 300)}`) + }, // ListObjectsV2, page by page. async list(prefix) { safeKey(prefix.replace(/\/$/, "")) diff --git a/ee/apps/headless-runner/src/saved-files.ts b/ee/apps/headless-runner/src/saved-files.ts index 96ace0a92f6..242241f14b2 100644 --- a/ee/apps/headless-runner/src/saved-files.ts +++ b/ee/apps/headless-runner/src/saved-files.ts @@ -104,8 +104,8 @@ export class SavedFiles { } } - async add(sessionId: string, input: { name: string; mediaType?: string | null; bytes: Uint8Array; source: SavedFile["source"] }) { - this.checkFits(sessionId, input.bytes.byteLength) + /** A new kept file's row and where its bytes go; nothing is written yet. */ + private newFile(sessionId: string, input: { name: string; mediaType?: string | null; size: number; source: SavedFile["source"] }) { const id = `fl_${randomUUID().replaceAll("-", "")}` const name = cleanFileName(input.name) const file: SavedFile = { @@ -113,12 +113,17 @@ export class SavedFiles { sessionId, name, mediaType: mediaTypeFor(name, input.mediaType), - size: input.bytes.byteLength, + size: input.size, source: input.source, createdAt: this.now(), updatedAt: this.now(), } - const storageKey = `sessions/${sessionId}/${id}` + return { file, storageKey: `sessions/${sessionId}/${id}` } + } + + async add(sessionId: string, input: { name: string; mediaType?: string | null; bytes: Uint8Array; source: SavedFile["source"] }) { + this.checkFits(sessionId, input.bytes.byteLength) + const { file, storageKey } = this.newFile(sessionId, { ...input, size: input.bytes.byteLength }) // Bytes first: a row never points at an object that was not written. await this.blobs.put(storageKey, input.bytes, file.mediaType) this.store.addSavedFile({ ...file, storageKey }) @@ -152,6 +157,57 @@ export class SavedFiles { return bytes ? { file: entry.file, bytes } : null } + /** + * A link the conversation's computer fetches a kept file from, so the bytes never pass through the runner. It + * reaches that one file and stops working after `expiresSeconds`. Null when the store can't make links (a disk + * store) or the file is gone. + */ + async downloadUrl(sessionId: string, id: string, expiresSeconds: number): Promise { + const entry = this.blobs.link ? this.store.getSavedFile(sessionId, id) : null + return entry && this.blobs.link ? this.blobs.link("GET", entry.storageKey, expiresSeconds) : null + } + + /** + * Takes a file the conversation's computer sends straight to the store. The link writes one holding object and + * stops working after `expiresSeconds`; it never reaches a kept file. `finish` then measures what arrived, checks + * it against the limits, and only then copies it into place as a new kept file (or, with `replaces`, as that + * file's new version, like `replace`), so a kept file is never larger than its row says. `finish` gives null when + * nothing arrived. Throws like `add` when `size` doesn't fit; null when the store can't make links. + */ + async receive(sessionId: string, input: { name: string; size: number; source: SavedFile["source"]; replaces?: string; expiresSeconds: number }) { + const blobs = this.blobs + const { link, size: sizeOf, copy } = blobs + if (!link || !sizeOf || !copy) return null + const earlier = () => (input.replaces ? this.store.getSavedFile(sessionId, input.replaces) : null) + this.checkFits(sessionId, input.size, earlier()?.file.size ?? 0) + const holding = `sessions/${sessionId}/incoming-${randomUUID().replaceAll("-", "")}` + const url = await link.call(blobs, "PUT", holding, input.expiresSeconds) + const finish = async (): Promise => { + try { + const size = await sizeOf.call(blobs, holding) + if (size === null) return null + const entry = earlier() + this.checkFits(sessionId, size, entry?.file.size ?? 0) + if (entry) { + const mediaType = mediaTypeFor(entry.file.name, entry.file.mediaType) + await this.deletePreview(entry.storageKey) + await copy.call(blobs, holding, entry.storageKey, mediaType) + const updatedAt = this.now() + this.store.updateSavedFile(sessionId, entry.file.id, { size, mediaType, updatedAt }) + return { ...entry.file, size, mediaType, updatedAt } + } + const { file, storageKey } = this.newFile(sessionId, { name: input.name, size, source: input.source }) + // Bytes first, as in `add`: a row never points at an object that was not written. + await copy.call(blobs, holding, storageKey, file.mediaType) + this.store.addSavedFile({ ...file, storageKey }) + return file + } finally { + await blobs.delete(holding).catch(() => undefined) + } + } + return { url, finish } + } + async delete(sessionId: string, id: string) { const entry = this.store.getSavedFile(sessionId, id) if (!entry) return false diff --git a/ee/packages/headless-computer/src/computer-files.test.ts b/ee/packages/headless-computer/src/computer-files.test.ts new file mode 100644 index 00000000000..8a1f3d1fe55 --- /dev/null +++ b/ee/packages/headless-computer/src/computer-files.test.ts @@ -0,0 +1,351 @@ +/** How the person's files reach the computer and its outputs reach their Files (computer.ts): one command to see what is there, links for the bytes, the runner as the fallback. */ +import assert from "node:assert/strict" +import { execFile, execFileSync } from "node:child_process" +import { chmodSync, existsSync, mkdirSync, mkdtempSync, readdirSync, readFileSync, rmSync, statSync, writeFileSync } from "node:fs" +import { createServer } from "node:http" +import type { AddressInfo } from "node:net" +import { tmpdir } from "node:os" +import { dirname, join } from "node:path" +import { test } from "node:test" +import type { RunResult, RunSpec, SandboxProvider } from "@openwork/sandbox" +import { createFakeProvider } from "@openwork/sandbox/testing" +import { Computers } from "./computer.js" +import type { ComputerFile, ComputerFiles } from "./types.js" + +const SESSION = "conversation" +const encode = (text: string) => new Uint8Array(Buffer.from(text)) + +/** The conversation's kept files in memory. `carried` lists what went through the runner: `in:` and `out:`. */ +function keptFiles(links?: Pick) { + const rows: ComputerFile[] = [] + const bytes = new Map>() + const carried: string[] = [] + const keep = (name: string, content: Uint8Array, source: ComputerFile["source"]) => { + const file: ComputerFile = { id: `fl_${rows.length + 1}`, name, size: content.byteLength, source, updatedAt: rows.length + 1 } + rows.unshift(file) + bytes.set(file.id, content) + return file + } + const files: ComputerFiles = { + maxFileBytes: 1_000_000, + list: () => [...rows], + async read(_session, id) { + carried.push(`in:${id}`) + const file = rows.find((row) => row.id === id) + const content = bytes.get(id) + return file && content ? { file, bytes: content } : null + }, + async add(_session, input) { + carried.push(`out:${input.name}`) + return keep(input.name, input.bytes, "agent") + }, + async replace() { + return null + }, + async putPreview() { + return false + }, + ...links, + } + return { files, rows, bytes, carried, keep, sent: (name: string, text: string) => keep(name, encode(text), "user") } +} + +/** A computer whose commands answer from `answer` (by the command's text), recording commands, file writes and provider calls. */ +function scripted(files: ComputerFiles, answer: (command: string) => string | undefined, options: { desktop?: boolean } = {}) { + const commands: string[] = [] + const written = new Map() + const calls: string[] = [] + const fake = createFakeProvider({ + onOperation: (operation) => void calls.push(operation.name), + onRun: (spec) => { + commands.push(spec.command) + return { exitCode: 0, stdout: answer(spec.command) ?? "", stderr: "" } + }, + }) + const inner = fake.blocks?.files + assert.ok(inner) + const provider: SandboxProvider = { + ...fake, + blocks: { + ...fake.blocks, + files: { + ...inner, + async write(handle, path, content, opts) { + written.set(path, { text: Buffer.from(content).toString(), mode: opts.mode }) + return inner.write(handle, path, content, opts) + }, + }, + }, + } + const computers = new Computers({ provider, idlePauseMs: 60_000, keepDays: 1, desktop: options.desktop }, { files, readFile: async () => ({ text: "file" }) }) + return { computers, commands, written, calls, provider } +} + +const isSeeing = (command: string) => command.startsWith("find /workspace/files ") +const isListing = (command: string) => command.startsWith("find /workspace/out ") +const isTransfer = (command: string) => command.includes("curl ") +const SIGNATURE = "X-Amz-Signature=0123abcd" +const link = (id: string) => `https://store.example/files/${id}?${SIGNATURE}` + +test("a new process sees what is on the computer in one command and copies nothing again", async () => { + const kept = keptFiles({ downloadUrl: async (_session, id) => link(id) }) + kept.sent("report.pdf", "pdf") + kept.sent("data.csv", "rows") + const { computers, commands, calls, written } = scripted(kept.files, (command) => + isSeeing(command) ? `report.pdf\t3\ndata.csv\t4\nnotes of mine.txt\t9\n---computer-copied---\n{"fl_1":"report.pdf","fl_2":"data.csv"}\n` : undefined) + const result = await computers.bash(SESSION, "true", 10) + assert.equal(result.output, "[exit 0]") + // One command to see the folder and the record, then the command itself and the look at its outputs. + assert.equal(commands.length, 3) + assert.deepEqual(calls.filter((name) => name.startsWith("files.")), []) + assert.deepEqual(kept.carried, []) + assert.equal(written.size, 0) + // The same process doesn't look again. + await computers.bash(SESSION, "true", 10) + assert.equal(commands.filter(isSeeing).length, 1) + await computers.delete(SESSION) +}) + +test("a replacement computer has the record but not the files: they are fetched again by link, in one command", async () => { + for (const desktop of [false, true]) { + const kept = keptFiles({ downloadUrl: async (_session, id) => link(id) }) + kept.sent("report.pdf", "pdf") + kept.sent("it's data.csv", "rows") + // The record was carried over; /workspace/files was not, and holds something else called report.pdf. + const { computers, commands, written, calls } = scripted(kept.files, (command) => + isSeeing(command) ? `report.pdf\t999\n---computer-copied---\n{"fl_1":"report.pdf","fl_2":"it's data.csv"}\n` : isTransfer(command) ? "copied 0\ncopied 1\n" : undefined, { desktop }) + const result = await computers.look(SESSION, ["nothing.txt"]) + assert.match(result.output, /no such file/) + + const folder = desktop ? "/home/ubuntu" : "/workspace/.jobs" + const [[planPath, plan] = ["", { text: "", mode: undefined }]] = [...written] + assert.equal(written.size, 1) + assert.match(planPath, new RegExp(`^${folder}/\\.openwork-transfer-[0-9a-f]{32}$`)) + // On a desktop that folder is the provider account's own; the command closes the file to others first. + assert.equal(plan.mode, undefined) + const part = (index: number) => new RegExp(`output = "(/workspace/files/\\.openwork-part-[0-9a-f]{32}-${index})"`).exec(plan.text)?.[1] ?? "missing" + assert.ok(plan.text.includes(`url = "${link("fl_2")}"\noutput = "${part(0)}"\nfail\n`)) + assert.ok(plan.text.includes(`\nnext\nurl = "${link("fl_1")}"\noutput = "${part(1)}"\n`)) + + // The links are in that file only: never in a command. + assert.ok(commands.every((command) => !command.includes("store.example") && !command.includes("X-Amz"))) + const transfer = commands.find(isTransfer) ?? "" + assert.ok(transfer.includes(`--config ${planPath} `)) + const permissions = desktop ? `chmod 660 ${part(0)} && ` : "" + assert.ok(transfer.includes(`[ "$(stat -c %s ${part(0)} 2>/dev/null)" = 4 ] && ${permissions}mv -fT -- ${part(0)} '/workspace/files/it'"'"'s data.csv' && echo copied 0`)) + // The other report.pdf is not this one, so the person's file gets its own name; the record says so. + assert.ok(transfer.includes(`mv -fT -- ${part(1)} '/workspace/files/report (2).pdf' && echo copied 1`)) + assert.ok(transfer.includes(`printf %s '{"fl_1":"report (2).pdf","fl_2":"it'"'"'s data.csv"}' > /workspace/.jobs/.copied.json`)) + assert.ok(transfer.startsWith(`trap 'rm -rf ${planPath} ${planPath}.d ${part(0)} ${part(1)}' EXIT\nchmod 600 ${planPath}\n`)) + assert.equal(commands.length, 2) + assert.deepEqual(kept.carried, []) + assert.deepEqual(calls.filter((name) => name === "files.write" || name === "files.read"), ["files.write"]) + await computers.delete(SESSION) + } +}) + +test("when nothing arrives by link, files go through the runner, and links rest for a while", async (t) => { + const warned = t.mock.method(console, "warn", () => undefined) + const kept = keptFiles({ downloadUrl: async (_session, id) => link(id) }) + kept.sent("report.pdf", "pdf") + const { computers, commands, written, provider } = scripted(kept.files, (command) => (isSeeing(command) ? `---computer-copied---\n{}\n` : undefined)) + const first = await computers.bash(SESSION, "true", 10) + assert.equal(first.output, "[Copied the person's files to: /workspace/files/report.pdf]\n[exit 0]") + assert.deepEqual(kept.carried, ["in:fl_1"]) + assert.equal(written.get("/workspace/files/report.pdf")?.text, "pdf") + assert.equal(written.get("/workspace/.jobs/.copied.json")?.text, `{"fl_1":"report.pdf"}`) + assert.equal(commands.filter(isTransfer).length, 1) + // Said once, without the link. + assert.equal(warned.mock.callCount(), 1) + assert.ok(!JSON.stringify(warned.mock.calls[0]?.arguments).includes("store.example")) + + // The next file doesn't try a link again so soon. + kept.sent("more.csv", "rows") + const second = await computers.bash(SESSION, "true", 10) + assert.match(second.output, /^\[Copied the person's files to: \/workspace\/files\/more\.csv\]/) + assert.deepEqual(kept.carried, ["in:fl_1", "in:fl_2"]) + assert.equal(commands.filter(isTransfer).length, 1) + assert.equal(Buffer.from(await provider.blocks?.files?.read(await computers.vm(SESSION), "/workspace/files/more.csv", { timeoutMs: 1_000 }) ?? []).toString(), "rows") + await computers.delete(SESSION) +}) + +test("changed outputs are sent by link in one command; the store decides what is kept", async () => { + const asked: Array<{ name: string; size: number; replaces: string | undefined }> = [] + const kept = keptFiles({ + async receive(_session, input) { + asked.push({ name: input.name, size: input.size, replaces: input.replaces }) + if (input.name === "full.bin") throw Object.assign(new Error("This conversation's files are full."), { code: "files_full" }) + return { + url: link(`incoming-${asked.length}`), + // The chart arrived; the report didn't; the grown file turned out larger than the limit once measured. + finish: async () => { + if (input.name === "grown.bin") throw Object.assign(new Error("That file is too large."), { code: "file_too_large" }) + return input.name === "chart.png" ? kept.keep("chart.png", encode("12345"), "agent") : null + }, + } + }, + }) + const listing = [`chart.png\t5\t1.5`, `huge.bin\t2000000\t1.5`, `full.bin\t9\t1.5`, `grown.bin\t9\t1.5`, `drafts/it's "final".pdf\t7\t2.5`, `same.txt\t3\t7.5`].join("\n") + const { computers, commands, written, provider } = scripted(kept.files, (command) => + isListing(command) ? `${listing}\n---computer-exported---\n{"same.txt":"3:7.5|fl_old","chart.png":"4:1.0|fl_before"}\n` : undefined) + await provider.blocks?.files?.write(await computers.vm(SESSION), `/workspace/out/drafts/it's "final".pdf`, encode("report!"), { timeoutMs: 1_000 }) + written.clear() + + const result = await computers.bash(SESSION, "make things", 10) + assert.equal(result.output, [ + "[exit 0]", + `[Added to their Files: chart.png, it's "final".pdf]`, + "[Not added to their Files: huge.bin (1.9 MB; the most a kept file can be is 977 KB); full.bin (This conversation's files are full.); grown.bin (That file is too large.). Make it smaller, or tell them it is too big to hand over.]", + ].join("\n")) + // Asked for what changed and could fit; the chart is a new version of the file it was before. + assert.deepEqual(asked, [ + { name: "chart.png", size: 5, replaces: "fl_before" }, + { name: "full.bin", size: 9, replaces: undefined }, + { name: "grown.bin", size: 9, replaces: undefined }, + { name: `it's "final".pdf`, size: 7, replaces: undefined }, + ]) + // Only what didn't arrive went through the runner. + assert.deepEqual(kept.carried, [`out:it's "final".pdf`]) + + const planPath = [...written.keys()].find((path) => path.includes(".openwork-transfer-")) ?? "" + const links = `${planPath}.d` + assert.equal(written.get(planPath)?.text.split("\nnext\n").length, 3) + assert.ok(written.get(planPath)?.text.includes(`url = "${link("incoming-1")}"\nupload-file = "${links}/0"\nfail\ngloboff\n`)) + const transfer = commands.find(isTransfer) ?? "" + assert.ok(transfer.includes(`umask 077\nmkdir ${links} || exit 1\nln -s -- '/workspace/out/chart.png' ${links}/0\nln -s -- '/workspace/out/grown.bin' ${links}/1\nln -s -- '/workspace/out/drafts/it'"'"'s "final".pdf' ${links}/2\n`)) + assert.ok(commands.every((command) => !command.includes("store.example") && !command.includes("X-Amz"))) + assert.equal(commands.filter(isTransfer).length, 1) + + // Each version is recorded with what it became, so none is looked at again. + assert.deepEqual(JSON.parse(written.get("/workspace/.jobs/.exported.json")?.text ?? "{}"), { + "same.txt": "3:7.5|fl_old", + "chart.png": "5:1.5|fl_1", + "huge.bin": "2000000:1.5|", + "full.bin": "9:1.5|", + "grown.bin": "9:1.5|", + [`drafts/it's "final".pdf`]: "7:2.5|fl_2", + }) + await computers.delete(SESSION) +}) + +/** The computer's own tools: GNU stat, find and mv, and curl. Without them the last test has nothing to run. */ +function hasComputerTools() { + try { + execFileSync("sh", ["-c", "stat -c %s / && find / -maxdepth 0 -printf x && mv --version && curl --version"], { stdio: "ignore" }) + return true + } catch { + return false + } +} + +test("on a real shell with curl: files arrive whole, with their permissions and record, and outputs reach the store", { skip: hasComputerTools() ? false : "needs GNU stat, find and mv, and curl (run it on Linux)" }, async () => { + const root = mkdtempSync(join(tmpdir(), "computer-files-")) + const local = (text: string) => text.replaceAll("/workspace", `${root}/workspace`).replaceAll("/home/ubuntu", `${root}/home/ubuntu`) + for (const folder of ["workspace/files", "workspace/out", "workspace/.jobs", "home/ubuntu"]) mkdirSync(join(root, folder), { recursive: true }) + + // The store: hands out links to this server, and keeps what is sent to them. `refused` answers 403. + const stored = new Map() + const refused = new Set() + const requests: string[] = [] + const server = createServer((request, response) => { + const url = new URL(request.url ?? "/", "http://store") + const key = decodeURIComponent(url.pathname.slice(1)) + requests.push(`${request.method} ${key}`) + if (!url.searchParams.has("X-Amz-Signature") || refused.has(key)) return void response.writeHead(403).end("") + if (request.method === "GET") return void (stored.has(key) ? response.writeHead(200).end(stored.get(key)) : response.writeHead(404).end()) + const chunks: Buffer[] = [] + request.on("data", (chunk: Buffer) => chunks.push(chunk)).on("end", () => { + stored.set(key, Buffer.concat(chunks)) + response.writeHead(200).end() + }) + }) + await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve)) + const address = server.address() as AddressInfo + const linkTo = (key: string) => `http://127.0.0.1:${address.port}/${encodeURIComponent(key)}?X-Amz-Signature=abc&X-Amz-Expires=120` + + let incoming = 0 + const kept = keptFiles({ + downloadUrl: async (_session, id) => linkTo(id), + async receive(_session, input) { + const key = `incoming-${++incoming}` + return { url: linkTo(key), finish: async () => (stored.has(key) ? kept.keep(input.name, new Uint8Array(stored.get(key) ?? []), "agent") : null) } + }, + }) + + // A computer on this machine: its files under `root`, its commands in a real shell. Commands of a desktop go + // through its helper, which here just runs them. + const run = (spec: RunSpec) => new Promise((resolve) => { + const bridged = / \| sudo -n \/usr\/local\/sbin\/openwork-desktop-bridge command$/.exec(spec.command) + const command = bridged + ? `umask 007; cd ${root}/workspace; exec bash -c "$(${local(spec.command.slice(0, bridged.index))} | node -e 'process.stdout.write(JSON.parse(require("fs").readFileSync(0, "utf8")).command)')"` + : local(spec.command) + execFile("sh", ["-c", command], { env: { ...process.env, ...spec.env }, encoding: "utf8" }, (error, stdout, stderr) => + resolve({ exitCode: error ? (typeof error.code === "number" ? error.code : 1) : 0, stdout, stderr })) + }) + const fake = createFakeProvider() + const provider: SandboxProvider = { + ...fake, + blocks: { + ...fake.blocks, + run: (_handle, spec) => run(spec), + files: { + async read(_handle, path) { return new Uint8Array(readFileSync(local(path))) }, + async write(_handle, path, content, opts) { + mkdirSync(dirname(local(path)), { recursive: true }) + writeFileSync(local(path), path.includes(".openwork-transfer-") ? local(Buffer.from(content).toString()) : content) + if (opts.mode !== undefined) chmodSync(local(path), opts.mode) + }, + async stat(_handle, path) { return existsSync(local(path)) ? { size: statSync(local(path)).size, kind: "file" } : null }, + }, + }, + } + + try { + const computers = new Computers({ provider, idlePauseMs: 60_000, keepDays: 1, desktop: true }, { files: kept.files, readFile: async () => ({ text: "file" }) }) + const report = kept.sent("report.pdf", "%PDF the report") + const odd = kept.sent(`it's a "quoted" $HOME {name}.txt`, "odd name, plain content") + const big = kept.keep("clip.bin", new Uint8Array(3 * 1024 * 1024).fill(7), "user") + for (const file of [report, odd, big]) stored.set(file.id, Buffer.from(kept.bytes.get(file.id) ?? [])) + + const first = await computers.bash(SESSION, `cp files/report.pdf "out/copy of report.pdf" && mkdir -p out/charts && printf 'a chart' > "out/charts/q3's [final] {chart}.png" && ls files | wc -l`, 30) + assert.equal(first.isError, false) + assert.match(first.output, /^\[Copied the person's files to: \/workspace\/files\/clip\.bin, \/workspace\/files\/it's a "quoted" \$HOME \{name\}\.txt, \/workspace\/files\/report\.pdf\]\n\[exit 0\]\n3\n/) + assert.match(first.output, /\[Added to their Files: (copy of report\.pdf, q3's \[final\] \{chart\}\.png|q3's \[final\] \{chart\}\.png, copy of report\.pdf)\]$/) + + // The person's files: whole, readable and writable by the desktop's shared group, and recorded. + const filesDir = join(root, "workspace/files") + assert.deepEqual(readdirSync(filesDir).sort(), ["clip.bin", `it's a "quoted" $HOME {name}.txt`, "report.pdf"]) + assert.equal(readFileSync(join(filesDir, "report.pdf"), "utf8"), "%PDF the report") + assert.equal(readFileSync(join(filesDir, odd.name), "utf8"), "odd name, plain content") + assert.ok(readFileSync(join(filesDir, "clip.bin")).equals(Buffer.from(kept.bytes.get(big.id) ?? []))) + for (const name of readdirSync(filesDir)) assert.equal(statSync(join(filesDir, name)).mode & 0o777, 0o660, name) + assert.deepEqual(JSON.parse(readFileSync(join(root, "workspace/.jobs/.copied.json"), "utf8")), { [report.id]: "report.pdf", [odd.id]: odd.name, [big.id]: "clip.bin" }) + + // The outputs: sent by the computer itself, kept from what the store received. + assert.deepEqual(kept.rows.filter((file) => file.source === "agent").map((file) => [file.name, Buffer.from(kept.bytes.get(file.id) ?? []).toString()]).sort(), [ + ["copy of report.pdf", "%PDF the report"], + ["q3's [final] {chart}.png", "a chart"], + ]) + // Nothing went through the runner, and nothing with a link in it is left on the computer. + assert.deepEqual(kept.carried, []) + assert.deepEqual(readdirSync(join(root, "home/ubuntu")), []) + assert.deepEqual(requests.filter((request) => request.startsWith("GET")).sort(), [`GET ${big.id}`, `GET ${odd.id}`, `GET ${report.id}`].sort()) + assert.equal(requests.filter((request) => request.startsWith("PUT")).length, 2) + + // A store that refuses one file: that one goes through the runner, the others still by link. + const later = kept.sent("later.txt", "sent later") + const refusedFile = kept.sent("refused.txt", "the store said no") + stored.set(later.id, Buffer.from("sent later")) + refused.add(refusedFile.id) + const second = await computers.bash(SESSION, "cat files/later.txt files/refused.txt", 30) + assert.match(second.output, /\[exit 0\]\nsent laterthe store said no$/) + assert.deepEqual(kept.carried, [`in:${refusedFile.id}`]) + assert.equal(statSync(join(filesDir, "refused.txt")).mode & 0o777, 0o660) + assert.deepEqual(readdirSync(filesDir).filter((name) => name.startsWith(".")), []) + assert.deepEqual(readdirSync(join(root, "home/ubuntu")), []) + await computers.delete(SESSION) + } finally { + await new Promise((resolve) => server.close(resolve)) + rmSync(root, { recursive: true, force: true }) + } +}) diff --git a/ee/packages/headless-computer/src/computer.ts b/ee/packages/headless-computer/src/computer.ts index 06cf971ec22..3ce84ef2031 100644 --- a/ee/packages/headless-computer/src/computer.ts +++ b/ee/packages/headless-computer/src/computer.ts @@ -1,4 +1,4 @@ -import { createHash } from "node:crypto" +import { createHash, randomUUID } from "node:crypto" import { posix } from "node:path" import { ensureRunning, isSandboxError, shellQuote, SCOPE_LABEL, withBlocks, type RunSpec, type SandboxHandle, type SandboxProvider, type SandboxSpec, type SandboxState } from "@openwork/sandbox" import { z } from "zod" @@ -53,6 +53,23 @@ const OUT_DIR = `${WORKSPACE}/out` /** Which /workspace/out files were already added to the person's Files (path → size:mtime), kept in the VM itself. */ const EXPORTED_MANIFEST = `${WORKSPACE}/.jobs/.exported.json` const MANIFEST_MARKER = "---computer-exported---" +/** Which of the person's files were copied in (file id → its name in /workspace/files), kept in the VM next to it. */ +const COPIED_MANIFEST = `${WORKSPACE}/.jobs/.copied.json` +const COPIED_MARKER = "---computer-copied---" +/** + * Files move between the store and the computer directly when the store hands out links (ComputerFiles): the + * computer fetches or sends them with curl in one command, so the bytes never pass through the runner. The time a + * transfer gets grows with its size, at half a megabyte a second; its links stop working a minute after that. + */ +const TRANSFER_BASE_MS = 60_000 +const TRANSFER_MS_PER_MB = 2_000 +const MAX_TRANSFER_MS = 280_000 +const TRANSFER_LINK_GRACE_SECONDS = 60 +const TRANSFERS_AT_ONCE = 4 +/** After a transfer in which nothing arrived, files go through the runner for this long before links are tried again. */ +const DIRECT_RETRY_MS = 10 * 60_000 +const TRANSFER_PLAN = ".openwork-transfer-" +const COPIED_LINE = "copied" /** What a login shell would set, so `pip install --user` and npm globals are found without the slow profile. */ function computerBlocks(provider: SandboxProvider) { return withBlocks(provider, ["run", "files"], "Workbot computer") } const GUEST_PATH = "/home/ubuntu/.local/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" @@ -89,6 +106,8 @@ export class Computers { private readonly checkpointing = new Set() /** Saved file ids already copied into each session's VM by this process. */ private readonly copied = new Map>() + /** Until when files go through the runner because the last direct transfer brought nothing. */ + private directOffUntil = 0 /** The desktop request being timed for each VM, by its slug. */ private readonly timings = new Map() /** How each session's desktop last started in this process; a `stale` one is replaced the next time it is idle. */ @@ -602,7 +621,12 @@ export class Computers { }) } - /** Copies files the person sent into /workspace/files. Returns the paths it copied this time. */ + /** + * Copies files the person sent into /workspace/files. Returns the paths it copied this time. What is already there + * is read in one command: the folder's names and sizes, and the record of earlier copies kept on the computer, so + * a new process doesn't ask about each file. Missing files are fetched by the computer itself where the store + * hands out links, all in one command; otherwise, or when that doesn't arrive, they go through the runner. + */ private async copyUploads(sessionId: string): Promise { if (!this.files) return [] const done = this.copied.get(sessionId) ?? new Set() @@ -610,44 +634,148 @@ export class Computers { const pending = this.files.list(sessionId).filter((file) => file.source === "user" && !done.has(file.id)) if (pending.length === 0) return [] const vm = await this.vm(sessionId) - const copied: string[] = [] + const seen = await this.execute(vm, { + command: `find ${FILES_DIR} -mindepth 1 -maxdepth 1 -printf '%f\\t%s\\n' 2>/dev/null; echo '${COPIED_MARKER}'; cat ${COPIED_MANIFEST} 2>/dev/null || echo '{}'`, + timeoutMs: 30_000, + }) + const [listed = "", recordText = "{}"] = (seen.stdout ?? "").split(`${COPIED_MARKER}\n`) + const present = new Map(listed.split("\n").filter(Boolean).map((line) => { + const tab = line.lastIndexOf("\t") + return [line.slice(0, tab), Number(line.slice(tab + 1))] as const + })) + // File id → the name it was copied under. It only counts while that file is still there at that size: a + // replacement computer carries this record over but not /workspace/files. + const record = z.record(z.string(), z.string()).catch({}).parse(safeJson(recordText.trim())) + const missing: Array<{ file: ComputerFile; path: string }> = [] + let recorded = false for (const file of pending) { - const path = await this.uploadPath(vm, file.name, file.size) - if (path.exists) { + const earlier = record[file.id] + if (earlier !== undefined && present.get(earlier) === file.size) { done.add(file.id) continue } - const found = await this.files.read(sessionId, file.id) - if (!found) { + const place = uploadPath(present, file.name, file.size) + record[file.id] = place.name + recorded = true + if (place.exists) { done.add(file.id) continue } - await this.blocks.files.write(vm, path.path, found.bytes, { timeoutMs: 30_000, ...(this.desktop ? { mode: 0o660 } : {}) }) + // Taken from here on, so another file of this name sent with it gets its own. + present.set(place.name, file.size) + missing.push({ file, path: `${FILES_DIR}/${place.name}` }) + } + const recordJson = JSON.stringify(record) + const fetched = await this.fetchUploads(sessionId, vm, missing, recordJson) + const copied: string[] = [] + for (const { file, path } of missing) { + if (!fetched.has(file.id)) { + const found = await this.files.read(sessionId, file.id) + if (!found) { + done.add(file.id) + continue + } + await this.blocks.files.write(vm, path, found.bytes, { timeoutMs: 30_000, ...(this.desktop ? { mode: 0o660 } : {}) }) + } done.add(file.id) - copied.push(path.path) + copied.push(path) } + // The fetching command writes the record itself; otherwise it is written here, when it changed. Without it the + // next process still finds the files by name and size. + if (recorded && fetched.size === 0) await this.blocks.files.write(vm, COPIED_MANIFEST, new TextEncoder().encode(recordJson), { timeoutMs: 30_000 }).catch(() => undefined) return copied } - /** `/workspace/files/`, or a variant when a different file already has that name. */ - private async uploadPath(vm: SandboxHandle, name: string, size: number) { - const base = posix.basename(name) || "file" - const extension = posix.extname(base) - const stem = base.slice(0, base.length - extension.length) - for (let attempt = 0; attempt < 20; attempt += 1) { - const path = `${FILES_DIR}/${attempt === 0 ? base : `${stem} (${attempt + 1})${extension}`}` - const stat = await this.blocks.files.stat(vm, path, { timeoutMs: 10_000 }) - if (!stat) return { path, exists: false } - if (stat.size === size) return { path, exists: true } + /** + * Has the computer fetch the person's files from the store itself, all in one command. Each is fetched under a + * temporary name, checked against its size, given the permissions a copy through the runner gets, and only then + * moved into place. Returns the ids of the files that arrived; the rest are left for the runner to carry. + */ + private async fetchUploads(sessionId: string, vm: SandboxHandle, missing: Array<{ file: ComputerFile; path: string }>, recordJson: string): Promise> { + const arrived = new Set() + const files = this.files + if (!files?.downloadUrl || missing.length === 0 || Date.now() < this.directOffUntil) return arrived + const timeoutMs = transferTimeoutMs(missing.reduce((total, { file }) => total + file.size, 0)) + const token = randomUUID().replaceAll("-", "") + let outcome: unknown + try { + const fetching: Array<{ file: ComputerFile; path: string; part: string; url: string }> = [] + for (const [index, entry] of missing.entries()) { + const url = await files.downloadUrl(sessionId, entry.file.id, linkSeconds(timeoutMs)) + if (url) fetching.push({ ...entry, url, part: `${FILES_DIR}/.openwork-part-${token}-${index}` }) + } + if (fetching.length === 0) return arrived + const result = await this.transfer(vm, { + token, + blocks: fetching.map(({ url, part }) => [`url = "${url}"`, `output = "${part}"`]), + timeoutMs, + cleanup: fetching.map(({ part }) => part), + after: [ + ...fetching.map(({ file, path, part }, index) => + `[ "$(stat -c %s ${part} 2>/dev/null)" = ${file.size} ] && ${this.desktop ? `chmod 660 ${part} && ` : ""}mv -fT -- ${part} ${shellQuote(path)} && echo ${COPIED_LINE} ${index}`), + `printf %s ${shellQuote(recordJson)} > ${COPIED_MANIFEST}`, + ], + }) + for (const line of result.stdout.split("\n")) { + const [word, index] = line.trim().split(" ") + const entry = word === COPIED_LINE ? fetching[Number(index)] : undefined + if (entry) arrived.add(entry.file.id) + } + outcome = result.stderr + } catch (error) { + outcome = error } - return { path: `${FILES_DIR}/${Date.now()}-${base}`, exists: false } + if (arrived.size === 0) this.directFailed(sessionId, "to the computer", missing.length, outcome) + return arrived + } + + /** + * Runs curl on the computer over `blocks`, one transfer each, then the `after` commands. The links are written to a + * file, never to a command line, and the file is removed when the command ends. On a desktop, commands run as a + * separate account that can't enter the folder it is in; elsewhere there is one account. + */ + private async transfer(vm: SandboxHandle, input: { token: string; blocks: string[][]; before?: string[]; after?: string[]; cleanup?: string[]; timeoutMs: number }) { + const folder = this.transferFolder + const plan = `${folder}/${TRANSFER_PLAN}${input.token}` + const seconds = Math.max(10, Math.floor(input.timeoutMs / 1000) - 10) + const config = input.blocks.map((block) => [...block, "fail", "globoff", "connect-timeout = 10", `max-time = ${seconds}`].join("\n")).join("\nnext\n") + await this.blocks.files.write(vm, plan, new TextEncoder().encode(`${config}\n`), { timeoutMs: 10_000 }) + return this.execute(vm, { + command: [ + `trap 'rm -rf ${[plan, `${plan}.d`, ...(input.cleanup ?? [])].join(" ")}' EXIT`, + `chmod 600 ${plan}`, + // Whatever an interrupted transfer left behind; its links have long stopped working. + `find ${folder} -maxdepth 1 -name '${TRANSFER_PLAN}*' -mmin +15 -exec rm -rf {} + 2>/dev/null`, + ...(input.before ?? []), + `curl --silent --show-error --parallel --parallel-max ${TRANSFERS_AT_ONCE} --config ${plan} >/dev/null`, + ...(input.after ?? []), + "true", + ].join("\n"), + timeoutMs: input.timeoutMs, + }) + } + + /** Where a transfer's links are written: the provider account's own folder on a desktop, the jobs folder elsewhere. */ + private get transferFolder() { + return this.desktop ? "/home/ubuntu" : `${WORKSPACE}/.jobs` + } + + /** Files go through the runner for a while after a direct transfer in which nothing arrived. */ + private directFailed(sessionId: string, direction: string, files: number, outcome: unknown) { + this.directOffUntil = Date.now() + DIRECT_RETRY_MS + const detail = (outcome instanceof Error ? outcome.message : typeof outcome === "string" ? outcome : "").replace(/https?:\/\/\S+/g, "[link]").trim().slice(-300) + console.warn("[headless-computer] a direct file transfer brought nothing; files go through the runner for now", { sessionId, direction, files, ...(detail ? { detail } : {}) }) } - /** Adds new or changed files under /workspace/out to the person's Files. Returns their names. */ - /** Adds new or changed files in /workspace/out to the person's files. Returns what was added, and what wasn't. */ + /** + * Adds new or changed files in /workspace/out to the person's files. Returns what was added, and what wasn't. The + * computer sends them to the store itself where the store hands out links, all in one command; what doesn't + * arrive that way goes through the runner. + */ private async exportOutputs(sessionId: string): Promise<{ saved: string[]; skipped: string[] }> { if (!this.files) return { saved: [], skipped: [] } - const maxBytes = Math.min(MAX_EXPORT_BYTES, this.files.maxFileBytes) + const files = this.files + const maxBytes = Math.min(MAX_EXPORT_BYTES, files.maxFileBytes) const vm = await this.vm(sessionId) const listing = await this.execute(vm, { command: `find ${OUT_DIR} -type f -not -name '.*' -printf '%P\\t%s\\t%T@\\n' 2>/dev/null; echo '${MANIFEST_MARKER}'; cat ${EXPORTED_MANIFEST} 2>/dev/null || echo '{}'`, @@ -659,6 +787,13 @@ export class Computers { const saved: string[] = [] const skipped: string[] = [] let changed = false + // A version that doesn't fit is recorded as seen, so it is reported once rather than after every command. + const notKept = (output: Output, reason: string) => { + skipped.push(`${output.relative} (${reason})`) + manifest[output.relative] = `${output.signature}|${output.previousId ?? ""}` + changed = true + } + const outputs: Output[] = [] for (const line of listed.split("\n").filter(Boolean)) { const [relative, sizeText, mtime] = line.split("\t") if (!relative || !sizeText) continue @@ -667,35 +802,39 @@ export class Computers { const [previousSignature, recordedId] = (manifest[relative] ?? "").split("|") if (previousSignature === signature) continue // Outputs recorded before ids were kept: the newest file it made under that name is the earlier version. - const previousId = recordedId || (previousSignature ? this.files.list(sessionId).find((file) => file.source === "agent" && file.name === posix.basename(relative))?.id : undefined) - // A version that doesn't fit is recorded as seen, so it is reported once rather than after every command. - const notKept = (reason: string) => { - skipped.push(`${relative} (${reason})`) - manifest[relative] = `${signature}|${previousId ?? ""}` - changed = true - } - if (size > maxBytes) { - notKept(`${formatBytes(size)}; the most a kept file can be is ${formatBytes(maxBytes)}`) + const previousId = recordedId || (previousSignature ? files.list(sessionId).find((file) => file.source === "agent" && file.name === posix.basename(relative))?.id : undefined) + const output = { relative, size, signature, previousId } + if (size > maxBytes) notKept(output, `${formatBytes(size)}; the most a kept file can be is ${formatBytes(maxBytes)}`) + else outputs.push(output) + } + const sent = await this.sendOutputs(sessionId, vm, outputs) + for (const output of outputs) { + const { relative, previousId } = output + const direct = sent.get(relative) + if (direct instanceof Error) { + notKept(output, direct.message) continue } - const bytes = new Uint8Array(await this.blocks.files.read(vm, `${OUT_DIR}/${relative}`, { timeoutMs: 30_000 })) // A new version of a file already handed over updates it in place: one file in their Files, and anything // showing it (the answer's card, an open preview) moves to the new version. - let file: ComputerFile | null = null - try { - file = - (previousId ? await this.files.replace(sessionId, previousId, { bytes }) : null) ?? - (await this.files.add(sessionId, { name: posix.basename(relative), bytes, source: "agent" })) - } catch (error) { - // Over the conversation's limits: final for this version. Anything else (storage briefly unreachable) - // is left unrecorded, so the next command tries again. - if (isLimitError(error)) { - notKept(error.message) - continue + let file = direct ?? null + if (!file) { + const bytes = new Uint8Array(await this.blocks.files.read(vm, `${OUT_DIR}/${relative}`, { timeoutMs: 30_000 })) + try { + file = + (previousId ? await files.replace(sessionId, previousId, { bytes }) : null) ?? + (await files.add(sessionId, { name: posix.basename(relative), bytes, source: "agent" })) + } catch (error) { + // Over the conversation's limits: final for this version. Anything else (storage briefly unreachable) + // is left unrecorded, so the next command tries again. + if (isLimitError(error)) { + notKept(output, error.message) + continue + } + throw error } - throw error } - manifest[relative] = `${signature}|${file.id}` + manifest[relative] = `${output.signature}|${file.id}` changed = true saved.push(file.name) // Slides, documents and PDFs get page images for the preview panel, made in the background. @@ -705,6 +844,53 @@ export class Computers { return { saved, skipped } } + /** + * Has the computer send changed outputs to the store itself, all in one command. The store measures what arrived + * before keeping it, so the limits hold whatever was sent. Returns, by path, the kept file or the limit it ran + * into; an output that isn't in the result didn't arrive and is left for the runner to carry. + */ + private async sendOutputs(sessionId: string, vm: SandboxHandle, outputs: Output[]): Promise> { + const settled = new Map() + const files = this.files + if (!files?.receive || outputs.length === 0 || Date.now() < this.directOffUntil) return settled + const timeoutMs = transferTimeoutMs(outputs.reduce((total, output) => total + output.size, 0)) + const sending: Array<{ output: Output; url: string; finish(): Promise }> = [] + for (const output of outputs) { + try { + const place = await files.receive(sessionId, { name: posix.basename(output.relative), size: output.size, source: "agent", replaces: output.previousId, expiresSeconds: linkSeconds(timeoutMs) }) + if (place) sending.push({ output, ...place }) + } catch (error) { + // Anything but a limit is the store's trouble: the runner's own way reports it. + if (isLimitError(error)) settled.set(output.relative, error) + } + } + if (sending.length === 0) return settled + const token = randomUUID().replaceAll("-", "") + const links = `${this.transferFolder}/${TRANSFER_PLAN}${token}.d` + // curl reads each output through a link with a plain name, so a file's own name never has to be quoted for it. + const outcome: unknown = await this.transfer(vm, { + token, + blocks: sending.map(({ url }, index) => [`url = "${url}"`, `upload-file = "${links}/${index}"`]), + timeoutMs, + before: ["umask 077", `mkdir ${links} || exit 1`, ...sending.map(({ output }, index) => `ln -s -- ${shellQuote(`${OUT_DIR}/${output.relative}`)} ${links}/${index}`)], + }).then((result) => result.stderr, (error: unknown) => error) + let arrived = 0 + for (const { output, finish } of sending) { + try { + const file = await finish() + if (!file) continue + settled.set(output.relative, file) + arrived += 1 + } catch (error) { + if (!isLimitError(error)) continue + settled.set(output.relative, error) + arrived += 1 + } + } + if (arrived === 0) this.directFailed(sessionId, "from the computer", sending.length, outcome) + return settled + } + /** Renders a kept file's preview on demand (files from before previews, or ones the person sent). */ async previewSavedFile(sessionId: string, fileId: string) { if (!this.files) return @@ -816,6 +1002,31 @@ export class Computers { } } +/** A changed file in /workspace/out: its path there, the version seen (size:mtime), and the kept file it updates. */ +type Output = { relative: string; size: number; signature: string; previousId: string | undefined } + +/** + * A person's file's name in /workspace/files: its own, or a variant when a different file already has that name. + * `present` is what the folder holds (name → size); a file of the same name and size is taken to be this one. + */ +function uploadPath(present: ReadonlyMap, name: string, size: number) { + const plain = posix.basename(name) + const base = plain && plain !== "." && plain !== ".." ? plain : "file" + const extension = posix.extname(base) + const stem = base.slice(0, base.length - extension.length) + for (let attempt = 0; attempt < 20; attempt += 1) { + const candidate = attempt === 0 ? base : `${stem} (${attempt + 1})${extension}` + const taken = present.get(candidate) + if (taken === undefined) return { name: candidate, exists: false } + if (taken === size) return { name: candidate, exists: true } + } + return { name: `${Date.now()}-${base}`, exists: false } +} + +/** How long a direct transfer of this many bytes may take, and how long its links work. */ +const transferTimeoutMs = (bytes: number) => Math.min(MAX_TRANSFER_MS, TRANSFER_BASE_MS + Math.ceil(bytes / (1024 * 1024)) * TRANSFER_MS_PER_MB) +const linkSeconds = (timeoutMs: number) => Math.ceil(timeoutMs / 1000) + TRANSFER_LINK_GRACE_SECONDS + /** Width and height from a PNG's header. */ function pngSize(bytes: Uint8Array | undefined) { if (!bytes || bytes.byteLength < 24) return null diff --git a/ee/packages/headless-computer/src/types.ts b/ee/packages/headless-computer/src/types.ts index 0ed3c88694e..2ed782bc471 100644 --- a/ee/packages/headless-computer/src/types.ts +++ b/ee/packages/headless-computer/src/types.ts @@ -24,6 +24,21 @@ export interface ComputerFiles { replace(sessionId: string, id: string, input: { bytes: Uint8Array }): Promise /** Page images of a slide deck, document or PDF, for previews. */ putPreview(sessionId: string, id: string, pages: Array>, size: { width: number; height: number }): Promise + /** + * A link the computer fetches a kept file from; null when there is none to give. This and `receive` exist only + * with a store that can hand out links: bytes then go between the store and the computer directly and never pass + * through the runner. A link reaches one file and stops working after `expiresSeconds`. Without them, or when a + * transfer doesn't arrive, files go through `read`, `add` and `replace`. + */ + downloadUrl?(sessionId: string, id: string, expiresSeconds: number): Promise + /** + * A link the computer sends a file to. `finish` checks what arrived and keeps it as a new file, or as the new + * version of `replaces`; null when nothing arrived. Both throw like `add` when the file doesn't fit. + */ + receive?( + sessionId: string, + input: { name: string; size: number; source: "agent"; replaces?: string; expiresSeconds: number }, + ): Promise<{ url: string; finish(): Promise } | null> } /**