@kubb/studio 5.3.5 → 5.3.6

package/dist/index.js

@@ -1,5 +1,5 @@
 import { t as __name } from "./rolldown-runtime-CRm0XQPb.js";
-import { createJobId, isCommandMessage, isDisconnectMessage, isStudioPongMessage, isStudioReadyMessage } from "./protocol.js";
+import { generationEventTypes } from "./protocol.js";
 import { createRequire } from "node:module";
 import { promisify, styleText } from "node:util";
 import process$1 from "node:process";
@@ -19,6 +19,7 @@ import { mergeDeep } from "remeda";
 import { tmpdir } from "node:os";
 import { gzip } from "node:zlib";
 import { build } from "tsdown";
+import { RpcTarget, newWebSocketRpcSession } from "capnweb";
 import WebSocket from "ws";
 import { setTimeout as setTimeout$1 } from "node:timers/promises";
 //#region src/constants.ts
@@ -41,6 +42,8 @@ const agentDefaults = {
 	* a slower cadence would make a healthy agent look dead after a single missed ping.
 	*/
 	maxHeartbeatIntervalMs: 6e4,
+	/** How long a heartbeat ping may take before the session is treated as dead. */
+	heartbeatTimeoutMs: 1e4,
 	poolSize: 1
 };
 //#endregion
@@ -618,7 +621,7 @@ const INITIAL_POLL_DELAY_MS = 2e3;
 */
 const MAX_POLL_INTERVAL_MS = 3e4;
 /**
-* Polls `GET /api/jobs/{id}` until the job reaches `success` or `failed`, waiting
+* Polls `GET /api/jobs/{id}` until the job reaches a terminal status, waiting
 * {@link INITIAL_POLL_DELAY_MS} first and doubling up to {@link MAX_POLL_INTERVAL_MS} so a long
 * job stays inside the API key's rate limit.
 *
@@ -637,7 +640,7 @@ async function waitForJob({ studioUrl, token, id, timeoutMs = 6e4 }) {
 				headers: { "x-api-key": token },
 				retry: false
 			});
-			if (job.status === "success" || job.status === "failed") return job;
+			if (job.status === "success" || job.status === "failed" || job.status === "canceled") return job;
 		} catch (error) {
 			const response = error.response;
 			if (response?.status !== 429) throw error;
@@ -673,7 +676,7 @@ async function createAgent({ studioUrl, token, name, machineToken }) {
 }
 //#endregion
 //#region package.json
-var version = "5.3.5";
+var version = "5.3.6";
 //#endregion
 //#region src/hooks.ts
 /**
@@ -684,16 +687,19 @@ var version = "5.3.5";
 * Returns a remover, so a session that runs one generation after another on the same emitter does
 * not stack a listener per run.
 */
