From 229b19a6a726d47d149bb646bedcb17aed65ebfb Mon Sep 17 00:00:00 2001 From: Carl Calaquian Date: Wed, 5 Aug 2026 03:27:31 +0800 Subject: [PATCH] fix(sandbox-tenki): address coderabbit/cubic review feedback - Validate `publicKey` as a non-empty, single-line authorized_keys entry. - Release the lifecycle queue when a canceled execute abandons its resume, so an RPC with no client-side deadline cannot wedge later transitions. - Prefer the run's aggregate output when an output pump dies mid-stream. - Collapse CR/LF in a Tenki `reason` before composing the one-line diagnostic. - Document `TENKI_AUTH_TOKEN` alongside `TENKI_API_KEY` in the workspace docs. --- packages/sandbox-tenki/src/sandbox.spec.ts | 95 +++++++++++++++++++++- packages/sandbox-tenki/src/sandbox.ts | 87 +++++++++++++++++--- packages/sandbox-tenki/src/tools.spec.ts | 36 ++++++-- packages/sandbox-tenki/src/tools.ts | 4 + packages/sandbox-tenki/src/utils.spec.ts | 43 ++++++++++ packages/sandbox-tenki/src/utils.ts | 34 +++++++- website/docs/workspaces/sandbox.md | 2 +- 7 files changed, 277 insertions(+), 24 deletions(-) diff --git a/packages/sandbox-tenki/src/sandbox.spec.ts b/packages/sandbox-tenki/src/sandbox.spec.ts index 0096a0648..26756bff6 100644 --- a/packages/sandbox-tenki/src/sandbox.spec.ts +++ b/packages/sandbox-tenki/src/sandbox.spec.ts @@ -50,6 +50,12 @@ type HandleOptions = { killThrowsWith?: unknown; rejectWith?: unknown; errorStdout?: boolean; + /** + * Deliver this many stdout chunks over the stream, then kill it with a + * transport error. The resolved result still carries the full aggregate, so + * this models a mid-stream failure that loses bytes only on the stream. + */ + errorStdoutAfter?: number; }; const makeHandle = (options: HandleOptions = {}) => { @@ -68,6 +74,7 @@ const makeHandle = (options: HandleOptions = {}) => { killThrowsWith, rejectWith, errorStdout = false, + errorStdoutAfter, } = options; let stdoutCtl!: ReadableStreamDefaultController; @@ -89,9 +96,12 @@ const makeHandle = (options: HandleOptions = {}) => { stderrCtl = controller; }, }); - for (const chunk of stdout) { + stdout.forEach((chunk, index) => { + if (errorStdoutAfter !== undefined && index >= errorStdoutAfter) { + return; + } stdoutCtl.enqueue(typeof chunk === "string" ? enc.encode(chunk) : chunk); - } + }); for (const chunk of stderr) { stderrCtl.enqueue(typeof chunk === "string" ? enc.encode(chunk) : chunk); } @@ -121,13 +131,30 @@ const makeHandle = (options: HandleOptions = {}) => { } }; + // Kill the stream only once the pump has drained what was enqueued — + // `controller.error()` discards a queued chunk, which would leave the buffer + // empty and mask the partial-read the test is after. + const failStdoutAfterDrain = () => { + setTimeout(() => { + try { + stdoutCtl.error(new Error("stream boom")); + } catch { + // controller may already be closed + } + }, 0); + }; + if (rejectWith !== undefined) { safeClose(stdoutCtl); safeClose(stderrCtl); rejectResult(rejectWith); } else if (!hangUntilKill) { if (!keepStreamsOpen) { - safeClose(stdoutCtl); + if (errorStdoutAfter === undefined) { + safeClose(stdoutCtl); + } else { + failStdoutAfterDrain(); + } safeClose(stderrCtl); } resolveResult(result); @@ -413,6 +440,23 @@ describe("TenkiSandbox.execute", () => { expect(result.stderr).toBe("side"); }); + // A transport failure part-way through leaves the pump holding a prefix of + // the output. Preferring it would silently drop the rest with `truncated` + // unset, so the completed run's aggregate wins instead. + it("prefers the run's aggregate stdout when the stream dies mid-read", async () => { + const session = makeSession({ + run: vi.fn(() => + makeHandle({ stdout: ["first ", "second"], errorStdoutAfter: 1, exitCode: 0 }), + ), + }); + const sandbox = new TenkiSandbox({ session: session as never }); + + const result = await sandbox.execute({ command: "flaky" }); + + expect(result.stdout).toBe("first second"); + expect(result.stdoutTruncated).toBe(false); + }); + it("returns an aborted result when the signal is already aborted", async () => { const session = makeSession(); const sandbox = new TenkiSandbox({ session: session as never }); @@ -715,6 +759,51 @@ describe("TenkiSandbox.execute", () => { expect(session.resume).toHaveBeenCalledOnce(); expect(session.run).not.toHaveBeenCalled(); }); + + // Tenki's resume RPC carries no client-side deadline, so a caller that walked + // away on its own timeout must not leave its abandoned transition pinned at + // the head of the lifecycle queue — every later transition would wait on it. + it("releases the lifecycle queue when a canceled execute abandons its resume", async () => { + const session = makeSession({ + resume: vi.fn(() => new Promise(() => {})), + }); + const sandbox = new TenkiSandbox({ session: session as never }); + await sandbox.stop(); + + const result = await settlesWithin(sandbox.execute({ command: "sleep 100", timeoutMs: 10 })); + expect(result.timedOut).toBe(true); + + // The abandoned resume is still in flight; these must not queue behind it. + await expect(settlesWithin(sandbox.stop())).resolves.toBeUndefined(); + await expect(settlesWithin(sandbox.destroy())).resolves.toBeUndefined(); + }); + + it("ignores an abandoned resume that lands after a later transition", async () => { + let releaseResume!: () => void; + const session = makeSession({ + resume: vi.fn( + () => + new Promise((resolve) => { + releaseResume = resolve; + }), + ), + }); + const sandbox = new TenkiSandbox({ session: session as never }); + await sandbox.stop(); + expect(sandbox.status).toBe("idle"); + + const result = await settlesWithin(sandbox.execute({ command: "sleep 100", timeoutMs: 10 })); + expect(result.timedOut).toBe(true); + + // A later transition runs to completion while the resume is still pending. + await settlesWithin(sandbox.stop()); + + // The straggler must not flip the sandbox back to ready behind its back. + releaseResume(); + await new Promise((resolve) => setTimeout(resolve, 0)); + + expect(sandbox.status).toBe("idle"); + }); }); describe("TenkiSandbox lifecycle", () => { diff --git a/packages/sandbox-tenki/src/sandbox.ts b/packages/sandbox-tenki/src/sandbox.ts index 7978169a1..dc58f8dd5 100644 --- a/packages/sandbox-tenki/src/sandbox.ts +++ b/packages/sandbox-tenki/src/sandbox.ts @@ -36,6 +36,18 @@ import { */ export type TenkiSandboxInstance = Session; +/** + * Resolve once `signal` aborts. Used to release the lifecycle queue when a + * caller stops waiting on its own transition; the listener is dropped as soon + * as it fires, and an already-aborted signal resolves immediately. + */ +const abortedPromise = (signal: AbortSignal): Promise => + signal.aborted + ? Promise.resolve() + : new Promise((resolve) => { + signal.addEventListener("abort", () => resolve(), { once: true }); + }); + /** * Public constructor options for {@link TenkiSandbox}. * @@ -159,9 +171,23 @@ export class TenkiSandbox implements WorkspaceSandbox { * Serializes pause/resume RPCs so concurrent lifecycle callers observe one * transition at a time. The chain itself always resolves; the promise * returned to a caller still carries that caller's transition failure. + * + * The one case where a transition does not hold the chain to completion is a + * caller that cancels: it releases its slot at that point, so an RPC with no + * deadline cannot wedge the queue. See {@link serializeLifecycleTransition}. */ private lifecycleTransition: Promise = Promise.resolve(); + /** + * Monotonic id handed to each lifecycle transition as it starts running. + * A transition whose caller gave up (see {@link serializeLifecycleTransition}) + * releases the queue while its RPC is still in flight, so it can still settle + * after a later transition has already moved the sandbox. Comparing the token + * it was issued against the current value tells such a straggler that it no + * longer owns the state and must not write it back. + */ + private lifecycleToken = 0; + /** * SSH keys successfully applied via {@link authorizeSshKey}. Tenki's * `updateSshAuthorizedKeys` replaces the whole set and the SDK has no API to @@ -349,8 +375,8 @@ export class TenkiSandbox implements WorkspaceSandbox { * Resume the microVM when a previous {@link stop} paused it, returning the * sandbox to `ready` so commands can run again. No-op otherwise. */ - private async resumeIfPaused(session: Session): Promise { - return this.serializeLifecycleTransition(async () => { + private async resumeIfPaused(session: Session, signal?: AbortSignal): Promise { + return this.serializeLifecycleTransition(async (token) => { if (this.status === "destroyed" || this.session !== session) { throw new Error("Sandbox has been destroyed"); } @@ -367,17 +393,46 @@ export class TenkiSandbox implements WorkspaceSandbox { if (generation !== this.generation || this.session !== session) { throw new Error("Sandbox has been destroyed"); } + // An abandoned transition can land after a later one already ran; the + // newest transition owns the state. Leaving `paused` set is the safe + // direction — the next caller re-issues resume rather than trusting a + // sandbox this one no longer speaks for. + if (token !== this.lifecycleToken) { + return; + } this.paused = false; this.status = "ready"; - }); + }, signal); } - private serializeLifecycleTransition(operation: () => Promise): Promise { - const transition = this.lifecycleTransition.then(operation, operation); - this.lifecycleTransition = transition.then( + /** + * Queue `operation` behind any lifecycle transition already in flight. + * + * When the caller supplies a `signal`, the queue is released as soon as that + * signal aborts rather than only when the operation settles. Tenki's + * pause/resume RPCs carry no client-side deadline, so without this a caller + * that walked away on its own timeout — `execute()` racing cancellation — + * would leave its abandoned transition pinned at the head of the queue and + * every later lifecycle operation would wait on it indefinitely. The + * abandoned RPC still runs to completion in the background; its `token` guard + * keeps it from writing state that a subsequent transition now owns. + */ + private serializeLifecycleTransition( + operation: (token: number) => Promise, + signal?: AbortSignal, + ): Promise { + // The token is issued when the operation actually starts, not when it is + // queued, so it orders transitions by execution rather than arrival. + const run = () => { + this.lifecycleToken += 1; + return operation(this.lifecycleToken); + }; + const transition = this.lifecycleTransition.then(run, run); + const settled = transition.then( () => undefined, () => undefined, ); + this.lifecycleTransition = signal ? Promise.race([settled, abortedPromise(signal)]) : settled; return transition; } @@ -409,7 +464,7 @@ export class TenkiSandbox implements WorkspaceSandbox { return; } - await this.serializeLifecycleTransition(async () => { + await this.serializeLifecycleTransition(async (token) => { // Destroy supersedes a queued/in-flight stop and owns session teardown. if (this.status === "destroyed" || this.session !== session || this.paused) { return; @@ -417,7 +472,11 @@ export class TenkiSandbox implements WorkspaceSandbox { const generation = this.generation; await session.pause(); - if (generation !== this.generation || this.session !== session) { + if ( + generation !== this.generation || + this.session !== session || + token !== this.lifecycleToken + ) { return; } this.paused = true; @@ -574,7 +633,13 @@ export class TenkiSandbox implements WorkspaceSandbox { } } } catch { - // ignore stream errors; result bytes are used as a fallback + // Ignore stream errors, but record that this buffer stopped short of + // end-of-stream so `resolveOutput` can prefer the run's aggregate bytes + // instead of returning a silently truncated prefix. A canceled read is a + // deliberate stop, not a transport failure, so it does not count. + if (!canceled) { + buffer.failed = true; + } } finally { // Flush any bytes the decoder is still holding for a partial code point. if (decoder) { @@ -731,7 +796,9 @@ export class TenkiSandbox implements WorkspaceSandbox { // does not exist yet, so `requestKill()` would have nothing to kill. This // is the only await left between the guards and `session.run()`; the // `runOptions` build below is synchronous. - const resumed = await raceCancellation(this.resumeIfPaused(session)); + const resumed = await raceCancellation( + this.resumeIfPaused(session, cancellationController.signal), + ); if (resumed === cancellationMarker) { return cancellationResult(); } diff --git a/packages/sandbox-tenki/src/tools.spec.ts b/packages/sandbox-tenki/src/tools.spec.ts index 27d477746..ea640b6c5 100644 --- a/packages/sandbox-tenki/src/tools.spec.ts +++ b/packages/sandbox-tenki/src/tools.spec.ts @@ -6,22 +6,24 @@ type PreviewParameters = { safeParse: (input: unknown) => { success: boolean }; }; -const getPreviewParameters = (): PreviewParameters => { +const getToolParameters = (name: string): PreviewParameters => { const sandbox = { getSandbox: vi.fn(), authorizeSshKey: vi.fn(), } as unknown as TenkiSandbox; - const previewTool = createTenkiToolkit(sandbox).tools.find( - (tool) => (tool as { name?: string }).name === "expose_preview_url", + const tool = createTenkiToolkit(sandbox).tools.find( + (candidate) => (candidate as { name?: string }).name === name, ); - if (!previewTool) { - throw new Error("expose_preview_url tool not found"); + if (!tool) { + throw new Error(`${name} tool not found`); } - return (previewTool as { parameters: PreviewParameters }).parameters; + return (tool as { parameters: PreviewParameters }).parameters; }; +const getPreviewParameters = (): PreviewParameters => getToolParameters("expose_preview_url"); + describe("createTenkiToolkit preview input schema", () => { it.each([1, 65535])("accepts boundary port %s", (port) => { expect(getPreviewParameters().safeParse({ port }).success).toBe(true); @@ -42,3 +44,25 @@ describe("createTenkiToolkit preview input schema", () => { expect(getPreviewParameters().safeParse({ port: 3000, ttlMs }).success).toBe(false); }); }); + +describe("createTenkiToolkit ssh input schema", () => { + const getSshParameters = () => getToolParameters("authorize_ssh_key"); + + it("accepts a single-line authorized_keys entry", () => { + expect( + getSshParameters().safeParse({ publicKey: "ssh-ed25519 AAAAC3NzaC1lZDI1 user@host" }).success, + ).toBe(true); + }); + + // The value is forwarded verbatim into the authorized_keys set, where a blank + // entry is meaningless and an embedded newline would author a second key. + it.each([ + ["empty", ""], + ["whitespace only", " "], + ["newline only", "\n"], + ["two keys on separate lines", "ssh-ed25519 AAAA a\nssh-ed25519 BBBB b"], + ["a trailing CRLF line", "ssh-ed25519 AAAA a\r\nssh-ed25519 BBBB b"], + ])("rejects %s", (_label, publicKey) => { + expect(getSshParameters().safeParse({ publicKey }).success).toBe(false); + }); +}); diff --git a/packages/sandbox-tenki/src/tools.ts b/packages/sandbox-tenki/src/tools.ts index 4193c3e6b..4213c977a 100644 --- a/packages/sandbox-tenki/src/tools.ts +++ b/packages/sandbox-tenki/src/tools.ts @@ -52,6 +52,10 @@ export function createTenkiToolkit(sandbox: TenkiSandbox): Toolkit { parameters: z.object({ publicKey: z .string() + .refine( + (value) => value.trim().length > 0 && !/[\r\n]/.test(value), + "publicKey must be a non-empty, single-line authorized_keys entry", + ) .describe("SSH public key in authorized_keys format (e.g. 'ssh-ed25519 AAAA... user')"), }), outputSchema: z.object({ diff --git a/packages/sandbox-tenki/src/utils.spec.ts b/packages/sandbox-tenki/src/utils.spec.ts index e333d8e96..ea1b904cc 100644 --- a/packages/sandbox-tenki/src/utils.spec.ts +++ b/packages/sandbox-tenki/src/utils.spec.ts @@ -114,6 +114,35 @@ describe("OutputBuffer", () => { expect(resolved.content).toBe("fallback"); }); + // A pump that died mid-stream holds only a prefix of the real output, so the + // completed run's aggregate is the more complete source. + it("prefers the aggregate when the pump failed mid-stream", () => { + const buffer = initOutputBuffer(); + appendOutput(buffer, "partial", 100); + buffer.failed = true; + expect(resolveOutput(buffer, "partial output, complete", 100)).toEqual({ + content: "partial output, complete", + truncated: false, + }); + }); + + it("still applies the byte cap to the aggregate after a failed pump", () => { + const buffer = initOutputBuffer(); + appendOutput(buffer, "par", 100); + buffer.failed = true; + expect(resolveOutput(buffer, "partial output", 3)).toEqual({ content: "par", truncated: true }); + }); + + it("flags truncation when the pump failed and there is no aggregate", () => { + const buffer = initOutputBuffer(); + appendOutput(buffer, "partial", 100); + buffer.failed = true; + expect(resolveOutput(buffer, undefined, 100)).toEqual({ + content: "partial", + truncated: true, + }); + }); + it("resolves to empty when every captured chunk was zero-length", () => { const buffer = initOutputBuffer(); appendOutput(buffer, "", 100); @@ -245,6 +274,20 @@ describe("formatRunDiagnostic", () => { ); }); + // The diagnostic is documented as one line and is folded into stderr, so a + // multi-line guest-agent reason must not break it across lines. + it("collapses a multiline reason into a single line", () => { + expect( + formatRunDiagnostic({ exitCode: 127, errno: 2, reason: "exec failed:\n no such file\r\nx" }), + ).toBe("tenki: exec failed: ENOENT (errno 2), reason=exec failed: no such file x"); + }); + + it("collapses a multiline reason with no errno", () => { + expect(formatRunDiagnostic({ exitCode: 137, errno: 0, reason: "oom\nkilled" })).toBe( + "tenki: run ended: reason=oom killed", + ); + }); + it("surfaces an unrecognized reason when the run was signaled", () => { expect(formatRunDiagnostic({ exitCode: 0, signal: "KILL", reason: "engine_terminated" })).toBe( "tenki: run ended: reason=engine_terminated", diff --git a/packages/sandbox-tenki/src/utils.ts b/packages/sandbox-tenki/src/utils.ts index 5d0b1d2f7..b36f65bcc 100644 --- a/packages/sandbox-tenki/src/utils.ts +++ b/packages/sandbox-tenki/src/utils.ts @@ -36,9 +36,21 @@ export type OutputBuffer = { chunks: Buffer[]; size: number; truncated: boolean; + /** + * Set when the pump feeding this buffer died on a stream error rather than + * reaching end-of-stream. The bytes collected so far are then only a prefix + * of the real output, so {@link resolveOutput} prefers the completed run's + * aggregate instead of silently returning the short read. + */ + failed: boolean; }; -export const initOutputBuffer = (): OutputBuffer => ({ chunks: [], size: 0, truncated: false }); +export const initOutputBuffer = (): OutputBuffer => ({ + chunks: [], + size: 0, + truncated: false, + failed: false, +}); export const appendOutput = (buffer: OutputBuffer, chunk: unknown, maxBytes: number): void => { if (maxBytes <= 0) { @@ -143,17 +155,28 @@ export const truncateOutput = ( /** * Resolve the final output for a stream: prefer the streamed buffer when we * captured anything, otherwise fall back to the aggregated bytes on the result. + * + * A pump that died mid-stream is the exception. Its buffer holds a prefix of + * the real output, so preferring it would hand back silently truncated content + * with `truncated` unset; the completed run's aggregate is the more complete + * source and wins. When there is no aggregate to fall back on, the partial + * bytes are kept but flagged as truncated so the loss is visible. Byte-limit + * truncation still applies to whichever source is used. */ export const resolveOutput = ( buffer: OutputBuffer, fallback: string | undefined, maxBytes: number, ): { content: string; truncated: boolean } => { + if (buffer.failed && fallback) { + const resolved = truncateOutput(fallback, maxBytes); + return { content: resolved.content, truncated: resolved.truncated || buffer.truncated }; + } if (buffer.size > 0 || buffer.truncated) { - return { content: toOutputString(buffer), truncated: buffer.truncated }; + return { content: toOutputString(buffer), truncated: buffer.truncated || buffer.failed }; } if (!fallback) { - return { content: "", truncated: false }; + return { content: "", truncated: buffer.failed }; } return truncateOutput(fallback, maxBytes); }; @@ -259,7 +282,10 @@ export const formatRunDiagnostic = (result: unknown): string | undefined => { // `errno` is a non-optional proto scalar, so it arrives as 0 (not absent) // whenever the guest agent has nothing to report. const errno = typeof record.errno === "number" && record.errno !== 0 ? record.errno : undefined; - const rawReason = typeof record.reason === "string" ? record.reason.trim() : ""; + // The diagnostic is documented as a single line and is folded into stderr, + // so a multi-line guest-agent `reason` must not smuggle in line breaks. + const rawReason = + typeof record.reason === "string" ? record.reason.replace(/\s*[\r\n]+\s*/g, " ").trim() : ""; const reason = BENIGN_RUN_REASONS.has(rawReason.toLowerCase()) ? undefined : rawReason || undefined; diff --git a/website/docs/workspaces/sandbox.md b/website/docs/workspaces/sandbox.md index 68234ab1f..318a760b4 100644 --- a/website/docs/workspaces/sandbox.md +++ b/website/docs/workspaces/sandbox.md @@ -445,7 +445,7 @@ const workspace = new Workspace({ }); ``` -The API key defaults to the `TENKI_API_KEY` environment variable when omitted. Ordinary workspace API keys infer their workspace scope server-side, so omit `workspaceId` for normal usage. If you use trusted service credentials that can access multiple workspaces, pass `workspaceId` to select one explicitly: +The API key defaults to the `TENKI_API_KEY` (or `TENKI_AUTH_TOKEN`) environment variable when omitted. Ordinary workspace API keys infer their workspace scope server-side, so omit `workspaceId` for normal usage. If you use trusted service credentials that can access multiple workspaces, pass `workspaceId` to select one explicitly: ```ts const sandbox = new TenkiSandbox({