js-sdk/tests/integration/capsule-features.integration.test.ts

import { describe, expect, it } from "vitest";

import { Capsule, type CapsuleCreateOptions } from "../../src/capsule.js";
import type { CommandStreamEvent } from "../../src/commands.js";
import { DEFAULT_BASE_URL } from "../../src/config.js";
import type { PtyEvent } from "../../src/pty.js";

const baseUrl = process.env.WRENN_BASE_URL ?? DEFAULT_BASE_URL;
const apiKey = process.env.WRENN_API_KEY;
const template = process.env.WRENN_TEST_TEMPLATE ?? "minimal";
const waitTimeoutMs = Number(process.env.WRENN_TEST_WAIT_TIMEOUT_MS ?? 120_000);
const testTimeoutMs = waitTimeoutMs + 45_000;
const describeWithApiKey = apiKey ? describe : describe.skip;

const clientOpts: CapsuleCreateOptions = { baseUrl };
if (apiKey) clientOpts.apiKey = apiKey;

async function withLiveCapsule(
	fn: (capsule: Capsule) => Promise<void>,
): Promise<void> {
	let capsule: Capsule | undefined;
	try {
		capsule = await Capsule.create(template, {
			...clientOpts,
			timeout_sec: 120,
		});
		await capsule.waitForReady({
			intervalMs: 2_000,
			timeoutMs: waitTimeoutMs,
		});
		await fn(capsule);
	} finally {
		if (capsule) {
			await capsule.destroy().catch(() => undefined);
		}
	}
}

async function collectCommandEvents(
	events: AsyncGenerator<CommandStreamEvent>,
): Promise<CommandStreamEvent[]> {
	const collected: CommandStreamEvent[] = [];
	for await (const event of events) {
		collected.push(event);
		if (event.type === "exit" || event.type === "error") break;
	}
	return collected;
}

async function nextPtyEvent(
	events: AsyncIterableIterator<PtyEvent>,
): Promise<PtyEvent> {
	const timeout = new Promise<never>((_, reject) => {
		setTimeout(
			() => reject(new Error("Timed out waiting for PTY event")),
			15_000,
		);
	});
	const result = await Promise.race([events.next(), timeout]);
	if (result.done) throw new Error("PTY closed before next event");
	return result.value;
}

describeWithApiKey("Phase 4 live integration", () => {
	it(
		"executes foreground, streaming, and background commands",
		async () => {
			await withLiveCapsule(async (capsule) => {
				const result = await capsule.commands.exec("printf", {
					args: ["phase4-command"],
					timeoutSec: 10,
				});
				expect(result.exit_code).toBe(0);
				expect(result.stdout).toContain("phase4-command");

				const streamEvents = await collectCommandEvents(
					capsule.commands.stream("printf", { args: ["phase4-stream"] }),
				);
				expect(streamEvents.some((event) => event.type === "stdout")).toBe(
					true,
				);
				expect(streamEvents.at(-1)).toMatchObject({ type: "exit" });

				const tag = `phase4-${Date.now()}`;
				const process = await capsule.commands.start("sleep", {
					args: ["60"],
					tag,
				});
				expect(process.tag).toBe(tag);

				const processes = await capsule.commands.list();
				expect(processes.processes?.some((entry) => entry.tag === tag)).toBe(
					true,
				);

				await expect(
					capsule.commands.kill(tag, "SIGTERM"),
				).resolves.toBeUndefined();
			});
		},
		testTimeoutMs,
	);

	it(
		"performs file read/write/list/remove and streaming transfers",
		async () => {
			await withLiveCapsule(async (capsule) => {
				const dir = `/tmp/wrenn-phase4-${Date.now()}`;
				const file = `${dir}/hello.txt`;
				const streamFile = `${dir}/stream.txt`;

				await capsule.files.mkdir(dir);
				await capsule.files.write(file, "phase4-file");

				const content = await capsule.files.read(file);
				expect(content.toString()).toBe("phase4-file");

				const listing = await capsule.files.list(dir, { depth: 1 });
				expect(
					listing.entries?.some((entry) => entry.name === "hello.txt"),
				).toBe(true);

				await capsule.files.uploadStream(
					streamFile,
					Buffer.from("phase4-stream-file"),
				);
				const chunks: Buffer[] = [];
				for await (const chunk of capsule.files.downloadStream(streamFile)) {
					chunks.push(chunk);
				}
				expect(Buffer.concat(chunks).toString()).toBe("phase4-stream-file");

				await expect(capsule.files.remove(file)).resolves.toBeUndefined();
				await expect(capsule.files.remove(streamFile)).resolves.toBeUndefined();
				await expect(capsule.files.remove(dir)).resolves.toBeUndefined();
			});
		},
		testTimeoutMs,
	);

	it(
		"starts and controls an interactive PTY session",
		async () => {
			await withLiveCapsule(async (capsule) => {
				const session = await capsule.pty.start({
					cmd: "/bin/sh",
					cols: 80,
					rows: 24,
				});

				const started = await nextPtyEvent(session.events);
				expect(started.type).toBe("started");

				session.resize(100, 30);
				session.input("printf phase4-pty\\n\nexit\n");

				const events: PtyEvent[] = [];
				for (let i = 0; i < 10; i += 1) {
					const event = await nextPtyEvent(session.events);
					events.push(event);
					if (event.type === "exit" || event.type === "error") break;
				}

				const output = events
					.filter(
						(event) =>
							event.type === "output" && typeof event.data === "string",
					)
					.map((event) =>
						Buffer.from(event.data as string, "base64").toString(),
					)
					.join("");

				expect(output).toContain("phase4-pty");
				expect(events.at(-1)).toMatchObject({ type: "exit" });
				session.close();
			});
		},
		testTimeoutMs,
	);
});