-function setupHookListener(hooks, root) {
+function setupHookListener(hooks, root, signal) {
 	return hooks.hook("kubb:hook:start", async (ctx) => {
 		const { id, command, args } = ctx;
 		if (!id) return;
 		const commandWithArgs = args?.length ? `${command} ${args.join(" ")}` : command;
 		try {
-			const proc = x(command, [...args ?? []], { nodeOptions: {
-				cwd: root,
-				detached: true
-			} });
+			const proc = x(command, [...args ?? []], {
+				signal,
+				nodeOptions: {
+					cwd: root,
+					detached: true
+				}
+			});
 			for await (const line of proc) await hooks.callHook("kubb:hook:line", {
 				id,
 				line
@@ -742,8 +748,11 @@ function waitForHookEnd(hooks, hookId) {
 		const handleHookEnd = (ctx) => {
 			if (ctx.id !== hookId) return;
 			hooks.removeHook("kubb:hook:end", handleHookEnd);
-			if (ctx.success) resolve();
-			else reject(ctx.error);
+			if (ctx.success) {
+				resolve();
+				return;
+			}
+			reject(ctx.error);
 		};
 		hooks.hook("kubb:hook:end", handleHookEnd);
 	});
@@ -1282,14 +1291,14 @@ function applyRemove(call, path) {
 	});
 }
 /**
-* Adds a `pluginX(...)` call to a config's plugins array. Refuses when the plugin is already
-* present, or when its import name collides with an unrelated existing import.
+* Adds a `pluginX(...)` call to a config's plugins array. Replaying the same add is a no-op, while
+* an import name collision with an unrelated package remains an error.
 */
 function applyAddPlugin(mod, config, edit) {
 	if (!isKubbPluginSpecifier(edit.plugin)) return { reason: `"${edit.plugin}" is not a @kubb/plugin-* package` };
 	const importName = edit.importName ?? toExportName(edit.plugin);
 	if (!IDENTIFIER.test(importName)) return { reason: `"${importName}" is not a valid import name` };
-	if (pluginCalls(mod, config).some((plugin) => plugin.packageName === edit.plugin)) return { reason: `${edit.plugin} is already in the plugins array` };
+	if (pluginCalls(mod, config).some((plugin) => plugin.packageName === edit.plugin)) return { noop: true };
 	const taken = importedFrom(mod).get(importName);
 	if (taken && taken !== edit.plugin) return { reason: `${importName} is already imported from ${taken}` };
 	const options = edit.options ?? {};
@@ -1410,6 +1419,10 @@ function applyConfigEdits(source, edits) {
 				applied: false,
 				reason: result.reason
 			};
+			if ("noop" in result) return {
+				edit,
+				applied: true
+			};
 			const afterLine = lastImportEndLine(mod);
 			let next = generateCode(mod, { format }).code;
 			if (result.addImport) next = insertImportLine({
@@ -1543,14 +1556,19 @@ function formatGenerationFailure(diagnostics) {
 * can forward progress to connected clients. After a successful build, auto-formatting and
 * linting are applied when configured, followed by any user-defined `hooks.done` commands.
 */
-async function generate({ config, hooks }) {
+async function generate({ config, hooks, signal }) {
+	signal?.throwIfAborted();
 	const hrStart = process$1.hrtime();
 	await hooks.callHook("kubb:generation:start", { config });
 	await hooks.callHook("kubb:info", { message: config.name ? `Setup generation ${config.name}` : "Setup generation" });
-	const kubb = createKubb(config, { hooks });
+	const kubb = createKubb(config, {
+		hooks,
+		signal
+	});
 	await kubb.setup();
 	await hooks.callHook("kubb:info", { message: config.name ? `Build generation ${config.name}` : "Build generation" });
 	const { files, diagnostics, storage } = await kubb.safeBuild();
+	signal?.throwIfAborted();
 	await hooks.callHook("kubb:info", { message: "Load summary" });
 	for (const diagnostic of diagnostics.filter(isProblemErrorDiagnostic)) await hooks.callHook("kubb:error", { error: new Error(diagnostic.plugin ? `${diagnostic.plugin}: ${diagnostic.message}` : diagnostic.message) });
 	const status = Diagnostics.hasError(diagnostics) ? "failed" : "success";
@@ -1585,6 +1603,7 @@ async function generate({ config, hooks }) {
 			await hooks.callHook("kubb:success", { message: `${step.verbing} with ${tool} successfully` });
 		} catch (caughtError) {
 			await hooks.callHook("kubb:error", { error: new Error(command.errorMessage, { cause: caughtError }) });
+			signal?.throwIfAborted();
 		}
 		await hooks.callHook(`kubb:${step.kind}:end`);
 	}
@@ -1740,13 +1759,6 @@ async function createSnapshotPackage(files, packageInfo) {
 * takes over.
 */
 const CONNECT_TIMEOUT_MS = 5e3;
-/**
-* Per-socket event counter. Every data message carries the next value so Studio can restore the
-* agent's emission order even when the relay delivers frames out of order. Keyed by the socket so
-* the count stays monotonic across every generation run on one connection, and is dropped
-* automatically once the socket is collected.
-*/
-const eventSeqCounters = /* @__PURE__ */ new WeakMap();
 const require = createRequire(import.meta.url);
 function relativeStoragePath(root, filePath) {
 	return (isAbsolute(filePath) ? relative(resolve(root), filePath) : filePath).replaceAll("\\", "/");
@@ -1771,19 +1783,17 @@ async function resolvePeerDependencies(names) {
 	}));
 	for (const [index, name] of uniqueNames.entries()) {
 		const version = versions[index];
-		if (version) peerDependencies[name] = version;
-		else missingDependencies.push(name);
+		if (version) {
+			peerDependencies[name] = version;
+			continue;
+		}
+		missingDependencies.push(name);
 	}
 	return {
 		peerDependencies,
 		missingDependencies
 	};
 }
-function nextEventSeq(ws) {
-	const seq = eventSeqCounters.get(ws) ?? 0;
-	eventSeqCounters.set(ws, seq + 1);
-	return seq;
-}
 /**
 * Opens a Studio WebSocket connection and closes it when the initial handshake exceeds the configured timeout.
 */
@@ -1796,41 +1806,15 @@ function createWebsocket(url, options) {
 	ws.once("close", () => clearTimeout(timer));
 	return ws;
 }
-/**
-* Sends a serialized agent message when the Studio socket is ready to accept frames.
-*/
-function sendAgentMessage(ws, message) {
-	try {
-		if (ws.readyState !== WebSocket.OPEN) return;
-		ws.send(JSON.stringify(message));
-	} catch (error) {
-		throw new Error("Failed to send message to Kubb Studio", { cause: error });
-	}
-}
-/**
-* Sends a single `kubb:error` payload to Studio, stamped from the same per-socket counter the event stream
-* uses so Studio can still order it against the generation events around it.
-*/
-function sendErrorMessage(ws, error, jobId) {
-	sendAgentMessage(ws, {
-		type: "agent:data",
-		jobId,
-		payload: {
-			type: "kubb:error",
-			data: [{
-				message: error.message,
-				stack: error.stack
-			}],
-			timestamp: Date.now(),
-			seq: nextEventSeq(ws)
-		}
-	});
-}
-/**
-* Forwards selected Kubb lifecycle events to Studio as data messages for the active session.
-*/
-function setupEventsStream(ws, hooks, jobId, options = {}) {
+/** Forwards selected Kubb lifecycle events to a native Cap'n Web stream. */
+function createGenerationStream(hooks, jobId, options = {}) {
 	const unhooks = [];
+	let root = "";
+	const transform = new TransformStream(void 0, void 0, { highWaterMark: Infinity });
+	const writer = transform.writable.getWriter();
+	let writes = Promise.resolve();
+	let closed = false;
+	let streamError;
 	/**
 	* Registers a listener and keeps its remover, so one generation's listeners come off the session
 	* emitter again when that generation ends.
@@ -1838,98 +1822,73 @@ function setupEventsStream(ws, hooks, jobId, options = {}) {
 	function on(name, handler) {
 		unhooks.push(hooks.hook(name, handler));
 	}
-	function sendDataMessage(payload) {
-		sendAgentMessage(ws, {
-			type: "agent:data",
+	function emitEvent(type, data) {
+		const event = {
 			jobId,
-			payload: {
-				...payload,
-				timestamp: Date.now(),
-				seq: nextEventSeq(ws)
-			}
+			type,
+			data,
+			version: 1,
+			timestamp: Date.now()
+		};
+		writes = writes.then(() => writer.write(event)).catch((error) => {
+			streamError = error;
 		});
 	}
 	on("kubb:plugin:start", (ctx) => {
-		sendDataMessage({
-			type: "kubb:plugin:start",
-			data: [{ plugin: ctx.plugin }]
-		});
+		emitEvent("kubb:plugin:start", [{ plugin: { name: ctx.plugin.name } }]);
 	});
 	on("kubb:plugin:end", (ctx) => {
-		sendDataMessage({
-			type: "kubb:plugin:end",
-			data: [{
-				plugin: ctx.plugin,
-				duration: ctx.duration,
-				success: ctx.success
-			}]
-		});
+		emitEvent("kubb:plugin:end", [{
+			plugin: { name: ctx.plugin.name },
+			duration: ctx.duration,
+			success: ctx.success
+		}]);
 	});
 	on("kubb:build:start", ({ config, adapter }) => {
-		sendDataMessage({
-			type: "kubb:build:start",
-			data: [{
-				config: { name: config.name },
-				adapter: { name: adapter.name }
-			}]
-		});
+		root = config.root;
+		emitEvent("kubb:build:start", [{
+			config: { name: config.name },
+			adapter: { name: adapter.name }
+		}]);
 	});
 	on("kubb:build:end", ({ files, config, outputDir }) => {
-		sendDataMessage({
-			type: "kubb:build:end",
-			data: [{
-				files: files.map((file) => ({
-					path: relativeStoragePath(config.root, file.path),
-					name: file.name
-				})),
-				outputDir
-			}]
-		});
+		emitEvent("kubb:build:end", [{
+			files: files.map((file) => ({
+				path: relativeStoragePath(config.root, file.path),
+				name: file.name
+			})),
+			outputDir
+		}]);
 	});
 	on("kubb:files:processing:start", ({ files }) => {
-		sendDataMessage({
-			type: "kubb:files:processing:start",
-			data: [{ total: files.length }]
-		});
+		emitEvent("kubb:files:processing:start", [{ total: files.length }]);
 	});
 	on("kubb:files:processing:update", ({ files }) => {
-		sendDataMessage({
-			type: "kubb:files:processing:update",
-			data: [{ files: files.map(({ file, processed, total, percentage }) => ({
-				file: file.path,
-				processed,
-				total,
-				percentage
-			})) }]
-		});
+		emitEvent("kubb:files:processing:update", [{ files: files.map(({ file, processed, total, percentage }) => ({
+			file: relativeStoragePath(root, file.path),
+			processed,
+			total,
+			percentage
+		})) }]);
 	});
 	on("kubb:files:processing:end", ({ files }) => {
-		sendDataMessage({
-			type: "kubb:files:processing:end",
-			data: [{ total: files.length }]
-		});
+		emitEvent("kubb:files:processing:end", [{ total: files.length }]);
 	});
 	for (const type of [
 		"kubb:info",
 		"kubb:success",
 		"kubb:warn"
 	]) on(type, ({ message, info }) => {
-		sendDataMessage({
-			type,
-			data: [{
-				message,
-				info
-			}]
-		});
+		emitEvent(type, [{
+			message,
+			info
+		}]);
 	});
 	on("kubb:generation:start", ({ config }) => {
-		sendDataMessage({
-			type: "kubb:generation:start",
-			data: [{
-				name: config.name,
-				plugins: config.plugins.length
-			}]
-		});
+		emitEvent("kubb:generation:start", [{
+			name: config.name,
+			plugins: config.plugins.length
+		}]);
 	});
 	on("kubb:generation:end", async ({ config, storage, diagnostics = [], status, hrStart, filesCreated }) => {
 		const { peerDependencies, missingDependencies } = await resolvePeerDependencies(config.plugins.map(({ name }) => name));
@@ -1942,29 +1901,32 @@ function setupEventsStream(ws, hooks, jobId, options = {}) {
 			peerDependencies,
 			missingDependencies
 		});
-		sendDataMessage({
-			type: "kubb:generation:end",
-			data: []
-		});
+		emitEvent("kubb:generation:end", []);
 		if (!hrStart) return;
-		sendDataMessage({
-			type: "kubb:generation:summary",
-			data: [{
-				duration: Math.round(getElapsedMs(hrStart)),
-				fileCount: filesCreated ?? 0,
-				failedPlugins: Diagnostics.failedPlugins(diagnostics).length,
-				status: status ?? "success"
-			}]
-		});
+		emitEvent("kubb:generation:summary", [{
+			duration: Math.round(getElapsedMs(hrStart)),
+			fileCount: filesCreated ?? 0,
+			failedPlugins: Diagnostics.failedPlugins(diagnostics).length,
+			status: status ?? "success"
+		}]);
 	});
 	on("kubb:error", ({ error }) => {
-		sendDataMessage({
-			type: "kubb:error",
-			data: [{
-				message: error.message,
-				stack: error.stack
-			}]
-		});
+		emitEvent("kubb:error", [{
+			message: error.message,
+			stack: error.stack
+		}]);
+	});
+	on("kubb:diagnostic", ({ diagnostic }) => {
+		const cause = "cause" in diagnostic ? diagnostic.cause : void 0;
+		emitEvent("kubb:diagnostic", [{
+			code: diagnostic.code,
+			message: diagnostic.message,
+			severity: diagnostic.severity,
+			location: "location" in diagnostic ? diagnostic.location : void 0,
+			help: "help" in diagnostic ? diagnostic.help : void 0,
+			plugin: "plugin" in diagnostic ? diagnostic.plugin : void 0,
+			stack: cause?.stack
+		}]);
 	});
 	for (const type of [
 		"kubb:lifecycle:start",
@@ -1976,62 +1938,144 @@ function setupEventsStream(ws, hooks, jobId, options = {}) {
 		"kubb:hooks:start",
 		"kubb:hooks:end"
 	]) on(type, () => {
-		sendDataMessage({
-			type,
-			data: []
-		});
+		emitEvent(type, []);
 	});
 	on("kubb:hook:start", ({ id, command, args }) => {
-		sendDataMessage({
-			type: "kubb:hook:start",
-			data: [{
-				id,
-				command,
-				args: args ? [...args] : void 0
-			}]
-		});
+		emitEvent("kubb:hook:start", [{
+			id,
+			command,
+			args: args ? [...args] : void 0
+		}]);
 	});
 	on("kubb:hook:line", ({ id, line }) => {
-		sendDataMessage({
-			type: "kubb:hook:line",
-			data: [{
-				id,
-				line
-			}]
-		});
+		emitEvent("kubb:hook:line", [{
+			id,
+			line
+		}]);
 	});
 	on("kubb:hook:end", ({ id, command, args, success, error }) => {
-		sendDataMessage({
-			type: "kubb:hook:end",
-			data: [{
-				id,
-				command,
-				args: args ? [...args] : void 0,
-				success,
-				error: error ? {
-					message: error.message,
-					stack: error.stack
-				} : void 0
-			}]
-		});
+		emitEvent("kubb:hook:end", [{
+			id,
+			command,
+			args: args ? [...args] : void 0,
+			success,
+			error: error ? {
+				message: error.message,
+				stack: error.stack
+			} : void 0
+		}]);
 	});
-	return () => {
+	/**
+	* Takes this generation's listeners off the session emitter. Safe to call twice.
+	*/
+	function detach() {
 		for (const unhook of unhooks) unhook();
+		unhooks.length = 0;
+	}
+	async function close() {
+		if (closed) return;
+		closed = true;
+		detach();
+		await writes;
+		if (streamError) return;
+		await writer.close().catch(() => void 0);
+	}
+	function fail(error) {
+		detach();
+		if (closed) return;
+		closed = true;
+		writer.abort(error).catch(() => void 0);
+	}
+	return {
+		stream: transform.readable,
+		close,
+		dispose: () => fail(),
+		fail
 	};
 }
 //#endregion
-//#region src/StudioSession.ts
+//#region src/rpc.ts
 /**
-* How many files are read from storage at once when serving `studio:files` or packing a
-* `studio:snapshot`.
+* The only methods Studio may call on an agent. A `StudioSession` carries far more than
+* {@link AgentApi}, so it is wrapped rather than exposed: what Cap'n Web can reach is exactly what
+* this class re-declares.
 */
-const FILE_READ_CONCURRENCY = 50;
+var AgentRpcTarget = class extends RpcTarget {
+	api;
+	constructor(api) {
+		super();
+		this.api = api;
+	}
+	connect() {
+		return this.api.connect();
+	}
+	startGeneration(input) {
+		return this.api.startGeneration(input);
+	}
+	saveConfig(input) {
+		return this.api.saveConfig(input);
+	}
+	publishSnapshot(input) {
+		return this.api.publishSnapshot(input);
+	}
+	readFiles(input) {
+		return this.api.readFiles(input);
+	}
+};
+/**
+* Opens an authenticated Cap'n Web session to Studio over a WebSocket. Rejects an unencrypted URL
+* before opening the socket, so a bearer token never reaches a plaintext host.
+*
+* @example
+* ```ts
+* const rpc = await connectWebSocketRpc({ url: 'wss://studio.kubb.dev/s/1', token, local: session })
+* await rpc.studio.ping()
+* ```
+*/
+const connectWebSocketRpc = async ({ url, token, local }) => {
+	const { protocol, hostname, host } = new URL(url);
+	if (protocol !== "wss:" && !(protocol === "ws:" && (hostname === "localhost" || hostname === "127.0.0.1" || hostname === "[::1]"))) throw new Error(`Refusing unencrypted WebSocket to ${host}`);
+	const socket = createWebsocket(url, { headers: { Authorization: `Bearer ${token}` } });
+	const closed = new Promise((resolve) => socket.once("close", resolve));
+	const studio = newWebSocketRpcSession(socket, new AgentRpcTarget(local));
+	studio.onRpcBroken(() => socket.close());
+	return {
+		studio,
+		closed,
+		close: () => studio[Symbol.dispose]()
+	};
+};
+//#endregion
+//#region src/StudioSession.ts
 /**
-* How long the `agent:connect` handshake may go unacknowledged before `studio:ready` is given up
-* on for this open. A Studio that predates the ack never sends one, so this only ever produces a
-* warning, not a reconnect.
+* How many files are read from storage at once when serving `readFiles` or packing a snapshot.
 */
-const READY_TIMEOUT_MS = 1e4;
+const FILE_READ_CONCURRENCY = 50;
+var GenerationRunTarget = class extends RpcTarget {
+	generationStream;
+	generationResult;
+	cancelGeneration;
+	stopGeneration;
+	constructor(generationStream, generationResult, cancelGeneration, stopGeneration) {
+		super();
+		this.generationStream = generationStream;
+		this.generationResult = generationResult;
+		this.cancelGeneration = cancelGeneration;
+		this.stopGeneration = stopGeneration;
+	}
+	async events() {
+		return this.generationStream;
+	}
+	result() {
+		return this.generationResult;
+	}
+	cancel() {
+		return this.cancelGeneration();
+	}
+	[Symbol.dispose]() {
+		this.stopGeneration();
+	}
+};
 /**
 * Fills in a host's options: the hosted Studio URL, the current working directory, and every
 * permission off unless granted. Idempotent, so a reconnect can pass an already-resolved bag
@@ -2071,7 +2115,7 @@ function reconnect(options) {
 	const timer = setTimeout(() => {
 		signal?.removeEventListener("abort", cancel);
 		if (signal?.aborted) return;
-		new StudioSession(options).connect().catch((error) => {
+		new StudioSession(options).start().catch((error) => {
 			if (logLevel$1 !== void 0 && logLevel$1 > logLevel.silent) console.error(styleText("red", `Reconnect attempt to Kubb Studio failed: ${getErrorMessage(error)}`));
 			if (error instanceof InvalidAgentTokenError) {
 				onTokenRejected?.(error);
@@ -2083,8 +2127,8 @@ function reconnect(options) {
 	signal?.addEventListener("abort", cancel, { once: true });
 }
 /**
-* One WebSocket session with Studio: opening it, keeping it alive, and running the commands it
-* sends. `createClient` opens one per pool slot and is the only caller.
+* One agent-to-Studio RPC transport: opening it, keeping it alive, and serving remote methods.
+* `createClient` opens one per pool slot and is the only caller.
 */
 var StudioSession = class {
 	#options;
@@ -2099,17 +2143,20 @@ var StudioSession = class {
 	* Before it resolves there is nothing to disconnect and no sandbox flag to read.
 	*/
 	#session;
-	#ws;
+	#rpc;
 	#studioVersion;
-	#activeJobId = null;
 	#disposed = false;
 	#isGenerating = false;
 	#heartbeatTimer;
-	#readyTimer;
-	#lastPongAt = Date.now();
 	#lastGeneration;
+	/**
+	* Resolves when Studio calls {@link StudioSession.connect}. `studio:ready` waits on this so the
+	* host does not queue jobs before the agent session is registered.
+	*/
+	#connectAck = Promise.withResolvers();
 	constructor(options) {
 		this.#options = applyStudioDefaults(options);
+		this.#connectAck.promise.catch(() => {});
 	}
 	/**
 	* A sandbox agent runs on Studio's own infrastructure, so it has no user project to touch.
@@ -2136,7 +2183,7 @@ var StudioSession = class {
 	get #canRead() {
 		return this.#isSandbox || this.#options.permissions.allowRead;
 	}
-	async connect() {
+	async start() {
 		const { token, studioUrl, signal, heartbeatInterval, installLogger } = this.#options;
 		await installLogger?.(this.#hooks);
 		try {
@@ -2147,48 +2194,65 @@ var StudioSession = class {
 			});
 			this.#session = session;
 			this.#studioVersion = session.version;
-			const ws = createWebsocket(session.wsUrl, { headers: { Authorization: `Bearer ${token}` } });
-			this.#ws = ws;
-			this.#listen(ws, "open", this.#onOpen);
-			this.#listen(ws, "close", this.#onClose);
-			this.#listen(ws, "error", this.#onError);
-			this.#listen(ws, "message", this.#onMessage);
+			const rpc = await (this.#options.connector ?? connectWebSocketRpc)({
+				url: session.url,
+				token,
+				local: this
+			});
+			this.#rpc = rpc;
+			rpc.closed.then(this.#onClose);
 			signal?.addEventListener("abort", this.#onAbort, { once: true });
 			this.#unhooks.push(() => signal?.removeEventListener("abort", this.#onAbort));
-			this.#heartbeatTimer = setInterval(() => this.#sendHeartbeat(), heartbeatInterval);
-			this.#unhooks.push(this.#hooks.hook("kubb:error", ({ error }) => sendErrorMessage(ws, error, this.#activeJobId ?? "connection")));
+			this.#scheduleHeartbeat(heartbeatInterval);
+			await this.#hooks.callHook("studio:connected", {
+				url: studioUrl,
+				versions: {
+					studio: this.#studioVersion,
+					kubb: version,
+					agent: this.#options.version
+				}
+			});
+			await this.#connectAck.promise;
+			await this.#hooks.callHook("studio:ready", {});
 		} catch (error) {
+			this.#disposed = true;
+			this.dispose();
 			await this.#hooks.callHook("studio:error", { error: toError(error) });
 			if (error instanceof InvalidAgentTokenError) throw error;
 			reconnect(this.#options);
 		}
 	}
-	/**
-	* Adds a socket listener and tracks its remover, so `dispose` detaches every listener at once.
-	*/
-	#listen(ws, event, listener) {
-		ws.addEventListener(event, listener);
-		this.#unhooks.push(() => ws.removeEventListener(event, listener));
-	}
 	#warn(message) {
 		return this.#hooks.callHook("studio:warn", { message });
 	}
 	/**
-	* Forwards a failure to Studio over the connection emitter, which `connect` wired to this
-	* socket. Swallows a listener's own failure, since this is already the error path.
+	* Declines a request: logs why locally, then tells Studio. The two wordings differ on purpose,
+	* since the log names the request that was ignored and the error names what the caller can do.
 	*/
-	#emitError(error) {
-		return Promise.resolve(this.#hooks.callHook("kubb:error", { error })).catch(() => {});
-	}
-	#sendHeartbeat() {
-		if (Date.now() - this.#lastPongAt > this.#options.heartbeatInterval * 2) {
-			this.#warn("No reply from Kubb Studio, terminating the stale connection");
-			clearInterval(this.#heartbeatTimer);
-			this.#heartbeatTimer = void 0;
-			this.#ws?.terminate();
-			return;
-		}
-		if (this.#ws) sendAgentMessage(this.#ws, { type: "agent:ping" });
+	async #refuse(reason, message) {
+		await this.#warn(reason);
+		throw new Error(message);
+	}
+	#scheduleHeartbeat(interval) {
+		const rpc = this.#rpc;
+		if (!rpc) return;
+		this.#heartbeatTimer = setTimeout(async () => {
+			try {
+				await this.#ping(rpc);
+			} catch {
+				if (this.#rpc === rpc) rpc.close();
+				return;
+			}
+			if (this.#rpc === rpc) this.#scheduleHeartbeat(interval);
+		}, interval);
+	}
+	/**
+	* Races `studio.ping()` against a deadline, so a half-open socket can't hang it forever.
+	* */
+	#ping(rpc) {
+		const { promise: timedOut, reject: onTimeout } = Promise.withResolvers();
+		const timer = setTimeout(() => onTimeout(/* @__PURE__ */ new Error("Heartbeat ping timed out")), agentDefaults.heartbeatTimeoutMs);
+		return Promise.race([rpc.studio.ping(), timedOut]).finally(() => clearTimeout(timer));
 	}
 	/**
 	* Reads `kubb.config.ts` and reports which plugin options Studio may edit.
@@ -2205,90 +2269,48 @@ var StudioSession = class {
 			return;
 		}
 	}
-	async #sendConnectedPayload() {
+	async connect() {
 		const { configPath, root, version: version$1, loadConfig, permissions } = this.#options;
-		if (!this.#ws) return;
-		const config = await loadConfig();
-		sendAgentMessage(this.#ws, {
-			type: "agent:connect",
-			payload: {
-				versions: {
-					kubb: version,
-					agent: version$1
-				},
-				root,
-				config: {
-					path: configPath,
-					file: await this.#readConfigFileView(),
-					plugins: config.plugins.map((plugin) => ({
-						name: toPackageName(plugin.name),
-						options: plugin.options ?? {}
-					}))
-				},
-				permissions: {
-					...permissions,
-					allowWrite: this.#canWrite,
-					allowInput: this.#canUseInput,
-					allowConfigEdit: this.#canEditConfig,
-					allowRead: this.#canRead
-				}
-			}
-		});
-	}
-	async #handleOpen() {
-		this.#lastPongAt = Date.now();
-		await this.#hooks.callHook("studio:connected", {
-			url: this.#options.studioUrl,
+		const [config, file] = await Promise.all([loadConfig(), this.#readConfigFileView()]);
+		const payload = {
 			versions: {
-				studio: this.#studioVersion,
 				kubb: version,
-				agent: this.#options.version
+				agent: version$1
+			},
+			root,
+			config: {
+				path: configPath,
+				file,
+				plugins: config.plugins.map((plugin) => ({
+					name: toPackageName(plugin.name),
+					options: plugin.options ?? {}
+				}))
+			},
+			permissions: {
+				...permissions,
+				allowWrite: this.#canWrite,
+				allowInput: this.#canUseInput,
+				allowConfigEdit: this.#canEditConfig,
+				allowRead: this.#canRead
 			}
-		});
-		try {
-			await this.#sendConnectedPayload();
-			this.#waitForReady();
-		} catch (error) {
-			await this.#warn(`Failed to send the connect payload: ${getErrorMessage(error)}`);
-		}
+		};
+		this.#connectAck.resolve();
+		return payload;
 	}
-	/**
-	* Arms the timeout for Studio's `studio:ready` acknowledgement. Re-armed on every open, since a
-	* command sent while the agent is not attached is lost either way (see `#sendConnectedPayload`),
-	* so each reconnect needs its own fresh wait.
-	*/
-	#waitForReady() {
-		if (this.#disposed) return;
-		clearTimeout(this.#readyTimer);
-		this.#readyTimer = setTimeout(() => {
-			this.#readyTimer = void 0;
-			Promise.resolve(this.#warn(`Kubb Studio did not confirm the connection was ready within ${READY_TIMEOUT_MS}ms`)).catch(() => {});
-		}, READY_TIMEOUT_MS);
-	}
-	#onOpen = () => void this.#handleOpen().catch(() => {});
-	#onAbort = () => void this.#end({
-		reason: "shutdown",
-		retry: false
-	});
+	#onAbort = () => void this.#end({ retry: false });
 	#onClose = () => void this.#end({ retry: true });
-	#onError = () => {
-		this.#hooks.callHook("studio:error", { error: /* @__PURE__ */ new Error("Failed to connect to Kubb Studio") });
-		this.#onClose();
-	};
 	/**
 	* Drops the socket and detaches every listener and timer this session added. Idempotent, and
 	* safe before `connect` opened anything.
 	*
 	* @internal
 	*/
-	dispose(reason = "cleanup") {
-		clearInterval(this.#heartbeatTimer);
+	dispose() {
+		clearTimeout(this.#heartbeatTimer);
 		this.#heartbeatTimer = void 0;
-		clearTimeout(this.#readyTimer);
-		this.#readyTimer = void 0;
-		try {
-			this.#ws?.close(1e3, reason);
-		} catch {}
+		this.#rpc?.close();
+		this.#rpc = void 0;
+		this.#connectAck.reject(/* @__PURE__ */ new Error("Session ended before Studio called connect()"));
 		for (const unhook of this.#unhooks) unhook();
 		this.#unhooks.length = 0;
 	}
@@ -2296,15 +2318,12 @@ var StudioSession = class {
 	* Ends the session: tells Studio it is over, drops the socket, and optionally reconnects.
 	* `#disposed` keeps the close event from running this twice, and a shutdown from reconnecting.
 	*/
-	async #end({ reason, retry }) {
+	async #end({ retry }) {
 		const { studioUrl, token, logLevel } = this.#options;
 		if (this.#disposed) return;
 		this.#disposed = true;
-		if (reason === "shutdown" && this.#ws) sendAgentMessage(this.#ws, {
-			type: "agent:disconnect",
-			reason: "shutdown"
-		});
-		this.dispose(reason);
+		this.dispose();
+		await this.#hooks.callHook("studio:disconnected", { reason: retry ? "connection closed" : "shutdown" });
 		if (this.#session) await disconnect({
 			sessionId: this.#session.sessionId,
 			studioUrl,
@@ -2314,81 +2333,35 @@ var StudioSession = class {
 		}).catch(() => {});
 		if (retry) reconnect(this.#options);
 	}
-	#onMessage = async (message) => {
-		try {
-			const data = JSON.parse(message.data);
-			if (isStudioPongMessage(data)) {
-				this.#lastPongAt = Date.now();
-				return;
-			}
-			if (isStudioReadyMessage(data)) {
-				clearTimeout(this.#readyTimer);
-				this.#readyTimer = void 0;
-				await this.#hooks.callHook("studio:ready", {});
-				return;
-			}
-			if (isDisconnectMessage(data)) {
-				await this.#handleDisconnect(data.reason);
-				return;
-			}
-			if (isCommandMessage(data)) {
-				await this.#handleCommand(data);
-				return;
-			}
-			await this.#warn(`Ignored an unknown message from Kubb Studio: ${data.type}`);
-		} catch (error) {
-			await this.#hooks.callHook("studio:error", { error: toError(error) });
-			await this.#emitError(toError(error));
-		}
-	};
-	/**
-	* Studio ended the session itself. A revoked one stays ended, an expired one gets a fresh
-	* session, and anything else is left to the socket's own close event.
-	*/
-	async #handleDisconnect(reason) {
-		await this.#hooks.callHook("studio:disconnected", { reason });
-		if (reason !== "revoked" && reason !== "expired") return;
-		this.#disposed = true;
-		this.dispose(`session_${reason}`);
-		if (reason === "expired") reconnect(this.#options);
-	}
-	async #handleCommand(data) {
-		const ws = this.#ws;
-		if (!ws) return;
-		const command = data.type.slice(7);
-		await this.#hooks.callHook("studio:command:start", { command });
-		switch (data.type) {
-			case "studio:generate":
-				await this.#handleGenerate(ws, data, command);
-				return;
-			case "studio:connect":
-				this.#studioVersion = data.version ?? this.#studioVersion;
-				await this.#sendConnectedPayload();
-				await this.#hooks.callHook("studio:command:end", { command });
-				return;
-			case "studio:save":
-				await this.#handleSave(ws, data, command);
-				return;
-			case "studio:snapshot":
-				await this.#handleSnapshot(ws, data, command);
-				return;
-			case "studio:files":
-				await this.#handleFiles(ws, data, command);
-				return;
-		}
+	startGeneration(data) {
+		const generationStream = createGenerationStream(this.#hooks, data.jobId, { onGenerationEnd: (result) => {
+			this.#lastGeneration = result;
+		} });
+		const controller = new AbortController();
+		const result = this.#runGeneration(data, controller).then(async (value) => {
+			await generationStream.close();
+			return value;
+		}).catch((error) => {
+			generationStream.fail(error);
+			throw error;
+		});
+		result.catch(() => {});
+		return new GenerationRunTarget(generationStream.stream, result, async () => {
+			controller.abort(/* @__PURE__ */ new Error("Generation canceled"));
+		}, () => {
+			controller.abort(/* @__PURE__ */ new Error("Generation canceled"));
+			generationStream.dispose();
+		});
 	}
-	async #handleGenerate(ws, data, command) {
-		const { root, loadConfig, permissions, client } = this.#options;
-		if (this.#isGenerating) {
-			await this.#warn("Ignored generate: a generation is already in progress");
-			await this.#emitError(/* @__PURE__ */ new Error("A generation is already in progress, please wait for it to finish"));
-			return;
-		}
+	async #runGeneration(data, controller) {
+		if (this.#isGenerating) return this.#refuse("Ignored generate: a generation is already in progress", "A generation is already in progress, please wait for it to finish");
 		this.#isGenerating = true;
-		this.#activeJobId = data.jobId;
+		const command = "generate";
+		const { root, loadConfig, permissions, client } = this.#options;
 		try {
+			await this.#hooks.callHook("studio:command:start", { command });
 			const config = await loadConfig();
-			const patch = data.payload;
+			const patch = data.config;
 			const plugins = await mergePlugins(config.plugins, patch?.plugins);
 			const adapter = await mergeAdapter(config.adapter, patch?.adapter);
 			const inputOverride = this.#isSandbox ? patch?.input ?? "" : permissions.allowInput && patch?.input || void 0;
@@ -2399,9 +2372,7 @@ var StudioSession = class {
 			}
 			const resolvedPlugins = plugins ?? config.plugins;
 			this.#lastGeneration = void 0;
-			const detach = [setupHookListener(this.#hooks, root), setupEventsStream(ws, this.#hooks, data.jobId, { onGenerationEnd: (result) => {
-				this.#lastGeneration = result;
-			} })];
+			const detach = [setupHookListener(this.#hooks, root, controller.signal)];
 			try {
 				await generate({
 					config: {
@@ -2418,7 +2389,8 @@ var StudioSession = class {
 						plugins: resolvedPlugins,
 						adapter
 					},
-					hooks: this.#hooks
+					hooks: this.#hooks,
+					signal: controller.signal
 				});
 			} finally {
 				for (const remove of detach) remove();
@@ -2427,102 +2399,70 @@ var StudioSession = class {
 				command,
 				info: `${resolvedPlugins.length} plugin${resolvedPlugins.length === 1 ? "" : "s"}, ${this.#canWrite ? "written to disk" : "in memory"}${inputOverride !== void 0 ? ", from a Studio spec" : ""}`
 			});
+			const files = [...this.#lastGeneration?.paths ?? []];
+			return {
+				status: "success",
+				files,
+				fileCount: files.length
+			};
 		} finally {
 			this.#isGenerating = false;
-			this.#activeJobId = null;
 		}
 	}
-	async #handleSave(ws, data, command) {
+	async saveConfig(data) {
+		const command = "saveConfig";
+		await this.#hooks.callHook("studio:command:start", { command });
 		const { configPath, configFile } = this.#options;
 		if (!Array.isArray(data.edits)) {
 			await this.#warn("Ignored save: the message carried no edits");
-			sendAgentMessage(ws, {
-				type: "agent:save",
-				jobId: data.jobId,
-				payload: {
-					outcomes: [],
-					changed: false
-				}
-			});
-			return;
+			return {
+				outcomes: [],
+				changed: false
+			};
 		}
 		const edits = data.edits;
-		const refuse = (reason) => sendAgentMessage(ws, {
-			type: "agent:save",
-			jobId: data.jobId,
-			payload: {
-				outcomes: edits.map((edit) => ({
-					edit,
-					applied: false,
-					reason
-				})),
-				changed: false
-			}
+		const refuse = (reason) => ({
+			outcomes: edits.map((edit) => ({
+				edit,
+				applied: false,
+				reason
+			})),
+			changed: false
 		});
 		if (!this.#canEditConfig) {
 			await this.#warn("Ignored save: editing kubb.config.ts was not granted");
-			refuse("the agent was not granted permission to edit kubb.config.ts");
-			return;
-		}
-		if (this.#isGenerating) {
-			refuse("a generation is in progress");
-			return;
+			return refuse("the agent was not granted permission to edit kubb.config.ts");
 		}
+		if (this.#isGenerating) return refuse("a generation is in progress");
 		try {
 			const { source: patched, outcomes, changed } = applyConfigEdits(await read(configFile), edits);
 			if (changed) await writeFile(configFile, patched, "utf-8");
-			sendAgentMessage(ws, {
-				type: "agent:save",
-				jobId: data.jobId,
-				payload: {
-					outcomes,
-					changed,
-					file: changed ? await this.#readConfigFileView(patched) : void 0
-				}
-			});
 			const applied = outcomes.filter((outcome) => outcome.applied).length;
 			await this.#hooks.callHook("studio:command:end", {
 				command,
 				info: `applied ${applied}/${outcomes.length} edits to ${configPath}`
 			});
+			return {
+				outcomes,
+				changed,
+				file: changed ? await this.#readConfigFileView(patched) : void 0
+			};
 		} catch (error) {
 			await this.#hooks.callHook("studio:error", { error: toError(error) });
-			refuse(getErrorMessage(error));
+			return refuse(getErrorMessage(error));
 		}
 	}
-	async #handleSnapshot(ws, data, command) {
-		const refuse = (message) => sendAgentMessage(ws, {
-			type: "agent:snapshot",
-			jobId: data.jobId,
-			payload: {
-				status: "error",
-				message
-			}
-		});
-		if (this.#isSandbox) {
-			await this.#warn("Ignored snapshot: a sandbox agent has no project to build a package from");
-			refuse("a sandbox agent has no project to build a package from");
-			return;
-		}
-		const { name, version, bundledDependencies, uploadPath } = data.payload;
-		if (!name || !version || !uploadPath) {
-			await this.#warn("Ignored snapshot: the message was missing required fields");
-			refuse("the message was missing required fields");
-			return;
-		}
+	async publishSnapshot(data) {
+		const command = "snapshot";
+		await this.#hooks.callHook("studio:command:start", { command });
+		if (this.#isSandbox) return this.#refuse("Ignored snapshot: a sandbox agent has no project to build a package from", "A sandbox agent has no project to build a package from");
+		const { name, version, bundledDependencies, uploadPath } = data;
+		if (!name || !version || !uploadPath) return this.#refuse("Ignored snapshot: the message was missing required fields", "The request was missing required fields");
 		const generation = this.#lastGeneration;
-		if (!generation) {
-			await this.#warn("Ignored snapshot: no prior generation to pack");
-			refuse("no prior generation exists to pack, run a generation first");
-			return;
-		}
+		if (!generation) return this.#refuse("Ignored snapshot: no prior generation to pack", "No prior generation exists to pack, run a generation first");
 		const bundled = new Set(bundledDependencies ?? []);
 		const missing = generation.missingDependencies.filter((dependency) => !bundled.has(dependency));
-		if (missing.length) {
-			await this.#warn(`Ignored snapshot: missing dependencies: ${missing.join(", ")}`);
-			refuse(`missing dependencies: ${missing.join(", ")}`);
-			return;
-		}
+		if (missing.length) return this.#refuse(`Ignored snapshot: missing dependencies: ${missing.join(", ")}`, `Missing dependencies: ${missing.join(", ")}`);
 		try {
 			const files = {};
 			await inParallel({
@@ -2539,68 +2479,50 @@ var StudioSession = class {
 				peerDependencies: generation.peerDependencies
 			});
 			const { token, studioUrl } = this.#options;
-			const redirect = await fetch(new URL(uploadPath, studioUrl), {
+			const uploadUrl = new URL(uploadPath, studioUrl);
+			if (uploadUrl.origin !== new URL(studioUrl).origin) throw new Error("Snapshot upload path must stay on the Studio origin");
+			const redirect = await fetch(uploadUrl, {
 				method: "PUT",
 				headers: { Authorization: `Bearer ${token}` },
 				redirect: "manual"
 			});
 			const storageUrl = redirect.headers.get("location");
 			if (redirect.status !== 307 || !storageUrl) throw new Error(`Studio did not provide a storage URL (status ${redirect.status})`);
-			const response = await fetch(storageUrl, {
+			const storage = new URL(storageUrl);
+			if (storage.protocol !== "https:" && storage.hostname !== "localhost" && storage.hostname !== "127.0.0.1") throw new Error(`Refusing snapshot upload to ${storage.origin}`);
+			const response = await fetch(storage, {
 				method: "PUT",
-				body: new Uint8Array(bytes)
+				body: new Uint8Array(bytes),
+				redirect: "error"
 			});
 			if (!response.ok) throw new Error(`Snapshot upload failed with status ${response.status}`);
-			sendAgentMessage(ws, {
-				type: "agent:snapshot",
-				jobId: data.jobId,
-				payload: {
-					status: "ok",
-					integrity,
-					peerDependencies: generation.peerDependencies
-				}
-			});
 			await this.#hooks.callHook("studio:command:end", {
 				command,
 				info: `packed ${Object.keys(files).length} file${Object.keys(files).length === 1 ? "" : "s"}`
 			});
+			return {
+				integrity,
+				peerDependencies: generation.peerDependencies
+			};
 		} catch (error) {
 			await this.#hooks.callHook("studio:error", { error: toError(error) });
-			refuse(getErrorMessage(error));
+			throw error;
 		}
 	}
-	async #handleFiles(ws, data, command) {
+	async readFiles(data) {
+		const command = "readFiles";
+		await this.#hooks.callHook("studio:command:start", { command });
 		const { client } = this.#options;
-		const refuse = (message) => sendAgentMessage(ws, {
-			type: "agent:files",
-			jobId: data.jobId,
-			payload: {
-				status: "error",
-				message
-			}
-		});
 		if (!this.#canRead) {
 			await this.#warn("Ignored files: reading generated files was not granted");
-			refuse(`the agent was not granted permission to read generated files; set ${client?.kind === "cli" ? "--allow-read, or answer yes when kubb studio asks," : "KUBB_AGENT_ALLOW_READ=true"} to allow it`);
-			return;
-		}
-		if (!Array.isArray(data.payload?.paths)) {
-			await this.#warn("Ignored files: the message carried no paths");
-			refuse("the message carried no paths");
-			return;
-		}
-		const { paths } = data.payload;
-		if (paths.length > 50) {
-			await this.#warn(`Ignored files: requested ${paths.length} paths, more than the 50 allowed per request`);
-			refuse(`at most 50 paths may be requested at once`);
-			return;
+			const remedy = client?.kind === "cli" ? "--allow-read, or answer yes when kubb studio asks," : "KUBB_AGENT_ALLOW_READ=true";
+			throw new Error(`The agent was not granted permission to read generated files; set ${remedy} to allow it`);
 		}
+		if (!Array.isArray(data.paths)) return this.#refuse("Ignored files: the message carried no paths", "The request carried no paths");
+		const { paths } = data;
+		if (paths.length > 50) return this.#refuse(`Ignored files: requested ${paths.length} paths, more than the 50 allowed per request`, `At most 50 paths may be requested at once`);
 		const generation = this.#lastGeneration;
-		if (!generation) {
-			await this.#warn("Ignored files: no prior generation to read from");
-			refuse("no prior generation to read from, run a generation first");
-			return;
-		}
+		if (!generation) return this.#refuse("Ignored files: no prior generation to read from", "No prior generation to read from, run a generation first");
 		const requested = paths.filter((path) => generation.paths.has(path));
 		const files = {};
 		await inParallel({
@@ -2611,18 +2533,11 @@ var StudioSession = class {
 				if (content !== null) files[path] = content;
 			}
 		});
-		sendAgentMessage(ws, {
-			type: "agent:files",
-			jobId: data.jobId,
-			payload: {
-				status: "ok",
-				files
-			}
-		});
 		await this.#hooks.callHook("studio:command:end", {
 			command,
 			info: `read ${Object.keys(files).length}/${paths.length} requested file${paths.length === 1 ? "" : "s"}`
 		});
+		return { files };
 	}
 };
 //#endregion
@@ -2659,7 +2574,7 @@ function createClient({ storage, onAuthRequired, ...options }) {
 				...options,
 				signal: controller.signal,
 				onTokenRejected: notifyAuthRequired
-			}).connect()));
+			}).start()));
 		},
 		disconnect() {
 			controller.abort();
@@ -2829,6 +2744,6 @@ async function pollForPairingToken({ studioUrl = agentDefaults.studioUrl, sessio
 	throw new Error("The pairing code expired, pair again");
 }
 //#endregion
-export { InvalidAgentTokenError, PairingCanceledError, createAgent, createClient, createFileStorage, createJob, createJobId, defaultStudioUrl, machineTokenFrom, pollForPairingToken, runConnection, setStorage, startPairing, waitForJob };
+export { InvalidAgentTokenError, PairingCanceledError, connectWebSocketRpc, createAgent, createClient, createFileStorage, createJob, defaultStudioUrl, generationEventTypes, machineTokenFrom, pollForPairingToken, runConnection, setStorage, startPairing, waitForJob };
 
 //# sourceMappingURL=index.js.map