experimental-a2 0.20.0 → 0.22.0

package/dist/ai-client.d.ts

@@ -1,5 +1,5 @@
-import { s as CreateClientOptions, t as A2Client } from "./client-h_DKFAjj.js";
-import { a as AI_STATE_REDUCER, i as AIState, r as AIMessageOf, t as AIEventDefs } from "./ai-Cz3hqVhk.js";
+import { s as CreateClientOptions, t as A2Client } from "./client-CUMbS_38.js";
+import { a as AI_STATE_REDUCER, i as AIState, r as AIMessageOf, t as AIEventDefs } from "./ai-DAu77CDr.js";
 //#region src/ai-client.d.ts
 /**
  * The structural shape of `typeof` an `agent({ name })` server: one whose

package/dist/ai-client.js

@@ -1,5 +1,5 @@
-import { t as createClient } from "./client-BhHLoNKF.js";
-import { n as agentContract } from "./ai-mMiWul1T.js";
+import { t as createClient } from "./client-BQS4i_V9.js";
+import { n as agentContract } from "./ai-EEfKyc-z.js";
 //#region src/ai-client.ts
 /**
 * experimental-a2/ai/client — a client for an `agent()` route that needs no

package/dist/ai-fx.d.ts

@@ -1,8 +1,10 @@
-import { t as Harness } from "./ai-harness-CV6rXbB3.js";
+import { t as Harness } from "./ai-harness-DTjkAr2T.js";
 //#region src/ai-fx.d.ts
 /**
  * The FX harness. Every model call FX makes is one durable step; A2 rebuilds
- * the kernel for each step from its latest checkpoint and the log since.
+ * the kernel for each step from its latest checkpoint and the log since,
+ * unless this process holds the kernel the step before left waiting on its
+ * tool calls.
  */
 declare function fx(): Harness<"fx", never>;
 //#endregion

package/dist/ai-fx.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"ai-fx.d.ts","names":[],"sources":["../src/ai-fx.ts"],"mappings":";;;;;;iBA4CgB,MAAM"}
+{"version":3,"file":"ai-fx.d.ts","names":[],"sources":["../src/ai-fx.ts"],"mappings":";;;;;;;;iBA+CgB,MAAM"}

package/dist/ai-fx.js

@@ -1,5 +1,6 @@
+import { n as HARNESS_PARKING } from "./internal-CvsrprnU.js";
 import { t as boundedCache } from "./bounded-cache-CPHWF-8h.js";
-import { a as gatewaySessionAffinity, d as readHarness, i as gatewayCredential, n as deferred, p as saveHarness, r as defineHarness, s as harnessImages, t as Divergence, u as pending } from "./ai-harness-Clbg5ei7.js";
+import { a as gatewaySessionAffinity, d as readHarness, i as gatewayCredential, n as deferred, p as saveHarness, r as defineHarness, s as harnessImages, t as Divergence, u as pending } from "./ai-harness-dQeJaR1t.js";
 import { t as schemaInput } from "./ai-tool-input-DmSqMUQi.js";
 import { asSchema } from "ai";
 //#region src/ai-fx.ts
@@ -9,12 +10,27 @@ import { asSchema } from "ai";
 */
 /**
 * The FX harness. Every model call FX makes is one durable step; A2 rebuilds
-* the kernel for each step from its latest checkpoint and the log since.
+* the kernel for each step from its latest checkpoint and the log since,
+* unless this process holds the kernel the step before left waiting on its
+* tool calls.
 */
 function fx() {
+	const lots = /* @__PURE__ */ new WeakMap();
+	const lotOf = (scope) => {
+		let lot = lots.get(scope);
+		if (lot === void 0) {
+			lot = parkingLot();
+			lots.set(scope, lot);
+		}
+		return lot;
+	};
 	return defineHarness("fx", {
 		tools: {},
-		step: runStep
+		step: (step) => runStep(step, lotOf(step.scope)),
+		release: ({ scope, sessionId }) => {
+			const entry = lots.get(scope)?.take(sessionId);
+			if (entry !== void 0) closeParked(entry);
+		}
 	});
 }
 const SNAPSHOT_VERSION = 1;
