js-sdk/tests/integration/higher-level-abstractions.integration.test.ts

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

import { Capsule, type CapsuleCreateOptions } from "../../src/capsule.js";
import { CodeInterpreter } from "../../src/code-interpreter/index.js";
import { DEFAULT_BASE_URL } from "../../src/config.js";
import { Git } from "../../src/git/index.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 codeTemplate = process.env.WRENN_TEST_CODE_TEMPLATE ?? "jupyter";
const codeCapsuleId = process.env.WRENN_TEST_CODE_CAPSULE_ID;
const waitTimeoutMs = Number(process.env.WRENN_TEST_WAIT_TIMEOUT_MS ?? 120_000);
const testTimeoutMs = waitTimeoutMs + 60_000;
const describeWithApiKey = apiKey ? describe : describe.skip;

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

function delay(ms: number): Promise<void> {
	return new Promise((resolve) => setTimeout(resolve, ms));
}

function isRetryableServerError(error: unknown): boolean {
	return (
		typeof error === "object" &&
		error !== null &&
		"statusCode" in error &&
		typeof error.statusCode === "number" &&
		error.statusCode >= 500
	);
}

async function expectCommandOk(
	promise: Promise<{ exit_code?: number; stderr?: string }>,
): Promise<void> {
	const result = await promise;
	expect(result.exit_code, result.stderr).toBe(0);
}

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 createCodeInterpreterWithRetry(): Promise<CodeInterpreter> {
	let lastError: unknown;
	for (let attempt = 1; attempt <= 3; attempt += 1) {
		try {
			return await CodeInterpreter.create({
				...clientOpts,
				template: codeTemplate,
				timeout_sec: 120,
			});
		} catch (error) {
			lastError = error;
			if (!isRetryableServerError(error) || attempt === 3) break;
			await delay(2_000);
		}
	}
	throw lastError;
}

async function getCodeInterpreter(): Promise<{
	interpreter: CodeInterpreter;
	shouldDestroy: boolean;
}> {
	if (!codeCapsuleId) {
		return {
			interpreter: await createCodeInterpreterWithRetry(),
			shouldDestroy: true,
		};
	}

	const interpreter = CodeInterpreter.connect(codeCapsuleId, clientOpts);
	const info = await interpreter.capsule.getInfo();
	if (info.status === "paused") {
		await interpreter.capsule.resume({
			intervalMs: 2_000,
			timeoutMs: waitTimeoutMs,
			wait: true,
		});
	} else {
		await interpreter.capsule.waitForReady({
			intervalMs: 2_000,
			timeoutMs: waitTimeoutMs,
		});
	}

	return { interpreter, shouldDestroy: false };
}

describeWithApiKey("Higher-level abstractions live integration", () => {
	it(
		"runs git workflows inside a live capsule",
		async () => {
			await withLiveCapsule(async (capsule) => {
				const root = `/tmp/wrenn-git-${Date.now()}`;
				const origin = `${root}/origin.git`;
				const work = `${root}/work`;

				await expectCommandOk(
					capsule.commands.exec("mkdir", { args: ["-p", root] }),
				);
				await expectCommandOk(
					capsule.commands.exec("git", { args: ["init", "--bare", origin] }),
				);

				await expectCommandOk(capsule.git.clone(origin, { path: work }));
				const git = new Git(capsule, { cwd: work });
				const pwd = await capsule.commands.exec("pwd", { cwd: work });
				expect(pwd.stdout?.trim(), pwd.stderr).toBe(work);

				await expectCommandOk(
					capsule.commands.exec("git", {
						args: ["config", "user.email", "integration@example.com"],
						cwd: work,
					}),
				);
				await expectCommandOk(
					capsule.commands.exec("git", {
						args: ["config", "user.name", "Integration Test"],
						cwd: work,
					}),
				);

				await expectCommandOk(
					capsule.commands.exec("sh", {
						args: ["-lc", "printf 'phase5-git\\n' > README.md"],
						cwd: work,
					}),
				);
				const dirty = await git.status();
				expect(dirty.exit_code, dirty.stderr).toBe(0);
				expect(dirty.stdout, dirty.stderr).toContain("README.md");

				await expectCommandOk(git.add("README.md"));
				const commit = await git.commit("initial commit");
				expect(commit.exit_code, commit.stderr).toBe(0);

				const log = await git.log({ maxCount: 1 });
				expect(log.exit_code, log.stderr).toBe(0);
				expect(log.stdout, log.stderr).toContain("initial commit");

				await expectCommandOk(git.branch());
				await expectCommandOk(git.checkout("feature", { create: true }));
				const branch = await git.branch();
				expect(branch.exit_code, branch.stderr).toBe(0);
				expect(branch.stdout?.trim(), branch.stderr).toBe("feature");

				await expectCommandOk(
					git.push({ branch: "feature", remote: "origin" }),
				);
				await expectCommandOk(
					git.pull({ branch: "feature", remote: "origin" }),
				);
			});
		},
		testTimeoutMs,
	);

	it(
		"executes a code interpreter cell in a live capsule",
		async () => {
			let interpreter: CodeInterpreter | undefined;
			let shouldDestroy = false;
			try {
				const codeInterpreter = await getCodeInterpreter();
				interpreter = codeInterpreter.interpreter;
				shouldDestroy = codeInterpreter.shouldDestroy;

				const result = await interpreter.notebook.execCell("print(6 * 7)", {
					timeoutSec: 15,
				});

				expect(result.exitCode).toBe(0);
				expect(result.stdout.trim()).toBe("42");
				expect(result.stderr).toBe("");
			} finally {
				if (interpreter && shouldDestroy) {
					await interpreter.capsule.destroy().catch(() => undefined);
				}
			}
		},
		testTimeoutMs,
	);

	it(
		"destroys an owned code interpreter capsule when async disposed",
		async () => {
			let capsuleId: string | undefined;

			try {
				const interpreter = await createCodeInterpreterWithRetry();
				capsuleId = interpreter.capsule.id;

				await interpreter.capsule.waitForReady({
					intervalMs: 2_000,
					timeoutMs: waitTimeoutMs,
				});

				await interpreter[Symbol.asyncDispose]();

				await expect(
					Capsule.connect(capsuleId, clientOpts).getInfo(),
				).resolves.toMatchObject({ id: capsuleId, status: "stopped" });
				capsuleId = undefined;
			} finally {
				if (capsuleId) {
					await Capsule.destroy(capsuleId, clientOpts).catch(() => undefined);
				}
			}
		},
		testTimeoutMs,
	);
});