@@ -27,10 +43,52 @@ function fail(state, error) {
 	throw state.failure;
 }
 const diverged = (message) => new Divergence("FX", message);
-async function runStep(step) {
+function closeParked(entry) {
+	entry.fxTurn.cancel();
+	entry.kernel.agent.close().catch(() => {});
+}
+/**
+* One agent's parked kernels: one per session, at most `maxKernels`. Each is
+* closed when its session has nothing left to run in this process (the
+* agent's release), when its turn ends, or when another takes its place.
+*/
+function parkingLot() {
+	const held = /* @__PURE__ */ new Map();
+	const take = (sessionId) => {
+		const entry = held.get(sessionId);
+		held.delete(sessionId);
+		return entry;
+	};
+	return {
+		take,
+		park(sessionId, entry) {
+			const replaced = take(sessionId);
+			if (replaced !== void 0) closeParked(replaced);
+			held.set(sessionId, entry);
+			const ended = () => {
+				if (held.get(sessionId) !== entry) return;
+				held.delete(sessionId);
+				closeParked(entry);
+			};
+			entry.fxTurn.result.then(ended, ended);
+			const oldest = held.keys().next().value;
+			if (held.size > HARNESS_PARKING.maxKernels && oldest !== void 0) {
+				const evicted = take(oldest);
+				if (evicted !== void 0) closeParked(evicted);
+			}
+		}
+	};
+}
+async function runStep(step, lot) {
 	const { createFxAgent } = await loadLibfx();
 	const credential = gatewayCredential("fx");
 	const name = `${step.agentName}.a2.harness.fx`;
+	const definitions = await toolDefinitions(step.tools);
+	const settings = JSON.stringify([
+		kernelModel(step),
+		step.instructions,
+		definitions
+	]);
 	const { replayAll, ...first } = await readHarness({
 		step,
 		name,
@@ -41,16 +99,22 @@ async function runStep(step) {
 	let stopped = false;
 	const released = new AbortController();
 	const network = AbortSignal.any([step.signal, released.signal]);
-	let agent;
+	let kernel;
+	/** Whether this step left its kernel for the next one. */
+	let parked = false;
 	let turn;
-	const run = async (controller, turns) => {
+	const run = async (controller, turns, adopted) => {
 		controller.enqueue({
 			type: "start",
 			messageId: step.responseMessageId
 		});
+		if (adopted !== void 0) {
+			await resume(controller, turns.at(-1), adopted);
+			return;
+		}
 		for (const plan of turns) {
-			if (stopped) return;
-			const fxTurn = agent.prompt(plan.prompt, { signal: step.signal });
+			if (stopped || step.signal.aborted) return;
+			const fxTurn = kernel.agent.prompt(plan.prompt);
 			const state = {
 				plan,
 				fxTurn,
@@ -76,29 +140,82 @@ async function runStep(step) {
 				if (result.stopReason !== expected) throw diverged(`a recorded turn ended with ${result.stopReason}, not ${expected}`);
 				continue;
 			}
-			const outcome = await Promise.race([fxTurn.result.then((result) => ({
-				ended: true,
-				result
-			})), state.toolsRequested.promise.then(() => ({ ended: false }))]).catch((error) => {
-				throw state.failure ?? error;
-			});
-			if (state.failure) throw state.failure;
-			await state.parsed;
-			if (outcome.ended) {
-				await drained;
-				if (outcome.result.stopReason === "end_turn" && !step.signal.aborted) await saveCheckpoint();
-			} else fxTurn.cancel();
-			controller.enqueue({ type: "finish-step" });
-			controller.enqueue({
-				type: "finish",
-				...completion.finishReason === void 0 ? {} : { finishReason: completion.finishReason }
-			});
+			await finishLive(controller, state, drained);
 		}
 	};
+	/** Continues a parked kernel: the results it waits on are in the log. */
+	const resume = async (controller, plan, adopted) => {
+		const state = {
+			plan,
+			fxTurn: adopted.fxTurn,
+			next: plan.calls.length,
+			served: plan.calls.at(-1).toolCalls,
+			answered: /* @__PURE__ */ new Set(),
+			live: false,
+			liveStep: false,
+			toolsRequested: deferred()
+		};
+		turn = state;
+		for (const text of plan.steers.get(plan.calls.length - 1) ?? []) await adopted.fxTurn.steer(text);
+		for (const waiting of kernel.waiting.splice(0)) try {
+			const answered = answer(state, waiting.toolName, waiting.input);
+			if (answered === void 0) fail(state, diverged(`no recorded result for ${waiting.toolName}`));
+			waiting.resolve(answered.output);
+		} catch (error) {
+			waiting.reject(error);
+		}
+		await finishLive(controller, state, adopted.drained);
+	};
+	const finishLive = async (controller, state, drained) => {
+		const outcome = await Promise.race([state.fxTurn.result.then((result) => ({
+			ended: true,
+			result
+		})), state.toolsRequested.promise.then(() => ({ ended: false }))]).catch((error) => {
+			throw state.failure ?? error;
+		});
+		if (state.failure) throw state.failure;
+		await state.parsed;
+		if (outcome.ended) {
+			if (outcome.result.stopReason === "cancelled" && !stopped && !step.signal.aborted) throw diverged("the turn was cancelled while this step ran it");
+			await drained;
+			if (outcome.result.stopReason === "end_turn" && !step.signal.aborted) await saveCheckpoint();
+		} else if (kernel.waiting.length > 0 && !stopped && !closed) {
+			lot.park(step.sessionId, {
+				kernel,
+				fxTurn: state.fxTurn,
+				drained,
+				generationId: step.generationId,
+				settings
+			});
+			parked = true;
+		} else state.fxTurn.cancel();
+		controller.enqueue({ type: "finish-step" });
+		controller.enqueue({
+			type: "finish",
+			...completion.finishReason === void 0 ? {} : { finishReason: completion.finishReason }
+		});
+	};
+	/**
+	* Whether a parked kernel stands where this step's log does: made the same
+	* way, its model call the last one the log's live turn committed, and a
+	* recorded result for each tool call it waits on. A steered response
+	* continues the same turn, so its steps qualify too.
+	*/
+	const adoptable = (entry, turns) => {
+		const plan = turns.at(-1);
+		if (entry.settings !== settings || plan?.end !== "live" || plan.generationId !== entry.generationId) return false;
+		const unanswered = [...plan.calls.at(-1)?.toolCalls ?? []];
+		return entry.kernel.waiting.every((waiting) => {
+			const at = unanswered.findIndex((call) => call.toolName === waiting.toolName && sameJson(call.input, waiting.input) && results.has(call.toolCallId));
+			if (at === -1) return false;
+			unanswered.splice(at, 1);
+			return true;
+		});
+	};
 	const saveCheckpoint = async () => {
 		let checkpoint;
 		try {
-			checkpoint = await agent.checkpoint();
+			checkpoint = await kernel.agent.checkpoint();
 		} catch {
 			return;
 		}
@@ -119,7 +236,8 @@ async function runStep(step) {
 		const state = turn;
 		if (state.failure) throw state.failure;
 		const body = decodeBody(init.body);
-		const compaction = isCompaction(body);
+		const request = JSON.parse(body);
+		const compaction = request.toolChoice?.type === "none";
 		const recorded = state.plan.calls[state.next];
 		if (recorded !== void 0) {
 			if (recorded.kind === "compaction" !== compaction) fail(state, diverged(`request ${state.next} is not a ${recorded.kind}`));
@@ -145,7 +263,7 @@ async function runStep(step) {
 		state.live = true;
 		return live(state, url, {
 			...init,
-			body
+			body: withGatewayCaching(body, request)
 		}, compaction);
 	};
 	const affinity = gatewaySessionAffinity({
@@ -158,7 +276,6 @@ async function runStep(step) {
 		headers.set("x-session-affinity", affinity);
 		const response = await fetch(url, {
 			...init,
-			body: withGatewayCaching(init.body),
 			headers,
 			signal: network
 		});
@@ -202,42 +319,67 @@ async function runStep(step) {
 			headers: response.headers
 		});
 	};
-	const hostTools = await kernelTools(step.tools, (toolName, toolInput) => {
-		const state = turn;
-		if (state.failure) throw state.failure;
+	/** The recorded result of a call from the step the kernel was served. */
+	const answer = (state, toolName, toolInput) => {
 		const call = state.served.find((candidate) => candidate.toolName === toolName && !state.answered.has(candidate.toolCallId) && sameJson(candidate.input, toolInput));
 		const result = call === void 0 ? void 0 : results.get(call.toolCallId);
-		if (call !== void 0 && result !== void 0) {
-			state.answered.add(call.toolCallId);
-			if (result.error !== void 0) throw new Error(result.error);
-			if (result.denied === true) throw new Error("The tool call was denied.");
-			return result.output ?? null;
-		}
+		if (call === void 0 || result === void 0) return void 0;
+		state.answered.add(call.toolCallId);
+		if (result.error !== void 0) throw new Error(result.error);
+		if (result.denied === true) throw new Error("The tool call was denied.");
+		return { output: result.output ?? null };
+	};
+	const hostTool = (toolName, toolInput) => {
+		const state = turn;
+		if (state.failure) throw state.failure;
+		const answered = answer(state, toolName, toolInput);
+		if (answered !== void 0) return answered.output;
 		if (state.live) {
 			state.toolsRequested.resolve();
-			return pending(void 0);
+			return new Promise((resolve, reject) => {
+				kernel.waiting.push({
+					toolName,
+					input: toolInput,
+					resolve,
+					reject
+				});
+			});
 		}
 		if (state.plan.end === "interrupted") {
 			state.fxTurn.cancel();
 			return pending(void 0);
 		}
 		return fail(state, diverged(`no recorded result for ${toolName}`));
-	});
-	const plans = (turns) => Promise.all(turns.map((recorded) => turnPlan(recorded, step.signal)));
+	};
+	const hooks = {
+		fetch: hostFetch,
+		tool: hostTool
+	};
+	const planTurns = (turns) => Promise.all(turns.map((recorded) => turnPlan(recorded, step.signal)));
 	let results = read.results;
 	let controller;
 	let closed = false;
 	const emit = (chunk) => {
 		if (!closed) controller.enqueue(chunk);
 	};
-	const create = (checkpoint) => createFxAgent({
-		apiKey: credential.token,
-		model: kernelModel(step),
-		...step.instructions.length === 0 ? {} : { instructions: step.instructions },
-		tools: hostTools,
-		...checkpoint === void 0 ? {} : { checkpoint },
-		fetch: hostFetch
-	});
+	const create = async (checkpoint) => {
+		const created = {
+			hooks,
+			waiting: []
+		};
+		const agent = await createFxAgent({
+			apiKey: credential.token,
+			model: kernelModel(step),
+			...step.instructions.length === 0 ? {} : { instructions: step.instructions },
+			tools: definitions.map((definition) => ({
+				...definition,
+				execute: (toolInput) => created.hooks.tool(definition.name, toolInput)
+			})),
+			...checkpoint === void 0 ? {} : { checkpoint },
+			fetch: (input, init) => created.hooks.fetch(input, init)
+		});
+		return Object.assign(created, { agent });
+	};
 	return {
 		stream: new ReadableStream({
 			start(streamController) {
@@ -249,15 +391,23 @@ async function runStep(step) {
 				step.signal.addEventListener("abort", abort, { once: true });
 				(async () => {
 					try {
-						try {
-							agent = await create(read.saved);
+						let turns = await planTurns(read.turns);
+						const candidate = step.signal.aborted ? void 0 : lot.take(step.sessionId);
+						const adopted = candidate !== void 0 && adoptable(candidate, turns) ? candidate : void 0;
+						if (candidate !== void 0 && adopted === void 0) closeParked(candidate);
+						if (adopted !== void 0) {
+							kernel = adopted.kernel;
+							kernel.hooks = hooks;
+						} else try {
+							kernel = await create(read.saved);
 						} catch (error) {
 							if (read.saved === void 0) throw error;
 							read = await replayAll();
 							results = read.results;
-							agent = await create(void 0);
+							kernel = await create(void 0);
+							turns = await planTurns(read.turns);
 						}
-						await run(streamController, await plans(read.turns));
+						await run(streamController, turns, adopted);
 						closed = true;
 						streamController.close();
 					} catch (error) {
@@ -266,7 +416,7 @@ async function runStep(step) {
 					} finally {
 						released.abort();
 						step.signal.removeEventListener("abort", abort);
-						await agent?.close().catch(() => {});
+						if (!parked) await kernel?.agent.close().catch(() => {});
 					}
 				})();
 			},
@@ -283,20 +433,19 @@ async function runStep(step) {
 async function turnPlan(turn, signal) {
 	const calls = [];
 	const steers = /* @__PURE__ */ new Map();
-	const steer = (at, inputs) => {
-		if (inputs !== void 0) steers.set(Math.max(0, at), inputs.map(steerText));
-	};
-	steer(0, turn.generations.length === 0 ? turn.steers.get(0) : void 0);
 	turn.generations.forEach((generation, index) => {
 		calls.push(...recordedCalls(generation.chunks, generation.completed));
-		steer(calls.length - 1, turn.steers.get(index));
+		const inputs = turn.steers.get(index);
+		if (inputs !== void 0) steers.set(calls.length - 1, inputs.map(steerText));
 	});
 	const images = await harnessImages({
 		harness: "fx()",
 		input: turn.prompt,
 		signal
 	});
+	const generationId = turn.generations.at(-1)?.completed.generationId;
 	return {
+		...generationId === void 0 ? {} : { generationId },
 		prompt: images.length === 0 ? turn.prompt.text : [...images.map((image) => ({
 			type: "image",
 			data: image.data,
@@ -605,8 +754,8 @@ async function loadLibfx() {
 * asks for it; options the request already has take precedence. Every other
 * byte stays as the kernel wrote it, a tool input's spacing included.
 */
-function withGatewayCaching(body) {
-	const options = JSON.parse(body).providerOptions;
+function withGatewayCaching(body, request) {
+	const options = request.providerOptions;
 	if (options === void 0) return body.replace(/^\{/, "{\"providerOptions\":{\"gateway\":{\"caching\":\"auto\"}},");
 	const own = `"providerOptions":${JSON.stringify(options)}}`;
 	if (!body.endsWith(own)) return body;
@@ -622,6 +771,7 @@ function withGatewayCaching(body) {
 function withAuth(headers, credential) {
 	const result = new Headers(headers);
 	result.set("ai-gateway-auth-method", credential.method);
+	if (result.has("authorization")) result.set("authorization", `Bearer ${credential.token}`);
 	return result;
 }
 function kernelModel(step) {
@@ -634,12 +784,12 @@ function kernelModel(step) {
 		...fast ? { fast } : {}
 	};
 }
-async function kernelTools(tools, execute) {
+/** The agent's tools as the kernel takes them, without their execution. */
+async function toolDefinitions(tools) {
 	return Promise.all(Object.entries(tools).map(async ([toolName, tool]) => ({
 		name: toolName,
 		description: typeof tool.description === "string" ? tool.description : toolName,
-		inputSchema: await asSchema(tool.inputSchema).jsonSchema,
-		execute: (input) => execute(toolName, input)
+		inputSchema: await asSchema(tool.inputSchema).jsonSchema
 	})));
 }
 const catalogCache = boundedCache({
@@ -674,9 +824,6 @@ function savedCheckpoint(state) {
 	if (version !== SNAPSHOT_VERSION || typeof checkpoint !== "string") return void 0;
 	return new Uint8Array(Buffer.from(checkpoint, "base64"));
 }
-function isCompaction(body) {
-	return JSON.parse(body).toolChoice?.type === "none";
-}
 function decodeBody(body) {
 	if (typeof body === "string") return body;
 	if (body instanceof Uint8Array || body instanceof ArrayBuffer) return new TextDecoder().decode(body);